add main.py

main.py

@@ -0,0 +1,115 @@
+"""Stream-режим AI-агента: потоковый вывод токенов и событий updates."""
+from __future__ import annotations
+
+import os
+
+from dotenv import load_dotenv
+from langchain.agents import create_agent
+from langchain.tools import tool
+from langchain_openai import ChatOpenAI
+from pydantic import SecretStr
+
+load_dotenv()
+
+llm = ChatOpenAI(
+    model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"),
+    base_url=os.getenv("OPENAI_BASE_URL", "https://openrouter.ai/api/v1"),
+    api_key=SecretStr(os.getenv("OPENAI_API_KEY", "fake")),
+    temperature=0.7,
+)
+
+step = 1
+
+
+@tool
+def get_price(product: str, city: str) -> str:
+    """Узнать примерную цену продукта в указанном городе. Возвращает строку таблицы."""
+    price_agent = create_agent(
+        model=llm,
+        system_prompt=(
+            "Ты эксперт по розничным ценам в России. "
+            "Ответ — одна строка таблицы: | Продукт | Цена (руб.) | Магазин |"
+        ),
+    )
+    result = price_agent.invoke(
+        {
+            "messages": [
+                {
+                    "role": "human",
+                    "content": (
+                        f"Какая примерная цена на «{product}» в городе {city}? "
+                        "Верни строку | Продукт | Цена (руб.) | Магазин |"
+                    ),
+                }
+            ]
+        }
+    )
+    return result["messages"][-1].content
+
+
+shopping_agent = create_agent(
+    model=llm,
+    tools=[get_price],
+    system_prompt="Ты помощник по планированию покупок",
+)
+
+
+def format_message(message) -> str:
+    """Текст сообщения или вызов инструмента."""
+    if message.content:
+        return str(message.content)
+    tool_calls = getattr(message, "tool_calls", None) or []
+    if tool_calls:
+        tc = tool_calls[0]
+        name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", "?")
+        args = tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {})
+        return f"{name}({args})"
+    return str(message)
+
+
+def format_chunk_message(chunk_data: tuple) -> None:
+    """Потоковый вывод токенов с разделителем при смене шага."""
+    global step
+    message, meta = chunk_data
+    current_step = meta.get("langgraph_step", step)
+
+    if current_step != step:
+        step = current_step
+        print("\n --- --- --- \n")
+
+    if message.content:
+        print(message.content, end="", flush=True)
+
+
+def main() -> None:
+    global step
+    step = 1
+
+    question = (
+        "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
+    )
+
+    stream = shopping_agent.stream(
+        {"messages": [{"role": "human", "content": question}]},
+        stream_mode=["messages", "updates"],
+    )
+
+    for chunk in stream:
+        chunk_type, chunk_data = chunk
+
+        if chunk_type == "messages":
+            format_chunk_message(chunk_data)
+
+        if chunk_type == "updates":
+            model_update = chunk_data.get("model")
+            if model_update:
+                last_message = model_update["messages"][-1]
+                formatted = format_message(last_message)
+                if formatted:
+                    print(formatted)
+
+    print()
+
+
+if __name__ == "__main__":
+    main()