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

kalibr/instrumentation/base.py

@@ -0,0 +1,108 @@
+"""
+Base instrumentation class for LLM SDKs
+
+Provides common functionality for monkey-patching LLM SDKs and
+emitting OpenTelemetry-compatible spans.
+"""
+
+import time
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+from opentelemetry import trace
+from opentelemetry.trace import SpanKind, Status, StatusCode
+
+
+class BaseInstrumentation(ABC):
+    """Base class for LLM SDK instrumentation"""
+
+    def __init__(self, tracer_name: str):
+        self.tracer = trace.get_tracer(tracer_name)
+        self._is_instrumented = False
+
+    @abstractmethod
+    def instrument(self) -> bool:
+        """
+        Apply monkey-patching to instrument the SDK
+
+        Returns:
+            bool: True if instrumentation succeeded, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    def uninstrument(self) -> bool:
+        """
+        Remove monkey-patching to restore original SDK behavior
+
+        Returns:
+            bool: True if uninstrumentation succeeded, False otherwise
+        """
+        pass
+
+    @property
+    def is_instrumented(self) -> bool:
+        """Check if SDK is currently instrumented"""
+        return self._is_instrumented
+
+    def create_span(self, name: str, attributes: Dict[str, Any], kind: SpanKind = SpanKind.CLIENT):
+        """
+        Create an OpenTelemetry span with standardized attributes
+
+        Args:
+            name: Span name (e.g., "openai.chat.completions.create")
+            attributes: Span attributes following OTel semantic conventions
+            kind: Span kind (default: CLIENT for LLM API calls)
+
+        Returns:
+            Context manager for the span
+        """
+        return self.tracer.start_as_current_span(name, kind=kind, attributes=attributes)
+
+    @staticmethod
+    def set_error(span: trace.Span, error: Exception) -> None:
+        """
+        Set error status and attributes on a span
+
+        Args:
+            span: The span to update
+            error: The exception that occurred
+        """
+        span.set_status(Status(StatusCode.ERROR))
+        span.set_attribute("error.type", type(error).__name__)
+        span.set_attribute("error.message", str(error))
+        span.record_exception(error)
+
+
+class BaseCostAdapter(ABC):
+    """Base class for cost calculation adapters"""
+
+    PRICING: Dict[str, Dict[str, float]] = {}
+
+    @abstractmethod
+    def calculate_cost(self, model: str, usage: Dict[str, int]) -> float:
+        """
+        Calculate cost in USD for an LLM API call
+
+        Args:
+            model: Model identifier (e.g., "gpt-4")
+            usage: Token usage dictionary with prompt_tokens, completion_tokens
+
+        Returns:
+            Cost in USD (rounded to 6 decimal places)
+        """
+        pass
+
+    def get_pricing(self, model: str) -> Optional[Dict[str, float]]:
+        """
+        Get pricing for a specific model
+
+        Args:
+            model: Model identifier
+
+        Returns:
+            Dictionary with "input" and "output" prices per 1K tokens,
+            or None if model not found
+        """
+        return self.PRICING.get(model)

kalibr/instrumentation/google_instr.py

@@ -0,0 +1,281 @@
+"""
+Google Generative AI SDK Instrumentation
+
+Monkey-patches the Google Generative AI SDK to automatically emit OpenTelemetry spans
+for all content generation API calls.
+"""
+
+import time
+from functools import wraps
+from typing import Any, Dict, Optional
+
+from opentelemetry.trace import SpanKind
+
+from .base import BaseCostAdapter, BaseInstrumentation
+
+
+class GoogleCostAdapter(BaseCostAdapter):
+    """Cost calculation adapter for Google Generative AI models"""
+
+    # Pricing per 1K tokens (USD) - Updated November 2025
+    PRICING = {
+        # Gemini 2.5 models
+        "gemini-2.5-pro": {"input": 0.00125, "output": 0.005},
+        "gemini-2.5-flash": {"input": 0.000075, "output": 0.0003},
+        # Gemini 2.0 models
+        "gemini-2.0-flash": {"input": 0.000075, "output": 0.0003},
+        "gemini-2.0-flash-thinking": {"input": 0.000075, "output": 0.0003},
+        # Gemini 1.5 models
+        "gemini-1.5-pro": {"input": 0.00125, "output": 0.005},
+        "gemini-1.5-flash": {"input": 0.000075, "output": 0.0003},
+        "gemini-1.5-flash-8b": {"input": 0.0000375, "output": 0.00015},
+        # Gemini 1.0 models
+        "gemini-1.0-pro": {"input": 0.0005, "output": 0.0015},
+        "gemini-pro": {"input": 0.0005, "output": 0.0015},  # Alias
+    }
+
+    def calculate_cost(self, model: str, usage: Dict[str, int]) -> float:
+        """Calculate cost in USD for a Google Generative AI API call"""
+        # Normalize model name
+        base_model = model.lower()
+
+        # Try exact match first
+        pricing = self.get_pricing(base_model)
+
+        # Try fuzzy matching for versioned models
+        if not pricing:
+            for known_model in self.PRICING.keys():
+                if known_model in base_model or base_model in known_model:
+                    pricing = self.PRICING[known_model]
+                    break
+
+        if not pricing:
+            # Default to Gemini 1.5 Pro pricing if unknown
+            pricing = {"input": 0.00125, "output": 0.005}
+
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+
+        input_cost = (prompt_tokens / 1000) * pricing["input"]
+        output_cost = (completion_tokens / 1000) * pricing["output"]
+
+        return round(input_cost + output_cost, 6)
+
+
+class GoogleInstrumentation(BaseInstrumentation):
+    """Instrumentation for Google Generative AI SDK"""
+
+    def __init__(self):
+        super().__init__("kalibr.google")
+        self._original_generate_content = None
+        self._original_async_generate_content = None
+        self.cost_adapter = GoogleCostAdapter()
+
+    def instrument(self) -> bool:
+        """Apply monkey-patching to Google Generative AI SDK"""
+        if self._is_instrumented:
+            return True
+
+        try:
+            import google.generativeai as genai
+            from google.generativeai.generative_models import GenerativeModel
+
+            # Patch sync method
+            if hasattr(GenerativeModel, "generate_content"):
+                self._original_generate_content = GenerativeModel.generate_content
+                GenerativeModel.generate_content = self._traced_generate_wrapper(
+                    GenerativeModel.generate_content
+                )
+
+            # Patch async method (if available)
+            if hasattr(GenerativeModel, "generate_content_async"):
+                self._original_async_generate_content = GenerativeModel.generate_content_async
+                GenerativeModel.generate_content_async = self._traced_async_generate_wrapper(
+                    GenerativeModel.generate_content_async
+                )
+
+            self._is_instrumented = True
+            return True
+
+        except ImportError:
+            print("⚠️  Google Generative AI SDK not installed, skipping instrumentation")
+            return False
+        except Exception as e:
+            print(f"❌ Failed to instrument Google Generative AI SDK: {e}")
+            return False
+
+    def uninstrument(self) -> bool:
+        """Remove monkey-patching from Google Generative AI SDK"""
+        if not self._is_instrumented:
+            return True
+
+        try:
+            import google.generativeai as genai
+            from google.generativeai.generative_models import GenerativeModel
+
+            # Restore sync method
+            if self._original_generate_content:
+                GenerativeModel.generate_content = self._original_generate_content
+
+            # Restore async method
+            if self._original_async_generate_content:
+                GenerativeModel.generate_content_async = self._original_async_generate_content
+
+            self._is_instrumented = False
+            return True
+
+        except Exception as e:
+            print(f"❌ Failed to uninstrument Google Generative AI SDK: {e}")
+            return False
+
+    def _traced_generate_wrapper(self, original_func):
+        """Wrapper for sync generate_content method"""
+
+        @wraps(original_func)
+        def wrapper(self_instance, *args, **kwargs):
+            # Extract model name from instance
+            model = getattr(self_instance, "_model_name", "unknown")
+
+            # Create span with initial attributes
+            with self.tracer.start_as_current_span(
+                "google.generativeai.generate_content",
+                kind=SpanKind.CLIENT,
+                attributes={
+                    "llm.vendor": "google",
+                    "llm.request.model": model,
+                    "llm.system": "google.generativeai",
+                },
+            ) as span:
+                start_time = time.time()
+
+                # Phase 3: Inject Kalibr context for HTTP→SDK linking
+                try:
+                    from kalibr.context import inject_kalibr_context_into_span
+
+                    inject_kalibr_context_into_span(span)
+                except Exception:
+                    pass  # Fail silently if context not available
+
+                try:
+                    # Call original method
+                    result = original_func(self_instance, *args, **kwargs)
+
+                    # Extract and set response metadata
+                    self._set_response_attributes(span, result, model, start_time)
+
+                    return result
+
+                except Exception as e:
+                    self.set_error(span, e)
+                    raise
+
+        return wrapper
+
+    def _traced_async_generate_wrapper(self, original_func):
+        """Wrapper for async generate_content method"""
+
+        @wraps(original_func)
+        async def wrapper(self_instance, *args, **kwargs):
+            # Extract model name from instance
+            model = getattr(self_instance, "_model_name", "unknown")
+
+            # Create span with initial attributes
+            with self.tracer.start_as_current_span(
+                "google.generativeai.generate_content",
+                kind=SpanKind.CLIENT,
+                attributes={
+                    "llm.vendor": "google",
+                    "llm.request.model": model,
+                    "llm.system": "google.generativeai",
+                },
+            ) as span:
+                start_time = time.time()
+
+                # Phase 3: Inject Kalibr context for HTTP→SDK linking
+                try:
+                    from kalibr.context import inject_kalibr_context_into_span
+
+                    inject_kalibr_context_into_span(span)
+                except Exception:
+                    pass  # Fail silently if context not available
+
+                try:
+                    # Call original async method
+                    result = await original_func(self_instance, *args, **kwargs)
+
+                    # Extract and set response metadata
+                    self._set_response_attributes(span, result, model, start_time)
+
+                    return result
+
+                except Exception as e:
+                    self.set_error(span, e)
+                    raise
+
+        return wrapper
+
+    def _set_response_attributes(self, span, result, model: str, start_time: float) -> None:
+        """Extract metadata from response and set span attributes"""
+        try:
+            # Model (from instance)
+            span.set_attribute("llm.response.model", model)
+
+            # Token usage
+            if hasattr(result, "usage_metadata") and result.usage_metadata:
+                usage = result.usage_metadata
+
+                prompt_tokens = getattr(usage, "prompt_token_count", 0)
+                completion_tokens = getattr(usage, "candidates_token_count", 0)
+                total_tokens = getattr(
+                    usage, "total_token_count", prompt_tokens + completion_tokens
+                )
+
+                span.set_attribute("llm.usage.prompt_tokens", prompt_tokens)
+                span.set_attribute("llm.usage.completion_tokens", completion_tokens)
+                span.set_attribute("llm.usage.total_tokens", total_tokens)
+
+                # Calculate cost
+                cost = self.cost_adapter.calculate_cost(
+                    model,
+                    {
+                        "prompt_tokens": prompt_tokens,
+                        "completion_tokens": completion_tokens,
+                    },
+                )
+                span.set_attribute("llm.cost_usd", cost)
+
+            # Latency
+            latency_ms = (time.time() - start_time) * 1000
+            span.set_attribute("llm.latency_ms", round(latency_ms, 2))
+
+            # Finish reason (if available)
+            if hasattr(result, "candidates") and result.candidates:
+                candidate = result.candidates[0]
+                if hasattr(candidate, "finish_reason"):
+                    span.set_attribute("llm.response.finish_reason", str(candidate.finish_reason))
+
+        except Exception as e:
+            # Don't fail the call if metadata extraction fails
+            span.set_attribute("llm.metadata_extraction_error", str(e))
+
+
+# Singleton instance
+_google_instrumentation = None
+
+
+def get_instrumentation() -> GoogleInstrumentation:
+    """Get or create the Google instrumentation singleton"""
+    global _google_instrumentation
+    if _google_instrumentation is None:
+        _google_instrumentation = GoogleInstrumentation()
+    return _google_instrumentation
+
+
+def instrument() -> bool:
+    """Instrument Google Generative AI SDK"""
+    return get_instrumentation().instrument()
+
+
+def uninstrument() -> bool:
+    """Uninstrument Google Generative AI SDK"""
+    return get_instrumentation().uninstrument()

kalibr/instrumentation/openai_instr.py

@@ -0,0 +1,265 @@
+"""
+OpenAI SDK Instrumentation
+
+Monkey-patches the OpenAI SDK to automatically emit OpenTelemetry spans
+for all chat completion API calls.
+"""
+
+import time
+from functools import wraps
+from typing import Any, Dict, Optional
+
+from opentelemetry.trace import SpanKind
+
+from .base import BaseCostAdapter, BaseInstrumentation
+
+
+class OpenAICostAdapter(BaseCostAdapter):
+    """Cost calculation adapter for OpenAI models"""
+
+    # Pricing per 1K tokens (USD) - Updated November 2025
+    PRICING = {
+        # GPT-5 models
+        "gpt-5": {"input": 0.005, "output": 0.015},
+        "gpt-5-turbo": {"input": 0.0025, "output": 0.0075},
+        # GPT-4 models
+        "gpt-4": {"input": 0.03, "output": 0.06},
+        "gpt-4-turbo": {"input": 0.01, "output": 0.03},
+        "gpt-4o": {"input": 0.0025, "output": 0.01},
+        "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
+        # GPT-3.5 models
+        "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
+        "gpt-3.5-turbo-16k": {"input": 0.001, "output": 0.002},
+    }
+
+    def calculate_cost(self, model: str, usage: Dict[str, int]) -> float:
+        """Calculate cost in USD for an OpenAI API call"""
+        # Normalize model name (remove version suffixes)
+        base_model = model.split("-2")[0]  # Remove date suffixes like -20240101
+
+        pricing = self.get_pricing(base_model)
+        if not pricing:
+            # Default to GPT-4 pricing if unknown
+            pricing = {"input": 0.03, "output": 0.06}
+
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+
+        input_cost = (prompt_tokens / 1000) * pricing["input"]
+        output_cost = (completion_tokens / 1000) * pricing["output"]
+
+        return round(input_cost + output_cost, 6)
+
+
+class OpenAIInstrumentation(BaseInstrumentation):
+    """Instrumentation for OpenAI SDK"""
+
+    def __init__(self):
+        super().__init__("kalibr.openai")
+        self._original_create = None
+        self._original_async_create = None
+        self.cost_adapter = OpenAICostAdapter()
+
+    def instrument(self) -> bool:
+        """Apply monkey-patching to OpenAI SDK"""
+        if self._is_instrumented:
+            return True
+
+        try:
+            import openai
+            from openai.resources.chat import completions
+
+            # Patch sync method
+            if hasattr(completions.Completions, "create"):
+                self._original_create = completions.Completions.create
+                completions.Completions.create = self._traced_create_wrapper(
+                    completions.Completions.create
+                )
+
+            # Patch async method
+            if hasattr(completions.AsyncCompletions, "create"):
+                self._original_async_create = completions.AsyncCompletions.create
+                completions.AsyncCompletions.create = self._traced_async_create_wrapper(
+                    completions.AsyncCompletions.create
+                )
+
+            self._is_instrumented = True
+            return True
+
+        except ImportError:
+            print("⚠️  OpenAI SDK not installed, skipping instrumentation")
+            return False
+        except Exception as e:
+            print(f"❌ Failed to instrument OpenAI SDK: {e}")
+            return False
+
+    def uninstrument(self) -> bool:
+        """Remove monkey-patching from OpenAI SDK"""
+        if not self._is_instrumented:
+            return True
+
+        try:
+            import openai
+            from openai.resources.chat import completions
+
+            # Restore sync method
+            if self._original_create:
+                completions.Completions.create = self._original_create
+
+            # Restore async method
+            if self._original_async_create:
+                completions.AsyncCompletions.create = self._original_async_create
+
+            self._is_instrumented = False
+            return True
+
+        except Exception as e:
+            print(f"❌ Failed to uninstrument OpenAI SDK: {e}")
+            return False
+
+    def _traced_create_wrapper(self, original_func):
+        """Wrapper for sync create method"""
+
+        @wraps(original_func)
+        def wrapper(self_instance, *args, **kwargs):
+            # Extract model from kwargs
+            model = kwargs.get("model", "unknown")
+
+            # Create span with initial attributes
+            with self.tracer.start_as_current_span(
+                "openai.chat.completions.create",
+                kind=SpanKind.CLIENT,
+                attributes={
+                    "llm.vendor": "openai",
+                    "llm.request.model": model,
+                    "llm.system": "openai",
+                },
+            ) as span:
+                start_time = time.time()
+
+                # Phase 3: Inject Kalibr context for HTTP→SDK linking
+                try:
+                    from kalibr.context import inject_kalibr_context_into_span
+
+                    inject_kalibr_context_into_span(span)
+                except Exception:
+                    pass  # Fail silently if context not available
+
+                try:
+                    # Call original method
+                    result = original_func(self_instance, *args, **kwargs)
+
+                    # Extract and set response metadata
+                    self._set_response_attributes(span, result, start_time)
+
+                    return result
+
+                except Exception as e:
+                    self.set_error(span, e)
+                    raise
+
+        return wrapper
+
+    def _traced_async_create_wrapper(self, original_func):
+        """Wrapper for async create method"""
+
+        @wraps(original_func)
+        async def wrapper(self_instance, *args, **kwargs):
+            # Extract model from kwargs
+            model = kwargs.get("model", "unknown")
+
+            # Create span with initial attributes
+            with self.tracer.start_as_current_span(
+                "openai.chat.completions.create",
+                kind=SpanKind.CLIENT,
+                attributes={
+                    "llm.vendor": "openai",
+                    "llm.request.model": model,
+                    "llm.system": "openai",
+                },
+            ) as span:
+                start_time = time.time()
+
+                # Phase 3: Inject Kalibr context for HTTP→SDK linking
+                try:
+                    from kalibr.context import inject_kalibr_context_into_span
+
+                    inject_kalibr_context_into_span(span)
+                except Exception:
+                    pass  # Fail silently if context not available
+
+                try:
+                    # Call original async method
+                    result = await original_func(self_instance, *args, **kwargs)
+
+                    # Extract and set response metadata
+                    self._set_response_attributes(span, result, start_time)
+
+                    return result
+
+                except Exception as e:
+                    self.set_error(span, e)
+                    raise
+
+        return wrapper
+
+    def _set_response_attributes(self, span, result, start_time: float) -> None:
+        """Extract metadata from response and set span attributes"""
+        try:
+            # Model
+            if hasattr(result, "model"):
+                span.set_attribute("llm.response.model", result.model)
+
+            # Token usage
+            if hasattr(result, "usage") and result.usage:
+                usage = result.usage
+                if hasattr(usage, "prompt_tokens"):
+                    span.set_attribute("llm.usage.prompt_tokens", usage.prompt_tokens)
+                if hasattr(usage, "completion_tokens"):
+                    span.set_attribute("llm.usage.completion_tokens", usage.completion_tokens)
+                if hasattr(usage, "total_tokens"):
+                    span.set_attribute("llm.usage.total_tokens", usage.total_tokens)
+
+                # Calculate cost
+                cost = self.cost_adapter.calculate_cost(
+                    result.model,
+                    {
+                        "prompt_tokens": usage.prompt_tokens,
+                        "completion_tokens": usage.completion_tokens,
+                    },
+                )
+                span.set_attribute("llm.cost_usd", cost)
+
+            # Latency
+            latency_ms = (time.time() - start_time) * 1000
+            span.set_attribute("llm.latency_ms", round(latency_ms, 2))
+
+            # Response ID
+            if hasattr(result, "id"):
+                span.set_attribute("llm.response.id", result.id)
+
+        except Exception as e:
+            # Don't fail the call if metadata extraction fails
+            span.set_attribute("llm.metadata_extraction_error", str(e))
+
+
+# Singleton instance
+_openai_instrumentation = None
+
+
+def get_instrumentation() -> OpenAIInstrumentation:
+    """Get or create the OpenAI instrumentation singleton"""
+    global _openai_instrumentation
+    if _openai_instrumentation is None:
+        _openai_instrumentation = OpenAIInstrumentation()
+    return _openai_instrumentation
+
+
+def instrument() -> bool:
+    """Instrument OpenAI SDK"""
+    return get_instrumentation().instrument()
+
+
+def uninstrument() -> bool:
+    """Uninstrument OpenAI SDK"""
+    return get_instrumentation().uninstrument()