aiqa-client 0.4.1__py3-none-any.whl → 0.4.7__py3-none-any.whl

aiqa/tracing.py

@@ -7,6 +7,7 @@ import json
 import logging
 import inspect
 import os
+import copy
 from typing import Any, Callable, Optional, List
 from functools import wraps
 from opentelemetry import trace
@@ -15,10 +16,12 @@ from opentelemetry.trace import Status, StatusCode, SpanContext, TraceFlags
 from opentelemetry.propagate import inject, extract
 from .aiqa_exporter import AIQASpanExporter
 from .client import get_aiqa_client, get_component_tag, set_component_tag as _set_component_tag, get_aiqa_tracer
-from .constants import AIQA_TRACER_NAME
+from .constants import AIQA_TRACER_NAME, LOG_TAG
 from .object_serialiser import serialize_for_span
+from .http_utils import build_headers, get_server_url, get_api_key
+from .tracing_llm_utils import _extract_and_set_token_usage, _extract_and_set_provider_and_model
 
-logger = logging.getLogger("AIQA")
+logger = logging.getLogger(LOG_TAG)
 
 
 async def flush_tracing() -> None:
@@ -114,10 +117,12 @@ class TracingOptions:
         self.filter_output = filter_output
 
 
-
-
 def _prepare_input(args: tuple, kwargs: dict) -> Any:
-    """Prepare input for span attributes.
+    """Prepare input for span attributes. 
+    Aims to produce nice span attributes for the input, since {args, kwargs} is not a natural way to read function input. 
+    So can "unwrap" the args, kwargs.
+    
+    For single-arg-dicts or kwargs-only, returns a shallow copy of the input data.
     
     Note: This function does NOT serialize values - it just structures the data.
     Serialization happens later via serialize_for_span() to avoid double-encoding
@@ -125,14 +130,18 @@ def _prepare_input(args: tuple, kwargs: dict) -> Any:
     """
     if not args and not kwargs:
         return None
-    if len(args) == 1 and not kwargs:
-        return args[0]  # Don't serialize here - will be serialized later
-    # Multiple args or kwargs - combine into dict
-    result = {}
-    if args:
-        result["args"] = list(args)  # Keep as-is, will be serialized later
-    if kwargs:
-        result["kwargs"] = dict(kwargs)  # Keep as-is, will be serialized later
+    if not kwargs:
+        if len(args) == 1:
+            arg0 = args[0]
+            if isinstance(arg0, dict): # shallow copy to protect against mutating the input
+                return arg0.copy()
+            return arg0
+        return list(args)
+    if kwargs and len(args) == 0:
+        return kwargs.copy() # shallow copy to protect against mutating the input
+    # Multiple args and kwargs - combine into dict
+    result = kwargs.copy()
+    result["args"] = list(args)
     return result
 
 
@@ -142,7 +151,10 @@ def _prepare_and_filter_input(
     filter_input: Optional[Callable[[Any], Any]],
     ignore_input: Optional[List[str]],
 ) -> Any:
-    """Prepare and filter input for span attributes."""
+    """
+    Prepare and filter input for span attributes - applies the user's filter_input and ignore_input.
+    For single-arg-dicts or kwargs-only, returns a shallow copy of the input data.
+    """
     # Handle "self" in ignore_input by skipping the first argument
     filtered_args = args
     filtered_kwargs = kwargs.copy() if kwargs else {}
@@ -156,14 +168,17 @@ def _prepare_and_filter_input(
         # Also remove "self" from kwargs if present
         if "self" in filtered_kwargs:
             del filtered_kwargs["self"]
-    
+    # turn args, kwargs into one "nice" object
     input_data = _prepare_input(filtered_args, filtered_kwargs)
-    if filter_input:
+    if filter_input and input_data is not None:
         input_data = filter_input(input_data)
-    if filtered_ignore_input and isinstance(input_data, dict):
-        for key in filtered_ignore_input:
-            if key in input_data:
-                del input_data[key]
+    if filtered_ignore_input and len(filtered_ignore_input) > 0:
+        if not isinstance(input_data, dict):
+            logger.warning(f"_prepare_and_filter_input: skip: ignore_input is set beyond 'self': {filtered_ignore_input} but input_data is not a dict: {type(input_data)}")
+        else:
+            for key in filtered_ignore_input:
+                if key in input_data:
+                    del input_data[key]
     # Also handle case where input_data is just self (single value, not dict)
     # If we filtered out self and there are no remaining args/kwargs, return None
     if ignore_input and "self" in ignore_input and not filtered_args and not filtered_kwargs:
@@ -171,213 +186,67 @@ def _prepare_and_filter_input(
     return input_data
 
 
-def _prepare_and_filter_output(
+def _filter_and_serialize_output(
     result: Any,
     filter_output: Optional[Callable[[Any], Any]],
     ignore_output: Optional[List[str]],
 ) -> Any:
-    """Prepare and filter output for span attributes."""
+    """Filter and serialize output for span attributes."""
     output_data = result
     if filter_output:
+        if isinstance(output_data, dict):
+            output_data = output_data.copy() # copy to provide shallow protection against the user accidentally mutating the output with filter_output
         output_data = filter_output(output_data)
     if ignore_output and isinstance(output_data, dict):
         output_data = output_data.copy()
         for key in ignore_output:
             if key in output_data:
                 del output_data[key]
-    return output_data
+    
+    # Serialize immediately to create immutable result (removes mutable structures)
+    return serialize_for_span(output_data)
 
 
 def _handle_span_exception(span: trace.Span, exception: Exception) -> None:
     """Record exception on span and set error status."""
+    logger.info(f"span end: Handling span exception for {span.name}")
     error = exception if isinstance(exception, Exception) else Exception(str(exception))
     span.record_exception(error)
     span.set_status(Status(StatusCode.ERROR, str(error)))
 
 
-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:
+def _finalize_span_success_common(
+    span: trace.Span,
+    result_for_metadata: Any,
+    output_data: Any,
+    filter_output: Optional[Callable[[Any], Any]] = None,
+    ignore_output: Optional[List[str]] = None,
+) -> None:
     """
-    Extract OpenAI API style token usage from result and add to span attributes
-    using OpenTelemetry semantic conventions for gen_ai.
+    Common logic for finalizing a span with success status.
+    Extracts token usage and provider/model from result, sets output attribute, and sets status to OK.
     
-    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.
+    Serializes output immediately to capture its state when the function returns,
+    preventing mutations from affecting the trace.
     
-    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:
+    Args:
+        span: The span to finalize
+        result_for_metadata: Value to extract token usage and provider/model from
+        output_data: The output data to set on the span (will be filtered if needed)
+        filter_output: Optional function to filter output data
+        ignore_output: Optional list of keys to exclude from output
     """
-    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.
+    logger.info(f"span end: Finalizing for {span.name}")
+    _extract_and_set_token_usage(span, result_for_metadata)
+    _extract_and_set_provider_and_model(span, result_for_metadata)
     
-    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")
+    # Prepare, filter, and serialize output (serialization happens in _prepare_and_filter_output)
+    output_data = _filter_and_serialize_output(output_data, filter_output, ignore_output)
+    if output_data is not None:
+        # output_data is already serialized (immutable) from _prepare_and_filter_output
+        span.set_attribute("output", output_data)
+    span.set_status(Status(StatusCode.OK))
 
 
 class TracedGenerator:
@@ -410,7 +279,8 @@ class TracedGenerator:
         
         try:
             value = next(self._generator)
-            self._yielded_values.append(value)
+            # Serialize immediately to create immutable result (removes mutable structures)
+            self._yielded_values.append(serialize_for_span(value))
             return value
         except StopIteration:
             self._exhausted = True
@@ -428,10 +298,7 @@ class TracedGenerator:
     def _finalize_span_success(self):
         """Set output and success status on span."""
         # Check last yielded value for token usage (common pattern in streaming responses)
-        if self._yielded_values:
-            last_value = self._yielded_values[-1]
-            _extract_and_set_token_usage(self._span, last_value)
-            _extract_and_set_provider_and_model(self._span, last_value)
+        result_for_metadata = self._yielded_values[-1] if self._yielded_values else None
         
         # Record summary of yielded values
         output_data = {
@@ -448,10 +315,13 @@ class TracedGenerator:
             if len(self._yielded_values) > sample_size:
                 output_data["truncated"] = True
         
-        output_data = _prepare_and_filter_output(output_data, self._filter_output, self._ignore_output)
-        if output_data is not None:
-            self._span.set_attribute("output", serialize_for_span(output_data))
-        self._span.set_status(Status(StatusCode.OK))
+        _finalize_span_success_common(
+            self._span,
+            result_for_metadata,
+            output_data,
+            self._filter_output,
+            self._ignore_output,
+        )
 
 
 class TracedAsyncGenerator:
@@ -484,7 +354,8 @@ class TracedAsyncGenerator:
         
         try:
             value = await self._generator.__anext__()
-            self._yielded_values.append(value)
+            # Serialize immediately to create immutable result (removes mutable structures)
+            self._yielded_values.append(serialize_for_span(value))
             return value
         except StopAsyncIteration:
             self._exhausted = True
@@ -502,10 +373,7 @@ class TracedAsyncGenerator:
     def _finalize_span_success(self):
         """Set output and success status on span."""
         # Check last yielded value for token usage (common pattern in streaming responses)
-        if self._yielded_values:
-            last_value = self._yielded_values[-1]
-            _extract_and_set_token_usage(self._span, last_value)
-            _extract_and_set_provider_and_model(self._span, last_value)
+        result_for_metadata = self._yielded_values[-1] if self._yielded_values else None
         
         # Record summary of yielded values
         output_data = {
@@ -522,10 +390,13 @@ class TracedAsyncGenerator:
             if len(self._yielded_values) > sample_size:
                 output_data["truncated"] = True
         
-        output_data = _prepare_and_filter_output(output_data, self._filter_output, self._ignore_output)
-        if output_data is not None:
-            self._span.set_attribute("output", serialize_for_span(output_data))
-        self._span.set_status(Status(StatusCode.OK))
+        _finalize_span_success_common(
+            self._span,
+            result_for_metadata,
+            output_data,
+            self._filter_output,
+            self._ignore_output,
+        )
 
 
 def WithTracing(
@@ -560,20 +431,6 @@ def WithTracing(
         def my_function(x, y):
             return x + y
         
-        @WithTracing
-        async def my_async_function(x, y):
-            return x + y
-        
-        @WithTracing
-        def my_generator(n):
-            for i in range(n):
-                yield i * 2
-        
-        @WithTracing
-        async def my_async_generator(n):
-            for i in range(n):
-                yield i * 2
-        
         @WithTracing(name="custom_name")
         def another_function():
             pass
@@ -585,7 +442,7 @@ def WithTracing(
         if hasattr(fn, "_is_traced"):
             logger.warning(f"Function {fn_name} is already traced, skipping tracing again")
             return fn
-        
+        logger.info(f"WithTracing function {fn_name}")
         is_async = inspect.iscoroutinefunction(fn)
         is_generator = inspect.isgeneratorfunction(fn)
         is_async_generator = inspect.isasyncgenfunction(fn) if hasattr(inspect, 'isasyncgenfunction') else False
@@ -594,7 +451,12 @@ def WithTracing(
         # This ensures initialization only happens when tracing is actually used
         
         def _setup_span(span: trace.Span, input_data: Any) -> bool:
-            """Setup span with input data. Returns True if span is recording."""
+            """
+            Setup span with input data. Returns True if span is recording.
+            
+            Serializes input immediately to capture its state at function start,
+            preventing mutations from affecting the trace.
+            """
             if not span.is_recording():
                 logger.warning(f"Span {fn_name} is not recording - will not be exported")
                 return False
@@ -607,6 +469,8 @@ def WithTracing(
                 span.set_attribute("gen_ai.component.id", component_tag)
             
             if input_data is not None:
+                # Serialize input immediately to capture state at function start
+                # input_data has already been copied in _prepare_and_filter_input
                 span.set_attribute("input", serialize_for_span(input_data))
             
             trace_id = format(span.get_span_context().trace_id, "032x")
@@ -615,30 +479,28 @@ def WithTracing(
         
         def _finalize_span_success(span: trace.Span, result: Any) -> None:
             """Set output and success status on span."""
-            # Extract and set token usage if present (before filtering output)
-            _extract_and_set_token_usage(span, result)
-            # Extract and set provider/model if present (before filtering output)
-            _extract_and_set_provider_and_model(span, result)
-            
-            output_data = _prepare_and_filter_output(result, filter_output, ignore_output)
-            if output_data is not None:
-                span.set_attribute("output", serialize_for_span(output_data))
-            span.set_status(Status(StatusCode.OK))
+            _finalize_span_success_common(
+                span,
+                result,
+                result,
+                filter_output,
+                ignore_output,
+            )
         
         def _execute_with_span_sync(executor: Callable[[], Any], input_data: Any) -> Any:
-            """Execute sync function within span context, handling input/output and exceptions."""
+            """Execute sync function within span context, handling input/output and exceptions.
+            Note: input_data has already gone through _prepare_and_filter_input
+            """
             # Ensure tracer provider is initialized before creating spans
             # This is called lazily when the function runs, not at decorator definition time
             client = get_aiqa_client()
             if not client.enabled:
                 return executor()
-            
             # Get tracer after initialization (lazy)
             tracer = get_aiqa_tracer()
             with tracer.start_as_current_span(fn_name) as span:
                 if not _setup_span(span, input_data):
-                    return executor()
-                
+                    return executor() # span is not recording, so just execute the function and return the result                
                 try:
                     result = executor()
                     _finalize_span_success(span, result)
@@ -688,7 +550,7 @@ def WithTracing(
             
             try:
                 if not _setup_span(span, input_data):
-                    generator = executor()
+                    generator = executor() # span is not recording, so just execute the function and return the result
                     trace.context_api.detach(token)
                     span.end()
                     return generator
@@ -1168,8 +1030,8 @@ def get_span(span_id: str, organisation_id: Optional[str] = None, exclude: Optio
     import os
     import requests
     
-    server_url = os.getenv("AIQA_SERVER_URL", "").rstrip("/")
-    api_key = os.getenv("AIQA_API_KEY", "")
+    server_url = get_server_url()
+    api_key = get_api_key()
     org_id = organisation_id or os.getenv("AIQA_ORGANISATION_ID", "")
     
     if not server_url:
@@ -1190,9 +1052,7 @@ def get_span(span_id: str, organisation_id: Optional[str] = None, exclude: Optio
             "fields": "*" if not exclude else None,
         }
         
-        headers = {"Content-Type": "application/json"}
-        if api_key:
-            headers["Authorization"] = f"ApiKey {api_key}"
+        headers = build_headers(api_key)
         
         response = requests.get(url, params=params, headers=headers)
         if response.status_code == 200: