Add main.py

main.py

@@ -1,54 +1,52 @@
 from langchain_openai import ChatOpenAI
-from langchain.tools import tool
-from langchain.agents import create_agent, AgentExecutor
+from langchain.tools import tool, BaseTool
+from langchain.agents import create_agent
+from pydantic import SecretStr
 import json
 
-# Configure LLM
+# Connect to local LLM via OpenAI-compatible API
 llm = ChatOpenAI(
-    model="gpt-4o-mini",  # replace with your local model name
+    model="gpt-4o-mini",  # replace with your LM Studio model name
     base_url="http://localhost:1234/v1",
-    api_key="fake",
+    api_key=SecretStr("fake"),
     temperature=0.7,
 )
 
-@tool("Get price for a product in a city")
+# Sub-agent that generates a price table for a product in a city
+@tool
 def get_price(product: str, city: str) -> str:
-    """
-    Returns a table with product, price and store.
-    The function internally creates a sub-agent that generates realistic prices.
-    """
-    # Sub‑agent to generate price
-    sub_llm = ChatOpenAI(
-        model="gpt-4o-mini",
-        base_url="http://localhost:1234/v1",
-        api_key="fake",
-        temperature=0.5,
-    )
+    """Return a realistic price table for the given product and city."""
+    # Create a sub‑agent with a simple prompt to generate a table
     sub_agent = create_agent(
-        model=sub_llm,
+        model=llm,
         tools=[],
-        system_prompt=f"You are a price estimator for {city}. Provide a realistic price and store name for {product} in a table format.",
+        system_prompt=f"You are a market analyst. Provide a realistic price for {product} in {city}. Return the result as a markdown table with columns: Продукт, Цена (руб.), Магазин.",
     )
-    result = sub_agent.invoke({"messages": [{"role": "human", "content": f"Give me the price of {product} in {city}"}]})
-    return json.dumps(result["messages"][-1]["content"], ensure_ascii=False)
+    response = sub_agent.invoke({"messages": [{"role": "human", "content": f"Generate price for {product} in {city}"}]})
+    # The agent returns a dict with messages; the last message contains the table
+    return response["messages"][-1]["content"]
 
-# Main agent with get_price tool
+# Main agent that uses get_price to build shopping list
 main_agent = create_agent(
     model=llm,
     tools=[get_price],
-    system_prompt="You are a shopping assistant. Use the get_price tool to help users plan their purchases.",
+    system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цены каждого продукта.",
 )
 
-if __name__ == "__main__":
-    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
-    response = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
-    # Print all messages
-    for msg in response["messages"]:
-        if "content" in msg:
-            print(msg["content"])
-        elif "tool_calls" in msg:
-            for call in msg["tool_calls"]:
-                print(f"{call['name']}({json.dumps(call['args'])})")
-    # Final answer
-    final = response["messages"][-1]["content"] if "content" in response["messages"][-1] else ""
-    print("\nFinal answer:\n", final)
+# Example query
+query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
+result = main_agent.invoke({"messages": [{"role": "human", "content": query}]})
+
+# Pretty‑print all messages
+for msg in result["messages"]:
+    if "content" in msg and msg["content"]:
+        print(msg["content"])
+    elif "tool_calls" in msg and msg["tool_calls"]:
+        for call in msg["tool_calls"]:
+            name = call["name"]
+            args = json.dumps(call["args"], ensure_ascii=False)
+            print(f"{name}({args})")
+    else:
+        print(msg)
+
+print("\n--- End of conversation ---")