feat(llmo): support Vertex AI

Author: Radu-Raicea

posthog/ai/gemini/gemini.py

@@ -42,6 +42,12 @@ class Client:
     def __init__(
         self,
         api_key: Optional[str] = None,
+        vertexai: Optional[bool] = None,
+        credentials: Optional[Any] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        debug_config: Optional[Any] = None,
+        http_options: Optional[Any] = None,
         posthog_client: Optional[PostHogClient] = None,
         posthog_distinct_id: Optional[str] = None,
         posthog_properties: Optional[Dict[str, Any]] = None,
@@ -51,7 +57,13 @@ def __init__(
     ):
         """
         Args:
-            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
+            vertexai: Whether to use Vertex AI authentication
+            credentials: Vertex AI credentials object
+            project: GCP project ID for Vertex AI
+            location: GCP location for Vertex AI
+            debug_config: Debug configuration for the client
+            http_options: HTTP options for the client
             posthog_client: PostHog client for tracking usage
             posthog_distinct_id: Default distinct ID for all calls (can be overridden per call)
             posthog_properties: Default properties for all calls (can be overridden per call)
@@ -66,6 +78,12 @@ def __init__(
 
         self.models = Models(
             api_key=api_key,
+            vertexai=vertexai,
+            credentials=credentials,
+            project=project,
+            location=location,
+            debug_config=debug_config,
+            http_options=http_options,
             posthog_client=self._ph_client,
             posthog_distinct_id=posthog_distinct_id,
             posthog_properties=posthog_properties,
@@ -85,6 +103,12 @@ class Models:
     def __init__(
         self,
         api_key: Optional[str] = None,
+        vertexai: Optional[bool] = None,
+        credentials: Optional[Any] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        debug_config: Optional[Any] = None,
+        http_options: Optional[Any] = None,
         posthog_client: Optional[PostHogClient] = None,
         posthog_distinct_id: Optional[str] = None,
         posthog_properties: Optional[Dict[str, Any]] = None,
@@ -94,7 +118,13 @@ def __init__(
     ):
         """
         Args:
-            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
+            vertexai: Whether to use Vertex AI authentication
+            credentials: Vertex AI credentials object
+            project: GCP project ID for Vertex AI
+            location: GCP location for Vertex AI
+            debug_config: Debug configuration for the client
+            http_options: HTTP options for the client
             posthog_client: PostHog client for tracking usage
             posthog_distinct_id: Default distinct ID for all calls
             posthog_properties: Default properties for all calls
@@ -113,16 +143,40 @@ def __init__(
         self._default_privacy_mode = posthog_privacy_mode
         self._default_groups = posthog_groups
 
-        # Handle API key - try parameter first, then environment variables
-        if api_key is None:
-            api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY")
+        # Build genai.Client arguments
+        client_args = {}
+        
+        # Add Vertex AI parameters if provided
+        if vertexai is not None:
+            client_args["vertexai"] = vertexai
+        if credentials is not None:
+            client_args["credentials"] = credentials
+        if project is not None:
+            client_args["project"] = project
+        if location is not None:
+            client_args["location"] = location
+        if debug_config is not None:
+            client_args["debug_config"] = debug_config
+        if http_options is not None:
+            client_args["http_options"] = http_options
+
+        # Handle API key authentication
+        if vertexai:
+            # For Vertex AI, api_key is optional
+            if api_key is not None:
+                client_args["api_key"] = api_key
+        else:
+            # For non-Vertex AI mode, api_key is required (backwards compatibility)
+            if api_key is None:
+                api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY")
 
-        if api_key is None:
-            raise ValueError(
-                "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable"
-            )
+            if api_key is None:
+                raise ValueError(
+                    "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable"
+                )
+            client_args["api_key"] = api_key
 
-        self._client = genai.Client(api_key=api_key)
+        self._client = genai.Client(**client_args)
         self._base_url = "https://generativelanguage.googleapis.com"
 
     def _merge_posthog_params(

posthog/test/ai/gemini/test_gemini.py

@@ -318,3 +318,71 @@ def test_new_client_override_defaults(
     assert props["team"] == "ai"  # from defaults
     assert props["feature"] == "chat"  # from call
     assert props["urgent"] is True  # from call
+
+
+def test_vertex_ai_parameters_passed_through(mock_client, mock_google_genai_client, mock_gemini_response):
+    """Test that Vertex AI parameters are properly passed to genai.Client"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+    
+    # Mock credentials object
+    mock_credentials = MagicMock()
+    mock_debug_config = MagicMock()
+    mock_http_options = MagicMock()
+
+    # Create client with Vertex AI parameters
+    Client(
+        vertexai=True,
+        credentials=mock_credentials,
+        project="test-project",
+        location="us-central1",
+        debug_config=mock_debug_config,
+        http_options=mock_http_options,
+        posthog_client=mock_client,
+    )
+
+    # Verify genai.Client was called with correct parameters
+    google_genai.Client.assert_called_once_with(
+        vertexai=True,
+        credentials=mock_credentials,
+        project="test-project",
+        location="us-central1",
+        debug_config=mock_debug_config,
+        http_options=mock_http_options
+    )
+
+
+def test_api_key_mode(mock_client, mock_google_genai_client):
+    """Test API key authentication mode"""
+    
+    # Create client with just API key (traditional mode)
+    Client(
+        api_key="test-api-key",
+        posthog_client=mock_client,
+    )
+
+    # Verify genai.Client was called with only api_key
+    google_genai.Client.assert_called_once_with(api_key="test-api-key")
+
+
+def test_vertex_ai_mode_with_optional_api_key(mock_client, mock_google_genai_client, mock_gemini_response):
+    """Test Vertex AI mode with optional API key"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+    
+    mock_credentials = MagicMock()
+
+    # Create client with Vertex AI + API key
+    Client(
+        vertexai=True,
+        api_key="test-api-key",
+        credentials=mock_credentials,
+        project="test-project",
+        posthog_client=mock_client,
+    )
+
+    # Verify genai.Client was called with both Vertex AI params and API key
+    google_genai.Client.assert_called_once_with(
+        vertexai=True,
+        api_key="test-api-key",
+        credentials=mock_credentials,
+        project="test-project"
+    )