Add support for Anthropic LLM (#150)

* Add support for Anthopic LLM

* mypy, docs, changelog

* Fix test names

Author: stellasia

CHANGELOG.md

@@ -14,6 +14,7 @@
 - Introduced Vertex AI LLM class for integrating Vertex AI models.
 - Added unit tests for the Vertex AI LLM class.
 - Added support for Cohere LLM and embeddings - added optional dependency to `cohere`.
+- Added support for Anthropic LLM - added optional dependency to `anthropic`.
 
 ### Fixed
 - Resolved import issue with the Vertex AI Embeddings class.

docs/source/api.rst

@@ -172,8 +172,11 @@ CohereEmbeddings
 Generation
 **********
 
+LLM
+===
+
 LLMInterface
-============
+------------
 
 .. autoclass:: neo4j_graphrag.llm.LLMInterface
     :members:
@@ -187,12 +190,27 @@ OpenAILLM
     :undoc-members: get_messages, client_class, async_client_class
 
 
+AzureOpenAILLM
+--------------
+
+.. autoclass:: neo4j_graphrag.llm.openai_llm.AzureOpenAILLM
+    :members:
+    :undoc-members: get_messages, client_class, async_client_class
+
+
 VertexAILLM
 -----------
 
 .. autoclass:: neo4j_graphrag.llm.vertexai_llm.VertexAILLM
     :members:
 
+AnthropicLLM
+------------
+
+.. autoclass:: neo4j_graphrag.llm.anthropic_llm.AnthropicLLM
+    :members:
+
+
 CohereLLM
 ---------
 

docs/source/user_guide_rag.rst

@@ -75,14 +75,15 @@ Using Another LLM Model
 
 If OpenAI cannot be used directly, there are a few available alternatives:
 
-- Use Azure OpenAI.
-- Use Google VertexAI.
+- Use Azure OpenAI (GPT...).
+- Use Google VertexAI (Gemini...).
+- Use Anthropic LLM (Claude...).
 - Use Cohere.
 - Use a local Ollama model.
 - Implement a custom interface.
 - Utilize any LangChain chat model.
 
-All options are illustrated below, using a local Ollama model as an example.
+All options are illustrated below.
 
 Using Azure Open AI LLM
 -----------------------
@@ -109,6 +110,9 @@ to learn more about the configuration.
     `pip install openai`
 
 
+See :ref:`azureopenaillm`.
+
+
 Using VertexAI LLM
 ------------------
 
@@ -132,6 +136,59 @@ To use VertexAI, instantiate the `VertexAILLM` class:
     `pip install google-cloud-aiplatform`
 
 
+See :ref:`vertexaillm`.
+
+
+Using Anthropic LLM
+-------------------
+
+To use Anthropic, instantiate the `AnthropicLLM` class:
+
+.. code:: python
+
+    from neo4j_graphrag.llm import AnthropicLLM
+
+    llm = AnthropicLLM(
+        model_name="claude-3-opus-20240229",
+        model_params={"max_tokens": 1000},  # max_tokens must be specified
+        api_key=api_key,  # can also set `ANTHROPIC_API_KEY` in env vars
+    )
+    llm.invoke("say something")
+
+
+.. note::
+
+    In order to run this code, the `anthropic` Python package needs to be installed:
+    `pip install anthropic`
+
+See :ref:`anthropicllm`.
+
+
+Using Cohere LLM
+----------------
+
+To use Cohere, instantiate the `CohereLLM` class:
+
+.. code:: python
+
+    from neo4j_graphrag.llm import CohereLLM
+
+    llm = CohereLLM(
+        model_name="command-r",
+        api_key=api_key,  # can also set `CO_API_KEY` in env vars
+    )
+    llm.invoke("say something")
+
+
+.. note::
+
+    In order to run this code, the `cohere` Python package needs to be installed:
+    `pip install cohere`
+
+
+See :ref:`coherellm`.
+
+
 Using a Local Model via Ollama
 -------------------------------
 

poetry.lock

@@ -43,6 +43,7 @@ pygraphviz = [
 ]
 google-cloud-aiplatform = {version = "^1.66.0", optional = true}
 cohere = {version = "^5.9.0", optional = true}
+anthropic = { version = "^0.34.2", optional = true}
 
 [tool.poetry.group.dev.dependencies]
 pylint = "^3.1.0"
@@ -77,9 +78,10 @@ pygraphviz = [
 ]
 google-cloud-aiplatform = {version = "^1.66.0"}
 cohere = {version = "^5.9.0"}
+anthropic = { version = "^0.34.2"}
 
 [tool.poetry.extras]
-external_clients = ["weaviate-client", "pinecone-client", "google-cloud-aiplatform", "cohere"]
+external_clients = ["weaviate-client", "pinecone-client", "google-cloud-aiplatform", "cohere", "anthropic"]
 kg_creation_tools = ["pygraphviz"]
 
 [build-system]

pyproject.toml

@@ -12,9 +12,19 @@
 #  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.
+from .anthropic_llm import AnthropicLLM
 from .base import LLMInterface
+from .cohere_llm import CohereLLM
 from .openai_llm import AzureOpenAILLM, OpenAILLM
 from .types import LLMResponse
 from .vertexai_llm import VertexAILLM
 
-__all__ = ["LLMResponse", "LLMInterface", "OpenAILLM", "VertexAILLM", "AzureOpenAILLM"]
+__all__ = [
+    "AnthropicLLM",
+    "CohereLLM",
+    "LLMResponse",
+    "LLMInterface",
+    "OpenAILLM",
+    "VertexAILLM",
+    "AzureOpenAILLM",
+]

src/neo4j_graphrag/llm/__init__.py

@@ -0,0 +1,116 @@
+#  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.
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm.base import LLMInterface
+from neo4j_graphrag.llm.types import LLMResponse
+
+try:
+    import anthropic
+    from anthropic import APIError
+except ImportError:
+    anthropic = None  # type: ignore
+    APIError = None  # type: ignore
+
+
+class AnthropicLLM(LLMInterface):
+    """Interface for large language models on Anthropic
+
+    Args:
+        model_name (str, optional): Name of the LLM to use. Defaults to "gemini-1.5-flash-001".
+        model_params (Optional[dict], optional): Additional parameters passed to the model when text is sent to it. 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,
+        **kwargs: Any,
+    ):
+        if anthropic is None:
+            raise ImportError(
+                "Could not import Anthropic Python client. "
+                "Please install it with `pip install anthropic`."
+            )
+        super().__init__(model_name, model_params)
+        self.client = anthropic.Anthropic(**kwargs)
+        self.async_client = anthropic.AsyncAnthropic(**kwargs)
+
+    def invoke(self, input: str) -> LLMResponse:
+        """Sends text to the LLM and returns a response.
+
+        Args:
+            input (str): The text to send to the LLM.
+
+        Returns:
+            LLMResponse: The response from the LLM.
+        """
+        try:
+            response = self.client.messages.create(
+                model=self.model_name,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": input,
+                    }
+                ],
+                **self.model_params,
+            )
+            return LLMResponse(content=response.content)
+        except APIError as e:
+            raise LLMGenerationError(e)
+
+    async def ainvoke(self, input: str) -> LLMResponse:
+        """Asynchronously sends text to the LLM and returns a response.
+
+        Args:
+            input (str): The text to send to the LLM.
+
+        Returns:
+            LLMResponse: The response from the LLM.
+        """
+        try:
+            response = await self.async_client.messages.create(
+                model=self.model_name,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": input,
+                    }
+                ],
+                **self.model_params,
+            )
+            return LLMResponse(content=response.content)
+        except APIError as e:
+            raise LLMGenerationError(e)

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -17,7 +17,8 @@
 from typing import Any, Optional
 
 from neo4j_graphrag.exceptions import LLMGenerationError
-from neo4j_graphrag.llm import LLMInterface, LLMResponse
+from neo4j_graphrag.llm.base import LLMInterface
+from neo4j_graphrag.llm.types import LLMResponse
 
 try:
     import cohere