SupplyAgent/find_agent_v2/runtime.py

"""Thin ReAct host used by each deterministic workflow node."""

from __future__ import annotations

from collections.abc import Iterable
from typing import Any

from find_agent_v2.observability import InputSlot, ObagentObserver
from find_agent_v2.state import NodeRun
from find_agent_v2.tools import ToolFn, build_tool_registry
from supply_agent import Agent
from supply_agent.config import Settings


class _ObagentEventCollector:
    """Collect ReAct events in memory for obagent; never writes project log artifacts."""

    def __init__(self) -> None:
        self.events: list[dict[str, Any]] = []

    def start_run(self, *_args, **_kwargs) -> str:
        return ""

    def log_llm_input(self, iteration, model, messages, tools, temperature) -> None:
        self.events.append({
            "type": "llm_input",
            "iteration": iteration,
            "model": model,
            "temperature": temperature,
            "messages": [message.model_dump(mode="json") for message in messages],
            "tools": [item.model_dump(mode="json") for item in (tools or [])],
        })

    def log_llm_output(
        self, iteration, response, raw_response=None, *, model=None, provider="openrouter",
    ) -> None:
        usage = getattr(raw_response, "usage", None)
        if hasattr(usage, "model_dump"):
            usage = usage.model_dump()
        self.events.append({
            "type": "llm_output",
            "iteration": iteration,
            "model": model,
            "provider": provider,
            "response": response.model_dump(mode="json"),
            "usage": usage,
        })

    def log_tool_call(
        self, iteration, name, arguments, result, is_error=False, *, tool_call_id=None,
    ) -> None:
        self.events.append({
            "type": "tool_call",
            "iteration": iteration,
            "name": name,
            "arguments": arguments,
            "result": result,
            "is_error": bool(is_error),
            "tool_call_id": tool_call_id,
        })

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

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


class FindAgentNodeHost:
    """Build a fresh, physically capability-limited Agent for every node."""

    def __init__(
        self,
        *,
        settings: Settings | None = None,
        models_by_role: dict[str, str] | None = None,
        default_model: str = "google/gemini-3-flash-preview",
        observer: ObagentObserver | None = None,
    ) -> None:
        self.settings = settings
        self.models_by_role = dict(models_by_role or {})
        self.default_model = default_model
        self.observer = observer or ObagentObserver()

    async def run_node(
        self,
        *,
        node: str,
        round_index: int,
        system_prompt: str,
        user_content: str,
        tools: Iterable[ToolFn] = (),
        max_iterations: int = 12,
        slots: tuple[InputSlot, ...] = (),
    ) -> NodeRun:
        tool_functions = tuple(tools)
        model = self.models_by_role.get(node, self.default_model)
        event_collector = _ObagentEventCollector()
        agent = Agent(
            settings=self.settings,
            name=f"find_agent_v2.{node}",
            model=model,
            system_prompt=system_prompt,
            tools=build_tool_registry(tool_functions),
            max_iterations=max_iterations,
            temperature=0.2,
            # v2 must not write the project's JSONL/log/OSS visualization artifacts.
            logger=event_collector,
        )
        # Stage capabilities are exact. The generic load_skill tool is not part of this workflow.
        agent.tools.unregister("load_skill")
        try:
            with self.observer.node(node=node) as observation:
                try:
                    actual_user_content = observation.declare(
                        fallback=user_content,
                        system_prompt=system_prompt,
                        slots=slots,
                        tools=tool_functions,
                        model=model,
                    )
                    result = await agent.arun_core(actual_user_content)
                    messages = [message.model_dump(mode="json") for message in result.messages]
                    process = {
                        "messages": messages,
                        "events": event_collector.events,
                        "iterations": result.iterations,
                        "tool_calls_made": result.tool_calls_made,
                    }
                    observation.record_react(output=process, ok=True)
                    observation.set_output(
                        {"agent输出": result.content, **process}, ok=True,
                    )
                    return NodeRun.from_agent_result(node, round_index, result)
                except Exception as exc:
                    error = {"error": f"{type(exc).__name__}: {exc}"}
                    observation.record_react(output=error, ok=False)
                    observation.set_output(error, ok=False)
                    raise
        finally:
            client = getattr(agent.llm, "_async_client", None)
            if client is not None:
                await client.close()


def normalize_models(
    *,
    model: str | None = None,
    planning: str | None = None,
    search: str | None = None,
    evidence: str | None = None,
    evaluation: str | None = None,
    report: str | None = None,
) -> dict[str, str]:
    base = model or "google/gemini-3-flash-preview"
    return {
        "planner": planning or base,
        "search": search or base,
        "evidence": evidence or base,
        "evaluator": evaluation or base,
        "report": report or base,
    }