diff --git a/CHANGELOG.md b/CHANGELOG.md index e080385fd..749f8b2dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,9 @@ ### 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 `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`, passed through to both the sync `anthropic.Anthropic` and async `anthropic.AsyncAnthropic` clients. +- Added a new docs page, `docs/source/llm_extensibility.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 069d21ffe..1ae22414d 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..b57c50fc0 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_extensibility.rst api.rst types.rst diff --git a/docs/source/llm_extensibility.rst b/docs/source/llm_extensibility.rst new file mode 100644 index 000000000..e9730002a --- /dev/null +++ b/docs/source/llm_extensibility.rst @@ -0,0 +1,78 @@ +.. _llm-extensibility: + +*********************************************** +Extending LLMs: BaseAnthropicLLM/BaseOpenAILLM +*********************************************** + +``BaseAnthropicLLM`` and ``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 arguments: + +- ``base_url`` (``Optional[str]``): 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. 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 logs a warning and falls back to the SDK's default client. + +Both arguments can be used together: ``base_url`` changes where requests go, +while ``http_client`` changes how they're sent. + +Subclassing example +==================== + +Because ``BaseAnthropicLLM``/``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 + + + class MyCustomAnthropicLLM(BaseAnthropicLLM): + """Talks to a self-hosted, Anthropic-compatible endpoint.""" + + 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) + self.client = anthropic.Anthropic( + base_url="https://my-custom-endpoint.example.com", + api_key="my-custom-api-key", + ) + self.async_client = anthropic.AsyncAnthropic( + base_url="https://my-custom-endpoint.example.com", + api_key="my-custom-api-key", + ) + + + 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 ``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..10c9c97c3 100644 --- a/src/neo4j_graphrag/llm/__init__.py +++ b/src/neo4j_graphrag/llm/__init__.py @@ -15,19 +15,20 @@ 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 .vertexai_llm import VertexAILLM __all__ = [ "AnthropicLLM", + "BaseAnthropicLLM", "BedrockLLM", "CohereLLM", "GeminiLLM", @@ -38,6 +39,7 @@ "LLMInterfaceV2", "OllamaLLM", "OpenAILLM", + "BaseOpenAILLM", "VertexAILLM", "AzureOpenAILLM", "MistralAILLM", diff --git a/src/neo4j_graphrag/llm/anthropic_llm.py b/src/neo4j_graphrag/llm/anthropic_llm.py index 556b39199..1b25f0879 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 import warnings from typing import ( @@ -56,8 +57,11 @@ ) if TYPE_CHECKING: - from anthropic import Omit + from anthropic import AsyncAnthropic, Anthropic, Omit from anthropic.types.message_param import MessageParam +else: + Anthropic = Any + AsyncAnthropic = Any # --------------------------------------------------------------------------- @@ -174,35 +178,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 +class BaseAnthropicLLM(LLMBase, abc.ABC): + """Base class for Anthropic LLMs. - 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?") + 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, @@ -217,6 +205,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, @@ -224,21 +213,6 @@ def __init__( rate_limit_handler=rate_limit_handler, **kwargs, ) - self.anthropic = anthropic - http_client = kwargs.pop("http_client", None) - sync_params = kwargs.copy() - async_params = kwargs.copy() - if httpx is not None and isinstance(http_client, httpx.Client): - sync_params["http_client"] = http_client - elif httpx is not None and isinstance(http_client, httpx.AsyncClient): - async_params["http_client"] = http_client - elif http_client is not None: - warnings.warn( - f"Invalid http_client type (got {type(http_client)}, expected httpx.Client or httpx.AsyncClient). Using default client.", - stacklevel=2, - ) - self.client = anthropic.Anthropic(**sync_params) - self.async_client = anthropic.AsyncAnthropic(**async_params) def invoke( self, @@ -547,3 +521,67 @@ 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. 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, + ) + anthropic = self.anthropic + http_client = kwargs.pop("http_client", None) + sync_params = kwargs.copy() + async_params = kwargs.copy() + if base_url is not None: + sync_params["base_url"] = base_url + async_params["base_url"] = base_url + if httpx is not None and isinstance(http_client, httpx.Client): + sync_params["http_client"] = http_client + elif httpx is not None and isinstance(http_client, httpx.AsyncClient): + async_params["http_client"] = http_client + elif http_client is not None: + warnings.warn( + f"Invalid http_client type (got {type(http_client)}, expected httpx.Client or httpx.AsyncClient). Using default client.", + stacklevel=2, + ) + self.client = anthropic.Anthropic(**sync_params) + self.async_client = anthropic.AsyncAnthropic(**async_params) diff --git a/tests/unit/llm/test_anthropic_llm.py b/tests/unit/llm/test_anthropic_llm.py index 6e9f4bae4..75d7fa18e 100644 --- a/tests/unit/llm/test_anthropic_llm.py +++ b/tests/unit/llm/test_anthropic_llm.py @@ -25,6 +25,7 @@ from neo4j_graphrag.experimental.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..d651a9e88 --- /dev/null +++ b/tests/unit/llm/test_llm_init.py @@ -0,0 +1,30 @@ +# 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__