task3-memory-hitl/agent.py

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

Агент с памятью разговора (MemorySaver + thread_id) и механизмом
подтверждения каждого вызова инструмента (interrupt_before=['tools']).
"""

import os
from typing import Optional
from rich.console import Console
from rich.panel import Panel
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver

console = Console()


# Инструмент для демонстрации
@tool
def get_weather(city: str) -> str:
    """Получить текущую погоду в указанном городе."""
    weather_data = {
        "Казань": "🌤 Казань: +18°C, переменная облачность",
        "Москва": "☀️ Москва: +22°C, ясно",
        "Санкт-Петербург": "🌧 Санкт-Петербург: +14°C, дождь",
    }
    return weather_data.get(city, f"🌡 {city}: данные о погоде недоступны")


@tool
def search_info(query: str) -> str:
    """Поиск информации по запросу."""
    return f"Результат поиска по запросу: '{query}' — найдено 5 релевантных источников."


def create_agent():
    """Создание агента с памятью и подтверждением вызова инструментов."""

    llm = ChatOpenAI(
        model=os.environ.get("OPENAI_MODEL", "gpt-4o-mini"),
        openai_api_base=os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
        openai_api_key=os.environ.get("OPENAI_API_KEY", ""),
        temperature=0,
    )

    checkpointer = MemorySaver()

    agent = create_react_agent(
        model=llm,
        tools=[get_weather, search_info],
        checkpointer=checkpointer,
        interrupt_before=["tools"],  # Пауза перед вызовом инструмента
    )

    return agent


def ask_and_run(agent, user_input: str, config: dict):
    """Запуск агента с обработкой прерываний и стримингом."""

    final_response = ""

    for chunk in agent.stream(
        {"messages": [{"role": "user", "content": user_input}]},
        config,
        stream_mode=["messages", "updates"],
    ):
        chunk_type, chunk_data = chunk

        # Обработка текстовых чанков (токены)
        if chunk_type == "messages":
            message, meta = chunk_data
            if hasattr(message, "content") and message.content:
                console.print(message.content, end="")
                final_response += message.content

        # Обработка событий (вызовы инструментов, прерывания)
        elif chunk_type == "updates":
            if "__interrupt__" in chunk_data:
                intr = chunk_data["__interrupt__"]
                state = agent.get_state(config)

                # Показываем вызов инструмента
                last_msg = state.values["messages"][-1]
                if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
                    tool_call = last_msg.tool_calls[0]
                    console.print()
                    console.print(Panel(
                        f"[bold yellow]Вызов инструмента:[/bold yellow]\n"
                        f"  Инструмент: [cyan]{tool_call['name']}[/cyan]\n"
                        f"  Аргументы:  [green]{tool_call['args']}[/green]",
                        title="⏸ Подтверждение",
                        border_style="yellow",
                    ))

                    # Спрашиваем пользователя
                    answer = input("\nРазрешить вызов? (Y/n): ").strip().lower()

                    if answer in ("", "y", "yes", "да"):
                        console.print("[green]✓ Разрешено, продолжаем...[/green]\n")
                        # Возобновляем без изменений
                        for resume_chunk in agent.stream(None, config, stream_mode=["messages", "updates"]):
                            rt, rd = resume_chunk
                            if rt == "messages":
                                msg, meta = rd
                                if hasattr(msg, "content") and msg.content:
                                    console.print(msg.content, end="")
                                    final_response += msg.content
                            elif rt == "updates" and "__interrupt__" in rd:
                                # Рекурсивная обработка следующего прерывания
                                pass
                    else:
                        console.print("[red]✗ Отменено пользователем[/red]")
                        # Отменяем — отправляем сообщение об отмене
                        for cancel_chunk in agent.stream(
                            {"messages": [{"role": "user", "content": "Отменено пользователем. Прекрати выполнение."}]},
                            config,
                            stream_mode=["messages", "updates"],
                        ):
                            ct, cd = cancel_chunk
                            if ct == "messages":
                                msg, _ = cd
                                if hasattr(msg, "content") and msg.content:
                                    console.print(msg.content, end="")

    console.print()  # Финальный перевод строки
    return final_response


def chat_loop():
    """Интерактивный чат-цикл с агентом."""

    agent = create_agent()
    config = {"configurable": {"thread_id": "разговор-1"}}

    console.print(Panel(
        "[bold blue]Агент с памятью и подтверждением действий[/bold blue]\n"
        "Введите 'выход' или 'exit' для завершения.\n"
        "Агент запоминает контекст и спрашивает подтверждение перед каждым вызовом инструмента.",
        title="🤖 Чат",
        border_style="blue",
    ))

    while True:
        console.print()
        user_input = input("[Вы] > ").strip()

        if user_input.lower() in ("выход", "exit", "quit", "q"):
            console.print("[dim]До свидания![/dim]")
            break

        if not user_input:
            continue

        console.print("[Агент] > ", end="")
        ask_and_run(agent, user_input, config)


if __name__ == "__main__":
    chat_loop()