Source code for taimoe.platform.types.observability

"""Observability payload models sent from runtime SDKs to the platform.

Trace and span IDs are W3C Trace Context format (lowercase hex). Pydantic
validators reject anything else at the boundary, so server-side audit
storage can rely on the format invariant without re-checking.
"""

from __future__ import annotations

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator

from taimoe.platform._ids import is_valid_span_id, is_valid_trace_id

SpanStatus = Literal["ok", "error"]
SpanType = Literal["agent", "llm", "tool", "policy", "workflow", "app", "custom"]


def _validate_trace_id(value: str) -> str:
    if not is_valid_trace_id(value):
        raise ValueError(
            f"Invalid trace_id {value!r}: expected 32 lowercase hex chars (W3C format)."
        )
    return value


def _validate_span_id(value: str) -> str:
    if not is_valid_span_id(value):
        raise ValueError(
            f"Invalid span_id {value!r}: expected 16 lowercase hex chars (W3C format)."
        )
    return value




class SpanEvent(BaseModel):
    """A single trace span emitted by a managed runtime."""

    model_config = ConfigDict(frozen=True)

    trace_id: str
    span_id: str
    parent_span_id: str | None = None
    session_id: str | None = None
    runtime_id: str
    agent_id: str | None = None
    team_id: str | None = None
    span_type: SpanType = "custom"
    name: str
    status: SpanStatus = "ok"
    start_time: datetime
    end_time: datetime
    latency_ms: float
    input: Any = None
    output: Any = None
    error: str | None = None
    prompt_tokens: int | None = None
    completion_tokens: int | None = None
    cost_usd: float | None = None
    attributes: dict[str, Any] = Field(default_factory=dict)

    @field_validator("trace_id")
    @classmethod
    def _check_trace_id(cls, value: str) -> str:
        return _validate_trace_id(value)

    @field_validator("span_id")
    @classmethod
    def _check_span_id(cls, value: str) -> str:
        return _validate_span_id(value)

    @field_validator("parent_span_id")
    @classmethod
    def _check_parent_span_id(cls, value: str | None) -> str | None:
        return None if value is None else _validate_span_id(value)





class SpanBatch(BaseModel):
    """Batch upload payload for runtime observability spans."""

    model_config = ConfigDict(frozen=True)

    runtime_id: str
    trace_id: str
    session_id: str | None = None
    spans: tuple[SpanEvent, ...] = ()

    @field_validator("trace_id")
    @classmethod
    def _check_trace_id(cls, value: str) -> str:
        return _validate_trace_id(value)