auditi 0.1.0__py3-none-any.whl

auditi/decorators.py

@@ -0,0 +1,1441 @@
+"""
+Tracing decorators for instrumenting AI agents, tools, and LLM calls.
+
+This version uses the provider abstraction layer for robust multi-provider support.
+"""
+
+import functools
+import json
+import os
+from datetime import datetime
+from uuid import uuid4
+import asyncio
+import inspect
+from typing import Optional, Callable, Any
+
+from .types import TraceInput, SpanInput
+from .context import (
+    get_current_trace,
+    set_current_trace,
+    get_current_span,
+    push_span,
+    pop_span,
+    get_context,
+)
+from .client import get_client
+from .evaluator import BaseEvaluator
+from .providers import detect_provider
+from .events import EventType, INTERNAL_EVENTS, CONTENT_EVENTS
+
+# Debug flag - set via environment variable
+DEBUG = os.getenv("AUDITI_DEBUG", "false").lower() in ("true", "1", "yes")
+
+
+def _debug_log(message: str, data: Any = None) -> None:
+    """Helper function to conditionally log debug information."""
+    if DEBUG:
+        print(f"[Auditi Debug] {message}")
+        if data is not None:
+            try:
+                print(json.dumps(data, indent=2, default=str))
+            except Exception as e:
+                print(f"[Auditi Debug] Could not serialize data: {e}")
+                print(f"Raw data: {data}")
+
+
+def _apply_usage_to_span(span: SpanInput, usage: Any, response: Any = None) -> None:
+    """
+    Apply usage metrics to a span using provider abstraction.
+
+    Args:
+        span: The span to update
+        usage: Raw usage object/dict from API response
+        response: Optional full response object for provider detection
+    """
+    # Detect provider
+    provider = detect_provider(model=span.model, response=response)
+
+    _debug_log(
+        f"Detected provider '{provider.name}' for span '{span.name}'",
+        {"model": span.model, "provider": provider.name},
+    )
+
+    # Extract usage using provider-specific logic
+    input_tokens, output_tokens, total_tokens = provider.extract_usage(usage)
+
+    _debug_log(
+        f"Extracted usage for span '{span.name}':",
+        {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "total_tokens": total_tokens,
+            "provider": provider.name,
+        },
+    )
+
+    # Update span
+    if input_tokens is not None:
+        span.input_tokens = input_tokens
+    if output_tokens is not None:
+        span.output_tokens = output_tokens
+    if total_tokens is not None:
+        span.tokens = total_tokens
+
+    # Calculate cost using provider-specific pricing
+    span.cost = provider.calculate_cost(span.model, input_tokens, output_tokens)
+
+    _debug_log(
+        f"Calculated cost for span '{span.name}':",
+        {"cost": span.cost, "model": span.model, "provider": provider.name},
+    )
+
+
+def _apply_usage_to_trace(trace: TraceInput, usage: Any, model: Optional[str] = None) -> None:
+    """
+    Apply usage metrics to a trace using provider abstraction.
+
+    Args:
+        trace: The trace to update
+        usage: Raw usage object/dict from API response
+        model: Optional model name for provider detection
+    """
+    # Detect provider
+    provider = detect_provider(model=model)
+
+    # Extract usage
+    input_tokens, output_tokens, total_tokens = provider.extract_usage(usage)
+
+    _debug_log(
+        f"Applying usage to trace '{trace.name}':",
+        {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "total_tokens": total_tokens,
+            "current_trace_tokens": trace.total_tokens,
+            "provider": provider.name,
+        },
+    )
+
+    if total_tokens is None:
+        return
+
+    # Accumulate tokens
+    trace.total_tokens = (trace.total_tokens or 0) + total_tokens
+
+    # Calculate incremental cost
+    incremental_cost = provider.calculate_cost(model, input_tokens, output_tokens)
+    trace.cost = (trace.cost or 0.0) + incremental_cost
+
+
+def _execute_as_standalone_trace(
+    func: Callable,
+    args: tuple,
+    kwargs: dict,
+    span_type: str,
+    name: Optional[str],
+    model: Optional[str],
+) -> Any:
+    """
+    Execute a function as a standalone trace (no parent trace required).
+
+    This creates a simple trace with the function result as the only content.
+    Used when @trace_llm or @trace_tool is called with standalone=True
+    outside of a @trace_agent context.
+    """
+    client = get_client()
+    trace_id = uuid4()
+    span_id = uuid4()
+    start_time = datetime.utcnow()
+    func_name = name or func.__name__
+
+    # Extract user input from args/kwargs
+    user_input = ""
+    if args:
+        first_arg = args[0]
+        if isinstance(first_arg, str):
+            user_input = first_arg
+        elif isinstance(first_arg, list):
+            # Could be messages array
+            user_input = str(first_arg)
+        elif isinstance(first_arg, dict):
+            user_input = first_arg.get("content") or first_arg.get("message") or str(first_arg)
+        else:
+            user_input = str(first_arg)
+
+    if not user_input:
+        user_input = kwargs.get("prompt") or kwargs.get("message") or kwargs.get("query") or ""
+
+    # Create trace
+    trace = TraceInput(
+        id=trace_id,
+        start_time=start_time,
+        user_input=user_input,
+        name=func_name,
+        tags=["standalone", span_type],
+    )
+    set_current_trace(trace)
+
+    # Create span
+    span = SpanInput(
+        id=span_id,
+        trace_id=trace_id,
+        name=func_name,
+        span_type=span_type,
+        start_time=start_time,
+        inputs={"prompt": user_input} if user_input else {},
+        model=model,
+    )
+    push_span(span)
+
+    result = None
+    error_msg = None
+
+    try:
+        result = func(*args, **kwargs)
+        print(f"[Auditi] Standalone {span_type} trace captured.")
+
+        # Extract model from response if not set
+        if not span.model:
+            provider = detect_provider(response=result)
+            extracted_model = provider.extract_model(result)
+            if extracted_model:
+                span.model = extracted_model
+
+        # Extract output and usage based on response type
+        if hasattr(result, "choices") and result.choices:
+            # OpenAI-style response
+            try:
+                choice = result.choices[0]
+                if hasattr(choice, "message") and hasattr(choice.message, "content"):
+                    output = str(choice.message.content)
+                elif hasattr(choice, "text"):
+                    output = str(choice.text)
+                else:
+                    output = str(result)
+            except (IndexError, AttributeError):
+                output = str(result)
+
+            trace.assistant_output = output
+            span.outputs = output
+
+            if hasattr(result, "usage") and result.usage:
+                _apply_usage_to_span(span, result.usage, response=result)
+
+        elif hasattr(result, "data"):
+            # Embedding response (OpenAI embeddings have .data)
+            if hasattr(result, "usage"):
+                _apply_usage_to_span(span, result.usage, response=result)
+            span.outputs = f"Generated {len(result.data)} embeddings"
+            trace.assistant_output = span.outputs
+
+        elif isinstance(result, str):
+            trace.assistant_output = result
+            span.outputs = result
+
+        elif isinstance(result, dict):
+            content = (
+                result.get("content") or result.get("text") or result.get("response") or str(result)
+            )
+            trace.assistant_output = str(content)
+            span.outputs = str(content)
+            if "usage" in result:
+                _apply_usage_to_span(span, result["usage"], response=result)
+
+        elif hasattr(result, "content"):
+            trace.assistant_output = str(result.content)
+            span.outputs = str(result.content)
+            if hasattr(result, "usage"):
+                _apply_usage_to_span(span, result.usage, response=result)
+
+        else:
+            trace.assistant_output = str(result) if result else ""
+            span.outputs = str(result) if result else ""
+
+        span.status = "ok"
+
+    except Exception as e:
+        error_msg = str(e)
+        trace.assistant_output = f"Error: {e}"
+        trace.error = error_msg
+        span.error = error_msg
+        span.status = "error"
+        raise
+
+    finally:
+        end_time = datetime.utcnow()
+        span.end_time = end_time
+        trace.end_time = end_time
+
+        pop_span()
+        trace.spans.append(span)
+
+        # Aggregate metrics from span to trace
+        if span.tokens:
+            trace.total_tokens = span.tokens
+        if span.cost:
+            trace.cost = span.cost
+
+        # Send trace
+        trace_payload = trace.model_dump(mode="json")
+        _debug_log(f"Sending standalone trace payload for '{func_name}':", trace_payload)
+        client.transport.send_trace(trace_payload)
+
+        # Clear context
+        from .context import clear_current_trace
+
+        clear_current_trace()
+
+    return result
+
+
+def trace_agent(
+    name: Optional[str] = None,
+    user_id: Optional[str] = None,
+    evaluator: Optional[BaseEvaluator] = None,
+    capture_input: bool = True,
+) -> Callable:
+    """
+    Decorator to trace an entire agent interaction.
+    Supports both sync and async functions, and handles instance methods.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        # Check if the function is async
+        is_async = asyncio.iscoroutinefunction(func)
+
+        if is_async:
+
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs) -> Any:
+                client = get_client()
+                trace_id = uuid4()
+                start_time = datetime.utcnow()
+
+                trace_name = name or func.__name__
+
+                # Smart extraction of context from kwargs or global context
+                ctx = get_context()
+                session_id = (
+                    kwargs.get("session_id")
+                    or kwargs.get("conversation_id")
+                    or ctx.get("session_id")
+                )
+                resolved_user_id = kwargs.get("user_id") or user_id or ctx.get("user_id")
+
+                _debug_log(
+                    f"Starting trace '{trace_name}':",
+                    {
+                        "trace_id": str(trace_id),
+                        "user_id": resolved_user_id,
+                        "session_id": session_id,
+                        "args": [str(arg)[:100] for arg in args],
+                        "kwargs_keys": list(kwargs.keys()),
+                    },
+                )
+
+                # Smart extraction of user input - SKIP SELF/CLS
+                user_input = ""
+                if capture_input and args:
+                    # Determine starting index (skip 'self' or 'cls')
+                    start_idx = 0
+                    if args:
+                        first_arg = args[0]
+                        # Check if first arg is likely 'self' or 'cls'
+                        # It's self/cls if it's an object that's not a basic type
+                        if not isinstance(
+                            first_arg, (str, int, float, bool, list, dict, tuple, type(None))
+                        ):
+                            start_idx = 1
+                            _debug_log(
+                                f"Skipping first argument (detected as self/cls): {type(first_arg)}"
+                            )
+
+                    # Get the actual user input from the correct position
+                    if len(args) > start_idx:
+                        first_arg = args[start_idx]
+                        if isinstance(first_arg, str):
+                            user_input = first_arg
+                        elif isinstance(first_arg, dict) and "message" in first_arg:
+                            user_input = first_arg["message"]
+                        elif isinstance(first_arg, dict) and "content" in first_arg:
+                            user_input = first_arg["content"]
+                        else:
+                            user_input = str(first_arg)
+
+                # Also check for user_input/message/query in kwargs
+                if not user_input:
+                    user_input = (
+                        kwargs.get("user_input")
+                        or kwargs.get("message")
+                        or kwargs.get("query")
+                        or kwargs.get("prompt")
+                        or ""
+                    )
+
+                _debug_log(f"Captured user input:", {"user_input": user_input[:200]})
+
+                # Create Trace
+                trace = TraceInput(
+                    id=trace_id,
+                    user_id=resolved_user_id,
+                    conversation_id=session_id,
+                    start_time=start_time,
+                    user_input=user_input,
+                    name=trace_name,
+                    tags=kwargs.get("tags", []),
+                )
+                set_current_trace(trace)
+
+                result = None
+                error_msg = None
+
+                try:
+                    result = await func(*args, **kwargs)
+                    print("[Auditi] Trace captured.")
+                    if DEBUG:
+                        print(f"result type: {type(result)}")
+                        print(f"result: {str(result)[:200]}")
+
+                    # Smart extraction of assistant output
+                    if isinstance(result, str):
+                        trace.assistant_output = result
+                    elif isinstance(result, dict):
+                        trace.assistant_output = (
+                            result.get("content")
+                            or result.get("message")
+                            or result.get("response")
+                            or str(result)
+                        )
+
+                        # Extract model for provider detection
+                        model = result.get("model")
+
+                        # EXTRACT METRICS from result dict if available
+                        if "usage" in result:
+                            _apply_usage_to_trace(trace, result["usage"], model=model)
+
+                    elif hasattr(result, "content"):
+                        trace.assistant_output = str(result.content)
+
+                        # Extract model from response
+                        model = getattr(result, "model", None)
+                        if model:
+                            model = str(model)
+
+                        # Check for usage on object
+                        if hasattr(result, "usage"):
+                            try:
+                                _apply_usage_to_trace(trace, result.usage, model=model)
+                            except Exception as e:
+                                _debug_log(f"Failed to extract usage from result object: {e}")
+
+                    else:
+                        trace.assistant_output = str(result) if result else ""
+
+                        # Try to extract model and usage
+                        model = getattr(result, "model", None)
+                        if model:
+                            model = str(model)
+
+                        if hasattr(result, "usage"):
+                            try:
+                                _apply_usage_to_trace(trace, result.usage, model=model)
+                            except Exception as e:
+                                _debug_log(f"Failed to extract usage: {e}")
+
+                    _debug_log(
+                        f"Captured assistant output:", {"output": str(trace.assistant_output)[:200]}
+                    )
+
+                except Exception as e:
+                    error_msg = str(e)
+                    trace.assistant_output = f"Error: {e}"
+                    trace.error = error_msg
+                    _debug_log(f"Error in trace '{trace_name}':", {"error": error_msg})
+                    raise
+                finally:
+                    trace.end_time = datetime.utcnow()
+
+                    # Aggregate metrics from spans if not set on trace
+                    if (trace.total_tokens is None or trace.total_tokens == 0) and trace.spans:
+                        calculated_tokens = 0
+                        calculated_cost = 0.0
+                        for s in trace.spans:
+                            if s.tokens:
+                                calculated_tokens += s.tokens
+                            if s.cost:
+                                calculated_cost += s.cost
+
+                        if calculated_tokens > 0:
+                            trace.total_tokens = calculated_tokens
+                            if trace.cost is None or trace.cost == 0.0:
+                                trace.cost = calculated_cost
+
+                    # Run Evaluator if provided and no error
+                    if evaluator and not error_msg:
+                        try:
+                            eval_result = evaluator.evaluate(trace)
+                            # FIX: Map evaluation result to TraceInput fields
+                            trace.status = eval_result.status
+                            trace.score = eval_result.score
+                            if eval_result.reason:
+                                trace.eval_reason = eval_result.reason
+                            if eval_result.failure_mode:
+                                trace.failure_mode = eval_result.failure_mode
+
+                            _debug_log("Evaluation result:", eval_result)
+                        except Exception as e:
+                            print(f"[Auditi] Evaluator failed: {e}")
+
+                    # Prepare and log payload
+                    trace_payload = trace.model_dump(mode="json")
+                    _debug_log(f"Sending trace payload for '{trace_name}':", trace_payload)
+
+                    # Send Trace
+                    client.transport.send_trace(trace_payload)
+
+                return result
+
+            return async_wrapper
+
+        elif inspect.isasyncgenfunction(func):
+
+            @functools.wraps(func)
+            async def async_gen_wrapper(*args, **kwargs) -> Any:
+                client = get_client()
+                trace_id = uuid4()
+                start_time = datetime.utcnow()
+
+                trace_name = name or func.__name__
+
+                # Smart extraction of context from kwargs or global context
+                ctx = get_context()
+                session_id = (
+                    kwargs.get("session_id")
+                    or kwargs.get("conversation_id")
+                    or ctx.get("session_id")
+                )
+                resolved_user_id = kwargs.get("user_id") or user_id or ctx.get("user_id")
+
+                _debug_log(
+                    f"Starting async generator trace '{trace_name}':",
+                    {
+                        "trace_id": str(trace_id),
+                        "user_id": resolved_user_id,
+                        "session_id": session_id,
+                        "args": [str(arg)[:100] for arg in args],
+                        "kwargs_keys": list(kwargs.keys()),
+                    },
+                )
+
+                # Smart extraction of user input - SKIP SELF/CLS
+                user_input = ""
+                if capture_input and args:
+                    # Determine starting index (skip 'self' or 'cls')
+                    start_idx = 0
+                    if args:
+                        first_arg = args[0]
+                        # Check if first arg is likely 'self' or 'cls'
+                        # It's self/cls if it's an object that's not a basic type
+                        if not isinstance(
+                            first_arg, (str, int, float, bool, list, dict, tuple, type(None))
+                        ):
+                            start_idx = 1
+
+                    # Get the actual user input from the correct position
+                    if len(args) > start_idx:
+                        first_arg = args[start_idx]
+                        if isinstance(first_arg, str):
+                            user_input = first_arg
+                        elif isinstance(first_arg, dict) and "message" in first_arg:
+                            user_input = first_arg["message"]
+                        elif isinstance(first_arg, dict) and "content" in first_arg:
+                            user_input = first_arg["content"]
+                        else:
+                            user_input = str(first_arg)
+
+                # Also check for user_input/message/query in kwargs
+                if not user_input:
+                    user_input = (
+                        kwargs.get("user_input")
+                        or kwargs.get("message")
+                        or kwargs.get("query")
+                        or kwargs.get("prompt")
+                        or ""
+                    )
+
+                _debug_log(f"Captured user input:", {"user_input": user_input[:200]})
+
+                # Create Trace
+                trace = TraceInput(
+                    id=trace_id,
+                    user_id=resolved_user_id,
+                    conversation_id=session_id,
+                    start_time=start_time,
+                    user_input=user_input,
+                    name=trace_name,
+                    tags=kwargs.get("tags", []),
+                )
+                set_current_trace(trace)
+
+                error_msg = None
+
+                # Accumulators for streaming response
+                accumulated_content = []
+                final_content = None
+                accumulated_model = None
+                accumulated_usage = None
+
+                try:
+                    # Call the function to get the generator
+                    gen = func(*args, **kwargs)
+
+                    # Iterate through generator
+                    async for item in gen:
+                        yield item
+
+                        # Process item for trace
+                        if isinstance(item, str):
+                            accumulated_content.append(item)
+                        elif isinstance(item, dict):
+                            # Handle different event types
+                            evt_type = item.get("type")
+
+                            # Token event
+                            if evt_type == "token" and "content" in item:
+                                accumulated_content.append(str(item["content"]))
+
+                            # Complete event (often holds final answer/usage)
+                            elif evt_type == "complete":
+                                if "content" in item:
+                                    final_content = str(item["content"])
+                                if "usage" in item:
+                                    accumulated_usage = item["usage"]
+                                if "model" in item:
+                                    accumulated_model = item["model"]
+
+                            # Other events with content (fallback)
+                            elif "content" in item:
+                                # Start of new phase or just info, maybe not content to accumulate?
+                                # Generative agents often send content chunks in 'content' field
+                                if evt_type not in (
+                                    "tool_exec_start",
+                                    "tool_exec_end",
+                                    "phase_start",
+                                    "phase_end",
+                                ):
+                                    if isinstance(item["content"], str):
+                                        # Be careful not to double count if 'token' events are used
+                                        pass
+
+                        elif hasattr(item, "content"):
+                            accumulated_content.append(str(item.content))
+
+                    # Determine final output and usage
+                    if final_content is not None:
+                        trace.assistant_output = final_content
+                    else:
+                        trace.assistant_output = "".join(accumulated_content)
+
+                    # Apply usage if found
+                    if accumulated_usage:
+                        try:
+                            _apply_usage_to_trace(trace, accumulated_usage, model=accumulated_model)
+                        except Exception:
+                            pass
+
+                    print("[Auditi] Async generator trace captured.")
+                    _debug_log(
+                        f"Captured assistant output:", {"output": str(trace.assistant_output)[:200]}
+                    )
+
+                except Exception as e:
+                    error_msg = str(e)
+                    trace.assistant_output = f"Error: {e}"
+                    trace.error = error_msg
+                    _debug_log(f"Error in trace '{trace_name}':", {"error": error_msg})
+                    raise
+                finally:
+                    trace.end_time = datetime.utcnow()
+
+                    # Aggregate metrics from spans if not set on trace
+                    if (trace.total_tokens is None or trace.total_tokens == 0) and trace.spans:
+                        calculated_tokens = 0
+                        calculated_cost = 0.0
+                        for s in trace.spans:
+                            if s.tokens:
+                                calculated_tokens += s.tokens
+                            if s.cost:
+                                calculated_cost += s.cost
+
+                        if calculated_tokens > 0:
+                            trace.total_tokens = calculated_tokens
+                            if trace.cost is None or trace.cost == 0.0:
+                                trace.cost = calculated_cost
+
+                    # Run Evaluator if provided and no error
+                    if evaluator and not error_msg:
+                        try:
+                            eval_result = evaluator.evaluate(trace)
+                            # FIX: Map evaluation result to TraceInput fields
+                            trace.status = eval_result.status
+                            trace.score = eval_result.score
+                            if eval_result.reason:
+                                trace.eval_reason = eval_result.reason
+                            if eval_result.failure_mode:
+                                trace.failure_mode = eval_result.failure_mode
+
+                            _debug_log("Evaluation result:", eval_result)
+                        except Exception as e:
+                            print(f"[Auditi] Evaluator failed: {e}")
+
+                    # Prepare and log payload
+                    trace_payload = trace.model_dump(mode="json")
+                    _debug_log(f"Sending trace payload for '{trace_name}':", trace_payload)
+
+                    # Send Trace
+                    client.transport.send_trace(trace_payload)
+
+            return async_gen_wrapper
+
+        else:
+            # Original sync wrapper (adapted with new improvements)
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs) -> Any:
+                client = get_client()
+                trace_id = uuid4()
+                start_time = datetime.utcnow()
+
+                trace_name = name or func.__name__
+
+                # Smart extraction of context from kwargs or global context
+                ctx = get_context()
+                session_id = (
+                    kwargs.get("session_id")
+                    or kwargs.get("conversation_id")
+                    or ctx.get("session_id")
+                )
+                resolved_user_id = kwargs.get("user_id") or user_id or ctx.get("user_id")
+
+                _debug_log(
+                    f"Starting trace '{trace_name}':",
+                    {
+                        "trace_id": str(trace_id),
+                        "user_id": resolved_user_id,
+                        "session_id": session_id,
+                        "args": [str(arg)[:100] for arg in args],
+                        "kwargs_keys": list(kwargs.keys()),
+                    },
+                )
+
+                # Smart extraction of user input - SKIP SELF/CLS
+                user_input = ""
+                if capture_input and args:
+                    # Determine starting index (skip 'self' or 'cls')
+                    start_idx = 0
+                    if args:
+                        first_arg = args[0]
+                        # Check if first arg is likely 'self' or 'cls'
+                        # It's self/cls if it's an object that's not a basic type
+                        if not isinstance(
+                            first_arg, (str, int, float, bool, list, dict, tuple, type(None))
+                        ):
+                            start_idx = 1
+                            _debug_log(
+                                f"Skipping first argument (detected as self/cls): {type(first_arg)}"
+                            )
+
+                    # Get the actual user input from the correct position
+                    if len(args) > start_idx:
+                        first_arg = args[start_idx]
+                        if isinstance(first_arg, str):
+                            user_input = first_arg
+                        elif isinstance(first_arg, dict) and "message" in first_arg:
+                            user_input = first_arg["message"]
+                        elif isinstance(first_arg, dict) and "content" in first_arg:
+                            user_input = first_arg["content"]
+                        else:
+                            user_input = str(first_arg)
+
+                # Also check for user_input/message/query in kwargs
+                if not user_input:
+                    user_input = (
+                        kwargs.get("user_input")
+                        or kwargs.get("message")
+                        or kwargs.get("query")
+                        or kwargs.get("prompt")
+                        or ""
+                    )
+
+                _debug_log(f"Captured user input:", {"user_input": user_input[:200]})
+
+                # Create Trace
+                trace = TraceInput(
+                    id=trace_id,
+                    user_id=resolved_user_id,
+                    conversation_id=session_id,
+                    start_time=start_time,
+                    user_input=user_input,
+                    name=trace_name,
+                    tags=kwargs.get("tags", []),
+                )
+                set_current_trace(trace)
+
+                result = None
+                error_msg = None
+
+                try:
+                    result = func(*args, **kwargs)
+                    print("[Auditi] Trace captured.")
+                    if DEBUG:
+                        print(f"result: {result}")
+
+                    # Smart extraction of assistant output
+                    if isinstance(result, str):
+                        trace.assistant_output = result
+                    elif isinstance(result, dict):
+                        trace.assistant_output = (
+                            result.get("content")
+                            or result.get("message")
+                            or result.get("response")
+                            or str(result)
+                        )
+
+                        # Extract model for provider detection
+                        model = result.get("model")
+
+                        # EXTRACT METRICS from result dict if available
+                        if "usage" in result:
+                            _apply_usage_to_trace(trace, result["usage"], model=model)
+
+                    elif hasattr(result, "content"):
+                        trace.assistant_output = str(result.content)
+
+                        # Extract model from response
+                        model = getattr(result, "model", None)
+                        if model:
+                            model = str(model)
+
+                        # Check for usage on object
+                        if hasattr(result, "usage"):
+                            try:
+                                _apply_usage_to_trace(trace, result.usage, model=model)
+                            except Exception as e:
+                                _debug_log(f"Failed to extract usage from result object: {e}")
+
+                    else:
+                        trace.assistant_output = str(result) if result else ""
+
+                        # Try to extract model and usage
+                        model = getattr(result, "model", None)
+                        if model:
+                            model = str(model)
+
+                        if hasattr(result, "usage"):
+                            try:
+                                _apply_usage_to_trace(trace, result.usage, model=model)
+                            except Exception as e:
+                                _debug_log(f"Failed to extract usage: {e}")
+
+                    _debug_log(
+                        f"Captured assistant output:", {"output": str(trace.assistant_output)[:200]}
+                    )
+
+                except Exception as e:
+                    error_msg = str(e)
+                    trace.assistant_output = f"Error: {e}"
+                    trace.error = error_msg
+                    _debug_log(f"Error in trace '{trace_name}':", {"error": error_msg})
+                    raise
+                finally:
+                    trace.end_time = datetime.utcnow()
+
+                    # Aggregate metrics from spans if not set on trace
+                    if (trace.total_tokens is None or trace.total_tokens == 0) and trace.spans:
+                        calculated_tokens = 0
+                        calculated_cost = 0.0
+                        for s in trace.spans:
+                            if s.tokens:
+                                calculated_tokens += s.tokens
+                            if s.cost:
+                                calculated_cost += s.cost
+
+                        if calculated_tokens > 0:
+                            trace.total_tokens = calculated_tokens
+                            if trace.cost is None or trace.cost == 0.0:
+                                trace.cost = calculated_cost
+
+                    # Run Evaluator if provided and no error
+                    if evaluator and not error_msg:
+                        try:
+                            eval_result = evaluator.evaluate(trace)
+                            # FIX: Map evaluation result to TraceInput fields
+                            trace.status = eval_result.status
+                            trace.score = eval_result.score
+                            if eval_result.reason:
+                                trace.eval_reason = eval_result.reason
+                            if eval_result.failure_mode:
+                                trace.failure_mode = eval_result.failure_mode
+
+                            _debug_log("Evaluation result:", eval_result)
+                        except Exception as e:
+                            print(f"[Auditi] Evaluator failed: {e}")
+
+                    # Prepare and log payload
+                    trace_payload = trace.model_dump(mode="json")
+                    _debug_log(f"Sending trace payload for '{trace_name}':", trace_payload)
+
+                    # Send Trace
+                    client.transport.send_trace(trace_payload)
+
+                return result
+
+            return wrapper
+
+    return decorator
+
+
+def _capture_inputs(func, args, kwargs, model_hint=None):
+    """
+    Helper to capture and normalize inputs from function arguments.
+    Handles 'self' skipping, kwarg flattening, and fallback str conversion.
+    """
+    inputs = {}
+    effective_model = model_hint
+
+    try:
+        import inspect
+
+        sig = inspect.signature(func)
+        bound = sig.bind(*args, **kwargs)
+        bound.apply_defaults()
+
+        # Capture all arguments as inputs
+        for arg_name, value in bound.arguments.items():
+            # Skip 'self' or 'cls' typically found in methods
+            if arg_name in ("self", "cls"):
+                continue
+
+            # Flatten **kwargs if they exist and are a dict
+            param = sig.parameters.get(arg_name)
+            if param and param.kind == inspect.Parameter.VAR_KEYWORD and isinstance(value, dict):
+                inputs.update({k: str(v) for k, v in value.items()})
+            else:
+                inputs[arg_name] = str(value)
+
+        # Check for model in arguments
+        if not effective_model:
+            if "model" in inputs:
+                effective_model = inputs["model"]
+            elif "model_id" in inputs:  # ADD THIS
+                effective_model = inputs["model_id"]
+            # Also check original bound args just in case normalization changed something
+            elif "model" in bound.arguments:
+                effective_model = str(bound.arguments["model"])
+            elif "model_id" in bound.arguments:  # ADD THIS
+                effective_model = str(bound.arguments["model_id"])
+
+    except Exception:
+        # Binding failed, ignore
+        pass
+
+    # Fallback model detection
+    if not effective_model:
+        if "model" in kwargs:
+            effective_model = str(kwargs["model"])
+        elif "model_id" in kwargs:  # ADD THIS
+            effective_model = str(kwargs["model_id"])
+        elif args and hasattr(args[0], "model"):
+            effective_model = str(args[0].model)
+        elif args and hasattr(args[0], "model_id"):  # ADD THIS
+            effective_model = str(args[0].model_id)
+
+    # Fallback input detection
+    if not inputs:
+        if args:
+            first_arg = args[0]
+            # Skip likely self/cls in fallback mode too (basic heuristic check)
+            start_idx = 0
+            if not isinstance(first_arg, (str, int, float, bool, list, dict, type(None))):
+                start_idx = 1
+
+            if len(args) > start_idx:
+                val = args[start_idx]
+                if isinstance(val, str):
+                    inputs["prompt"] = val
+                elif isinstance(val, (dict, list)):
+                    inputs["data"] = val
+                else:
+                    inputs["input"] = str(val)
+
+        if kwargs:
+            inputs.update({k: str(v) for k, v in kwargs.items()})
+
+    return inputs, effective_model
+
+
+def _trace_span(
+    span_type: str,
+    name: Optional[str] = None,
+    model: Optional[str] = None,
+    standalone: bool = False,
+) -> Callable:
+    """
+    Internal helper to create span tracing decorators.
+
+    Args:
+        span_type: Type of span ("llm", "tool", "embedding", etc.)
+        name: Optional custom name for the span
+        model: Optional model name
+        standalone: If True, creates a standalone trace when no parent trace exists
+    """
+
+    def decorator(func: Callable) -> Callable:
+        # Check for async generator first
+        if inspect.isasyncgenfunction(func):
+
+            @functools.wraps(func)
+            async def async_gen_span_wrapper(*args, **kwargs) -> Any:
+                trace = get_current_trace()
+                if not trace:
+                    if standalone:
+                        # Create a standalone trace for this call
+                        # Note: Simple detection for standalone async gen not fully implemented in _execute_as_standalone_trace yet
+                        # For now, falling back to treating it as regular func (may need enhancement)
+                        async for item in func(*args, **kwargs):
+                            yield item
+                    else:
+                        async for item in func(*args, **kwargs):
+                            yield item
+                    return
+
+                parent = get_current_span()
+                span_id = uuid4()
+                start_time = datetime.utcnow()
+                span_name = name or func.__name__
+
+                # Use robust input capture
+                inputs, effective_model = _capture_inputs(func, args, kwargs, model)
+
+                _debug_log(
+                    f"Starting async generator span '{span_name}' (type: {span_type}):",
+                    {
+                        "span_id": str(span_id),
+                        "parent_id": str(parent.id) if parent else None,
+                        "model": effective_model,
+                        "trace_id": str(trace.id),
+                    },
+                )
+
+                span = SpanInput(
+                    id=span_id,
+                    trace_id=trace.id,
+                    parent_id=parent.id if parent else None,
+                    name=span_name,
+                    span_type=span_type,
+                    start_time=start_time,
+                    inputs=inputs,
+                    model=effective_model,
+                )
+
+                push_span(span)
+
+                accumulated_outputs = []
+
+                try:
+                    # Execute async generator
+                    gen = func(*args, **kwargs)
+
+                    async for item in gen:
+                        yield item
+                        # Only accumulate relevant content
+                        if isinstance(item, str):
+                            accumulated_outputs.append(item)
+                        elif isinstance(item, dict):
+                            evt_type = item.get("type")
+
+                            # Filter out internal events using standardized set
+                            if evt_type in {e.value for e in INTERNAL_EVENTS}:
+                                # Handle turn_metadata to extract usage before skipping
+                                if evt_type == EventType.TURN_METADATA.value:
+                                    if "usage" in item and item["usage"]:
+                                        _apply_usage_to_span(span, item["usage"])
+                                continue
+
+                            # Accumulate content from content events
+                            if evt_type == EventType.TOKEN.value and "content" in item:
+                                accumulated_outputs.append(str(item["content"]))
+                            elif evt_type == EventType.COMPLETE.value and "content" in item:
+                                pass  # Complete event content handled elsewhere
+                            elif "content" in item:
+                                accumulated_outputs.append(str(item["content"]))
+
+                        elif hasattr(item, "content"):
+                            accumulated_outputs.append(str(item.content))
+                        else:
+                            accumulated_outputs.append(str(item))
+
+                    # If usage was not found in stream but captured via other means (e.g. side channel), we rely on it being yields.
+
+                    span.outputs = "".join(accumulated_outputs)
+                    span.status = "ok"
+
+                except Exception as e:
+                    span.error = str(e)
+                    span.status = "error"
+                    _debug_log(f"Span '{span_name}' failed:", {"error": str(e)})
+                    # Don't raise here if we are yielding? Actually we should raise to propagate error
+                    raise
+                finally:
+                    span.end_time = datetime.utcnow()
+                    pop_span()
+
+                    span_payload = span.model_dump(mode="json")
+                    _debug_log(f"Adding span '{span_name}' to trace:", span_payload)
+                    trace.spans.append(span)
+
+            return async_gen_span_wrapper
+
+        # Check for regular async function
+        elif asyncio.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_span_wrapper(*args, **kwargs) -> Any:
+                trace = get_current_trace()
+                if not trace:
+                    if standalone:
+                        # Standalone async trace
+                        return _execute_as_standalone_trace(
+                            func, args, kwargs, span_type, name, model
+                        )
+                    else:
+                        return await func(*args, **kwargs)
+
+                parent = get_current_span()
+                span_id = uuid4()
+                start_time = datetime.utcnow()
+                span_name = name or func.__name__
+
+                # Use robust input capture
+                inputs, effective_model = _capture_inputs(func, args, kwargs, model)
+
+                span = SpanInput(
+                    id=span_id,
+                    trace_id=trace.id,
+                    parent_id=parent.id if parent else None,
+                    name=span_name,
+                    span_type=span_type,
+                    start_time=start_time,
+                    inputs=inputs,
+                    model=effective_model,
+                )
+
+                push_span(span)
+
+                try:
+                    result = await func(*args, **kwargs)
+
+                    # Capture output
+                    if isinstance(result, str):
+                        span.outputs = result
+                    elif hasattr(result, "content"):
+                        span.outputs = str(result.content)
+                    else:
+                        span.outputs = str(result)
+
+                    span.status = "ok"
+                    return result
+
+                except Exception as e:
+                    span.error = str(e)
+                    span.status = "error"
+                    raise
+                finally:
+                    span.end_time = datetime.utcnow()
+                    pop_span()
+
+                    span_payload = span.model_dump(mode="json")
+                    _debug_log(f"Adding span '{span_name}' to trace:", span_payload)
+                    trace.spans.append(span)
+
+            return async_span_wrapper
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs) -> Any:
+            trace = get_current_trace()
+            if not trace:
+                if standalone:
+                    # Create a standalone trace for this call
+                    return _execute_as_standalone_trace(func, args, kwargs, span_type, name, model)
+                else:
+                    # Original behavior: just run the function normally
+                    return func(*args, **kwargs)
+
+            parent = get_current_span()
+            span_id = uuid4()
+            start_time = datetime.utcnow()
+            span_name = name or func.__name__
+
+            # Auto-detect model and capture inputs using signature binding if available
+            effective_model = model
+            inputs = {}
+
+            try:
+                import inspect
+
+                sig = inspect.signature(func)
+                bound = sig.bind(*args, **kwargs)
+                bound.apply_defaults()
+
+                # Capture all arguments as inputs
+                for arg_name, value in bound.arguments.items():
+                    # Skip 'self' or 'cls' typically found in methods
+                    if arg_name in ("self", "cls"):
+                        continue
+                    # Flatten **kwargs if they exist and are a dict
+                    param = sig.parameters.get(arg_name)
+                    if (
+                        param
+                        and param.kind == inspect.Parameter.VAR_KEYWORD
+                        and isinstance(value, dict)
+                    ):
+                        inputs.update({k: str(v) for k, v in value.items()})
+                    else:
+                        inputs[arg_name] = str(value)
+
+                # Check for model in arguments
+                if not effective_model:
+                    if "model" in kwargs:
+                        effective_model = str(kwargs["model"])
+                    elif "model_id" in kwargs:  # ADD THIS
+                        effective_model = str(kwargs["model_id"])
+                    elif args and hasattr(args[0], "model"):
+                        effective_model = str(args[0].model)
+                    elif args and hasattr(args[0], "model_id"):  # ADD THIS
+                        effective_model = str(args[0].model_id)
+
+            except Exception:
+                # If binding fails, valid case for built-ins or certain wrappers
+                pass
+
+            # Fallback model detection (if signature extraction failed or didn't find model)
+            if not effective_model:
+                if "model" in kwargs:
+                    effective_model = str(kwargs["model"])
+                elif args and hasattr(args[0], "model"):
+                    effective_model = str(args[0].model)
+
+            # Fallback input detection (if signature extraction failed or resulted in empty inputs)
+            if not inputs:
+                if args:
+                    first_arg = args[0]
+                    if isinstance(first_arg, str):
+                        inputs["prompt"] = first_arg
+                    elif isinstance(first_arg, (dict, list)):
+                        inputs["data"] = first_arg
+                    else:
+                        inputs["input"] = str(first_arg)
+                if kwargs:
+                    inputs.update({k: str(v) for k, v in kwargs.items()})
+
+            _debug_log(
+                f"Starting span '{span_name}' (type: {span_type}):",
+                {
+                    "span_id": str(span_id),
+                    "parent_id": str(parent.id) if parent else None,
+                    "model": effective_model,
+                    "trace_id": str(trace.id),
+                    "inputs_keys": list(inputs.keys()),
+                },
+            )
+
+            span = SpanInput(
+                id=span_id,
+                trace_id=trace.id,
+                parent_id=parent.id if parent else None,
+                name=span_name,
+                span_type=span_type,
+                start_time=start_time,
+                inputs=inputs,
+                model=effective_model,
+            )
+
+            push_span(span)
+
+            try:
+                result = func(*args, **kwargs)
+
+                # Use provider abstraction to extract model if not set
+                if not span.model:
+                    provider = detect_provider(response=result)
+                    extracted_model = provider.extract_model(result)
+                    if extracted_model:
+                        span.model = extracted_model
+
+                # Smart output capture for various LLM response types
+
+                # OpenAI ChatCompletion object (has choices array)
+                if hasattr(result, "choices") and result.choices:
+                    try:
+                        choice = result.choices[0]
+                        if hasattr(choice, "message") and hasattr(choice.message, "content"):
+                            span.outputs = str(choice.message.content)
+                        elif hasattr(choice, "text"):  # Legacy completions API
+                            span.outputs = str(choice.text)
+                    except (IndexError, AttributeError):
+                        span.outputs = str(result)
+
+                    # Extract usage from response using provider abstraction
+                    if hasattr(result, "usage") and result.usage:
+                        _apply_usage_to_span(span, result.usage, response=result)
+
+                # Simple string result
+                elif isinstance(result, str):
+                    span.outputs = result
+
+                # Object with .content attribute (e.g., Anthropic)
+                elif hasattr(result, "content"):
+                    span.outputs = str(result.content)
+
+                    # Extract usage
+                    if hasattr(result, "usage"):
+                        try:
+                            _apply_usage_to_span(span, result.usage, response=result)
+                        except Exception as e:
+                            _debug_log(f"Failed to extract usage from content object: {e}")
+
+                # Dict result (e.g., from custom LLM wrappers)
+                elif isinstance(result, dict):
+                    # Try to extract content from common keys
+                    content = (
+                        result.get("content")
+                        or result.get("text")
+                        or result.get("message")
+                        or str(result)
+                    )
+                    span.outputs = str(content)
+
+                    # Extract usage
+                    if "usage" in result:
+                        _apply_usage_to_span(span, result["usage"], response=result)
+
+                # Fallback for unknown types
+                else:
+                    span.outputs = str(result)
+
+                    # Try to extract usage
+                    if hasattr(result, "usage"):
+                        try:
+                            _apply_usage_to_span(span, result.usage, response=result)
+                        except Exception as e:
+                            _debug_log(f"Failed to extract usage from unknown type: {e}")
+
+                span.status = "ok"
+
+                _debug_log(
+                    f"Span '{span_name}' completed successfully:",
+                    {
+                        "output_length": len(str(span.outputs)) if span.outputs else 0,
+                        "model": span.model,
+                        "tokens": span.tokens if hasattr(span, "tokens") and span.tokens else None,
+                        "cost": span.cost if hasattr(span, "cost") else None,
+                    },
+                )
+
+                return result
+            except Exception as e:
+                span.error = str(e)
+                span.status = "error"
+                _debug_log(f"Span '{span_name}' failed:", {"error": str(e)})
+                raise
+            finally:
+                span.end_time = datetime.utcnow()
+                pop_span()
+
+                # Log span payload before adding to trace
+                span_payload = span.model_dump(mode="json")
+                _debug_log(f"Adding span '{span_name}' to trace:", span_payload)
+
+                trace.spans.append(span)
+
+        return wrapper
+
+    return decorator
+
+
+def trace_tool(name: Optional[str] = None, standalone: bool = False) -> Callable:
+    """
+    Decorator to trace a tool/function call within an agent.
+
+    Args:
+        name: Optional custom name for the tool span
+        standalone: If True, creates its own trace when not inside @trace_agent
+
+    Example:
+        >>> @trace_tool("database_search")
+        ... def search_db(query: str) -> list:
+        ...     return db.search(query)
+
+        >>> # Standalone tool call (creates its own trace)
+        >>> @trace_tool(standalone=True)
+        ... def standalone_tool(data: str) -> str:
+        ...     return process(data)
+    """
+    return _trace_span("tool", name, standalone=standalone)
+
+
+def trace_llm(
+    name: Optional[str] = None, model: Optional[str] = None, standalone: bool = False
+) -> Callable:
+    """
+    Decorator to trace an LLM call within an agent.
+
+    Args:
+        name: Optional custom name for the LLM span
+        model: Optional model name (auto-detected from response if not provided)
+        standalone: If True, creates its own trace when not inside @trace_agent
+
+    Example:
+        >>> @trace_llm("generate_response", model="gpt-4")
+        ... def call_gpt(prompt: str) -> str:
+        ...     return openai.chat(prompt)
+
+        >>> # Standalone LLM call (creates its own trace)
+        >>> @trace_llm(standalone=True)
+        ... def simple_chat(prompt: str) -> str:
+        ...     return openai.chat(prompt)
+    """
+    return _trace_span("llm", name, model, standalone=standalone)
+
+
+def trace_embedding(name: Optional[str] = None, model: Optional[str] = None) -> Callable:
+    """
+    Decorator to trace an embedding call. Always creates a standalone trace
+    when not inside @trace_agent (embeddings are typically standalone operations).
+
+    Args:
+        name: Optional custom name for the embedding span
+        model: Optional model name (auto-detected from response if not provided)
+
+    Example:
+        >>> @trace_embedding()
+        ... def embed_text(text: str) -> list:
+        ...     return openai.embeddings.create(input=text, model="text-embedding-3-small")
+    """
+    return _trace_span("embedding", name, model, standalone=True)
+
+
+def trace_retrieval(name: Optional[str] = None) -> Callable:
+    """
+    Decorator to trace a retrieval/search operation (e.g., vector DB search).
+    Always creates a standalone trace when not inside @trace_agent.
+
+    Args:
+        name: Optional custom name for the retrieval span
+
+    Example:
+        >>> @trace_retrieval("vector_search")
+        ... def search_docs(query: str) -> list:
+        ...     return vector_db.similarity_search(query)
+    """
+    return _trace_span("retrieval", name, standalone=True)