aiqa-client 0.4.3__py3-none-any.whl → 0.5.2__py3-none-any.whl

aiqa/tracing_llm_utils.py

@@ -0,0 +1,191 @@
+# Functions for extracting and setting LLM-specific attributes on a span.
+import logging
+from .constants import LOG_TAG
+from opentelemetry import trace
+from typing import Any
+
+logger = logging.getLogger(LOG_TAG)
+
+
+def _is_attribute_set(span: trace.Span, attribute_name: str) -> bool:
+    """
+    Check if an attribute is already set on a span.
+    Returns True if the attribute exists, False otherwise.
+    Safe against exceptions.
+    """
+    try:
+        # Try multiple ways to access span attributes (SDK spans may store them differently)
+        # Check public 'attributes' property
+        if hasattr(span, "attributes"):
+            attrs = span.attributes
+            if attrs and attribute_name in attrs:
+                return True
+        
+        # Check private '_attributes' (common in OpenTelemetry SDK)
+        if hasattr(span, "_attributes"):
+            attrs = span._attributes
+            if attrs and attribute_name in attrs:
+                return True
+        
+        # If we can't find the attribute, assume not set (conservative approach)
+        return False
+    except Exception:
+        # If anything goes wrong, assume not set (conservative approach)
+        return False
+
+def _extract_and_set_token_usage(span: trace.Span, result: Any) -> None:
+    """
+    Extract OpenAI API style token usage from result and add to span attributes
+    using OpenTelemetry semantic conventions for gen_ai.
+    
+    Looks for usage dict with prompt_tokens, completion_tokens, and total_tokens.
+    Sets gen_ai.usage.input_tokens, gen_ai.usage.output_tokens, and gen_ai.usage.total_tokens.
+    Only sets attributes that are not already set.
+    
+    This function detects token usage from OpenAI API response patterns:
+    - OpenAI Chat Completions API: The 'usage' object contains 'prompt_tokens', 'completion_tokens', and 'total_tokens'.
+      See https://platform.openai.com/docs/api-reference/chat/object (usage field)
+    - OpenAI Completions API: The 'usage' object contains 'prompt_tokens', 'completion_tokens', and 'total_tokens'.
+      See https://platform.openai.com/docs/api-reference/completions/object (usage field)
+    
+    This function is safe against exceptions and will not derail tracing or program execution.
+    """
+    try:
+        if not span.is_recording():
+            return
+        
+        usage = None
+        
+        # Check if result is a dict with 'usage' key
+        try:
+            if isinstance(result, dict):
+                usage = result.get("usage")
+                # Also check if result itself is a usage dict (OpenAI format)
+                if usage is None and all(key in result for key in ("prompt_tokens", "completion_tokens", "total_tokens")):
+                    usage = result
+                # Also check if result itself is a usage dict (Bedrock format)
+                elif usage is None and all(key in result for key in ("input_tokens", "output_tokens")):
+                    usage = result
+            
+            # Check if result has a 'usage' attribute (e.g., OpenAI response object)
+            elif hasattr(result, "usage"):
+                usage = result.usage
+        except Exception:
+            # If accessing result properties fails, just return silently
+            return
+        
+        # Extract token usage if found
+        if isinstance(usage, dict):
+            try:
+                # Support both OpenAI format (prompt_tokens/completion_tokens) and Bedrock format (input_tokens/output_tokens)
+                prompt_tokens = usage.get("prompt_tokens") or usage.get("PromptTokens")
+                completion_tokens = usage.get("completion_tokens") or usage.get("CompletionTokens")
+                input_tokens = usage.get("input_tokens") or usage.get("InputTokens")
+                output_tokens = usage.get("output_tokens") or usage.get("OutputTokens")
+                total_tokens = usage.get("total_tokens") or usage.get("TotalTokens")
+                
+                # Use Bedrock format if OpenAI format not available
+                if prompt_tokens is None:
+                    prompt_tokens = input_tokens
+                if completion_tokens is None:
+                    completion_tokens = output_tokens
+                
+                # Calculate total_tokens if not provided but we have input and output
+                if total_tokens is None and prompt_tokens is not None and completion_tokens is not None:
+                    total_tokens = prompt_tokens + completion_tokens
+                
+                # Only set attributes that are not already set
+                if prompt_tokens is not None and not _is_attribute_set(span, "gen_ai.usage.input_tokens"):
+                    span.set_attribute("gen_ai.usage.input_tokens", prompt_tokens)
+                if completion_tokens is not None and not _is_attribute_set(span, "gen_ai.usage.output_tokens"):
+                    span.set_attribute("gen_ai.usage.output_tokens", completion_tokens)
+                if total_tokens is not None and not _is_attribute_set(span, "gen_ai.usage.total_tokens"):
+                    span.set_attribute("gen_ai.usage.total_tokens", total_tokens)
+            except Exception:
+                # If setting attributes fails, log but don't raise
+                logger.debug(f"Failed to set token usage attributes on span")
+    except Exception:
+        # Catch any other exceptions to ensure this never derails tracing
+        logger.debug(f"Error in _extract_and_set_token_usage")
+
+
+def _extract_and_set_provider_and_model(span: trace.Span, result: Any) -> None:
+    """
+    Extract provider and model information from result and add to span attributes
+    using OpenTelemetry semantic conventions for gen_ai.
+    
+    Looks for 'model', 'provider', 'provider_name' fields in the result.
+    Sets gen_ai.provider.name and gen_ai.request.model.
+    Only sets attributes that are not already set.
+    
+    This function detects model information from common API response patterns:
+    - OpenAI Chat Completions API: The 'model' field is at the top level of the response.
+      See https://platform.openai.com/docs/api-reference/chat/object
+    - OpenAI Completions API: The 'model' field is at the top level of the response.
+      See https://platform.openai.com/docs/api-reference/completions/object
+    
+    This function is safe against exceptions and will not derail tracing or program execution.
+    """
+    try:
+        if not span.is_recording():
+            return
+        
+        model = None
+        provider = None
+        
+        # Check if result is a dict
+        try:
+            if isinstance(result, dict):
+                model = result.get("model") or result.get("Model")
+                provider = result.get("provider") or result.get("Provider") or result.get("provider_name") or result.get("providerName")
+            
+            # Check if result has attributes (e.g., OpenAI response object)
+            elif hasattr(result, "model"):
+                model = result.model
+            if hasattr(result, "provider"):
+                provider = result.provider
+            elif hasattr(result, "provider_name"):
+                provider = result.provider_name
+            elif hasattr(result, "providerName"):
+                provider = result.providerName
+            
+            # Check nested structures (e.g., response.data.model)
+            if model is None and hasattr(result, "data"):
+                data = result.data
+                if isinstance(data, dict):
+                    model = data.get("model") or data.get("Model")
+                elif hasattr(data, "model"):
+                    model = data.model
+            
+            # Check for model in choices (OpenAI pattern)
+            if model is None and isinstance(result, dict):
+                choices = result.get("choices")
+                if choices and isinstance(choices, list) and len(choices) > 0:
+                    first_choice = choices[0]
+                    if isinstance(first_choice, dict):
+                        model = first_choice.get("model")
+                    elif hasattr(first_choice, "model"):
+                        model = first_choice.model
+        except Exception:
+            # If accessing result properties fails, just return silently
+            return
+        
+        # Set attributes if found and not already set
+        try:
+            if model is not None and not _is_attribute_set(span, "gen_ai.request.model"):
+                # Convert to string if needed
+                model_str = str(model) if model is not None else None
+                if model_str:
+                    span.set_attribute("gen_ai.request.model", model_str)
+            
+            if provider is not None and not _is_attribute_set(span, "gen_ai.provider.name"):
+                # Convert to string if needed
+                provider_str = str(provider) if provider is not None else None
+                if provider_str:
+                    span.set_attribute("gen_ai.provider.name", provider_str)
+        except Exception:
+            # If setting attributes fails, log but don't raise
+            logger.debug(f"Failed to set provider/model attributes on span")
+    except Exception:
+        # Catch any other exceptions to ensure this never derails tracing
+        logger.debug(f"Error in _extract_and_set_provider_and_model")

{aiqa_client-0.4.3.dist-info → aiqa_client-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiqa-client
-Version: 0.4.3
+Version: 0.5.2
 Summary: OpenTelemetry-based Python client for tracing functions and sending traces to the AIQA server
 Author-email: AIQA <info@aiqa.dev>
 License: MIT

aiqa_client-0.5.2.dist-info/RECORD

@@ -0,0 +1,14 @@
+aiqa/__init__.py,sha256=V1VLfpxU_BXfkvKyhGckZsPYF43NJhoXeVX6FkeHr9g,1721
+aiqa/client.py,sha256=Ba3v-voBlTSCr-RU88INLXsF_5vqp42QiQWCFciSJbU,12542
+aiqa/constants.py,sha256=tZuh7XvKs6hFvWc-YnQ5Na6uogJMsRrMy-rWOauvcIA,226
+aiqa/experiment_runner.py,sha256=XAZsjVP70UH_QTk5ANSOQYAhmozuGXwKB5qWWHs-zeE,11186
+aiqa/http_utils.py,sha256=OIB4tRI2TiDl4VKDmtbLWg9Q7TicMBeL7scLYEhVPXI,4944
+aiqa/object_serialiser.py,sha256=DBv7EyXIwfwjwXHDsIwdZNFmQffRb5fKAE0r8qhoqgc,16958
+aiqa/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aiqa/tracing.py,sha256=gdmtpvBVbvc_HrJHgtr78_XH9sIWRjUoEkXuCuNmuc0,45662
+aiqa/tracing_llm_utils.py,sha256=rNx6v6Wh_Mhv-_DPU9_aWS7YQcO46oiv0YPdBK1KVL8,9338
+aiqa_client-0.5.2.dist-info/licenses/LICENSE.txt,sha256=kIzkzLuzG0HHaWYm4F4W5FeJ1Yxut3Ec6bhLWyw798A,1062
+aiqa_client-0.5.2.dist-info/METADATA,sha256=xMaQSnI3AiNE6lYs2vM6BV9VxQWMHXyDoIl6JXwdi3I,7705
+aiqa_client-0.5.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aiqa_client-0.5.2.dist-info/top_level.txt,sha256=nwcsuVVSuWu27iLxZd4n1evVzv1W6FVTrSnCXCc-NQs,5
+aiqa_client-0.5.2.dist-info/RECORD,,