128 lines
4.5 KiB
Python
128 lines
4.5 KiB
Python
# solution.py
|
|
|
|
import os
|
|
from typing import Any, Dict, Iterable, Tuple
|
|
|
|
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()
|
|
|
|
# ---------- 1. Настройка модели и инструментов ----------
|
|
llm = ChatOpenAI(
|
|
model="gpt-4o-mini",
|
|
temperature=0,
|
|
)
|
|
|
|
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:
|
|
"""
|
|
Запускает потоковый ответ агента с возможностью подтверждения
|
|
вызова инструмента пользователем.
|
|
"""
|
|
# Инициализируем сообщение пользователя в состоянии
|
|
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)
|
|
|
|
# 3.1 Вывод токенов LLM
|
|
if chunk_type == "messages":
|
|
token = chunk_data["content"] or ""
|
|
console.print(token, end="", style="cyan")
|
|
continue
|
|
|
|
# 3.2 Обработка обновлений (вызовы инструментов)
|
|
if chunk_type == "updates":
|
|
for update in chunk_data:
|
|
console.print(f"\n[bold magenta]Обновление:{update}")
|
|
continue
|
|
|
|
# 3.3 Пауза перед вызовом инструмента
|
|
if "__interrupt__" in chunk_data and state.next == ("tools",):
|
|
tool_call = state.values["messages"][-1]["tool_calls"][0]
|
|
name = tool_call["name"]
|
|
args = tool_call.get("args", {})
|
|
console.print(f"\n[bold yellow]Пауза: [{name}]({args})")
|
|
answer = input("Разрешить? (Y/n): ").strip().lower()
|
|
if answer in ("y", ""):
|
|
# Возобновляем с того места, где остановились
|
|
ask_and_run(None, config)
|
|
else:
|
|
console.print("[red]Отменено пользователем.")
|
|
break
|
|
|
|
|
|
# ---------- 4. Основной цикл ----------
|
|
def main() -> None:
|
|
thread_id = "thread-1"
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
|
|
console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n")
|
|
|
|
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() |