Add GoogleModel support

This PR adds support to Google Models via:
- Gemini API
- VertexAI

The actual backend used depends on the parameters passed to GoogleModel

Author: szykol

pyproject.toml

@@ -36,6 +36,7 @@ compat = ["six>=1.17.0"]
 ai = ["httpx==0.28.1", "langchain>=1.2.15", "mcp>=1.27.0", "pydantic>=2.13.1"]
 anthropic = ["splunk-sdk[ai]>=2.1.1", "langchain-anthropic>=1.4.0"]
 openai = ["splunk-sdk[ai]>=2.1.1", "langchain-openai>=1.1.13"]
+google = ["splunk-sdk[ai]>=2.1.1", "langchain-google-genai>=4.2.2", "google-auth>=2.0.0"]
 
 # Treat the same as NPM's `devDependencies`
 [dependency-groups]
@@ -50,7 +51,7 @@ release = ["build>=1.4.3", "jinja2>=3.1.6", "sphinx>=9.1.0", "twine>=6.2.0"]
 lint = ["basedpyright>=1.39.0", "ruff>=0.15.10"]
 dev = [
   "rich>=14.3.3",
-  "splunk-sdk[openai, anthropic]",
+  "splunk-sdk[openai, anthropic, google]",
   { include-group = "test" },
   { include-group = "lint" },
   { include-group = "release" },

splunklib/ai/README.md

@@ -47,6 +47,7 @@ We support following predefined models:
 
 - `OpenAIModel` - works with OpenAI and any [OpenAI-compatible API](https://platform.openai.com/docs/api-reference).
 - `AnthropicModel` - works with Anthropic and any [Anthropic-compatible API](https://docs.anthropic.com/en/api).
+- `GoogleModel` - works with Google's Gemini models via the [Gemini API](https://ai.google.dev/gemini-api/docs) or [Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/overview).
 
 ### OpenAI
 
@@ -76,6 +77,88 @@ model = AnthropicModel(
 async with Agent(model=model) as agent: ....
 ```
 
+### Google
+
+`GoogleModel` supports two backends: the [Gemini API](https://ai.google.dev/gemini-api/docs) and [Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/overview).
+The backend is selected automatically based on the parameters you provide, or you can
+force it with the `vertexai` flag.
+
+Requires the `google` optional extra:
+
+```sh
+pip install "splunk-sdk[google]"
+# or with uv:
+uv add splunk-sdk[google]
+```
+
+#### Gemini API
+
+Use this when you have a Google AI Studio API key and do not need Vertex AI infrastructure.
+Only `model` and `api_key` are required.
+
+```py
+from splunklib.ai import Agent, GoogleModel
+
+model = GoogleModel(
+    model="gemini-2.0-flash",
+    api_key="YOUR_GOOGLE_API_KEY",
+)
+
+async with Agent(model=model) as agent: ...
+```
+
+#### Vertex AI - API key
+
+Use this to route requests through Vertex AI with an API key. Providing `project` is enough
+for the SDK to switch to the Vertex AI backend automatically.
+
+```py
+from splunklib.ai import Agent, GoogleModel
+
+model = GoogleModel(
+    model="gemini-2.0-flash",
+    api_key="YOUR_VERTEX_API_KEY",
+    project="your-gcp-project-id",
+    # location="us-central1",  # optional, defaults to us-central1
+)
+
+async with Agent(model=model) as agent: ...
+```
+
+#### Vertex AI - service account credentials
+
+Use this when authenticating with a service account key file (or any
+`google.auth.credentials.Credentials`-compatible object). No `api_key` is needed.
+
+```py
+from google.oauth2 import service_account
+from splunklib.ai import Agent, GoogleModel
+
+credentials = service_account.Credentials.from_service_account_file(
+    "path/to/service-account.json",
+    scopes=["https://www.googleapis.com/auth/cloud-platform"],
+)
+
+model = GoogleModel(
+    model="gemini-2.0-flash",
+    project="your-gcp-project-id",
+    credentials=credentials,
+    # location="us-central1",  # optional, defaults to us-central1
+)
+
+async with Agent(model=model) as agent: ...
+```
+
+#### Backend selection rules
+
+| `project` | `credentials` | `vertexai` | Backend used |
+|---|---|---|---|
+| not set | not set | `None` (default) | Gemini API |
+| set | - | `None` (default) | Vertex AI |
+| - | set | `None` (default) | Vertex AI |
+| any | any | `True` | Vertex AI (forced) |
+| any | any | `False` | Gemini API (forced) |
+
 ### Self-hosted models via Ollama
 
 [Ollama](https://ollama.com/) can serve local models with both OpenAI and Anthropic-compatible endpoints, so either model class works.

splunklib/ai/engines/langchain.py

@@ -109,7 +109,7 @@
     subagent_middleware,
     tool_middleware,
 )
-from splunklib.ai.model import AnthropicModel, OpenAIModel, PredefinedModel
+from splunklib.ai.model import AnthropicModel, GoogleModel, OpenAIModel, PredefinedModel
 from splunklib.ai.security import create_structured_prompt
 from splunklib.ai.structured_output import (
     StructuredOutputGenerationException,
@@ -1694,6 +1694,33 @@ def _create_langchain_model(model: PredefinedModel) -> BaseChatModel:
                     + "# or if using uv:\n"
                     + "uv add splunk-sdk[anthropic]"
                 )
+        case GoogleModel():
+            try:
+                from langchain_google_genai import ChatGoogleGenerativeAI
+
+                google_kwargs: dict[str, Any] = {"model": model.model}
+                if model.api_key is not None:
+                    google_kwargs["google_api_key"] = model.api_key
+                if model.project is not None:
+                    google_kwargs["project"] = model.project
+                if model.location is not None:
+                    google_kwargs["location"] = model.location
+                if model.credentials is not None:
+                    google_kwargs["credentials"] = model.credentials
+                if model.vertexai is not None:
+                    google_kwargs["vertexai"] = model.vertexai
+                if model.temperature is not None:
+                    google_kwargs["temperature"] = model.temperature
+
+                return ChatGoogleGenerativeAI(**google_kwargs)
+            except ImportError:
+                raise ImportError(
+                    "Google GenAI support is not installed.\n"
+                    + "To enable Google / Gemini models, install the optional extra:\n"
+                    + 'pip install "splunk-sdk[google]"\n'
+                    + "# or if using uv:\n"
+                    + "uv add splunk-sdk[google]"
+                )
         case _:
             raise InvalidModelError(
                 "Cannot create langchain model - invalid SDK model provided"

splunklib/ai/model.py

@@ -13,10 +13,13 @@
 # under the License.
 
 from dataclasses import dataclass
-from typing import Any, Mapping
+from typing import TYPE_CHECKING, Any, Mapping
 
 import httpx
 
+if TYPE_CHECKING:
+    from google.oauth2 import service_account
+
 
 @dataclass(frozen=True)
 class PredefinedModel:
@@ -63,8 +66,40 @@ class AnthropicModel(PredefinedModel):
     temperature: float | None = None
 
 
+@dataclass(frozen=True)
+class GoogleModel(PredefinedModel):
+    """Predefined Google Model
+
+    Supports the Gemini API and Vertex AI. The backend is chosen
+    automatically: Vertex AI when ``project`` or ``credentials`` is set,
+    otherwise the Gemini API. Override with ``vertexai=True/False``.
+
+    See the README for full usage examples and authentication options.
+    """
+
+    model: str
+    api_key: str | None = None
+    """API key for the Gemini API or Vertex AI."""
+
+    project: str | None = None
+    """Google Cloud project ID (Vertex AI only)."""
+
+    location: str | None = None
+    """Vertex AI region, e.g. ``"us-central1"`` or ``"europe-west4"``."""
+
+    credentials: "service_account.Credentials | None" = None
+    """Service account credentials for Vertex AI. When set, ``api_key`` is not required."""
+
+    vertexai: bool | None = None
+    """Force backend selection: ``True`` for Vertex AI, ``False`` for Gemini API, ``None`` to auto-detect."""
+
+    temperature: float | None = None
+    """Sampling temperature in the range ``[0.0, 2.0]``."""
+
+
 __all__ = [
     "AnthropicModel",
+    "GoogleModel",
     "OpenAIModel",
     "PredefinedModel",
 ]

tests/unit/ai/engine/test_langchain_backend.py

@@ -39,7 +39,7 @@
     ToolMessage,
     ToolResult,
 )
-from splunklib.ai.model import AnthropicModel, OpenAIModel, PredefinedModel
+from splunklib.ai.model import AnthropicModel, GoogleModel, OpenAIModel, PredefinedModel
 from splunklib.ai.tools import ToolType
 
 
@@ -387,6 +387,56 @@ def test_create_langchain_model_anthropic_with_base_url(self) -> None:
         # ChatAnthropic stores base_url in anthropic_api_url
         assert result.anthropic_api_url == model.base_url
 
+    def test_create_langchain_model_google_gemini_api(self) -> None:
+        pytest.importorskip("langchain_google_genai")
+        import langchain_google_genai
+
+        model = GoogleModel(model="gemini-2.0-flash", api_key="test-key")
+        result = lc._create_langchain_model(model)
+
+        assert isinstance(result, langchain_google_genai.ChatGoogleGenerativeAI)
+        assert result.model == model.model
+        assert result._use_vertexai is False  # pyright: ignore[reportAttributeAccessIssue]
+
+    def test_create_langchain_model_google_vertex_ai_via_project(self) -> None:
+        pytest.importorskip("langchain_google_genai")
+        import langchain_google_genai
+
+        model = GoogleModel(
+            model="gemini-2.0-flash",
+            api_key="test-key",
+            project="my-project",
+        )
+        result = lc._create_langchain_model(model)
+
+        assert isinstance(result, langchain_google_genai.ChatGoogleGenerativeAI)
+        assert result.project == model.project
+        assert result._use_vertexai is True  # pyright: ignore[reportAttributeAccessIssue]
+
+    def test_create_langchain_model_google_vertex_ai_explicit_flag(self) -> None:
+        pytest.importorskip("langchain_google_genai")
+        import langchain_google_genai
+
+        model = GoogleModel(
+            model="gemini-2.0-flash",
+            api_key="test-key",
+            vertexai=True,
+        )
+        result = lc._create_langchain_model(model)
+
+        assert isinstance(result, langchain_google_genai.ChatGoogleGenerativeAI)
+        assert result._use_vertexai is True  # pyright: ignore[reportAttributeAccessIssue]
+
+    def test_create_langchain_model_google_temperature(self) -> None:
+        pytest.importorskip("langchain_google_genai")
+        import langchain_google_genai
+
+        model = GoogleModel(model="gemini-2.0-flash", api_key="test-key", temperature=0.5)
+        result = lc._create_langchain_model(model)
+
+        assert isinstance(result, langchain_google_genai.ChatGoogleGenerativeAI)
+        assert result.temperature == model.temperature
+
 
 @pytest.mark.parametrize(
     ("name", "tool_type", "expected_name"),