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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: paid-python
-Version: 0.2.1
+Version: 0.3.1
 Summary: 
 Requires-Python: >=3.9,<3.14
 Classifier: Intended Audience :: Developers
@@ -20,11 +20,11 @@ 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)
 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)

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

@@ -3,7 +3,7 @@ name = "paid-python"
 
 [tool.poetry]
 name = "paid-python"
-version = "0.2.1"
+version = "0.3.1"
 description = ""
 readme = "README.md"
 authors = []
@@ -44,11 +44,11 @@ opentelemetry-exporter-otlp-proto-http = ">=1.23.0"
 mutagen = ">=1.47.0"
 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"
+openinference-instrumentation-google-genai = ">=0.1.8"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "==1.13.0"

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

@@ -21,13 +21,6 @@ try:
 except ImportError:
     ANTHROPIC_AVAILABLE = False
 
-try:
-    from opentelemetry.instrumentation.google_generativeai import GoogleGenerativeAiInstrumentor
-
-    GEMINI_AVAILABLE = True
-except ImportError:
-    GEMINI_AVAILABLE = False
-
 try:
     from opentelemetry.instrumentation.openai import OpenAIInstrumentor
 
@@ -56,6 +49,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] = []
@@ -72,11 +72,11 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
     Args:
         libraries: List of library names to instrument. Currently supported:
                   - "anthropic": Anthropic library
-                  - "gemini": Gemini library
                   - "openai": OpenAI library
                   - "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 +106,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", "openai", "openai-agents", "bedrock", "langchain", "google-genai"]
 
     for library in libraries:
         if library in _initialized_instrumentors:
@@ -115,8 +115,6 @@ def paid_autoinstrument(libraries: Optional[List[str]] = None) -> None:
 
         if library == "anthropic":
             _instrument_anthropic()
-        elif library == "gemini":
-            _instrument_gemini()
         elif library == "openai":
             _instrument_openai()
         elif library == "openai-agents":
@@ -125,6 +123,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"
@@ -148,21 +148,6 @@ def _instrument_anthropic() -> None:
     logger.info("Anthropic auto-instrumentation enabled")
 
 
-def _instrument_gemini() -> None:
-    """
-    Instrument the Google Gemini library using opentelemetry-instrumentation-google-generativeai.
-    """
-    if not GEMINI_AVAILABLE:
-        logger.warning("Gemini library not available, skipping instrumentation")
-        return
-
-    # Instrument Gemini with Paid's tracer provider
-    GoogleGenerativeAiInstrumentor().instrument(tracer_provider=tracing.paid_tracer_provider)
-
-    _initialized_instrumentors.append("gemini")
-    logger.info("Gemini auto-instrumentation enabled")
-
-
 def _instrument_openai() -> None:
     """
     Instrument the OpenAI library using opentelemetry-instrumentation-openai.
@@ -217,8 +202,22 @@ def _instrument_langchain() -> None:
         return
 
     # Instrument LangChain with Paid's tracer provider
-    LangchainInstrumentor(disable_trace_context_propagation=True).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.2.1 → paid_python-0.3.1}/src/paid/tracing/context_data.py

@@ -44,6 +44,8 @@ class ContextData:
 
     @classmethod
     def set_context_key(cls, key: str, value: Any) -> None:
+        if value is None:
+            return
         if key not in cls._context:
             logger.warning(f"Invalid context key: {key}")
             return
@@ -51,6 +53,18 @@ class ContextData:
         reset_tokens = cls._get_or_create_reset_tokens()
         reset_tokens[key] = reset_token
 
+    @classmethod
+    def unset_context_key(cls, key: str) -> None:
+        """Unset a specific context key"""
+        if key not in cls._context:
+            logger.warning(f"Invalid context key: {key}")
+            return
+        reset_tokens = cls._reset_tokens.get()
+        if reset_tokens:
+            _ = cls._context[key].set(None)
+            if key in reset_tokens:
+                del reset_tokens[key]
+
     @classmethod
     def reset_context(cls) -> None:
         reset_tokens = cls._reset_tokens.get()
@@ -58,3 +72,14 @@ class ContextData:
             for key, reset_token in reset_tokens.items():
                 cls._context[key].reset(reset_token)
             reset_tokens.clear()
+
+    @classmethod
+    def reset_context_key(cls, key: str) -> None:
+        """Reset a specific context key to its previous value using the stored reset token."""
+        if key not in cls._context:
+            logger.warning(f"Invalid context key: {key}")
+            return
+        reset_tokens = cls._reset_tokens.get()
+        if reset_tokens and key in reset_tokens:
+            cls._context[key].reset(reset_tokens[key])
+            del reset_tokens[key]

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

@@ -110,4 +110,4 @@ def unset_tracing_token():
         DeprecationWarning,
         stacklevel=2,
     )
-    ContextData.set_context_key("trace_id", None)
+    ContextData.unset_context_key("trace_id")

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

@@ -137,12 +137,11 @@ 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
-            ):
+            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}")
@@ -169,8 +168,10 @@ def setup_graceful_termination():
         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")
+        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):
@@ -220,7 +221,7 @@ def initialize_tracing(api_key: Optional[str] = None, collector_endpoint: Option
         span_processor = SimpleSpanProcessor(otlp_exporter)
         paid_tracer_provider.add_span_processor(span_processor)
 
-        setup_graceful_termination() # doesn't throw
+        setup_graceful_termination()  # doesn't throw
 
         logger.info("Paid tracing initialized successfully - collector at %s", collector_endpoint)
     except Exception as e: