ai-pipeline-core 0.2.6__py3-none-any.whl → 0.4.1__py3-none-any.whl

ai_pipeline_core/observability/_tracking/_models.py

@@ -0,0 +1,138 @@
+"""Pydantic row models and enums for ClickHouse tracking tables."""
+
+from datetime import datetime
+from enum import StrEnum
+from typing import Protocol
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class RunStatus(StrEnum):
+    """Pipeline run status."""
+
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+class SpanType(StrEnum):
+    """Span type classification."""
+
+    TASK = "task"
+    FLOW = "flow"
+    LLM = "llm"
+    TRACE = "trace"
+
+
+class DocumentEventType(StrEnum):
+    """Document lifecycle event types."""
+
+    TASK_INPUT = "task_input"
+    TASK_OUTPUT = "task_output"
+    FLOW_INPUT = "flow_input"
+    FLOW_OUTPUT = "flow_output"
+    LLM_CONTEXT = "llm_context"
+    LLM_MESSAGE = "llm_message"
+    STORE_SAVED = "store_saved"
+    STORE_SAVE_FAILED = "store_save_failed"
+
+
+# --- Table names ---
+
+TABLE_PIPELINE_RUNS = "pipeline_runs"
+TABLE_TRACKED_SPANS = "tracked_spans"
+TABLE_DOCUMENT_EVENTS = "document_events"
+TABLE_SPAN_EVENTS = "span_events"
+
+# --- OTel span attribute names for document lineage ---
+
+ATTR_INPUT_DOCUMENT_SHA256S = "pipeline.input_document_sha256s"
+ATTR_OUTPUT_DOCUMENT_SHA256S = "pipeline.output_document_sha256s"
+
+
+# --- Row models ---
+
+
+class PipelineRunRow(BaseModel):
+    """Row model for pipeline_runs table."""
+
+    model_config = ConfigDict(frozen=True)
+
+    run_id: UUID
+    project_name: str
+    flow_name: str
+    run_scope: str = ""
+    status: RunStatus
+    start_time: datetime
+    end_time: datetime | None = None
+    total_cost: float = 0.0
+    total_tokens: int = 0
+    metadata: str = "{}"
+    version: int = 1
+
+
+class TrackedSpanRow(BaseModel):
+    """Row model for tracked_spans table."""
+
+    model_config = ConfigDict(frozen=True)
+
+    span_id: str
+    trace_id: str
+    run_id: UUID
+    parent_span_id: str | None = None
+    name: str
+    span_type: SpanType
+    status: str
+    start_time: datetime
+    end_time: datetime | None = None
+    duration_ms: int = 0
+    cost: float = 0.0
+    tokens_input: int = 0
+    tokens_output: int = 0
+    llm_model: str | None = None
+    user_summary: str | None = None
+    user_visible: bool = False
+    user_label: str | None = None
+    input_document_sha256s: tuple[str, ...] = Field(default_factory=tuple)
+    output_document_sha256s: tuple[str, ...] = Field(default_factory=tuple)
+    version: int = 1
+
+
+class DocumentEventRow(BaseModel):
+    """Row model for document_events table."""
+
+    model_config = ConfigDict(frozen=True)
+
+    event_id: UUID
+    run_id: UUID
+    document_sha256: str
+    span_id: str
+    event_type: DocumentEventType
+    timestamp: datetime
+    metadata: str = "{}"
+
+
+class SpanEventRow(BaseModel):
+    """Row model for span_events table."""
+
+    model_config = ConfigDict(frozen=True)
+
+    event_id: UUID
+    run_id: UUID
+    span_id: str
+    name: str
+    timestamp: datetime
+    attributes: str = "{}"
+    level: str | None = None
+
+
+# --- Protocol for circular import resolution ---
+
+
+class SummaryRowBuilder(Protocol):
+    """Protocol satisfied by TrackingService for writer callback."""
+
+    def build_span_summary_update(self, span_id: str, summary: str) -> TrackedSpanRow | None:
+        """Build a replacement row with summary filled."""
+        ...

ai_pipeline_core/observability/_tracking/_processor.py

@@ -0,0 +1,158 @@
+"""OpenTelemetry SpanProcessor that feeds the tracking system."""
+
+from datetime import UTC, datetime
+from typing import Any
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+from opentelemetry.trace import StatusCode
+
+from ai_pipeline_core.logging import get_pipeline_logger
+
+from ._internal import is_internal_tracking
+from ._models import ATTR_INPUT_DOCUMENT_SHA256S, ATTR_OUTPUT_DOCUMENT_SHA256S, SpanType
+from ._service import TrackingService
+
+logger = get_pipeline_logger(__name__)
+
+
+def _hex_span_id(span_id: int) -> str:
+    """Convert integer span ID to hex string."""
+    return format(span_id, "016x")
+
+
+def _hex_trace_id(trace_id: int) -> str:
+    """Convert integer trace ID to hex string."""
+    return format(trace_id, "032x")
+
+
+def _ns_to_datetime(ns: int) -> datetime:
+    """Convert nanosecond timestamp to datetime."""
+    return datetime.fromtimestamp(ns / 1e9, tz=UTC)
+
+
+def _classify_span(attrs: dict[str, Any]) -> SpanType:
+    """Determine span type from attributes."""
+    span_type_str = str(attrs.get("lmnr.span.type", ""))
+    if span_type_str == "LLM":
+        return SpanType.LLM
+    if attrs.get("prefect.flow.name"):
+        return SpanType.FLOW
+    if attrs.get("prefect.task.name"):
+        return SpanType.TASK
+    return SpanType.TRACE
+
+
+class TrackingSpanProcessor(SpanProcessor):
+    """Forwards completed spans to TrackingService.
+
+    Skips internal tracking spans (summary LLM calls) to prevent recursion.
+    """
+
+    def __init__(self, service: TrackingService) -> None:
+        """Initialize with tracking service."""
+        self._service = service
+
+    @staticmethod
+    def _parent_span_id(span: Span | ReadableSpan) -> str | None:
+        """Extract parent span ID as hex string, or None."""
+        parent = span.parent
+        if parent is None:
+            return None
+        return _hex_span_id(parent.span_id)
+
+    def on_start(self, span: Span, parent_context: Context | None = None) -> None:
+        """Record span start."""
+        if is_internal_tracking():
+            return
+        try:
+            ctx = span.get_span_context()
+            if ctx is None:
+                return
+            attrs: dict[str, Any] = dict(span.attributes or {})
+            self._service.track_span_start(
+                span_id=_hex_span_id(ctx.span_id),
+                trace_id=_hex_trace_id(ctx.trace_id),
+                parent_span_id=self._parent_span_id(span),
+                name=span.name,
+                span_type=_classify_span(attrs),
+            )
+        except Exception as e:
+            logger.debug(f"TrackingSpanProcessor.on_start failed: {e}")
+
+    def on_end(self, span: ReadableSpan) -> None:  # noqa: PLR0914
+        """Record span completion with full details."""
+        if is_internal_tracking():
+            return
+        try:
+            ctx = span.get_span_context()
+            if ctx is None:
+                return
+            attrs: dict[str, Any] = dict(span.attributes or {})
+
+            start_ns = span.start_time or 0
+            end_ns = span.end_time or 0
+            start_time = _ns_to_datetime(start_ns)
+            end_time = _ns_to_datetime(end_ns)
+            duration_ms = max(0, (end_ns - start_ns) // 1_000_000)
+
+            status = "failed" if span.status.status_code == StatusCode.ERROR else "completed"
+
+            # Extract LLM-specific attributes
+            cost = float(attrs.get("gen_ai.usage.cost", 0.0))
+            tokens_input = int(attrs.get("gen_ai.usage.input_tokens", 0))
+            tokens_output = int(attrs.get("gen_ai.usage.output_tokens", 0))
+            llm_model = str(attrs.get("gen_ai.request.model", "")) or None
+
+            # Extract document SHA256 arrays set by track_task_io
+            raw_input_sha256s = attrs.get(ATTR_INPUT_DOCUMENT_SHA256S)
+            input_doc_sha256s = list(raw_input_sha256s) if raw_input_sha256s else None
+            raw_output_sha256s = attrs.get(ATTR_OUTPUT_DOCUMENT_SHA256S)
+            output_doc_sha256s = list(raw_output_sha256s) if raw_output_sha256s else None
+
+            span_id = _hex_span_id(ctx.span_id)
+            self._service.track_span_end(
+                span_id=span_id,
+                trace_id=_hex_trace_id(ctx.trace_id),
+                parent_span_id=self._parent_span_id(span),
+                name=span.name,
+                span_type=_classify_span(attrs),
+                status=status,
+                start_time=start_time,
+                end_time=end_time,
+                duration_ms=duration_ms,
+                cost=cost,
+                tokens_input=tokens_input,
+                tokens_output=tokens_output,
+                llm_model=llm_model,
+                input_document_sha256s=input_doc_sha256s,
+                output_document_sha256s=output_doc_sha256s,
+            )
+
+            # Forward span events
+            if span.events:
+                events: list[tuple[str, datetime, dict[str, str], str | None]] = []
+                for event in span.events:
+                    event_attrs = dict(event.attributes) if event.attributes else {}
+                    level = str(event_attrs.pop("log.level", "")) or None
+                    events.append((
+                        event.name,
+                        _ns_to_datetime(event.timestamp),
+                        {k: str(v) for k, v in event_attrs.items()},
+                        level,
+                    ))
+                self._service.track_span_events(
+                    span_id=span_id,
+                    events=events,
+                )
+        except Exception as e:
+            logger.debug(f"TrackingSpanProcessor.on_end failed: {e}")
+
+    def shutdown(self) -> None:
+        """Shutdown the tracking service."""
+        self._service.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:  # noqa: PLR6301
+        """Force flush is a no-op — the writer flushes on its own schedule."""
+        _ = timeout_millis
+        return True

ai_pipeline_core/observability/_tracking/_service.py

@@ -0,0 +1,311 @@
+"""TrackingService — central coordinator for pipeline observability.
+
+Manages run context, version counters, row caching for summary updates,
+and coordinates the ClickHouse writer thread.
+"""
+
+import json
+from datetime import UTC, datetime
+from threading import Lock
+from uuid import UUID, uuid4
+
+from lmnr.opentelemetry_lib.tracing import context as laminar_context
+from opentelemetry import context as otel_context
+
+from ai_pipeline_core.logging import get_pipeline_logger
+
+from ._client import ClickHouseClient
+from ._models import (
+    TABLE_DOCUMENT_EVENTS,
+    TABLE_PIPELINE_RUNS,
+    TABLE_SPAN_EVENTS,
+    TABLE_TRACKED_SPANS,
+    DocumentEventRow,
+    DocumentEventType,
+    PipelineRunRow,
+    RunStatus,
+    SpanEventRow,
+    SpanType,
+    TrackedSpanRow,
+)
+from ._writer import ClickHouseWriter, SpanSummaryFn, SummaryJob
+
+logger = get_pipeline_logger(__name__)
+
+
+class TrackingService:
+    """Central tracking coordinator.
+
+    Thread-safe — all mutable state is protected by ``_lock``.
+    """
+
+    def __init__(
+        self,
+        client: ClickHouseClient,
+        *,
+        summary_model: str = "gemini-3-flash",
+        span_summary_fn: SpanSummaryFn | None = None,
+    ) -> None:
+        """Initialize tracking service and start writer thread."""
+        self._client = client
+        self._summary_model = summary_model
+
+        self._writer = ClickHouseWriter(
+            client,
+            summary_row_builder=self,
+            span_summary_fn=span_summary_fn,
+        )
+        self._writer.start()
+
+        # Run context
+        self._run_id: UUID | None = None
+        self._project_name: str = ""
+        self._flow_name: str = ""
+        self._run_scope: str = ""
+        self._run_start_time: datetime | None = None
+
+        # Version counters
+        self._versions: dict[str, int] = {}
+        self._lock = Lock()
+
+        # Row caches for summary updates
+        self._span_cache: dict[str, TrackedSpanRow] = {}
+
+    # --- Run context ---
+
+    def set_run_context(self, *, run_id: UUID, project_name: str, flow_name: str, run_scope: str = "") -> None:
+        """Set the current run context. Called at pipeline start."""
+        with self._lock:
+            self._run_id = run_id
+            self._project_name = project_name
+            self._flow_name = flow_name
+            self._run_scope = run_scope
+
+    def clear_run_context(self) -> None:
+        """Clear run context and caches. Called by flush() and shutdown()."""
+        with self._lock:
+            self._run_id = None
+            self._project_name = ""
+            self._flow_name = ""
+            self._run_scope = ""
+            self._run_start_time = None
+            self._versions.clear()
+            self._span_cache.clear()
+
+    # --- Version management ---
+
+    def _next_version(self, key: str) -> int:
+        """Increment and return version counter. Must be called under _lock."""
+        v = self._versions.get(key, 0) + 1
+        self._versions[key] = v
+        return v
+
+    # --- Run tracking ---
+
+    def track_run_start(self, *, run_id: UUID, project_name: str, flow_name: str, run_scope: str = "") -> None:
+        """Record pipeline run start."""
+        now = datetime.now(UTC)
+        with self._lock:
+            self._run_start_time = now
+            version = self._next_version(f"run:{run_id}")
+        row = PipelineRunRow(
+            run_id=run_id,
+            project_name=project_name,
+            flow_name=flow_name,
+            run_scope=run_scope,
+            status=RunStatus.RUNNING,
+            start_time=now,
+            version=version,
+        )
+        self._writer.write(TABLE_PIPELINE_RUNS, [row])
+
+    def track_run_end(
+        self,
+        *,
+        run_id: UUID,
+        status: RunStatus,
+        total_cost: float = 0.0,
+        total_tokens: int = 0,
+        metadata: dict[str, object] | None = None,
+    ) -> None:
+        """Record pipeline run completion or failure."""
+        now = datetime.now(UTC)
+        with self._lock:
+            version = self._next_version(f"run:{run_id}")
+            start_time = self._run_start_time or now
+        row = PipelineRunRow(
+            run_id=run_id,
+            project_name=self._project_name,
+            flow_name=self._flow_name,
+            run_scope=self._run_scope,
+            status=status,
+            start_time=start_time,
+            end_time=now,
+            total_cost=total_cost,
+            total_tokens=total_tokens,
+            metadata=json.dumps(metadata) if metadata else "{}",
+            version=version,
+        )
+        self._writer.write(TABLE_PIPELINE_RUNS, [row])
+
+    # --- Span tracking ---
+
+    def track_span_start(self, *, span_id: str, trace_id: str, parent_span_id: str | None, name: str, span_type: SpanType) -> None:
+        """Record span start."""
+        if self._run_id is None:
+            return
+        now = datetime.now(UTC)
+        with self._lock:
+            version = self._next_version(f"span:{span_id}")
+        row = TrackedSpanRow(
+            span_id=span_id,
+            trace_id=trace_id,
+            run_id=self._run_id,
+            parent_span_id=parent_span_id,
+            name=name,
+            span_type=span_type,
+            status="running",
+            start_time=now,
+            version=version,
+        )
+        self._writer.write(TABLE_TRACKED_SPANS, [row])
+
+    def track_span_end(
+        self,
+        *,
+        span_id: str,
+        trace_id: str,
+        parent_span_id: str | None,
+        name: str,
+        span_type: SpanType,
+        status: str,
+        start_time: datetime,
+        end_time: datetime,
+        duration_ms: int,
+        cost: float = 0.0,
+        tokens_input: int = 0,
+        tokens_output: int = 0,
+        llm_model: str | None = None,
+        user_summary: str | None = None,
+        user_visible: bool = False,
+        user_label: str | None = None,
+        input_document_sha256s: list[str] | None = None,
+        output_document_sha256s: list[str] | None = None,
+    ) -> None:
+        """Record span completion with full details."""
+        if self._run_id is None:
+            return
+        with self._lock:
+            version = self._next_version(f"span:{span_id}")
+        row = TrackedSpanRow(
+            span_id=span_id,
+            trace_id=trace_id,
+            run_id=self._run_id,
+            parent_span_id=parent_span_id,
+            name=name,
+            span_type=span_type,
+            status=status,
+            start_time=start_time,
+            end_time=end_time,
+            duration_ms=duration_ms,
+            cost=cost,
+            tokens_input=tokens_input,
+            tokens_output=tokens_output,
+            llm_model=llm_model,
+            user_summary=user_summary,
+            user_visible=user_visible,
+            user_label=user_label,
+            input_document_sha256s=tuple(input_document_sha256s) if input_document_sha256s else (),
+            output_document_sha256s=tuple(output_document_sha256s) if output_document_sha256s else (),
+            version=version,
+        )
+        self._writer.write(TABLE_TRACKED_SPANS, [row])
+        with self._lock:
+            self._span_cache[span_id] = row
+
+    def track_span_events(self, *, span_id: str, events: list[tuple[str, datetime, dict[str, str], str | None]]) -> None:
+        """Record span events (including bridged log events)."""
+        if self._run_id is None or not events:
+            return
+        rows = [
+            SpanEventRow(
+                event_id=uuid4(),
+                run_id=self._run_id,
+                span_id=span_id,
+                name=name,
+                timestamp=ts,
+                attributes=json.dumps(attrs) if attrs else "{}",
+                level=level,
+            )
+            for name, ts, attrs, level in events
+        ]
+        self._writer.write(TABLE_SPAN_EVENTS, list(rows))
+
+    # --- Document tracking ---
+
+    def track_document_event(
+        self,
+        *,
+        document_sha256: str,
+        span_id: str,
+        event_type: DocumentEventType,
+        metadata: dict[str, str] | None = None,
+    ) -> None:
+        """Record a document lifecycle event."""
+        if self._run_id is None:
+            return
+        row = DocumentEventRow(
+            event_id=uuid4(),
+            run_id=self._run_id,
+            document_sha256=document_sha256,
+            span_id=span_id,
+            event_type=event_type,
+            timestamp=datetime.now(UTC),
+            metadata=json.dumps(metadata) if metadata else "{}",
+        )
+        self._writer.write(TABLE_DOCUMENT_EVENTS, [row])
+
+    # --- Summary scheduling ---
+
+    def schedule_summary(self, span_id: str, label: str, output_hint: str) -> None:
+        """Schedule LLM summary generation for a span."""
+        self._writer.write_job(
+            SummaryJob(
+                span_id=span_id,
+                label=label,
+                output_hint=output_hint,
+                summary_model=self._summary_model,
+                parent_otel_context=otel_context.get_current(),
+                parent_laminar_context=laminar_context.get_current_context(),
+            )
+        )
+
+    # --- Summary row builders (SummaryRowBuilder protocol) ---
+
+    def build_span_summary_update(self, span_id: str, summary: str) -> TrackedSpanRow | None:
+        """Build a replacement row with summary filled. Called from writer thread."""
+        with self._lock:
+            cached = self._span_cache.get(span_id)
+            if cached is None:
+                return None
+            version = self._next_version(f"span:{span_id}")
+        return cached.model_copy(update={"user_summary": summary, "version": version})
+
+    # --- Lifecycle ---
+
+    def flush(self, timeout: float = 30.0) -> None:
+        """Wait for all pending items (including summary LLM jobs) to complete, then clear run context.
+
+        Use between runs in long-lived processes to prevent unbounded cache growth.
+        """
+        self._writer.flush(timeout=timeout)
+        self.clear_run_context()
+
+    def shutdown(self, timeout: float = 30.0) -> None:
+        """Shutdown the writer thread and clear run context.
+
+        Writer drains all pending items (including summary LLM jobs)
+        before caches are cleared, ensuring summaries can look up span data.
+        """
+        self._writer.shutdown(timeout=timeout)
+        self.clear_run_context()