kalibr 1.0.28__py3-none-any.whl → 1.1.0__py3-none-any.whl

kalibr/client.py

@@ -0,0 +1,293 @@
+"""Kalibr client with context propagation and batch flushing.
+
+Refactored for Phase 1 - SDK Stabilization:
+- Uses modular tracer and cost adapters
+- Supports new event schema v1.0
+- Enhanced error handling with stack traces
+- Environment-based configuration
+"""
+
+import atexit
+import hashlib
+import hmac
+import queue
+import threading
+import time
+from typing import Any, Callable, Dict, Optional
+
+import httpx
+
+from .decorators import create_trace_decorator
+from .tracer import Tracer
+from .utils import (
+    load_config_from_env,
+    log_error,
+    log_info,
+    log_success,
+    log_warning,
+    serialize_event,
+    validate_event,
+)
+
+
+class KalibrClient:
+    """Kalibr observability client with Phase 1 enhancements.
+
+    Features:
+    - Modular tracer with cost adapters
+    - New event schema v1.0
+    - Enhanced error handling
+    - Environment-based configuration
+
+    Args:
+        api_key: Kalibr API key
+        endpoint: Collector endpoint URL (optional, can load from env)
+        tenant_id: Tenant identifier (optional, can load from env)
+        environment: Environment (prod/staging/dev, can load from env)
+        service: Service name (optional, can load from env)
+        secret: HMAC secret for request signing
+        batch_size: Max events per batch
+        flush_interval: Flush interval in seconds
+        max_queue_size: Max queue size (drops oldest on overflow)
+    """
+
+    def __init__(
+        self,
+        api_key: str = None,
+        endpoint: str = None,
+        tenant_id: str = None,
+        environment: str = None,
+        service: str = None,
+        workflow_id: str = None,
+        workflow_version: str = None,
+        secret: Optional[str] = None,
+        batch_size: int = 100,
+        flush_interval: float = 2.0,
+        max_queue_size: int = 5000,
+    ):
+        # Load config from environment if not provided
+        env_config = load_config_from_env()
+
+        self.api_key = api_key or env_config.get("auth_token", "")
+        self.endpoint = endpoint or env_config.get(
+            "api_endpoint", "http://localhost:8001/api/v1/traces"
+        )
+        self.tenant_id = tenant_id or env_config.get("tenant_id", "default")
+        self.environment = environment or env_config.get("environment", "prod")
+        self.service = service or env_config.get("project_name", "kalibr-app")
+        self.secret = secret
+
+        # Workflow configuration (can override env)
+        if workflow_id:
+            env_config["workflow_id"] = workflow_id
+        if workflow_version:
+            env_config["workflow_version"] = workflow_version
+        self.batch_size = batch_size
+        self.flush_interval = flush_interval
+
+        # Extract workflow and runtime fields from config
+        self.workflow_id = env_config.get("workflow_id", "default-workflow")
+        self.workflow_version = env_config.get("workflow_version", "1.0")
+        self.sandbox_id = env_config.get("sandbox_id", "local")
+        self.runtime_env = env_config.get("runtime_env", "local")
+        self.parent_trace_id = env_config.get("parent_trace_id")
+
+        # Create tracer instance
+        self.tracer = Tracer(
+            tenant_id=self.tenant_id,
+            environment=self.environment,
+            service=self.service,
+            workflow_id=self.workflow_id,
+            workflow_version=self.workflow_version,
+            sandbox_id=self.sandbox_id,
+            runtime_env=self.runtime_env,
+            parent_trace_id=self.parent_trace_id,
+        )
+
+        # Event queue
+        self.queue = queue.Queue(maxsize=max_queue_size)
+        self.max_queue_size = max_queue_size
+
+        # HTTP client
+        self.client = httpx.Client(timeout=10.0)
+
+        # Background flusher thread
+        self._shutdown = False
+        self._flush_thread = threading.Thread(target=self._flush_loop, daemon=True)
+        self._flush_thread.start()
+
+        # Register cleanup
+        atexit.register(self.shutdown)
+
+        log_success(
+            f"Kalibr client initialized: {self.tenant_id} @ {self.environment} (service: {self.service})"
+        )
+
+    def _flush_loop(self):
+        """Background thread to flush events periodically."""
+        batch = []
+        last_flush = time.time()
+
+        while not self._shutdown:
+            try:
+                # Get event with timeout
+                try:
+                    event = self.queue.get(timeout=0.1)
+                    batch.append(event)
+                except queue.Empty:
+                    pass
+
+                # Flush if batch is full or interval elapsed
+                now = time.time()
+                should_flush = len(batch) >= self.batch_size or (
+                    batch and now - last_flush >= self.flush_interval
+                )
+
+                if should_flush:
+                    self._send_batch(batch)
+                    batch = []
+                    last_flush = now
+
+            except Exception as e:
+                log_warning(f"Flush loop error: {e}")
+
+        # Final flush on shutdown
+        if batch:
+            self._send_batch(batch)
+
+    def _send_batch(self, batch: list):
+        """Send batch to collector with validation."""
+        if not batch:
+            return
+
+        try:
+            # Validate events
+            valid_events = []
+            for event in batch:
+                if validate_event(event):
+                    valid_events.append(event)
+                else:
+                    log_warning(f"Invalid event skipped: {event.get('span_id', 'unknown')}")
+
+            if not valid_events:
+                log_warning("No valid events in batch")
+                return
+
+            # ✅ Fixed Bug 2: Send as JSON dict instead of NDJSON string
+            # Backend expects: {"events": [event_dict]}
+            payload = {"events": valid_events}
+
+            # Create HMAC signature if secret provided
+            headers = {"X-API-Key": self.api_key} if self.api_key else {}
+            if self.secret:
+                # Sign the JSON payload
+                body = serialize_event(payload).encode("utf-8")
+                signature = hmac.new(self.secret.encode(), body, hashlib.sha256).hexdigest()
+                headers["X-Signature"] = signature
+
+            # Send request with json parameter (automatically serializes dict)
+            response = self.client.post(
+                self.endpoint,
+                json=payload,  # ✅ Sends as JSON object, not string
+                headers=headers,
+            )
+            response.raise_for_status()
+
+            result = response.json()
+            accepted = result.get("accepted", 0)
+            rejected = result.get("rejected", 0)
+
+            if rejected > 0:
+                log_warning(f"Batch: {accepted} accepted, {rejected} rejected")
+
+        except Exception as e:
+            log_error(f"Batch send failed: {e}")
+
+    def _enqueue(self, event: Dict):
+        """Add event to queue (drop oldest if full)."""
+        try:
+            self.queue.put_nowait(event)
+        except queue.Full:
+            # Drop oldest event
+            try:
+                self.queue.get_nowait()
+                self.queue.put_nowait(event)
+            except:
+                pass
+
+    def shutdown(self):
+        """Shutdown client and flush remaining events."""
+        if self._shutdown:
+            return
+
+        log_info("Shutting down Kalibr client...")
+        self._shutdown = True
+
+        # Wait for flush thread
+        if self._flush_thread.is_alive():
+            self._flush_thread.join(timeout=5.0)
+
+        # Close HTTP client
+        self.client.close()
+        log_success("Kalibr client shutdown complete")
+
+    def trace(
+        self,
+        operation: str = "model_call",
+        vendor: str = "unknown",
+        model: str = "unknown",
+        endpoint: Optional[str] = None,
+    ):
+        """Decorator to trace function calls using new Phase 1 tracer.
+
+        Args:
+            operation: Operation type (chat_completion, embedding, etc.)
+            vendor: Vendor name (openai, anthropic, etc.)
+            model: Model identifier (gpt-4, claude-3-sonnet, etc.)
+            endpoint: API endpoint or function name
+
+        Example:
+            @kalibr.trace(operation="chat_completion", vendor="openai", model="gpt-4")
+            def call_openai(prompt):
+                return openai.chat.completions.create(...)
+        """
+        # Create decorator using the new tracer system
+        trace_decorator = create_trace_decorator(self.tracer)
+        decorator_func = trace_decorator(
+            operation=operation, vendor=vendor, model=model, endpoint=endpoint
+        )
+
+        # Wrap to enqueue events from tracer
+        def enhanced_decorator(func: Callable) -> Callable:
+            wrapped = decorator_func(func)
+
+            def wrapper(*args, **kwargs):
+                from .context import trace_context
+
+                # Clear events before call
+                ctx = trace_context.get()
+                ctx["events"] = []
+                trace_context.set(ctx)
+
+                try:
+                    result = wrapped(*args, **kwargs)
+
+                    # Enqueue events created by tracer
+                    ctx = trace_context.get()
+                    events = ctx.get("events", [])
+                    for event in events:
+                        self._enqueue(event)
+
+                    return result
+
+                except Exception as e:
+                    # Enqueue error events too
+                    ctx = trace_context.get()
+                    events = ctx.get("events", [])
+                    for event in events:
+                        self._enqueue(event)
+                    raise
+
+            return wrapper
+
+        return enhanced_decorator

kalibr/collector.py

@@ -0,0 +1,173 @@
+"""
+OpenTelemetry Collector Setup
+
+Configures OpenTelemetry tracer provider with multiple exporters:
+1. OTLP exporter for sending to OpenTelemetry collectors
+2. File exporter for local JSONL fallback
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Optional
+
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import (
+    BatchSpanProcessor,
+    ConsoleSpanExporter,
+    SpanExporter,
+    SpanExportResult,
+)
+
+try:
+    from opentelemetry.sdk.trace import ReadableSpan
+except ImportError:
+    ReadableSpan = None
+
+
+class FileSpanExporter(SpanExporter):
+    """Export spans to a JSONL file"""
+
+    def __init__(self, file_path: str = "/tmp/kalibr_otel_spans.jsonl"):
+        self.file_path = Path(file_path)
+        self.file_path.parent.mkdir(parents=True, exist_ok=True)
+
+    def export(self, spans) -> SpanExportResult:
+        """Export spans to JSONL file"""
+        try:
+            with open(self.file_path, "a") as f:
+                for span in spans:
+                    span_dict = self._span_to_dict(span)
+                    f.write(json.dumps(span_dict) + "\n")
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            print(f"❌ Failed to export spans to file: {e}")
+            return SpanExportResult.FAILURE
+
+    def shutdown(self):
+        """Shutdown the exporter"""
+        pass
+
+    def _span_to_dict(self, span) -> dict:
+        """Convert span to dictionary format"""
+        return {
+            "trace_id": format(span.context.trace_id, "032x"),
+            "span_id": format(span.context.span_id, "016x"),
+            "parent_span_id": format(span.parent.span_id, "016x") if span.parent else None,
+            "name": span.name,
+            "kind": span.kind.name if hasattr(span.kind, "name") else str(span.kind),
+            "start_time_unix_nano": span.start_time,
+            "end_time_unix_nano": span.end_time,
+            "attributes": dict(span.attributes) if span.attributes else {},
+            "status": {
+                "code": (
+                    span.status.status_code.name
+                    if hasattr(span.status.status_code, "name")
+                    else str(span.status.status_code)
+                ),
+                "description": getattr(span.status, "description", ""),
+            },
+            "events": [
+                {
+                    "name": event.name,
+                    "timestamp": event.timestamp,
+                    "attributes": dict(event.attributes) if event.attributes else {},
+                }
+                for event in (span.events or [])
+            ],
+        }
+
+
+_tracer_provider: Optional[TracerProvider] = None
+_is_configured = False
+
+
+def setup_collector(
+    service_name: str = "kalibr",
+    otlp_endpoint: Optional[str] = None,
+    file_export: bool = True,
+    console_export: bool = False,
+) -> TracerProvider:
+    """
+    Setup OpenTelemetry collector with multiple exporters
+
+    Args:
+        service_name: Service name for the tracer provider
+        otlp_endpoint: OTLP collector endpoint (e.g., "http://localhost:4317")
+                       If None, reads from OTEL_EXPORTER_OTLP_ENDPOINT env var
+        file_export: Whether to export spans to local JSONL file
+        console_export: Whether to export spans to console (for debugging)
+
+    Returns:
+        Configured TracerProvider instance
+    """
+    global _tracer_provider, _is_configured
+
+    if _is_configured and _tracer_provider:
+        return _tracer_provider
+
+    # Create resource with service name
+    resource = Resource(attributes={SERVICE_NAME: service_name})
+
+    # Create tracer provider
+    provider = TracerProvider(resource=resource)
+
+    # Add OTLP exporter if endpoint is configured
+    otlp_endpoint = otlp_endpoint or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
+    if otlp_endpoint:
+        try:
+            otlp_exporter = OTLPSpanExporter(endpoint=otlp_endpoint)
+            provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
+            print(f"✅ OTLP exporter configured: {otlp_endpoint}")
+        except Exception as e:
+            print(f"⚠️  Failed to configure OTLP exporter: {e}")
+
+    # Add file exporter for local fallback
+    if file_export:
+        try:
+            file_exporter = FileSpanExporter("/tmp/kalibr_otel_spans.jsonl")
+            provider.add_span_processor(BatchSpanProcessor(file_exporter))
+            print("✅ File exporter configured: /tmp/kalibr_otel_spans.jsonl")
+        except Exception as e:
+            print(f"⚠️  Failed to configure file exporter: {e}")
+
+    # Add console exporter for debugging
+    if console_export:
+        try:
+            console_exporter = ConsoleSpanExporter()
+            provider.add_span_processor(BatchSpanProcessor(console_exporter))
+            print("✅ Console exporter configured")
+        except Exception as e:
+            print(f"⚠️  Failed to configure console exporter: {e}")
+
+    # Set as global tracer provider
+    trace.set_tracer_provider(provider)
+
+    _tracer_provider = provider
+    _is_configured = True
+
+    return provider
+
+
+def get_tracer_provider() -> Optional[TracerProvider]:
+    """Get the current tracer provider"""
+    return _tracer_provider
+
+
+def is_configured() -> bool:
+    """Check if collector is configured"""
+    return _is_configured
+
+
+def shutdown_collector():
+    """Shutdown the tracer provider and flush all spans"""
+    global _tracer_provider, _is_configured
+
+    if _tracer_provider:
+        _tracer_provider.shutdown()
+        _tracer_provider = None
+        _is_configured = False
+        print("✅ Tracer provider shutdown")

kalibr/context.py

@@ -0,0 +1,132 @@
+"""
+Context propagation for distributed tracing.
+
+Phase 3: Enhanced with OpenTelemetry context propagation for linking
+HTTP requests to SDK calls (OpenAI, Anthropic, Google).
+"""
+
+import random
+import string
+import uuid
+from contextvars import ContextVar
+from typing import Dict, Optional
+
+from opentelemetry import trace
+from opentelemetry.trace import Span
+
+# Thread-local context for trace and span propagation
+trace_context: ContextVar[Dict] = ContextVar("trace_context", default={})
+
+# Phase 3: OpenTelemetry context for HTTP→SDK linking
+_otel_request_context: ContextVar[Optional[Dict]] = ContextVar("otel_request_context", default=None)
+
+
+def get_trace_id() -> Optional[str]:
+    """Get current trace ID from context."""
+    ctx = trace_context.get()
+    return ctx.get("trace_id")
+
+
+def get_parent_span_id() -> Optional[str]:
+    """Get parent span ID from context."""
+    ctx = trace_context.get()
+    span_stack = ctx.get("span_stack", [])
+    return span_stack[-1] if span_stack else None
+
+
+def new_span_id() -> str:
+    """Generate new span ID (UUIDv4 for consistency)."""
+    return str(uuid.uuid4())
+
+
+def new_trace_id() -> str:
+    """Generate new trace ID (UUID v4)."""
+    return str(uuid.uuid4())
+
+
+# ============================================================================
+# Phase 3: OpenTelemetry Context Propagation
+# ============================================================================
+
+
+def set_otel_request_context(context_token: str, trace_id: str, span_id: str):
+    """
+    Set OpenTelemetry request context for HTTP→SDK span linking.
+
+    Called by AutoTracerMiddleware to establish parent context for
+    all SDK calls within the HTTP request.
+
+    Args:
+        context_token: Kalibr context token for chaining
+        trace_id: OpenTelemetry trace ID (hex format)
+        span_id: OpenTelemetry span ID (hex format)
+    """
+    _otel_request_context.set(
+        {
+            "context_token": context_token,
+            "trace_id": trace_id,
+            "span_id": span_id,
+        }
+    )
+
+
+def get_otel_request_context() -> Optional[Dict]:
+    """
+    Get current OpenTelemetry request context.
+
+    Returns:
+        Dictionary with context_token, trace_id, span_id or None
+    """
+    return _otel_request_context.get()
+
+
+def clear_otel_request_context():
+    """Clear OpenTelemetry request context (called at end of request)"""
+    _otel_request_context.set(None)
+
+
+def get_current_otel_span() -> Optional[Span]:
+    """
+    Get the currently active OpenTelemetry span.
+
+    Returns:
+        Current span or None if no active span
+    """
+    return trace.get_current_span()
+
+
+def get_otel_trace_context() -> Dict:
+    """
+    Get current OpenTelemetry trace context from active span.
+
+    Returns:
+        Dictionary with trace_id, span_id, or empty dict
+    """
+    span = get_current_otel_span()
+    if span and span.get_span_context().is_valid:
+        ctx = span.get_span_context()
+        return {
+            "trace_id": format(ctx.trace_id, "032x"),
+            "span_id": format(ctx.span_id, "016x"),
+            "is_valid": True,
+        }
+    return {}
+
+
+def inject_kalibr_context_into_span(span: Span):
+    """
+    Inject Kalibr-specific context into OpenTelemetry span attributes.
+
+    This links the OTel span back to Kalibr's context_token and HTTP trace.
+
+    Args:
+        span: OpenTelemetry span to annotate
+    """
+    ctx = get_otel_request_context()
+    if ctx:
+        if ctx.get("context_token"):
+            span.set_attribute("kalibr.context_token", ctx["context_token"])
+        if ctx.get("trace_id"):
+            span.set_attribute("kalibr.http_trace_id", ctx["trace_id"])
+        if ctx.get("span_id"):
+            span.set_attribute("kalibr.http_span_id", ctx["span_id"])