paid-python 0.1.1__tar.gz → 0.2.1__tar.gz

{paid_python-0.1.1 → paid_python-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: paid-python
-Version: 0.1.1
+Version: 0.2.1
 Summary: 
 Requires-Python: >=3.9,<3.14
 Classifier: Intended Audience :: Developers
@@ -25,6 +25,7 @@ Requires-Dist: opentelemetry-api (>=1.23.0)
 Requires-Dist: opentelemetry-exporter-otlp-proto-http (>=1.23.0)
 Requires-Dist: opentelemetry-instrumentation-anthropic (>=0.47.0)
 Requires-Dist: opentelemetry-instrumentation-google-generativeai (>=0.47.0)
+Requires-Dist: opentelemetry-instrumentation-langchain (>=0.47.0)
 Requires-Dist: opentelemetry-instrumentation-openai (>=0.47.0)
 Requires-Dist: opentelemetry-sdk (>=1.23.0)
 Requires-Dist: pydantic (>=1.9.0)
@@ -162,6 +163,13 @@ Both approaches:
 - Gracefully fall back to normal execution if tracing fails
 - Support the same parameters: `external_customer_id`, `external_agent_id`, `tracing_token`, `store_prompt`, `metadata`
 
+* Note - if it happens that you're calling `paid_tracing` from non-main thread, then it's advised to initialize from main thread:
+```python
+from paid.tracing import initialize_tracing
+initialize_tracing()
+```
+* `initialize_tracing` also accepts optional arguments like OTEL collector endpoint and api key if you want to reroute your tracing somewhere else :)
+
 ### Using the Paid wrappers
 
 You can track usage costs by using Paid wrappers around your AI provider's SDK.
@@ -182,9 +190,11 @@ Example usage:
 
 ```python
 from openai import OpenAI
-from paid.tracing import paid_tracing
+from paid.tracing import paid_tracing, initialize_tracing
 from paid.tracing.wrappers.openai import PaidOpenAI
 
+initialize_tracing()
+
 openAIClient = PaidOpenAI(OpenAI(
     # This is the default and can be omitted
     api_key="<OPENAI_API_KEY>",
@@ -212,10 +222,12 @@ You can attach custom metadata to your traces by passing a `metadata` dictionary
 <Tabs>
   <Tab title="Python - Decorator">
     ```python
-    from paid.tracing import paid_tracing, signal
+    from paid.tracing import paid_tracing, signal, initialize_tracing
     from paid.tracing.wrappers import PaidOpenAI
     from openai import OpenAI
 
+    initialize_tracing()
+
     openai_client = PaidOpenAI(OpenAI(api_key="<OPENAI_API_KEY>"))
 
     @paid_tracing(
@@ -243,10 +255,12 @@ You can attach custom metadata to your traces by passing a `metadata` dictionary
 
   <Tab title="Python - Context Manager">
     ```python
-    from paid.tracing import paid_tracing, signal
+    from paid.tracing import paid_tracing, signal, initialize_tracing
     from paid.tracing.wrappers import PaidOpenAI
     from openai import OpenAI
 
+    initialize_tracing()
+
     openai_client = PaidOpenAI(OpenAI(api_key="<OPENAI_API_KEY>"))
 
     def process_event(event):
@@ -305,14 +319,14 @@ For maximum convenience, you can use OpenTelemetry auto-instrumentation to autom
 
 ```python
 from paid import Paid
-from paid.tracing import paid_autoinstrument
+from paid.tracing import paid_autoinstrument, initialize_tracing
 from openai import OpenAI
 
 # Initialize Paid SDK
 client = Paid(token="PAID_API_KEY")
+initialize_tracing()
 
-# Enable auto-instrumentation for all supported libraries
-paid_autoinstrument()  # instruments all available: anthropic, gemini, openai, openai-agents, bedrock
+paid_autoinstrument()  # instruments all available: anthropic, gemini, openai, openai-agents, bedrock, langchain
 
 # Now all OpenAI calls will be automatically traced
 openai_client = OpenAI(api_key="<OPENAI_API_KEY>")
@@ -338,6 +352,7 @@ gemini             - Google Generative AI (google-generativeai)
 openai             - OpenAI Python SDK
 openai-agents      - OpenAI Agents SDK
 bedrock            - AWS Bedrock (boto3)
+langchain          - LangChain framework
 ```
 
 #### Selective Instrumentation
@@ -437,10 +452,12 @@ For such cases, you can pass a tracing token directly to `@paid_tracing()` or co
 The simplest way to implement distributed tracing is to pass the token directly to the decorator or context manager:
 
 ```python
-from paid.tracing import paid_tracing, signal, generate_tracing_token
+from paid.tracing import paid_tracing, signal, generate_tracing_token, initialize_tracing
 from paid.tracing.wrappers.openai import PaidOpenAI
 from openai import OpenAI
 
+initialize_tracing()
+
 openai_client = PaidOpenAI(OpenAI(api_key="<OPENAI_API_KEY>"))
 
 # Process 1: Generate token and do initial work
@@ -482,10 +499,12 @@ process_part_2()
 Using context manager instead of decorator:
 
 ```python
-from paid.tracing import paid_tracing, signal, generate_tracing_token
+from paid.tracing import paid_tracing, signal, generate_tracing_token, initialize_tracing
 from paid.tracing.wrappers.openai import PaidOpenAI
 from openai import OpenAI
 
+initialize_tracing()
+
 # Initialize
 openai_client = PaidOpenAI(OpenAI(api_key="<OPENAI_API_KEY>"))
 
@@ -643,9 +662,11 @@ The `@paid_tracing` decorator automatically handles both sync and async function
 
 ```python
 from openai import AsyncOpenAI
-from paid.tracing import paid_tracing
+from paid.tracing import paid_tracing, initialize_tracing
 from paid.tracing.wrappers.openai import PaidAsyncOpenAI
 
+initialize_tracing()
+
 # Wrap the async OpenAI client
 openai_client = PaidAsyncOpenAI(AsyncOpenAI(api_key="<OPENAI_API_KEY>"))
 
@@ -669,10 +690,12 @@ await generate_image()
 The `signal()` function works seamlessly in async contexts:
 
 ```python
-from paid.tracing import paid_tracing, signal
+from paid.tracing import paid_tracing, signal, initialize_tracing
 from paid.tracing.wrappers.openai import PaidAsyncOpenAI
 from openai import AsyncOpenAI
 
+initialize_tracing()
+
 openai_client = PaidAsyncOpenAI(AsyncOpenAI(api_key="<OPENAI_API_KEY>"))
 
 @paid_tracing("your_external_customer_id", "your_external_agent_id")

{paid_python-0.1.1 → paid_python-0.2.1}/README.md

@@ -127,6 +127,13 @@ Both approaches:
 - Gracefully fall back to normal execution if tracing fails
 - Support the same parameters: `external_customer_id`, `external_agent_id`, `tracing_token`, `store_prompt`, `metadata`
 
+* Note - if it happens that you're calling `paid_tracing` from non-main thread, then it's advised to initialize from main thread:
+```python
+from paid.tracing import initialize_tracing
+initialize_tracing()
+```
+* `initialize_tracing` also accepts optional arguments like OTEL collector endpoint and api key if you want to reroute your tracing somewhere else :)
+
 ### Using the Paid wrappers
 
 You can track usage costs by using Paid wrappers around your AI provider's SDK.
@@ -147,9 +154,11 @@ Example usage:
 
 ```python
 from openai import OpenAI
-from paid.tracing import paid_tracing
+from paid.tracing import paid_tracing, initialize_tracing
 from paid.tracing.wrappers.openai import PaidOpenAI
 
+initialize_tracing()
+
 openAIClient = PaidOpenAI(OpenAI(
     # This is the default and can be omitted
     api_key="<OPENAI_API_KEY>",
@@ -177,10 +186,12 @@ You can attach custom metadata to your traces by passing a `metadata` dictionary
 <Tabs>
   <Tab title="Python - Decorator">
     ```python
-    from paid.tracing import paid_tracing, signal
+    from paid.tracing import paid_tracing, signal, initialize_tracing
     from paid.tracing.wrappers import PaidOpenAI
     from openai import OpenAI
 
+    initialize_tracing()
+
     openai_client = PaidOpenAI(OpenAI(api_key="<OPENAI_API_KEY>"))
 
     @paid_tracing(
@@ -208,10 +219,12 @@ You can attach custom metadata to your traces by passing a `metadata` dictionary
 
   <Tab title="Python - Context Manager">
     ```python
-    from paid.tracing import paid_tracing, signal
+    from paid.tracing import paid_tracing, signal, initialize_tracing
     from paid.tracing.wrappers import PaidOpenAI
     from openai import OpenAI
 
+    initialize_tracing()
+
     openai_client = PaidOpenAI(OpenAI(api_key="<OPENAI_API_KEY>"))
 
     def process_event(event):
@@ -270,14 +283,14 @@ For maximum convenience, you can use OpenTelemetry auto-instrumentation to autom
 
 ```python
 from paid import Paid
-from paid.tracing import paid_autoinstrument
+from paid.tracing import paid_autoinstrument, initialize_tracing
 from openai import OpenAI
 
 # Initialize Paid SDK
 client = Paid(token="PAID_API_KEY")
+initialize_tracing()
 
-# Enable auto-instrumentation for all supported libraries
-paid_autoinstrument()  # instruments all available: anthropic, gemini, openai, openai-agents, bedrock
+paid_autoinstrument()  # instruments all available: anthropic, gemini, openai, openai-agents, bedrock, langchain
 
 # Now all OpenAI calls will be automatically traced
 openai_client = OpenAI(api_key="<OPENAI_API_KEY>")
@@ -303,6 +316,7 @@ gemini             - Google Generative AI (google-generativeai)
 openai             - OpenAI Python SDK
 openai-agents      - OpenAI Agents SDK
 bedrock            - AWS Bedrock (boto3)
+langchain          - LangChain framework
 ```
 
 #### Selective Instrumentation
@@ -402,10 +416,12 @@ For such cases, you can pass a tracing token directly to `@paid_tracing()` or co
 The simplest way to implement distributed tracing is to pass the token directly to the decorator or context manager:
 
 ```python
-from paid.tracing import paid_tracing, signal, generate_tracing_token
+from paid.tracing import paid_tracing, signal, generate_tracing_token, initialize_tracing
 from paid.tracing.wrappers.openai import PaidOpenAI
 from openai import OpenAI
 
+initialize_tracing()
+
 openai_client = PaidOpenAI(OpenAI(api_key="<OPENAI_API_KEY>"))
 
 # Process 1: Generate token and do initial work
@@ -447,10 +463,12 @@ process_part_2()
 Using context manager instead of decorator:
 
 ```python
-from paid.tracing import paid_tracing, signal, generate_tracing_token
+from paid.tracing import paid_tracing, signal, generate_tracing_token, initialize_tracing
 from paid.tracing.wrappers.openai import PaidOpenAI
 from openai import OpenAI
 
+initialize_tracing()
+
 # Initialize
 openai_client = PaidOpenAI(OpenAI(api_key="<OPENAI_API_KEY>"))
 
@@ -608,9 +626,11 @@ The `@paid_tracing` decorator automatically handles both sync and async function
 
 ```python
 from openai import AsyncOpenAI
-from paid.tracing import paid_tracing
+from paid.tracing import paid_tracing, initialize_tracing
 from paid.tracing.wrappers.openai import PaidAsyncOpenAI
 
+initialize_tracing()
+
 # Wrap the async OpenAI client
 openai_client = PaidAsyncOpenAI(AsyncOpenAI(api_key="<OPENAI_API_KEY>"))
 
@@ -634,10 +654,12 @@ await generate_image()
 The `signal()` function works seamlessly in async contexts:
 
 ```python
-from paid.tracing import paid_tracing, signal
+from paid.tracing import paid_tracing, signal, initialize_tracing
 from paid.tracing.wrappers.openai import PaidAsyncOpenAI
 from openai import AsyncOpenAI
 
+initialize_tracing()
+
 openai_client = PaidAsyncOpenAI(AsyncOpenAI(api_key="<OPENAI_API_KEY>"))
 
 @paid_tracing("your_external_customer_id", "your_external_agent_id")

{paid_python-0.1.1 → paid_python-0.2.1}/pyproject.toml

@@ -3,7 +3,7 @@ name = "paid-python"
 
 [tool.poetry]
 name = "paid-python"
-version = "0.1.1"
+version = "0.2.1"
 description = ""
 readme = "README.md"
 authors = []
@@ -46,6 +46,7 @@ python-dotenv = ">=0.15.0"
 opentelemetry-instrumentation-anthropic = ">=0.47.0"
 opentelemetry-instrumentation-google-generativeai = ">=0.47.0"
 opentelemetry-instrumentation-openai = ">=0.47.0"
+opentelemetry-instrumentation-langchain = ">=0.47.0"
 openinference-instrumentation-openai-agents = ">=1.0.0"
 openinference-instrumentation-bedrock = ">=0.1.0"
 

{paid_python-0.1.1 → paid_python-0.2.1}/src/paid/client.py

@@ -17,7 +17,7 @@ from .tracing.distributed_tracing import (
 from .tracing.signal import signal
 from .tracing.tracing import (
     DEFAULT_COLLECTOR_ENDPOINT,
-    initialize_tracing_,
+    initialize_tracing,
     trace_async_,
     trace_sync_,
 )
@@ -114,7 +114,7 @@ class Paid:
             stacklevel=2,
         )
         token = self._client_wrapper._get_token()
-        initialize_tracing_(token, collector_endpoint=collector_endpoint)
+        initialize_tracing(token, collector_endpoint=collector_endpoint)
 
     def generate_tracing_token(self) -> int:
         """
@@ -394,7 +394,7 @@ class AsyncPaid:
             stacklevel=2,
         )
         token = self._client_wrapper._get_token()
-        initialize_tracing_(token, collector_endpoint=collector_endpoint)
+        initialize_tracing(token, collector_endpoint=collector_endpoint)
 
     def generate_tracing_token(self) -> int:
         """

{paid_python-0.1.1 → paid_python-0.2.1}/src/paid/tracing/__init__.py

@@ -7,11 +7,13 @@ from .distributed_tracing import (
     unset_tracing_token,
 )
 from .signal import signal
+from .tracing import initialize_tracing
 
 __all__ = [
     "generate_tracing_token",
     "paid_autoinstrument",
     "paid_tracing",
+    "initialize_tracing",
     "set_tracing_token",
     "unset_tracing_token",
     "signal",

{paid_python-0.1.1 → paid_python-0.2.1}/src/paid/tracing/autoinstrumentation.py

@@ -8,7 +8,7 @@ sending traces to the Paid collector endpoint.
 from typing import List, Optional
 
 from . import tracing
-from .tracing import initialize_tracing_
+from .tracing import initialize_tracing
 from opentelemetry.trace import NoOpTracerProvider
 
 from paid.logger import logger
@@ -49,6 +49,13 @@ try:
 except ImportError:
     BEDROCK_AVAILABLE = False
 
+try:
+    from opentelemetry.instrumentation.langchain import LangchainInstrumentor
+
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+
 
 # Track which instrumentors have been initialized
 _initialized_instrumentors: List[str] = []
@@ -69,6 +76,7 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
                   - "openai": OpenAI library
                   - "openai-agents": OpenAI Agents SDK
                   - "bedrock": AWS Bedrock
+                  - "langchain": LangChain library
                   If None, all supported libraries that are installed will be instrumented.
 
     Note:
@@ -94,11 +102,11 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
     # Initialize tracing if not already initialized
     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:
-        libraries = ["anthropic", "gemini", "openai", "openai-agents", "bedrock"]
+        libraries = ["anthropic", "gemini", "openai", "openai-agents", "bedrock", "langchain"]
 
     for library in libraries:
         if library in _initialized_instrumentors:
@@ -115,9 +123,11 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
             _instrument_openai_agents()
         elif library == "bedrock":
             _instrument_bedrock()
+        elif library == "langchain":
+            _instrument_langchain()
         else:
             logger.warning(
-                f"Unknown library '{library}' - supported libraries: anthropic, gemini, openai, openai-agents, bedrock"
+                f"Unknown library '{library}' - supported libraries: anthropic, gemini, openai, openai-agents, bedrock, langchain"
             )
 
     logger.info(f"Auto-instrumentation enabled for: {', '.join(_initialized_instrumentors)}")
@@ -196,3 +206,19 @@ def _instrument_bedrock() -> None:
 
     _initialized_instrumentors.append("bedrock")
     logger.info("Bedrock auto-instrumentation enabled")
+
+
+def _instrument_langchain() -> None:
+    """
+    Instrument LangChain using opentelemetry-instrumentation-langchain.
+    """
+    if not LANGCHAIN_AVAILABLE:
+        logger.warning("LangChain instrumentation library not available, skipping instrumentation")
+        return
+
+    # Instrument LangChain with Paid's tracer provider
+    LangchainInstrumentor(disable_trace_context_propagation=True).instrument(tracer_provider=tracing.paid_tracer_provider)
+
+
+    _initialized_instrumentors.append("langchain")
+    logger.info("LangChain auto-instrumentation enabled")

paid_python-0.2.1/src/paid/tracing/context_data.py

@@ -0,0 +1,60 @@
+import contextvars
+from typing import Any, Optional
+
+from paid.logger import logger
+
+
+# this class is used like a namespace, it's not for instantiation
+class ContextData:
+    _EXTERNAL_CUSTOMER_ID = contextvars.ContextVar[Optional[str]]("external_customer_id", default=None)
+    _EXTERNAL_AGENT_ID = contextvars.ContextVar[Optional[str]]("external_agent_id", default=None)
+    _TRACE_ID = contextvars.ContextVar[Optional[int]]("trace_id", default=None)
+    _STORE_PROMPT = contextvars.ContextVar[Optional[bool]]("store_prompt", default=False)
+    _USER_METADATA = contextvars.ContextVar[Optional[dict[str, Any]]]("user_metadata", default=None)
+
+    _context: dict[str, contextvars.ContextVar] = {
+        "external_customer_id": _EXTERNAL_CUSTOMER_ID,
+        "external_agent_id": _EXTERNAL_AGENT_ID,
+        "trace_id": _TRACE_ID,
+        "store_prompt": _STORE_PROMPT,
+        "user_metadata": _USER_METADATA,
+    }
+
+    # Use ContextVar for reset tokens to avoid race conditions in async/concurrent scenarios
+    _reset_tokens: contextvars.ContextVar[Optional[dict[str, Any]]] = contextvars.ContextVar(
+        "reset_tokens", default=None
+    )
+
+    @classmethod
+    def _get_or_create_reset_tokens(cls) -> dict[str, Any]:
+        """Get the reset tokens dict for this context, creating a new one if needed."""
+        reset_tokens = cls._reset_tokens.get()
+        if reset_tokens is None:
+            reset_tokens = {}
+            cls._reset_tokens.set(reset_tokens)
+        return reset_tokens
+
+    @classmethod
+    def get_context(cls) -> dict[str, Any]:
+        return {key: var.get() for key, var in cls._context.items()}
+
+    @classmethod
+    def get_context_key(cls, key: str) -> Any:
+        return cls._context[key].get() if key in cls._context else None
+
+    @classmethod
+    def set_context_key(cls, key: str, value: Any) -> None:
+        if key not in cls._context:
+            logger.warning(f"Invalid context key: {key}")
+            return
+        reset_token = cls._context[key].set(value)
+        reset_tokens = cls._get_or_create_reset_tokens()
+        reset_tokens[key] = reset_token
+
+    @classmethod
+    def reset_context(cls) -> None:
+        reset_tokens = cls._reset_tokens.get()
+        if reset_tokens:
+            for key, reset_token in reset_tokens.items():
+                cls._context[key].reset(reset_token)
+            reset_tokens.clear()

{paid_python-0.1.1 → paid_python-0.2.1}/src/paid/tracing/context_manager.py

@@ -1,10 +1,10 @@
 import asyncio
-import contextvars
 import functools
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Any, Callable, Dict, Optional
 
 from . import distributed_tracing, tracing
-from .tracing import get_paid_tracer, get_token, initialize_tracing_, trace_async_, trace_sync_
+from .context_data import ContextData
+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
@@ -78,39 +78,23 @@ class paid_tracing:
         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)
+            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,
-        )
+        ContextData.set_context_key("external_customer_id", self.external_customer_id)
+        ContextData.set_context_key("external_agent_id", self.external_agent_id)
+        ContextData.set_context_key("store_prompt", self.store_prompt)
+        ContextData.set_context_key("user_metadata", self.metadata)
 
         # Handle distributed tracing token
         override_trace_id = self.tracing_token
         if not override_trace_id:
-            override_trace_id = tracing.paid_trace_id_var.get()
+            override_trace_id = ContextData.get_context_key("trace_id")
 
         ctx: Optional[Context] = None
         if override_trace_id is not None:
@@ -126,18 +110,7 @@ class paid_tracing:
 
     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
+        ContextData.reset_context()
 
     # Context manager methods for sync
     def __enter__(self):
@@ -190,7 +163,7 @@ class paid_tracing:
                 # Auto-initialize tracing if not done
                 if get_token() is None:
                     try:
-                        initialize_tracing_(None, self.collector_endpoint)
+                        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
@@ -219,7 +192,7 @@ class paid_tracing:
                 # Auto-initialize tracing if not done
                 if get_token() is None:
                     try:
-                        initialize_tracing_(None, self.collector_endpoint)
+                        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

{paid_python-0.1.1 → paid_python-0.2.1}/src/paid/tracing/distributed_tracing.py

@@ -1,6 +1,6 @@
 import warnings
 
-from . import tracing
+from .context_data import ContextData
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 
 otel_id_generator = RandomIdGenerator()
@@ -85,7 +85,7 @@ def set_tracing_token(token: int):
         DeprecationWarning,
         stacklevel=2,
     )
-    _ = tracing.paid_trace_id_var.set(token)
+    ContextData.set_context_key("trace_id", token)
 
 
 def unset_tracing_token():
@@ -110,4 +110,4 @@ def unset_tracing_token():
         DeprecationWarning,
         stacklevel=2,
     )
-    _ = tracing.paid_trace_id_var.set(None)
+    ContextData.set_context_key("trace_id", None)

{paid_python-0.1.1 → paid_python-0.2.1}/src/paid/tracing/tracing.py

@@ -1,13 +1,13 @@
 # Initializing tracing for OTLP
 import asyncio
 import atexit
-import contextvars
 import os
 import signal
 from typing import Any, Awaitable, Callable, Dict, Optional, Tuple, TypeVar, Union
 
 import dotenv
 from . import distributed_tracing
+from .context_data import ContextData
 from opentelemetry import trace
 from opentelemetry.context import Context
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
@@ -23,24 +23,6 @@ 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(
-    "paid_external_customer_id", default=None
-)
-paid_external_agent_id_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
-    "paid_external_agent_id", default=None
-)
-# trace id storage (generated from token)
-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
-)
-# user metadata
-paid_user_metadata_var: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar(
-    "paid_user_metadata", default=None
-)
-
 T = TypeVar("T")
 
 
@@ -102,15 +84,15 @@ class PaidSpanProcessor(SpanProcessor):
             span.update_name(f"{self.SPAN_NAME_PREFIX}{span.name}")
 
         # Add customer and agent IDs from context
-        customer_id = paid_external_customer_id_var.get()
+        customer_id = ContextData.get_context_key("external_customer_id")
         if customer_id:
             span.set_attribute("external_customer_id", customer_id)
 
-        agent_id = paid_external_agent_id_var.get()
+        agent_id = ContextData.get_context_key("external_agent_id")
         if agent_id:
             span.set_attribute("external_agent_id", agent_id)
 
-        metadata = paid_user_metadata_var.get()
+        metadata = ContextData.get_context_key("user_metadata")
         if metadata:
             metadata_attributes: dict[str, Any] = {}
 
@@ -131,7 +113,7 @@ class PaidSpanProcessor(SpanProcessor):
 
     def on_end(self, span: ReadableSpan) -> None:
         """Filter out prompt and response contents unless explicitly asked to store"""
-        store_prompt = paid_store_prompt_var.get()
+        store_prompt = ContextData.get_context_key("store_prompt")
         if store_prompt:
             return
 
@@ -155,8 +137,43 @@ class PaidSpanProcessor(SpanProcessor):
         """Called to force flush. Always returns True since there's nothing to flush."""
         return True
 
+def setup_graceful_termination():
+    def flush_traces():
+        try:
+            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}")
+
+    def create_chained_signal_handler(signum: int):
+        current_handler = signal.getsignal(signum)
+
+        def chained_handler(_signum, frame):
+            logger.warning(f"Received signal {_signum}, flushing traces")
+            flush_traces()
+            # Restore the original handler
+            signal.signal(_signum, current_handler)
+            # Re-raise the signal to let the original handler (or default) handle it
+            os.kill(os.getpid(), _signum)
+
+        return chained_handler
+
+    try:
+        # This is already done by default OTEL shutdown,
+        # but user might turn that off - so register it explicitly
+        atexit.register(flush_traces)
 
-def initialize_tracing_(api_key: Optional[str] = None, collector_endpoint: Optional[str] = DEFAULT_COLLECTOR_ENDPOINT):
+        # signal handlers
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            signal.signal(sig, create_chained_signal_handler(sig))
+    except Exception as e:
+        logger.warning(f"Could not set up termination handlers: {e}"
+                       "\nConsider calling initialize_tracing() from the main thread during app initialization if you don't already")
+
+
+def initialize_tracing(api_key: Optional[str] = None, collector_endpoint: Optional[str] = DEFAULT_COLLECTOR_ENDPOINT):
     """
     Initialize OpenTelemetry with OTLP exporter for Paid backend.
 
@@ -203,36 +220,7 @@ def initialize_tracing_(api_key: Optional[str] = None, collector_endpoint: Optio
         span_processor = SimpleSpanProcessor(otlp_exporter)
         paid_tracer_provider.add_span_processor(span_processor)
 
-        # Terminate gracefully and don't lose traces
-        def flush_traces():
-            try:
-                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}")
-
-        def create_chained_signal_handler(signum: int):
-            current_handler = signal.getsignal(signum)
-
-            def chained_handler(_signum, frame):
-                logger.warning(f"Received signal {_signum}, flushing traces")
-                flush_traces()
-                # Restore the original handler
-                signal.signal(_signum, current_handler)
-                # Re-raise the signal to let the original handler (or default) handle it
-                os.kill(os.getpid(), _signum)
-
-            return chained_handler
-
-        # This is already done by default OTEL shutdown,
-        # but user might turn that off - so register it explicitly
-        atexit.register(flush_traces)
-
-        # Handle signals
-        for sig in (signal.SIGINT, signal.SIGTERM):
-            signal.signal(sig, create_chained_signal_handler(sig))
+        setup_graceful_termination() # doesn't throw
 
         logger.info("Paid tracing initialized successfully - collector at %s", collector_endpoint)
     except Exception as e:
@@ -293,15 +281,15 @@ def trace_sync_(
     kwargs = kwargs or {}
 
     # Set context variables for access by nested spans
-    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_store_prompt_ctx_token = paid_store_prompt_var.set(store_prompt)
-    reset_user_metadata_ctx_token = paid_user_metadata_var.set(metadata)
+    ContextData.set_context_key("external_customer_id", external_customer_id)
+    ContextData.set_context_key("external_agent_id", external_agent_id)
+    ContextData.set_context_key("store_prompt", store_prompt)
+    ContextData.set_context_key("user_metadata", metadata)
 
     # If user set trace context manually
     override_trace_id = tracing_token
     if not override_trace_id:
-        override_trace_id = paid_trace_id_var.get()
+        override_trace_id = ContextData.get_context_key("trace_id")
     ctx: Optional[Context] = None
     if override_trace_id is not None:
         span_context = SpanContext(
@@ -325,10 +313,7 @@ def trace_sync_(
                 span.set_status(Status(StatusCode.ERROR, str(error)))
                 raise
     finally:
-        paid_external_customer_id_var.reset(reset_customer_id_ctx_token)
-        paid_external_agent_id_var.reset(reset_agent_id_ctx_token)
-        paid_store_prompt_var.reset(reset_store_prompt_ctx_token)
-        paid_user_metadata_var.reset(reset_user_metadata_ctx_token)
+        ContextData.reset_context()
 
 
 async def trace_async_(
@@ -367,15 +352,15 @@ async def trace_async_(
     kwargs = kwargs or {}
 
     # Set context variables for access by nested spans
-    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_store_prompt_ctx_token = paid_store_prompt_var.set(store_prompt)
-    reset_user_metadata_ctx_token = paid_user_metadata_var.set(metadata)
+    ContextData.set_context_key("external_customer_id", external_customer_id)
+    ContextData.set_context_key("external_agent_id", external_agent_id)
+    ContextData.set_context_key("store_prompt", store_prompt)
+    ContextData.set_context_key("user_metadata", metadata)
 
     # If user set trace context manually
     override_trace_id = tracing_token
     if not override_trace_id:
-        override_trace_id = paid_trace_id_var.get()
+        override_trace_id = ContextData.get_context_key("trace_id")
     ctx: Optional[Context] = None
     if override_trace_id is not None:
         span_context = SpanContext(
@@ -402,7 +387,4 @@ async def trace_async_(
                 span.set_status(Status(StatusCode.ERROR, str(error)))
                 raise
     finally:
-        paid_external_customer_id_var.reset(reset_customer_id_ctx_token)
-        paid_external_agent_id_var.reset(reset_agent_id_ctx_token)
-        paid_store_prompt_var.reset(reset_store_prompt_ctx_token)
-        paid_user_metadata_var.reset(reset_user_metadata_ctx_token)
+        ContextData.reset_context()