Добавлен solution.py

solution.py

@@ -0,0 +1,95 @@
+# solution.py
+
+"""
+Основной скрипт агента с потоковым выводом.
+Используется стек: LangChain, OpenAI и rich.
+"""
+
+from langchain import hub
+from langchain.agents import initialize_agent, AgentType
+from langchain.tools import tool
+from langchain_community.llms.openai import OpenAI
+from rich import print as rprint
+
+# ---------- 1. Подключаем LLM ----------
+llm = OpenAI(temperature=0, model="gpt-4o-mini", streaming=True)
+
+# ---------- 2. Определяем инструмент ----------
+@tool
+def get_price(product: str, city: str) -> str:
+    """
+    Возвращает таблицу цен на указанный продукт в заданном городе.
+    """
+    # В реальном приложении здесь будет запрос к базе данных или API.
+    # Для демонстрации возвращаем статический ответ.
+    if product.lower() == "молоко":
+        return "| Продукт | Цена (руб.) | Магазин |\n| Молоко  | 89          | Магнит   |"
+    elif product.lower() == "хлеб":
+        return "| Продукт | Цена (руб.) | Магазин |\n| Хлеб    | 25          | Пятёрочка |"
+    else:
+        return f"Информация о {product} в городе {city} недоступна."
+
+# ---------- 3. Создаём агента ----------
+tools = [get_price]
+agent = initialize_agent(
+    tools,
+    llm,
+    agent=AgentType.OPENAI_FUNCTIONS,
+    verbose=False,          # отключаем стандартный вывод
+)
+
+# ---------- 4. Функции форматирования ----------
+def format_message(message) -> str:
+    """
+    Форматируем сообщение для вывода в консоль.
+    Если есть content – возвращаем его,
+    иначе формируем строку вызова инструмента.
+    """
+    if message.content:
+        return message.content
+    # В случае tool_calls: берём первый вызов
+    call = message.tool_calls[0]
+    name = call["name"]
+    args = call.get("args", {})
+    return f"{name}({args})"
+
+def format_chunk_message(chunk):
+    """
+    Обрабатываем чанк типа 'messages'.
+    Выводим токен без перевода строки, добавляем разделитель при смене шага.
+    """
+    message, meta = chunk
+    global step
+
+    if meta.get("langgraph_step") != step:
+        step = meta["langgraph_step"]
+        rprint("\n --- --- --- \n")
+
+    content = format_message(message)
+    # Выводим без перевода строки, чтобы токены «текли»
+    print(content, end="", flush=True)
+
+# ---------- 5. Запускаем потоковый вывод ----------
+if __name__ == "__main__":
+    # Пример запроса
+    user_prompt = "Сколько стоит молоко в Казани? А как насчёт хлеба?"
+
+    stream = agent.stream(
+        {"messages": [{"role": "human", "content": user_prompt}]},
+        stream_mode=["messages", "updates"],
+    )
+
+    step = 1
+
+    for chunk in stream:
+        chunk_type, chunk_data = chunk
+        if chunk_type == "messages":
+            format_chunk_message(chunk_data)
+        elif chunk_type == "updates" and chunk_data.get("model"):
+            # После завершения шага выводим итоговое сообщение модели
+            last_msg = chunk_data["model"]["messages"][-1]
+            rprint("\n")
+            print(format_message(last_msg), flush=True)
+
+    # Завершающий перевод строки, чтобы консоль оставалась аккуратной
+    print()