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")