paid-python 0.2.0__tar.gz → 0.3.0__tar.gz

{paid_python-0.2.0 → paid_python-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: paid-python
-Version: 0.2.0
+Version: 0.3.0
 Summary: 
 Requires-Python: >=3.9,<3.14
 Classifier: Intended Audience :: Developers
@@ -20,6 +20,7 @@ Classifier: Typing :: Typed
 Requires-Dist: httpx (>=0.21.0)
 Requires-Dist: mutagen (>=1.47.0)
 Requires-Dist: openinference-instrumentation-bedrock (>=0.1.0)
+Requires-Dist: openinference-instrumentation-google-genai (>=0.1.8)
 Requires-Dist: openinference-instrumentation-openai-agents (>=1.0.0)
 Requires-Dist: opentelemetry-api (>=1.23.0)
 Requires-Dist: opentelemetry-exporter-otlp-proto-http (>=1.23.0)

{paid_python-0.2.0 → paid_python-0.3.0}/pyproject.toml

@@ -3,7 +3,7 @@ name = "paid-python"
 
 [tool.poetry]
 name = "paid-python"
-version = "0.2.0"
+version = "0.3.0"
 description = ""
 readme = "README.md"
 authors = []
@@ -49,6 +49,7 @@ 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"
+openinference-instrumentation-google-genai = ">=0.1.8"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "==1.13.0"

{paid_python-0.2.0 → paid_python-0.3.0}/src/paid/tracing/autoinstrumentation.py

@@ -56,6 +56,13 @@ try:
 except ImportError:
     LANGCHAIN_AVAILABLE = False
 
+try:
+    from openinference.instrumentation.google_genai import GoogleGenAIInstrumentor
+
+    GOOGLE_GENAI_AVAILABLE = True
+except ImportError:
+    GOOGLE_GENAI_AVAILABLE = False
+
 
 # Track which instrumentors have been initialized
 _initialized_instrumentors: List[str] = []
@@ -77,6 +84,7 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
                   - "openai-agents": OpenAI Agents SDK
                   - "bedrock": AWS Bedrock
                   - "langchain": LangChain library
+                  - "google-genai": Google GenAI library
                   If None, all supported libraries that are installed will be instrumented.
 
     Note:
@@ -106,7 +114,7 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
 
     # Default to all supported libraries if none specified
     if libraries is None:
-        libraries = ["anthropic", "gemini", "openai", "openai-agents", "bedrock", "langchain"]
+        libraries = ["anthropic", "gemini", "openai", "openai-agents", "bedrock", "langchain", "google-genai"]
 
     for library in libraries:
         if library in _initialized_instrumentors:
@@ -125,6 +133,8 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
             _instrument_bedrock()
         elif library == "langchain":
             _instrument_langchain()
+        elif library == "google-genai":
+            _instrument_google_genai()
         else:
             logger.warning(
                 f"Unknown library '{library}' - supported libraries: anthropic, gemini, openai, openai-agents, bedrock, langchain"
@@ -217,7 +227,21 @@ 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")
+
+def _instrument_google_genai() -> None:
+    """
+    Instrument Google GenAI using openinference-instrumentation-google-genai.
+    """
+    if not GOOGLE_GENAI_AVAILABLE:
+        logger.warning("Google GenAI instrumentation library not available, skipping instrumentation")
+        return
+
+    
+    GoogleGenAIInstrumentor().instrument(tracer_provider=tracing.paid_tracer_provider)
+    _initialized_instrumentors.append("google-genai")
+    logger.info("Google GenAI auto-instrumentation enabled")

paid_python-0.3.0/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.2.0 → paid_python-0.3.0}/src/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_python-0.2.0 → paid_python-0.3.0}/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.2.0 → paid_python-0.3.0}/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
 
@@ -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()