RouteKitAI 0.1.0__py3-none-any.whl

routekitai/observability/exporters/otel.py

@@ -0,0 +1,119 @@
+"""OpenTelemetry trace exporter."""
+
+import json
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from routekitai.core.errors import RuntimeError as RouteKitRuntimeError
+from routekitai.observability.trace import Trace
+
+
+class OTELExporterError(RouteKitRuntimeError):
+    """Error raised by OTEL exporter operations."""
+
+    pass
+
+
+class OTELExporter(BaseModel):
+    """Exports traces to OpenTelemetry format.
+
+    Converts routkitai traces to OpenTelemetry format and exports them.
+    For MVP, exports in JSON format compatible with OTEL collectors.
+    """
+
+    endpoint: str | None = Field(default=None, description="OTEL collector endpoint")
+    headers: dict[str, str] = Field(default_factory=dict, description="Export headers")
+
+    def _convert_trace_to_otel(self, trace: Trace) -> dict[str, Any]:
+        """Convert routkitai trace to OpenTelemetry format.
+
+        Args:
+            trace: routkitai trace
+
+        Returns:
+            OTEL-compatible trace data
+        """
+        spans = []
+        for event in trace.events:
+            span = {
+                "traceId": trace.trace_id,
+                "spanId": event.data.get("step_id", "unknown"),
+                "name": event.type,
+                "startTimeUnixNano": int(event.timestamp * 1_000_000_000),
+                "endTimeUnixNano": int(event.timestamp * 1_000_000_000),
+                "attributes": [
+                    {"key": "event.type", "value": {"stringValue": event.type}},
+                    {"key": "trace.id", "value": {"stringValue": trace.trace_id}},
+                ],
+            }
+
+            # Add event data as attributes
+            for key, value in event.data.items():
+                if isinstance(value, (str, int, float, bool)):
+                    span["attributes"].append({"key": key, "value": {"stringValue": str(value)}})
+                elif isinstance(value, dict):
+                    span["attributes"].append(
+                        {"key": key, "value": {"stringValue": json.dumps(value)}}
+                    )
+
+            spans.append(span)
+
+        return {
+            "resourceSpans": [
+                {
+                    "resource": {
+                        "attributes": [
+                            {"key": "service.name", "value": {"stringValue": "routekit"}},
+                        ]
+                    },
+                    "scopeSpans": [{"spans": spans}],
+                }
+            ]
+        }
+
+    async def export(self, trace: Trace) -> None:
+        """Export trace to OpenTelemetry.
+
+        Args:
+            trace: Trace to export
+
+        Raises:
+            OTELExporterError: If export fails
+        """
+        try:
+            otel_data = self._convert_trace_to_otel(trace)
+
+            if self.endpoint:
+                # Export to OTEL collector endpoint
+                try:
+                    import httpx
+
+                    async with httpx.AsyncClient() as client:
+                        response = await client.post(
+                            self.endpoint,
+                            json=otel_data,
+                            headers=self.headers,
+                            timeout=10.0,
+                        )
+                        response.raise_for_status()
+                except ImportError:
+                    raise OTELExporterError(
+                        "httpx is required for OTEL export. Install with: pip install httpx",
+                        context={"endpoint": self.endpoint},
+                    ) from None
+                except Exception as e:
+                    raise OTELExporterError(
+                        f"Failed to export trace to OTEL endpoint: {e}",
+                        context={"endpoint": self.endpoint, "trace_id": trace.trace_id},
+                    ) from e
+            else:
+                # No endpoint specified - just log the OTEL format (for debugging)
+                import logging
+
+                logger = logging.getLogger(__name__)
+                logger.debug(f"OTEL trace (no endpoint): {json.dumps(otel_data, indent=2)}")
+        except Exception as e:
+            raise OTELExporterError(
+                f"OTEL export failed: {e}", context={"trace_id": trace.trace_id}
+            ) from e

routekitai/observability/spans.py

@@ -0,0 +1,111 @@
+"""Span management for distributed tracing."""
+
+import contextlib
+import time
+from collections.abc import AsyncGenerator, Generator
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class SpanContext(BaseModel):
+    """Context for span propagation."""
+
+    trace_id: str = Field(..., description="Trace ID")
+    span_id: str = Field(..., description="Span ID")
+    parent_span_id: str | None = Field(default=None, description="Parent span ID")
+
+
+class Span(BaseModel):
+    """Execution span for tracing."""
+
+    span_id: str = Field(..., description="Span ID")
+    name: str = Field(..., description="Span name")
+    start_time: float = Field(..., description="Start timestamp")
+    end_time: float | None = Field(default=None, description="End timestamp")
+    attributes: dict[str, Any] = Field(default_factory=dict, description="Span attributes")
+    events: list[dict[str, Any]] = Field(default_factory=list, description="Span events")
+
+    def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None:
+        """Add event to span.
+
+        Args:
+            name: Event name
+            attributes: Optional event attributes
+        """
+        self.events.append(
+            {
+                "name": name,
+                "timestamp": time.time(),
+                "attributes": attributes or {},
+            }
+        )
+
+    @property
+    def duration(self) -> float | None:
+        """Get span duration in seconds.
+
+        Returns:
+            Duration if span is ended, None otherwise
+        """
+        if self.end_time is None:
+            return None
+        return self.end_time - self.start_time
+
+
+@contextlib.contextmanager
+def span(name: str, attributes: dict[str, Any] | None = None) -> Generator[Span, None, None]:
+    """Context manager for creating a span.
+
+    Args:
+        name: Span name
+        attributes: Optional span attributes
+
+    Yields:
+        Span instance
+
+    Example:
+        >>> with span("operation") as s:
+        ...     # do work
+        ...     s.add_event("milestone")
+    """
+    s = Span(
+        span_id=f"span_{int(time.time() * 1000000)}",
+        name=name,
+        start_time=time.time(),
+        attributes=attributes or {},
+    )
+    try:
+        yield s
+    finally:
+        s.end_time = time.time()
+
+
+@contextlib.asynccontextmanager
+async def async_span(
+    name: str, attributes: dict[str, Any] | None = None
+) -> AsyncGenerator[Span, None]:
+    """Async context manager for creating a span.
+
+    Args:
+        name: Span name
+        attributes: Optional span attributes
+
+    Yields:
+        Span instance
+
+    Example:
+        >>> async with async_span("async_operation") as s:
+        ...     # do async work
+        ...     s.add_event("milestone")
+    """
+    s = Span(
+        span_id=f"span_{int(time.time() * 1000000)}",
+        name=name,
+        start_time=time.time(),
+        attributes=attributes or {},
+    )
+    try:
+        yield s
+    finally:
+        s.end_time = time.time()

routekitai/observability/streaming.py

@@ -0,0 +1,117 @@
+"""Streaming support for trace events."""
+
+import asyncio
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+
+class TraceEventBroadcaster:
+    """Broadcasts trace events to multiple subscribers via WebSocket/SSE."""
+
+    def __init__(self) -> None:
+        """Initialize broadcaster."""
+        self._subscribers: set[asyncio.Queue[Any]] = set()
+        self._lock = asyncio.Lock()
+
+    async def subscribe(self) -> asyncio.Queue[Any]:
+        """Subscribe to trace events.
+
+        Returns:
+            Queue that will receive trace events
+        """
+        queue: asyncio.Queue[Any] = asyncio.Queue()
+        async with self._lock:
+            self._subscribers.add(queue)
+        return queue
+
+    async def unsubscribe(self, queue: asyncio.Queue[Any]) -> None:
+        """Unsubscribe from trace events.
+
+        Args:
+            queue: Queue to remove from subscribers
+        """
+        async with self._lock:
+            self._subscribers.discard(queue)
+
+    async def broadcast(self, event: Any) -> None:
+        """Broadcast an event to all subscribers.
+
+        Args:
+            event: Trace event to broadcast
+        """
+        async with self._lock:
+            # Create a copy of subscribers to avoid modification during iteration
+            subscribers = list(self._subscribers)
+
+        # Send to all subscribers, removing dead ones
+        dead_subscribers = []
+        for queue in subscribers:
+            try:
+                queue.put_nowait(event)
+            except Exception:
+                dead_subscribers.append(queue)
+
+        # Clean up dead subscribers
+        if dead_subscribers:
+            async with self._lock:
+                for queue in dead_subscribers:
+                    self._subscribers.discard(queue)
+
+    async def stream_events(
+        self, queue: asyncio.Queue[Any], trace_id: str | None = None
+    ) -> AsyncIterator[str]:
+        """Stream events as SSE format.
+
+        Args:
+            queue: Queue to read events from
+            trace_id: Optional trace ID to filter events
+
+        Yields:
+            SSE-formatted strings
+        """
+        try:
+            while True:
+                try:
+                    # Wait for event with timeout to allow periodic checks
+                    event = await asyncio.wait_for(queue.get(), timeout=1.0)
+                    queue.task_done()
+
+                    # Filter by trace_id if provided
+                    if (
+                        trace_id
+                        and hasattr(event, "data")
+                        and event.data.get("trace_id") != trace_id
+                    ):
+                        continue
+
+                    # Format as SSE
+                    event_data = {
+                        "type": event.type,
+                        "timestamp": event.timestamp,
+                        "data": event.data,
+                    }
+                    yield f"data: {json.dumps(event_data)}\n\n"
+                except TimeoutError:
+                    # Send keepalive
+                    yield ": keepalive\n\n"
+                except asyncio.CancelledError:
+                    break
+        finally:
+            await self.unsubscribe(queue)
+
+
+# Global broadcaster instance
+_broadcaster: TraceEventBroadcaster | None = None
+
+
+def get_broadcaster() -> TraceEventBroadcaster:
+    """Get the global trace event broadcaster.
+
+    Returns:
+        TraceEventBroadcaster instance
+    """
+    global _broadcaster
+    if _broadcaster is None:
+        _broadcaster = TraceEventBroadcaster()
+    return _broadcaster

routekitai/observability/trace.py

@@ -0,0 +1,144 @@
+"""Trace collection and management."""
+
+import asyncio
+import time
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class TraceEvent(BaseModel):
+    """Immutable event in a trace."""
+
+    type: str = Field(..., description="Event type")
+    timestamp: float = Field(..., description="Event timestamp")
+    data: dict[str, Any] = Field(..., description="Event data")
+
+
+TraceEventCallback = Callable[[TraceEvent], Awaitable[None] | None]
+
+
+class Trace(BaseModel):
+    """Immutable trace of agent execution.
+
+    A trace is an append-only event log that records all execution events.
+    """
+
+    trace_id: str = Field(..., description="Trace ID")
+    events: list[TraceEvent] = Field(default_factory=list, description="Trace events")
+    metadata: dict[str, Any] = Field(default_factory=dict, description="Trace metadata")
+
+    def __init__(self, **data: Any) -> None:
+        """Initialize trace with callbacks."""
+        super().__init__(**data)
+        self._event_callbacks: list[TraceEventCallback] = []
+
+    def add_event_callback(self, callback: TraceEventCallback) -> None:
+        """Add a callback to be notified of new events.
+
+        Args:
+            callback: Callback function that receives TraceEvent
+        """
+        self._event_callbacks.append(callback)
+
+    def remove_event_callback(self, callback: TraceEventCallback) -> None:
+        """Remove an event callback.
+
+        Args:
+            callback: Callback function to remove
+        """
+        if callback in self._event_callbacks:
+            self._event_callbacks.remove(callback)
+
+    def add_event(self, event_type: str, data: dict[str, Any] | None = None) -> None:
+        """Add an event to the trace.
+
+        Args:
+            event_type: Type of event (run_started, model_called, tool_called, etc.)
+            data: Event data
+        """
+        event = TraceEvent(
+            type=event_type,
+            timestamp=time.time(),
+            data=data or {},
+        )
+        self.events.append(event)
+
+        # Notify callbacks
+        for callback in self._event_callbacks:
+            try:
+                result = callback(event)
+                # If callback is async, schedule it
+                if asyncio.iscoroutine(result):
+                    # Create task if we're in an event loop, otherwise this will be handled by caller
+                    try:
+                        loop = asyncio.get_event_loop()
+                        if loop.is_running():
+                            asyncio.create_task(result)
+                    except RuntimeError:
+                        # No event loop, skip async callback
+                        pass
+            except Exception:
+                # Don't let callback errors break trace collection
+                pass
+
+        # Broadcast to streaming subscribers (lazy import to avoid circular dependency)
+        try:
+            from routekitai.observability.streaming import get_broadcaster
+
+            broadcaster = get_broadcaster()
+            # Schedule broadcast in event loop if available
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_running():
+                    asyncio.create_task(broadcaster.broadcast(event))
+            except RuntimeError:
+                # No event loop, skip broadcasting
+                pass
+        except Exception:
+            # Don't let broadcasting errors break trace collection
+            pass
+
+    def get_events_by_type(self, event_type: str) -> list[TraceEvent]:
+        """Get all events of a specific type.
+
+        Args:
+            event_type: Event type to filter by
+
+        Returns:
+            List of matching events
+        """
+        return [event for event in self.events if event.type == event_type]
+
+
+class TraceCollector(BaseModel):
+    """Collects and manages traces."""
+
+    traces: dict[str, Trace] = Field(default_factory=dict, description="Collected traces")
+
+    def start_trace(self, trace_id: str, metadata: dict[str, Any] | None = None) -> Trace:
+        """Start a new trace.
+
+        Args:
+            trace_id: Unique trace ID
+            metadata: Optional trace metadata
+
+        Returns:
+            New trace instance
+        """
+        trace = Trace(trace_id=trace_id, metadata=metadata or {})
+        trace.add_event("run_started", {"trace_id": trace_id, "metadata": metadata or {}})
+        self.traces[trace_id] = trace
+        return trace
+
+    def get_trace(self, trace_id: str) -> Trace | None:
+        """Get trace by ID.
+
+        Args:
+            trace_id: Trace ID
+
+        Returns:
+            Trace if found, None otherwise
+        """
+        return self.traces.get(trace_id)

routekitai/providers/__init__.py

@@ -0,0 +1,9 @@
+"""Model providers for RouteKit."""
+
+from routekitai.providers.local import FakeModel
+from routekitai.providers.openai import OpenAIChatModel
+
+__all__ = [
+    "OpenAIChatModel",
+    "FakeModel",
+]