SupplyAgent/find_agent_v2/context.py

"""Stage-specific input contracts backed only by ``find_agent_v2_*`` tables."""

from __future__ import annotations

from pydantic import BaseModel

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:
    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:
    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_assignment(assignment: BaseModel) -> str:
    """Serialize the exact object sent to a model and shown in observability."""
    return assignment.model_dump_json(exclude_none=True)


def assignment_slots(assignment: BaseModel, *, source: str) -> tuple[InputSlot, ...]:
    """Use one canonical payload so observation can never change model semantics."""
    return (
        InputSlot(
            "阶段任务",
            render_assignment(assignment),
            "stage_assignment",
            source,
            False,
        ),
    )