genai-otel-instrument 0.1.1.dev0__py3-none-any.whl

genai_otel/instrumentors/openai_instrumentor.py

@@ -0,0 +1,241 @@
+"""OpenTelemetry instrumentor for the OpenAI Python SDK.
+
+This instrumentor automatically traces chat completion calls made using the
+OpenAI SDK, capturing relevant attributes such as the model name, message count,
+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 OpenAIInstrumentor(BaseInstrumentor):
+    """Instrumentor for OpenAI SDK"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._openai_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if OpenAI library is available."""
+        try:
+            import openai
+
+            self._openai_available = True
+            logger.debug("OpenAI library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("OpenAI library not installed, instrumentation will be skipped")
+            self._openai_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument OpenAI SDK if available.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        if not self._openai_available:
+            logger.debug("Skipping OpenAI instrumentation - library not available")
+            return
+
+        self.config = config
+
+        try:
+            import openai
+            import wrapt
+
+            # Instrument OpenAI client initialization
+            if hasattr(openai, "OpenAI"):
+                original_init = openai.OpenAI.__init__
+
+                def wrapped_init(wrapped, instance, args, kwargs):
+                    result = wrapped(*args, **kwargs)
+                    self._instrument_client(instance)
+                    return result
+
+                openai.OpenAI.__init__ = wrapt.FunctionWrapper(original_init, wrapped_init)
+                self._instrumented = True
+                logger.info("OpenAI instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument OpenAI: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _instrument_client(self, client):
+        """Instrument OpenAI client methods.
+
+        Args:
+            client: The OpenAI 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="openai.chat.completion",
+                extract_attributes=self._extract_openai_attributes,
+            )(original_create)
+            client.chat.completions.create = instrumented_create_method
+
+    def _extract_openai_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from OpenAI 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"] = "openai"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "chat"  # NEW: operation name
+        attrs["gen_ai.request.message_count"] = len(messages)
+
+        # Request parameters (NEW)
+        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"]
+        if "frequency_penalty" in kwargs:
+            attrs["gen_ai.request.frequency_penalty"] = kwargs["frequency_penalty"]
+        if "presence_penalty" in kwargs:
+            attrs["gen_ai.request.presence_penalty"] = kwargs["presence_penalty"]
+
+        # Tool/function definitions (Phase 3.1)
+        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 OpenAI 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
+            usage_dict = {
+                "prompt_tokens": getattr(usage, "prompt_tokens", 0),
+                "completion_tokens": getattr(usage, "completion_tokens", 0),
+                "total_tokens": getattr(usage, "total_tokens", 0),
+            }
+
+            # Extract reasoning tokens for o1 models (Phase 3.2)
+            if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
+                details = usage.completion_tokens_details
+                usage_dict["completion_tokens_details"] = {
+                    "reasoning_tokens": getattr(details, "reasoning_tokens", 0)
+                }
+
+            return usage_dict
+        return None
+
+    def _extract_response_attributes(self, result) -> Dict[str, Any]:
+        """Extract response attributes from OpenAI 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
+
+            # Tool calls extraction (Phase 3.1)
+            for choice_idx, choice in enumerate(result.choices):
+                message = getattr(choice, "message", None)
+                if message and hasattr(message, "tool_calls") and message.tool_calls:
+                    for tc_idx, tool_call in enumerate(message.tool_calls):
+                        prefix = f"llm.output_messages.{choice_idx}.message.tool_calls.{tc_idx}"
+                        if hasattr(tool_call, "id"):
+                            attrs[f"{prefix}.tool_call.id"] = tool_call.id
+                        if hasattr(tool_call, "function"):
+                            if hasattr(tool_call.function, "name"):
+                                attrs[f"{prefix}.tool_call.function.name"] = tool_call.function.name
+                            if hasattr(tool_call.function, "arguments"):
+                                attrs[f"{prefix}.tool_call.function.arguments"] = (
+                                    tool_call.function.arguments
+                                )
+
+        return attrs
+
+    def _add_content_events(self, span, result, request_kwargs: dict):
+        """Add prompt and completion content as span events.
+
+        Args:
+            span: The OpenTelemetry span.
+            result: The API response object.
+            request_kwargs: The original request kwargs.
+        """
+        # Add prompt content events
+        messages = request_kwargs.get("messages", [])
+        for idx, message in enumerate(messages):
+            if isinstance(message, dict):
+                role = message.get("role", "unknown")
+                content = message.get("content", "")
+                span.add_event(
+                    f"gen_ai.prompt.{idx}",
+                    attributes={"gen_ai.prompt.role": role, "gen_ai.prompt.content": str(content)},
+                )
+
+        # Add completion content events
+        if hasattr(result, "choices") and result.choices:
+            for idx, choice in enumerate(result.choices):
+                if hasattr(choice, "message") and hasattr(choice.message, "content"):
+                    content = choice.message.content
+                    span.add_event(
+                        f"gen_ai.completion.{idx}",
+                        attributes={
+                            "gen_ai.completion.role": "assistant",
+                            "gen_ai.completion.content": str(content),
+                        },
+                    )

genai_otel/instrumentors/replicate_instrumentor.py

@@ -0,0 +1,42 @@
+"""OpenTelemetry instrumentor for the Replicate API client.
+
+This instrumentor automatically traces calls to Replicate models, capturing
+relevant attributes such as the model name.
+"""
+
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+
+class ReplicateInstrumentor(BaseInstrumentor):
+    """Instrumentor for Replicate"""
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+        try:
+            import replicate
+
+            original_run = replicate.run
+
+            def wrapped_run(*args, **kwargs):
+                with self.tracer.start_as_current_span("replicate.run") as span:
+                    model = args[0] if args else "unknown"
+
+                    span.set_attribute("gen_ai.system", "replicate")
+                    span.set_attribute("gen_ai.request.model", model)
+
+                    if self.request_counter:
+                        self.request_counter.add(1, {"model": model, "provider": "replicate"})
+
+                    result = original_run(*args, **kwargs)
+                    return result
+
+            replicate.run = wrapped_run
+
+        except ImportError:
+            pass
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        return None

genai_otel/instrumentors/togetherai_instrumentor.py

@@ -0,0 +1,42 @@
+"""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.
+"""
+
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+
+class TogetherAIInstrumentor(BaseInstrumentor):
+    """Instrumentor for Together AI"""
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+        try:
+            import together
+
+            original_complete = together.Complete.create
+
+            def wrapped_complete(*args, **kwargs):
+                with self.tracer.start_as_current_span("together.complete") as span:
+                    model = kwargs.get("model", "unknown")
+
+                    span.set_attribute("gen_ai.system", "together")
+                    span.set_attribute("gen_ai.request.model", model)
+
+                    if self.request_counter:
+                        self.request_counter.add(1, {"model": model, "provider": "together"})
+
+                    result = original_complete(*args, **kwargs)
+                    return result
+
+            together.Complete.create = wrapped_complete
+
+        except ImportError:
+            pass
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        return None

genai_otel/instrumentors/vertexai_instrumentor.py

@@ -0,0 +1,42 @@
+"""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.
+"""
+
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+
+class VertexAIInstrumentor(BaseInstrumentor):
+    """Instrumentor for Google Vertex AI"""
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+        try:
+            from vertexai.preview.generative_models import GenerativeModel
+
+            original_generate = GenerativeModel.generate_content
+
+            def wrapped_generate(instance, *args, **kwargs):
+                with self.tracer.start_as_current_span("vertexai.generate_content") as span:
+                    model_name = getattr(instance, "_model_name", "unknown")
+
+                    span.set_attribute("gen_ai.system", "vertexai")
+                    span.set_attribute("gen_ai.request.model", model_name)
+
+                    if self.request_counter:
+                        self.request_counter.add(1, {"model": model_name, "provider": "vertexai"})
+
+                    result = original_generate(instance, *args, **kwargs)
+                    return result
+
+            GenerativeModel.generate_content = wrapped_generate
+
+        except ImportError:
+            pass
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        return None