From d7b827593e1b3d82e082c242d1b862857907158c Mon Sep 17 00:00:00 2001
From: marat <maratfazylov19@gmail.com>
Date: Tue, 12 May 2026 11:04:06 +0300
Subject: [PATCH] Add stream-mode agent with messages and updates handling

---
 README.md        |  17 ++++++++
 requirements.txt |   4 ++
 stream_agent.py  | 107 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 128 insertions(+)
 create mode 100644 README.md
 create mode 100644 requirements.txt
 create mode 100644 stream_agent.py

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..d9e37e0
--- /dev/null
+++ b/README.md
@@ -0,0 +1,17 @@
+# Stream-режим AI-агента
+
+Агент с потоковым выводом в консоль.
+
+## Что реализовано
+
+1. **`.stream()` вместо `.invoke()`** — `agent.stream(..., stream_mode=['messages', 'updates'])`
+2. **Обработка `'messages'`** — вывод токенов текста сразу (`end=''`), разделители `--- --- ---` при смене шага
+3. **Обработка `'updates'`** — отображение вызовов инструментов через `format_message()`
+
+## Запуск
+
+```bash
+pip install -r requirements.txt
+export OPENAI_API_KEY="your-key"
+python stream_agent.py
+```
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..5a4d337
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,4 @@
+langchain>=1.0.0
+langchain-openai>=0.2.0
+langgraph>=0.2.0
+rich>=13.0.0
diff --git a/stream_agent.py b/stream_agent.py
new file mode 100644
index 0000000..22d530b
--- /dev/null
+++ b/stream_agent.py
@@ -0,0 +1,107 @@
+"""
+Stream-режим AI-агента
+
+Агент с потоковым выводом: замена .invoke() на .stream()
+с stream_mode=['messages', 'updates'].
+Текст появляется в консоли по мере генерации, с разделителями между шагами.
+"""
+
+import os
+from rich.console import Console
+from langchain_core.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.prebuilt import create_react_agent
+
+console = Console()
+
+
+@tool
+def get_weather(city: str) -> str:
+    """Получить текущую погоду в указанном городе."""
+    weather_data = {
+        "Казань": "🌤 Казань: +18°C, переменная облачность",
+        "Москва": "☀️ Москва: +22°C, ясно",
+        "Санкт-Петербург": "🌧 Санкт-Петербург: +14°C, дождь",
+    }
+    return weather_data.get(city, f"🌡 {city}: данные о погоде недоступны")
+
+
+@tool
+def search_info(query: str) -> str:
+    """Поиск информации по запросу."""
+    return f"Результат поиска по запросу: '{query}' — найдено 5 релевантных источников."
+
+
+def format_message(message):
+    """Форматирование сообщения для вывода."""
+    if hasattr(message, "content") and message.content:
+        return message.content
+    return str(message)
+
+
+def run_stream_agent():
+    """Запуск агента в stream-режиме."""
+
+    llm = ChatOpenAI(
+        model=os.environ.get("OPENAI_MODEL", "gpt-4o-mini"),
+        openai_api_base=os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
+        openai_api_key=os.environ.get("OPENAI_API_KEY", ""),
+        temperature=0,
+    )
+
+    agent = create_react_agent(
+        model=llm,
+        tools=[get_weather, search_info],
+        prompt="Ты полезный ассистент. Отвечай на вопросы пользователя.",
+    )
+
+    user_input = "Какая погода в Казани и Москве? Также найди информацию о LangGraph."
+
+    console.print("=" * 60)
+    console.print(f"  Запрос: {user_input}")
+    console.print("=" * 60)
+    console.print()
+
+    prev_step = None
+
+    # Заменили .invoke() на .stream() с stream_mode=['messages', 'updates']
+    for chunk in agent.stream(
+        {"messages": [{"role": "user", "content": user_input}]},
+        stream_mode=["messages", "updates"],
+    ):
+        chunk_type, chunk_data = chunk
+
+        # Обработка 'messages' — токены текста
+        if chunk_type == "messages":
+            message, meta = chunk_data
+            if hasattr(message, "content") and message.content:
+                # При смене langgraph_step — печатаем разделитель
+                current_step = meta.get("langgraph_step")
+                if prev_step is not None and current_step != prev_step:
+                    console.print()
+                    console.print("--- --- ---")
+                    console.print()
+                prev_step = current_step
+
+                # Выводим контент сразу без перевода строки
+                console.print(message.content, end="", style="white")
+
+        # Обработка 'updates' — события вызова инструментов
+        elif chunk_type == "updates":
+            if "model" in chunk_data:
+                model_output = chunk_data["model"]
+                if "messages" in model_output and model_output["messages"]:
+                    last_msg = model_output["messages"][-1]
+                    formatted = format_message(last_msg)
+                    if formatted:
+                        console.print(f"\n[bold cyan]Шаг модели:[/bold cyan] {formatted}")
+
+    console.print()
+    console.print()
+    console.print("=" * 60)
+    console.print("  Ответ получен (stream-режим)")
+    console.print("=" * 60)
+
+
+if __name__ == "__main__":
+    run_stream_agent()