add main.py

main.py

@@ -1,60 +1,54 @@
-"""
-Simple hierarchical LangChain agent for shopping list.
-"""
-
 from langchain_openai import ChatOpenAI
-from pydantic import SecretStr
 from langchain.tools import tool
 from langchain.agents import create_agent, AgentExecutor
 import json
 
-# LLM configuration – replace <model_name> with your LM Studio model name
+# Configure LLM
 llm = ChatOpenAI(
-    model="<model_name>",
+    model="gpt-4o-mini",  # replace with your local model name
     base_url="http://localhost:1234/v1",
-    api_key=SecretStr("fake"),
+    api_key="fake",
     temperature=0.7,
 )
 
-# Sub‑agent that generates a price table for one product in a city
-@tool(name="get_price", description="Return realistic price for a product in a city as a markdown table.")
+@tool("Get price for a product in a city")
 def get_price(product: str, city: str) -> str:
     """
-    Generates a markdown table with columns: Product | Price (руб.) | Store.
-    The sub‑agent uses the same LLM to produce realistic values.
+    Returns a table with product, price and store.
+    The function internally creates a sub-agent that generates realistic prices.
     """
-    # Create a tiny agent that only returns the price table
-    prompt = (
-        f"You are a local market assistant. Provide a markdown table with columns:\n"
-        f"| Продукт | Цена (руб.) | Магазин |\n"
-        f"For product '{product}' in city '{city}'. Use realistic Russian prices and store names.")
-    sub_agent = create_agent(
-        model=llm,
-        tools=[],
-        system_prompt=prompt,
+    # Sub‑agent to generate price
+    sub_llm = ChatOpenAI(
+        model="gpt-4o-mini",
+        base_url="http://localhost:1234/v1",
+        api_key="fake",
+        temperature=0.5,
     )
-    result = sub_agent.invoke({"messages": [{"role": "human", "content": "Generate table"}]})
-    return result["messages"][-1]["content"]
+    sub_agent = create_agent(
+        model=sub_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.",
+    )
+    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)
 
 # Main agent with get_price tool
 main_agent = create_agent(
     model=llm,
     tools=[get_price],
-    system_prompt="Ты помощник по планированию покупок.",
+    system_prompt="You are a shopping assistant. Use the get_price tool to help users plan their purchases.",
 )
 
 if __name__ == "__main__":
-    user_query = (
-        "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
-    )
+    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
     response = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
     # Print all messages
     for msg in response["messages"]:
-        if msg.get("content"):
+        if "content" in msg:
             print(msg["content"])
-        elif msg.get("tool_calls"):
-            call = msg["tool_calls"][0]
+        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"]
-    print("\n---\n", final)
+    final = response["messages"][-1]["content"] if "content" in response["messages"][-1] else ""
+    print("\nFinal answer:\n", final)