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

aiqa/tracing.py

@@ -7,18 +7,21 @@ import json
 import logging
 import inspect
 import os
+import copy
+import requests
 from typing import Any, Callable, Optional, List
 from functools import wraps
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 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:
@@ -28,13 +31,11 @@ async def flush_tracing() -> None:
     if you want to flush immediately, e.g. before exiting a process.
     A common use is if you are tracing unit tests or experiment runs.
 
-    This flushes both the BatchSpanProcessor and the exporter buffer.
+    This flushes the BatchSpanProcessor (OTLP exporter doesn't have a separate flush method).
     """
     client = get_aiqa_client()
     if client.provider:
         client.provider.force_flush()  # Synchronous method
-    if client.exporter:    
-        await client.exporter.flush()
 
 
 # Export provider and exporter accessors for advanced usage
@@ -114,10 +115,10 @@ class TracingOptions:
         self.filter_output = filter_output
 
 
-
-
-def _prepare_input(args: tuple, kwargs: dict) -> Any:
-    """Prepare input for span attributes.
+def _prepare_input(args: tuple, kwargs: dict, sig: Optional[inspect.Signature] = None) -> Any:
+    """Prepare input for span attributes. 
+    Converts args and kwargs into a unified dict structure using function signature when available.
+    Falls back to legacy behavior for functions without inspectable signatures.
     
     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 +126,34 @@ 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
+    
+    # Try to bind args to parameter names using function signature
+    if sig is not None:
+        try:
+            bound = sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+            # Return dict of all arguments (positional args are now named)
+            result = bound.arguments.copy()
+            # Shallow copy to protect against mutating the input
+            return result
+        except (TypeError, ValueError):
+            # Binding failed (e.g., wrong number of args, *args/**kwargs issues)
+            # Fall through to legacy behavior
+            pass
+    
+    # in case binding fails
+    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
 
 
@@ -141,12 +162,17 @@ def _prepare_and_filter_input(
     kwargs: dict,
     filter_input: Optional[Callable[[Any], Any]],
     ignore_input: Optional[List[str]],
+    sig: Optional[inspect.Signature] = None,
 ) -> Any:
-    """Prepare and filter input for span attributes."""
+    """
+    Prepare and filter input for span attributes - applies the user's filter_input and ignore_input.
+    Converts all args to a dict using function signature when available.
+    """
     # Handle "self" in ignore_input by skipping the first argument
     filtered_args = args
     filtered_kwargs = kwargs.copy() if kwargs else {}
     filtered_ignore_input = ignore_input
+    filtered_sig = sig
     if ignore_input and "self" in ignore_input:
         # Remove "self" from ignore_input list (we'll handle it specially)
         filtered_ignore_input = [key for key in ignore_input if key != "self"]
@@ -156,14 +182,23 @@ def _prepare_and_filter_input(
         # Also remove "self" from kwargs if present
         if "self" in filtered_kwargs:
             del filtered_kwargs["self"]
-    
-    input_data = _prepare_input(filtered_args, filtered_kwargs)
-    if filter_input:
+        # Adjust signature to remove "self" parameter if present
+        # This is needed because we removed self from args, so signature binding will fail otherwise
+        if filtered_sig is not None:
+            params = list(filtered_sig.parameters.values())
+            if params and params[0].name == "self":
+                filtered_sig = filtered_sig.replace(parameters=params[1:])
+    # turn args, kwargs into one "nice" object (now always a dict when signature is available)
+    input_data = _prepare_input(filtered_args, filtered_kwargs, filtered_sig)
+    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 +206,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.
-    
-    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.
+    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.
     
-    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)
+    Serializes output immediately to capture its state when the function returns,
+    preventing mutations from affecting the trace.
     
-    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.
-    
-    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
+    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 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 +299,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 +318,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 +335,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 +374,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 +393,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 +410,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 +451,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,16 +462,30 @@ 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
         
+        # Get function signature once at decoration time for efficient arg name resolution
+        fn_sig: Optional[inspect.Signature] = None
+        try:
+            fn_sig = inspect.signature(fn)
+        except (ValueError, TypeError):
+            # Some callables (e.g., builtins, C extensions) don't have inspectable signatures
+            # Will fall back to legacy behavior
+            pass
+        
         # Don't get tracer here - get it lazily when function is called
         # 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 +498,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 +508,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 +579,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
@@ -733,7 +624,7 @@ def WithTracing(
         if is_async_generator:
             @wraps(fn)
             async def async_gen_traced_fn(*args, **kwargs):
-                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
+                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input, fn_sig)
                 return await _execute_generator_async(
                     lambda: fn(*args, **kwargs),
                     input_data
@@ -745,7 +636,7 @@ def WithTracing(
         elif is_generator:
             @wraps(fn)
             def gen_traced_fn(*args, **kwargs):
-                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
+                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input, fn_sig)
                 return _execute_generator_sync(
                     lambda: fn(*args, **kwargs),
                     input_data
@@ -757,7 +648,7 @@ def WithTracing(
         elif is_async:
             @wraps(fn)
             async def async_traced_fn(*args, **kwargs):
-                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
+                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input, fn_sig)
                 return await _execute_with_span_async(
                     lambda: fn(*args, **kwargs),
                     input_data
@@ -769,7 +660,7 @@ def WithTracing(
         else:
             @wraps(fn)
             def sync_traced_fn(*args, **kwargs):
-                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
+                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input, fn_sig)
                 return _execute_with_span_sync(
                     lambda: fn(*args, **kwargs),
                     input_data
@@ -816,6 +707,7 @@ def get_active_span() -> Optional[trace.Span]:
 
 def set_conversation_id(conversation_id: str) -> bool:
     """
+    Naturally a conversation might span several traces. 
     Set the gen_ai.conversation.id attribute on the active span.
     This allows you to group multiple traces together that are part of the same conversation.
     See https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/ for more details.
@@ -1165,14 +1057,12 @@ def get_span(span_id: str, organisation_id: Optional[str] = None, exclude: Optio
             print(f"Found span: {span['name']}")
             my_function(**span['input'])
     """
-    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:
+    # Check if server_url is the default (meaning AIQA_SERVER_URL was not set)
+    if not os.getenv("AIQA_SERVER_URL"):
         raise ValueError("AIQA_SERVER_URL is not set. Cannot retrieve span.")    
     if not org_id:
         raise ValueError("Organisation ID is required. Provide it as parameter or set AIQA_ORGANISATION_ID environment variable.")
@@ -1190,9 +1080,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: