script-build-and-framework-tmp/visualization/backend/app/runtime_event_index.py

from __future__ import annotations

import json
from collections import defaultdict
from dataclasses import dataclass
from typing import Any, Iterable, Iterator

from .runtime_tool_catalog import MAIN_DECISION_READ_TOOLS


_SUCCESS_STATUSES = {"ok", "success", "succeeded", "completed", "complete"}
_FAILURE_STATUSES = {"error", "failed", "failure"}


@dataclass(frozen=True)
class EventBounds:
    """Sequence bounds used by projectors; both ends are exclusive by default."""

    after_seq: int | None = None
    before_seq: int | None = None


class RuntimeEventIndex:
    """Read-only structural index over ``script_build_event`` projections.

    This class intentionally knows nothing about card wording.  It establishes
    event identity, order, hierarchy, scope and actor once, then exposes strict
    queries used by the three main-Agent decision projectors.
    """

    def __init__(self, events: Iterable[dict[str, Any]]):
        ordered = sorted((dict(item) for item in events), key=event_order)
        self._ordered: tuple[dict[str, Any], ...] = tuple(ordered)
        self._by_id: dict[int, dict[str, Any]] = {}
        self._by_name: dict[str, list[dict[str, Any]]] = defaultdict(list)
        self._by_scope: dict[int, list[dict[str, Any]]] = defaultdict(list)
        self._children: dict[int, list[dict[str, Any]]] = defaultdict(list)
        self._by_actor: dict[str, list[dict[str, Any]]] = defaultdict(list)
        self._by_round: dict[int, list[dict[str, Any]]] = defaultdict(list)

        for event in ordered:
            event_id = integer(event.get("id"))
            if event_id is not None:
                self._by_id[event_id] = event
            self._by_name[str(event.get("event_name") or "")].append(event)
            scope_id = integer(event.get("scope_event_id"))
            if scope_id is not None:
                self._by_scope[scope_id].append(event)
            parent_id = integer(event.get("parent_event_id"))
            if parent_id is not None:
                self._children[parent_id].append(event)
            self._by_actor[event_actor(event)].append(event)
            round_index = integer(event.get("round_index"))
            if round_index is not None:
                self._by_round[round_index].append(event)

        self._main_scope_ids = frozenset(
            event_id
            for event in ordered
            if is_main_scope(event)
            and (event_id := integer(event.get("scope_event_id")) or integer(event.get("id")))
            is not None
        )

    @property
    def ordered(self) -> tuple[dict[str, Any], ...]:
        return self._ordered

    @property
    def main_scope_ids(self) -> frozenset[int]:
        return self._main_scope_ids

    def by_id(self, event_id: int | None) -> dict[str, Any] | None:
        return self._by_id.get(int(event_id)) if event_id is not None else None

    def children_of(self, event_id: int) -> tuple[dict[str, Any], ...]:
        return tuple(self._children.get(int(event_id), ()))

    def events_in_scope(self, scope_id: int) -> tuple[dict[str, Any], ...]:
        return tuple(self._by_scope.get(int(scope_id), ()))

    def events_for_actor(self, actor: str) -> tuple[dict[str, Any], ...]:
        return tuple(self._by_actor.get(str(actor), ()))

    def events_for_round(self, round_index: int) -> tuple[dict[str, Any], ...]:
        return tuple(self._by_round.get(int(round_index), ()))

    def tools_in_scope(
        self, scope_id: int, *, names: set[str] | frozenset[str] | None = None
    ) -> list[dict[str, Any]]:
        return [
            event
            for event in self.events_in_scope(scope_id)
            if str(event.get("event_type") or "") == "tool_call"
            and (names is None or str(event.get("event_name") or "") in names)
        ]

    def named(
        self,
        name: str,
        *,
        event_type: str | None = None,
        successful: bool | None = None,
        bounds: EventBounds | None = None,
    ) -> list[dict[str, Any]]:
        events = self._by_name.get(str(name), ())
        return [
            event
            for event in events
            if (event_type is None or str(event.get("event_type") or "") == event_type)
            and (successful is None or event_succeeded(event) is successful)
            and _within_bounds(event, bounds)
        ]

    def anchors(
        self,
        name: str,
        *,
        successful: bool | None = None,
        bounds: EventBounds | None = None,
    ) -> list[dict[str, Any]]:
        """Return only anchor calls made directly by a real main scope."""

        return [
            event
            for event in self.named(
                name,
                event_type="tool_call",
                successful=successful,
                bounds=bounds,
            )
            if self.is_main_direct_event(event)
        ]

    def is_main_direct_event(self, event: dict[str, Any]) -> bool:
        return (
            str(event.get("agent_role") or "") == "main"
            and integer(event.get("agent_depth")) == 0
            and integer(event.get("scope_event_id")) in self._main_scope_ids
        )

    def main_direct_reads(
        self,
        *,
        scope_id: int | None = None,
        names: set[str] | frozenset[str] | None = None,
        bounds: EventBounds | None = None,
        successful: bool | None = True,
    ) -> list[dict[str, Any]]:
        allowed = names or MAIN_DECISION_READ_TOOLS
        return [
            event
            for event in self._ordered
            if str(event.get("event_type") or "") == "tool_call"
            and str(event.get("event_name") or "") in allowed
            and self.is_main_direct_event(event)
            and (scope_id is None or integer(event.get("scope_event_id")) == scope_id)
            and _within_bounds(event, bounds)
            and (successful is None or event_succeeded(event) is successful)
        ]

    def agent_returns(
        self,
        name: str,
        *,
        round_index: int | None = None,
        before_seq: int | None = None,
    ) -> list[tuple[dict[str, Any], str]]:
        """Return evaluator calls and their honest association strength.

        ``returned-report`` requires an explicit parent main scope.  A matching
        structured round without that hierarchy is only ``runtime-associated``.
        Report text is never inspected to infer its round.
        """

        results: list[tuple[dict[str, Any], str]] = []
        for event in self.named(name, event_type="agent_invoke", successful=True):
            if before_seq is not None and event_sequence(event) >= before_seq:
                continue
            if round_index is not None and integer(event.get("round_index")) != round_index:
                continue
            parent_id = integer(event.get("parent_event_id"))
            relation = (
                "returned-report"
                if parent_id in self._main_scope_ids
                else "runtime-associated"
            )
            results.append((event, relation))
        return results

    def resolved_begin_round(self, event: dict[str, Any]) -> int | None:
        """Resolve the newly opened round, never trusting the pre-call context first."""

        output = event_output(event)
        if isinstance(output, dict):
            round_index = integer(output.get("round_index"))
            if output.get("success") is not False and round_index is not None:
                return round_index

        seq = event_sequence(event)
        scope_id = integer(event.get("scope_event_id"))
        for marker in self._ordered:
            marker_seq = event_sequence(marker)
            if marker_seq <= seq:
                continue
            if str(marker.get("event_type") or "") == "tool_call" and str(
                marker.get("event_name") or ""
            ) == "begin_round":
                break
            if str(marker.get("event_type") or "") != "round_begin":
                continue
            marker_scope = integer(marker.get("scope_event_id"))
            if scope_id is not None and marker_scope not in {None, scope_id}:
                continue
            payload = marker.get("payload")
            round_index = integer(payload.get("round_index")) if isinstance(payload, dict) else None
            return round_index or integer(marker.get("event_name"))
        return None

    def successful_begin_rounds(self) -> dict[int, list[dict[str, Any]]]:
        result: dict[int, list[dict[str, Any]]] = defaultdict(list)
        for event in self.anchors("begin_round", successful=True):
            round_index = self.resolved_begin_round(event)
            if round_index is not None:
                result[round_index].append(event)
        return dict(result)


def is_main_scope(event: dict[str, Any]) -> bool:
    return (
        str(event.get("event_type") or "") == "agent_scope"
        and str(event.get("event_name") or "") == "main"
        and str(event.get("agent_role") or "") == "main"
        and integer(event.get("agent_depth")) == 0
    )


def event_actor(event: dict[str, Any]) -> str:
    role = str(event.get("agent_role") or "")
    name = str(event.get("event_name") or "")
    if role == "main" and integer(event.get("agent_depth")) == 0:
        return "main"
    if role == "script_evaluator" or name == "script_evaluator":
        return "overall-evaluator"
    if role == "script_multipath_evaluator" or name == "script_multipath_evaluator":
        return "multipath-evaluator"
    return role or "unknown"


def event_order(event: dict[str, Any]) -> tuple[int, int]:
    return (event_sequence(event), integer(event.get("id")) or 0)


def event_sequence(event: dict[str, Any]) -> int:
    return integer(event.get("event_seq")) or 0


def event_succeeded(event: dict[str, Any]) -> bool:
    status = str(event.get("status") or "").lower()
    if status in _FAILURE_STATUSES or status == "running":
        return False
    output = event_output(event)
    if isinstance(output, dict):
        if output.get("success") is False:
            return False
        if output.get("error") not in (None, "", False, [], {}):
            return False
    return status in _SUCCESS_STATUSES or bool(event.get("ended_at"))


def event_input(event: dict[str, Any]) -> Any:
    if "inputData" in event:
        return event.get("inputData")
    return _json_value(event.get("input_preview"))


def event_output(event: dict[str, Any]) -> Any:
    for key in ("agentOutputData", "outputData"):
        if key in event and event.get(key) is not None:
            return _unwrap_output(event.get(key))
    return _unwrap_output(_json_value(event.get("output_preview")))


def event_text_output(event: dict[str, Any]) -> str | None:
    output = event_output(event)
    if isinstance(output, dict):
        for key in ("summary", "report", "content", "message"):
            value = output.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()
        return None
    return output.strip() if isinstance(output, str) and output.strip() else None


def detail_ref(event: dict[str, Any]) -> str | None:
    event_id = integer(event.get("id"))
    return f"event:{event_id}" if event_id is not None else None


def event_summary(
    event: dict[str, Any], association: str | None
) -> dict[str, Any]:
    """Return a lightweight shared event reference without full body data."""

    event_id = integer(event.get("id"))
    return {
        "id": f"event:{event_id}" if event_id is not None else "event:unknown",
        "eventId": event_id,
        "eventSeq": event.get("event_seq"),
        "type": event.get("event_type"),
        "name": event.get("event_name"),
        "status": event.get("status"),
        "title": event.get("title"),
        "inputPreview": event.get("input_preview"),
        "outputPreview": event.get("output_preview"),
        "startedAt": event.get("started_at"),
        "endedAt": event.get("ended_at"),
        "association": association,
        "detailRef": f"event:{event_id}" if event_id is not None else None,
    }


def integer(value: Any) -> int | None:
    try:
        return int(value) if value is not None else None
    except (TypeError, ValueError):
        return None


def _within_bounds(event: dict[str, Any], bounds: EventBounds | None) -> bool:
    if bounds is None:
        return True
    seq = event_sequence(event)
    if bounds.after_seq is not None and seq <= bounds.after_seq:
        return False
    if bounds.before_seq is not None and seq >= bounds.before_seq:
        return False
    return True


def _json_value(value: Any) -> Any:
    if not isinstance(value, str):
        return value
    text = value.strip()
    if not text:
        return None
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return value


def _unwrap_output(value: Any) -> Any:
    # Tool bodies currently store many JSON results as a JSON-looking string.
    # Unwrap at most twice so malformed or ordinary prose remains untouched.
    result = value
    for _ in range(2):
        parsed = _json_value(result)
        if parsed is result or parsed == result:
            break
        result = parsed
    return result