task-69a86305c46fd26feae6bcaa/main.py

"""Human-in-the-Loop через HumanInTheLoopMiddleware (LangChain)."""
from __future__ import annotations

import os
from typing import Any

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command
from pydantic import SecretStr

load_dotenv()

WEATHER_DB: dict[str, str] = {
    "казань": "В Казани сегодня около +5°C, облачно, без осадков.",
    "москва": "В Москве сегодня около +3°C, пасмурно, слабый снег.",
    "санкт-петербург": "В Санкт-Петербурге сегодня около +1°C, ветрено, дождь со снегом.",
}


@tool
def get_weather(city: str, date: str = "сегодня") -> str:
    """Возвращает краткий прогноз погоды для города на указанную дату."""
    key = city.strip().lower()
    if key in WEATHER_DB:
        base = WEATHER_DB[key]
        if date and date.lower() != "сегодня":
            return base.replace("сегодня", date)
        return base
    return f"В городе {city} на {date}: переменная облачность, около +10°C."


def build_llm() -> ChatOpenAI:
    return ChatOpenAI(
        model=os.getenv("OPENROUTER_MODEL", "poolside/laguna-m.1:free"),
        base_url=os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"),
        api_key=SecretStr(os.getenv("OPENAI_API_KEY", "")),
        temperature=0.2,
    )


def build_agent():
    memory = MemorySaver()
    return create_agent(
        model=build_llm(),
        tools=[get_weather],
        system_prompt="Ты полезный ассистент. Для вопросов о погоде вызывай get_weather.",
        middleware=[
            HumanInTheLoopMiddleware(
                interrupt_on={
                    "get_weather": {"allowed_decisions": ["approve", "reject"]},
                },
                description_prefix="Подтвердите вызов инструмента",
            ),
        ],
        checkpointer=memory,
    )


def _has_interrupt(result: Any) -> bool:
    if isinstance(result, dict):
        return bool(result.get("__interrupt__"))
    interrupts = getattr(result, "interrupts", None)
    return bool(interrupts)


def _interrupt_payload(result: Any) -> dict[str, Any]:
    if isinstance(result, dict) and result.get("__interrupt__"):
        item = result["__interrupt__"][0]
        value = item.value if hasattr(item, "value") else item
        return dict(value)
    interrupts = getattr(result, "interrupts", None)
    if interrupts:
        item = interrupts[0]
        value = item.value if hasattr(item, "value") else item
        return dict(value)
    return {}


def _action_args(action: dict[str, Any]) -> dict[str, Any]:
    return action.get("args") or action.get("arguments") or {}


def _prompt_decisions(
    action_requests: list[dict[str, Any]],
    review_configs: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    allowed_by_tool = {
        cfg.get("action_name", ""): cfg.get("allowed_decisions", ["approve", "reject"])
        for cfg in review_configs
    }
    decisions: list[dict[str, Any]] = []

    for action in action_requests:
        name = action.get("name", "")
        args = _action_args(action)
        allowed = allowed_by_tool.get(name, ["approve", "reject"])

        print("\n--- Подтверждение ---")
        print(f"Инструмент: {name}")
        print(f"Аргументы: {args}")
        if action.get("description"):
            print(f"Описание: {action['description']}")

        while True:
            choice = input("a = approve, r = reject: ").strip().lower()
            if choice in ("a", "approve") and "approve" in allowed:
                decisions.append({"type": "approve"})
                break
            if choice in ("r", "reject") and "reject" in allowed:
                message = input("Сообщение для агента (причина отказа): ").strip()
                decision: dict[str, Any] = {"type": "reject"}
                if message:
                    decision["message"] = message
                decisions.append(decision)
                break
            print("Неверный ввод. Допустимо: a (approve) или r (reject).")

    return decisions


def run_with_hitl(agent, user_text: str, config: dict) -> str:
    result = agent.invoke(
        {"messages": [{"role": "human", "content": user_text}]},
        config=config,
    )

    while _has_interrupt(result):
        payload = _interrupt_payload(result)
        action_requests = payload.get("action_requests", [])
        review_configs = payload.get("review_configs", [])
        decisions = _prompt_decisions(action_requests, review_configs)
        result = agent.invoke(
            Command(resume={"decisions": decisions}),
            config=config,
        )

    messages = result.get("messages", []) if isinstance(result, dict) else []
    if not messages and hasattr(result, "value"):
        messages = result.value.get("messages", [])

    if not messages:
        return "(агент завершил без текстового ответа)"

    last = messages[-1]
    content = getattr(last, "content", None)
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = [p.get("text", "") for p in content if isinstance(p, dict)]
        return "".join(parts) or str(content)
    return str(content)


def main() -> None:
    agent = build_agent()
    config = {"configurable": {"thread_id": "сессия-1"}}

    print("Human-in-the-Loop (middleware). Введите 'exit' для выхода.\n")

    while True:
        user_text = input("Вы: ").strip()
        if not user_text:
            continue
        if user_text.lower() in {"exit", "quit", "выход"}:
            break

        answer = run_with_hitl(agent, user_text, config)
        print(f"\nАгент: {answer}\n")


if __name__ == "__main__":
    main()