Update LLMInterface to restore LC compatibility

stellasia · stellasia · commit 7b69305e11ac · 2025-08-25T14:54:54.000+02:00
diff --git a/examples/customize/llms/openai_llm.py b/examples/customize/llms/openai_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)
diff --git a/src/neo4j_graphrag/llm/base.py b/src/neo4j_graphrag/llm/base.py
@@ -19,15 +19,18 @@
 
 from neo4j_graphrag.message_history import MessageHistory
 from neo4j_graphrag.types import LLMMessage
+from .rate_limit import rate_limit_handler
 
 from .types import LLMResponse, ToolCallResponse
 from .rate_limit import (
     DEFAULT_RATE_LIMIT_HANDLER,
+    async_rate_limit_handler,
 )
 
 from neo4j_graphrag.tool import Tool
 
 from .rate_limit import RateLimitHandler
+from .utils import legacy_inputs_to_message_history
 
 
 class LLMInterface(ABC):
@@ -55,20 +58,27 @@ def __init__(
         else:
             self._rate_limit_handler = DEFAULT_RATE_LIMIT_HANDLER
 
-    @abstractmethod
+    @rate_limit_handler
     def invoke(
         self,
-        input: str,
+        input: Union[str, List[LLMMessage], MessageHistory],
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
+        message_history = legacy_inputs_to_message_history(
+            input, message_history, system_instruction
+        )
+        return self._invoke(message_history.messages)
+
+    @abstractmethod
+    def _invoke(
+        self,
+        input: list[LLMMessage],
     ) -> LLMResponse:
         """Sends a text input to the LLM and retrieves a response.
 
         Args:
-            input (str): Text sent to the LLM.
-            message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
-                with each message having a specific role assigned.
-            system_instruction (Optional[str]): An option to override the llm system message for this invocation.
+            input (MessageHistory): Text sent to the LLM.
 
         Returns:
             LLMResponse: The response from the LLM.
@@ -77,20 +87,27 @@ def invoke(
             LLMGenerationError: If anything goes wrong.
         """
 
-    @abstractmethod
+    @async_rate_limit_handler
     async def ainvoke(
         self,
-        input: str,
+        input: Union[str, List[LLMMessage], MessageHistory],
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
+        message_history = legacy_inputs_to_message_history(
+            input, message_history, system_instruction
+        )
+        return await self._ainvoke(message_history.messages)
+
+    @abstractmethod
+    async def _ainvoke(
+        self,
+        input: list[LLMMessage],
     ) -> LLMResponse:
         """Asynchronously sends a text input to the LLM and retrieves a response.
 
         Args:
             input (str): Text sent to the LLM.
-            message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
-                with each message having a specific role assigned.
-            system_instruction (Optional[str]): An option to override the llm system message for this invocation.
 
         Returns:
             LLMResponse: The response from the LLM.
diff --git a/src/neo4j_graphrag/llm/openai_llm.py b/src/neo4j_graphrag/llm/openai_llm.py
@@ -28,37 +28,33 @@
     cast,
 )
 
-from pydantic import ValidationError
-
 from neo4j_graphrag.message_history import MessageHistory
 from neo4j_graphrag.types import LLMMessage
 
 from ..exceptions import LLMGenerationError
 from .base import LLMInterface
-from .rate_limit import RateLimitHandler, rate_limit_handler, async_rate_limit_handler
 from .types import (
-    BaseMessage,
     LLMResponse,
-    MessageList,
     ToolCall,
     ToolCallResponse,
-    SystemMessage,
-    UserMessage,
 )
 
 from neo4j_graphrag.tool import Tool
 
 if TYPE_CHECKING:
     from openai.types.chat import (
         ChatCompletionMessageParam,
-        ChatCompletionToolParam,
-    )
+        ChatCompletionToolParam, ChatCompletionUserMessageParam,
+        ChatCompletionSystemMessageParam, ChatCompletionAssistantMessageParam,
+)
     from openai import OpenAI, AsyncOpenAI
+    from .rate_limit import RateLimitHandler
 else:
     ChatCompletionMessageParam = Any
     ChatCompletionToolParam = Any
     OpenAI = Any
     AsyncOpenAI = Any
+    RateLimitHandler = Any
 
 
 class BaseOpenAILLM(LLMInterface, abc.ABC):
@@ -93,23 +89,26 @@ def __init__(
 
     def get_messages(
         self,
-        input: str,
-        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
-        system_instruction: Optional[str] = None,
+        messages: list[LLMMessage],
     ) -> Iterable[ChatCompletionMessageParam]:
-        messages = []
-        if system_instruction:
-            messages.append(SystemMessage(content=system_instruction).model_dump())
-        if message_history:
-            if isinstance(message_history, MessageHistory):
-                message_history = message_history.messages
-            try:
-                MessageList(messages=cast(list[BaseMessage], message_history))
-            except ValidationError as e:
-                raise LLMGenerationError(e.errors()) from e
-            messages.extend(cast(Iterable[dict[str, Any]], message_history))
-        messages.append(UserMessage(content=input).model_dump())
-        return messages  # type: ignore
+        chat_messages = []
+        for m in messages:
+            message_type: ChatCompletionMessageParam
+            if m["role"] == "system":
+                message_type = ChatCompletionSystemMessageParam
+            elif m["role"] == "user":
+                message_type = ChatCompletionUserMessageParam
+            elif m["role"] == "assistant":
+                message_type = ChatCompletionAssistantMessageParam
+            else:
+                raise ValueError(f"Unknown message type: {m['role']}")
+            chat_messages.append(
+                message_type(
+                    role=m["role"],
+                    content=m["content"],
+                )
+            )
+        return chat_messages
 
     def _convert_tool_to_openai_format(self, tool: Tool) -> Dict[str, Any]:
         """Convert a Tool object to OpenAI's expected format.
@@ -132,21 +131,15 @@ def _convert_tool_to_openai_format(self, tool: Tool) -> Dict[str, Any]:
         except AttributeError:
             raise LLMGenerationError(f"Tool {tool} is not a valid Tool object")
 
-    @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:
         """Sends a text input to the OpenAI chat completion model
         and returns the response's content.
 
         Args:
             input (str): Text sent to the LLM.
-            message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
-                with each message having a specific role assigned.
-            system_instruction (Optional[str]): An option to override the llm system message for this invocation.
 
         Returns:
             LLMResponse: The response from OpenAI.
@@ -155,10 +148,8 @@ def invoke(
             LLMGenerationError: If anything goes wrong.
         """
         try:
-            if isinstance(message_history, MessageHistory):
-                message_history = message_history.messages
             response = self.client.chat.completions.create(
-                messages=self.get_messages(input, message_history, system_instruction),
+                messages=self.get_messages(input),
                 model=self.model_name,
                 **self.model_params,
             )
@@ -167,7 +158,6 @@ def invoke(
         except self.openai.OpenAIError as e:
             raise LLMGenerationError(e)
 
-    @rate_limit_handler
     def invoke_with_tools(
         self,
         input: str,
@@ -242,21 +232,15 @@ def invoke_with_tools(
         except self.openai.OpenAIError as e:
             raise LLMGenerationError(e)
 
-    @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:
         """Asynchronously sends a text input to the OpenAI chat
         completion model and returns the response's content.
 
         Args:
             input (str): Text sent to the LLM.
-            message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
-                with each message having a specific role assigned.
-            system_instruction (Optional[str]): An option to override the llm system message for this invocation.
 
         Returns:
             LLMResponse: The response from OpenAI.
@@ -265,10 +249,8 @@ async def ainvoke(
             LLMGenerationError: If anything goes wrong.
         """
         try:
-            if isinstance(message_history, MessageHistory):
-                message_history = message_history.messages
             response = await self.async_client.chat.completions.create(
-                messages=self.get_messages(input, message_history, system_instruction),
+                messages=self.get_messages(input),
                 model=self.model_name,
                 **self.model_params,
             )
@@ -277,7 +259,6 @@ async def ainvoke(
         except self.openai.OpenAIError as e:
             raise LLMGenerationError(e)
 
-    @async_rate_limit_handler
     async def ainvoke_with_tools(
         self,
         input: str,
diff --git a/src/neo4j_graphrag/message_history.py b/src/neo4j_graphrag/message_history.py
@@ -74,6 +74,9 @@ class MessageHistory(ABC):
     @abstractmethod
     def messages(self) -> List[LLMMessage]: ...
 
+    def is_empty(self) -> bool:
+        return len(self.messages) == 0
+
     @abstractmethod
     def add_message(self, message: LLMMessage) -> None: ...
 
