diff --git a/AGENTS.md b/AGENTS.md index f461517dd..cd4be80a1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -97,6 +97,7 @@ ENVIRONMENT MANAGEMENT: - `WeakKeyDictionary` works for caching per-driver since `neo4j.Driver` is hashable - Neo4j 2026 CREATE VECTOR INDEX syntax: WITH clause must come BEFORE OPTIONS, not after - E2E tests for SEARCH clause: use `docker compose -f tests/e2e/docker-compose.neo4j2026.yml up -d` +- If `tests/unit/llm/test_anthropic_llm.py` fails with `AttributeError: module 'anthropic' has no attribute 'omit'`, or other unit test files error on missing optional-dependency imports (openai, cohere, etc.) at collection time, the local `.venv` is stale relative to `pyproject.toml` extras. Run `uv sync --all-extras` to fix. --- diff --git a/CHANGELOG.md b/CHANGELOG.md index e38462e56..a94f3cfad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,11 @@ ### Added - `AnthropicLLM` now supports structured output via the `response_format` argument, accepting a Pydantic model or an Anthropic `output_config` dict, alongside `OpenAILLM` and `VertexAILLM`. -- 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. +- 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. +- `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. ### Changed diff --git a/docs/source/api.rst b/docs/source/api.rst index b2e7be894..cdaa42087 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -339,6 +339,16 @@ LLMBase :members: +BaseOpenAILLM +------------- + +See :ref:`llm-extensibility` for the ``http_client``/``base_url`` extension +contract. + +.. autoclass:: neo4j_graphrag.llm.openai_llm.BaseOpenAILLM + :members: + + OpenAILLM --------- @@ -374,6 +384,16 @@ VertexAILLM .. autoclass:: neo4j_graphrag.llm.vertexai_llm.VertexAILLM :members: +BaseAnthropicLLM +---------------- + +See :ref:`llm-extensibility` for the ``http_client``/``base_url`` extension +contract. + +.. autoclass:: neo4j_graphrag.llm.anthropic_llm.BaseAnthropicLLM + :members: + + AnthropicLLM ------------ diff --git a/docs/source/index.rst b/docs/source/index.rst index b1241a722..d1711cefd 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -38,6 +38,7 @@ Topics + :ref:`user-guide-rag` + :ref:`user-guide-kg-builder` + :ref:`user-guide-pipeline` ++ :ref:`llm-extensibility` + :ref:`api-documentation` + :ref:`types-documentation` @@ -50,6 +51,7 @@ Topics user_guide_rag.rst user_guide_kg_builder.rst user_guide_pipeline.rst + llm.rst api.rst types.rst diff --git a/docs/source/llm.rst b/docs/source/llm.rst new file mode 100644 index 000000000..1a8ed3892 --- /dev/null +++ b/docs/source/llm.rst @@ -0,0 +1,126 @@ +.. _llm-page: + +*********************************************** +LLMs +*********************************************** + +This page gathers LLM-related documentation: configuring the built-in +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 +constructor settings: + +- ``base_url`` (``Optional[str]``): an explicit constructor parameter on both + classes that overrides the default API endpoint. Passed through to both the + sync and async SDK clients (``anthropic.Anthropic`` / + ``anthropic.AsyncAnthropic`` for Anthropic, and the equivalent OpenAI SDK + clients for OpenAI). +- ``http_client`` (``Optional[httpx.Client | httpx.AsyncClient]``): an + already-configured ``httpx`` client to use for requests, e.g. to add custom + TLS settings, proxies, or timeouts. Accepted through ``**kwargs`` by both + classes. The concrete class inspects the type of the object you pass: an + ``httpx.Client`` is routed to the sync SDK client, and an + ``httpx.AsyncClient`` is routed to the async SDK client. Passing something + else emits a ``UserWarning`` (via ``warnings.warn``) and falls back to the + SDK's default client. + +Both settings can be used together: ``base_url`` changes where requests go, +while ``http_client`` changes how they're sent. Note that a single +``http_client`` only customizes one direction: calling the other one (e.g. +``ainvoke`` after passing a sync ``httpx.Client``) still works and targets +``base_url``, but through the SDK's default transport. + +.. note:: + + A ``base_url`` configured on the ``httpx`` client itself is ignored (a + ``UserWarning`` is emitted when one is detected): both + SDKs build absolute request URLs from their own ``base_url`` and only use + the ``httpx`` client as transport. To change the endpoint, always use the + ``base_url`` constructor parameter: + + .. code-block:: python + + # IGNORED -- the httpx base_url is never used; requests still go to + # the SDK's default endpoint + AnthropicLLM( + model_name="...", + http_client=httpx.Client(base_url="https://my-endpoint"), + ) + + # WORKS -- requests go to the custom endpoint, with the custom transport + AnthropicLLM( + model_name="...", + base_url="https://my-endpoint", + http_client=httpx.Client(proxy="http://my-proxy:8080"), + ) + +The sync/async routing is implemented by +:func:`neo4j_graphrag.llm.utils.split_http_client_kwargs`, which is exported +for exactly one reason: a custom subclass that constructs its own SDK clients +should call it too, so it preserves the same routing contract instead of +reintroducing the type-mismatch bug the helper fixes. + +Subclassing example +------------------- + +Because :class:`~neo4j_graphrag.llm.anthropic_llm.BaseAnthropicLLM` / +:class:`~neo4j_graphrag.llm.openai_llm.BaseOpenAILLM` only require the concrete +subclass to assign ``self.client``/``self.async_client``, you can build your +own thin subclass to reach a custom Anthropic-compatible endpoint with +different defaults or credential handling than the built-in ``AnthropicLLM``: + +.. code:: python + + from typing import Any, Optional + + import anthropic + + from neo4j_graphrag.llm import BaseAnthropicLLM + from neo4j_graphrag.llm.utils import split_http_client_kwargs + + + class MyCustomAnthropicLLM(BaseAnthropicLLM): + """Talks to a self-hosted, Anthropic-compatible endpoint.""" + + DEFAULT_ENDPOINT = "https://my-custom-endpoint.example.com" + + def __init__( + self, + model_name: str, + model_params: Optional[dict[str, Any]] = None, + **kwargs: Any, + ): + super().__init__(model_name=model_name, model_params=model_params, **kwargs) + # Route an optional http_client kwarg to the matching sync/async + # client, exactly as the built-in AnthropicLLM does. + sync_params, async_params = split_http_client_kwargs(kwargs) + sync_params.setdefault("base_url", self.DEFAULT_ENDPOINT) + async_params.setdefault("base_url", self.DEFAULT_ENDPOINT) + self.client = anthropic.Anthropic(**sync_params) + self.async_client = anthropic.AsyncAnthropic(**async_params) + + + llm = MyCustomAnthropicLLM(model_name="claude-3-opus-20240229") + llm.invoke("Who is the mother of Paul Atreides?") + +All of ``invoke``/``ainvoke``, structured-output handling, and message +building are inherited from :class:`~neo4j_graphrag.llm.anthropic_llm.BaseAnthropicLLM` unchanged; the subclass only +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. diff --git a/examples/customize/llms/anthropic_llm.py b/examples/customize/llms/anthropic_llm.py index 265e6ea12..a67bf32a8 100644 --- a/examples/customize/llms/anthropic_llm.py +++ b/examples/customize/llms/anthropic_llm.py @@ -10,3 +10,15 @@ ) as llm: res: LLMResponse = llm.invoke("say something") print(res.content) + +# To reach a custom or self-hosted, Anthropic-compatible endpoint instead of +# Anthropic's default API, pass `base_url`. It's forwarded to both the sync +# and async SDK clients. +with AnthropicLLM( + model_name="claude-3-opus-20240229", + model_params={"max_tokens": 1000}, + api_key=api_key, + base_url="https://my-custom-endpoint.example.com", +) as custom_llm: + 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 e47965328..11a16f4c8 100644 --- a/src/neo4j_graphrag/llm/__init__.py +++ b/src/neo4j_graphrag/llm/__init__.py @@ -15,19 +15,21 @@ import warnings from typing import Any -from .anthropic_llm import AnthropicLLM +from .anthropic_llm import AnthropicLLM, BaseAnthropicLLM from .base import LLMBase, LLMInterface, LLMInterfaceV2 from .bedrock_llm import BedrockLLM from .cohere_llm import CohereLLM from .google_genai_llm import GeminiLLM from .mistralai_llm import MistralAILLM from .ollama_llm import OllamaLLM -from .openai_llm import AzureOpenAILLM, OpenAILLM +from .openai_llm import AzureOpenAILLM, BaseOpenAILLM, OpenAILLM from .types import LLMResponse, LLMUsage +from .utils import split_http_client_kwargs from .vertexai_llm import VertexAILLM __all__ = [ "AnthropicLLM", + "BaseAnthropicLLM", "BedrockLLM", "CohereLLM", "GeminiLLM", @@ -38,9 +40,11 @@ "LLMInterfaceV2", "OllamaLLM", "OpenAILLM", + "BaseOpenAILLM", "VertexAILLM", "AzureOpenAILLM", "MistralAILLM", + "split_http_client_kwargs", ] diff --git a/src/neo4j_graphrag/llm/anthropic_llm.py b/src/neo4j_graphrag/llm/anthropic_llm.py index cf8c62ba3..6fd235cbd 100644 --- a/src/neo4j_graphrag/llm/anthropic_llm.py +++ b/src/neo4j_graphrag/llm/anthropic_llm.py @@ -13,6 +13,7 @@ # limitations under the License. from __future__ import annotations +import abc import json from typing import ( TYPE_CHECKING, @@ -50,7 +51,7 @@ ) if TYPE_CHECKING: - from anthropic import Omit + from anthropic import AsyncAnthropic, Anthropic, Omit from anthropic.types.message_param import MessageParam @@ -168,35 +169,19 @@ def _restore_open_maps(value: Any, schema: dict[str, Any], defs: dict[str, Any]) # pylint: disable=redefined-builtin, arguments-differ, raise-missing-from, no-else-return, import-outside-toplevel -class AnthropicLLM(LLMBase): - """Interface for large language models on Anthropic - - Args: - model_name (str): Name of the LLM to use. - model_params (Optional[dict], optional): Additional parameters for LLMInterface(V1) passed to the model when text is sent to it. Defaults to None. - system_instruction: Optional[str], optional): Additional instructions for setting the behavior and context for the model in a conversation. Defaults to None. - rate_limit_handler (Optional[RateLimitHandler], optional): Handler for managing rate limits for LLMInterface(V1). Defaults to None. - **kwargs (Any): Arguments passed to the model when for the class is initialised. Defaults to None. - - Raises: - LLMGenerationError: If there's an error generating the response from the model. - - Example: - - .. code-block:: python - - from neo4j_graphrag.llm import AnthropicLLM +class BaseAnthropicLLM(LLMBase, abc.ABC): + """Base class for Anthropic LLMs. - llm = AnthropicLLM( - model_name="claude-3-opus-20240229", - model_params={"max_tokens": 1000}, - api_key="sk...", # can also be read from env vars - ) - llm.invoke("Who is the mother of Paul Atreides?") + Holds all the shared message-building, schema-conversion, and + response-parsing logic. Subclasses are only responsible for + constructing the ``client``/``async_client`` SDK instances. """ supports_structured_output: bool = True + client: Anthropic + async_client: AsyncAnthropic + def __init__( self, model_name: str, @@ -211,6 +196,7 @@ def __init__( """Could not import Anthropic Python client. Please install it with `pip install "neo4j-graphrag[anthropic]"`.""" ) + self.anthropic = anthropic LLMBase.__init__( self, model_name=model_name, @@ -218,10 +204,6 @@ def __init__( rate_limit_handler=rate_limit_handler, **kwargs, ) - self.anthropic = anthropic - sync_params, async_params = split_http_client_kwargs(kwargs) - self.client = anthropic.Anthropic(**sync_params) - self.async_client = anthropic.AsyncAnthropic(**async_params) def invoke( self, @@ -530,3 +512,58 @@ def get_messages_v2( ) ) return system_instruction, messages + + +class AnthropicLLM(BaseAnthropicLLM): + """Interface for large language models on Anthropic + + Args: + model_name (str): Name of the LLM to use. + model_params (Optional[dict], optional): Additional parameters for LLMInterface(V1) passed to the model when text is sent to it. Defaults to None. + system_instruction: Optional[str], optional): Additional instructions for setting the behavior and context for the model in a conversation. Defaults to None. + rate_limit_handler (Optional[RateLimitHandler], optional): Handler for managing rate limits for LLMInterface(V1). Defaults to None. + base_url (Optional[str], optional): Base URL to use instead of Anthropic's default API + endpoint, e.g. to reach a custom Anthropic-compatible endpoint. Passed through to + both the sync and async SDK clients. Can be combined with an ``http_client`` + passed via kwargs (``base_url`` sets where requests go, ``http_client`` how they + are sent); a base URL configured on the httpx client itself is ignored by the + SDK — use this parameter instead. Defaults to None. + **kwargs (Any): Arguments passed to the model when for the class is initialised. Defaults to None. + + Raises: + LLMGenerationError: If there's an error generating the response from the model. + + Example: + + .. code-block:: python + + from neo4j_graphrag.llm import AnthropicLLM + + llm = AnthropicLLM( + model_name="claude-3-opus-20240229", + model_params={"max_tokens": 1000}, + api_key="sk...", # can also be read from env vars + ) + llm.invoke("Who is the mother of Paul Atreides?") + """ + + def __init__( + self, + model_name: str, + model_params: Optional[dict[str, Any]] = None, + rate_limit_handler: Optional[RateLimitHandler] = None, + base_url: Optional[str] = None, + **kwargs: Any, + ): + super().__init__( + model_name=model_name, + model_params=model_params, + rate_limit_handler=rate_limit_handler, + **kwargs, + ) + sync_params, async_params = split_http_client_kwargs(kwargs) + if base_url is not None: + sync_params["base_url"] = base_url + async_params["base_url"] = base_url + self.client = self.anthropic.Anthropic(**sync_params) + self.async_client = self.anthropic.AsyncAnthropic(**async_params) diff --git a/src/neo4j_graphrag/llm/openai_llm.py b/src/neo4j_graphrag/llm/openai_llm.py index 47645d4e3..46a95b3d5 100644 --- a/src/neo4j_graphrag/llm/openai_llm.py +++ b/src/neo4j_graphrag/llm/openai_llm.py @@ -637,6 +637,7 @@ def __init__( model_name: str, model_params: Optional[dict[str, Any]] = None, rate_limit_handler: Optional[RateLimitHandler] = None, + base_url: Optional[str] = None, **kwargs: Any, ): """OpenAI LLM @@ -647,6 +648,12 @@ def __init__( model_name (str): model_params (str): Parameters for LLMInterface(V1) like temperature that will be passed to the model when text is sent to it. Defaults to None. rate_limit_handler (Optional[RateLimitHandler]): Handler for rate limiting for LLMInterface(V1). Defaults to retry with exponential backoff. + base_url (Optional[str], optional): Base URL to use instead of OpenAI's default API + endpoint, e.g. to reach an OpenAI-compatible server. Passed through to both the + sync and async SDK clients. Can be combined with an ``http_client`` passed via + kwargs (``base_url`` sets where requests go, ``http_client`` how they are + sent); a base URL configured on the httpx client itself is ignored by the + SDK — use this parameter instead. Defaults to None. kwargs: All other parameters will be passed to the openai.OpenAI init. """ super().__init__( @@ -655,6 +662,9 @@ def __init__( rate_limit_handler=rate_limit_handler, ) sync_params, async_params = split_http_client_kwargs(kwargs) + if base_url is not None: + sync_params["base_url"] = base_url + async_params["base_url"] = base_url self.client = self.openai.OpenAI(**sync_params) self.async_client = self.openai.AsyncOpenAI(**async_params) diff --git a/src/neo4j_graphrag/llm/utils.py b/src/neo4j_graphrag/llm/utils.py index d2fe9d28b..c3af3a029 100644 --- a/src/neo4j_graphrag/llm/utils.py +++ b/src/neo4j_graphrag/llm/utils.py @@ -108,10 +108,21 @@ def split_http_client_kwargs( http_client = kwargs.pop("http_client", None) sync_kwargs = kwargs.copy() async_kwargs = kwargs.copy() - if httpx is not None and isinstance(http_client, httpx.Client): - sync_kwargs["http_client"] = http_client - elif httpx is not None and isinstance(http_client, httpx.AsyncClient): - async_kwargs["http_client"] = http_client + if httpx is not None and isinstance(http_client, (httpx.Client, httpx.AsyncClient)): + if str(http_client.base_url): + # stacklevel=3 attributes the warning to the caller of the LLM + # constructor, not to the constructor's own call into this helper. + warnings.warn( + "The base_url configured on the provided http_client is ignored: " + "the SDK builds request URLs from its own base_url and uses the " + "http_client as transport only. Pass base_url to the LLM " + "constructor instead.", + stacklevel=3, + ) + if isinstance(http_client, httpx.Client): + sync_kwargs["http_client"] = http_client + else: + async_kwargs["http_client"] = http_client elif http_client is not None: # stacklevel=3 attributes the warning to the caller of the LLM # constructor, not to the constructor's own call into this helper. diff --git a/tests/unit/llm/test_anthropic_llm.py b/tests/unit/llm/test_anthropic_llm.py index 93ab0224c..ed33ac864 100644 --- a/tests/unit/llm/test_anthropic_llm.py +++ b/tests/unit/llm/test_anthropic_llm.py @@ -25,6 +25,7 @@ from neo4j_graphrag.components.types import Neo4jGraph from neo4j_graphrag.llm.anthropic_llm import ( AnthropicLLM, + BaseAnthropicLLM, _is_open_map, _resolve_ref, _restore_open_maps, @@ -539,6 +540,79 @@ def test_anthropic_llm_close(mock_anthropic: Mock) -> None: mock_anthropic.AsyncAnthropic.return_value.close.assert_called_once() +# --------------------------------------------------------------------------- +# BaseAnthropicLLM / thin subclass contract tests +# --------------------------------------------------------------------------- + + +def test_minimal_base_anthropic_llm_subclass_exercises_invoke( + mock_anthropic: Mock, +) -> None: + """BaseAnthropicLLM does not construct SDK clients itself: a minimal + subclass that only assigns client/async_client should exercise the shared + invoke/schema logic correctly.""" + mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock( + content=[MagicMock(text="minimal subclass response")] + ) + + class MinimalAnthropicLLM(BaseAnthropicLLM): + def __init__(self, model_name: str) -> None: + super().__init__(model_name=model_name) + self.client = self.anthropic.Anthropic() + self.async_client = self.anthropic.AsyncAnthropic() + + llm = MinimalAnthropicLLM(model_name="claude-3-opus-20240229") + response = llm.invoke("hello") + assert response.content == "minimal subclass response" + + +def test_anthropic_llm_is_subclass_of_base_anthropic_llm(mock_anthropic: Mock) -> None: + """AnthropicLLM is the thin, concrete subclass of BaseAnthropicLLM.""" + llm = AnthropicLLM(model_name="claude-3-opus-20240229") + assert isinstance(llm, BaseAnthropicLLM) + + +def test_anthropic_llm_base_url_reaches_both_clients(mock_anthropic: Mock) -> None: + """base_url must be forwarded to both the sync and async SDK clients.""" + base_url = "https://custom-anthropic-endpoint.example.com" + AnthropicLLM(model_name="claude-3-opus-20240229", base_url=base_url) + + _, sync_kwargs = mock_anthropic.Anthropic.call_args + assert sync_kwargs.get("base_url") == base_url + _, async_kwargs = mock_anthropic.AsyncAnthropic.call_args + assert async_kwargs.get("base_url") == base_url + + +def test_anthropic_llm_no_base_url_not_passed_to_clients(mock_anthropic: Mock) -> None: + """Omitting base_url should not pass it (or None) to either client.""" + AnthropicLLM(model_name="claude-3-opus-20240229") + + _, sync_kwargs = mock_anthropic.Anthropic.call_args + assert "base_url" not in sync_kwargs + _, async_kwargs = mock_anthropic.AsyncAnthropic.call_args + assert "base_url" not in async_kwargs + + +def test_anthropic_llm_base_url_with_http_client(mock_anthropic: Mock) -> None: + """base_url and http_client can be combined; both reach the expected client.""" + http_client = httpx.Client() + base_url = "https://custom-anthropic-endpoint.example.com" + try: + AnthropicLLM( + model_name="claude-3-opus-20240229", + base_url=base_url, + http_client=http_client, + ) + _, sync_kwargs = mock_anthropic.Anthropic.call_args + assert sync_kwargs.get("base_url") == base_url + assert sync_kwargs.get("http_client") is http_client + _, async_kwargs = mock_anthropic.AsyncAnthropic.call_args + assert async_kwargs.get("base_url") == base_url + assert "http_client" not in async_kwargs + finally: + http_client.close() + + # --------------------------------------------------------------------------- # http_client sync/async routing tests # --------------------------------------------------------------------------- diff --git a/tests/unit/llm/test_llm_init.py b/tests/unit/llm/test_llm_init.py new file mode 100644 index 000000000..c8c115ae0 --- /dev/null +++ b/tests/unit/llm/test_llm_init.py @@ -0,0 +1,37 @@ +# Neo4j Sweden AB [https://neo4j.com] +# # +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# # +# https://www.apache.org/licenses/LICENSE-2.0 +# # +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for public exports of neo4j_graphrag.llm.""" + +import neo4j_graphrag.llm as llm_module + + +def test_base_anthropic_llm_is_exported() -> None: + from neo4j_graphrag.llm import BaseAnthropicLLM + + assert BaseAnthropicLLM is not None + assert "BaseAnthropicLLM" in llm_module.__all__ + + +def test_base_openai_llm_is_exported() -> None: + from neo4j_graphrag.llm import BaseOpenAILLM + + assert BaseOpenAILLM is not None + assert "BaseOpenAILLM" in llm_module.__all__ + + +def test_split_http_client_kwargs_is_exported() -> None: + from neo4j_graphrag.llm import split_http_client_kwargs + + assert callable(split_http_client_kwargs) + assert "split_http_client_kwargs" in llm_module.__all__ diff --git a/tests/unit/llm/test_llm_utils.py b/tests/unit/llm/test_llm_utils.py index 29c417cdb..3677b0eb7 100644 --- a/tests/unit/llm/test_llm_utils.py +++ b/tests/unit/llm/test_llm_utils.py @@ -14,6 +14,8 @@ # limitations under the License. from typing import AsyncGenerator, Generator +import warnings + import httpx import pytest import pytest_asyncio @@ -117,6 +119,36 @@ async def test_split_http_client_kwargs_routes_async_client( assert sync_kwargs["api_key"] == "sk-test" +def test_split_http_client_kwargs_warns_on_client_with_base_url() -> None: + client = httpx.Client(base_url="https://my-endpoint.example.com") + try: + with pytest.warns(UserWarning, match="base_url configured on the provided"): + sync_kwargs, _ = split_http_client_kwargs({"http_client": client}) + # the client is still routed; only the base_url is flagged + assert sync_kwargs["http_client"] is client + finally: + client.close() + + +@pytest.mark.asyncio +async def test_split_http_client_kwargs_warns_on_async_client_with_base_url() -> None: + client = httpx.AsyncClient(base_url="https://my-endpoint.example.com") + try: + with pytest.warns(UserWarning, match="base_url configured on the provided"): + _, async_kwargs = split_http_client_kwargs({"http_client": client}) + assert async_kwargs["http_client"] is client + finally: + await client.aclose() + + +def test_split_http_client_kwargs_no_warning_without_base_url( + httpx_sync_client: httpx.Client, +) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("error") + split_http_client_kwargs({"http_client": httpx_sync_client}) + + def test_split_http_client_kwargs_invalid_type_warns_and_drops() -> None: with pytest.warns(UserWarning, match="Invalid http_client type"): sync_kwargs, async_kwargs = split_http_client_kwargs( diff --git a/tests/unit/llm/test_openai_llm.py b/tests/unit/llm/test_openai_llm.py index da4757960..23fda3db1 100644 --- a/tests/unit/llm/test_openai_llm.py +++ b/tests/unit/llm/test_openai_llm.py @@ -980,6 +980,35 @@ def test_openai_llm_with_httpx_client(mock_import: Mock) -> None: assert async_kwargs.get("http_client") is None +@patch("builtins.__import__") +def test_openai_llm_base_url_reaches_both_clients(mock_import: Mock) -> None: + """base_url must be forwarded to both the sync and async OpenAI clients.""" + mock_openai = get_mock_openai() + mock_import.return_value = mock_openai + + base_url = "https://custom-openai-endpoint.example.com/v1" + OpenAILLM(model_name="gpt", api_key="my key", base_url=base_url) + + _, sync_kwargs = mock_openai.OpenAI.call_args + assert sync_kwargs.get("base_url") == base_url + _, async_kwargs = mock_openai.AsyncOpenAI.call_args + assert async_kwargs.get("base_url") == base_url + + +@patch("builtins.__import__") +def test_openai_llm_no_base_url_not_passed_to_clients(mock_import: Mock) -> None: + """Omitting base_url should not pass it (or None) to either client.""" + mock_openai = get_mock_openai() + mock_import.return_value = mock_openai + + OpenAILLM(model_name="gpt", api_key="my key") + + _, sync_kwargs = mock_openai.OpenAI.call_args + assert "base_url" not in sync_kwargs + _, async_kwargs = mock_openai.AsyncOpenAI.call_args + assert "base_url" not in async_kwargs + + @patch("builtins.__import__") def test_openai_llm_with_httpx_async_client(mock_import: Mock) -> None: """Test that httpx.AsyncClient is forwarded only to the async OpenAI client without warning."""