paid-python 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

paid/tracing/autoinstrumentation.py

@@ -217,7 +217,8 @@ def _instrument_langchain() -> None:
         return
 
     # Instrument LangChain with Paid's tracer provider
-    LangchainInstrumentor().instrument(tracer_provider=tracing.paid_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/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/tracing/context_manager.py

@@ -1,9 +1,9 @@
 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 .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
@@ -78,14 +78,6 @@ 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)
@@ -94,23 +86,15 @@ class paid_tracing:
         """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):

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/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
 
@@ -299,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(
@@ -331,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_(
@@ -373,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(
@@ -408,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()

{paid_python-0.2.0.dist-info → paid_python-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: paid-python
-Version: 0.2.0
+Version: 0.2.1
 Summary: 
 Requires-Python: >=3.9,<3.14
 Classifier: Intended Audience :: Developers

{paid_python-0.2.0.dist-info → paid_python-0.2.1.dist-info}/RECORD

@@ -37,11 +37,12 @@ paid/orders/lines/raw_client.py,sha256=KZN_yBokCOkf1lUb4ZJtX_NZbqmTqCdJNoaIOdWar
 paid/orders/raw_client.py,sha256=650e1Sj2vi9KVJc15M3ENXIKYoth0qMz66dzvXy1Sb4,16245
 paid/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 paid/tracing/__init__.py,sha256=Pe55koIwqJ6Vv5-9Wqi8xIdwCS2BbxZds-MK5fD-F5Y,506
-paid/tracing/autoinstrumentation.py,sha256=qlLLiP9cWNIXX226dLDLvvlioLCPO-0F81mU_0fBb9s,7604
-paid/tracing/context_manager.py,sha256=Qtl59mjDsacX53LPFFAVXsfeGCEA6CIfxPsvjO-Kmx8,9729
-paid/tracing/distributed_tracing.py,sha256=CpUWpHai-4LxLLHbGxz41r9h5wLG-dC83YL5Vg29OBI,3967
+paid/tracing/autoinstrumentation.py,sha256=NUVvzYwbrI9rDs0sWEgOk0jYYLQPtjo3Cfk_OsZX6-A,7643
+paid/tracing/context_data.py,sha256=UFCZxX6zGa9w5lGLQK5tXdMM69aV52nReUkxwfN9S6A,2378
+paid/tracing/context_manager.py,sha256=ZQtsJ9JPxTwn2t4AW26WpYboaOEZdI2T1Sw0Rwsbf-E,8470
+paid/tracing/distributed_tracing.py,sha256=XCUGFBB2lksrR19Mry10cTZRFSCSztUbUjE_sKCwjX8,3995
 paid/tracing/signal.py,sha256=PfYxF6EFQS8j7RY5_C5NXrCBVu9Hq2E2tyG4fdQScJk,3252
-paid/tracing/tracing.py,sha256=w07I7KQ4O8yJdvLs9RDAQVVofVWeR4xNei3zC5-EE6c,16080
+paid/tracing/tracing.py,sha256=CxumDRwQHJ5S4FnzyTfUi_s557Ka7LDxVYSId0jStSc,14783
 paid/tracing/wrappers/__init__.py,sha256=IIleLB_JUbzLw7FshrU2VHZAKF3dZHMGy1O5zCBwwqM,1588
 paid/tracing/wrappers/anthropic/__init__.py,sha256=_x1fjySAQxuT5cIGO_jU09LiGcZH-WQLqKg8mUFAu2w,115
 paid/tracing/wrappers/anthropic/anthropicWrapper.py,sha256=pGchbOb41CbTxc7H8xXoM-LjR085spqrzXqCVC_rrFk,4913
@@ -98,7 +99,7 @@ paid/usage/__init__.py,sha256=_VhToAyIt_5axN6CLJwtxg3-CO7THa_23pbUzqhXJa4,85
 paid/usage/client.py,sha256=280WJuepoovk3BAVbAx2yN2Q_qBdvx3CcPkLu8lXslc,3030
 paid/usage/raw_client.py,sha256=2acg5C4lxuZodZjepU9QYF0fmBxgG-3ZgXs1zUJG-wM,3709
 paid/version.py,sha256=QIpDFnOrxMxrs86eL0iNH0mSZ1DO078wWHYY9TYAoew,78
-paid_python-0.2.0.dist-info/LICENSE,sha256=Nz4baY1zvv0Qy7lqrQtbaiMhmEeGr2Q7A93aqzpml4c,1071
-paid_python-0.2.0.dist-info/METADATA,sha256=tFvHkGpKqlXPnxZhDhxFm-vETz6RbVihbJZOEdjGPUw,22343
-paid_python-0.2.0.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
-paid_python-0.2.0.dist-info/RECORD,,
+paid_python-0.2.1.dist-info/LICENSE,sha256=Nz4baY1zvv0Qy7lqrQtbaiMhmEeGr2Q7A93aqzpml4c,1071
+paid_python-0.2.1.dist-info/METADATA,sha256=kmG2tSGACEXkk1jBxsOu8DePxUxfu4JKq6aXGe_worE,22343
+paid_python-0.2.1.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
+paid_python-0.2.1.dist-info/RECORD,,