paid-python 0.0.5a40__py3-none-any.whl → 0.1.1__py3-none-any.whl

paid/tracing/tracing.py

@@ -2,37 +2,26 @@
 import asyncio
 import atexit
 import contextvars
-import functools
-import logging
 import os
 import signal
 from typing import Any, Awaitable, Callable, Dict, Optional, Tuple, TypeVar, Union
 
 import dotenv
+from . import distributed_tracing
 from opentelemetry import trace
 from opentelemetry.context import Context
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor, TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
-from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
-from opentelemetry.trace import NonRecordingSpan, SpanContext, Status, StatusCode, TraceFlags
-
-# Configure logging
-dotenv.load_dotenv()
-log_level_name = os.environ.get("PAID_LOG_LEVEL")
-if log_level_name is not None:
-    log_level = getattr(logging, log_level_name.upper())
-else:
-    log_level = logging.ERROR  # Default to show errors
-logger = logging.getLogger(__name__)
-logger.setLevel(log_level)
-if not logger.hasHandlers():
-    handler = logging.StreamHandler()
-    handler.setLevel(log_level)
-    formatter = logging.Formatter("%(levelname)s:%(name)s:%(message)s")
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
+from opentelemetry.trace import NonRecordingSpan, NoOpTracerProvider, SpanContext, Status, StatusCode, TraceFlags
+
+from paid.logger import logger
+
+_ = dotenv.load_dotenv()
+DEFAULT_COLLECTOR_ENDPOINT = (
+    os.environ.get("PAID_OTEL_COLLECTOR_ENDPOINT") or "https://collector.agentpaid.io:4318/v1/traces"
+)
 
 # Context variables for passing data to nested spans (e.g., in openAiWrapper)
 paid_external_customer_id_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
@@ -41,10 +30,8 @@ paid_external_customer_id_var: contextvars.ContextVar[Optional[str]] = contextva
 paid_external_agent_id_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
     "paid_external_agent_id", default=None
 )
-# api_key storage
-paid_token_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar("paid_token", default=None)
 # trace id storage (generated from token)
-paid_trace_id: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar("paid_trace_id", default=None)
+paid_trace_id_var: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar("paid_trace_id", default=None)
 # flag to enable storing prompt contents
 paid_store_prompt_var: contextvars.ContextVar[Optional[bool]] = contextvars.ContextVar(
     "paid_store_prompt", default=False
@@ -56,25 +43,36 @@ paid_user_metadata_var: contextvars.ContextVar[Optional[Dict[str, Any]]] = conte
 
 T = TypeVar("T")
 
-_token: Optional[str] = None
+
+class _TokenStore:
+    """Private token storage to enforce access through getter/setter."""
+
+    __token: Optional[str] = None
+
+    @classmethod
+    def get(cls) -> Optional[str]:
+        """Get the stored API token."""
+        return cls.__token
+
+    @classmethod
+    def set(cls, token: str) -> None:
+        """Set the API token."""
+        cls.__token = token
 
 
 def get_token() -> Optional[str]:
     """Get the stored API token."""
-    global _token
-    return _token
+    return _TokenStore.get()
 
 
 def set_token(token: str) -> None:
     """Set the API token."""
-    global _token
-    _token = token
+    _TokenStore.set(token)
 
 
-otel_id_generator = RandomIdGenerator()
-
 # Isolated tracer provider for Paid - separate from any user OTEL setup
-paid_tracer_provider: Optional[TracerProvider] = None
+# Initialized at module load with defaults, never None (uses no-op provider if not initialized or API key isn't available)
+paid_tracer_provider: Union[TracerProvider, NoOpTracerProvider] = NoOpTracerProvider()
 
 
 class PaidSpanProcessor(SpanProcessor):
@@ -146,7 +144,6 @@ class PaidSpanProcessor(SpanProcessor):
                 for k, v in original_attributes.items()
                 if not any(k.startswith(prefix) for prefix in self.PROMPT_ATTRIBUTES_PREFIXES)
             }
-            # Temporarily replace attributes for export
             # This works because the exporter reads attributes during serialization
             object.__setattr__(span, "_attributes", filtered_attrs)
 
@@ -159,9 +156,7 @@ class PaidSpanProcessor(SpanProcessor):
         return True
 
 
-def _initialize_tracing(
-    api_key: Optional[str] = None, collector_endpoint: Optional[str] = "https://collector.agentpaid.io:4318/v1/traces"
-):
+def initialize_tracing_(api_key: Optional[str] = None, collector_endpoint: Optional[str] = DEFAULT_COLLECTOR_ENDPOINT):
     """
     Initialize OpenTelemetry with OTLP exporter for Paid backend.
 
@@ -171,20 +166,21 @@ def _initialize_tracing(
     """
     global paid_tracer_provider
 
+    if not collector_endpoint:
+        collector_endpoint = DEFAULT_COLLECTOR_ENDPOINT
+
     try:
-        if _token is not None:
-            raise RuntimeError("Tracing is already initialized.")
+        if get_token() is not None:
+            logger.warning("Tracing is already initialized - skipping re-initialization")
+            return
 
         # Get API key from parameter or environment
         if api_key is None:
-            import dotenv
-
-            dotenv.load_dotenv()
             api_key = os.environ.get("PAID_API_KEY")
             if api_key is None:
-                raise ValueError(
-                    "API key must be provided either as parameter or via PAID_API_KEY environment variable"
-                )
+                logger.error("API key must be provided via PAID_API_KEY environment variable")
+                # don't throw - tracing should not break the app
+                return
 
         set_token(api_key)
 
@@ -210,7 +206,9 @@ def _initialize_tracing(
         # Terminate gracefully and don't lose traces
         def flush_traces():
             try:
-                if paid_tracer_provider is not None and not paid_tracer_provider.force_flush(10000):
+                if not isinstance(paid_tracer_provider, NoOpTracerProvider) and not paid_tracer_provider.force_flush(
+                    10000
+                ):
                     logger.error("OTEL force flush : timeout reached")
             except Exception as e:
                 logger.error(f"Error flushing traces: {e}")
@@ -237,9 +235,9 @@ def _initialize_tracing(
             signal.signal(sig, create_chained_signal_handler(sig))
 
         logger.info("Paid tracing initialized successfully - collector at %s", collector_endpoint)
-    except Exception:
-        logger.exception("Failed to initialize Paid tracing")
-        raise
+    except Exception as e:
+        logger.error(f"Failed to initialize Paid tracing: {e}")
+        # don't throw - tracing should not break the app
 
 
 def get_paid_tracer() -> trace.Tracer:
@@ -251,14 +249,16 @@ def get_paid_tracer() -> trace.Tracer:
 
     Raises:
         RuntimeError: If the tracer provider is not initialized.
+
+    Notes:
+        Tracing is automatically initialized when using @paid_tracing decorator or context manager.
     """
-    if paid_tracer_provider is None:
-        raise RuntimeError("Paid tracer provider is not initialized. Call Paid.initialize_tracing() first.")
+    global paid_tracer_provider
     return paid_tracer_provider.get_tracer("paid.python")
 
 
-def _trace_sync(
-    external_customer_id: str,
+def trace_sync_(
+    external_customer_id: Optional[str],
     fn: Callable[..., T],
     external_agent_id: Optional[str] = None,
     tracing_token: Optional[int] = None,
@@ -267,30 +267,46 @@ def _trace_sync(
     args: Optional[Tuple] = None,
     kwargs: Optional[Dict] = None,
 ) -> T:
+    """
+    Internal function for synchronous tracing. Use @paid_tracing decorator instead.
+
+    This is a low-level internal function. Users should use the @paid_tracing decorator
+    or context manager for a more Pythonic interface.
+
+    Parameters:
+        external_customer_id: The external customer ID to associate with the trace.
+        fn: The function to execute and trace.
+        external_agent_id: Optional external agent ID.
+        tracing_token: Optional token for distributed tracing.
+        store_prompt: Whether to store prompt/completion contents.
+        metadata: Optional metadata to attach to the trace.
+        args: Positional arguments for the function.
+        kwargs: Keyword arguments for the function.
+
+    Returns:
+        The result of executing fn(*args, **kwargs).
+
+    Raises:
+        Only when user callback raises.
+    """
     args = args or ()
     kwargs = kwargs or {}
-    token = get_token()
-    if not token:
-        raise RuntimeError(
-            "No token found - tracing is not initialized and will not be captured. Call Paid.initialize_tracing() first."
-        )
 
     # Set context variables for access by nested spans
-    reset_id_ctx_token = paid_external_customer_id_var.set(external_customer_id)
+    reset_customer_id_ctx_token = paid_external_customer_id_var.set(external_customer_id)
     reset_agent_id_ctx_token = paid_external_agent_id_var.set(external_agent_id)
-    reset_token_ctx_token = paid_token_var.set(token)
     reset_store_prompt_ctx_token = paid_store_prompt_var.set(store_prompt)
     reset_user_metadata_ctx_token = paid_user_metadata_var.set(metadata)
 
     # If user set trace context manually
     override_trace_id = tracing_token
     if not override_trace_id:
-        override_trace_id = paid_trace_id.get()
+        override_trace_id = paid_trace_id_var.get()
     ctx: Optional[Context] = None
     if override_trace_id is not None:
         span_context = SpanContext(
             trace_id=override_trace_id,
-            span_id=otel_id_generator.generate_span_id(),
+            span_id=distributed_tracing.otel_id_generator.generate_span_id(),
             is_remote=True,
             trace_flags=TraceFlags(TraceFlags.SAMPLED),
         )
@@ -300,9 +316,6 @@ def _trace_sync(
         tracer = get_paid_tracer()
         logger.info(f"Creating span for external_customer_id: {external_customer_id}")
         with tracer.start_as_current_span("parent_span", context=ctx) as span:
-            span.set_attribute("external_customer_id", external_customer_id)
-            if external_agent_id:
-                span.set_attribute("external_agent_id", external_agent_id)
             try:
                 result = fn(*args, **kwargs)
                 span.set_status(Status(StatusCode.OK))
@@ -312,15 +325,14 @@ def _trace_sync(
                 span.set_status(Status(StatusCode.ERROR, str(error)))
                 raise
     finally:
-        paid_external_customer_id_var.reset(reset_id_ctx_token)
+        paid_external_customer_id_var.reset(reset_customer_id_ctx_token)
         paid_external_agent_id_var.reset(reset_agent_id_ctx_token)
-        paid_token_var.reset(reset_token_ctx_token)
         paid_store_prompt_var.reset(reset_store_prompt_ctx_token)
         paid_user_metadata_var.reset(reset_user_metadata_ctx_token)
 
 
-async def _trace_async(
-    external_customer_id: str,
+async def trace_async_(
+    external_customer_id: Optional[str],
     fn: Callable[..., Union[T, Awaitable[T]]],
     external_agent_id: Optional[str] = None,
     tracing_token: Optional[int] = None,
@@ -329,30 +341,46 @@ async def _trace_async(
     args: Optional[Tuple] = None,
     kwargs: Optional[Dict] = None,
 ) -> Union[T, Awaitable[T]]:
+    """
+    Internal function for asynchronous tracing. Use @paid_tracing decorator instead.
+
+    This is a low-level internal function. Users should use the @paid_tracing decorator
+    or context manager for a more Pythonic interface.
+
+    Parameters:
+        external_customer_id: The external customer ID to associate with the trace.
+        fn: The async function to execute and trace.
+        external_agent_id: Optional external agent ID.
+        tracing_token: Optional token for distributed tracing.
+        store_prompt: Whether to store prompt/completion contents.
+        metadata: Optional metadata to attach to the trace.
+        args: Positional arguments for the function.
+        kwargs: Keyword arguments for the function.
+
+    Returns:
+        The result of executing fn(*args, **kwargs).
+
+    Raises:
+        Only when user callback raises.
+    """
     args = args or ()
     kwargs = kwargs or {}
-    token = get_token()
-    if not token:
-        raise RuntimeError(
-            "No token found - tracing is not initialized and will not be captured. Call Paid.initialize_tracing() first."
-        )
 
     # Set context variables for access by nested spans
-    reset_id_ctx_token = paid_external_customer_id_var.set(external_customer_id)
+    reset_customer_id_ctx_token = paid_external_customer_id_var.set(external_customer_id)
     reset_agent_id_ctx_token = paid_external_agent_id_var.set(external_agent_id)
-    reset_token_ctx_token = paid_token_var.set(token)
     reset_store_prompt_ctx_token = paid_store_prompt_var.set(store_prompt)
     reset_user_metadata_ctx_token = paid_user_metadata_var.set(metadata)
 
     # If user set trace context manually
     override_trace_id = tracing_token
     if not override_trace_id:
-        override_trace_id = paid_trace_id.get()
+        override_trace_id = paid_trace_id_var.get()
     ctx: Optional[Context] = None
     if override_trace_id is not None:
         span_context = SpanContext(
             trace_id=override_trace_id,
-            span_id=otel_id_generator.generate_span_id(),
+            span_id=distributed_tracing.otel_id_generator.generate_span_id(),
             is_remote=True,
             trace_flags=TraceFlags(TraceFlags.SAMPLED),
         )
@@ -362,9 +390,6 @@ async def _trace_async(
         tracer = get_paid_tracer()
         logger.info(f"Creating span for external_customer_id: {external_customer_id}")
         with tracer.start_as_current_span("parent_span", context=ctx) as span:
-            span.set_attribute("external_customer_id", external_customer_id)
-            if external_agent_id:
-                span.set_attribute("external_agent_id", external_agent_id)
             try:
                 if asyncio.iscoroutinefunction(fn):
                     result = await fn(*args, **kwargs)
@@ -377,374 +402,7 @@ async def _trace_async(
                 span.set_status(Status(StatusCode.ERROR, str(error)))
                 raise
     finally:
-        paid_external_customer_id_var.reset(reset_id_ctx_token)
+        paid_external_customer_id_var.reset(reset_customer_id_ctx_token)
         paid_external_agent_id_var.reset(reset_agent_id_ctx_token)
-        paid_token_var.reset(reset_token_ctx_token)
         paid_store_prompt_var.reset(reset_store_prompt_ctx_token)
         paid_user_metadata_var.reset(reset_user_metadata_ctx_token)
-
-
-def generate_tracing_token() -> int:
-    """
-    This will generate and return a tracing token but it will not set it
-    for the tracing context. Needed when you only want to store or send a tracing token
-    somewhere else.
-    """
-    return otel_id_generator.generate_trace_id()
-
-
-def generate_and_set_tracing_token() -> int:
-    """
-    *Advanced feature*
-    In cases when you can't share the same Paid.trace() or @paid_tracing() context with
-    code that you want to track together (complex concurrency logic,
-    or disjoint workflows, or work is separated between processes),
-    then you can manually generate a tracing token with generate_and_set_tracing_token()
-    and share it with the other parts of your application or service using set_tracing_token().
-
-    This function returns tracing token and attaches it to all consequent
-    Paid.trace() or @paid_tracing() tracing contexts. So all the costs and signals that share this
-    tracing context are associated with each other.
-
-    To stop associating the traces one can either call
-    generate_and_set_tracing_token() once again or call unset_tracing_token().
-    The former is suitable if you still want to trace but in a fresh
-    context, and the latter will go back to unique traces per Paid.trace() or @paid_tracing().
-
-    Returns:
-        int: The tracing token (OpenTelemetry trace ID)
-
-    Example:
-        >>> from paid.tracing import generate_and_set_tracing_token, set_tracing_token, unset_tracing_token
-        >>> # Process 1: Generate token
-        >>> token = generate_and_set_tracing_token()
-        >>> save_to_redis("workflow_123", token)
-        >>>
-        >>> # Process 2: Use token
-        >>> token = load_from_redis("workflow_123")
-        >>> set_tracing_token(token)
-        >>> # ... do traced work ...
-        >>> unset_tracing_token()
-    """
-    random_trace_id = otel_id_generator.generate_trace_id()
-    _ = paid_trace_id.set(random_trace_id)
-    return random_trace_id
-
-
-def set_tracing_token(token: int):
-    """
-    *Advanced feature*
-    In cases when you can't share the same Paid.trace() or @paid_tracing() context with
-    code that you want to track together (complex concurrency logic,
-    or disjoint workflows, or work is separated between processes),
-    then you can manually generate a tracing token with generate_and_set_tracing_token()
-    and share it with the other parts of your application or service using set_tracing_token().
-
-    Sets tracing token. Provided token should come from generate_and_set_tracing_token().
-    Once set, the consequent traces will be related to each other.
-
-    Args:
-        token (int): A tracing token from generate_and_set_tracing_token()
-
-    Example:
-        >>> from paid.tracing import set_tracing_token, unset_tracing_token, paid_tracing
-        >>> # Retrieve token from storage
-        >>> token = get_from_redis("workflow_123")
-        >>> set_tracing_token(token)
-        >>>
-        >>> @paid_tracing("customer_123", "agent_123")
-        >>> def process_workflow():
-        ...     # This trace will be linked to the token
-        ...     pass
-        >>>
-        >>> process_workflow()
-        >>> unset_tracing_token()
-    """
-    _ = paid_trace_id.set(token)
-
-
-def unset_tracing_token():
-    """
-    Unsets the token previously set by generate_and_set_tracing_token()
-    or by set_tracing_token(token). Does nothing if the token was never set.
-    When tracing token is unset, traces are unique for a single Paid.trace() or @paid_tracing() context.
-
-    Example:
-        >>> from paid.tracing import set_tracing_token, unset_tracing_token
-        >>> set_tracing_token(stored_token)
-        >>> try:
-        ...     process_workflow()
-        ... finally:
-        ...     unset_tracing_token()  # Always clean up
-    """
-    _ = paid_trace_id.set(None)
-
-
-class paid_tracing:
-    """
-    Decorator and context manager for tracing with Paid.
-
-    This class can be used both as a decorator and as a context manager (with/async with),
-    providing flexible tracing capabilities for both functions and code blocks.
-
-    Parameters
-    ----------
-    external_customer_id : str
-        The external customer ID to associate with the trace.
-    external_agent_id : Optional[str], optional
-        The external agent ID to associate with the trace, by default None.
-    tracing_token : Optional[int], optional
-        Optional tracing token for distributed tracing, by default None.
-    store_prompt : bool, optional
-        Whether to store prompt contents in span attributes, by default False.
-    collector_endpoint: Optional[str], optional
-        OTEL collector HTTP endpoint, by default "https://collector.agentpaid.io:4318/v1/traces".
-    metadata : Optional[Dict[str, Any]], optional
-        Optional metadata to attach to the trace, by default None.
-
-    Examples
-    --------
-    As a decorator (sync):
-    >>> @paid_tracing(external_customer_id="customer123", external_agent_id="agent456")
-    ... def my_function(arg1, arg2):
-    ...     return arg1 + arg2
-
-    As a decorator (async):
-    >>> @paid_tracing(external_customer_id="customer123")
-    ... async def my_async_function(arg1, arg2):
-    ...     return arg1 + arg2
-
-    As a context manager (sync):
-    >>> with paid_tracing(external_customer_id="customer123", external_agent_id="agent456"):
-    ...     result = expensive_computation()
-
-    As a context manager (async):
-    >>> async with paid_tracing(external_customer_id="customer123"):
-    ...     result = await async_operation()
-
-    Notes
-    -----
-    If tracing is not already initialized, the decorator will automatically
-    initialize it using the PAID_API_KEY environment variable.
-    """
-
-    def __init__(
-        self,
-        external_customer_id: str,
-        *,
-        external_agent_id: Optional[str] = None,
-        tracing_token: Optional[int] = None,
-        store_prompt: bool = False,
-        collector_endpoint: Optional[str] = "https://collector.agentpaid.io:4318/v1/traces",
-        metadata: Optional[Dict[str, Any]] = None,
-    ):
-        self.external_customer_id = external_customer_id
-        self.external_agent_id = external_agent_id
-        self.tracing_token = tracing_token
-        self.store_prompt = store_prompt
-        self.collector_endpoint = collector_endpoint
-        self.metadata = metadata
-        self._span: Any = None
-        self._reset_tokens: Optional[
-            Tuple[
-                contextvars.Token[Optional[str]],
-                contextvars.Token[Optional[str]],
-                contextvars.Token[Optional[str]],
-                contextvars.Token[Optional[bool]],
-                contextvars.Token[Optional[Dict[str, Any]]],
-            ]
-        ] = None
-
-    def _setup_context(self) -> Optional[Context]:
-        """Set up context variables and return OTEL context if needed."""
-        token = get_token()
-        if not token:
-            raise RuntimeError("No token found - tracing is not initialized. Call Paid.initialize_tracing() first.")
-
-        # Set context variables
-        reset_id_ctx_token = paid_external_customer_id_var.set(self.external_customer_id)
-        reset_agent_id_ctx_token = paid_external_agent_id_var.set(self.external_agent_id)
-        reset_token_ctx_token = paid_token_var.set(token)
-        reset_store_prompt_ctx_token = paid_store_prompt_var.set(self.store_prompt)
-        reset_user_metadata_ctx_token = paid_user_metadata_var.set(self.metadata)
-
-        # Store reset tokens for cleanup
-        self._reset_tokens = (
-            reset_id_ctx_token,
-            reset_agent_id_ctx_token,
-            reset_token_ctx_token,
-            reset_store_prompt_ctx_token,
-            reset_user_metadata_ctx_token,
-        )
-
-        # Handle distributed tracing token
-        override_trace_id = self.tracing_token
-        if not override_trace_id:
-            override_trace_id = paid_trace_id.get()
-
-        ctx: Optional[Context] = None
-        if override_trace_id is not None:
-            span_context = SpanContext(
-                trace_id=override_trace_id,
-                span_id=otel_id_generator.generate_span_id(),
-                is_remote=True,
-                trace_flags=TraceFlags(TraceFlags.SAMPLED),
-            )
-            ctx = trace.set_span_in_context(NonRecordingSpan(span_context))
-
-        return ctx
-
-    def _cleanup_context(self):
-        """Reset all context variables."""
-        if self._reset_tokens:
-            (
-                reset_id_ctx_token,
-                reset_agent_id_ctx_token,
-                reset_token_ctx_token,
-                reset_store_prompt_ctx_token,
-                reset_user_metadata_ctx_token,
-            ) = self._reset_tokens
-            paid_external_customer_id_var.reset(reset_id_ctx_token)
-            paid_external_agent_id_var.reset(reset_agent_id_ctx_token)
-            paid_token_var.reset(reset_token_ctx_token)
-            paid_store_prompt_var.reset(reset_store_prompt_ctx_token)
-            paid_user_metadata_var.reset(reset_user_metadata_ctx_token)
-            self._reset_tokens = None
-
-    # Context manager methods for sync
-    def __enter__(self):
-        """Enter synchronous context."""
-        ctx = self._setup_context()
-
-        tracer = get_paid_tracer()
-        logger.info(f"Creating span for external_customer_id: {self.external_customer_id}")
-        self._span = tracer.start_as_current_span("parent_span", context=ctx)
-        span = self._span.__enter__()
-
-        span.set_attribute("external_customer_id", self.external_customer_id)
-        if self.external_agent_id:
-            span.set_attribute("external_agent_id", self.external_agent_id)
-
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Exit synchronous context."""
-        try:
-            if self._span:
-                if exc_type is not None:
-                    # Get the actual span object to set status
-                    span_obj = trace.get_current_span()
-                    if span_obj:
-                        span_obj.set_status(Status(StatusCode.ERROR, str(exc_val)))
-                else:
-                    span_obj = trace.get_current_span()
-                    if span_obj:
-                        span_obj.set_status(Status(StatusCode.OK))
-                        logger.info("Context block executed successfully")
-
-                self._span.__exit__(exc_type, exc_val, exc_tb)
-                self._span = None
-        finally:
-            self._cleanup_context()
-
-        return False  # Don't suppress exceptions
-
-    # Context manager methods for async
-    async def __aenter__(self):
-        """Enter asynchronous context."""
-        ctx = self._setup_context()
-
-        tracer = get_paid_tracer()
-        logger.info(f"Creating span for external_customer_id: {self.external_customer_id}")
-        self._span = tracer.start_as_current_span("parent_span", context=ctx)
-        span = self._span.__enter__()
-
-        span.set_attribute("external_customer_id", self.external_customer_id)
-        if self.external_agent_id:
-            span.set_attribute("external_agent_id", self.external_agent_id)
-
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        """Exit asynchronous context."""
-        try:
-            if self._span:
-                if exc_type is not None:
-                    # Get the actual span object to set status
-                    span_obj = trace.get_current_span()
-                    if span_obj:
-                        span_obj.set_status(Status(StatusCode.ERROR, str(exc_val)))
-                else:
-                    span_obj = trace.get_current_span()
-                    if span_obj:
-                        span_obj.set_status(Status(StatusCode.OK))
-                        logger.info("Async context block executed successfully")
-
-                self._span.__exit__(exc_type, exc_val, exc_tb)
-                self._span = None
-        finally:
-            self._cleanup_context()
-
-        return False  # Don't suppress exceptions
-
-    # Decorator functionality
-    def __call__(self, func: Callable) -> Callable:
-        """Use as a decorator."""
-        if asyncio.iscoroutinefunction(func):
-
-            @functools.wraps(func)
-            async def async_wrapper(*args, **kwargs):
-                # Auto-initialize tracing if not done
-                if get_token() is None:
-                    try:
-                        _initialize_tracing(None, self.collector_endpoint)
-                    except Exception as e:
-                        logger.error(f"Failed to auto-initialize tracing: {e}")
-                        # Fall back to executing function without tracing
-                        return await func(*args, **kwargs)
-
-                try:
-                    return await _trace_async(
-                        external_customer_id=self.external_customer_id,
-                        fn=func,
-                        external_agent_id=self.external_agent_id,
-                        tracing_token=self.tracing_token,
-                        store_prompt=self.store_prompt,
-                        metadata=self.metadata,
-                        args=args,
-                        kwargs=kwargs,
-                    )
-                except Exception as e:
-                    logger.error(f"Failed to trace async function {func.__name__}: {e}")
-                    raise e
-
-            return async_wrapper
-        else:
-
-            @functools.wraps(func)
-            def sync_wrapper(*args, **kwargs):
-                # Auto-initialize tracing if not done
-                if get_token() is None:
-                    try:
-                        _initialize_tracing(None, self.collector_endpoint)
-                    except Exception as e:
-                        logger.error(f"Failed to auto-initialize tracing: {e}")
-                        # Fall back to executing function without tracing
-                        return func(*args, **kwargs)
-
-                try:
-                    return _trace_sync(
-                        external_customer_id=self.external_customer_id,
-                        fn=func,
-                        external_agent_id=self.external_agent_id,
-                        tracing_token=self.tracing_token,
-                        store_prompt=self.store_prompt,
-                        metadata=self.metadata,
-                        args=args,
-                        kwargs=kwargs,
-                    )
-                except Exception as e:
-                    logger.error(f"Failed to trace sync function {func.__name__}: {e}")
-                    raise e
-
-            return sync_wrapper