Update LLMInterface to restore LC compatibility (#416)

* Update LLMInterface to restore LC compatibility

* Update AnthropicLLM

* Update MistralAILLM

* Update OllamaLLM

* Update CohereLLM

* Mypy / ruff

* Update VertexAILLM

* Update (a)invoke_with_tools methods in the same way

* Rename method and return directly list[LLMMessage]

* Update GraphRAG to restore full LC compatibility

* Test for the utils functions

* WIP: update tests

* Improve test coverage for utils and base modules

* Fix UT OpenAILLM

* Update Ollama tests

* Update Ollama/Anthropic

* WIP update cohere

* CHANGELOG.md

* Ruff after rebase

* More fixes on cohere tests

* Add tests for retry behavior

* Fix MistralAILLM tests

* Fix VertexAILLM tests

* mypy

* Address comments

* Fix mypy again

* Fix e2e

* Fix CI

* Comments

Author: stellasia

CHANGELOG.md

@@ -4,8 +4,15 @@
 
 ### Added
 
+- Document node is now always created when running SimpleKGPipeline, even if `from_pdf=False`.
+- Document metadata is exposed in SimpleKGPipeline run method.
 - Added automatic rate limiting with retry logic and exponential backoff for all Embedding providers using tenacity. The `RateLimitHandler` interface allows for custom rate limiting strategies, including the ability to disable rate limiting entirely.
 
+### Fixed
+
+- LangChain Chat models compatibility is now working again.
+
+
 ## 1.10.0
 
 ### Added

examples/README.md

@@ -69,7 +69,7 @@ are listed in [the last section of this file](#customize).
 - [OpenAI (GPT)](./customize/llms/openai_llm.py)
 - [Azure OpenAI]()
 - [VertexAI (Gemini)](./customize/llms/vertexai_llm.py)
-- [MistralAI](./customize/llms/mistalai_llm.py)
+- [MistralAI](customize/llms/mistralai_llm.py)
 - [Cohere](./customize/llms/cohere_llm.py)
 - [Anthropic (Claude)](./customize/llms/anthropic_llm.py)
 - [Ollama](./customize/llms/ollama_llm.py)
@@ -142,7 +142,7 @@ are listed in [the last section of this file](#customize).
 
 ### Answer: GraphRAG
 
-- [LangChain compatibility](./customize/answer/langchain_compatiblity.py)
+- [LangChain compatibility](customize/answer/langchain_compatibility.py)
 - [Use a custom prompt](./customize/answer/custom_prompt.py)
 
 

examples/customize/answer/langchain_compatiblity.py renamed to examples/customize/answer/langchain_compatibility.py

@@ -1,12 +1,28 @@
 from neo4j_graphrag.llm import AnthropicLLM, LLMResponse
+from neo4j_graphrag.types import LLMMessage
 
 # set api key here on in the ANTHROPIC_API_KEY env var
 api_key = None
 
+messages: list[LLMMessage] = [
+    {
+        "role": "system",
+        "content": "You are a seasoned actor and expert performer, renowned for your one-man shows and comedic talent.",
+    },
+    {
+        "role": "user",
+        "content": "say something",
+    },
+]
+
+
 llm = AnthropicLLM(
     model_name="claude-3-opus-20240229",
     model_params={"max_tokens": 1000},  # max_tokens must be specified
     api_key=api_key,
 )
-res: LLMResponse = llm.invoke("say something")
+res: LLMResponse = llm.invoke(
+    # "say something",
+    messages,
+)
 print(res.content)

examples/customize/llms/anthropic_llm.py

@@ -1,11 +1,23 @@
 from neo4j_graphrag.llm import CohereLLM, LLMResponse
+from neo4j_graphrag.types import LLMMessage
 
 # set api key here on in the CO_API_KEY env var
 api_key = None
 
+messages: list[LLMMessage] = [
+    {
+        "role": "system",
+        "content": "You are a seasoned actor and expert performer, renowned for your one-man shows and comedic talent.",
+    },
+    {
+        "role": "user",
+        "content": "say something",
+    },
+]
+
 llm = CohereLLM(
     model_name="command-r",
     api_key=api_key,
 )
-res: LLMResponse = llm.invoke("say something")
+res: LLMResponse = llm.invoke(input=messages)
 print(res.content)

examples/customize/llms/cohere_llm.py

@@ -1,14 +1,13 @@
 import random
 import string
-from typing import Any, Awaitable, Callable, List, Optional, TypeVar, Union
+from typing import Any, Awaitable, Callable, Optional, TypeVar
 
 from neo4j_graphrag.llm import LLMInterface, LLMResponse
 from neo4j_graphrag.utils.rate_limit import (
     RateLimitHandler,
     # rate_limit_handler,
     # async_rate_limit_handler,
 )
-from neo4j_graphrag.message_history import MessageHistory
 from neo4j_graphrag.types import LLMMessage
 
 
@@ -18,37 +17,27 @@ def __init__(
     ):
         super().__init__(model_name, **kwargs)
 
-    # Optional: Apply rate limit handling to synchronous invoke method
-    # @rate_limit_handler
-    def invoke(
+    def _invoke(
         self,
-        input: str,
-        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
-        system_instruction: Optional[str] = None,
+        input: list[LLMMessage],
     ) -> LLMResponse:
         content: str = (
             self.model_name + ": " + "".join(random.choices(string.ascii_letters, k=30))
         )
         return LLMResponse(content=content)
 
-    # Optional: Apply rate limit handling to asynchronous ainvoke method
-    # @async_rate_limit_handler
-    async def ainvoke(
+    async def _ainvoke(
         self,
-        input: str,
-        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
-        system_instruction: Optional[str] = None,
+        input: list[LLMMessage],
     ) -> LLMResponse:
         raise NotImplementedError()
 
 
-llm = CustomLLM(
-    ""
-)  # if rate_limit_handler and async_rate_limit_handler decorators are used, the default rate limit handler will be applied automatically (retry with exponential backoff)
+llm = CustomLLM("")
 res: LLMResponse = llm.invoke("text")
 print(res.content)
 
-# If rate_limit_handler and async_rate_limit_handler decorators are used and you want to use a custom rate limit handler
+# If  you want to use a custom rate limit handler
 # Type variables for function signatures used in rate limit handlers
 F = TypeVar("F", bound=Callable[..., Any])
 AF = TypeVar("AF", bound=Callable[..., Awaitable[Any]])

examples/customize/llms/custom_llm.py

@@ -0,0 +1,32 @@
+from neo4j_graphrag.llm import MistralAILLM, LLMResponse
+from neo4j_graphrag.message_history import InMemoryMessageHistory
+from neo4j_graphrag.types import LLMMessage
+
+# set api key here on in the MISTRAL_API_KEY env var
+api_key = None
+
+
+messages: list[LLMMessage] = [
+    {
+        "role": "system",
+        "content": "You are a seasoned actor and expert performer, renowned for your one-man shows and comedic talent.",
+    },
+    {
+        "role": "user",
+        "content": "say something",
+    },
+]
+
+
+llm = MistralAILLM(
+    model_name="mistral-small-latest",
+    api_key=api_key,
+)
+res: LLMResponse = llm.invoke(
+    # "say something",
+    # messages,
+    InMemoryMessageHistory(
+        messages=messages,
+    )
+)
+print(res.content)

examples/customize/llms/mistalai_llm.py

@@ -3,11 +3,26 @@
 """
 
 from neo4j_graphrag.llm import LLMResponse, OllamaLLM
+from neo4j_graphrag.types import LLMMessage
+
+messages: list[LLMMessage] = [
+    {
+        "role": "system",
+        "content": "You are a seasoned actor and expert performer, renowned for your one-man shows and comedic talent.",
+    },
+    {
+        "role": "user",
+        "content": "say something",
+    },
+]
+
 
 llm = OllamaLLM(
-    model_name="<model_name>",
+    model_name="orca-mini:latest",
     # model_params={"options": {"temperature": 0}, "format": "json"},
     # host="...",  # if using a remote server
 )
-res: LLMResponse = llm.invoke("What is the additive color model?")
+res: LLMResponse = llm.invoke(
+    messages,
+)
 print(res.content)

examples/customize/llms/mistralai_llm.py

@@ -1,8 +1,28 @@
 from neo4j_graphrag.llm import LLMResponse, OpenAILLM
+from neo4j_graphrag.message_history import InMemoryMessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 # set api key here on in the OPENAI_API_KEY env var
 api_key = None
 
+messages: list[LLMMessage] = [
+    {
+        "role": "system",
+        "content": "You are a seasoned actor and expert performer, renowned for your one-man shows and comedic talent.",
+    },
+    {
+        "role": "user",
+        "content": "say something",
+    },
+]
+
+
 llm = OpenAILLM(model_name="gpt-4o", api_key=api_key)
-res: LLMResponse = llm.invoke("say something")
+res: LLMResponse = llm.invoke(
+    # "say something",
+    # messages,
+    InMemoryMessageHistory(
+        messages=messages,
+    )
+)
 print(res.content)