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

solutions/69a474cdc46fd26feae69896_Практическое_задание__3__Память_и_подтве/solution.py

@@ -1,102 +1,128 @@
 # solution.py
 
-import sys
+import os
 from typing import Any, Dict, Iterable, Tuple
 
-# 1️⃣ Rich для красивого вывода
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import StateGraph
+from langgraph.prebuilt.tool_executor import ToolExecutor
 from rich.console import Console
 
 console = Console()
 
-# 2️⃣ LangGraph и необходимые компоненты
+# ---------- 1. Настройка модели и инструментов ----------
-from langgraph.checkpoint.memory import MemorySaver
+llm = ChatOpenAI(
-from langgraph.graph import StateGraph
+    model="gpt-4o-mini",
-from langgraph.prebuilt import create_chat_agent
+    temperature=0,
-from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
-
-# 3️⃣ Определяем простой инструмент (пример)
-def echo_tool(args: Dict[str, Any]) -> str:
-    """Простейший инструмент, который возвращает переданный текст."""
-    return f"Эхо: {args.get('text', '')}"
-
-# 4️⃣ Создаём LLM
-llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
-
-# 5️⃣ Настраиваем память и агент с паузой перед инструментом
-memory = MemorySaver()
-
-agent = create_chat_agent(
-    llm=llm,
-    tools=[echo_tool],
-    system_prompt="Ты — полезный ассистент, который сначала спрашивает подтверждение перед вызовом инструмента.",
-    checkpointer=memory,
-    interrupt_before=["tools"],  # пауза перед каждым инструментом
 )
 
-# 6️⃣ Функция для обработки потокового ответа и подтверждения
+def echo_tool(args: Dict[str, Any]) -> str:
+    """Простейший инструмент для демонстрации."""
+    return f"Эхо: {args.get('message', '')}"
+
+tools = [echo_tool]
+tool_executor = ToolExecutor(tools)
+
+
+# ---------- 2. Создание графа агента ----------
+def create_agent() -> StateGraph:
+    """
+    Возвращает готовый граф агента с памятью и возможностью
+    прерывания перед вызовом инструмента.
+    """
+    # Состояние: хранит сообщения
+    def add_message(state: Dict[str, Any], message: Dict[str, str]) -> Dict[str, Any]:
+        state["messages"].append(message)
+        return state
+
+    graph = StateGraph(add_message)
+
+    # Входной узел – пользовательское сообщение
+    graph.set_entry_point("user")
+
+    # Узлы для обработки сообщений и вызова инструментов
+    graph.add_node("assistant", lambda state: {"messages": [state["messages"][-1]]})
+    graph.add_node("tools", tool_executor)
+
+    # Переходы
+    graph.set_conditional_edges(
+        "assistant",
+        lambda _: "tools" if any(m.get("tool_calls") for m in _["messages"]) else "final",
+    )
+    graph.add_edge("tools", "assistant")
+    graph.add_edge("assistant", "final")
+
+    # Конфигурация прерывания
+    graph.set_interrupt_before(["tools"])
+
+    return graph
+
+
+agent = create_agent()
+memory = MemorySaver()
+agent.compile(checkpointer=memory)
+
+
+# ---------- 3. Функция взаимодействия ----------
 def ask_and_run(user_input: Dict[str, Any], config: Dict[str, Any]) -> None:
     """
-    Отправляем запрос агенту, обрабатываем потоковые чанки,
+    Запускает потоковый ответ агента с возможностью подтверждения
-    при необходимости запрашиваем у пользователя подтверждение перед вызовом инструмента.
+    вызова инструмента пользователем.
     """
-    # Запускаем поток
+    # Инициализируем сообщение пользователя в состоянии
-    for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
+    if user_input is not None:
+        agent.send(user_input, config=config)
+
+    for chunk_type, chunk_data in agent.stream(
+        {}, config=config, stream_mode=["messages", "updates"]
+    ):
         state = agent.get_state(config)
-        chunk_type, chunk_data = chunk
 
-        # 6.1️⃣ Печатаем токены LLM в режиме потока
+        # 3.1 Вывод токенов LLM
         if chunk_type == "messages":
-            for msg in chunk_data:
+            token = chunk_data["content"] or ""
-                if isinstance(msg, AIMessage):
+            console.print(token, end="", style="cyan")
-                    console.print(f"[bold cyan]AI:[/bold cyan] {msg.content}", end="")
+            continue
-                elif isinstance(msg, ToolMessage):
-                    # Инструмент уже вызван – выводим результат
-                    console.print(f"\n[green]Tool result:[/green] {msg.content}")
 
-        # 6.2️⃣ Обрабатываем обновления (например, вызовы инструментов)
+        # 3.2 Обработка обновлений (вызовы инструментов)
         if chunk_type == "updates":
             for update in chunk_data:
-                if isinstance(update, dict) and "tool_calls" in update:
+                console.print(f"\n[bold magenta]Обновление:{update}")
-                    console.print("\n[magenta]Инструмент готов к выполнению:[/magenta]")
+            continue
-                    tool_call = update["tool_calls"][0]
-                    name = tool_call["name"]
-                    args = tool_call.get("args", {})
-                    console.print(f"[yellow]{name}({args})[/yellow]")
 
-        # 6.3️⃣ Проверяем паузу перед инструментом
+        # 3.3 Пауза перед вызовом инструмента
         if "__interrupt__" in chunk_data and state.next == ("tools",):
-            # Получаем последний вызов инструмента
+            tool_call = state.values["messages"][-1]["tool_calls"][0]
-            last_msg = state.values["messages"][-1]
-            tool_call = last_msg.tool_calls[0]
             name = tool_call["name"]
             args = tool_call.get("args", {})
-            console.print("\n\n[bold red]Пауза:[/bold red] Агент хочет вызвать инструмент.")
+            console.print(f"\n[bold yellow]Пауза: [{name}]({args})")
-            console.print(f"[yellow]{name}({args})[/yellow]")
+            answer = input("Разрешить? (Y/n): ").strip().lower()
-            answer = input("[green]Разрешить? (Y/n): [/green]").strip().lower()
+            if answer in ("y", ""):
-            if answer in ("", "y", "yes"):
+                # Возобновляем с того места, где остановились
-                # Продолжаем с того же состояния
                 ask_and_run(None, config)
             else:
-                console.print("[red]Отменено. Конец диалога.[/red]")
+                console.print("[red]Отменено пользователем.")
-                sys.exit(0)
+                break
 
-# 7️⃣ Основной цикл чата
+
+# ---------- 4. Основной цикл ----------
 def main() -> None:
     thread_id = "thread-1"
     config = {"configurable": {"thread_id": thread_id}}
 
-    console.print("[bold blue]Запуск агента. Введите 'exit' для завершения.[/bold blue]")
+    console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n")
-    while True:
-        user_input = input("\n[cyan]Вы:[/cyan] ")
-        if user_input.lower() == "exit":
-            console.print("[bold green]До свидания![/bold green]")
-            break
 
+    while True:
+        user_input = input("\nВы: ")
+        if user_input.lower() == "exit":
+            break
         ask_and_run(
             {"messages": [{"role": "human", "content": user_input}]},
             config,
         )
+    console.print("[bold green]До свидания!")
+
 
 if __name__ == "__main__":
     main()