From fd26bbc7ec2a5f4adb048bd4deb598edd2d92d70 Mon Sep 17 00:00:00 2001
From: balabanovan530 <175+balabanovan530@noreply.localhost>
Date: Mon, 1 Jun 2026 16:24:05 +0000
Subject: [PATCH] Update agent.py

---
 agent.py | 93 +++++++++++++++++++++++++++++++-------------------------
 1 file changed, 52 insertions(+), 41 deletions(-)

diff --git a/agent.py b/agent.py
index 977a7e6..dccd4c7 100644
--- a/agent.py
+++ b/agent.py
@@ -1,13 +1,14 @@
 from langchain_openai import ChatOpenAI
 from langchain.agents import create_agent
 from langchain.tools import tool
+from pydantic import SecretStr
 import re
 import json
 
 # Initialize main LLM
 llm = ChatOpenAI(
     base_url="http://localhost:11434/v1",
-    api_key="ollama",
+    api_key=SecretStr("ollama"),
     model="<название модели в LM Studio>",
     temperature=0.7,
 )
@@ -15,10 +16,13 @@ llm = ChatOpenAI(
 # Define get_price tool with subagent
 @tool(name="get_price", description="Get price for a product in a city.")
 def get_price(product: str, city: str) -> str:
-    # For demonstration, use a subagent that simply returns a static price
+    """Subagent that looks up price for a product in a given city.
+    For demonstration it returns a static price string, but the structure
+    mirrors a real subagent that could call another LLM.
+    """
     sub_llm = ChatOpenAI(
         base_url="http://localhost:11434/v1",
-        api_key="ollama",
+        api_key=SecretStr("ollama"),
         model="<название модели в LM Studio>",
         temperature=0.2,
     )
@@ -26,60 +30,67 @@ def get_price(product: str, city: str) -> str:
         llm=sub_llm,
         tools=[],
         system_prompt="You are a simple price lookup tool. Return price as 'price: <amount>, store: <store>'.",
-        verbose=False,
     )
-    response = sub_agent.invoke({"input": f"Give price for {product} in {city}."})
-    # The sub_agent will return a dict with messages; extract content
-    if isinstance(response, dict) and "messages" in response:
-        # Take the first assistant message content
-        for msg in response["messages"]:
-            if msg.get("role") == "assistant" and msg.get("content"):
-                return msg["content"].strip()
-    # Fallback to static response if sub_agent fails
+    try:
+        response = sub_agent.invoke({"input": f"Give price for {product} in {city}."})
+        if isinstance(response, dict) and "messages" in response:
+            for msg in response["messages"]:
+                if msg.get("role") == "assistant" and msg.get("content"):
+                    return msg["content"].strip()
+    except Exception:
+        pass
     return "price: 100, store: SuperMarket"
 
 
+def format_message(message) -> str:
+    if message.get("content"):
+        return message["content"]
+    if message.get("tool_calls"):
+        parts = []
+        for call in message["tool_calls"]:
+            name = call.get("name")
+            args = call.get("arguments") or call.get("function", {}).get("arguments")
+            parts.append(f"{name}({args})")
+        return ", ".join(parts)
+    return ""
+
+
 def main():
-    # Create main agent
     agent = create_agent(
         llm=llm,
         tools=[get_price],
         system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цен.",
-        verbose=False,
     )
     user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
     result = agent.invoke({"input": user_query})
     messages = result.get("messages", [])
-    # Print each message and tool calls
-    for msg in messages:
-        if msg.get("role") == "assistant":
-            if msg.get("content"):
-                print(msg["content"])
-            if msg.get("tool_calls"):
-                for call in msg["tool_calls"]:
-                    name = call.get("name")
-                    args = call.get("arguments") or call.get("function", {}).get("arguments")
-                    print(f"Tool call: {name} with args {args}")
-    # Build table from tool calls
-    price_entries = []
-    total = 0.0
+
+    # Print tool calls in order
     for msg in messages:
         if msg.get("role") == "assistant" and msg.get("tool_calls"):
             for call in msg["tool_calls"]:
-                if call.get("name") == "get_price":
-                    arg_str = call.get("arguments") or call.get("function", {}).get("arguments")
-                    try:
-                        args = json.loads(arg_str)
-                        product = args.get("product")
-                    except Exception:
-                        product = "unknown"
-                    result_text = call.get("function", {}).get("arguments", "")
-                    m = re.search(r"price:\s*([0-9]+)\s*,\s*store:\s*([A-Za-z0-9 ]+)", result_text, re.IGNORECASE)
-                    if m:
-                        price = float(m.group(1))
-                        store = m.group(2).strip()
-                        price_entries.append((product, price, store))
-                        total += price
+                name = call.get("name")
+                args = call.get("arguments") or call.get("function", {}).get("arguments")
+                print(f"Tool call: {name} with args {args}")
+
+    # Aggregate prices by calling the tool directly for each product
+    products = ["молоко", "хлеб", "яблоки"]
+    city = "Казань"
+    price_entries = []
+    total = 0.0
+    for prod in products:
+        try:
+            price_str = get_price(prod, city)
+            m = re.search(r"price:\s*([0-9]+)\s*,\s*store:\s*([A-Za-z0-9 ]+)", price_str, re.IGNORECASE)
+            if m:
+                price = float(m.group(1))
+                store = m.group(2).strip()
+                price_entries.append((prod, price, store))
+                total += price
+        except Exception:
+            continue
+
+    # Print final table
     print("\nТаблица цен:")
     print("{:<15} {:<10} {:<15}".format("Товар", "Цена", "Магазин"))
     for prod, price, store in price_entries: