ref(google_genai): Use new integrations API (#144)

Author: AbhiPrasad

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()