123 lines
5.5 KiB
Python
123 lines
5.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 import create_react_agent
|
|
from rich.console import Console
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 1. Конфигурация LLM и инструмента
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
# Задаём ключ API OpenAI (можно задать переменной окружения)
|
|
os.environ.setdefault("OPENAI_API_KEY", "sk-...")
|
|
|
|
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
|
|
|
|
# Пример простого инструмента: поиск в интернете
|
|
def search(query: str) -> str:
|
|
"""Простейший инструмент поиска."""
|
|
return f"Результаты поиска по запросу '{query}': ..."
|
|
|
|
search_tool = {
|
|
"name": "search",
|
|
"description": "Используется для поиска информации в интернете.",
|
|
"func": search,
|
|
}
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 2. Создаём агент с памятью и возможностью паузы перед вызовом инструмента
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
memory = MemorySaver()
|
|
|
|
agent = create_react_agent(
|
|
model=llm,
|
|
tools=[search_tool],
|
|
prompt="You are a helpful assistant.",
|
|
checkpointer=memory, # подключаем память
|
|
interrupt_before=["tools"], # ставим паузу перед каждым вызовом инструмента
|
|
)
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 3. Вспомогательные функции
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
console = Console()
|
|
|
|
def ask_and_run(user_input: Dict[str, Any], config: Dict[str, Any]) -> None:
|
|
"""
|
|
Обрабатывает поток от агента, выводит токены и обрабатывает паузы.
|
|
"""
|
|
# Переменная для хранения последнего состояния (для возобновления)
|
|
last_state = None
|
|
|
|
for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
|
|
chunk_type, chunk_data = chunk
|
|
|
|
if chunk_type == "messages":
|
|
# Печатаем токены ответа агента
|
|
console.print(chunk_data["content"], end="")
|
|
|
|
elif chunk_type == "updates":
|
|
# Выводим информацию о вызове инструмента (если есть)
|
|
for update in chunk_data:
|
|
if isinstance(update, dict) and "tool_calls" in update.get("values", {}):
|
|
tool_call = update["values"]["messages"][-1].tool_calls[0]
|
|
console.print(f"\n[bold cyan]Инструмент:[/bold cyan] {tool_call['name']}({tool_call['args']})")
|
|
|
|
# Проверяем, пришёл ли сигнал о паузе
|
|
if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
|
|
# Получаем информацию о том, какой инструмент будет вызван
|
|
state = agent.get_state(config)
|
|
tool_call = state.values["messages"][-1].tool_calls[0]
|
|
console.print(f"\n[bold magenta]Пауза:[/bold magenta] Агент хочет вызвать инструмент "
|
|
f"[green]{tool_call['name']}[/green] с аргументами {tool_call['args']}")
|
|
# Спрашиваем подтверждение у пользователя
|
|
answer = input("Разрешить? (Y/n): ").strip().lower()
|
|
if answer in ("", "y", "yes"):
|
|
console.print("[bold green]Продолжаем...[/bold green]")
|
|
# Возобновляем работу агента с того места, где остановились
|
|
ask_and_run(None, config)
|
|
return
|
|
else:
|
|
console.print("[bold red]Отменено. Завершаем разговор.[/bold red]")
|
|
break
|
|
|
|
last_state = chunk_data
|
|
|
|
console.print("\n[bold blue]Финальный ответ получен.[/bold blue]\n")
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 4. Основной цикл чата
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def main() -> None:
|
|
"""
|
|
Запускает бесконечный чат с агентом.
|
|
"""
|
|
thread_id = "thread-1"
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
|
|
console.print("[bold underline]Запуск агента. Введите 'exit' для завершения.[/bold underline]")
|
|
while True:
|
|
user_input = input("\nВы: ")
|
|
if user_input.lower() == "exit":
|
|
console.print("[bold red]Завершение работы.[/bold red]")
|
|
break
|
|
|
|
# Формируем сообщение в формате, ожидаемом агентом
|
|
message = {"messages": [{"role": "human", "content": user_input}]}
|
|
ask_and_run(message, config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |