fenra 0.1.0__py3-none-any.whl

fenra/integrations/gemini/__init__.py

@@ -0,0 +1,529 @@
+"""Gemini integration with auto-instrumentation via monkey patching."""
+
+import functools
+import logging
+from typing import Any
+
+from fenra._context import get_context
+from fenra._core import enqueue_transaction
+
+logger = logging.getLogger(__name__)
+
+_patched_sync = False
+_patched_async = False
+_patched_stream_sync = False
+_patched_stream_async = False
+
+# Store original methods for unpatching
+_original_generate_content: Any = None
+_original_generate_content_async: Any = None
+_original_generate_content_stream: Any = None
+_original_generate_content_stream_async: Any = None
+
+
+def _track_generate_content_response(
+    response: Any,
+    model: str,
+    context: dict[str, Any],
+    config: Any = None,
+) -> None:
+    """
+    Extract usage from GenerateContentResponse and queue transaction.
+
+    Handles:
+    - Token usage (including thinking tokens for reasoning models)
+    - Google Search tool usage (web_search requests)
+    - Image generation output
+    """
+    if not hasattr(response, "usage_metadata") or response.usage_metadata is None:
+        return
+
+    usage = response.usage_metadata
+
+    # Build token metrics
+    # thoughts_token_count is used by thinking models like gemini-2.0-flash-thinking
+    output_tokens = (getattr(usage, "candidates_token_count", 0) or 0) + (
+        getattr(usage, "thoughts_token_count", None) or 0
+    )
+
+    metrics: dict[str, Any] = {
+        "input_tokens": getattr(usage, "prompt_token_count", 0) or 0,
+        "output_tokens": output_tokens,
+        "total_tokens": getattr(usage, "total_token_count", 0) or 0,
+    }
+
+    # Include cached tokens if present
+    cached = getattr(usage, "cached_content_token_count", None)
+    if cached:
+        metrics["cached_tokens"] = cached
+
+    usage_entries: list[dict[str, Any]] = [
+        {
+            "type": "tokens",
+            "metrics": metrics,
+        }
+    ]
+
+    # Detect Google Search tool usage
+    search_usage = _detect_tool_usage(response, config)
+    if search_usage:
+        usage_entries.append(search_usage)
+
+    # Detect image generation output
+    image_usage = _detect_image_output(response, config)
+    if image_usage:
+        usage_entries.append(image_usage)
+
+    transaction: dict[str, Any] = {
+        "provider": "gemini",
+        "model": model,
+        "usage": usage_entries,
+        "context": context,
+    }
+
+    # Include raw usage for debugging if available
+    try:
+        transaction["provider_usage_raw"] = {
+            "prompt_token_count": getattr(usage, "prompt_token_count", None),
+            "candidates_token_count": getattr(usage, "candidates_token_count", None),
+            "total_token_count": getattr(usage, "total_token_count", None),
+            "thoughts_token_count": getattr(usage, "thoughts_token_count", None),
+            "cached_content_token_count": getattr(
+                usage, "cached_content_token_count", None
+            ),
+        }
+    except Exception:
+        pass
+
+    enqueue_transaction(transaction)
+
+
+def _detect_tool_usage(response: Any, config: Any) -> dict[str, Any] | None:
+    """
+    Detect Google Search tool usage from grounding metadata.
+
+    Returns a requests usage entry if web search was used, None otherwise.
+    """
+    if config is None:
+        return None
+
+    # Check if GoogleSearch tool was configured
+    tools = getattr(config, "tools", None)
+    if not tools:
+        return None
+
+    has_google_search = False
+    for tool in tools:
+        if hasattr(tool, "google_search") and tool.google_search is not None:
+            has_google_search = True
+            break
+
+    if not has_google_search:
+        return None
+
+    # Extract search count from grounding metadata
+    search_count = 0
+    candidates = getattr(response, "candidates", None) or []
+    for candidate in candidates:
+        grounding_metadata = getattr(candidate, "grounding_metadata", None)
+        if grounding_metadata:
+            web_search_queries = getattr(grounding_metadata, "web_search_queries", None)
+            if web_search_queries:
+                search_count = len(web_search_queries)
+                break
+
+    # Default to 1 if tool was configured but we couldn't detect count
+    if search_count == 0:
+        search_count = 1
+
+    return {
+        "type": "requests",
+        "metrics": {"count": search_count, "request_type": "web_search"},
+    }
+
+
+def _detect_image_output(response: Any, config: Any) -> dict[str, Any] | None:
+    """
+    Detect image generation from response parts.
+
+    Returns an images usage entry if images were generated, None otherwise.
+    """
+    if config is None:
+        return None
+
+    # Check if IMAGE was in response modalities
+    response_modalities = getattr(config, "response_modalities", None)
+    if not response_modalities or "IMAGE" not in response_modalities:
+        return None
+
+    # Count images from response parts
+    image_count = 0
+    candidates = getattr(response, "candidates", None) or []
+    for candidate in candidates:
+        content = getattr(candidate, "content", None)
+        if content:
+            parts = getattr(content, "parts", None) or []
+            for part in parts:
+                if hasattr(part, "inline_data") and part.inline_data:
+                    image_count += 1
+
+    if image_count == 0:
+        return None
+
+    return {
+        "type": "images",
+        "metrics": {
+            "generated": image_count,
+            "size_px": 1024,  # Default Gemini image size
+            "quality": "standard",
+        },
+    }
+
+
+def _track_stream(
+    stream: Any,
+    model: str,
+    context: dict[str, Any],
+    config: Any = None,
+) -> Any:
+    """
+    Wrap a streaming response to track usage from the final chunk.
+
+    Returns a wrapped stream that tracks usage when the stream is exhausted.
+    """
+    last_chunk = None
+
+    def wrapped_stream() -> Any:
+        nonlocal last_chunk
+        try:
+            for chunk in stream:
+                last_chunk = chunk
+                yield chunk
+        finally:
+            # After stream is exhausted, track usage from final chunk
+            if last_chunk:
+                try:
+                    _track_generate_content_response(last_chunk, model, context, config)
+                except Exception as e:
+                    logger.error(
+                        f"Error tracking Gemini streaming usage: {e}", exc_info=True
+                    )
+
+    return wrapped_stream()
+
+
+async def _track_stream_async(
+    stream: Any,
+    model: str,
+    context: dict[str, Any],
+    config: Any = None,
+) -> Any:
+    """
+    Wrap an async streaming response to track usage from the final chunk.
+
+    Returns a wrapped async stream that tracks usage when the stream is exhausted.
+    """
+    last_chunk = None
+
+    try:
+        async for chunk in stream:
+            last_chunk = chunk
+            yield chunk
+    finally:
+        # After stream is exhausted, track usage from final chunk
+        if last_chunk:
+            try:
+                _track_generate_content_response(last_chunk, model, context, config)
+            except Exception as e:
+                logger.error(
+                    f"Error tracking Gemini streaming usage: {e}", exc_info=True
+                )
+
+
+# ============================================================================
+# Sync Patching
+# ============================================================================
+
+
+def patch_gemini() -> None:
+    """
+    Patch Gemini client to auto-track usage.
+
+    This patches the synchronous `Models.generate_content` method.
+    """
+    global _patched_sync, _original_generate_content
+
+    if _patched_sync:
+        logger.debug("Gemini sync already patched, skipping")
+        return
+
+    try:
+        from google.genai.models import Models
+    except ImportError:
+        logger.warning(
+            "Gemini SDK not installed. Install with: pip install 'fenra[gemini]'"
+        )
+        return
+
+    original_generate_content = Models.generate_content
+    _original_generate_content = original_generate_content
+
+    @functools.wraps(original_generate_content)
+    def patched_generate_content(self: Models, *args: Any, **kwargs: Any) -> Any:
+        response = original_generate_content(self, *args, **kwargs)
+
+        context = get_context()
+
+        model = kwargs.get("model") or (args[0] if args else "unknown")
+
+        config = kwargs.get("config")
+
+        try:
+            _track_generate_content_response(response, model, context, config)
+        except Exception as e:
+            logger.error(f"Error tracking Gemini usage: {e}", exc_info=True)
+
+        return response
+
+    Models.generate_content = patched_generate_content  # type: ignore[assignment]
+    _patched_sync = True
+    logger.info("Gemini SDK patched for auto-instrumentation")
+
+
+def patch_gemini_stream() -> None:
+    """
+    Patch Gemini client to auto-track streaming usage.
+
+    This patches the synchronous `Models.generate_content_stream` method.
+    """
+    global _patched_stream_sync, _original_generate_content_stream
+
+    if _patched_stream_sync:
+        logger.debug("Gemini sync stream already patched, skipping")
+        return
+
+    try:
+        from google.genai.models import Models
+    except ImportError:
+        logger.warning(
+            "Gemini SDK not installed. Install with: pip install 'fenra[gemini]'"
+        )
+        return
+
+    if not hasattr(Models, "generate_content_stream"):
+        logger.debug("Gemini Models has no generate_content_stream method, skipping")
+        return
+
+    original_generate_content_stream = Models.generate_content_stream
+    _original_generate_content_stream = original_generate_content_stream
+
+    @functools.wraps(original_generate_content_stream)
+    def patched_generate_content_stream(
+        self: Models, *args: Any, **kwargs: Any
+    ) -> Any:
+        stream = original_generate_content_stream(self, *args, **kwargs)
+
+        context = get_context()
+
+        model = kwargs.get("model") or (args[0] if args else "unknown")
+
+        config = kwargs.get("config")
+
+        try:
+            return _track_stream(stream, model, context, config)
+        except Exception as e:
+            logger.error(f"Error wrapping Gemini stream: {e}", exc_info=True)
+            return stream
+
+    Models.generate_content_stream = patched_generate_content_stream  # type: ignore[assignment]
+    _patched_stream_sync = True
+    logger.info("Gemini streaming SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# Async Patching
+# ============================================================================
+
+
+def patch_gemini_async() -> None:
+    """
+    Patch Gemini async client to auto-track usage.
+
+    This patches the asynchronous `AsyncModels.generate_content` method.
+    """
+    global _patched_async, _original_generate_content_async
+
+    if _patched_async:
+        logger.debug("Gemini async already patched, skipping")
+        return
+
+    try:
+        from google.genai.models import AsyncModels
+    except ImportError:
+        logger.warning(
+            "Gemini SDK not installed. Install with: pip install 'fenra[gemini]'"
+        )
+        return
+
+    original_generate_content = AsyncModels.generate_content
+    _original_generate_content_async = original_generate_content
+
+    @functools.wraps(original_generate_content)
+    async def patched_generate_content(
+        self: AsyncModels, *args: Any, **kwargs: Any
+    ) -> Any:
+        response = await original_generate_content(self, *args, **kwargs)
+
+        context = get_context()
+
+        model = kwargs.get("model") or (args[0] if args else "unknown")
+
+        config = kwargs.get("config")
+
+        try:
+            _track_generate_content_response(response, model, context, config)
+        except Exception as e:
+            logger.error(f"Error tracking Gemini async usage: {e}", exc_info=True)
+
+        return response
+
+    AsyncModels.generate_content = patched_generate_content  # type: ignore[assignment]
+    _patched_async = True
+    logger.info("Gemini async SDK patched for auto-instrumentation")
+
+
+def patch_gemini_stream_async() -> None:
+    """
+    Patch Gemini async client to auto-track streaming usage.
+
+    This patches the asynchronous `AsyncModels.generate_content_stream` method.
+    """
+    global _patched_stream_async, _original_generate_content_stream_async
+
+    if _patched_stream_async:
+        logger.debug("Gemini async stream already patched, skipping")
+        return
+
+    try:
+        from google.genai.models import AsyncModels
+    except ImportError:
+        logger.warning(
+            "Gemini SDK not installed. Install with: pip install 'fenra[gemini]'"
+        )
+        return
+
+    if not hasattr(AsyncModels, "generate_content_stream"):
+        logger.debug(
+            "Gemini AsyncModels has no generate_content_stream method, skipping"
+        )
+        return
+
+    original_generate_content_stream = AsyncModels.generate_content_stream
+    _original_generate_content_stream_async = original_generate_content_stream
+
+    @functools.wraps(original_generate_content_stream)
+    async def patched_generate_content_stream(
+        self: AsyncModels, *args: Any, **kwargs: Any
+    ) -> Any:
+        stream = await original_generate_content_stream(self, *args, **kwargs)
+
+        context = get_context()
+
+        model = kwargs.get("model") or (args[0] if args else "unknown")
+
+        config = kwargs.get("config")
+
+        try:
+            return _track_stream_async(stream, model, context, config)
+        except Exception as e:
+            logger.error(f"Error wrapping Gemini async stream: {e}", exc_info=True)
+            return stream
+
+    AsyncModels.generate_content_stream = patched_generate_content_stream  # type: ignore[assignment]
+    _patched_stream_async = True
+    logger.info("Gemini async streaming SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# Unpatch Functions
+# ============================================================================
+
+
+def unpatch_gemini() -> None:
+    """Restore original Gemini Models.generate_content method."""
+    global _patched_sync, _original_generate_content
+
+    if not _patched_sync or _original_generate_content is None:
+        return
+
+    try:
+        from google.genai.models import Models
+
+        Models.generate_content = _original_generate_content  # type: ignore[assignment]
+        _patched_sync = False
+        _original_generate_content = None
+        logger.info("Gemini SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_gemini_stream() -> None:
+    """Restore original Gemini Models.generate_content_stream method."""
+    global _patched_stream_sync, _original_generate_content_stream
+
+    if not _patched_stream_sync or _original_generate_content_stream is None:
+        return
+
+    try:
+        from google.genai.models import Models
+
+        Models.generate_content_stream = _original_generate_content_stream  # type: ignore[assignment]
+        _patched_stream_sync = False
+        _original_generate_content_stream = None
+        logger.info("Gemini streaming SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_gemini_async() -> None:
+    """Restore original Gemini AsyncModels.generate_content method."""
+    global _patched_async, _original_generate_content_async
+
+    if not _patched_async or _original_generate_content_async is None:
+        return
+
+    try:
+        from google.genai.models import AsyncModels
+
+        AsyncModels.generate_content = _original_generate_content_async  # type: ignore[assignment]
+        _patched_async = False
+        _original_generate_content_async = None
+        logger.info("Gemini async SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_gemini_stream_async() -> None:
+    """Restore original Gemini AsyncModels.generate_content_stream method."""
+    global _patched_stream_async, _original_generate_content_stream_async
+
+    if not _patched_stream_async or _original_generate_content_stream_async is None:
+        return
+
+    try:
+        from google.genai.models import AsyncModels
+
+        AsyncModels.generate_content_stream = _original_generate_content_stream_async  # type: ignore[assignment]
+        _patched_stream_async = False
+        _original_generate_content_stream_async = None
+        logger.info("Gemini async streaming SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_gemini_all() -> None:
+    """Restore all original Gemini SDK methods."""
+    unpatch_gemini()
+    unpatch_gemini_stream()
+    unpatch_gemini_async()
+    unpatch_gemini_stream_async()