Add stream mode implementation

2026-05-13 18:29:46 +03:00
parent db1bc2f862
commit 7a683c75c8
2 changed files with 106 additions and 0 deletions
@@ -0,0 +1,99 @@
+"""
+Практическое задание: 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()