aiqa-client 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

aiqa/tracing.py

@@ -3,24 +3,43 @@ OpenTelemetry tracing setup and utilities. Initializes tracer provider on import
 Provides WithTracing decorator to automatically trace function calls.
 """
 
+import os
 import json
 import logging
 import inspect
-import os
-from typing import Any, Callable, Optional, List
+from typing import Any, Callable, Optional, Dict
 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 opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.semconv.resource import ResourceAttributes
+from opentelemetry.trace import Status, StatusCode
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from .aiqa_exporter import AIQASpanExporter
-from .client import get_aiqa_client, AIQA_TRACER_NAME, get_component_tag, set_component_tag as _set_component_tag, get_aiqa_tracer
-from .object_serialiser import serialize_for_span
 
-logger = logging.getLogger("AIQA")
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+exporter = AIQASpanExporter()
+
+# Initialize OpenTelemetry
+provider = TracerProvider(
+    resource=Resource.create(
+        {
+            ResourceAttributes.SERVICE_NAME: os.getenv("OTEL_SERVICE_NAME", "aiqa-service"),
+        }
+    )
+)
+
+provider.add_span_processor(BatchSpanProcessor(exporter))
+trace.set_tracer_provider(provider)
+
+# Get a tracer instance
+tracer = trace.get_tracer("aiqa-tracer")
 
 
-async def flush_tracing() -> None:
+async def flush_spans() -> None:
     """
     Flush all pending spans to the server.
     Flushes also happen automatically every few seconds. So you only need to call this function
@@ -28,11 +47,8 @@ async def flush_tracing() -> None:
 
     This flushes both the BatchSpanProcessor and the exporter buffer.
     """
-    client = get_aiqa_client()
-    if client.get("provider"):
-        client["provider"].force_flush()  # Synchronous method
-    if client.get("exporter"):    
-        await client["exporter"].flush()
+    provider.force_flush()  # Synchronous method
+    await exporter.flush()
 
 
 async def shutdown_tracing() -> None:
@@ -40,86 +56,25 @@ async def shutdown_tracing() -> None:
     Shutdown the tracer provider and exporter.
     It is not necessary to call this function.
     """
-    try:
-        client = get_aiqa_client()
-        if client.get("provider"):
-            client["provider"].shutdown()  # Synchronous method
-        if client.get("exporter"):
-            client["exporter"].shutdown()  # Synchronous method
-    except Exception as e:
-        logger.error(f"Error shutting down tracing: {e}")
+    provider.shutdown()  # Synchronous method
+    await exporter.shutdown()
 
 
-# Export provider and exporter accessors for advanced usage
-
-__all__ = [
-    "get_provider", "get_exporter", "flush_tracing", "shutdown_tracing", "WithTracing",
-    "set_span_attribute", "set_span_name", "get_active_span",
-    "get_trace_id", "get_span_id", "create_span_from_trace_id", "inject_trace_context", "extract_trace_context",
-    "set_conversation_id", "set_component_tag", "set_token_usage", "set_provider_and_model", "get_span", "submit_feedback"
-]
+# Export provider and exporter for advanced usage
+__all__ = ["provider", "exporter", "flush_spans", "shutdown_tracing", "WithTracing", "set_span_attribute", "set_span_name", "get_active_span"]
 
 
 class TracingOptions:
-    """
-    Options for WithTracing decorator.
-    
-    This class is used to configure how function calls are traced and what data
-    is recorded in span attributes. All fields are optional.
-    """
+    """Options for WithTracing decorator"""
 
     def __init__(
         self,
         name: Optional[str] = None,
-        ignore_input: Optional[List[str]] = None,
-        ignore_output: Optional[List[str]] = None,
+        ignore_input: Optional[Any] = None,
+        ignore_output: Optional[Any] = None,
         filter_input: Optional[Callable[[Any], Any]] = None,
         filter_output: Optional[Callable[[Any], Any]] = None,
     ):
-        """
-        Initialize TracingOptions.
-        
-        Args:
-            name: Custom name for the span. If not provided, the function name
-                will be used. Useful for renaming spans or providing more
-                descriptive names.
-            
-            ignore_input: Iterable of keys (e.g., list, set) to exclude from
-                input data when recording span attributes. Only applies when
-                input is a dictionary. For example, use `["password", "api_key"]`
-                to exclude sensitive fields from being traced.
-            
-            ignore_output: Iterable of keys (e.g., list, set) to exclude from
-                output data when recording span attributes. Only applies when
-                output is a dictionary. Useful for excluding large or sensitive
-                fields from traces.
-            
-            filter_input: Callable function that receives the prepared input data
-                and returns a filtered/transformed version to be recorded in the
-                span. The function should accept one argument (the input data)
-                and return the transformed data. This is applied before
-                ignore_input filtering.
-            
-            filter_output: Callable function that receives the output data and
-                returns a filtered/transformed version to be recorded in the span.
-                The function should accept one argument (the output data) and
-                return the transformed data. This is applied before
-                ignore_output filtering.
-        
-        Example:
-            # Exclude sensitive fields from input
-            @WithTracing(ignore_input=["password", "secret_key"])
-            def authenticate(username, password):
-                return {"token": "..."}
-            
-            # Custom span name and filter output
-            @WithTracing(
-                name="data_processing",
-                filter_output=lambda x: {"count": len(x)} if isinstance(x, list) else x
-            )
-            def process_data(items):
-                return items
-        """
         self.name = name
         self.ignore_input = ignore_input
         self.ignore_output = ignore_output
@@ -127,426 +82,56 @@ class TracingOptions:
         self.filter_output = filter_output
 
 
+def _serialize_for_span(value: Any) -> Any:
+    """
+    Serialize a value for span attributes.
+    OpenTelemetry only accepts primitives (bool, str, bytes, int, float) or sequences of those.
+    Complex types (dicts, lists, objects) are converted to JSON strings.
+    """
+    # Keep primitives as is (including None)
+    if value is None or isinstance(value, (str, int, float, bool, bytes)):
+        return value
+    
+    # For sequences, check if all elements are primitives
+    if isinstance(value, (list, tuple)):
+        # If all elements are primitives, return as list
+        if all(isinstance(item, (str, int, float, bool, bytes, type(None))) for item in value):
+            return list(value)
+        # Otherwise serialize to JSON string
+        try:
+            return json.dumps(value)
+        except (TypeError, ValueError):
+            return str(value)
+    
+    # For dicts and other complex types, serialize to JSON string
+    try:
+        return json.dumps(value)
+    except (TypeError, ValueError):
+        # If JSON serialization fails, convert to string
+        return str(value)
 
 
 def _prepare_input(args: tuple, kwargs: dict) -> Any:
-    """Prepare input for span attributes.
-    
-    Note: This function does NOT serialize values - it just structures the data.
-    Serialization happens later via serialize_for_span() to avoid double-encoding
-    (e.g., converting messages to JSON string, then encoding that string again).
-    """
+    """Prepare input for span attributes."""
     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
+        return _serialize_for_span(args[0])
     # Multiple args or kwargs - combine into dict
     result = {}
     if args:
-        result["args"] = list(args)  # Keep as-is, will be serialized later
+        result["args"] = [_serialize_for_span(arg) for arg in args]
     if kwargs:
-        result["kwargs"] = dict(kwargs)  # Keep as-is, will be serialized later
+        result["kwargs"] = {k: _serialize_for_span(v) for k, v in kwargs.items()}
     return result
 
 
-def _prepare_and_filter_input(
-    args: tuple,
-    kwargs: dict,
-    filter_input: Optional[Callable[[Any], Any]],
-    ignore_input: Optional[List[str]],
-) -> Any:
-    """Prepare and filter input for span attributes."""
-    # 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
-    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"]
-        # Skip first arg if it exists (typically self for bound methods)
-        if args:
-            filtered_args = args[1:]
-        # 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:
-        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]
-    # 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:
-        return None
-    return input_data
-
-
-def _prepare_and_filter_output(
-    result: Any,
-    filter_output: Optional[Callable[[Any], Any]],
-    ignore_output: Optional[List[str]],
-) -> Any:
-    """Prepare and filter output for span attributes."""
-    output_data = result
-    if 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
-
-
-def _handle_span_exception(span: trace.Span, exception: Exception) -> None:
-    """Record exception on span and set error status."""
-    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:
-    """
-    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")
-
-
-class TracedGenerator:
-    """Wrapper for sync generators that traces iteration."""
-    
-    def __init__(
-        self,
-        generator: Any,
-        span: trace.Span,
-        fn_name: str,
-        filter_output: Optional[Callable[[Any], Any]],
-        ignore_output: Optional[List[str]],
-        context_token: Any,
-    ):
-        self._generator = generator
-        self._span = span
-        self._fn_name = fn_name
-        self._filter_output = filter_output
-        self._ignore_output = ignore_output
-        self._context_token = context_token
-        self._yielded_values = []
-        self._exhausted = False
-    
-    def __iter__(self):
-        return self
-    
-    def __next__(self):
-        if self._exhausted:
-            raise StopIteration
-        
-        try:
-            value = next(self._generator)
-            self._yielded_values.append(value)
-            return value
-        except StopIteration:
-            self._exhausted = True
-            self._finalize_span_success()
-            trace.context_api.detach(self._context_token)
-            self._span.end()
-            raise
-        except Exception as exception:
-            self._exhausted = True
-            _handle_span_exception(self._span, exception)
-            trace.context_api.detach(self._context_token)
-            self._span.end()
-            raise
-    
-    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)
-        
-        # Record summary of yielded values
-        output_data = {
-            "type": "generator",
-            "yielded_count": len(self._yielded_values),
-        }
-        
-        # Optionally include sample values (limit to avoid huge spans)
-        if self._yielded_values:
-            sample_size = min(10, len(self._yielded_values))
-            output_data["sample_values"] = [
-                serialize_for_span(v) for v in self._yielded_values[:sample_size]
-            ]
-            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))
-
-
-class TracedAsyncGenerator:
-    """Wrapper for async generators that traces iteration."""
-    
-    def __init__(
-        self,
-        generator: Any,
-        span: trace.Span,
-        fn_name: str,
-        filter_output: Optional[Callable[[Any], Any]],
-        ignore_output: Optional[List[str]],
-        context_token: Any,
-    ):
-        self._generator = generator
-        self._span = span
-        self._fn_name = fn_name
-        self._filter_output = filter_output
-        self._ignore_output = ignore_output
-        self._context_token = context_token
-        self._yielded_values = []
-        self._exhausted = False
-    
-    def __aiter__(self):
-        return self
-    
-    async def __anext__(self):
-        if self._exhausted:
-            raise StopAsyncIteration
-        
-        try:
-            value = await self._generator.__anext__()
-            self._yielded_values.append(value)
-            return value
-        except StopAsyncIteration:
-            self._exhausted = True
-            self._finalize_span_success()
-            trace.context_api.detach(self._context_token)
-            self._span.end()
-            raise
-        except Exception as exception:
-            self._exhausted = True
-            _handle_span_exception(self._span, exception)
-            trace.context_api.detach(self._context_token)
-            self._span.end()
-            raise
-    
-    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)
-        
-        # Record summary of yielded values
-        output_data = {
-            "type": "async_generator",
-            "yielded_count": len(self._yielded_values),
-        }
-        
-        # Optionally include sample values (limit to avoid huge spans)
-        if self._yielded_values:
-            sample_size = min(10, len(self._yielded_values))
-            output_data["sample_values"] = [
-                serialize_for_span(v) for v in self._yielded_values[:sample_size]
-            ]
-            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))
-
-
 def WithTracing(
     func: Optional[Callable] = None,
     *,
     name: Optional[str] = None,
-    ignore_input: Optional[List[str]] = None,
-    ignore_output: Optional[List[str]] = None,
+    ignore_input: Optional[Any] = None,
+    ignore_output: Optional[Any] = None,
     filter_input: Optional[Callable[[Any], Any]] = None,
     filter_output: Optional[Callable[[Any], Any]] = None,
 ):
@@ -554,17 +139,13 @@ def WithTracing(
     Decorator to automatically create spans for function calls.
     Records input/output as span attributes. Spans are automatically linked via OpenTelemetry context.
     
-    Works with synchronous functions, asynchronous functions, generator functions, and async generator functions.
+    Works with both synchronous and asynchronous functions.
     
     Args:
         func: The function to trace (when used as @WithTracing)
         name: Optional custom name for the span (defaults to function name)
-        ignore_input: List of keys to exclude from input data when recording span attributes.
-            Only applies when input is a dictionary. For example, use ["password", "api_key"]
-            to exclude sensitive fields from being traced.
-        ignore_output: List of keys to exclude from output data when recording span attributes.
-            Only applies when output is a dictionary. Useful for excluding large or sensitive
-            fields from traces.
+        ignore_input: Fields to ignore in input (not yet implemented)
+        ignore_output: Fields to ignore in output (not yet implemented)
         filter_input: Function to filter/transform input before recording
         filter_output: Function to filter/transform output before recording
     
@@ -577,16 +158,6 @@ def 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
@@ -600,156 +171,53 @@ def WithTracing(
             return fn
         
         is_async = inspect.iscoroutinefunction(fn)
-        is_generator = inspect.isgeneratorfunction(fn)
-        is_async_generator = inspect.isasyncgenfunction(fn) if hasattr(inspect, 'isasyncgenfunction') else False
-        
-        tracer = get_aiqa_tracer()
-        
-        def _setup_span(span: trace.Span, input_data: Any) -> bool:
-            """Setup span with input data. Returns True if span is recording."""
-            if not span.is_recording():
-                logger.warning(f"Span {fn_name} is not recording - will not be exported")
-                return False
-            
-            logger.debug(f"Span {fn_name} is recording, trace_id={format(span.get_span_context().trace_id, '032x')}")
-            
-            # Set component tag if configured
-            component_tag = get_component_tag()
-            if component_tag:
-                span.set_attribute("gen_ai.component.id", component_tag)
-            
-            if input_data is not None:
-                span.set_attribute("input", serialize_for_span(input_data))
-            
-            trace_id = format(span.get_span_context().trace_id, "032x")
-            logger.debug(f"do traceable stuff {fn_name} {trace_id}")
-            return True
-        
-        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))
-        
-        def _execute_with_span_sync(executor: Callable[[], Any], input_data: Any) -> Any:
-            """Execute sync function within span context, handling input/output and exceptions."""
-            # Ensure tracer provider is initialized before creating spans
-            get_aiqa_client()
-            with tracer.start_as_current_span(fn_name) as span:
-                if not _setup_span(span, input_data):
-                    return executor()
-                
-                try:
-                    result = executor()
-                    _finalize_span_success(span, result)
-                    return result
-                except Exception as exception:
-                    _handle_span_exception(span, exception)
-                    raise
-        
-        async def _execute_with_span_async(executor: Callable[[], Any], input_data: Any) -> Any:
-            """Execute async function within span context, handling input/output and exceptions."""
-            # Ensure tracer provider is initialized before creating spans
-            get_aiqa_client()
-            with tracer.start_as_current_span(fn_name) as span:
-                if not _setup_span(span, input_data):
-                    return await executor()
-                
-                try:
-                    result = await executor()
-                    _finalize_span_success(span, result)
-                    logger.debug(f"Span {fn_name} completed successfully, is_recording={span.is_recording()}")
-                    return result
-                except Exception as exception:
-                    _handle_span_exception(span, exception)
-                    raise
-                finally:
-                    logger.debug(f"Span {fn_name} context exiting, is_recording={span.is_recording()}")
         
-        def _execute_generator_sync(executor: Callable[[], Any], input_data: Any) -> Any:
-            """Execute sync generator function, returning a traced generator."""
-            # Ensure tracer provider is initialized before creating spans
-            get_aiqa_client()
-            # Create span but don't use 'with' - span will be closed by TracedGenerator
-            span = tracer.start_span(fn_name)
-            token = trace.context_api.attach(trace.context_api.set_span_in_context(span))
-            
-            try:
-                if not _setup_span(span, input_data):
-                    generator = executor()
-                    trace.context_api.detach(token)
-                    span.end()
-                    return generator
-                
-                generator = executor()
-                return TracedGenerator(generator, span, fn_name, filter_output, ignore_output, token)
-            except Exception as exception:
-                trace.context_api.detach(token)
-                _handle_span_exception(span, exception)
-                span.end()
-                raise
-        
-        async def _execute_generator_async(executor: Callable[[], Any], input_data: Any) -> Any:
-            """Execute async generator function, returning a traced async generator."""
-            # Ensure tracer provider is initialized before creating spans
-            get_aiqa_client()
-            # Create span but don't use 'with' - span will be closed by TracedAsyncGenerator
-            span = tracer.start_span(fn_name)
-            token = trace.context_api.attach(trace.context_api.set_span_in_context(span))
-            
-            try:
-                if not _setup_span(span, input_data):
-                    generator = executor()
-                    trace.context_api.detach(token)
-                    span.end()
-                    return generator
-                
-                generator = executor()
-                return TracedAsyncGenerator(generator, span, fn_name, filter_output, ignore_output, token)
-            except Exception as exception:
-                trace.context_api.detach(token)
-                _handle_span_exception(span, exception)
-                span.end()
-                raise
-        
-        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)
-                return await _execute_generator_async(
-                    lambda: fn(*args, **kwargs),
-                    input_data
-                )
-            
-            async_gen_traced_fn._is_traced = True
-            logger.debug(f"Function {fn_name} is now traced (async generator)")
-            return async_gen_traced_fn
-        elif is_generator:
-            @wraps(fn)
-            def gen_traced_fn(*args, **kwargs):
-                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
-                return _execute_generator_sync(
-                    lambda: fn(*args, **kwargs),
-                    input_data
-                )
-            
-            gen_traced_fn._is_traced = True
-            logger.debug(f"Function {fn_name} is now traced (generator)")
-            return gen_traced_fn
-        elif is_async:
+        if is_async:
             @wraps(fn)
             async def async_traced_fn(*args, **kwargs):
-                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
-                return await _execute_with_span_async(
-                    lambda: fn(*args, **kwargs),
-                    input_data
-                )
+                # Prepare input
+                input_data = _prepare_input(args, kwargs)
+                if filter_input:
+                    input_data = filter_input(input_data)
+                if ignore_input and isinstance(input_data, dict):                    
+                    for key in ignore_input:
+                        if key in input_data:
+                            del input_data[key]
+                
+                # Use start_as_current_span to ensure span is recorded by BatchSpanProcessor
+                with tracer.start_as_current_span(fn_name) as span:
+                    if input_data is not None:
+                        # Serialize for span attributes (OpenTelemetry only accepts primitives or JSON strings)
+                        serialized_input = _serialize_for_span(input_data)
+                        span.set_attribute("input", serialized_input)
+                    
+                    try:
+                        # Call the function within the span context
+                        trace_id = format(span.get_span_context().trace_id, "032x")
+                        logger.debug(f"do traceable stuff {fn_name} {trace_id}")
+                        
+                        result = await fn(*args, **kwargs)
+                        
+                        # Prepare output
+                        output_data = result
+                        if filter_output:
+                            output_data = filter_output(output_data)                    
+                        if ignore_output and isinstance(output_data, dict):
+                            # Make a copy of output_data to avoid modifying the original
+                            output_data = output_data.copy()
+                            for key in ignore_output:
+                                if key in output_data:
+                                    del output_data[key]
+                        
+                        span.set_attribute("output", _serialize_for_span(output_data))
+                        span.set_status(Status(StatusCode.OK))
+                        
+                        return result
+                    except Exception as exception:
+                        error = exception if isinstance(exception, Exception) else Exception(str(exception))
+                        span.record_exception(error)
+                        span.set_status(Status(StatusCode.ERROR, str(error)))
+                        raise
             
             async_traced_fn._is_traced = True
             logger.debug(f"Function {fn_name} is now traced (async)")
@@ -757,11 +225,49 @@ def WithTracing(
         else:
             @wraps(fn)
             def sync_traced_fn(*args, **kwargs):
-                input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
-                return _execute_with_span_sync(
-                    lambda: fn(*args, **kwargs),
-                    input_data
-                )
+                # Prepare input
+                input_data = _prepare_input(args, kwargs)
+                if filter_input:
+                    input_data = filter_input(input_data)
+                if ignore_input and isinstance(input_data, dict):
+                    for key in ignore_input:
+                        if key in input_data:
+                            del input_data[key]
+                
+                # Use start_as_current_span to ensure span is recorded by BatchSpanProcessor
+                with tracer.start_as_current_span(fn_name) as span:
+                    if input_data is not None:
+                        # Serialize for span attributes (OpenTelemetry only accepts primitives or JSON strings)
+                        serialized_input = _serialize_for_span(input_data)
+                        span.set_attribute("input", serialized_input)
+                    
+                    try:
+                        # Call the function within the span context
+                        trace_id = format(span.get_span_context().trace_id, "032x")
+                        logger.debug(f"do traceable stuff {fn_name} {trace_id}")
+                        
+                        result = fn(*args, **kwargs)
+                        
+                        # Prepare output
+                        output_data = result
+                        if filter_output:
+                            output_data = filter_output(output_data)
+                        if ignore_output and isinstance(output_data, dict):
+                            # Make a copy of output_data to avoid modifying the original
+                            output_data = output_data.copy()
+                            for key in ignore_output:
+                                if key in output_data:
+                                    del output_data[key]
+                        
+                        span.set_attribute("output", _serialize_for_span(output_data))
+                        span.set_status(Status(StatusCode.OK))
+                        
+                        return result
+                    except Exception as exception:
+                        error = exception if isinstance(exception, Exception) else Exception(str(exception))
+                        span.record_exception(error)
+                        span.set_status(Status(StatusCode.ERROR, str(error)))
+                        raise
             
             sync_traced_fn._is_traced = True
             logger.debug(f"Function {fn_name} is now traced (sync)")
@@ -783,7 +289,7 @@ def set_span_attribute(attribute_name: str, attribute_value: Any) -> bool:
     """
     span = trace.get_current_span()
     if span and span.is_recording():
-        span.set_attribute(attribute_name, serialize_for_span(attribute_value))
+        span.set_attribute(attribute_name, _serialize_for_span(attribute_value))
         return True
     return False
 
@@ -801,456 +307,3 @@ def get_active_span() -> Optional[trace.Span]:
     """Get the currently active span."""
     return trace.get_current_span()
 
-
-def set_conversation_id(conversation_id: str) -> bool:
-    """
-    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.
-    
-    Args:
-        conversation_id: A unique identifier for the conversation (e.g., user session ID, chat ID, etc.)
-    
-    Returns:
-        True if gen_ai.conversation.id was set, False if no active span found
-    
-    Example:
-        from aiqa import WithTracing, set_conversation_id
-        
-        @WithTracing
-        def handle_user_request(user_id: str, request: dict):
-            # Set conversation ID to group all traces for this user session
-            set_conversation_id(f"user_{user_id}_session_{request.get('session_id')}")
-            # ... rest of function
-    """
-    return set_span_attribute("gen_ai.conversation.id", conversation_id)
-
-
-def set_token_usage(
-    input_tokens: Optional[int] = None,
-    output_tokens: Optional[int] = None,
-    total_tokens: Optional[int] = None,
-) -> bool:
-    """
-    Set token usage attributes on the active span using OpenTelemetry semantic conventions for gen_ai.
-    This allows you to explicitly record token usage information.
-    AIQA tracing will automatically detect and set token usage from standard OpenAI-like API responses.
-    See https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/ for more details.
-    
-    Args:
-        input_tokens: Number of input tokens used (maps to gen_ai.usage.input_tokens)
-        output_tokens: Number of output tokens generated (maps to gen_ai.usage.output_tokens)
-        total_tokens: Total number of tokens used (maps to gen_ai.usage.total_tokens)
-    
-    Returns:
-        True if at least one token usage attribute was set, False if no active span found
-    
-    Example:
-        from aiqa import WithTracing, set_token_usage
-        
-        @WithTracing
-        def call_llm(prompt: str):
-            response = openai_client.chat.completions.create(...)
-            # Explicitly set token usage
-            set_token_usage(
-                input_tokens=response.usage.prompt_tokens,
-                output_tokens=response.usage.completion_tokens,
-                total_tokens=response.usage.total_tokens
-            )
-            return response
-    """
-    span = trace.get_current_span()
-    if not span or not span.is_recording():
-        return False
-    
-    set_count = 0
-    try:
-        if input_tokens is not None:
-            span.set_attribute("gen_ai.usage.input_tokens", input_tokens)
-            set_count += 1
-        if output_tokens is not None:
-            span.set_attribute("gen_ai.usage.output_tokens", output_tokens)
-            set_count += 1
-        if total_tokens is not None:
-            span.set_attribute("gen_ai.usage.total_tokens", total_tokens)
-            set_count += 1
-    except Exception as e:
-        logger.warning(f"Failed to set token usage attributes: {e}")
-        return False
-    
-    return set_count > 0
-
-
-def set_provider_and_model(
-    provider: Optional[str] = None,
-    model: Optional[str] = None,
-) -> bool:
-    """
-    Set provider and model attributes on the active span using OpenTelemetry semantic conventions for gen_ai.
-    This allows you to explicitly record provider and model information.
-    AIQA tracing will automatically detect and set provider/model from standard API responses.
-    See https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/ for more details.
-    
-    Args:
-        provider: Name of the AI provider (e.g., "openai", "anthropic", "google") (maps to gen_ai.provider.name)
-        model: Name of the model used (e.g., "gpt-4", "claude-3-5-sonnet") (maps to gen_ai.request.model)
-    
-    Returns:
-        True if at least one attribute was set, False if no active span found
-    
-    Example:
-        from aiqa import WithTracing, set_provider_and_model
-        
-        @WithTracing
-        def call_llm(prompt: str):
-            response = openai_client.chat.completions.create(...)
-            # Explicitly set provider and model
-            set_provider_and_model(
-                provider="openai",
-                model=response.model
-            )
-            return response
-    """
-    span = trace.get_current_span()
-    if not span or not span.is_recording():
-        return False
-    
-    set_count = 0
-    try:
-        if provider is not None:
-            span.set_attribute("gen_ai.provider.name", str(provider))
-            set_count += 1
-        if model is not None:
-            span.set_attribute("gen_ai.request.model", str(model))
-            set_count += 1
-    except Exception as e:
-        logger.warning(f"Failed to set provider/model attributes: {e}")
-        return False
-    
-    return set_count > 0
-
-
-def set_component_tag(tag: str) -> None:
-    """
-    Set the component tag that will be added to all spans created by AIQA.
-    This can also be set via the AIQA_COMPONENT_TAG environment variable.
-    The component tag allows you to identify which component/system generated the spans.
-    
-    Note: If using environment variables, ensure you call get_aiqa_client() first to initialize
-    the client and load environment variables.
-    
-    Args:
-        tag: A component identifier (e.g., "mynamespace.mysystem", "backend.api", etc.)
-    
-    Example:
-        from aiqa import get_aiqa_client, set_component_tag, WithTracing
-        
-        # Initialize client (loads env vars including AIQA_COMPONENT_TAG)
-        get_aiqa_client()
-        
-        # Or set component tag programmatically (overrides env var)
-        set_component_tag("mynamespace.mysystem")
-        
-        @WithTracing
-        def my_function():
-            pass
-    """
-    _set_component_tag(tag)
-
-def get_provider() -> Optional[TracerProvider]:
-    """Get the tracer provider for advanced usage."""
-    client = get_aiqa_client()
-    return client.get("provider")
-
-def get_exporter() -> Optional[AIQASpanExporter]:
-    """Get the exporter for advanced usage."""
-    client = get_aiqa_client()
-    return client.get("exporter")
-
-
-def get_trace_id() -> Optional[str]:
-    """
-    Get the current trace ID as a hexadecimal string (32 characters).
-    
-    Returns:
-        The trace ID as a hex string, or None if no active span exists.
-    
-    Example:
-        trace_id = get_trace_id()
-        # Pass trace_id to another service/agent
-        # e.g., include in HTTP headers, message queue metadata, etc.
-    """
-    span = trace.get_current_span()
-    if span and span.get_span_context().is_valid:
-        return format(span.get_span_context().trace_id, "032x")
-    return None
-
-
-def get_span_id() -> Optional[str]:
-    """
-    Get the current span ID as a hexadecimal string (16 characters).
-    
-    Returns:
-        The span ID as a hex string, or None if no active span exists.
-    
-    Example:
-        span_id = get_span_id()
-        # Can be used to create child spans in other services
-    """
-    span = trace.get_current_span()
-    if span and span.get_span_context().is_valid:
-        return format(span.get_span_context().span_id, "016x")
-    return None
-
-
-def create_span_from_trace_id(
-    trace_id: str,
-    parent_span_id: Optional[str] = None,
-    span_name: str = "continued_span",
-) -> trace.Span:
-    """
-    Create a new span that continues from an existing trace ID.
-    This is useful for linking traces across different services or agents.
-    
-    Args:
-        trace_id: The trace ID as a hexadecimal string (32 characters)
-        parent_span_id: Optional parent span ID as a hexadecimal string (16 characters).
-            If provided, the new span will be a child of this span.
-        span_name: Name for the new span (default: "continued_span")
-    
-    Returns:
-        A new span that continues the trace. Use it in a context manager or call end() manually.
-    
-    Example:
-        # In service A: get trace ID
-        trace_id = get_trace_id()
-        span_id = get_span_id()
-        
-        # Send to service B (e.g., via HTTP, message queue, etc.)
-        # ...
-        
-        # In service B: continue the trace
-        with create_span_from_trace_id(trace_id, parent_span_id=span_id, span_name="service_b_operation"):
-            # Your code here
-            pass
-    """
-    try:
-        # Parse trace ID from hex string
-        trace_id_int = int(trace_id, 16)
-        
-        # Parse parent span ID if provided
-        parent_span_id_int = None
-        if parent_span_id:
-            parent_span_id_int = int(parent_span_id, 16)
-        
-        # Create a parent span context
-        parent_span_context = SpanContext(
-            trace_id=trace_id_int,
-            span_id=parent_span_id_int if parent_span_id_int else 0,
-            is_remote=True,
-            trace_flags=TraceFlags(0x01),  # SAMPLED flag
-        )
-        
-        # Create a context with this span context as the parent
-        from opentelemetry.trace import set_span_in_context
-        parent_context = set_span_in_context(trace.NonRecordingSpan(parent_span_context))
-        
-        # Start a new span in this context (it will be a child of the parent span)
-        tracer = get_aiqa_tracer()
-        span = tracer.start_span(span_name, context=parent_context)
-        
-        # Set component tag if configured
-        component_tag = get_component_tag()
-        if component_tag:
-            span.set_attribute("gen_ai.component.id", component_tag)
-        
-        return span
-    except (ValueError, AttributeError) as e:
-        logger.error(f"Error creating span from trace_id: {e}")
-        # Fallback: create a new span
-        tracer = get_aiqa_tracer()
-        span = tracer.start_span(span_name)
-        component_tag = get_component_tag()
-        if component_tag:
-            span.set_attribute("gen_ai.component.id", component_tag)
-        return span
-
-
-def inject_trace_context(carrier: dict) -> None:
-    """
-    Inject the current trace context into a carrier (e.g., HTTP headers).
-    This allows you to pass trace context to another service.
-    
-    Args:
-        carrier: Dictionary to inject trace context into (e.g., HTTP headers dict)
-    
-    Example:
-        import requests
-        
-        headers = {}
-        inject_trace_context(headers)
-        response = requests.get("http://other-service/api", headers=headers)
-    """
-    try:
-        inject(carrier)
-    except Exception as e:
-        logger.warning(f"Error injecting trace context: {e}")
-
-
-def extract_trace_context(carrier: dict) -> Any:
-    """
-    Extract trace context from a carrier (e.g., HTTP headers).
-    Use this to continue a trace that was started in another service.
-    
-    Args:
-        carrier: Dictionary containing trace context (e.g., HTTP headers dict)
-    
-    Returns:
-        A context object that can be used with trace.use_span() or tracer.start_span()
-    
-    Example:
-        from opentelemetry.trace import use_span
-        
-        # Extract context from incoming request headers
-        ctx = extract_trace_context(request.headers)
-        
-        # Use the context to create a span
-        with use_span(ctx):
-            # Your code here
-            pass
-        
-        # Or create a span with the context
-        tracer = get_aiqa_tracer()
-        with tracer.start_as_current_span("operation", context=ctx):
-            # Your code here
-            pass
-    """
-    try:
-        return extract(carrier)
-    except Exception as e:
-        logger.warning(f"Error extracting trace context: {e}")
-        return None
-
-
-def get_span(span_id: str, organisation_id: Optional[str] = None, exclude: Optional[List[str]] = None) -> Optional[dict]:
-    """
-    Get a span by its ID from the AIQA server.
-    
-    Expected usage is: re-playing a specific function call in a unit test (either a developer debugging an issue, or as part of a test suite).
-    
-    Args:
-        span_id: The span ID as a hexadecimal string (16 characters) or client span ID
-        organisation_id: Optional organisation ID. If not provided, will try to get from
-            AIQA_ORGANISATION_ID environment variable. The organisation is typically
-            extracted from the API key during authentication, but the API requires it
-            as a query parameter.
-        exclude: Optional list of fields to exclude from the span data. By default this function WILL return 'attributes' (often large).
-    
-    Returns:
-        The span data as a dictionary, or None if not found
-    
-    Example:
-        from aiqa import get_span
-        
-        span = get_span('abc123...')
-        if span:
-            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", "")
-    org_id = organisation_id or os.getenv("AIQA_ORGANISATION_ID", "")
-    
-    if not 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.")
-    if not api_key:
-        raise ValueError("API key is required. Set AIQA_API_KEY environment variable.")
-    
-    # Try both spanId and clientSpanId queries
-    for query_field in ["spanId", "clientSpanId"]:
-        url = f"{server_url}/span"
-        params = {
-            "q": f"{query_field}:{span_id}",
-            "organisation": org_id,
-            "limit": "1",
-            "exclude": ",".join(exclude) if exclude else None,
-            "fields": "*" if not exclude else None,
-        }
-        
-        headers = {"Content-Type": "application/json"}
-        if api_key:
-            headers["Authorization"] = f"ApiKey {api_key}"
-        
-        response = requests.get(url, params=params, headers=headers)
-        if response.status_code == 200:
-            result = response.json()
-            hits = result.get("hits", [])
-            if hits and len(hits) > 0:
-                return hits[0]
-        elif response.status_code == 404:
-            # Try next query field
-            continue
-        else:
-            error_text = response.text
-            raise ValueError(f"Failed to get span: {response.status_code} - {error_text[:500]}")        
-    # not found
-    return None
-
-
-async def submit_feedback(
-    trace_id: str,
-    thumbs_up: Optional[bool] = None,
-    comment: Optional[str] = None,
-) -> None:
-    """
-    Submit feedback for a trace by creating a new span with the same trace ID.
-    This allows you to add feedback (thumbs-up, thumbs-down, comment) to a trace after it has completed.
-    
-    Args:
-        trace_id: The trace ID as a hexadecimal string (32 characters)
-        thumbs_up: True for positive feedback, False for negative feedback, None for neutral
-        comment: Optional text comment
-    
-    Example:
-        from aiqa import submit_feedback
-        
-        # Submit positive feedback
-        await submit_feedback('abc123...', thumbs_up=True, comment='Great response!')
-        
-        # Submit negative feedback
-        await submit_feedback('abc123...', thumbs_up=False, comment='Incorrect answer')
-    """
-    if not trace_id or len(trace_id) != 32:
-        raise ValueError('Invalid trace ID: must be 32 hexadecimal characters')
-    
-    # Create a span for feedback with the same trace ID
-    span = create_span_from_trace_id(trace_id, span_name='feedback')
-    
-    try:
-        # Set feedback attributes
-        if thumbs_up is not None:
-            span.set_attribute('feedback.thumbs_up', thumbs_up)
-            span.set_attribute('feedback.type', 'positive' if thumbs_up else 'negative')
-        else:
-            span.set_attribute('feedback.type', 'neutral')
-        
-        if comment:
-            span.set_attribute('feedback.comment', comment)
-        
-        # Mark as feedback span
-        span.set_attribute('aiqa.span_type', 'feedback')
-        
-        # End the span
-        span.end()
-        
-        # Flush to ensure it's sent immediately
-        await flush_tracing()
-    except Exception as e:
-        span.end()
-        raise e
-