Stream-режим AI-агента: solution.py

2026-05-28 10:02:46 +00:00
parent 1feb50c484
commit 4414aa3a10
1 changed files with 108 additions and 0 deletions
@@ -0,0 +1,108 @@
 # solution.py
 """
 Основной скрипт агента с потоковым выводом.
 Используются только библиотеки, указанные в условии:
    - langchain.agents.create_agent
    - langchain_ollama.ChatOllama
    - langchain.tools.tool
    - rich.print
 Потоковый вывод реализован через agent.stream().
 """
 # ────────────────────── подключение к LLM ──────────────────────
 from langchain_ollama import ChatOllama
 from langchain.agents import create_agent
 from langchain.tools import tool
 from rich import print as rprint
 # Создаём объект модели Ollama (используем локальный сервер)
 llm = ChatOllama(model="llama3.1")  # можно заменить на любой доступный модельный тег
 # ────────────────────── определение инструмента с субагентом ──────────────────────
@tool
 def get_price(product: str, city: str) -> str:
    """
    Симуляция запроса цены продукта в указанном городе.
    В реальном проекте здесь могла бы быть интеграция с API магазина.
    """
    # простая фиктивная логика для демонстрации работы инструмента
    prices = {
        ("молоко", "Казань"): 89,
        ("хлеб", "Казань"): 35,
        ("яблоки", "Казань"): 120,
    }
    price = prices.get((product.lower(), city.lower()), None)
    if price is None:
        return f"Цена на {product} в городе {city} не найдена."
    # форматируем таблицу
    table = (
        "| Продукт | Цена (руб.) | Магазин |\n"
        "|---------|-------------|---------|\n"
        f"| {product.title()}  | {price:<11} | Магнит   |\n"
    )
    return table
 # ────────────────────── создание главного агента ──────────────────────
 agent = create_agent(
    llm=llm,
    tools=[get_price],
    verbose=False,          # отключаем внутренний вывод LangChain
 )
 # ────────────────────── потоковый вывод ──────────────────────
 def format_message(message) -> str:
    """
    Форматируем сообщение для вывода в консоль.
    Если у сообщения есть content – возвращаем его,
    иначе формируем строку вызова инструмента.
    """
    if message.content:
        return message.content
    # если нет content, но есть tool_calls – выводим их как строку
    if message.tool_calls:
        call = message.tool_calls[0]
        name = call["name"]
        args = call.get("args", {})
        return f"{name}({args})"
    return ""
 def format_chunk_message(chunk):
    """
    Обрабатываем чанк типа 'messages'.
    Выводим токен сразу, добавляя разделитель при смене шага.
    """
    message, meta = chunk
    # Проверяем шаг агента
    if meta.get("langgraph_step") != format_chunk_message.step:
        format_chunk_message.step = meta["langgraph_step"]
        rprint("\n --- --- --- \n")
    content = message.content or ""
    if content:
        rprint(content, end="", flush=True)
 # Инициализируем переменную шага
 format_chunk_message.step = 0
 # Запускаем потоковый вывод
 stream = agent.stream(
    {
        "messages": [
            {"role": "human", "content": "Сколько стоит молоко в Казани?"}
        ]
    },
    stream_mode=["messages", "updates"],
 )
 for chunk_type, chunk_data in stream:
    if chunk_type == "messages":
        format_chunk_message(chunk_data)
    elif chunk_type == "updates":
        # Обрабатываем события обновления состояния (например, завершение шага модели)
        model_info = chunk_data.get("model")
        if model_info:
            last_msg = model_info["messages"][-1]
            rprint("\n" + format_message(last_msg) + "\n")
 # Завершаем вывод
 rprint("\n--- Конец ответа ---\n")