Add stream-mode agent with messages and updates handling

This commit is contained in:
2026-05-12 11:04:06 +03:00
commit d7b827593e
3 changed files with 128 additions and 0 deletions
+17
View File
@@ -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
```
+4
View File
@@ -0,0 +1,4 @@
langchain>=1.0.0
langchain-openai>=0.2.0
langgraph>=0.2.0
rich>=13.0.0
+107
View File
@@ -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()