stonks-oracle/services/intelligence_pipeline_v3/observability/tracing.py

"""Distributed tracing for the v3 intelligence pipeline.

Every document gets one trace ID that covers preprocessing, specialist
stages, routing, adjudication, impact prediction, and persistence.
"""

from __future__ import annotations

import enum
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from uuid import UUID, uuid4


class SpanStatus(str, enum.Enum):
    """Status of a trace span."""

    RUNNING = "running"
    SUCCEEDED = "succeeded"
    FAILED = "failed"
    SKIPPED = "skipped"


@dataclass
class StageSpan:
    """A single stage span within a pipeline trace."""

    span_id: UUID
    trace_id: UUID
    stage_name: str
    parent_span_id: UUID | None
    started_at: datetime
    ended_at: datetime | None = None
    status: SpanStatus = SpanStatus.RUNNING
    duration_ms: float = 0.0
    attributes: dict[str, Any] = field(default_factory=dict)
    error_message: str | None = None

    def finish(
        self,
        status: SpanStatus = SpanStatus.SUCCEEDED,
        error_message: str | None = None,
    ) -> None:
        """Mark the span as complete."""
        self.ended_at = datetime.now(timezone.utc)
        self.status = status
        self.error_message = error_message
        if self.started_at and self.ended_at:
            self.duration_ms = (
                self.ended_at - self.started_at
            ).total_seconds() * 1000

    def set_attribute(self, key: str, value: Any) -> None:
        """Add a span attribute."""
        self.attributes[key] = value


@dataclass
class PipelineTrace:
    """Complete distributed trace for one document through the pipeline."""

    trace_id: UUID
    document_id: str
    run_id: UUID
    started_at: datetime
    ended_at: datetime | None = None
    spans: list[StageSpan] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def create(cls, document_id: str, run_id: UUID) -> PipelineTrace:
        return cls(
            trace_id=uuid4(),
            document_id=document_id,
            run_id=run_id,
            started_at=datetime.now(timezone.utc),
        )

    def start_span(
        self,
        stage_name: str,
        parent_span_id: UUID | None = None,
        attributes: dict[str, Any] | None = None,
    ) -> StageSpan:
        """Start a new span for a pipeline stage."""
        span = StageSpan(
            span_id=uuid4(),
            trace_id=self.trace_id,
            stage_name=stage_name,
            parent_span_id=parent_span_id,
            started_at=datetime.now(timezone.utc),
            attributes=attributes or {},
        )
        self.spans.append(span)
        return span

    def finish(self) -> None:
        """Mark the trace as complete."""
        self.ended_at = datetime.now(timezone.utc)

    @property
    def total_duration_ms(self) -> float:
        if self.started_at and self.ended_at:
            return (self.ended_at - self.started_at).total_seconds() * 1000
        return 0.0

    @property
    def failed_spans(self) -> list[StageSpan]:
        return [s for s in self.spans if s.status == SpanStatus.FAILED]

    @property
    def is_complete(self) -> bool:
        return self.ended_at is not None

    def to_dict(self) -> dict[str, Any]:
        """Serialize trace for export/storage."""
        return {
            "trace_id": str(self.trace_id),
            "document_id": self.document_id,
            "run_id": str(self.run_id),
            "started_at": self.started_at.isoformat(),
            "ended_at": self.ended_at.isoformat() if self.ended_at else None,
            "total_duration_ms": self.total_duration_ms,
            "span_count": len(self.spans),
            "failed_span_count": len(self.failed_spans),
            "metadata": self.metadata,
            "spans": [
                {
                    "span_id": str(s.span_id),
                    "stage_name": s.stage_name,
                    "status": s.status.value,
                    "duration_ms": s.duration_ms,
                    "attributes": s.attributes,
                    "error_message": s.error_message,
                }
                for s in self.spans
            ],
        }


@dataclass
class TraceCollector:
    """Collects and stores pipeline traces.

    In production, this would export to an observability backend
    (Jaeger, Tempo, etc.). This implementation provides the collection
    logic for testing and local development.
    """

    _traces: dict[UUID, PipelineTrace] = field(default_factory=dict)
    max_stored: int = 10000

    def start_trace(self, document_id: str, run_id: UUID) -> PipelineTrace:
        """Create and store a new trace."""
        trace = PipelineTrace.create(document_id, run_id)
        self._traces[trace.trace_id] = trace
        # Evict oldest if over limit
        if len(self._traces) > self.max_stored:
            oldest_key = next(iter(self._traces))
            del self._traces[oldest_key]
        return trace

    def get_trace(self, trace_id: UUID) -> PipelineTrace | None:
        return self._traces.get(trace_id)

    def get_by_document(self, document_id: str) -> list[PipelineTrace]:
        return [
            t for t in self._traces.values() if t.document_id == document_id
        ]

    def get_by_run(self, run_id: UUID) -> PipelineTrace | None:
        for t in self._traces.values():
            if t.run_id == run_id:
                return t
        return None

    @property
    def trace_count(self) -> int:
        return len(self._traces)