genai-otel-instrument 0.1.2.dev0__py3-none-any.whl → 0.1.7.dev0__py3-none-any.whl

genai_otel/instrumentors/togetherai_instrumentor.py

@@ -1,42 +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.
+capturing relevant attributes such as the model name and token usage.
 """
 
-from typing import Dict, Optional
+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
 
-            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")
+            # Instrument chat completions (newer API)
+            if hasattr(together, "Together"):
+                # This is the newer Together SDK with client-based API
+                original_init = together.Together.__init__
 
-                    span.set_attribute("gen_ai.system", "together")
-                    span.set_attribute("gen_ai.request.model", model)
+                def wrapped_init(instance, *args, **kwargs):
+                    original_init(instance, *args, **kwargs)
+                    self._instrument_client(instance)
 
-                    if self.request_counter:
-                        self.request_counter.add(1, {"model": model, "provider": "together"})
+                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
 
-                    result = original_complete(*args, **kwargs)
-                    return result
+                wrapped_complete = self.create_span_wrapper(
+                    span_name="together.complete",
+                    extract_attributes=self._extract_complete_attributes,
+                )(original_complete)
 
-            together.Complete.create = wrapped_complete
+                together.Complete.create = wrapped_complete
+                self._instrumented = True
+                logger.info("Together AI instrumentation enabled (Complete API)")
 
         except ImportError:
-            pass
+            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]]:
-        return None
+        """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

@@ -1,42 +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.
+capturing relevant attributes such as the model name and token usage.
 """
 
-from typing import Dict, Optional
+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
 
-            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")
+            # Wrap using create_span_wrapper
+            wrapped_generate = self.create_span_wrapper(
+                span_name="vertexai.generate_content",
+                extract_attributes=self._extract_generate_attributes,
+            )(original_generate)
 
-                    span.set_attribute("gen_ai.system", "vertexai")
-                    span.set_attribute("gen_ai.request.model", model_name)
+            GenerativeModel.generate_content = wrapped_generate
+            self._instrumented = True
+            logger.info("Vertex AI instrumentation enabled")
 
-                    if self.request_counter:
-                        self.request_counter.add(1, {"model": model_name, "provider": "vertexai"})
+        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
 
-                    result = original_generate(instance, *args, **kwargs)
-                    return result
+    def _extract_generate_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Vertex AI generate_content call.
 
-            GenerativeModel.generate_content = wrapped_generate
+        Args:
+            instance: The GenerativeModel instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
 
-        except ImportError:
-            pass
+        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]]:
-        return None
+        """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