Add stream-mode agent with messages and updates handling
This commit is contained in:
@@ -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
|
||||
```
|
||||
@@ -0,0 +1,4 @@
|
||||
langchain>=1.0.0
|
||||
langchain-openai>=0.2.0
|
||||
langgraph>=0.2.0
|
||||
rich>=13.0.0
|
||||
+107
@@ -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()
|
||||
Reference in New Issue
Block a user