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