add main.py

main.py

@@ -1,52 +1,95 @@
+"""
+Simple hierarchical AI agent for shopping list using LangChain.
+
+The script demonstrates:
+* Connection to a local LLM via OpenAI-compatible API.
+* A tool that internally creates a sub‑agent to generate realistic prices.
+* A top‑level agent that orchestrates the price queries and aggregates results.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Dict, Any
+
 from langchain_openai import ChatOpenAI
-from langchain.tools import tool, BaseTool
+from langchain.tools import tool
 from langchain.agents import create_agent
 from pydantic import SecretStr
-import json
 
-# Connect to local LLM via OpenAI-compatible API
+# ---------------------------------------------------------------------------
+# 1. LLM configuration – adjust model name to your LM Studio instance.
+# ---------------------------------------------------------------------------
 llm = ChatOpenAI(
-    model="gpt-4o-mini",  # replace with your LM Studio model name
+    model="gpt-4o-mini",  # replace with the actual model name in LM Studio
     base_url="http://localhost:1234/v1",
     api_key=SecretStr("fake"),
     temperature=0.7,
 )
 
-# Sub-agent that generates a price table for a product in a city
+# ---------------------------------------------------------------------------
+# 2. Sub‑agent that generates a price table for a single product.
+# ---------------------------------------------------------------------------
 @tool
 def get_price(product: str, city: str) -> str:
-    """Return a realistic price table for the given product and city."""
+    """Return a realistic price table for *product* in *city*.
-    # Create a sub‑agent with a simple prompt to generate a table
+
+    The function internally creates a small LangChain agent that asks the LLM to
+    produce a markdown table with columns Product | Price (руб.) | Store.
+    """
+    # Create a tiny agent that only has the task of generating a price row.
     sub_agent = create_agent(
         model=llm,
         tools=[],
-        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: Продукт, Цена (руб.), Магазин.",
+        system_prompt=(
+            f"You are an assistant that provides realistic prices for products in {city}. "
+            "Respond with a markdown table containing columns: Product, Price (руб.), Store."
+        ),
     )
-    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 that uses get_price to build shopping list
+    # Ask the sub‑agent to produce the price.
-main_agent = create_agent(
+    response = sub_agent.invoke(
+        {
+            "messages": [
+                {"role": "human", "content": f"Provide a price for {product} in {city}."}
+            ]
+        }
+    )
+
+    # Extract the final message content.
+    messages = response.get("messages", [])
+    if not messages:
+        return f"| {product} | N/A | Unknown |
+"
+    last_msg = messages[-1]
+    content = last_msg.get("content") or ""
+    # Ensure the table has a header row.
+    if "|" in content and "Product" not in content:
+        content = f"| Product | Price (руб.) | Store |\n{content}"
+    return content
+
+# ---------------------------------------------------------------------------
+# 3. Main agent that orchestrates price queries for a shopping list.
+# ---------------------------------------------------------------------------
+shopping_agent = create_agent(
     model=llm,
     tools=[get_price],
-    system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цены каждого продукта.",
+    system_prompt="You are an assistant that helps plan a shopping list and calculates total cost.",
 )
 
-# Example query
+def main() -> None:
-query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
+    user_query = (
-result = main_agent.invoke({"messages": [{"role": "human", "content": query}]})
+        "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
+    )
+    result = shopping_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
 
-# Pretty‑print all messages
+    # Pretty‑print all messages.
-for msg in result["messages"]:
+    for msg in result.get("messages", []):
-    if "content" in msg and msg["content"]:
+        if msg.get("content"):
             print(msg["content"])
-    elif "tool_calls" in msg and msg["tool_calls"]:
+        elif msg.get("tool_calls"):
-        for call in msg["tool_calls"]:
+            call = msg["tool_calls"][0]
-            name = call["name"]
+            print(f"{call['name']}({json.dumps(call['args'], ensure_ascii=False)})")
-            args = json.dumps(call["args"], ensure_ascii=False)
-            print(f"{name}({args})")
-    else:
-        print(msg)
 
-print("\n--- End of conversation ---")
+if __name__ == "__main__":
+    main()