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

genai_otel/instrumentors/ollama_instrumentor.py

@@ -2,11 +2,11 @@
 
 This instrumentor automatically traces calls to Ollama models for both
 generation and chat functionalities, capturing relevant attributes such as
-the model name.
+the model name and token usage.
 """
 
 import logging
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from ..config import OTelConfig
 from .base import BaseInstrumentor
@@ -22,8 +22,8 @@ class OllamaInstrumentor(BaseInstrumentor):
         super().__init__()
         self._ollama_available = False
         self._ollama_module = None
-        self._original_generate = None  # Add this
-        self._original_chat = None  # Add this
+        self._original_generate = None
+        self._original_chat = None
         self._check_availability()
 
     def _check_availability(self):
@@ -46,38 +46,107 @@ class OllamaInstrumentor(BaseInstrumentor):
         if not self._ollama_available or self._ollama_module is None:
             return
 
-        # Store original methods
-        self._original_generate = self._ollama_module.generate
-        self._original_chat = self._ollama_module.chat
-
-        def wrapped_generate(*args, **kwargs):
-            with self.tracer.start_as_current_span("ollama.generate") as span:
-                model = kwargs.get("model", "unknown")
-
-                span.set_attribute("gen_ai.system", "ollama")
-                span.set_attribute("gen_ai.request.model", model)
-
-                if self.request_counter:
-                    self.request_counter.add(1, {"model": model, "provider": "ollama"})
-
-                result = self._original_generate(*args, **kwargs)
-                return result
-
-        def wrapped_chat(*args, **kwargs):
-            with self.tracer.start_as_current_span("ollama.chat") as span:
-                model = kwargs.get("model", "unknown")
-
-                span.set_attribute("gen_ai.system", "ollama")
-                span.set_attribute("gen_ai.request.model", model)
+        try:
+            # Store original methods and wrap them
+            self._original_generate = self._ollama_module.generate
+            self._original_chat = self._ollama_module.chat
+
+            # Wrap generate method
+            wrapped_generate = self.create_span_wrapper(
+                span_name="ollama.generate",
+                extract_attributes=self._extract_generate_attributes,
+            )(self._original_generate)
+            self._ollama_module.generate = wrapped_generate
+
+            # Wrap chat method
+            wrapped_chat = self.create_span_wrapper(
+                span_name="ollama.chat",
+                extract_attributes=self._extract_chat_attributes,
+            )(self._original_chat)
+            self._ollama_module.chat = wrapped_chat
+
+            self._instrumented = True
+            logger.info("Ollama instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument Ollama: %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 Ollama generate call.
+
+        Args:
+            instance: The client instance (None for module-level functions).
+            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"] = "ollama"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "generate"
+
+        return attrs
+
+    def _extract_chat_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Ollama chat call.
+
+        Args:
+            instance: The client instance (None for module-level functions).
+            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"] = "ollama"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "chat"
+        attrs["gen_ai.request.message_count"] = len(messages)
+
+        return attrs
 
-                if self.request_counter:
-                    self.request_counter.add(1, {"model": model, "provider": "ollama"})
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from Ollama response.
 
-                result = self._original_chat(*args, **kwargs)
-                return result
+        Ollama responses include:
+        - prompt_eval_count: Input tokens
+        - eval_count: Output tokens
 
-        self._ollama_module.generate = wrapped_generate
-        self._ollama_module.chat = wrapped_chat
+        Args:
+            result: The API response object or dictionary.
 
-    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
-        return None
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        try:
+            # Handle both dict and object responses
+            if isinstance(result, dict):
+                prompt_tokens = result.get("prompt_eval_count", 0)
+                completion_tokens = result.get("eval_count", 0)
+            elif hasattr(result, "prompt_eval_count") and hasattr(result, "eval_count"):
+                prompt_tokens = getattr(result, "prompt_eval_count", 0)
+                completion_tokens = getattr(result, "eval_count", 0)
+            else:
+                return None
+
+            if prompt_tokens == 0 and completion_tokens == 0:
+                return None
+
+            return {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+            }
+        except Exception as e:
+            logger.debug("Failed to extract usage from Ollama response: %s", e)
+            return None

genai_otel/instrumentors/replicate_instrumentor.py

@@ -2,41 +2,86 @@
 
 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.
 """
 
-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 ReplicateInstrumentor(BaseInstrumentor):
-    """Instrumentor for Replicate"""
+    """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
 
-            def wrapped_run(*args, **kwargs):
-                with self.tracer.start_as_current_span("replicate.run") as span:
-                    model = args[0] if args else "unknown"
+            # Wrap using create_span_wrapper
+            wrapped_run = self.create_span_wrapper(
+                span_name="replicate.run",
+                extract_attributes=self._extract_run_attributes,
+            )(original_run)
 
-                    span.set_attribute("gen_ai.system", "replicate")
-                    span.set_attribute("gen_ai.request.model", model)
+            replicate.run = wrapped_run
+            self._instrumented = True
+            logger.info("Replicate instrumentation enabled")
 
-                    if self.request_counter:
-                        self.request_counter.add(1, {"model": model, "provider": "replicate"})
+        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
 
-                    result = original_run(*args, **kwargs)
-                    return result
+    def _extract_run_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Replicate run call.
 
-            replicate.run = wrapped_run
+        Args:
+            instance: The instance (None for module-level functions).
+            args: Positional arguments (first arg is typically the model).
+            kwargs: Keyword arguments.
 
-        except ImportError:
-            pass
+        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/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