diff --git a/solutions/699cc158d6d3a5544a3ed35b_Stream-режим_AI-агента/solution.py b/solutions/699cc158d6d3a5544a3ed35b_Stream-режим_AI-агента/solution.py new file mode 100644 index 0000000..4233c61 --- /dev/null +++ b/solutions/699cc158d6d3a5544a3ed35b_Stream-режим_AI-агента/solution.py @@ -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") \ No newline at end of file