tracer-cloud--opensre/core/events.py

"""Typed event contract for the shared agent runtime."""

# ruff: noqa: UP040

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any, Literal, TypeAlias

# CodeQL's explicit-export query does not recognize Python 3.12 ``type``
# statements in __all__, so keep these exported aliases as TypeAlias assignments.
RuntimeEventType: TypeAlias = Literal[
    "agent_start",
    "turn_start",
    "message_start",
    "message_update",
    "provider_request_start",
    "provider_request_end",
    "tool_execution_start",
    "tool_execution_update",
    "tool_execution_end",
    "turn_end",
    "agent_end",
]
RuntimeEventKind: TypeAlias = RuntimeEventType


@dataclass(frozen=True)
class AgentStartEvent:
    type: Literal["agent_start"] = "agent_start"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class TurnStartEvent:
    iteration: int
    type: Literal["turn_start"] = "turn_start"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class MessageStartEvent:
    message: Any
    iteration: int | None = None
    type: Literal["message_start"] = "message_start"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class MessageUpdateEvent:
    message: Any
    delta: str | None = None
    iteration: int | None = None
    type: Literal["message_update"] = "message_update"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class ProviderRequestStartEvent:
    iteration: int
    message_count: int
    type: Literal["provider_request_start"] = "provider_request_start"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class ProviderRequestEndEvent:
    iteration: int
    has_tool_calls: bool
    type: Literal["provider_request_end"] = "provider_request_end"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class ToolExecutionStartEvent:
    tool_call_id: str
    tool_name: str
    args: dict[str, Any]
    iteration: int
    type: Literal["tool_execution_start"] = "tool_execution_start"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class ToolExecutionUpdateEvent:
    tool_call_id: str
    tool_name: str
    args: dict[str, Any]
    partial_result: Any
    iteration: int
    type: Literal["tool_execution_update"] = "tool_execution_update"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class ToolExecutionEndEvent:
    tool_call_id: str
    tool_name: str
    args: dict[str, Any]
    result: Any
    is_error: bool
    iteration: int
    type: Literal["tool_execution_end"] = "tool_execution_end"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class TurnEndEvent:
    iteration: int
    message: Any
    tool_results: tuple[Any, ...] = ()
    type: Literal["turn_end"] = "turn_end"
    data: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class AgentEndEvent:
    messages: tuple[Any, ...] = ()
    type: Literal["agent_end"] = "agent_end"
    data: dict[str, Any] = field(default_factory=dict)


RuntimeEvent: TypeAlias = (
    AgentStartEvent
    | TurnStartEvent
    | MessageStartEvent
    | MessageUpdateEvent
    | ProviderRequestStartEvent
    | ProviderRequestEndEvent
    | ToolExecutionStartEvent
    | ToolExecutionUpdateEvent
    | ToolExecutionEndEvent
    | TurnEndEvent
    | AgentEndEvent
)
RuntimeEventCallback: TypeAlias = Callable[[RuntimeEvent], None]
# Callback shape used by surfaces to receive events as ``(kind, data)`` tuples.
# The typed :data:`RuntimeEventCallback` is the internal shape ``Agent`` calls;
# ``TupleEventCallback`` is what surfaces provide from their observer hooks.
TupleEventCallback: TypeAlias = Callable[[str, dict[str, Any]], None]


def tool_result_is_error(result: Any) -> bool:
    return isinstance(result, dict) and "error" in result


def runtime_event_from_tuple(kind: str, data: dict[str, Any]) -> RuntimeEvent | None:
    """Convert a ``(kind, data)`` tuple payload to a typed :class:`RuntimeEvent` when possible."""
    payload = dict(data)
    if kind == "agent_start":
        return AgentStartEvent(data=payload)
    if kind == "llm_start":
        return TurnStartEvent(iteration=int(payload.get("iteration", 0)), data=payload)
    if kind == "provider_request_start":
        return ProviderRequestStartEvent(
            iteration=int(payload.get("iteration", 0)),
            message_count=int(payload.get("message_count", 0)),
            data=payload,
        )
    if kind == "provider_request_end":
        return ProviderRequestEndEvent(
            iteration=int(payload.get("iteration", 0)),
            has_tool_calls=bool(payload.get("has_tool_calls", False)),
            data=payload,
        )
    if kind == "tool_start":
        args = payload.get("input")
        return ToolExecutionStartEvent(
            tool_call_id=str(payload.get("id") or payload.get("tool_call_id") or ""),
            tool_name=str(payload.get("name") or payload.get("tool_name") or "tool"),
            args=dict(args) if isinstance(args, dict) else {},
            iteration=int(payload.get("iteration", -1)),
            data=payload,
        )
    if kind == "tool_update":
        args = payload.get("input")
        return ToolExecutionUpdateEvent(
            tool_call_id=str(payload.get("id") or payload.get("tool_call_id") or ""),
            tool_name=str(payload.get("name") or payload.get("tool_name") or "tool"),
            args=dict(args) if isinstance(args, dict) else {},
            partial_result=payload.get("update"),
            iteration=int(payload.get("iteration", -1)),
            data=payload,
        )
    if kind == "tool_end":
        args = payload.get("input")
        output = payload.get("output")
        return ToolExecutionEndEvent(
            tool_call_id=str(payload.get("id") or payload.get("tool_call_id") or ""),
            tool_name=str(payload.get("name") or payload.get("tool_name") or "tool"),
            args=dict(args) if isinstance(args, dict) else {},
            result=output,
            is_error=tool_result_is_error(output),
            iteration=int(payload.get("iteration", -1)),
            data=payload,
        )
    if kind == "agent_end":
        return AgentEndEvent(data=payload)
    return None


def tuple_payload_from_event(event: RuntimeEvent) -> tuple[str, dict[str, Any]] | None:
    """Map a typed :class:`RuntimeEvent` onto the ``(kind, data)`` tuple callback shape."""
    if isinstance(event, AgentStartEvent):
        return "agent_start", dict(event.data)
    if isinstance(event, TurnStartEvent):
        return "llm_start", {"iteration": event.iteration, **event.data}
    if isinstance(event, ToolExecutionStartEvent):
        return (
            "tool_start",
            {
                "id": event.tool_call_id,
                "name": event.tool_name,
                "input": event.args,
                **event.data,
            },
        )
    if isinstance(event, ToolExecutionUpdateEvent):
        return (
            "tool_update",
            {
                "id": event.tool_call_id,
                "name": event.tool_name,
                "input": event.args,
                "update": event.partial_result,
                **event.data,
            },
        )
    if isinstance(event, ToolExecutionEndEvent):
        return (
            "tool_end",
            {
                "id": event.tool_call_id,
                "name": event.tool_name,
                "input": event.args,
                "output": event.result,
                **event.data,
            },
        )
    if isinstance(event, AgentEndEvent):
        return "agent_end", dict(event.data)
    return None


def runtime_event_callback_from_observer(
    observer: Callable[..., None] | None,
) -> RuntimeEventCallback | None:
    """Adapt a ``(kind, data)``-tuple observer to a :class:`RuntimeEventCallback`.

    Returns ``None`` when ``observer`` is omitted so callers can pass the result
    straight to ``Agent(on_runtime_event=...)``.
    """
    if observer is None:
        return None

    def on_runtime_event(event: RuntimeEvent) -> None:
        payload = tuple_payload_from_event(event)
        if payload is not None:
            observer(*payload)

    return on_runtime_event


__all__ = [
    "AgentEndEvent",
    "AgentStartEvent",
    "MessageStartEvent",
    "MessageUpdateEvent",
    "ProviderRequestEndEvent",
    "ProviderRequestStartEvent",
    "RuntimeEvent",
    "RuntimeEventCallback",
    "RuntimeEventKind",
    "RuntimeEventType",
    "ToolExecutionEndEvent",
    "ToolExecutionStartEvent",
    "ToolExecutionUpdateEvent",
    "TupleEventCallback",
    "TurnEndEvent",
    "TurnStartEvent",
    "runtime_event_callback_from_observer",
    "runtime_event_from_tuple",
    "tool_result_is_error",
    "tuple_payload_from_event",
]