Add stream mode implementation
This commit is contained in:
@@ -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()
|
||||
@@ -0,0 +1,7 @@
|
||||
langchain>=0.1.0
|
||||
langchain-core>=0.1.0
|
||||
langchain-openai>=0.0.5
|
||||
langgraph>=0.0.15
|
||||
langgraph-checkpoint>=0.0.1
|
||||
pydantic>=2.0.0
|
||||
python-dotenv>=1.0.0
|
||||
Reference in New Issue
Block a user