diff --git a/CHANGELOG.md b/CHANGELOG.md index a94f3cfad..2cac4ee10 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/source/api.rst b/docs/source/api.rst index cdaa42087..5d216f8e1 100644 --- a/docs/source/api.rst +++ b/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 --------- diff --git a/docs/source/llm.rst b/docs/source/llm.rst index 1a8ed3892..bddbf68fa 100644 --- a/docs/source/llm.rst +++ b/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. 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/__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..d95b396e7 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,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) diff --git a/tests/unit/llm/test_google_genai_llm.py b/tests/unit/llm/test_google_genai_llm.py index 997d985a1..feeb5ca67 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,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()