108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
"""
|
||
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()
|