From f8647cdcb6e1427a5c1ba5d0442b5ad872914278 Mon Sep 17 00:00:00 2001 From: matteomedioli Date: Fri, 17 Jul 2026 13:35:32 +0200 Subject: [PATCH 1/4] Extract BaseGeminiLLM to hold shared message-building, config, and response-parsing logic GeminiLLM now subclasses BaseGeminiLLM and is responsible only for constructing the genai.Client; everything else (get_messages, get_messages_v2, config/schema building, tool handling, response parsing) moves to the base class unchanged. BaseGeminiLLM is exported from neo4j_graphrag.llm as a documented extension point, mirroring BaseAnthropicLLM/BaseOpenAILLM. Also sets supports_structured_output = True on the base class: GeminiLLM already implements structured output via response_schema/response_mime_type but never declared the capability flag, so callers had to patch it in manually. --- CHANGELOG.md | 2 + docs/source/api.rst | 7 ++++ src/neo4j_graphrag/llm/__init__.py | 3 +- src/neo4j_graphrag/llm/google_genai_llm.py | 48 +++++++++++++++++----- tests/unit/llm/test_google_genai_llm.py | 31 +++++++++++++- 5 files changed, 79 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a94f3cfad..b9d34c849 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ - Added an explicit `base_url` keyword parameter to `AnthropicLLM` and `OpenAILLM`, passed through to both the sync and async SDK clients of each. - `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 @@ -27,6 +28,7 @@ ### Fixed - Fixed a bug in `AnthropicLLM` where an `http_client` passed via kwargs (whether an `httpx.Client` or `httpx.AsyncClient`) was forwarded to both the sync `anthropic.Anthropic` and async `anthropic.AsyncAnthropic` clients, causing a type mismatch. `http_client` is now routed to the matching sync/async client only; other kwargs remain shared. An `http_client` of an unrecognized type now emits a warning and is ignored instead of raising, matching `OpenAILLM`'s existing behavior. +- `GeminiLLM` now correctly declares `supports_structured_output = True`. It already implemented structured-output handling via `response_schema`/`response_mime_type`, but never set the capability flag, so callers gating behavior on it saw `False` and had to work around it manually. ## 1.18.0 diff --git a/docs/source/api.rst b/docs/source/api.rst index cdaa42087..b35ae6fa1 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -357,6 +357,13 @@ OpenAILLM :undoc-members: get_messages, client_class, async_client_class +BaseGeminiLLM +------------- + +.. autoclass:: neo4j_graphrag.llm.google_genai_llm.BaseGeminiLLM + :members: + + GeminiLLM --------- diff --git a/src/neo4j_graphrag/llm/__init__.py b/src/neo4j_graphrag/llm/__init__.py index 11a16f4c8..9f3d1e21c 100644 --- a/src/neo4j_graphrag/llm/__init__.py +++ b/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", diff --git a/src/neo4j_graphrag/llm/google_genai_llm.py b/src/neo4j_graphrag/llm/google_genai_llm.py index 52cee38fd..f55a670bb 100644 --- a/src/neo4j_graphrag/llm/google_genai_llm.py +++ b/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,18 @@ # 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. """ + supports_structured_output: bool = True + + client: "genai.Client" + def __init__( self, model_name: str = "gemini-2.0-flash", @@ -84,14 +87,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( @@ -406,3 +408,29 @@ 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. + **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, + **kwargs: Any, + ) -> None: + super().__init__( + model_name=model_name, + model_params=model_params, + rate_limit_handler=rate_limit_handler, + **kwargs, + ) + self.client = genai.Client(**kwargs) diff --git a/tests/unit/llm/test_google_genai_llm.py b/tests/unit/llm/test_google_genai_llm.py index 997d985a1..374de4011 100644 --- a/tests/unit/llm/test_google_genai_llm.py +++ b/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,32 @@ 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_supports_structured_output( + mock_genai: Tuple[MagicMock, MagicMock], +) -> None: + llm = GeminiLLM("gemini-2.0-flash") + assert llm.supports_structured_output is True + assert BaseGeminiLLM.supports_structured_output is True + + +def test_gemini_llm_init_only_constructs_client( + mock_genai: Tuple[MagicMock, MagicMock], +) -> None: + """GeminiLLM's __init__ should only be responsible for client construction; + everything else (message building, config building, parsing) is inherited + from BaseGeminiLLM unchanged.""" + 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 + # inherited behavior, not reimplemented on GeminiLLM + assert llm.get_messages.__func__ is BaseGeminiLLM.get_messages + assert llm.get_messages_v2.__func__ is BaseGeminiLLM.get_messages_v2 + assert llm._build_config.__func__ is BaseGeminiLLM._build_config From 1f82ee64b37c99a602182bbe02d06df50dda5bc9 Mon Sep 17 00:00:00 2001 From: matteomedioli Date: Fri, 17 Jul 2026 15:25:28 +0200 Subject: [PATCH 2/4] refactor(llm): keep BaseGeminiLLM extraction behavior-neutral MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups: - Drop the supports_structured_output=True flag flip: it changes SimpleKGPipeline/extractor default behavior for existing GeminiLLM users and belongs in its own PR with a (breaking)-labeled changelog entry, mirroring how the AnthropicLLM flag change was documented. - Replace method-identity assertions (implementation-shape tests) with a minimal BaseGeminiLLM subclass test that exercises invoke() end to end — the actual exported extension contract. --- CHANGELOG.md | 1 - src/neo4j_graphrag/llm/google_genai_llm.py | 2 - tests/unit/llm/test_google_genai_llm.py | 43 ++++++++++++++-------- 3 files changed, 28 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b9d34c849..7a3c502f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,7 +28,6 @@ ### Fixed - Fixed a bug in `AnthropicLLM` where an `http_client` passed via kwargs (whether an `httpx.Client` or `httpx.AsyncClient`) was forwarded to both the sync `anthropic.Anthropic` and async `anthropic.AsyncAnthropic` clients, causing a type mismatch. `http_client` is now routed to the matching sync/async client only; other kwargs remain shared. An `http_client` of an unrecognized type now emits a warning and is ignored instead of raising, matching `OpenAILLM`'s existing behavior. -- `GeminiLLM` now correctly declares `supports_structured_output = True`. It already implemented structured-output handling via `response_schema`/`response_mime_type`, but never set the capability flag, so callers gating behavior on it saw `False` and had to work around it manually. ## 1.18.0 diff --git a/src/neo4j_graphrag/llm/google_genai_llm.py b/src/neo4j_graphrag/llm/google_genai_llm.py index f55a670bb..070358723 100644 --- a/src/neo4j_graphrag/llm/google_genai_llm.py +++ b/src/neo4j_graphrag/llm/google_genai_llm.py @@ -71,8 +71,6 @@ class BaseGeminiLLM(LLMBase, abc.ABC): constructing the ``client`` SDK instance. """ - supports_structured_output: bool = True - client: "genai.Client" def __init__( diff --git a/tests/unit/llm/test_google_genai_llm.py b/tests/unit/llm/test_google_genai_llm.py index 374de4011..960caa600 100644 --- a/tests/unit/llm/test_google_genai_llm.py +++ b/tests/unit/llm/test_google_genai_llm.py @@ -140,26 +140,39 @@ def test_gemini_llm_is_base_gemini_llm_subclass() -> None: assert issubclass(GeminiLLM, BaseGeminiLLM) -def test_gemini_llm_supports_structured_output( - mock_genai: Tuple[MagicMock, MagicMock], -) -> None: - llm = GeminiLLM("gemini-2.0-flash") - assert llm.supports_structured_output is True - assert BaseGeminiLLM.supports_structured_output is True - - def test_gemini_llm_init_only_constructs_client( mock_genai: Tuple[MagicMock, MagicMock], ) -> None: - """GeminiLLM's __init__ should only be responsible for client construction; - everything else (message building, config building, parsing) is inherited - from BaseGeminiLLM unchanged.""" + """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 - # inherited behavior, not reimplemented on GeminiLLM - assert llm.get_messages.__func__ is BaseGeminiLLM.get_messages - assert llm.get_messages_v2.__func__ is BaseGeminiLLM.get_messages_v2 - assert llm._build_config.__func__ is BaseGeminiLLM._build_config + + +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() From cb076982c4748dbe8a23911f419cac2eae41b3f0 Mon Sep 17 00:00:00 2001 From: matteomedioli Date: Wed, 22 Jul 2026 18:06:51 +0200 Subject: [PATCH 3/4] add base_url support and extensibility docs for GeminiLLM --- CHANGELOG.md | 1 + docs/source/api.rst | 3 + docs/source/llm.rst | 56 +++++++++++++---- examples/customize/llms/google_genai_llm.py | 11 ++++ src/neo4j_graphrag/llm/google_genai_llm.py | 18 ++++++ tests/unit/llm/test_google_genai_llm.py | 66 +++++++++++++++++++++ 6 files changed, 143 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a3c502f9..45ba72afe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ - 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. - `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. diff --git a/docs/source/api.rst b/docs/source/api.rst index b35ae6fa1..5d216f8e1 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -360,6 +360,9 @@ OpenAILLM BaseGeminiLLM ------------- +See :ref:`llm-extensibility` for the ``http_client``/``base_url`` extension +contract. + .. autoclass:: neo4j_graphrag.llm.google_genai_llm.BaseGeminiLLM :members: diff --git a/docs/source/llm.rst b/docs/source/llm.rst index 1a8ed3892..ba0690545 100644 --- a/docs/source/llm.rst +++ b/docs/source/llm.rst @@ -9,18 +9,20 @@ 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. +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. Both ``AnthropicLLM`` and ``OpenAILLM`` accept two related, independent constructor settings: @@ -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. diff --git a/examples/customize/llms/google_genai_llm.py b/examples/customize/llms/google_genai_llm.py index bb3505d12..4c537cfe1 100644 --- a/examples/customize/llms/google_genai_llm.py +++ b/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) diff --git a/src/neo4j_graphrag/llm/google_genai_llm.py b/src/neo4j_graphrag/llm/google_genai_llm.py index 070358723..3cdb4cf04 100644 --- a/src/neo4j_graphrag/llm/google_genai_llm.py +++ b/src/neo4j_graphrag/llm/google_genai_llm.py @@ -415,6 +415,13 @@ class GeminiLLM(BaseGeminiLLM): 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. """ @@ -423,6 +430,7 @@ def __init__( 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__( @@ -431,4 +439,14 @@ def __init__( 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) diff --git a/tests/unit/llm/test_google_genai_llm.py b/tests/unit/llm/test_google_genai_llm.py index 960caa600..d9790de9b 100644 --- a/tests/unit/llm/test_google_genai_llm.py +++ b/tests/unit/llm/test_google_genai_llm.py @@ -176,3 +176,69 @@ def __init__(self, model_name: str) -> None: 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 From 5ca3ff99ce5a22c991cf4ec320412e0b26917f7e Mon Sep 17 00:00:00 2001 From: matteomedioli Date: Wed, 22 Jul 2026 18:22:18 +0200 Subject: [PATCH 4/4] close genai client on aclose, fix llm docs wording --- CHANGELOG.md | 1 + docs/source/llm.rst | 2 +- src/neo4j_graphrag/llm/google_genai_llm.py | 4 ++++ tests/unit/llm/test_google_genai_llm.py | 16 ++++++++++++++++ 4 files changed, 22 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 45ba72afe..2cac4ee10 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ - 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. diff --git a/docs/source/llm.rst b/docs/source/llm.rst index ba0690545..bddbf68fa 100644 --- a/docs/source/llm.rst +++ b/docs/source/llm.rst @@ -24,7 +24,7 @@ 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. -Both ``AnthropicLLM`` and ``OpenAILLM`` accept two related, independent +``AnthropicLLM`` and ``OpenAILLM`` accept two related, independent constructor settings: - ``base_url`` (``Optional[str]``): an explicit constructor parameter on both diff --git a/src/neo4j_graphrag/llm/google_genai_llm.py b/src/neo4j_graphrag/llm/google_genai_llm.py index 3cdb4cf04..d95b396e7 100644 --- a/src/neo4j_graphrag/llm/google_genai_llm.py +++ b/src/neo4j_graphrag/llm/google_genai_llm.py @@ -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, diff --git a/tests/unit/llm/test_google_genai_llm.py b/tests/unit/llm/test_google_genai_llm.py index d9790de9b..feeb5ca67 100644 --- a/tests/unit/llm/test_google_genai_llm.py +++ b/tests/unit/llm/test_google_genai_llm.py @@ -242,3 +242,19 @@ def test_gemini_llm_no_base_url_not_passed_to_client( _, 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()