add: main.py — Stream-режим AI-агента

main.py

@@ -6,7 +6,9 @@ from langchain.tools import tool
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-# LLM - OpenRouter
+# -------------------------------------------------
+# LLM configuration (OpenRouter)
+# -------------------------------------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -14,7 +16,9 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# Simple backend - filesystem + local shell
+# -------------------------------------------------
+# Backend for tool execution
+# -------------------------------------------------
 backend = CompositeBackend(
     [
         LocalShellBackend(workspace_dir="./workspace"),
@@ -22,32 +26,52 @@ backend = CompositeBackend(
     ]
 )
 
-# Example tool - echo (replace with real logic if needed)
+# -------------------------------------------------
+# Example tool (can be replaced with any real tool)
+# -------------------------------------------------
 @tool
-def echo(query: str) -> str:
-    """Return the received query unchanged."""
-    return query
+def get_price(product: str, city: str) -> str:
+    """
+    Return a mock price table for the given product in the specified city.
+    """
+    # In a real scenario this could call an external API or run a script.
+    return f"""| Продукт | Цена (руб.) | Город |
+| {product} | 89 | {city} |
+"""
 
+# -------------------------------------------------
 # Create the deep agent
+# -------------------------------------------------
 agent = create_deep_agent(
     model=llm,
-    tools=[echo],
+    tools=[get_price],
     backend=backend,
-    system_prompt="You are a helpful assistant that streams its answer token by token.",
+    system_prompt="You are a helpful assistant that can call tools when needed.",
 )
 
+# -------------------------------------------------
+# Helper functions for streaming output
+# -------------------------------------------------
 def format_message(message) -> str:
-    """Convert a LangChain message to a printable string."""
+    """
+    Convert a LangChain message to a printable string.
+    If the message contains tool calls, format them as a function call.
+    """
     if getattr(message, "content", None):
         return message.content
-    # If the message is a tool call, show the call
+    # Tool call case
     if getattr(message, "tool_calls", None):
-        tc = message.tool_calls[0]
-        return f"{tc['name']}({tc['args']})"
+        tool_call = message.tool_calls[0]
+        name = tool_call["name"]
+        args = tool_call["args"]
+        return f"{name}({args})"
     return ""
 
 def format_chunk_message(chunk):
-    """Print a token chunk, adding a separator when the step changes."""
+    """
+    Print token fragments from 'messages' chunks.
+    Insert a separator when the LangGraph step changes.
+    """
     message, meta = chunk
     global current_step
     step = meta.get("langgraph_step", 0)
@@ -57,28 +81,29 @@ def format_chunk_message(chunk):
     if getattr(message, "content", None):
         print(message.content, end="", flush=True)
 
+# -------------------------------------------------
+# Main async entry point
+# -------------------------------------------------
 async def main():
-    # Prepare the input
-    user_input = "Расскажи, как приготовить борщ, используя инструмент echo для демонстрации."
+    user_query = "Сколько стоит молоко в Казани?"
     stream = agent.stream(
-        {"messages": [HumanMessage(content=user_input)]},
+        {"messages": [HumanMessage(content=user_query)]},
         stream_mode=["messages", "updates"],
     )
 
     global current_step
-    current_step = 0
+    current_step = -1  # initialize step counter
 
-    # Iterate over the stream
     for chunk_type, chunk_data in stream:
         if chunk_type == "messages":
             format_chunk_message(chunk_data)
         elif chunk_type == "updates":
-            # When a model update arrives, print the last model message (tool call or final answer)
+            # When a model update contains a finished message, print it nicely
             model_info = chunk_data.get("model")
             if model_info and "messages" in model_info:
                 last_msg = model_info["messages"][-1]
                 print("\n" + format_message(last_msg))
-    # Ensure the final newline
+    # Ensure final newline
     print()
 
 if __name__ == "__main__":