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

kalibr/simple_tracer.py

@@ -0,0 +1,255 @@
+"""Simplified Kalibr tracer with direct event emission and capsule support.
+
+Usage:
+    from kalibr.simple_tracer import trace
+
+    @trace(operation="summarize", provider="openai", model="gpt-4o")
+    def my_function(text):
+        return call_llm(text)
+
+Capsule Usage (automatic when middleware is active):
+    from fastapi import FastAPI, Request
+    from kalibr.capsule_middleware import add_capsule_middleware
+    from kalibr import trace
+
+    app = FastAPI()
+    add_capsule_middleware(app)
+
+    @trace(operation="chat", provider="openai", model="gpt-4o")
+    def process_request(request: Request, prompt: str):
+        # Capsule automatically updated with this hop
+        return llm_call(prompt)
+"""
+
+import json
+import os
+import random
+import string
+import time
+import uuid
+from datetime import datetime, timezone
+from functools import wraps
+from typing import Callable, Optional
+
+try:
+    import requests
+except ImportError:
+    print("[Kalibr SDK] ⚠️  requests not installed, install with: pip install requests")
+    requests = None
+
+
+def generate_span_id() -> str:
+    """Generate UUIDv4 span ID for consistency."""
+    return str(uuid.uuid4())
+
+
+def send_event(payload: dict):
+    """Send event directly to collector.
+
+    Args:
+        payload: Event data dict
+    """
+    if not requests:
+        print("[Kalibr SDK] ❌ requests library not available")
+        return
+
+    url = os.getenv("KALIBR_COLLECTOR_URL", "http://localhost:8001/api/ingest")
+    api_key = os.getenv("KALIBR_API_KEY", "test_key_12345")
+
+    format_pref = os.getenv("KALIBR_COLLECTOR_FORMAT", "ndjson").lower()
+    use_json_envelope = format_pref == "json"
+
+    headers = {"X-API-Key": api_key}
+    if use_json_envelope:
+        headers["Content-Type"] = "application/json"
+        body_cfg = {"events": [payload]}
+    else:
+        headers["Content-Type"] = "application/x-ndjson"
+        body_cfg = "\n".join(json.dumps(evt) for evt in [payload]) + "\n"
+
+    try:
+        if use_json_envelope:
+            response = requests.post(url, headers=headers, json=body_cfg, timeout=5)
+        else:
+            response = requests.post(url, headers=headers, data=body_cfg, timeout=5)
+        if not response.ok:
+            print(
+                f"[Kalibr SDK] ❌ Collector rejected event: {response.status_code} - {response.text}"
+            )
+        else:
+            duration_ms = payload.get("duration_ms") or payload.get("latency_ms") or 0
+            total_cost = payload.get("total_cost_usd") or payload.get("cost_usd") or 0.0
+            print(
+                f"[Kalibr SDK] ✅ Event sent: {payload.get('operation','event')} ({duration_ms}ms, ${total_cost:.6f})"
+            )
+    except Exception as e:
+        print(f"[Kalibr SDK] ❌ Failed to send event: {e}")
+
+
+def trace(
+    operation: str, provider: str, model: str, input_tokens: int = None, output_tokens: int = None
+):
+    """Decorator to trace function calls with full telemetry.
+
+    Captures:
+    - Duration (ms)
+    - Tokens (estimated if not provided)
+    - Cost (USD)
+    - Errors
+    - Runtime metadata
+
+    Args:
+        operation: Operation type (e.g., "summarize", "refine", "analyze")
+        provider: LLM provider (e.g., "openai", "anthropic", "google")
+        model: Model name (e.g., "gpt-4o", "claude-3-sonnet")
+        input_tokens: Input token count (optional, will estimate)
+        output_tokens: Output token count (optional, will estimate)
+
+    Example:
+        @trace(operation="summarize", provider="openai", model="gpt-4o")
+        def summarize_text(text: str) -> str:
+            return openai.chat.completions.create(...)
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Generate IDs
+            trace_id = str(uuid.uuid4())
+            span_id = generate_span_id()  # Base62, 16 chars
+            parent_span_id = kwargs.pop("parent_span_id", None)  # None or base62 string
+
+            # Load environment config
+            tenant_id = os.getenv("KALIBR_TENANT_ID", "emergent")
+            workflow_id = os.getenv("KALIBR_WORKFLOW_ID", "multi_agent_demo")
+            sandbox_id = os.getenv("SANDBOX_ID", "vercel_vm_001")
+            runtime_env = os.getenv("RUNTIME_ENV", "vercel_vm")
+
+            # Start timing
+            start_time = time.time()
+
+            # Execute function
+            result = None
+            status = "success"
+            error_type = None
+            error_message = None
+            exception_to_raise = None
+
+            try:
+                result = func(*args, **kwargs)
+            except Exception as e:
+                status = "error"
+                error_type = type(e).__name__
+                error_message = str(e)
+                exception_to_raise = e
+                print(f"[Kalibr SDK] ⚠️  Error in {func.__name__}: {error_type} - {error_message}")
+
+            # End timing
+            end_time = time.time()
+            duration_ms = int((end_time - start_time) * 1000)
+
+            # Token estimation
+            actual_input_tokens = input_tokens or kwargs.get("input_tokens", 1000)
+            actual_output_tokens = output_tokens or kwargs.get("output_tokens", 500)
+
+            # Cost calculation (simplified pricing)
+            # OpenAI GPT-4o: ~$2.50/1M input, ~$10/1M output
+            # Anthropic Claude-3-Sonnet: ~$3/1M input, ~$15/1M output
+            pricing_map = {
+                "openai": {"gpt-4o": 0.00000250, "gpt-4": 0.00003000},
+                "anthropic": {"claude-3-sonnet": 0.00000300, "claude-3-opus": 0.00001500},
+                "google": {"gemini-pro": 0.00000125},
+            }
+
+            # Get unit price
+            provider_pricing = pricing_map.get(provider, {})
+            unit_price_usd = provider_pricing.get(model, 0.00002000)  # Default $0.02/1M
+
+            # Calculate total cost
+            total_cost_usd = (actual_input_tokens + actual_output_tokens) * unit_price_usd
+
+            # Build payload
+            payload = {
+                "schema_version": "1.0",
+                "trace_id": trace_id,
+                "span_id": span_id,
+                "parent_id": parent_span_id,  # Note: parent_id not parent_span_id
+                "tenant_id": tenant_id,
+                "workflow_id": workflow_id,
+                "sandbox_id": sandbox_id,
+                "runtime_env": runtime_env,
+                "provider": provider,
+                "model_name": model,
+                "model_id": model,  # For backward compatibility
+                "operation": operation,
+                "endpoint": func.__name__,
+                "input_tokens": actual_input_tokens,
+                "output_tokens": actual_output_tokens,
+                "duration_ms": duration_ms,
+                "latency_ms": duration_ms,  # For backward compatibility
+                "unit_price_usd": unit_price_usd,
+                "total_cost_usd": round(total_cost_usd, 6),
+                "cost_usd": round(total_cost_usd, 6),  # For backward compatibility
+                "status": status,
+                "error_type": error_type,
+                "error_message": error_message,
+                "timestamp": datetime.now(timezone.utc).isoformat(),
+                "ts_start": datetime.now(timezone.utc).isoformat(),
+                "ts_end": datetime.now(timezone.utc).isoformat(),
+                "environment": os.getenv("KALIBR_ENVIRONMENT", "prod"),
+                "service": os.getenv("KALIBR_SERVICE", "kalibr-app"),
+                "vendor": provider,  # v0.2 compatibility
+                "data_class": "economic",
+            }
+
+            # Send event to collector
+            send_event(payload)
+
+            # ========================================================================
+            # PHASE 6: Append hop to capsule if available (from middleware)
+            # ========================================================================
+            try:
+                # Check if we're in a FastAPI context with capsule middleware
+                from starlette.requests import Request
+
+                # Try to get capsule from request context (if available)
+                capsule = kwargs.get("__kalibr_capsule")
+
+                if capsule:
+                    # Create hop from trace data
+                    hop = {
+                        "provider": provider,
+                        "operation": operation,
+                        "model": model,
+                        "duration_ms": duration_ms,
+                        "status": status,
+                        "cost_usd": round(total_cost_usd, 6),
+                        "input_tokens": actual_input_tokens,
+                        "output_tokens": actual_output_tokens,
+                    }
+
+                    if error_type:
+                        hop["error_type"] = error_type
+
+                    # Add agent name if available
+                    agent_name = os.getenv("KALIBR_AGENT_NAME", func.__name__)
+                    hop["agent_name"] = agent_name
+
+                    capsule.append_hop(hop)
+                    print(
+                        f"[Kalibr SDK] 📦 Appended hop to capsule: {operation} ({provider}/{model})"
+                    )
+            except Exception as e:
+                # Capsule update is non-critical, just log
+                print(f"[Kalibr SDK] ⚠️  Could not update capsule: {e}")
+            # ========================================================================
+
+            # Re-raise exception if there was one
+            if exception_to_raise:
+                raise exception_to_raise
+
+            return result
+
+        return wrapper
+
+    return decorator

kalibr/tokens.py

@@ -0,0 +1,52 @@
+"""Token counting utilities."""
+
+from typing import Optional
+
+import tiktoken
+
+# Cache for tokenizer instances
+_tokenizer_cache = {}
+
+
+def count_tokens(text: str, model_id: str) -> int:
+    """Count tokens for given text and model.
+
+    Args:
+        text: Input text
+        model_id: Model identifier
+
+    Returns:
+        Token count (approximate)
+    """
+    if not text:
+        return 0
+
+    # Try to get exact tokenizer for OpenAI models
+    if "gpt" in model_id.lower():
+        try:
+            encoding = get_openai_encoding(model_id)
+            return len(encoding.encode(text))
+        except Exception:
+            pass
+
+    # Fallback: approximate (1 token ~= 4 chars)
+    return len(text) // 4
+
+
+def get_openai_encoding(model_id: str):
+    """Get tiktoken encoding for OpenAI model."""
+    if model_id in _tokenizer_cache:
+        return _tokenizer_cache[model_id]
+
+    try:
+        # Map model to encoding
+        if "gpt-4" in model_id or "gpt-3.5" in model_id:
+            encoding = tiktoken.encoding_for_model("gpt-4")
+        else:
+            encoding = tiktoken.get_encoding("cl100k_base")
+
+        _tokenizer_cache[model_id] = encoding
+        return encoding
+    except Exception as e:
+        print(f"⚠️  Failed to load tokenizer for {model_id}: {e}")
+        raise

kalibr/trace_capsule.py

@@ -0,0 +1,296 @@
+"""
+Kalibr Trace Capsule - Portable JSON payload for cross-MCP trace propagation.
+
+A capsule carries observability context across agent hops, maintaining a rolling
+window of recent operations and aggregate metrics.
+
+Usage:
+    from kalibr.trace_capsule import TraceCapsule
+
+    # Create new capsule
+    capsule = TraceCapsule()
+
+    # Append hop
+    capsule.append_hop({
+        "provider": "openai",
+        "operation": "summarize",
+        "model": "gpt-4o",
+        "duration_ms": 1200,
+        "status": "success",
+        "cost_usd": 0.005
+    })
+
+    # Serialize for HTTP header
+    header_value = capsule.to_json()
+
+    # Deserialize from header
+    received_capsule = TraceCapsule.from_json(header_value)
+"""
+
+import json
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+
+
+class TraceCapsule:
+    """Portable JSON payload containing rolling trace history.
+
+    Attributes:
+        trace_id: Unique identifier for the trace chain
+        timestamp: ISO 8601 timestamp of last update
+        aggregate_cost_usd: Cumulative cost across all hops
+        aggregate_latency_ms: Cumulative latency across all hops
+        last_n_hops: Rolling window of last N hops (max 5)
+        tenant_id: Optional tenant identifier
+        workflow_id: Optional workflow identifier
+        metadata: Optional custom metadata
+    """
+
+    MAX_HOPS = 5  # Keep payload compact for HTTP headers
+
+    def __init__(
+        self,
+        trace_id: Optional[str] = None,
+        last_n_hops: Optional[List[Dict[str, Any]]] = None,
+        aggregate_cost_usd: float = 0.0,
+        aggregate_latency_ms: float = 0.0,
+        tenant_id: Optional[str] = None,
+        workflow_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        context_token: Optional[str] = None,
+        parent_context_token: Optional[str] = None,
+    ):
+        """Initialize a new TraceCapsule.
+
+        Args:
+            trace_id: Unique trace identifier (generates UUID if not provided)
+            last_n_hops: Existing hop history
+            aggregate_cost_usd: Starting cumulative cost
+            aggregate_latency_ms: Starting cumulative latency
+            tenant_id: Tenant identifier
+            workflow_id: Workflow identifier
+            metadata: Custom metadata
+            context_token: Context token for this runtime session (Phase 3C)
+            parent_context_token: Parent runtime's context token (Phase 3C)
+        """
+        self.trace_id = trace_id or str(uuid.uuid4())
+        self.timestamp = datetime.now(timezone.utc).isoformat()
+        self.aggregate_cost_usd = aggregate_cost_usd
+        self.aggregate_latency_ms = aggregate_latency_ms
+        self.last_n_hops: List[Dict[str, Any]] = last_n_hops or []
+        self.tenant_id = tenant_id
+        self.workflow_id = workflow_id
+        self.metadata = metadata or {}
+        # Phase 3C: Context token propagation (keep as UUID for consistency)
+        self.context_token = context_token or str(uuid.uuid4())
+        self.parent_context_token = parent_context_token
+
+    def append_hop(self, hop: Dict[str, Any]) -> None:
+        """Append a new hop to the capsule.
+
+        Maintains a rolling window of last N hops to keep payload compact.
+        Updates aggregate metrics automatically.
+
+        Args:
+            hop: Dictionary containing hop metadata
+                Required fields: provider, operation, status
+                Optional fields: model, duration_ms, cost_usd, input_tokens,
+                                output_tokens, error_type, agent_name
+
+        Example:
+            capsule.append_hop({
+                "provider": "openai",
+                "operation": "chat_completion",
+                "model": "gpt-4o",
+                "duration_ms": 1200,
+                "status": "success",
+                "cost_usd": 0.005,
+                "input_tokens": 150,
+                "output_tokens": 75,
+                "agent_name": "code-writer"
+            })
+        """
+        # Add hop_index
+        hop["hop_index"] = len(self.last_n_hops)
+
+        # Append to history
+        self.last_n_hops.append(hop)
+
+        # Maintain rolling window (keep last N hops)
+        if len(self.last_n_hops) > self.MAX_HOPS:
+            self.last_n_hops.pop(0)
+
+        # Update aggregates
+        self.aggregate_cost_usd += hop.get("cost_usd", 0.0)
+        self.aggregate_latency_ms += hop.get("duration_ms", 0.0)
+
+        # Update timestamp
+        self.timestamp = datetime.now(timezone.utc).isoformat()
+
+    def get_last_hop(self) -> Optional[Dict[str, Any]]:
+        """Get the most recent hop.
+
+        Returns:
+            Last hop dictionary or None if no hops exist
+        """
+        return self.last_n_hops[-1] if self.last_n_hops else None
+
+    def get_hop_count(self) -> int:
+        """Get total number of hops in capsule.
+
+        Returns:
+            Number of hops in the rolling window
+        """
+        return len(self.last_n_hops)
+
+    def to_json(self) -> str:
+        """Serialize capsule to JSON string for HTTP header transmission.
+
+        Returns:
+            Compact JSON string representation
+        """
+        data = {
+            "trace_id": self.trace_id,
+            "timestamp": self.timestamp,
+            "aggregate_cost_usd": round(self.aggregate_cost_usd, 6),
+            "aggregate_latency_ms": round(self.aggregate_latency_ms, 2),
+            "last_n_hops": self.last_n_hops,
+        }
+
+        # Add optional fields only if present
+        if self.tenant_id:
+            data["tenant_id"] = self.tenant_id
+        if self.workflow_id:
+            data["workflow_id"] = self.workflow_id
+        if self.metadata:
+            data["metadata"] = self.metadata
+
+        # Phase 3C: Include context tokens
+        if self.context_token:
+            data["context_token"] = self.context_token
+        if self.parent_context_token:
+            data["parent_context_token"] = self.parent_context_token
+
+        return json.dumps(data, separators=(",", ":"))  # Compact JSON
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert capsule to dictionary.
+
+        Returns:
+            Dictionary representation of capsule
+        """
+        data = {
+            "trace_id": self.trace_id,
+            "timestamp": self.timestamp,
+            "aggregate_cost_usd": self.aggregate_cost_usd,
+            "aggregate_latency_ms": self.aggregate_latency_ms,
+            "last_n_hops": self.last_n_hops,
+        }
+
+        if self.tenant_id:
+            data["tenant_id"] = self.tenant_id
+        if self.workflow_id:
+            data["workflow_id"] = self.workflow_id
+        if self.metadata:
+            data["metadata"] = self.metadata
+
+        # Phase 3C: Include context tokens
+        if self.context_token:
+            data["context_token"] = self.context_token
+        if self.parent_context_token:
+            data["parent_context_token"] = self.parent_context_token
+
+        return data
+
+    @classmethod
+    def from_json(cls, s: str) -> "TraceCapsule":
+        """Deserialize capsule from JSON string.
+
+        Args:
+            s: JSON string from HTTP header
+
+        Returns:
+            TraceCapsule instance
+
+        Raises:
+            json.JSONDecodeError: If JSON is invalid
+            KeyError: If required fields are missing
+        """
+        try:
+            data = json.loads(s)
+            return cls(
+                trace_id=data.get("trace_id"),
+                last_n_hops=data.get("last_n_hops", []),
+                aggregate_cost_usd=data.get("aggregate_cost_usd", 0.0),
+                aggregate_latency_ms=data.get("aggregate_latency_ms", 0.0),
+                tenant_id=data.get("tenant_id"),
+                workflow_id=data.get("workflow_id"),
+                metadata=data.get("metadata"),
+                # Phase 3C: Context token propagation
+                context_token=data.get("context_token"),
+                parent_context_token=data.get("parent_context_token"),
+            )
+        except json.JSONDecodeError as e:
+            # Return empty capsule if parsing fails (graceful degradation)
+            print(f"⚠️ Failed to parse TraceCapsule: {e}")
+            return cls()
+        except Exception as e:
+            print(f"⚠️ Error creating TraceCapsule: {e}")
+            return cls()
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "TraceCapsule":
+        """Create capsule from dictionary.
+
+        Args:
+            data: Dictionary containing capsule data
+
+        Returns:
+            TraceCapsule instance
+        """
+        return cls(
+            trace_id=data.get("trace_id"),
+            last_n_hops=data.get("last_n_hops", []),
+            aggregate_cost_usd=data.get("aggregate_cost_usd", 0.0),
+            aggregate_latency_ms=data.get("aggregate_latency_ms", 0.0),
+            tenant_id=data.get("tenant_id"),
+            workflow_id=data.get("workflow_id"),
+            metadata=data.get("metadata"),
+        )
+
+    def __repr__(self) -> str:
+        """String representation of capsule."""
+        return (
+            f"TraceCapsule(trace_id={self.trace_id}, "
+            f"hops={len(self.last_n_hops)}, "
+            f"cost=${self.aggregate_cost_usd:.6f}, "
+            f"latency={self.aggregate_latency_ms:.2f}ms)"
+        )
+
+    def __str__(self) -> str:
+        """Human-readable string representation."""
+        hops_summary = ", ".join(
+            [f"{hop.get('provider', '?')}/{hop.get('operation', '?')}" for hop in self.last_n_hops]
+        )
+        return (
+            f"TraceCapsule[{self.trace_id}]: "
+            f"{len(self.last_n_hops)} hops ({hops_summary}), "
+            f"${self.aggregate_cost_usd:.4f}, "
+            f"{self.aggregate_latency_ms:.0f}ms"
+        )
+
+
+# Convenience function for FastAPI integration
+def get_or_create_capsule(header_value: Optional[str] = None) -> TraceCapsule:
+    """Get existing capsule from header or create new one.
+
+    Args:
+        header_value: Value of X-Kalibr-Capsule header
+
+    Returns:
+        TraceCapsule instance
+    """
+    if header_value:
+        return TraceCapsule.from_json(header_value)
+    return TraceCapsule()