Практическое задание №3: Память и подтверждение действий: solution.py

This commit is contained in:
2026-05-27 14:30:21 +00:00
parent 117968930e
commit 3906779000
@@ -1,123 +1,84 @@
# solution.py
"""
Основной скрипт агента с памятью и подтверждением действий.
"""
import os
from typing import Any, Dict, Iterable, Tuple
from typing import Any, Dict, Iterable
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
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
# --------------------------------------------------------------------------- #
# 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
# Настройка модели LLM (потребуется переменная окружения OPENAI_API_KEY)
llm = ChatOpenAI(temperature=0.7)
for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
# Простой инструмент – эхо‑ответ
def echo_tool(text: str) -> str:
"""Возвращает введённый текст."""
return text
# Создаём агент с памятью и возможностью паузы перед вызовом инструмента
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"], # остановка перед каждым вызовом инструмента
)
# Конфигурация разговора (thread_id)
config = {"configurable": {"thread_id": "chat-1"}}
def ask_and_run(user_input: Dict[str, Any] | None, cfg: Dict[str, Any]) -> None:
"""
Отправляем запрос агенту и обрабатываем потоковые ответы.
Если агент останавливается перед вызовом инструмента – спрашиваем подтверждение.
"""
# Запускаем поток
for chunk in agent.stream(
user_input,
config=cfg,
stream_mode=["messages", "updates"],
):
state = agent.get_state(cfg)
chunk_type, chunk_data = chunk
if chunk_type == "messages":
# Печатаем токены ответа агента
console.print(chunk_data["content"], end="")
# Печатаем токены ответа LLM
for msg in chunk_data:
if isinstance(msg, AIMessage):
console.print(msg.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']})")
# Выводим информацию о вызовах инструментов (если они произошли)
for upd in chunk_data:
console.print(f"[bold cyan]Tool call:{upd}")
# Проверяем, пришёл ли сигнал о паузе
if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
# Получаем информацию о том, какой инструмент будет вызван
state = agent.get_state(config)
# Обнаруживаем паузу перед инструментом
if "__interrupt__" in chunk_data and state.next == ("tools",):
# Получаем детали последнего вызова инструмента
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']}")
# Спрашиваем подтверждение у пользователя
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"):
console.print("[bold green]Продолжаем...[/bold green]")
# Возобновляем работу агента с того места, где остановились
ask_and_run(None, config)
return
# Возобновляем работу с того места, где остановились
ask_and_run(None, cfg)
else:
console.print("[bold red]Отменено. Завершаем разговор.[/bold red]")
console.print("[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()
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)