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

paid/logger.py

@@ -0,0 +1,21 @@
+"""
+This file exports `logger` object for unified logging across the Paid SDK.
+"""
+
+import logging
+import os
+
+import dotenv
+
+# Configure logging
+_ = dotenv.load_dotenv()
+logger = logging.getLogger(__name__)
+
+# Set default log level to ERROR, allow override via PAID_LOG_LEVEL environment variable
+log_level_name = os.environ.get("PAID_LOG_LEVEL")
+if log_level_name is not None:
+    log_level = getattr(logging, log_level_name.upper(), logging.ERROR)
+else:
+    log_level = logging.ERROR
+
+logger.setLevel(log_level)

paid/tracing/__init__.py

@@ -1,18 +1,18 @@
 # Tracing module for OpenTelemetry integration
 from .autoinstrumentation import paid_autoinstrument
-from .tracing import (
-    generate_and_set_tracing_token,
+from .context_manager import paid_tracing
+from .distributed_tracing import (
     generate_tracing_token,
-    paid_tracing,
     set_tracing_token,
     unset_tracing_token,
 )
+from .signal import signal
 
 __all__ = [
-    "generate_and_set_tracing_token",
     "generate_tracing_token",
     "paid_autoinstrument",
     "paid_tracing",
     "set_tracing_token",
     "unset_tracing_token",
+    "signal",
 ]

paid/tracing/autoinstrumentation.py

@@ -8,7 +8,10 @@ sending traces to the Paid collector endpoint.
 from typing import List, Optional
 
 from . import tracing
-from .tracing import _initialize_tracing, logger
+from .tracing import initialize_tracing_
+from opentelemetry.trace import NoOpTracerProvider
+
+from paid.logger import logger
 
 # Safe imports for instrumentation libraries
 try:
@@ -89,9 +92,9 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
         >>> anthropic_client.messages.create(...)
     """
     # Initialize tracing if not already initialized
-    if not tracing.paid_tracer_provider:
+    if isinstance(tracing.paid_tracer_provider, NoOpTracerProvider):
         logger.info("Tracing not initialized, initializing automatically")
-        _initialize_tracing()
+        initialize_tracing_()
 
     # Default to all supported libraries if none specified
     if libraries is None:

paid/tracing/context_manager.py

@@ -0,0 +1,243 @@
+import asyncio
+import contextvars
+import functools
+from typing import Any, Callable, Dict, Optional, Tuple
+
+from . import distributed_tracing, tracing
+from .tracing import get_paid_tracer, get_token, initialize_tracing_, trace_async_, trace_sync_
+from opentelemetry import trace
+from opentelemetry.context import Context
+from opentelemetry.trace import NonRecordingSpan, Span, SpanContext, Status, StatusCode, TraceFlags
+
+from paid.logger import logger
+
+
+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] = tracing.DEFAULT_COLLECTOR_ENDPOINT,
+        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: Optional[Span] = None
+        self.span_ctx: Optional[Any] = None  # Context manager for the span
+        self.reset_tokens: Optional[
+            Tuple[
+                contextvars.Token[Optional[str]],  # external_customer_id
+                contextvars.Token[Optional[str]],  # external_agent_id
+                contextvars.Token[Optional[bool]],  # store_prompt
+                contextvars.Token[Optional[Dict[str, Any]]],  # metadata
+            ]
+        ] = None
+
+        if not get_token():
+            initialize_tracing_(None, self.collector_endpoint)
+
+    def _setup_context(self) -> Optional[Context]:
+        """Set up context variables and return OTEL context if needed."""
+
+        # Set context variables
+        reset_customer_id_ctx_token = tracing.paid_external_customer_id_var.set(self.external_customer_id)
+        reset_agent_id_ctx_token = tracing.paid_external_agent_id_var.set(self.external_agent_id)
+        reset_store_prompt_ctx_token = tracing.paid_store_prompt_var.set(self.store_prompt)
+        reset_user_metadata_ctx_token = tracing.paid_user_metadata_var.set(self.metadata)
+
+        # Store reset tokens for cleanup
+        self.reset_tokens = (
+            reset_customer_id_ctx_token,
+            reset_agent_id_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 = tracing.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=distributed_tracing.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_customer_id_ctx_token,
+                reset_agent_id_ctx_token,
+                reset_store_prompt_ctx_token,
+                reset_user_metadata_ctx_token,
+            ) = self.reset_tokens
+            tracing.paid_external_customer_id_var.reset(reset_customer_id_ctx_token)
+            tracing.paid_external_agent_id_var.reset(reset_agent_id_ctx_token)
+            tracing.paid_store_prompt_var.reset(reset_store_prompt_ctx_token)
+            tracing.paid_user_metadata_var.reset(reset_user_metadata_ctx_token)
+            self.reset_tokens = None
+
+    # Context manager methods for sync
+    def __enter__(self):
+        return self._enter_ctx()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._exit_ctx(exc_type, exc_val, exc_tb)
+
+    # Context manager methods for async
+    async def __aenter__(self):
+        return self._enter_ctx()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return self._exit_ctx(exc_type, exc_val, exc_tb)
+
+    def _enter_ctx(self):
+        ctx = self._setup_context()
+        tracer = get_paid_tracer()
+        logger.info(f"Creating span for external_customer_id: {self.external_customer_id}")
+        self.span_ctx = tracer.start_as_current_span("parent_span", context=ctx)
+        self.span = self.span_ctx.__enter__()
+        return self
+
+    def _exit_ctx(self, exc_type, exc_val, exc_tb):
+        """Exit synchronous context."""
+        try:
+            if self.span and self.span_ctx:
+                if exc_type is not None:
+                    self.span.set_status(Status(StatusCode.ERROR, str(exc_val)))
+                else:
+                    self.span.set_status(Status(StatusCode.OK))
+                    logger.info("Context block executed successfully")
+
+                self.span_ctx.__exit__(exc_type, exc_val, exc_tb)
+                self.span_ctx = None
+                self.span = None
+
+        finally:
+            self._cleanup_context()
+
+        return False  # Don't suppress exceptions
+
+    # Decorator functionality
+    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
+        """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

paid/tracing/distributed_tracing.py

@@ -0,0 +1,113 @@
+import warnings
+
+from . import tracing
+from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
+
+otel_id_generator = RandomIdGenerator()
+
+
+def generate_tracing_token() -> int:
+    """
+    Generate a unique tracing token without setting it in the context.
+
+    Use this when you want to generate a trace ID to store or pass to another
+    process/service without immediately associating it with the current tracing context.
+    The token can later be used with set_tracing_token() to link traces across
+    different execution contexts.
+
+    Returns:
+        int: A unique OpenTelemetry trace ID.
+
+    Notes:
+        - This function only generates and returns the token; it does NOT set it in the context.
+        - Use this when you need to store the token separately before setting it.
+
+    Examples:
+        Generate token to store for later use:
+
+            from paid.tracing import generate_tracing_token, set_tracing_token
+
+            # Process 1: Generate and store
+            token = generate_tracing_token()
+            save_to_database("task_123", token)
+
+            # Process 2: Retrieve and use
+            token = load_from_database("task_123")
+            set_tracing_token(token)
+
+            @paid_tracing(external_customer_id="cust_123", external_agent_id="agent_456")
+            def process_task():
+                # This trace is now linked to the same token
+                pass
+
+    See Also:
+        set_tracing_token: Set a previously generated token.
+    """
+    return otel_id_generator.generate_trace_id()
+
+
+def set_tracing_token(token: int):
+    """
+    Deprecated: Pass tracing_token directly to @paid_tracing() decorator instead.
+
+    This function is deprecated and will be removed in a future version.
+    Use the tracing_token parameter in @paid_tracing() to link traces across processes.
+
+    Instead of:
+        token = load_from_storage("workflow_123")
+        set_tracing_token(token)
+        @paid_tracing(external_customer_id="cust_123", external_agent_id="agent_456")
+        def process_workflow():
+            ...
+        unset_tracing_token()
+
+    Use:
+        token = load_from_storage("workflow_123")
+
+        @paid_tracing(
+            external_customer_id="cust_123",
+            external_agent_id="agent_456",
+            tracing_token=token
+        )
+        def process_workflow():
+            ...
+
+    Parameters:
+        token (int): A tracing token (for backward compatibility only).
+
+    Old behavior (for reference):
+        This function set a token in the context, so all subsequent @paid_tracing() calls
+        would use it automatically until unset_tracing_token() was called.
+    """
+    warnings.warn(
+        "set_tracing_token() is deprecated and will be removed in a future version. "
+        "Pass tracing_token directly to @paid_tracing(tracing_token=...) decorator instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    _ = tracing.paid_trace_id_var.set(token)
+
+
+def unset_tracing_token():
+    """
+    Deprecated: No longer needed. Use tracing_token parameter in @paid_tracing() instead.
+
+    This function is deprecated and will be removed in a future version.
+    Since tracing_token is now passed directly to @paid_tracing(), there's no need
+    to manually set/unset tokens in the context.
+
+    Old behavior (for reference):
+        This function unset a token previously set by set_tracing_token(), allowing subsequent @paid_tracing() calls
+        to have independent traces.
+
+    Migration:
+        If you were using set_tracing_token() + unset_tracing_token() pattern,
+        simply pass the token directly to @paid_tracing(tracing_token=...) instead.
+    """
+    warnings.warn(
+        "unset_tracing_token() is deprecated and will be removed in a future version. "
+        "Use tracing_token parameter in @paid_tracing(tracing_token=...) decorator instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    _ = tracing.paid_trace_id_var.set(None)

paid/tracing/signal.py

@@ -1,44 +1,69 @@
 import json
 import typing
 
-from .tracing import get_paid_tracer, logger, paid_external_agent_id_var, paid_external_customer_id_var, paid_token_var
-from opentelemetry import trace
+from .tracing import get_paid_tracer
 from opentelemetry.trace import Status, StatusCode
 
+from paid.logger import logger
 
-def _signal(event_name: str, enable_cost_tracing: bool, data: typing.Optional[typing.Dict] = None):
-    if not event_name:
-        logger.error("Event name is required for signal.")
-        return
-
-    # Check if there's an active span (from capture())
-    current_span = trace.get_current_span()
-    if current_span == trace.INVALID_SPAN:
-        logger.error("Cannot send signal: you should call signal() within tracing context")
-        return
-
-    external_customer_id = paid_external_customer_id_var.get()
-    external_agent_id = paid_external_agent_id_var.get()
-    token = paid_token_var.get()
-    if not (external_customer_id and external_agent_id and token):
-        logger.error(
-            f"Missing some of: external_customer_id: {external_customer_id}, "
-            f"external_agent_id: {external_agent_id}, or token. "
-            f"You should call signal() within a tracing context"
-        )
-        return
+
+def signal(event_name: str, enable_cost_tracing: bool = False, data: typing.Optional[dict[str, typing.Any]] = None):
+    """
+    Emit a signal within a tracing context.
+
+    This function must be called within an active @paid_tracing() context (decorator or context manager).
+
+    Parameters
+    ----------
+    event_name : str
+        The name of the signal (e.g., "user_signup", "payment_processed", "task_completed").
+    enable_cost_tracing : bool, optional
+        If True, associates this signal with cost/usage traces from the same tracing context.
+        Should only be called once per tracing context to avoid multiple signals referring to the same costs.
+        Default is False.
+    data : dict[str, Any], optional
+        Additional context data to attach to the signal. Will be JSON-serialized and stored
+        as a span attribute. Example: {"user_id": "123", "amount": 99.99}.
+
+    Notes
+    -----
+    - Signal must be called within a @paid_tracing() context; calling outside will log an error and return.
+    - Use enable_cost_tracing=True when you want to mark the point where costs were incurred
+      and link that signal to cost/usage data from the same trace.
+
+    Examples
+    --------
+    Basic signal within a tracing context:
+
+        from paid.tracing import paid_tracing, signal
+
+        @paid_tracing(external_customer_id="cust_123", external_agent_id="agent_456")
+        def process_order(order_id):
+            # ... do work ...
+            signal("order_processed", data={"order_id": order_id})
+
+    Signal with cost tracking:
+
+        @paid_tracing(external_customer_id="cust_123", external_agent_id="agent_456")
+        def call_ai_api():
+            # ... call AI provider ...
+            signal("ai_api_call_complete", enable_cost_tracing=True)
+
+    Using context manager:
+
+        with paid_tracing(external_customer_id="cust_123", external_agent_id="agent_456"):
+            # ... do work ...
+            signal("milestone_reached", data={"step": "validation_complete"})
+    """
 
     tracer = get_paid_tracer()
     with tracer.start_as_current_span("signal") as span:
-        attributes: typing.Dict[str, typing.Union[str, bool, int, float]] = {
-            "external_customer_id": external_customer_id,
-            "external_agent_id": external_agent_id,
+        attributes: dict[str, typing.Union[str, bool, int, float]] = {
             "event_name": event_name,
         }
 
         if enable_cost_tracing:
             # let the app know to associate this signal with cost traces
-            attributes["enable_cost_tracing"] = True
             if data is None:
                 data = {"paid": {"enable_cost_tracing": True}}
             else:
@@ -46,7 +71,12 @@ def _signal(event_name: str, enable_cost_tracing: bool, data: typing.Optional[ty
 
         # optional data (ex. manual cost tracking)
         if data:
-            attributes["data"] = json.dumps(data)
+            try:
+                attributes["data"] = json.dumps(data)
+            except (TypeError, ValueError) as e:
+                logger.error(f"Failed to serialize data into JSON for signal [{event_name}]: {e}")
+                if enable_cost_tracing:
+                    attributes["data"] = json.dumps({"paid": {"enable_cost_tracing": True}})
 
         span.set_attributes(attributes)
         # Mark span as successful