ai-pipeline-core 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

ai_pipeline_core/{tracing.py → observability/tracing.py}

@@ -6,18 +6,18 @@ This module centralizes:
   ``observe`` instrumentation, and optional support for test runs.
 """
 
+import contextlib
 import inspect
 import json
 import os
+from collections.abc import Callable
 from functools import wraps
-from typing import Any, Callable, Literal, ParamSpec, TypeVar, cast, overload
+from typing import Any, Literal, ParamSpec, TypeVar, cast, overload
 
 from lmnr import Attributes, Instruments, Laminar, observe
 from pydantic import BaseModel, Field
 
-# Import for document trimming - needed for isinstance checks
-# These are lazy imports only used when trim_documents is enabled
-from ai_pipeline_core.documents import Document, DocumentList
+from ai_pipeline_core.documents import Document
 from ai_pipeline_core.llm import AIMessages, ModelResponse
 from ai_pipeline_core.settings import settings
 
@@ -40,144 +40,112 @@ Values:
 # ---------------------------------------------------------------------------
 # Serialization helpers
 # ---------------------------------------------------------------------------
-def _serialize_for_tracing(obj: Any) -> Any:
-    """Convert objects to JSON-serializable format for tracing.
-
-    Handles Pydantic models, Documents, and other special types.
-    This is extracted for better testability.
-
-    Args:
-        obj: Object to serialize
-
-    Returns:
-        JSON-serializable representation of the object
-    """
-    # Our Document types - handle first to ensure serialize_model is used
+def _serialize_for_tracing(obj: Any) -> Any:  # noqa: PLR0911
+    """Convert objects to JSON-serializable format for tracing."""
     if isinstance(obj, Document):
         return obj.serialize_model()
-    # DocumentList
-    if isinstance(obj, DocumentList):
-        return [doc.serialize_model() for doc in obj]
-    # AIMessages
+    if isinstance(obj, list) and obj and isinstance(obj[0], Document):
+        return [doc.serialize_model() for doc in cast(list[Document], obj)]
     if isinstance(obj, AIMessages):
-        result = []
+        result: list[Any] = []
         for msg in obj:
             if isinstance(msg, Document):
                 result.append(msg.serialize_model())
             else:
                 result.append(msg)
         return result
-    # ModelResponse (special Pydantic model) - use standard model_dump
     if isinstance(obj, ModelResponse):
         return obj.model_dump()
-    # Pydantic models - use custom serializer that respects Document.serialize_model()
     if isinstance(obj, BaseModel):
-        # For Pydantic models, we need to handle Document fields specially
-        data = {}
+        data: dict[str, Any] = {}
         for field_name, field_value in obj.__dict__.items():
             if isinstance(field_value, Document):
-                # Use serialize_model for Documents to get base_type
                 data[field_name] = field_value.serialize_model()
             elif isinstance(field_value, BaseModel):
-                # Recursively handle nested Pydantic models
                 data[field_name] = _serialize_for_tracing(field_value)
             else:
-                # Let Pydantic handle other fields normally
                 data[field_name] = field_value
         return data
-    # Fallback to string representation
     try:
-        return str(obj)
+        return str(obj)  # pyright: ignore[reportUnknownArgumentType]
     except Exception:
-        return f"<{type(obj).__name__}>"
+        return f"<{type(obj).__name__}>"  # pyright: ignore[reportUnknownArgumentType]
 
 
 # ---------------------------------------------------------------------------
 # Document trimming utilities
 # ---------------------------------------------------------------------------
-def _trim_document_content(doc_dict: dict[str, Any]) -> dict[str, Any]:
-    """Trim document content based on document type and content type.
+def _trim_attachment_list(attachments: list[Any]) -> list[Any]:
+    """Trim attachment content in a serialized attachment list.
 
-    For non-FlowDocuments:
-    - Text content: Keep first 100 and last 100 chars (unless < 250 total)
-    - Binary content: Remove content entirely
-
-    For FlowDocuments:
-    - Text content: Keep full content
-    - Binary content: Remove content entirely
+    Always trims regardless of parent document type:
+    - Binary (base64): replace content with placeholder
+    - Text > 250 chars: keep first 100 + last 100
+    """
+    trimmed: list[Any] = []
+    for raw_att in attachments:
+        if not isinstance(raw_att, dict):
+            trimmed.append(raw_att)
+            continue
+        att: dict[str, Any] = cast(dict[str, Any], raw_att)
+        content_encoding: str = att.get("content_encoding", "utf-8")
+        if content_encoding == "base64":
+            att = att.copy()
+            att["content"] = "[binary content removed]"
+        elif isinstance(att.get("content"), str) and len(att["content"]) > 250:
+            att = att.copy()
+            c: str = att["content"]
+            trimmed_chars = len(c) - 200
+            att["content"] = c[:100] + f" ... [trimmed {trimmed_chars} chars] ... " + c[-100:]
+        trimmed.append(att)
+    return trimmed
 
-    Args:
-        doc_dict: Document dictionary with base_type, content, and content_encoding
 
-    Returns:
-        Modified document dictionary with trimmed content
-    """
-    # Check if this looks like a document (has required fields)
-    if not isinstance(doc_dict, dict):  # type: ignore[reportUnknownArgumentType]
-        return doc_dict
+def _trim_document_content(doc_dict: dict[str, Any]) -> dict[str, Any]:
+    """Trim document content for traces. All documents trimmed equally."""
+    if not isinstance(doc_dict, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
+        return doc_dict  # pyright: ignore[reportUnreachable]
 
-    if "base_type" not in doc_dict or "content" not in doc_dict:
+    if "content" not in doc_dict or "class_name" not in doc_dict:
         return doc_dict
 
-    base_type = doc_dict.get("base_type")
+    doc_dict = doc_dict.copy()
     content = doc_dict.get("content", "")
     content_encoding = doc_dict.get("content_encoding", "utf-8")
 
-    # For binary content (base64 encoded), remove content
+    # Trim attachments
+    if "attachments" in doc_dict and isinstance(doc_dict["attachments"], list):
+        doc_dict["attachments"] = _trim_attachment_list(cast(list[Any], doc_dict["attachments"]))
+
+    # Binary: remove content
     if content_encoding == "base64":
-        doc_dict = doc_dict.copy()
         doc_dict["content"] = "[binary content removed]"
         return doc_dict
 
-    # For FlowDocuments with text content, keep full content
-    if base_type == "flow":
-        return doc_dict
-
-    # For other documents (task, temporary), trim text content
+    # Text: trim if > 250 chars
     if isinstance(content, str) and len(content) > 250:
-        doc_dict = doc_dict.copy()
-        # Keep first 100 and last 100 characters
-        trimmed_chars = len(content) - 200  # Number of characters removed
-        doc_dict["content"] = (
-            content[:100] + f" ... [trimmed {trimmed_chars} chars] ... " + content[-100:]
-        )
+        trimmed_chars = len(content) - 200
+        doc_dict["content"] = content[:100] + f" ... [trimmed {trimmed_chars} chars] ... " + content[-100:]
 
     return doc_dict
 
 
 def _trim_documents_in_data(data: Any) -> Any:
-    """Recursively trim document content in nested data structures.
-
-    Processes dictionaries, lists, and nested structures to find and trim
-    documents based on their type and content.
-
-    Args:
-        data: Input data that may contain documents
-
-    Returns:
-        Data with document content trimmed according to rules
-    """
+    """Recursively trim document content in nested data structures."""
     if isinstance(data, dict):
-        # Check if this is a document
-        if "base_type" in data and "content" in data:
-            # This is a document, trim it
-            return _trim_document_content(data)
-        else:
-            # Recursively process dictionary values
-            return {k: _trim_documents_in_data(v) for k, v in data.items()}
-    elif isinstance(data, list):
-        # Process each item in list
-        return [_trim_documents_in_data(item) for item in data]
-    elif isinstance(data, tuple):
-        # Process tuples
-        return tuple(_trim_documents_in_data(item) for item in data)
-    else:
-        # Return other types unchanged
-        return data
+        data_dict = cast(dict[str, Any], data)
+        if "class_name" in data_dict and "content" in data_dict:
+            return _trim_document_content(data_dict)
+        return {k: _trim_documents_in_data(v) for k, v in data_dict.items()}
+    if isinstance(data, list):
+        return [_trim_documents_in_data(item) for item in cast(list[Any], data)]
+    if isinstance(data, tuple):
+        return tuple(_trim_documents_in_data(item) for item in cast(tuple[Any, ...], data))
+    return data
 
 
 # ---------------------------------------------------------------------------
-# ``TraceInfo`` – metadata container
+# ``TraceInfo`` - metadata container
 # ---------------------------------------------------------------------------
 class TraceInfo(BaseModel):
     """Container for propagating trace context through the pipeline.
@@ -196,30 +164,12 @@ class TraceInfo(BaseModel):
 
     Environment fallbacks:
         - LMNR_DEBUG: Controls debug-level tracing when set to "true"
-
-        Note: These variables are read directly by the tracing layer and are
+        These variables are read directly by the tracing layer and are
         not part of the Settings configuration.
 
-    Example:
-        >>> # Create trace context
-        >>> trace_info = TraceInfo(
-        ...     session_id="sess_123",
-        ...     user_id="user_456",
-        ...     metadata={"flow": "document_analysis", "version": "1.2"},
-        ...     tags=["production", "high_priority"]
-        ... )
-        >>>
-        >>> # Pass through function calls
-        >>> @trace
-        >>> async def process(data, trace_info: TraceInfo):
-        ...     # TraceInfo automatically propagates to nested calls
-        ...     result = await analyze(data, trace_info=trace_info)
-        ...     return result
-
-    Note:
-        TraceInfo is typically created at the entry point of a flow
-        and passed through all subsequent function calls for
-        consistent tracing context.
+    TraceInfo is typically created at the entry point of a flow
+    and passed through all subsequent function calls for
+    consistent tracing context.
     """
 
     session_id: str | None = None
@@ -243,14 +193,8 @@ class TraceInfo(BaseModel):
 
             Only non-empty values are included in the output.
 
-        Example:
-            >>> trace_info = TraceInfo(session_id="sess_123", tags=["test"])
-            >>> kwargs = trace_info.get_observe_kwargs()
-            >>> # Returns: {"session_id": "sess_123", "tags": ["test"]}
-
-        Note:
-            This method is called internally by the trace decorator
-            to configure Laminar observation parameters.
+        Called internally by the trace decorator to configure Laminar
+        observation parameters.
         """
         kwargs: dict[str, Any] = {}
 
@@ -276,9 +220,6 @@ class TraceInfo(BaseModel):
 # ---------------------------------------------------------------------------
 
 
-_debug_processor_initialized = False
-
-
 def _initialise_laminar() -> None:
     """Initialize Laminar SDK with project configuration.
 
@@ -286,70 +227,13 @@ def _initialise_laminar() -> None:
     from settings. Disables automatic OpenAI instrumentation to avoid
     conflicts with our custom tracing.
 
-    Configuration:
-        - Uses settings.lmnr_project_api_key for authentication
-        - Disables OPENAI instrument to prevent double-tracing
-        - Called automatically by trace decorator on first use
-        - Optionally adds local debug processor if TRACE_DEBUG_PATH is set
-
-    Note:
-        This is an internal function called once per process.
-        Multiple calls are safe (Laminar handles idempotency).
+    Called once per process. Multiple calls are safe (Laminar handles idempotency).
     """
-    global _debug_processor_initialized
-
     if settings.lmnr_project_api_key:
         Laminar.initialize(
-            project_api_key=settings.lmnr_project_api_key,
-            disabled_instruments=[Instruments.OPENAI] if Instruments.OPENAI else [],
-        )
-
-    # Add local debug processor if configured (only once)
-    if not _debug_processor_initialized:
-        _debug_processor_initialized = True
-        debug_path = os.environ.get("TRACE_DEBUG_PATH")
-        if debug_path:
-            _setup_debug_processor(debug_path)
-
-
-def _setup_debug_processor(debug_path: str) -> None:
-    """Set up local debug trace processor."""
-    try:
-        from pathlib import Path  # noqa: PLC0415
-
-        from opentelemetry import trace  # noqa: PLC0415
-
-        from ai_pipeline_core.debug import (  # noqa: PLC0415
-            LocalDebugSpanProcessor,
-            LocalTraceWriter,
-            TraceDebugConfig,
-        )
-
-        config = TraceDebugConfig(
-            path=Path(debug_path),
-            max_element_bytes=int(os.environ.get("TRACE_DEBUG_MAX_INLINE", 10000)),
-            max_traces=int(os.environ.get("TRACE_DEBUG_MAX_TRACES", 20)) or None,
+            project_api_key=settings.lmnr_project_api_key, disabled_instruments=[Instruments.OPENAI] if Instruments.OPENAI else [], export_timeout_seconds=15
         )
 
-        writer = LocalTraceWriter(config)
-        processor = LocalDebugSpanProcessor(writer)
-
-        # Add to tracer provider
-        provider = trace.get_tracer_provider()
-        add_processor = getattr(provider, "add_span_processor", None)
-        if add_processor is not None:
-            add_processor(processor)
-
-        # Register shutdown
-        import atexit  # noqa: PLC0415
-
-        atexit.register(processor.shutdown)
-
-    except Exception as e:
-        import logging  # noqa: PLC0415
-
-        logging.getLogger(__name__).warning(f"Failed to setup debug trace processor: {e}")
-
 
 # Overload for calls like @trace(name="...", level="debug")
 @overload
@@ -375,11 +259,11 @@ def trace(
 
 # Overload for the bare @trace call
 @overload
-def trace(func: Callable[P, R]) -> Callable[P, R]: ...
+def trace(func: Callable[P, R]) -> Callable[P, R]: ...  # noqa: UP047
 
 
 # Actual implementation
-def trace(
+def trace(  # noqa: UP047
     func: Callable[P, R] | None = None,
     *,
     level: TraceLevel = "always",
@@ -454,9 +338,10 @@ def trace(
                                 calls (default True). Set False for isolated traces.
 
         trim_documents: Automatically trim document content in traces (default True).
-                       When enabled, non-FlowDocument text content is trimmed to
+                       When enabled, text content is trimmed to
                        first/last 100 chars, and all binary content is removed.
-                       FlowDocuments keep full text content but binary is removed.
+                       Binary content is removed, text content is trimmed.
+                       Attachment content follows the same trimming rules.
                        Helps reduce trace size for large documents.
 
     Returns:
@@ -467,35 +352,6 @@ def trace(
         automatically creates or propagates a TraceInfo instance, ensuring
         consistent session/user tracking across the call chain.
 
-    Example:
-        >>> # RECOMMENDED - No parameters needed for most cases!
-        >>> @trace
-        >>> async def process_document(doc):
-        ...     return await analyze(doc)
-        >>>
-        >>> # With parameters (RARE - only when specifically needed):
-        >>> @trace(level="debug")  # Only for debug-specific tracing
-        >>> async def debug_operation():
-        ...     pass
-
-        >>> @trace(ignore_inputs=["api_key"])  # Only for sensitive data
-        >>> async def api_call(data, api_key):
-        ...     return await external_api(data, api_key)
-        >>>
-        >>> # AVOID unnecessary configuration - defaults handle:
-        >>> # - Automatic naming from function name
-        >>> # - Standard trace level ("always")
-        >>> # - Full input/output capture
-        >>> # - Proper span type inference
-        >>>
-        >>> # Custom formatting
-        >>> @trace(
-        ...     input_formatter=lambda doc: f"Document: {doc.id}",
-        ...     output_formatter=lambda res: f"Results: {len(res)} items"
-        >>> )
-        >>> def analyze(doc):
-        ...     return results
-
     Environment variables:
         - LMNR_DEBUG: Set to "true" to enable debug-level traces
         - LMNR_PROJECT_API_KEY: Required for trace submission
@@ -505,11 +361,9 @@ def trace(
         - When level="off", decorator returns original function unchanged
         - Large inputs/outputs can be excluded with ignore_* parameters
 
-    Note:
-        - Automatically initializes Laminar on first use
-        - Works with both sync and async functions
-        - Preserves function signature and metadata
-        - Thread-safe and async-safe
+    Automatically initializes Laminar on first use. Works with both sync and
+    async functions. Preserves function signature and metadata. Thread-safe
+    and async-safe.
     """
     if level == "off":
         if func:
@@ -550,27 +404,27 @@ def trace(
         sig = inspect.signature(f)
         is_coroutine = inspect.iscoroutinefunction(f)
         observe_name = name or f.__name__
-        _observe = observe
-
-        _session_id = session_id
-        _user_id = user_id
-        _metadata = metadata if metadata is not None else {}
-        _tags = tags if tags is not None else []
-        _span_type = span_type
-        _ignore_input = ignore_input
-        _ignore_output = ignore_output
-        _ignore_inputs = ignore_inputs
-        _input_formatter = input_formatter
-        _output_formatter = output_formatter
-        _ignore_exceptions = ignore_exceptions
-        _preserve_global_context = preserve_global_context
-        _trim_documents = trim_documents
+        bound_observe = observe
+
+        bound_session_id = session_id
+        bound_user_id = user_id
+        bound_metadata = metadata if metadata is not None else {}
+        bound_tags = tags if tags is not None else []
+        bound_span_type = span_type
+        bound_ignore_input = ignore_input
+        bound_ignore_output = ignore_output
+        bound_ignore_inputs = ignore_inputs
+        bound_input_formatter = input_formatter
+        bound_output_formatter = output_formatter
+        bound_ignore_exceptions = ignore_exceptions
+        bound_preserve_global_context = preserve_global_context
+        bound_trim_documents = trim_documents
 
         # Create document trimming formatters if needed
-        def _create_trimming_input_formatter(*args, **kwargs) -> str:
+        def _create_trimming_input_formatter(*args: Any, **kwargs: Any) -> str:
             # First, let any custom formatter process the data
-            if _input_formatter:
-                result = _input_formatter(*args, **kwargs)
+            if bound_input_formatter:
+                result = bound_input_formatter(*args, **kwargs)
                 # If formatter returns string, try to parse and trim
                 if isinstance(result, str):  # type: ignore[reportUnknownArgumentType]
                     try:
@@ -587,7 +441,7 @@ def trace(
                 # No custom formatter - mimic Laminar's get_input_from_func_args
                 # Build a dict with parameter names as keys (like Laminar does)
                 params = list(sig.parameters.keys())
-                data = {}
+                data: dict[str, Any] = {}
 
                 # Map args to parameter names
                 for i, arg in enumerate(args):
@@ -607,8 +461,8 @@ def trace(
 
         def _create_trimming_output_formatter(result: Any) -> str:
             # First, let any custom formatter process the data
-            if _output_formatter:
-                formatted = _output_formatter(result)
+            if bound_output_formatter:
+                formatted = bound_output_formatter(result)
                 # If formatter returns string, try to parse and trim
                 if isinstance(formatted, str):  # type: ignore[reportUnknownArgumentType]
                     try:
@@ -650,41 +504,41 @@ def trace(
             observe_params["name"] = observe_name
 
             # Override with decorator-level session_id and user_id if provided
-            if _session_id:
-                observe_params["session_id"] = _session_id
-            if _user_id:
-                observe_params["user_id"] = _user_id
-            if _metadata:
-                observe_params["metadata"] = _metadata
-            if _tags:
-                observe_params["tags"] = observe_params.get("tags", []) + _tags
-            if _span_type:
-                observe_params["span_type"] = _span_type
+            if bound_session_id:
+                observe_params["session_id"] = bound_session_id
+            if bound_user_id:
+                observe_params["user_id"] = bound_user_id
+            if bound_metadata:
+                observe_params["metadata"] = bound_metadata
+            if bound_tags:
+                observe_params["tags"] = observe_params.get("tags", []) + bound_tags
+            if bound_span_type:
+                observe_params["span_type"] = bound_span_type
 
             # Add the new Laminar parameters
-            if _ignore_input:
-                observe_params["ignore_input"] = _ignore_input
-            if _ignore_output:
-                observe_params["ignore_output"] = _ignore_output
-            if _ignore_inputs is not None:
-                observe_params["ignore_inputs"] = _ignore_inputs
+            if bound_ignore_input:
+                observe_params["ignore_input"] = bound_ignore_input
+            if bound_ignore_output:
+                observe_params["ignore_output"] = bound_ignore_output
+            if bound_ignore_inputs is not None:
+                observe_params["ignore_inputs"] = bound_ignore_inputs
 
             # Use trimming formatters if trim_documents is enabled
-            if _trim_documents:
+            if bound_trim_documents:
                 # Use the trimming formatters (which may wrap custom formatters)
                 observe_params["input_formatter"] = _create_trimming_input_formatter
                 observe_params["output_formatter"] = _create_trimming_output_formatter
             else:
                 # Use custom formatters directly if provided
-                if _input_formatter is not None:
-                    observe_params["input_formatter"] = _input_formatter
-                if _output_formatter is not None:
-                    observe_params["output_formatter"] = _output_formatter
+                if bound_input_formatter is not None:
+                    observe_params["input_formatter"] = bound_input_formatter
+                if bound_output_formatter is not None:
+                    observe_params["output_formatter"] = bound_output_formatter
 
-            if _ignore_exceptions:
-                observe_params["ignore_exceptions"] = _ignore_exceptions
-            if _preserve_global_context:
-                observe_params["preserve_global_context"] = _preserve_global_context
+            if bound_ignore_exceptions:
+                observe_params["ignore_exceptions"] = bound_ignore_exceptions
+            if bound_preserve_global_context:
+                observe_params["preserve_global_context"] = bound_preserve_global_context
 
             return observe_params
 
@@ -697,7 +551,7 @@ def trace(
                 The result of the wrapped function.
             """
             observe_params = _prepare_and_get_observe_params(kwargs)
-            observed_func = _observe(**observe_params)(f)
+            observed_func = bound_observe(**observe_params)(f)
             return observed_func(*args, **kwargs)
 
         @wraps(f)
@@ -708,8 +562,8 @@ def trace(
                 The result of the wrapped function.
             """
             observe_params = _prepare_and_get_observe_params(kwargs)
-            observed_func = _observe(**observe_params)(f)
-            return await observed_func(*args, **kwargs)  # pyright: ignore[reportGeneralTypeIssues]
+            observed_func = bound_observe(**observe_params)(f)
+            return await observed_func(*args, **kwargs)  # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType]
 
         wrapper = async_wrapper if is_coroutine else sync_wrapper
 
@@ -717,17 +571,14 @@ def trace(
         wrapper.__is_traced__ = True  # type: ignore[attr-defined]
 
         # Preserve the original function signature
-        try:
+        with contextlib.suppress(AttributeError, ValueError):
             wrapper.__signature__ = sig  # type: ignore[attr-defined]
-        except (AttributeError, ValueError):
-            pass
 
         return cast(Callable[P, R], wrapper)
 
     if func:
         return decorator(func)  # Called as @trace
-    else:
-        return decorator  # Called as @trace(...)
+    return decorator  # Called as @trace(...)
 
 
 def set_trace_cost(cost: float | str) -> None:
@@ -739,10 +590,10 @@ def set_trace_cost(cost: float | str) -> None:
     current operation. Particularly useful for tracking costs of external
     API calls, compute resources, or custom billing scenarios.
 
-    The cost is stored in three metadata fields for compatibility:
-    - gen_ai.usage.output_cost: Standard OpenAI cost field
-    - gen_ai.usage.cost: Alternative cost field
-    - cost: Simple cost field
+    The cost is stored in three metadata fields for observability tool consumption:
+    - gen_ai.usage.output_cost: OpenTelemetry GenAI semantic convention
+    - gen_ai.usage.cost: Aggregated cost field
+    - cost: Short-form cost field
 
     Args:
         cost: The cost value to set. Can be:
@@ -750,56 +601,14 @@ def set_trace_cost(cost: float | str) -> None:
               - str: USD format with dollar sign (e.g., "$0.05" or "$1.25")
               Only positive values will be set; zero or negative values are ignored.
 
-    Example:
-        >>> # Track cost of external API call
-        >>> @trace
-        >>> async def call_translation_api(text: str) -> str:
-        ...     # External API charges per character
-        ...     char_count = len(text)
-        ...     cost_per_char = 0.00001  # $0.00001 per character
-        ...
-        ...     result = await external_api.translate(text)
-        ...
-        ...     # Set the cost for this operation
-        ...     set_trace_cost(char_count * cost_per_char)
-        ...     return result
-        >>>
-        >>> # Track compute resource costs
-        >>> @trace
-        >>> def process_video(video_path: str) -> dict:
-        ...     duration = get_video_duration(video_path)
-        ...     cost_per_minute = 0.10  # $0.10 per minute
-        ...
-        ...     result = process_video_content(video_path)
-        ...
-        ...     # Set cost using string format
-        ...     set_trace_cost(f"${duration * cost_per_minute:.2f}")
-        ...     return result
-        >>>
-        >>> # Combine with LLM costs in pipeline
-        >>> @pipeline_task
-        >>> async def enriched_generation(prompt: str) -> str:
-        ...     # LLM cost tracked automatically via ModelResponse
-        ...     response = await llm.generate("gpt-5.1", messages=prompt)
-        ...
-        ...     # Add cost for post-processing
-        ...     processing_cost = 0.02  # Fixed cost for enrichment
-        ...     set_trace_cost(processing_cost)
-        ...
-        ...     return enrich_response(response.content)
-
     Raises:
         ValueError: If string format is invalid (not a valid USD amount).
 
-    Note:
-        - This function only works within a traced context (function decorated
-          with @trace, @pipeline_task, or @pipeline_flow)
-        - LLM costs are tracked automatically via ModelResponse; use this for non-LLM costs
-        - Cost should be a positive number representing actual monetary cost in USD
-        - The cost is added to the current span's attributes/metadata
-        - Multiple calls overwrite the previous cost (not cumulative)
-        - If called outside a traced context (no active span), it has no effect
-          and does not raise an error
+    Only works within a traced context (function decorated with @trace,
+    @pipeline_task, or @pipeline_flow). LLM costs are tracked automatically via
+    ModelResponse; use this for non-LLM costs. Multiple calls overwrite the
+    previous cost (not cumulative). If called outside a traced context, it has
+    no effect and does not raise an error.
     """
     # Parse string format if provided
     if isinstance(cost, str):
@@ -812,9 +621,7 @@ def set_trace_cost(cost: float | str) -> None:
             # Remove $ and convert to float
             cost_value = float(cost_str[1:])
         except ValueError as e:
-            raise ValueError(
-                f"Invalid USD format: {cost!r}. Must be a valid number after '$'"
-            ) from e
+            raise ValueError(f"Invalid USD format: {cost!r}. Must be a valid number after '$'") from e
     else:
         cost_value = cost
 
@@ -826,11 +633,8 @@ def set_trace_cost(cost: float | str) -> None:
             "cost": cost_value,
         }
 
-        try:
+        with contextlib.suppress(Exception):
             Laminar.set_span_attributes(attributes)
-        except Exception:
-            # Silently ignore if not in a traced context
-            pass
 
 
-__all__ = ["trace", "TraceLevel", "TraceInfo", "set_trace_cost"]
+__all__ = ["TraceInfo", "TraceLevel", "set_trace_cost", "trace"]