genai-otel-instrument 0.1.24__py3-none-any.whl

genai_otel/instrumentors/replicate_instrumentor.py

@@ -0,0 +1,87 @@
+"""OpenTelemetry instrumentor for the Replicate API client.
+
+This instrumentor automatically traces calls to Replicate models, capturing
+relevant attributes such as the model name.
+
+Note: Replicate uses hardware-based pricing (per second of GPU/CPU time),
+not token-based pricing. Cost tracking is not applicable as the pricing model
+is fundamentally different from token-based LLM APIs.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class ReplicateInstrumentor(BaseInstrumentor):
+    """Instrumentor for Replicate.
+
+    Note: Replicate uses hardware-based pricing ($/second), not token-based.
+    Cost tracking returns None as pricing is based on execution time and hardware type.
+    """
+
+    def instrument(self, config: OTelConfig):
+        """Instrument Replicate SDK if available."""
+        self.config = config
+        try:
+            import replicate
+
+            original_run = replicate.run
+
+            # Wrap using create_span_wrapper
+            wrapped_run = self.create_span_wrapper(
+                span_name="replicate.run",
+                extract_attributes=self._extract_run_attributes,
+            )(original_run)
+
+            replicate.run = wrapped_run
+            self._instrumented = True
+            logger.info("Replicate instrumentation enabled")
+
+        except ImportError:
+            logger.debug("Replicate library not installed, instrumentation will be skipped")
+        except Exception as e:
+            logger.error("Failed to instrument Replicate: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _extract_run_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Replicate run call.
+
+        Args:
+            instance: The instance (None for module-level functions).
+            args: Positional arguments (first arg is typically the model).
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = args[0] if args else kwargs.get("model", "unknown")
+
+        attrs["gen_ai.system"] = "replicate"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "run"
+
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from Replicate response.
+
+        Note: Replicate uses hardware-based pricing ($/second of GPU/CPU time),
+        not token-based pricing. Returns None as the pricing model is incompatible
+        with token-based cost calculation.
+
+        Args:
+            result: The API response.
+
+        Returns:
+            None: Replicate uses hardware-based pricing, not token-based.
+        """
+        # Replicate uses hardware-based pricing ($/second), not tokens
+        # Cannot track costs with token-based calculator
+        return None

genai_otel/instrumentors/sambanova_instrumentor.py

@@ -0,0 +1,196 @@
+"""OpenTelemetry instrumentor for the SambaNova SDK.
+
+This instrumentor automatically traces chat completion calls to SambaNova models,
+capturing relevant attributes such as the model name and token usage.
+"""
+
+import json
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class SambaNovaInstrumentor(BaseInstrumentor):
+    """Instrumentor for SambaNova"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._sambanova_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if SambaNova library is available."""
+        try:
+            import sambanova
+
+            self._sambanova_available = True
+            logger.debug("SambaNova library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("SambaNova library not installed, instrumentation will be skipped")
+            self._sambanova_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument SambaNova SDK if available.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        if not self._sambanova_available:
+            logger.debug("Skipping SambaNova instrumentation - library not available")
+            return
+
+        self.config = config
+
+        try:
+            import sambanova
+
+            original_init = sambanova.SambaNova.__init__
+
+            def wrapped_init(instance, *args, **kwargs):
+                original_init(instance, *args, **kwargs)
+                self._instrument_client(instance)
+                return instance
+
+            sambanova.SambaNova.__init__ = wrapped_init
+            self._instrumented = True
+            logger.info("SambaNova instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument SambaNova: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _instrument_client(self, client):
+        """Instrument SambaNova client methods.
+
+        Args:
+            client: The SambaNova client instance to instrument.
+        """
+        if (
+            hasattr(client, "chat")
+            and hasattr(client.chat, "completions")
+            and hasattr(client.chat.completions, "create")
+        ):
+            original_create = client.chat.completions.create
+            instrumented_create_method = self.create_span_wrapper(
+                span_name="sambanova.chat.completion",
+                extract_attributes=self._extract_sambanova_attributes,
+            )(original_create)
+            client.chat.completions.create = instrumented_create_method
+
+    def _extract_sambanova_attributes(
+        self, instance: Any, args: Any, kwargs: Any
+    ) -> Dict[str, Any]:
+        """Extract attributes from SambaNova API call.
+
+        Args:
+            instance: The client instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = kwargs.get("model", "unknown")
+        messages = kwargs.get("messages", [])
+
+        # Core attributes
+        attrs["gen_ai.system"] = "sambanova"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "chat"
+        attrs["gen_ai.request.message_count"] = len(messages)
+
+        # Request parameters
+        if "temperature" in kwargs:
+            attrs["gen_ai.request.temperature"] = kwargs["temperature"]
+        if "top_p" in kwargs:
+            attrs["gen_ai.request.top_p"] = kwargs["top_p"]
+        if "max_tokens" in kwargs:
+            attrs["gen_ai.request.max_tokens"] = kwargs["max_tokens"]
+
+        # Tool/function definitions
+        if "tools" in kwargs:
+            try:
+                attrs["llm.tools"] = json.dumps(kwargs["tools"])
+            except (TypeError, ValueError) as e:
+                logger.debug("Failed to serialize tools: %s", e)
+
+        if messages:
+            # Only capture first 200 chars to avoid sensitive data and span size issues
+            first_message = str(messages[0])[:200]
+            attrs["gen_ai.request.first_message"] = first_message
+
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from SambaNova response.
+
+        Args:
+            result: The API response object.
+
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        if hasattr(result, "usage") and result.usage:
+            usage = result.usage
+            return {
+                "prompt_tokens": getattr(usage, "prompt_tokens", 0),
+                "completion_tokens": getattr(usage, "completion_tokens", 0),
+                "total_tokens": getattr(usage, "total_tokens", 0),
+            }
+        return None
+
+    def _extract_response_attributes(self, result) -> Dict[str, Any]:
+        """Extract response attributes from SambaNova response.
+
+        Args:
+            result: The API response object.
+
+        Returns:
+            Dict[str, Any]: Dictionary of response attributes.
+        """
+        attrs = {}
+
+        # Response ID
+        if hasattr(result, "id"):
+            attrs["gen_ai.response.id"] = result.id
+
+        # Response model (actual model used, may differ from request)
+        if hasattr(result, "model"):
+            attrs["gen_ai.response.model"] = result.model
+
+        # Finish reasons
+        if hasattr(result, "choices") and result.choices:
+            finish_reasons = [
+                choice.finish_reason
+                for choice in result.choices
+                if hasattr(choice, "finish_reason")
+            ]
+            if finish_reasons:
+                attrs["gen_ai.response.finish_reasons"] = finish_reasons
+
+        return attrs
+
+    def _extract_finish_reason(self, result) -> Optional[str]:
+        """Extract finish reason from SambaNova response.
+
+        Args:
+            result: The SambaNova API response object.
+
+        Returns:
+            Optional[str]: The finish reason string or None if not available.
+        """
+        try:
+            if hasattr(result, "choices") and result.choices:
+                first_choice = result.choices[0]
+                if hasattr(first_choice, "finish_reason"):
+                    return first_choice.finish_reason
+        except Exception as e:
+            logger.debug("Failed to extract finish_reason: %s", e)
+        return None

genai_otel/instrumentors/togetherai_instrumentor.py

@@ -0,0 +1,146 @@
+"""OpenTelemetry instrumentor for the Together AI SDK.
+
+This instrumentor automatically traces completion calls to Together AI models,
+capturing relevant attributes such as the model name and token usage.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class TogetherAIInstrumentor(BaseInstrumentor):
+    """Instrumentor for Together AI"""
+
+    def instrument(self, config: OTelConfig):
+        """Instrument Together AI SDK if available."""
+        self.config = config
+        try:
+            import together
+
+            # Instrument chat completions (newer API)
+            if hasattr(together, "Together"):
+                # This is the newer Together SDK with client-based API
+                original_init = together.Together.__init__
+
+                def wrapped_init(instance, *args, **kwargs):
+                    original_init(instance, *args, **kwargs)
+                    self._instrument_client(instance)
+
+                together.Together.__init__ = wrapped_init
+                self._instrumented = True
+                logger.info("Together AI instrumentation enabled (client-based API)")
+            # Fallback to older Complete API if available
+            elif hasattr(together, "Complete"):
+                original_complete = together.Complete.create
+
+                wrapped_complete = self.create_span_wrapper(
+                    span_name="together.complete",
+                    extract_attributes=self._extract_complete_attributes,
+                )(original_complete)
+
+                together.Complete.create = wrapped_complete
+                self._instrumented = True
+                logger.info("Together AI instrumentation enabled (Complete API)")
+
+        except ImportError:
+            logger.debug("Together AI library not installed, instrumentation will be skipped")
+        except Exception as e:
+            logger.error("Failed to instrument Together AI: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _instrument_client(self, client):
+        """Instrument Together AI client methods."""
+        if hasattr(client, "chat") and hasattr(client.chat, "completions"):
+            original_create = client.chat.completions.create
+
+            wrapped_create = self.create_span_wrapper(
+                span_name="together.chat.completion",
+                extract_attributes=self._extract_chat_attributes,
+            )(original_create)
+
+            client.chat.completions.create = wrapped_create
+
+    def _extract_chat_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Together AI chat completion call.
+
+        Args:
+            instance: The client instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = kwargs.get("model", "unknown")
+        messages = kwargs.get("messages", [])
+
+        attrs["gen_ai.system"] = "together"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "chat"
+        attrs["gen_ai.request.message_count"] = len(messages)
+
+        # Optional parameters
+        if "temperature" in kwargs:
+            attrs["gen_ai.request.temperature"] = kwargs["temperature"]
+        if "top_p" in kwargs:
+            attrs["gen_ai.request.top_p"] = kwargs["top_p"]
+        if "max_tokens" in kwargs:
+            attrs["gen_ai.request.max_tokens"] = kwargs["max_tokens"]
+
+        return attrs
+
+    def _extract_complete_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Together AI complete call.
+
+        Args:
+            instance: The instance (None for class methods).
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = kwargs.get("model", "unknown")
+
+        attrs["gen_ai.system"] = "together"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "complete"
+
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from Together AI response.
+
+        Together AI uses OpenAI-compatible format with usage field containing:
+        - prompt_tokens: Input tokens
+        - completion_tokens: Output tokens
+        - total_tokens: Total tokens
+
+        Args:
+            result: The API response object.
+
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        try:
+            # Handle OpenAI-compatible response format
+            if hasattr(result, "usage") and result.usage:
+                usage = result.usage
+                return {
+                    "prompt_tokens": getattr(usage, "prompt_tokens", 0),
+                    "completion_tokens": getattr(usage, "completion_tokens", 0),
+                    "total_tokens": getattr(usage, "total_tokens", 0),
+                }
+
+            return None
+        except Exception as e:
+            logger.debug("Failed to extract usage from Together AI response: %s", e)
+            return None

genai_otel/instrumentors/vertexai_instrumentor.py

@@ -0,0 +1,106 @@
+"""OpenTelemetry instrumentor for Google Vertex AI SDK.
+
+This instrumentor automatically traces content generation calls to Vertex AI models,
+capturing relevant attributes such as the model name and token usage.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class VertexAIInstrumentor(BaseInstrumentor):
+    """Instrumentor for Google Vertex AI"""
+
+    def instrument(self, config: OTelConfig):
+        """Instrument Vertex AI SDK if available."""
+        self.config = config
+        try:
+            from vertexai.preview.generative_models import GenerativeModel
+
+            original_generate = GenerativeModel.generate_content
+
+            # Wrap using create_span_wrapper
+            wrapped_generate = self.create_span_wrapper(
+                span_name="vertexai.generate_content",
+                extract_attributes=self._extract_generate_attributes,
+            )(original_generate)
+
+            GenerativeModel.generate_content = wrapped_generate
+            self._instrumented = True
+            logger.info("Vertex AI instrumentation enabled")
+
+        except ImportError:
+            logger.debug("Vertex AI library not installed, instrumentation will be skipped")
+        except Exception as e:
+            logger.error("Failed to instrument Vertex AI: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _extract_generate_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Vertex AI generate_content call.
+
+        Args:
+            instance: The GenerativeModel instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model_name = getattr(instance, "_model_name", "unknown")
+
+        attrs["gen_ai.system"] = "vertexai"
+        attrs["gen_ai.request.model"] = model_name
+        attrs["gen_ai.operation.name"] = "generate_content"
+
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from Vertex AI response.
+
+        Vertex AI responses include usage_metadata with:
+        - prompt_token_count: Input tokens
+        - candidates_token_count: Output tokens
+        - total_token_count: Total tokens
+
+        Args:
+            result: The API response object.
+
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        try:
+            # Handle response with usage_metadata
+            if hasattr(result, "usage_metadata") and result.usage_metadata:
+                usage_metadata = result.usage_metadata
+
+                # Try snake_case first (Python SDK style)
+                prompt_tokens = getattr(usage_metadata, "prompt_token_count", None)
+                candidates_tokens = getattr(usage_metadata, "candidates_token_count", None)
+                total_tokens = getattr(usage_metadata, "total_token_count", None)
+
+                # Fallback to camelCase (REST API style)
+                if prompt_tokens is None:
+                    prompt_tokens = getattr(usage_metadata, "promptTokenCount", 0)
+                if candidates_tokens is None:
+                    candidates_tokens = getattr(usage_metadata, "candidatesTokenCount", 0)
+                if total_tokens is None:
+                    total_tokens = getattr(usage_metadata, "totalTokenCount", 0)
+
+                if prompt_tokens or candidates_tokens:
+                    return {
+                        "prompt_tokens": int(prompt_tokens or 0),
+                        "completion_tokens": int(candidates_tokens or 0),
+                        "total_tokens": int(total_tokens or 0),
+                    }
+
+            return None
+        except Exception as e:
+            logger.debug("Failed to extract usage from Vertex AI response: %s", e)
+            return None