116 lines
5.1 KiB
Python
116 lines
5.1 KiB
Python
# solution.py
|
|
|
|
import os
|
|
from typing import Any, Dict, Iterable, Tuple, Optional
|
|
|
|
from langchain.chat_models import ChatOpenAI
|
|
from langchain.tools import BaseTool
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langgraph.graph import StateGraph
|
|
from langgraph.prebuilt import create_agent
|
|
from rich.console import Console
|
|
|
|
# ------------------------------------------------------------
|
|
# Конфигурация LLM и инструмента
|
|
# ------------------------------------------------------------
|
|
|
|
# Для работы нужен ключ OpenAI. Если его нет – создаём фиктивный LLM.
|
|
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
|
if OPENAI_API_KEY:
|
|
llm = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
|
|
else:
|
|
# Минимальный заглушка‑LLM (не делает ничего полезного)
|
|
class DummyLLM:
|
|
def __call__(self, *args, **kwargs):
|
|
return {"content": "Dummy response"}
|
|
llm = DummyLLM()
|
|
|
|
# Пример простого инструмента
|
|
class EchoTool(BaseTool):
|
|
name: str = "echo"
|
|
description: str = "Возвращает переданный текст без изменений."
|
|
def _run(self, text: str) -> str:
|
|
return f"Echo: {text}"
|
|
|
|
my_tool = EchoTool()
|
|
|
|
# ------------------------------------------------------------
|
|
# Создание агента с памятью и паузой
|
|
# ------------------------------------------------------------
|
|
|
|
memory = MemorySaver()
|
|
agent = create_agent(
|
|
model=llm,
|
|
tools=[my_tool],
|
|
system_prompt="Ты полезный ассистент. Перед каждым вызовом инструмента спрашивай разрешение.",
|
|
checkpointer=memory, # нужна для сохранения состояния паузы
|
|
interrupt_before=["tools"], # остановка перед инструментом
|
|
)
|
|
|
|
console = Console()
|
|
|
|
# ------------------------------------------------------------
|
|
# Функция взаимодействия с пользователем
|
|
# ------------------------------------------------------------
|
|
|
|
def ask_and_run(user_input: Optional[str], config: Dict[str, Any]) -> None:
|
|
"""
|
|
Запускает потоковый диалог с агентом. Если агент хочет вызвать инструмент,
|
|
пользователь получает запрос на подтверждение.
|
|
"""
|
|
# Если user_input – None, продолжаем с того места, где остановились
|
|
stream = agent.stream(
|
|
{"messages": [] if user_input is None else [{"role": "user", "content": user_input}]},
|
|
config=config,
|
|
stream_mode=["messages", "updates"],
|
|
)
|
|
|
|
for chunk_type, chunk_data in stream:
|
|
# 1. Печать токенов (части сообщения)
|
|
if chunk_type == "messages":
|
|
console.print(chunk_data["content"], end="", style="cyan")
|
|
continue
|
|
|
|
# 2. Обработка вызова инструмента
|
|
if chunk_type == "updates" and "tool_calls" in chunk_data:
|
|
for tool_call in chunk_data["tool_calls"]:
|
|
name = tool_call["name"]
|
|
args = tool_call.get("args", {})
|
|
console.print(f"\n[bold magenta]Tool call:[/bold magenta] {name}({args})")
|
|
continue
|
|
|
|
# 3. Обнаружение паузы перед инструментом
|
|
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]
|
|
name = tool_call["name"]
|
|
args = tool_call.get("args", {})
|
|
console.print(f"\n[bold yellow]Пауза:[/bold yellow] {name}({args})")
|
|
answer = input("\nРазрешить вызов инструмента? (Y/n): ").strip().lower()
|
|
if answer in ("y", ""):
|
|
# Продолжаем с того же состояния
|
|
ask_and_run(None, config)
|
|
else:
|
|
console.print("[red]Вызов отменён.[/red]")
|
|
break
|
|
|
|
# ------------------------------------------------------------
|
|
# Точка входа
|
|
# ------------------------------------------------------------
|
|
|
|
if __name__ == "__main__":
|
|
thread_id = "thread-1"
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
console.print("[bold green]Запуск агента. Введите сообщение, а затем Enter.[/bold green]")
|
|
while True:
|
|
try:
|
|
user_msg = input("\nВы: ")
|
|
if user_msg.lower() in ("exit", "quit"):
|
|
console.print("[bold red]Завершение работы.[/bold red]")
|
|
break
|
|
ask_and_run(user_msg, config)
|
|
except KeyboardInterrupt:
|
|
console.print("\n[red]Прервано пользователем. Выход.[/red]")
|
|
break |