SupplyAgent/find_agent_v2/context.py

"""Context builders backed only by ``find_agent_v2_*`` tables."""

from __future__ import annotations

import json
from typing import Any

from find_agent_v2.observability import InputSlot
from find_agent_v2.service import get_find_agent_v2_service
from find_agent_v2.state import DiscoverySnapshot


def load_full_state(run_id: str, *, limit: int = 100) -> dict[str, Any]:
    return get_find_agent_v2_service().get_full_state(run_id, limit=limit)


def snapshot_run(run_id: str) -> DiscoverySnapshot:
    return get_find_agent_v2_service().snapshot(run_id)


def require_running_run(run_id: str) -> dict[str, Any]:
    run = get_find_agent_v2_service().require_run(run_id)
    if str(run.get("status") or "") not in {"running", "finished"}:
        raise ValueError(f"run_id={run_id} 当前状态不可执行: {run.get('status')}")
    return run


def render_node_context(
    *, user_input: str, full_state: dict[str, Any], round_index: int, plan: str = "",
) -> str:
    return (
        f"【当前轮次】\n{round_index}\n\n"
        f"【原始任务】\n{user_input}\n\n"
        f"【find_agent_v2 数据库状态快照】\n"
        f"{json.dumps(full_state, ensure_ascii=False, default=str)}\n\n"
        f"【本轮搜索计划】\n{plan}"
    )


def build_node_slots(
    *, user_input: str, full_state: dict[str, Any], round_index: int, plan: str = "",
) -> tuple[InputSlot, ...]:
    """Stable input structure shared by all node instances and obagent declarations."""
    return (
        InputSlot("当前轮次", str(round_index), "round_index", "FindAgentV2.begin_round", False),
        InputSlot("原始任务", user_input, "task", "find_agent_v2_run.input_json", False),
        InputSlot(
            "数据库状态快照",
            json.dumps(full_state, ensure_ascii=False, default=str),
            "state_snapshot",
            "FindAgentV2Service.get_full_state",
            False,
        ),
        InputSlot("本轮搜索计划", plan, "round_plan", "planner output"),
    )