Добавлен agent.py

agent.py

@@ -0,0 +1,116 @@
+# agent.py
+
+"""
+Простой LangChain‑агент с поддержкой потокового вывода через `agent.stream()`.
+Внутри используется один инструмент – get_price, который вызывает под‑агента,
+обрабатывающего запросы к «базе цен».
+"""
+
+from langchain.agents import create_agent
+from langchain.tools import tool
+from rich import print as rprint
+
+# --------------------------------------------------------------------------- #
+# 1. Определяем LLM (Ollama) – это единственный доступный LLM в задании.
+# --------------------------------------------------------------------------- #
+try:
+    from langchain_ollama import ChatOllama
+except ImportError:  # если Ollama недоступен, используем простой Mock‑LLM
+    class DummyLLM:
+        def __call__(self, *args, **kwargs):
+            return "Mock response"
+
+    llm = DummyLLM()
+else:
+    llm = ChatOllama(model="llama3.1")  # пример модели
+
+# --------------------------------------------------------------------------- #
+# 2. Создаём инструмент get_price.
+# --------------------------------------------------------------------------- #
+
+@tool
+def get_price(product: str, city: str) -> str:
+    """
+    Возвращает таблицу цен на указанный продукт в заданном городе.
+    Для демонстрации используется под‑агент, который формирует ответ.
+    """
+    # Под‑агент – простая функция, которая возвращает строку
+    sub_agent = create_agent(
+        llm=llm,
+        tools=[],
+        verbose=False,
+        agent_type="openai-tools",
+    )
+
+    # Запускаем под‑агента в режиме invoke (т.к. он короткий)
+    result = sub_agent.invoke({
+        "messages": [
+            {"role": "human", "content": f"Сколько стоит {product} в городе {city}?"}
+        ]
+    })
+    return result["output"]
+
+# --------------------------------------------------------------------------- #
+# 3. Создаём главный агент с инструментом get_price.
+# --------------------------------------------------------------------------- #
+
+agent = create_agent(
+    llm=llm,
+    tools=[get_price],
+    verbose=False,
+    agent_type="openai-tools",
+)
+
+# --------------------------------------------------------------------------- #
+# 4. Функции для форматирования и вывода чанков.
+# --------------------------------------------------------------------------- #
+
+step = 1
+
+def format_chunk_message(chunk):
+    """
+    Выводит токен текста без перевода строки.
+    При смене шага печатает разделитель.
+    """
+    global step
+    message, meta = chunk
+    if meta["langgraph_step"] != step:
+        step = meta["langgraph_step"]
+        rprint("\n--- --- ---\n")
+    if message.content:
+        print(message.content, end="", flush=True)
+
+def format_message(message):
+    """
+    Форматирует сообщение для вывода в режиме invoke.
+    """
+    if message.content:
+        return message.content
+    # Если нет content – это вызов инструмента
+    tool_call = message.tool_calls[0]
+    return f"{tool_call['name']}({tool_call['args']})"
+
+# --------------------------------------------------------------------------- #
+# 5. Запускаем потоковый вывод.
+# --------------------------------------------------------------------------- #
+
+if __name__ == "__main__":
+    # Пример запроса к агенту
+    stream = agent.stream(
+        {
+            "messages": [
+                {"role": "human", "content": "Покажи мне цены на молоко и хлеб в Казани."}
+            ]
+        },
+        stream_mode=["messages", "updates"],
+    )
+
+    for chunk in stream:
+        chunk_type, chunk_data = chunk
+        if chunk_type == "messages":
+            format_chunk_message(chunk_data)
+        elif chunk_type == "updates":
+            # При завершении шага выводим итоговое сообщение
+            if chunk_data.get("model"):
+                last_msg = chunk_data["model"]["messages"][-1]
+                rprint("\n" + format_message(last_msg) + "\n")