fix: main.py — Практическое задание №3: Память и подтверждение действий

main.py

@@ -1,18 +1,15 @@
 import os
+import uuid
 import asyncio
-from rich.console import Console
-from rich.text import Text
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage
 from langchain.tools import tool
+from langchain_core.messages import HumanMessage
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 from langgraph.checkpoint.memory import MemorySaver
+from rich.console import Console
 
-# ---------- Консоль ----------
-console = Console()
-
-# ---------- LLM ----------
+# Конфигурация LLM через OpenRouter
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -20,106 +17,82 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- Инструменты ----------
-@tool
-def get_price(city: str, date: str) -> str:
-    """
-    Возвращает условную цену погоды для указанного города и даты.
-    Параметры:
-        city: название города
-        date: дата в произвольном формате (например, "сегодня", "завтра")
-    """
-    # Здесь могла бы быть реальная логика обращения к API погоды.
-    return f"Цена погоды в {city} на {date}: 42 условных единиц."
-
-# ---------- Бэкенд ----------
-backend = CompositeBackend(
-    [
+# Backend для выполнения инструментов (необязательно, но удобно)
+backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
-    ]
-)
+])
 
-# ---------- Память ----------
+# Пример простого инструмента
+@tool
+def get_price(params: dict) -> str:
+    """
+    Получить цену товара в указанном городе и дате.
+    """
+    city = params.get("city", "неизвестный город")
+    date = params.get("date", "неизвестная дата")
+    return f"Цена в {city} на {date} составляет 100$"
+
+# Создание агента с памятью и паузой перед инструментом
 memory = MemorySaver()
-
-# ---------- Агент ----------
 agent = create_deep_agent(
     model=llm,
     tools=[get_price],
     backend=backend,
-    system_prompt="You are a helpful assistant. Use the provided tools when needed.",
+    system_prompt="You are a helpful agent that asks for confirmation before calling tools.",
     checkpointer=memory,
-    interrupt_before=["tools"],  # пауза перед каждым вызовом инструмента
+    interrupt_before=["tools"],
 )
 
-# ---------- Функция обработки одного запроса ----------
+console = Console()
+
 async def ask_and_run(user_input, config):
     """
-    user_input: dict вида {"messages": [{"role": "human", "content": "..."}]}
-                или None (для возобновления после подтверждения)
-    config: dict с ключом "configurable" -> {"thread_id": "..."}
+    Запускает потоковое выполнение агента, обрабатывает паузы и подтверждения.
     """
-    # Если передан None, просто продолжаем без новых сообщений
-    input_payload = user_input if user_input is not None else {}
-
-    async for chunk in agent.astream(
-        input_payload,
-        config,
-        stream_mode=["messages", "updates"],
-    ):
-        # Получаем актуальное состояние после каждого чанка
-        state = agent.get_state(config)
-
+    async for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
         chunk_type, chunk_data = chunk
 
-        # 1. Вывод сообщений (поток токенов)
+        # Потоковый вывод токенов
         if chunk_type == "messages":
-            # chunk_data содержит список сообщений, но мы выводим только новые токены
-            # Для простоты выводим полное содержание последнего сообщения, если оно от ассистента
-            last_msg = chunk_data[-1]
-            if getattr(last_msg, "role", None) == "assistant":
-                console.print(Text(last_msg.content, style="green"))
+            console.print(chunk_data, end="", style="cyan")
+            continue
 
-        # 2. Вывод обновлений (инструменты)
+        # Вывод вызовов инструментов
         if chunk_type == "updates":
-            # Здесь можно вывести информацию о вызове инструмента, если она есть
-            pass  # в данном простом примере ничего не делаем
+            console.print(chunk_data, style="magenta")
+            continue
 
-        # 3. Обработка прерывания перед инструментом
-        if "__interrupt__" in chunk_data and state.next == ("tools",):
-            # Получаем информацию о запланированном вызове инструмента
-            tool_msg = state.values["messages"][-1]
-            tool_call = tool_msg.tool_calls[0]
+        # Обнаружение паузы перед инструментом
+        if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
+            state = agent.get_state(config)
+            # Последнее сообщение должно содержать вызов инструмента
+            last_msg = state.values["messages"][-1]
+            tool_call = last_msg.tool_calls[0]
             tool_name = tool_call["name"]
             tool_args = tool_call["args"]
-            console.print(f"[bold yellow]Agent wants to call {tool_name}{tool_args}[/]")
-
-            answer = input("Разрешить? (Y/n): ").strip().lower()
-            if answer == "n" or answer == "no":
-                console.print("[red]Отменено[/]")
-                # Прерываем текущий поток, не вызывая инструмент
-                break
-            else:
-                # Возобновляем поток без нового пользовательского сообщения
+            console.print(f"\nАгент хочет вызвать утилиту {tool_name}({tool_args})")
+            answer = input("Разрешить? (Y/n): ")
+            if answer.lower().strip() == "y":
+                # Возобновляем выполнение с того места, где остановились
                 await ask_and_run(None, config)
-                # После возврата продолжаем обработку (обычно поток уже завершён)
-                return
+            else:
+                console.print("Отменено", style="red")
+                break
 
-# ---------- Основной цикл ----------
 async def main():
-    thread_id = "conversation-1"
+    thread_id = f"session-{uuid.uuid4()}"
     config = {"configurable": {"thread_id": thread_id}}
 
-    console.print("[bold cyan]=== Чат с подтверждением инструментов ===[/]")
+    console.print("Введите 'exit' для выхода.", style="bold green")
     while True:
-        user_text = input("\nВы: ").strip()
-        if user_text.lower() in {"exit", "quit"}:
-            console.print("[bold magenta]Завершение работы.[/]")
+        user_input = input("\nВы: ")
+        if user_input.lower() == "exit":
             break
-
-        user_message = {"messages": [{"role": "human", "content": user_text}]}
-        await ask_and_run(user_message, config)
+        await ask_and_run(
+            {"messages": [HumanMessage(content=user_input)]},
+            config,
+        )
 
 if __name__ == "__main__":
     asyncio.run(main())