Обновить solution.py

solution.py

@@ -1,87 +1,123 @@
+# solution.py
+
 import os
-from typing import Any, Dict, Iterable
+from typing import Any, Dict, Optional
 
-# Библиотеки для вывода и работы с агентом
-from rich.console import Console
-from langgraph.checkpoint.memory import MemorySaver
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
-from langgraph.prebuilt import create_react_agent
+from langchain_core.messages import AIMessage
+from langchain_core.tools import tool
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import StateGraph, MessagesState, END, START
+from langgraph.prebuilt import ToolNode
+from rich.console import Console
+
+# --------------------------------------------------------------------------- #
+# 1. Консоль для вывода
+# --------------------------------------------------------------------------- #
 
-# Инициализация консоли для красивого вывода
 console = Console()
 
-# Настройка модели LLM (потребуется переменная окружения OPENAI_API_KEY)
-llm = ChatOpenAI(temperature=0.7)
+# --------------------------------------------------------------------------- #
+# 2. LLM
+# --------------------------------------------------------------------------- #
 
-# Простой инструмент – эхо‑ответ
-from langchain_core.tools import tool
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+# --------------------------------------------------------------------------- #
+# 3. Инструмент
+# --------------------------------------------------------------------------- #
 
 @tool
 def echo_tool(text: str) -> str:
-    """Возвращает введённый текст."""
-    return text
+    """Возвращает введённый текст без изменений."""
+    return f"Эхо: {text}"
+
+tools = [echo_tool]
+
+# --------------------------------------------------------------------------- #
+# 4. Граф агента (актуальный подход LangGraph + LangChain v1.0+)
+# --------------------------------------------------------------------------- #
+
+llm_with_tools = llm.bind_tools(tools)
+
+def call_model(state: MessagesState) -> Dict[str, Any]:
+    """Вызывает LLM с историей сообщений."""
+    return {"messages": [llm_with_tools.invoke(state["messages"])]}
+
+def should_continue(state: MessagesState) -> str:
+    """Определяет следующий шаг: вызов инструмента или завершение."""
+    if state["messages"][-1].tool_calls:
+        return "tools"
+    return END
+
+graph = StateGraph(MessagesState)
+graph.add_node("agent", call_model)
+graph.add_node("tools", ToolNode(tools))
+graph.add_edge(START, "agent")
+graph.add_conditional_edges("agent", should_continue)
+graph.add_edge("tools", "agent")
 
-# Создаём агент с памятью и возможностью паузы перед вызовом инструмента
 memory = MemorySaver()
-agent = create_react_agent(
-    model=llm,
-    tools=[echo_tool],
-    prompt="You are a helpful assistant. Use the echo tool when needed.",
-    checkpointer=memory,          # обязательно для памяти и паузы
-    interrupt_before=["tools"],   # остановка перед каждым вызовом инструмента
+agent = graph.compile(
+    checkpointer=memory,
+    interrupt_before=["tools"],  # пауза перед каждым вызовом инструмента
 )
 
-# Конфигурация разговора (thread_id)
-config = {"configurable": {"thread_id": "chat-1"}}
+# --------------------------------------------------------------------------- #
+# 5. Функция взаимодействия
+# --------------------------------------------------------------------------- #
 
-def ask_and_run(user_input: Dict[str, Any] | None, cfg: Dict[str, Any]) -> None:
+def ask_and_run(user_input: Optional[Dict[str, Any]], config: Dict[str, Any]) -> None:
     """
-    Отправляем запрос агенту и обрабатываем потоковые ответы.
-    Если агент останавливается перед вызовом инструмента – спрашиваем подтверждение.
+    Запускает потоковый диалог с агентом.
+    При паузе перед инструментом запрашивает подтверждение пользователя.
     """
-    # Запускаем поток
     for chunk in agent.stream(
         user_input,
-        config=cfg,
+        config=config,
         stream_mode=["messages", "updates"],
     ):
-        state = agent.get_state(cfg)
+        state = agent.get_state(config)
         chunk_type, chunk_data = chunk
 
+        # 5.1 Потоковый вывод токенов LLM
         if chunk_type == "messages":
-            # Печатаем токены ответа LLM
             for msg in chunk_data:
                 if isinstance(msg, AIMessage):
                     console.print(msg.content, end="")
 
+        # 5.2 Вывод вызовов инструментов
         elif chunk_type == "updates":
-            # Выводим информацию о вызовах инструментов (если они произошли)
             for upd in chunk_data:
-                console.print(f"[bold cyan]Tool call:{upd}")
+                console.print(f"[bold cyan]Tool call: {upd}")
 
-        # Обнаруживаем паузу перед инструментом
-        if "__interrupt__" in chunk_data and state.next == ("tools",):
-            # Получаем детали последнего вызова инструмента
+        # 5.3 Обнаружение паузы перед инструментом
+        if chunk_type == "updates" and "__interrupt__" in chunk_data and state.next == ("tools",):
             tool_call = state.values["messages"][-1].tool_calls[0]
             name, args = tool_call["name"], tool_call["args"]
             console.print(f"\n\n[bold yellow]Пауза: агент хочет вызвать инструмент {name}({args})")
             answer = input("Разрешить? (Y/n): ").strip().lower()
             if answer in ("", "y", "yes"):
-                # Возобновляем работу с того места, где остановились
-                ask_and_run(None, cfg)
+                ask_and_run(None, config)  # возобновляем с того же места
             else:
                 console.print("[bold red]Отменено пользователем.")
             break
 
-# Основной цикл чата
+# --------------------------------------------------------------------------- #
+# 6. Основной цикл чата
+# --------------------------------------------------------------------------- #
+
 if __name__ == "__main__":
+    config = {"configurable": {"thread_id": "chat-1"}}
     console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n")
+
     while True:
         user_text = input("\nВы: ")
         if user_text.lower() == "exit":
             console.print("[bold magenta]До свидания!")
             break
-        # Формируем сообщение пользователя
-        msg = {"messages": [{"role": "human", "content": user_text}]}
-        ask_and_run(msg, config)
+
+        ask_and_run(
+            {"messages": [{"role": "human", "content": user_text}]},
+            config,
+        )