SupplyAgent/find_agent_v2/observability.py

"""obagent is the only visualization backend for find_agent_v2."""

from __future__ import annotations

import os
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Iterator

OBAGENT_PROJECT = "find_agent_v2"
OBAGENT_AGENT = "find_agent_v2"
OBAGENT_ROUND_ANCHOR = {"in": "run", "on": ["graph"]}
DEFAULT_ENDPOINT = "http://ob.aiddit.com"

MODULE_TITLES = {
    "planner": "搜索规划",
    "search": "候选搜索",
    "evidence": "证据补全",
    "evaluator": "评估与分池",
    "report": "结果报告",
}

GRAPH_SPEC = {
    "nodes": [
        {"key": key, "title": title,
         "module_key": f"{OBAGENT_PROJECT}.{OBAGENT_AGENT}.{key}"}
        for key, title in MODULE_TITLES.items()
        if key != "report"
    ]
}

_configured = False


@dataclass(frozen=True)
class InputSlot:
    title: str
    value: str
    key: str
    source: str
    optional: bool = True


class NullRunHandle:
    run_uid: str | None = None

    def finish(self, final_output: Any = None) -> None:
        del final_output


class NullModuleHandle:
    def declare(self, *, fallback: str, **_kwargs) -> str:
        return fallback

    def record_react(self, **_kwargs) -> None:
        return None

    def set_output(self, _output: Any, *, ok: bool = True) -> None:
        del ok


class NullObserver:
    """Test-only observer. Production construction uses :class:`ObagentObserver`."""

    @contextmanager
    def run(self, **_kwargs) -> Iterator[NullRunHandle]:
        yield NullRunHandle()

    @contextmanager
    def round(self, **_kwargs) -> Iterator[NullModuleHandle]:
        yield NullModuleHandle()

    @contextmanager
    def node(self, **_kwargs) -> Iterator[NullModuleHandle]:
        yield NullModuleHandle()


def console_endpoint() -> str:
    return (os.getenv("OBAGENT_ENDPOINT") or DEFAULT_ENDPOINT).strip().rstrip("/")


def run_url(run_uid: str) -> str:
    return f"{console_endpoint()}/#client_uid={run_uid}"


def ensure_configured() -> None:
    global _configured
    if _configured:
        return
    from obagent_sdk import configure

    kwargs: dict[str, Any] = {
        "endpoint": console_endpoint(),
        "project": OBAGENT_PROJECT,
        "timeout": float(os.getenv("OBAGENT_TIMEOUT", "120")),
        "timeout_per_op": float(os.getenv("OBAGENT_TIMEOUT_PER_OP", "1.0")),
    }
    api_key = os.getenv("OBAGENT_API_KEY", "").strip()
    if api_key:
        kwargs["api_key"] = api_key
    wal_dir = os.getenv("OBAGENT_WAL_DIR", "").strip()
    if wal_dir:
        kwargs["wal_dir"] = wal_dir
    enabled = os.getenv("OBAGENT_ENABLED", "1").strip().lower()
    if enabled in {"0", "false", "no", "off"}:
        kwargs["enabled"] = False
    configure(**kwargs)
    _configured = True


def _blocks(slots: tuple[InputSlot, ...]):
    from obagent_sdk.observe import InputBlock

    return [
        InputBlock(
            slot.title,
            slot.value,
            key=slot.key,
            source=slot.source,
            optional=slot.optional,
        )
        for slot in slots
    ]


class _ModuleHandle:
    def __init__(self, ctx) -> None:
        self.ctx = ctx

    def declare(
        self,
        *,
        fallback: str,
        system_prompt: str,
        slots: tuple[InputSlot, ...],
        tools: tuple,
        model: str,
    ) -> str:
        del fallback
        return self.ctx.declare(
            system_prompt=system_prompt,
            blocks=_blocks(slots),
            tools=list(tools),
            model=model,
        )

    def record_react(self, *, output: dict[str, Any], ok: bool) -> None:
        # One stable code stage per agent module; the payload holds the complete
        # custom ReAct message chain because this project does not use LangChain hooks.
        self.ctx.record_stage(
            "react",
            fn=_react_stage_identity,
            title="ReAct 运行",
            output=output,
            ok=ok,
        )

    def set_output(self, output: Any, *, ok: bool = True) -> None:
        self.ctx.set_output(output, ok=ok)


def _react_stage_identity(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Stable source identity for the custom ReAct runtime stage."""
    return messages


class ObagentObserver:
    @contextmanager
    def run(
        self,
        *,
        run_id: str,
        demand_word: str,
        model: str,
        models_by_role: dict[str, str],
    ):
        ensure_configured()
        from obagent_sdk import observe

        with observe.run(
            agent=OBAGENT_AGENT,
            objective=f"寻找视频 · {demand_word} · run_id={run_id}",
            project=OBAGENT_PROJECT,
            model_name=model,
            payload={"run_id": run_id, "demand_word": demand_word,
                     "models_by_role": models_by_role},
            tags={"engine": "staged-react", "version": "v2"},
            meta={"run_name": f"寻找 Agent v2 · {demand_word}",
                  "run_id": run_id, "models_by_role": models_by_role},
            round_anchor=OBAGENT_ROUND_ANCHOR,
        ) as handle:
            yield handle

    @contextmanager
    def round(self, *, round_index: int):
        from obagent_sdk import observe

        with observe.module(
            "graph",
            kind="workflow",
            title=f"寻找 Agent · 第 {round_index} 轮",
            module_key="graph",
            spec=GRAPH_SPEC,
        ) as ctx:
            from obagent_sdk.observe import InputBlock

            ctx.declare(blocks=[InputBlock(
                "当前轮次", str(round_index), key="round_index",
                source="FindAgentV2.begin_round", optional=False,
            )])
            yield _ModuleHandle(ctx)

    @contextmanager
    def node(self, *, node: str):
        from obagent_sdk import observe

        with observe.module(
            node,
            kind="agent",
            title=MODULE_TITLES[node],
            module_key=node,
        ) as ctx:
            yield _ModuleHandle(ctx)