Обновить solution.py

This commit is contained in:
2026-05-28 08:41:55 +00:00
parent 0b195119b3
commit d11951cba3
+75 -39
View File
@@ -1,87 +1,123 @@
# solution.py
import os 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_openai import ChatOpenAI
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage from langchain_core.messages import AIMessage
from langgraph.prebuilt import create_react_agent 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() console = Console()
# Настройка модели LLM (потребуется переменная окружения OPENAI_API_KEY) # --------------------------------------------------------------------------- #
llm = ChatOpenAI(temperature=0.7) # 2. LLM
# --------------------------------------------------------------------------- #
# Простой инструмент – эхо‑ответ llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
from langchain_core.tools import tool
# --------------------------------------------------------------------------- #
# 3. Инструмент
# --------------------------------------------------------------------------- #
@tool @tool
def echo_tool(text: str) -> str: 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() memory = MemorySaver()
agent = create_react_agent( agent = graph.compile(
model=llm, checkpointer=memory,
tools=[echo_tool], interrupt_before=["tools"], # пауза перед каждым вызовом инструмента
prompt="You are a helpful assistant. Use the echo tool when needed.",
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( for chunk in agent.stream(
user_input, user_input,
config=cfg, config=config,
stream_mode=["messages", "updates"], stream_mode=["messages", "updates"],
): ):
state = agent.get_state(cfg) state = agent.get_state(config)
chunk_type, chunk_data = chunk chunk_type, chunk_data = chunk
# 5.1 Потоковый вывод токенов LLM
if chunk_type == "messages": if chunk_type == "messages":
# Печатаем токены ответа LLM
for msg in chunk_data: for msg in chunk_data:
if isinstance(msg, AIMessage): if isinstance(msg, AIMessage):
console.print(msg.content, end="") console.print(msg.content, end="")
# 5.2 Вывод вызовов инструментов
elif chunk_type == "updates": elif chunk_type == "updates":
# Выводим информацию о вызовах инструментов (если они произошли)
for upd in chunk_data: for upd in chunk_data:
console.print(f"[bold cyan]Tool call: {upd}") console.print(f"[bold cyan]Tool call: {upd}")
# Обнаруживаем паузу перед инструментом # 5.3 Обнаружение паузы перед инструментом
if "__interrupt__" in chunk_data and state.next == ("tools",): if chunk_type == "updates" and "__interrupt__" in chunk_data and state.next == ("tools",):
# Получаем детали последнего вызова инструмента
tool_call = state.values["messages"][-1].tool_calls[0] tool_call = state.values["messages"][-1].tool_calls[0]
name, args = tool_call["name"], tool_call["args"] name, args = tool_call["name"], tool_call["args"]
console.print(f"\n\n[bold yellow]Пауза: агент хочет вызвать инструмент {name}({args})") console.print(f"\n\n[bold yellow]Пауза: агент хочет вызвать инструмент {name}({args})")
answer = input("Разрешить? (Y/n): ").strip().lower() answer = input("Разрешить? (Y/n): ").strip().lower()
if answer in ("", "y", "yes"): if answer in ("", "y", "yes"):
# Возобновляем работу с того места, где остановились ask_and_run(None, config) # возобновляем с того же места
ask_and_run(None, cfg)
else: else:
console.print("[bold red]Отменено пользователем.") console.print("[bold red]Отменено пользователем.")
break break
# Основной цикл чата # --------------------------------------------------------------------------- #
# 6. Основной цикл чата
# --------------------------------------------------------------------------- #
if __name__ == "__main__": if __name__ == "__main__":
config = {"configurable": {"thread_id": "chat-1"}}
console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n") console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n")
while True: while True:
user_text = input("\nВы: ") user_text = input("\nВы: ")
if user_text.lower() == "exit": if user_text.lower() == "exit":
console.print("[bold magenta]До свидания!") console.print("[bold magenta]До свидания!")
break break
# Формируем сообщение пользователя
msg = {"messages": [{"role": "human", "content": user_text}]} ask_and_run(
ask_and_run(msg, config) {"messages": [{"role": "human", "content": user_text}]},
config,
)