ref(google_genai): Use new integrations API

py/noxfile.py

@@ -187,7 +187,7 @@ def test_anthropic(session, version):
 def test_google_genai(session, version):
     _install_test_deps(session)
     _install(session, "google-genai", version)
-    _run_tests(session, f"{WRAPPER_DIR}/test_google_genai.py")
+    _run_tests(session, f"{INTEGRATION_DIR}/google_genai/test_google_genai.py")
     _run_core_tests(session)
 
 

py/src/braintrust/auto.py

@@ -7,7 +7,13 @@
 import logging
 from contextlib import contextmanager
 
-from braintrust.integrations import ADKIntegration, AgnoIntegration, AnthropicIntegration, ClaudeAgentSDKIntegration
+from braintrust.integrations import (
+    ADKIntegration,
+    AgnoIntegration,
+    AnthropicIntegration,
+    ClaudeAgentSDKIntegration,
+    GoogleGenAIIntegration,
+)
 
 
 __all__ = ["auto_instrument"]
@@ -113,7 +119,7 @@ def auto_instrument(
     if pydantic_ai:
         results["pydantic_ai"] = _instrument_pydantic_ai()
     if google_genai:
-        results["google_genai"] = _instrument_google_genai()
+        results["google_genai"] = _instrument_integration(GoogleGenAIIntegration)
     if agno:
         results["agno"] = _instrument_integration(AgnoIntegration)
     if claude_agent_sdk:
@@ -156,14 +162,6 @@ def _instrument_pydantic_ai() -> bool:
     return False
 
 
-def _instrument_google_genai() -> bool:
-    with _try_patch():
-        from braintrust.wrappers.google_genai import setup_genai
-
-        return setup_genai()
-    return False
-
-
 def _instrument_dspy() -> bool:
     with _try_patch():
         from braintrust.wrappers.dspy import patch_dspy

py/src/braintrust/integrations/__init__.py

@@ -2,6 +2,13 @@
 from .agno import AgnoIntegration
 from .anthropic import AnthropicIntegration
 from .claude_agent_sdk import ClaudeAgentSDKIntegration
+from .google_genai import GoogleGenAIIntegration
 
 
-__all__ = ["ADKIntegration", "AgnoIntegration", "AnthropicIntegration", "ClaudeAgentSDKIntegration"]
+__all__ = [
+    "ADKIntegration",
+    "AgnoIntegration",
+    "AnthropicIntegration",
+    "ClaudeAgentSDKIntegration",
+    "GoogleGenAIIntegration",
+]

py/src/braintrust/integrations/google_genai/__init__.py

@@ -0,0 +1,39 @@
+"""Braintrust integration for Google GenAI."""
+
+import logging
+
+from braintrust.logger import NOOP_SPAN, current_span, init_logger
+
+from .integration import GoogleGenAIIntegration
+
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "GoogleGenAIIntegration",
+    "setup_genai",
+]
+
+
+def setup_genai(
+    api_key: str | None = None,
+    project_id: str | None = None,
+    project_name: str | None = None,
+) -> bool:
+    """Setup Braintrust integration with Google GenAI.
+
+    Will automatically patch Google GenAI models for automatic tracing.
+
+    Args:
+        api_key: Braintrust API key.
+        project_id: Braintrust project ID.
+        project_name: Braintrust project name.
+
+    Returns:
+        True if setup was successful, False if google-genai is not installed.
+    """
+    span = current_span()
+    if span == NOOP_SPAN:
+        init_logger(project=project_name, api_key=api_key, project_id=project_id)
+
+    return GoogleGenAIIntegration.setup()

py/src/braintrust/integrations/google_genai/integration.py

@@ -0,0 +1,32 @@
+"""Google GenAI integration — orchestration class and setup entry-point."""
+
+import logging
+
+from braintrust.integrations.base import BaseIntegration
+
+from .patchers import (
+    AsyncModelsEmbedContentPatcher,
+    AsyncModelsGenerateContentPatcher,
+    AsyncModelsGenerateContentStreamPatcher,
+    ModelsEmbedContentPatcher,
+    ModelsGenerateContentPatcher,
+    ModelsGenerateContentStreamPatcher,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class GoogleGenAIIntegration(BaseIntegration):
+    """Braintrust instrumentation for the Google GenAI Python SDK."""
+
+    name = "google_genai"
+    import_names = ("google.genai",)
+    patchers = (
+        ModelsGenerateContentPatcher,
+        ModelsGenerateContentStreamPatcher,
+        ModelsEmbedContentPatcher,
+        AsyncModelsGenerateContentPatcher,
+        AsyncModelsGenerateContentStreamPatcher,
+        AsyncModelsEmbedContentPatcher,
+    )

py/src/braintrust/integrations/google_genai/patchers.py

@@ -0,0 +1,76 @@
+"""Google GenAI patchers — one patcher per coherent patch target."""
+
+from braintrust.integrations.base import FunctionWrapperPatcher
+
+from .tracing import (
+    _async_embed_content_wrapper,
+    _async_generate_content_stream_wrapper,
+    _async_generate_content_wrapper,
+    _embed_content_wrapper,
+    _generate_content_stream_wrapper,
+    _generate_content_wrapper,
+)
+
+
+# ---------------------------------------------------------------------------
+# Sync Models patchers
+# ---------------------------------------------------------------------------
+
+
+class ModelsGenerateContentPatcher(FunctionWrapperPatcher):
+    """Patch ``Models._generate_content`` for tracing."""
+
+    name = "google_genai.models.generate_content"
+    target_module = "google.genai.models"
+    target_path = "Models._generate_content"
+    wrapper = _generate_content_wrapper
+
+
+class ModelsGenerateContentStreamPatcher(FunctionWrapperPatcher):
+    """Patch ``Models.generate_content_stream`` for tracing."""
+
+    name = "google_genai.models.generate_content_stream"
+    target_module = "google.genai.models"
+    target_path = "Models.generate_content_stream"
+    wrapper = _generate_content_stream_wrapper
+
+
+class ModelsEmbedContentPatcher(FunctionWrapperPatcher):
+    """Patch ``Models.embed_content`` for tracing."""
+
+    name = "google_genai.models.embed_content"
+    target_module = "google.genai.models"
+    target_path = "Models.embed_content"
+    wrapper = _embed_content_wrapper
+
+
+# ---------------------------------------------------------------------------
+# Async Models patchers
+# ---------------------------------------------------------------------------
+
+
+class AsyncModelsGenerateContentPatcher(FunctionWrapperPatcher):
+    """Patch ``AsyncModels.generate_content`` for tracing."""
+
+    name = "google_genai.async_models.generate_content"
+    target_module = "google.genai.models"
+    target_path = "AsyncModels.generate_content"
+    wrapper = _async_generate_content_wrapper
+
+
+class AsyncModelsGenerateContentStreamPatcher(FunctionWrapperPatcher):
+    """Patch ``AsyncModels.generate_content_stream`` for tracing."""
+
+    name = "google_genai.async_models.generate_content_stream"
+    target_module = "google.genai.models"
+    target_path = "AsyncModels.generate_content_stream"
+    wrapper = _async_generate_content_stream_wrapper
+
+
+class AsyncModelsEmbedContentPatcher(FunctionWrapperPatcher):
+    """Patch ``AsyncModels.embed_content`` for tracing."""
+
+    name = "google_genai.async_models.embed_content"
+    target_module = "google.genai.models"
+    target_path = "AsyncModels.embed_content"
+    wrapper = _async_embed_content_wrapper

py/src/braintrust/wrappers/test_google_genai.py → py/src/braintrust/integrations/google_genai/test_google_genai.py

@@ -4,8 +4,8 @@
 
 import pytest
 from braintrust import logger
+from braintrust.integrations.google_genai import setup_genai
 from braintrust.test_helpers import init_test_logger
-from braintrust.wrappers.google_genai import setup_genai
 from braintrust.wrappers.test_utils import verify_autoinstrument_script
 from google.genai import types
 from google.genai.client import Client