task-69a474cdc46fd26feae69896/main.py

"""Практическое задание №3: память разговора + подтверждение вызовов инструментов."""
from __future__ import annotations

import os
from typing import Any

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from rich.console import Console

load_dotenv()

console = Console()

BROJS_INFERENCE_URL = "https://platform.brojs.ru/jrnl-bh/api/inference/v1"
DEFAULT_MODEL = "openai/gpt-oss-20b:free"
THREAD_ID = "разговор-1"


def _api_key() -> str:
    return (
        os.getenv("OPENAI_API_KEY")
        or os.getenv("JOURNAL_MCP_PAT")
        or os.getenv("JOURNAL_TOKEN")
        or ""
    )


def _base_url() -> str:
    if os.getenv("OPENAI_BASE_URL"):
        return os.environ["OPENAI_BASE_URL"]
    if os.getenv("OPENAI_API_KEY"):
        return os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
    return BROJS_INFERENCE_URL


def _model() -> str:
    return os.getenv("OPENAI_MODEL") or os.getenv("OPENROUTER_MODEL") or DEFAULT_MODEL


def build_llm() -> ChatOpenAI:
    return ChatOpenAI(
        model=_model(),
        base_url=_base_url(),
        api_key=_api_key(),
        temperature=0.0,
    )


@tool
def get_price(city: str, date: str = "сегодня") -> str:
    """Узнать примерную цену/стоимость покупок или погодные условия в городе на дату."""
    return f"{city}, {date}: ориентировочно 150–300 руб. за базовую корзину."


def build_agent(llm: ChatOpenAI | None = None):
    """Агент с памятью и interrupt_before tools — точка входа для автопроверки."""
    memory = MemorySaver()
    return create_agent(
        model=llm or build_llm(),
        tools=[get_price],
        system_prompt=(
            "Ты помощник по планированию покупок и погоде. "
            "Помни контекст разговора. Для цен вызывай get_price."
        ),
        checkpointer=memory,
        interrupt_before=["tools"],
    )


def _print_pending_tool(agent, config: dict) -> str:
    state = agent.get_state(config)
    messages = state.values.get("messages", [])
    if not messages:
        return "unknown_tool()"
    last = messages[-1]
    tool_calls = getattr(last, "tool_calls", None) or []
    if not tool_calls:
        return "unknown_tool()"
    tc = tool_calls[0]
    name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", "?")
    args = tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {})
    repr_str = f"{name}({args})"
    console.print("\n --- --- --- ")
    console.print(repr_str)
    console.print(f"Агент хочет вызвать утилиту {repr_str}")
    return repr_str


def _print_updates(chunk_data: dict) -> None:
    for node_name, update in chunk_data.items():
        if node_name == "__interrupt__":
            continue
        if not isinstance(update, dict):
            continue
        messages = update.get("messages", [])
        for msg in messages:
            tool_calls = getattr(msg, "tool_calls", None) or []
            for tc in tool_calls:
                name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", "?")
                args = tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {})
                console.print(f"\n --- --- --- ")
                console.print(f"{name}({args})")


def ask_and_run(
    agent,
    user_input: dict | None,
    config: dict,
    *,
    auto_approve: bool = False,
) -> None:
    """Запуск или возобновление агента с обработкой паузы перед tools."""
    for chunk in agent.stream(
        user_input,
        config=config,
        stream_mode=["messages", "updates"],
    ):
        state = agent.get_state(config)

        if not isinstance(chunk, tuple) or len(chunk) != 2:
            continue

        chunk_type, chunk_data = chunk

        if chunk_type == "messages":
            if isinstance(chunk_data, tuple) and len(chunk_data) >= 1:
                token = chunk_data[0]
                content = getattr(token, "content", None)
                if content:
                    console.print(content, end="")

        if chunk_type == "updates" and isinstance(chunk_data, dict):
            _print_updates(chunk_data)

            if "__interrupt__" in chunk_data and state.next == ("tools",):
                _print_pending_tool(agent, config)
                if auto_approve:
                    ask_and_run(agent, None, config, auto_approve=True)
                else:
                    answer = input("Разрешить? (Y/n): ").strip().lower()
                    if answer in ("", "y", "yes", "д", "да"):
                        ask_and_run(agent, None, config, auto_approve=auto_approve)
                    else:
                        console.print("Отменено")
                return

    console.print()


def run_chat() -> None:
    agent = build_agent()
    config = {"configurable": {"thread_id": THREAD_ID}}
    console.print(
        "Чат с агентом (память + подтверждение tools). "
        "Введите 'exit' для выхода.\n"
    )

    while True:
        user_text = input("\nВы: ").strip()
        if user_text.lower() == "exit":
            break
        if not user_text:
            continue

        ask_and_run(
            agent,
            {"messages": [{"role": "human", "content": user_text}]},
            config,
        )


def main() -> None:
    run_chat()


if __name__ == "__main__":
    main()