feat(llm): extract BaseGeminiLLM, add base_url support

CHANGELOG.md

@@ -8,8 +8,11 @@
 - Added `neo4j_graphrag.llm.utils.split_http_client_kwargs`, a shared helper that routes a constructor's `http_client` kwarg to whichever of the sync/async SDK clients it matches. `AnthropicLLM`, `OpenAILLM`, and `AzureOpenAILLM` now all use this single implementation instead of three separately maintained copies of the same logic. Custom subclasses that construct their own SDK clients can call it to get the same behavior; it is exported from `neo4j_graphrag.llm` for that purpose.
 - Added `BaseAnthropicLLM`, a new base class holding all of `AnthropicLLM`'s shared message-building, schema-conversion, and response-parsing logic, mirroring `BaseOpenAILLM`. Both `BaseAnthropicLLM` and `BaseOpenAILLM` are now exported from `neo4j_graphrag.llm` as documented, supported extension points for subclassing to reach custom Anthropic/OpenAI-compatible endpoints.
 - Added an explicit `base_url` keyword parameter to `AnthropicLLM` and `OpenAILLM`, passed through to both the sync and async SDK clients of each.
+- Added an explicit `base_url` keyword parameter to `GeminiLLM`. The `google-genai` SDK has no top-level `base_url` argument, so the value is applied through `http_options`; if `http_options` is also provided, only its `base_url` field is overridden.
+- `BaseGeminiLLM` now implements `aclose()`, closing the underlying `genai.Client` (sync and `aio` surfaces), so `with GeminiLLM(...)` / `async with` actually release resources like the other providers.
 - `split_http_client_kwargs` now emits a `UserWarning` when the provided `http_client` has a `base_url` configured, since the SDKs ignore it — the LLM constructor's `base_url` parameter is the supported way to change the endpoint.
 - Added a new docs page, `docs/source/llm.rst`, describing the `http_client`/`base_url` injection contract shared by `AnthropicLLM` and `OpenAILLM`, with a worked example of subclassing `BaseAnthropicLLM` to reach a custom endpoint.
+- Added `BaseGeminiLLM`, a new base class holding all of `GeminiLLM`'s shared message-building, config/schema-building, and response-parsing logic. `GeminiLLM` is now a thin subclass responsible only for constructing the `genai.Client`. `BaseGeminiLLM` is exported from `neo4j_graphrag.llm` as a documented, supported extension point for subclassing to reach a custom Gemini-compatible endpoint.
 
 ### Changed
 

docs/source/api.rst

@@ -357,6 +357,16 @@ OpenAILLM
     :undoc-members: get_messages, client_class, async_client_class
 
 
+BaseGeminiLLM
+-------------
+
+See :ref:`llm-extensibility` for the ``http_client``/``base_url`` extension
+contract.
+
+.. autoclass:: neo4j_graphrag.llm.google_genai_llm.BaseGeminiLLM
+    :members:
+
+
 GeminiLLM
 ---------
 

docs/source/llm.rst

@@ -9,20 +9,22 @@ providers and extending them to reach custom endpoints.
 
 .. _llm-extensibility:
 
-Extending LLMs: BaseAnthropicLLM/BaseOpenAILLM
-==============================================
-
-:class:`~neo4j_graphrag.llm.anthropic_llm.BaseAnthropicLLM` and
-:class:`~neo4j_graphrag.llm.openai_llm.BaseOpenAILLM` are the shared base classes behind
-:class:`~neo4j_graphrag.llm.anthropic_llm.AnthropicLLM` and
-:class:`~neo4j_graphrag.llm.openai_llm.OpenAILLM` respectively. They hold all
-the provider-agnostic logic (message building, schema conversion, response
-parsing, structured output handling) and leave SDK client construction to
-their subclasses. This page documents the ``http_client``/``base_url``
-injection contract they expose, so you can point either provider at a custom
-or self-hosted, API-compatible endpoint.
-
-Both ``AnthropicLLM`` and ``OpenAILLM`` accept two related, independent
+Extending LLMs: the Base* classes
+=================================
+
+:class:`~neo4j_graphrag.llm.anthropic_llm.BaseAnthropicLLM`,
+:class:`~neo4j_graphrag.llm.openai_llm.BaseOpenAILLM`, and
+:class:`~neo4j_graphrag.llm.google_genai_llm.BaseGeminiLLM` are the shared
+base classes behind :class:`~neo4j_graphrag.llm.anthropic_llm.AnthropicLLM`,
+:class:`~neo4j_graphrag.llm.openai_llm.OpenAILLM`, and
+:class:`~neo4j_graphrag.llm.google_genai_llm.GeminiLLM` respectively. They
+hold all the provider-agnostic logic (message building, schema conversion,
+response parsing, structured output handling) and leave SDK client
+construction to their subclasses. This page documents the
+``http_client``/``base_url`` injection contract they expose, so you can point
+each provider at a custom or self-hosted, API-compatible endpoint.
+
+``AnthropicLLM`` and ``OpenAILLM`` accept two related, independent
 constructor settings:
 
 - ``base_url`` (``Optional[str]``): an explicit constructor parameter on both
@@ -124,3 +126,33 @@ needs to decide how ``client``/``async_client`` get constructed.
 
 The same pattern applies to :class:`~neo4j_graphrag.llm.openai_llm.BaseOpenAILLM`
 for OpenAI-compatible endpoints.
+
+Gemini: single client, ``http_options`` instead of ``http_client``
+------------------------------------------------------------------
+
+:class:`~neo4j_graphrag.llm.google_genai_llm.GeminiLLM` follows the same
+contract with two SDK-specific differences:
+
+- The ``google.genai`` SDK uses a **single** ``genai.Client`` for both
+  directions (async calls go through ``client.aio``), so there is no
+  ``http_client`` sync/async routing and ``split_http_client_kwargs`` does
+  not apply.
+- ``genai.Client`` has no top-level ``base_url`` argument: the endpoint lives
+  in ``http_options``. ``GeminiLLM`` still accepts the same explicit
+  ``base_url`` constructor parameter and applies it through ``http_options``
+  for you — if you also pass ``http_options`` (as a dict or
+  ``types.HttpOptions``), only its ``base_url`` field is overridden.
+
+.. code-block:: python
+
+    from neo4j_graphrag.llm import GeminiLLM
+
+    llm = GeminiLLM(
+        model_name="gemini-2.0-flash",
+        base_url="https://my-gemini-endpoint.example.com",
+    )
+
+A custom subclass of
+:class:`~neo4j_graphrag.llm.google_genai_llm.BaseGeminiLLM` only needs to
+assign ``self.client`` (a ``genai.Client``); all message building, config
+handling, and response parsing are inherited.

examples/customize/llms/google_genai_llm.py

@@ -11,3 +11,14 @@
 )
 res = llm.invoke("say something")
 print(res.content)
+
+# To reach a custom or self-hosted, Gemini-compatible endpoint instead of
+# Google's default API, pass `base_url`. The genai SDK has no top-level
+# base_url argument, so GeminiLLM applies it through `http_options` for you.
+custom_llm = GeminiLLM(
+    model_name="gemini-2.5-flash",
+    api_key=api_key,
+    base_url="https://my-custom-endpoint.example.com",
+)
+res = custom_llm.invoke("say something")
+print(res.content)

src/neo4j_graphrag/llm/__init__.py

@@ -19,7 +19,7 @@
 from .base import LLMBase, LLMInterface, LLMInterfaceV2
 from .bedrock_llm import BedrockLLM
 from .cohere_llm import CohereLLM
-from .google_genai_llm import GeminiLLM
+from .google_genai_llm import BaseGeminiLLM, GeminiLLM
 from .mistralai_llm import MistralAILLM
 from .ollama_llm import OllamaLLM
 from .openai_llm import AzureOpenAILLM, BaseOpenAILLM, OpenAILLM
@@ -30,6 +30,7 @@
 __all__ = [
     "AnthropicLLM",
     "BaseAnthropicLLM",
+    "BaseGeminiLLM",
     "BedrockLLM",
     "CohereLLM",
     "GeminiLLM",

src/neo4j_graphrag/llm/google_genai_llm.py

@@ -16,6 +16,7 @@
 # built-in dependencies
 from __future__ import annotations
 
+import abc
 from typing import (
     Any,
     List,
@@ -32,7 +33,7 @@
 
 # project dependencies
 from neo4j_graphrag.exceptions import LLMGenerationError
-from neo4j_graphrag.llm.base import LLMInterface, LLMInterfaceV2
+from neo4j_graphrag.llm.base import LLMBase
 from neo4j_graphrag.llm.types import (
     BaseMessage,
     LLMResponse,
@@ -62,16 +63,16 @@
 
 
 # pylint: disable=redefined-builtin, arguments-differ, raise-missing-from, no-else-return, import-outside-toplevel
-class GeminiLLM(LLMInterface, LLMInterfaceV2):
-    """LLM interface for Google Gemini via the google.genai SDK.
+class BaseGeminiLLM(LLMBase, abc.ABC):
+    """Base class for Google Gemini LLMs (google.genai SDK).
 
-    Args:
-        model_name (str): Model name. Defaults to "gemini-2.0-flash".
-        model_params (Optional[dict]): Additional parameters passed to the model.
-        rate_limit_handler (Optional[RateLimitHandler]): Handler for rate limiting.
-        **kwargs (Any): Arguments passed to the genai.Client.
+    Holds all the shared message-building, config/schema-building, and
+    response-parsing logic. Subclasses are only responsible for
+    constructing the ``client`` SDK instance.
     """
 
+    client: "genai.Client"
+
     def __init__(
         self,
         model_name: str = "gemini-2.0-flash",
@@ -84,14 +85,13 @@ def __init__(
                 "Could not import google-genai python client. "
                 'Please install it with `pip install "neo4j-graphrag[google-genai]"`.'
             )
-        LLMInterfaceV2.__init__(
+        LLMBase.__init__(
             self,
             model_name=model_name,
             model_params=model_params or {},
             rate_limit_handler=rate_limit_handler,
             **kwargs,
         )
-        self.client = genai.Client(**kwargs)
 
     @overload  # type: ignore[no-overload-impl]
     def invoke(
@@ -275,6 +275,10 @@ async def ainvoke_with_tools(
         except Exception as e:
             raise LLMGenerationError(f"Error calling GeminiLLM with tools: {e}") from e
 
+    async def aclose(self) -> None:
+        self.client.close()
+        await self.client.aio.aclose()
+
     def get_messages(
         self,
         input: str,
@@ -406,3 +410,47 @@ def _parse_tool_response(
                         )
                     )
         return ToolCallResponse(tool_calls=tool_calls, content=None)
+
+
+class GeminiLLM(BaseGeminiLLM):
+    """LLM interface for Google Gemini via the google.genai SDK.
+
+    Args:
+        model_name (str): Model name. Defaults to "gemini-2.0-flash".
+        model_params (Optional[dict]): Additional parameters passed to the model.
+        rate_limit_handler (Optional[RateLimitHandler]): Handler for rate limiting.
+        base_url (Optional[str], optional): Base URL to use instead of Google's default
+            API endpoint, e.g. to reach a custom Gemini-compatible endpoint. Unlike the
+            Anthropic/OpenAI SDKs, ``genai.Client`` has no top-level ``base_url``
+            argument: the value is applied through ``http_options``. If ``http_options``
+            is also passed via kwargs (as a dict or ``types.HttpOptions``), this
+            parameter overrides its ``base_url`` field and leaves the rest untouched.
+            Defaults to None.
+        **kwargs (Any): Arguments passed to the genai.Client.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "gemini-2.0-flash",
+        model_params: Optional[dict[str, Any]] = None,
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+        base_url: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            model_name=model_name,
+            model_params=model_params,
+            rate_limit_handler=rate_limit_handler,
+            **kwargs,
+        )
+        if base_url is not None:
+            http_options = kwargs.get("http_options")
+            if http_options is None:
+                kwargs["http_options"] = types.HttpOptions(base_url=base_url)
+            elif isinstance(http_options, dict):
+                kwargs["http_options"] = {**http_options, "base_url": base_url}
+            else:  # types.HttpOptions
+                kwargs["http_options"] = http_options.model_copy(
+                    update={"base_url": base_url}
+                )
+        self.client = genai.Client(**kwargs)

tests/unit/llm/test_google_genai_llm.py

@@ -20,7 +20,7 @@
 import pytest
 
 from neo4j_graphrag.exceptions import LLMGenerationError
-from neo4j_graphrag.llm import GeminiLLM
+from neo4j_graphrag.llm import BaseGeminiLLM, GeminiLLM
 from neo4j_graphrag.types import LLMMessage
 
 
@@ -134,3 +134,127 @@ def test_gemini_invoke_error(mock_genai: Tuple[MagicMock, MagicMock]) -> None:
     llm = GeminiLLM("gemini-2.0-flash")
     with pytest.raises(LLMGenerationError):
         llm.invoke("hello")
+
+
+def test_gemini_llm_is_base_gemini_llm_subclass() -> None:
+    assert issubclass(GeminiLLM, BaseGeminiLLM)
+
+
+def test_gemini_llm_init_only_constructs_client(
+    mock_genai: Tuple[MagicMock, MagicMock],
+) -> None:
+    """GeminiLLM's __init__ should only be responsible for client construction."""
+    mock_gen, _ = mock_genai
+    llm = GeminiLLM("gemini-2.0-flash", api_key="my-key")
+
+    mock_gen.Client.assert_called_once_with(api_key="my-key")
+    assert llm.client is mock_gen.Client.return_value
+
+
+def test_minimal_base_gemini_llm_subclass_exercises_invoke(
+    mock_genai: Tuple[MagicMock, MagicMock],
+) -> None:
+    """BaseGeminiLLM does not construct the SDK client itself: a minimal
+    subclass that only assigns ``client`` should exercise the shared
+    invoke/message-building logic correctly. This is the extension contract
+    exported to custom subclasses."""
+    mock_gen, _ = mock_genai
+    custom_client = MagicMock()
+    mock_response = MagicMock()
+    mock_response.text = "minimal subclass response"
+    custom_client.models.generate_content.return_value = mock_response
+
+    class MinimalGeminiLLM(BaseGeminiLLM):
+        def __init__(self, model_name: str) -> None:
+            super().__init__(model_name=model_name)
+            self.client = custom_client
+
+    llm = MinimalGeminiLLM("gemini-2.0-flash")
+    response = llm.invoke("hello")
+
+    assert response.content == "minimal subclass response"
+    custom_client.models.generate_content.assert_called_once()
+    # the default genai.Client was never constructed
+    mock_gen.Client.assert_not_called()
+
+
+def test_gemini_llm_base_url_builds_http_options(
+    mock_genai: Tuple[MagicMock, MagicMock],
+) -> None:
+    """base_url alone must be applied through a new types.HttpOptions."""
+    mock_gen, mock_types = mock_genai
+    base_url = "https://custom-gemini-endpoint.example.com"
+
+    GeminiLLM(model_name="gemini-2.0-flash", base_url=base_url)
+
+    mock_types.HttpOptions.assert_called_once_with(base_url=base_url)
+    _, client_kwargs = mock_gen.Client.call_args
+    assert client_kwargs["http_options"] is mock_types.HttpOptions.return_value
+
+
+def test_gemini_llm_base_url_merges_into_http_options_dict(
+    mock_genai: Tuple[MagicMock, MagicMock],
+) -> None:
+    """base_url must override the base_url field of a dict http_options,
+    preserving the other fields."""
+    mock_gen, _ = mock_genai
+    base_url = "https://custom-gemini-endpoint.example.com"
+
+    GeminiLLM(
+        model_name="gemini-2.0-flash",
+        base_url=base_url,
+        http_options={"api_version": "v1", "base_url": "https://overridden"},
+    )
+
+    _, client_kwargs = mock_gen.Client.call_args
+    assert client_kwargs["http_options"] == {
+        "api_version": "v1",
+        "base_url": base_url,
+    }
+
+
+def test_gemini_llm_base_url_updates_http_options_object(
+    mock_genai: Tuple[MagicMock, MagicMock],
+) -> None:
+    """base_url must override the base_url field of a types.HttpOptions object
+    via model_copy, preserving the other fields."""
+    mock_gen, _ = mock_genai
+    base_url = "https://custom-gemini-endpoint.example.com"
+    http_options = MagicMock()  # stands in for a types.HttpOptions instance
+
+    GeminiLLM(
+        model_name="gemini-2.0-flash", base_url=base_url, http_options=http_options
+    )
+
+    http_options.model_copy.assert_called_once_with(update={"base_url": base_url})
+    _, client_kwargs = mock_gen.Client.call_args
+    assert client_kwargs["http_options"] is http_options.model_copy.return_value
+
+
+def test_gemini_llm_no_base_url_not_passed_to_client(
+    mock_genai: Tuple[MagicMock, MagicMock],
+) -> None:
+    """Omitting base_url should not pass http_options (or None) to the client."""
+    mock_gen, _ = mock_genai
+
+    GeminiLLM(model_name="gemini-2.0-flash")
+
+    _, client_kwargs = mock_gen.Client.call_args
+    assert "http_options" not in client_kwargs
+    assert "base_url" not in client_kwargs
+
+
+@pytest.mark.asyncio
+async def test_gemini_llm_aclose_closes_sync_and_async_clients(
+    mock_genai: Tuple[MagicMock, MagicMock],
+) -> None:
+    """aclose must release both the sync client and its async (aio) surface."""
+    mock_gen, _ = mock_genai
+    mock_client = mock_gen.Client.return_value
+    mock_client.aio.aclose = AsyncMock()
+
+    llm = GeminiLLM(model_name="gemini-2.0-flash")
+    await llm.aclose()
+
+    mock_client.close.assert_called_once()
+    mock_client.aio.aclose.assert_awaited_once()