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

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

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

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

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

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

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

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