feat(llm): extract BaseAnthropicLLM, add base_url, export extension points

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
 

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
 ------------
 

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
 

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.

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)

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",

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,28 +205,14 @@ 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,
             model_params=model_params or {},
             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)