Добавлен agent.py

agent.py

@@ -1,116 +1,80 @@
 # agent.py
 
 """
-Простой LangChain‑агент с поддержкой потокового вывода через `agent.stream()`.
-Внутри используется один инструмент – get_price, который вызывает под‑агента,
-обрабатывающего запросы к «базе цен».
+Простой пример агента на LangChain с поддержкой потокового вывода.
+Используем только указанный стек:
+- create_agent из langchain.agents
+- @tool из langchain.tools
+- rich.print как rprint
 """
 
 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.
-# --------------------------------------------------------------------------- #
+# ------------------------------
+# Определяем инструменты (tools)
+# ------------------------------
 
 @tool
-def get_price(product: str, city: str) -> str:
+def echo(text: str) -> str:
     """
-    Возвращает таблицу цен на указанный продукт в заданном городе.
-    Для демонстрации используется под‑агент, который формирует ответ.
+    Простейший инструмент, который просто возвращает переданный текст.
     """
-    # Под‑агент – простая функция, которая возвращает строку
-    sub_agent = create_agent(
-        llm=llm,
-        tools=[],
-        verbose=False,
-        agent_type="openai-tools",
-    )
+    return f"Echo: {text}"
 
-    # Запускаем под‑агента в режиме invoke (т.к. он короткий)
-    result = sub_agent.invoke({
-        "messages": [
-            {"role": "human", "content": f"Сколько стоит {product} в городе {city}?"}
-        ]
-    })
-    return result["output"]
 
-# --------------------------------------------------------------------------- #
-# 3. Создаём главный агент с инструментом get_price.
-# --------------------------------------------------------------------------- #
+@tool
+def add(a: int, b: int) -> int:
+    """
+    Возвращает сумму двух целых чисел.
+    """
+    return a + b
+
+
+# ------------------------------
+# Создаём агента с инструментами
+# ------------------------------
 
 agent = create_agent(
-    llm=llm,
-    tools=[get_price],
-    verbose=False,
-    agent_type="openai-tools",
+    tools=[echo, add],
+    # Можно задать любой шаблон диалога; здесь используем простой
+    system_message="Ты помощник, который может отвечать и выполнять инструменты.",
 )
 
-# --------------------------------------------------------------------------- #
-# 4. Функции для форматирования и вывода чанков.
-# --------------------------------------------------------------------------- #
+# ------------------------------------
+# Функция запуска агента в stream‑режиме
+# ------------------------------------
 
-step = 1
-
-def format_chunk_message(chunk):
+def run_agent(prompt: str):
     """
-    Выводит токен текста без перевода строки.
-    При смене шага печатает разделитель.
+    Запускает агента с заданным запросом и выводит результат токен за токеном.
     """
-    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": "Покажи мне цены на молоко и хлеб в Казани."}
-            ]
-        },
+        {"messages": [{"role": "human", "content": prompt}]},
         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")
+    # Перебираем события из потока и печатаем их в реальном времени
+    for event in stream:
+        if isinstance(event, dict):
+            # В режиме 'updates' событие может быть словарём с информацией о шаге
+            rprint(f"[bold cyan]Update:[/bold cyan] {event}")
+        else:
+            # В режиме 'messages' это строка‑токен
+            rprint(event, end="", flush=True)
+
+    rprint("\n[green]--- Конец ответа ---[/green]")
+
+
+# ------------------------------
+# Пример использования (можно удалить)
+# ------------------------------
+
+if __name__ == "__main__":
+    sample_prompt = (
+        "Привет! Скажи мне сумму 7 и 5, а потом эхо 'Hello World'."
+    )
+    run_agent(sample_prompt)