Alexandr
100 lines
3.6 KiB
Python
100 lines
3.6 KiB
Python
"""
|
|
Практическое задание: Stream-режим AI-агента
|
|
Заменить единовременный вызов .invoke() на потоковый вывод через .stream(),
|
|
чтобы ответ появлялся в консоли по мере генерации токен за токеном.
|
|
"""
|
|
|
|
import os
|
|
from typing import Literal
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain.agents import create_agent, AgentExecutor
|
|
from langchain_core.tools import tool
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
|
|
# Настройка API ключа
|
|
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your-api-key-here")
|
|
|
|
# Пример инструмента
|
|
@tool
|
|
def get_weather(city: str) -> str:
|
|
"""Получить погоду в городе"""
|
|
weather_data = {
|
|
"Казань": "В Казани сегодня +15°C, облачно",
|
|
"Москва": "В Москве сегодня +12°C, дождь",
|
|
"Санкт-Петербург": "В Санкт-Петербурге сегодня +10°C, солнечно",
|
|
}
|
|
return weather_data.get(city, f"Погода в {city} неизвестна")
|
|
|
|
|
|
@tool
|
|
def get_time(city: str) -> str:
|
|
"""Получить текущее время в городе"""
|
|
return f"Сейчас в {city} 14:30 (примерное время)"
|
|
|
|
|
|
# Создание агента с MemorySaver
|
|
memory = MemorySaver()
|
|
llm = ChatOpenAI(model="gpt-4o", temperature=0)
|
|
|
|
agent = create_agent(
|
|
model=llm,
|
|
tools=[get_weather, get_time],
|
|
system_prompt='Ты полезный ассистент. Отвечай на вопросы о погоде и времени.',
|
|
checkpointer=memory,
|
|
)
|
|
|
|
# Конфигурация с thread_id
|
|
config = {"configurable": {"thread_id": "test-thread-1"}}
|
|
|
|
def run_stream_agent(question: str):
|
|
"""
|
|
Запускает агента в stream-режиме.
|
|
Обрабатывает чанки типов 'messages' и 'updates'.
|
|
"""
|
|
print(f"\nВы: {question}")
|
|
print("-" * 40)
|
|
|
|
# Потоковый вызов агента
|
|
for chunk_type, chunk_data in agent.stream(
|
|
{"messages": [("user", question)]},
|
|
config,
|
|
stream_mode=['messages', 'updates']
|
|
):
|
|
if chunk_type == 'messages':
|
|
# Потоковый вывод токенов без переноса строки
|
|
if hasattr(chunk_data, 'content'):
|
|
print(chunk_data.content, end="", flush=True)
|
|
elif chunk_type == 'updates':
|
|
# Вывод событий об изменении состояния
|
|
if chunk_data.get('langgraph_step'):
|
|
print("\n" + "-" * 40)
|
|
step = chunk_data['langgraph_step']
|
|
print(f"[Шаг: {step}]")
|
|
# Вывод информации о вызове инструмента
|
|
if 'tool_calls' in chunk_data:
|
|
for call in chunk_data['tool_calls']:
|
|
print(f"\nВызов инструмента: {call.get('name')}")
|
|
print(f"Аргументы: {call.get('args')}")
|
|
|
|
print("\n" + "=" * 40)
|
|
|
|
|
|
def main():
|
|
"""Главная функция для демонстрации stream-режима."""
|
|
print("=" * 40)
|
|
print("Stream-режим AI-агента")
|
|
print("=" * 40)
|
|
|
|
questions = [
|
|
"Какая погода в Казани?",
|
|
"Сколько времени в Москве?",
|
|
"Какая погода будет в Санкт-Петербурге?"
|
|
]
|
|
|
|
for q in questions:
|
|
run_stream_agent(q)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|