llm interface v2 supported in mistralai

Author: serengil

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -15,12 +15,12 @@
 from __future__ import annotations
 
 import os
-from typing import Any, Iterable, List, Optional, Union, cast
+from typing import Any, Iterable, List, Optional, Union, cast, overload
 
 from pydantic import ValidationError
 
 from neo4j_graphrag.exceptions import LLMGenerationError
-from neo4j_graphrag.llm.base import LLMInterface
+from neo4j_graphrag.llm.base import LLMInterface, LLMInterfaceV2
 from neo4j_graphrag.utils.rate_limit import (
     RateLimitHandler,
     rate_limit_handler,
@@ -37,14 +37,22 @@
 from neo4j_graphrag.types import LLMMessage
 
 try:
-    from mistralai import Messages, Mistral
+    from mistralai import (
+        Messages,
+        UserMessage,
+        AssistantMessage,
+        SystemMessage,
+        Mistral,
+    )
     from mistralai.models.sdkerror import SDKError
 except ImportError:
     Mistral = None  # type: ignore
     SDKError = None  # type: ignore
 
 
+# pylint: disable=redefined-builtin, arguments-differ, raise-missing-from, no-else-return
 class MistralAILLM(LLMInterface):
+
     def __init__(
         self,
         model_name: str,
@@ -73,28 +81,67 @@ def __init__(
             api_key = os.getenv("MISTRAL_API_KEY", "")
         self.client = Mistral(api_key=api_key, **kwargs)
 
-    def get_messages(
+    # overloads for LLMInterface and LLMInterfaceV2 methods
+    @overload
+    def invoke(
         self,
         input: str,
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         system_instruction: Optional[str] = None,
-    ) -> list[Messages]:
-        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 cast(list[Messages], messages)
+    ) -> LLMResponse: ...
 
-    @rate_limit_handler
+    @overload
     def invoke(
+        self,
+        input: List[LLMMessage],
+    ) -> LLMResponse: ...
+
+    @overload
+    async def ainvoke(
+        self,
+        input: str,
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse: ...
+
+    @overload
+    async def ainvoke(
+        self,
+        input: List[LLMMessage],
+    ) -> LLMResponse: ...
+
+    # switching logics to LLMInterface or LLMInterfaceV2
+    def invoke(
+        self,
+        input: Union[str, List[LLMMessage]],
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
+        if isinstance(input, str):
+            return self.__legacy_invoke(input, message_history, system_instruction)
+        elif isinstance(input, list):
+            return self.__brand_new_invoke(input)
+        else:
+            raise ValueError(f"Invalid input type for invoke method - {type(input)}")
+
+    async def ainvoke(
+        self,
+        input: Union[str, List[LLMMessage]],
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
+        if isinstance(input, str):
+            return await self.__legacy_ainvoke(
+                input, message_history, system_instruction
+            )
+        elif isinstance(input, list):
+            return await self.__brand_new_ainvoke(input)
+        else:
+            raise ValueError(f"Invalid input type for ainvoke method - {type(input)}")
+
+    # implementations
+    @rate_limit_handler
+    def __legacy_invoke(
         self,
         input: str,
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
@@ -132,8 +179,40 @@ def invoke(
         except SDKError as e:
             raise LLMGenerationError(e)
 
+    def __brand_new_invoke(
+        self,
+        input: List[LLMMessage],
+    ) -> LLMResponse:
+        """Sends a text input to the Mistral chat completion model
+        and returns the response's content.
+
+        Args:
+            input (str): Text sent to the LLM.
+
+        Returns:
+            LLMResponse: The response from MistralAI.
+
+        Raises:
+            LLMGenerationError: If anything goes wrong.
+        """
+        try:
+            messages = self.get_brand_new_messages(input)
+            response = self.client.chat.complete(
+                model=self.model_name,
+                messages=messages,
+                **self.model_params,
+            )
+            content: str = ""
+            if response and response.choices:
+                possible_content = response.choices[0].message.content
+                if isinstance(possible_content, str):
+                    content = possible_content
+            return LLMResponse(content=content)
+        except SDKError as e:
+            raise LLMGenerationError(e)
+
     @async_rate_limit_handler
-    async def ainvoke(
+    async def __legacy_ainvoke(
         self,
         input: str,
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
@@ -171,3 +250,76 @@ async def ainvoke(
             return LLMResponse(content=content)
         except SDKError as e:
             raise LLMGenerationError(e)
+
+    async def __brand_new_ainvoke(
+        self,
+        input: List[LLMMessage],
+    ) -> LLMResponse:
+        """Asynchronously sends a text input to the MistralAI chat
+        completion model and returns the response's content.
+
+        Args:
+            input (str): Text sent to the LLM.
+
+        Returns:
+            LLMResponse: The response from MistralAI.
+
+        Raises:
+            LLMGenerationError: If anything goes wrong.
+        """
+        try:
+            messages = self.get_brand_new_messages(input)
+            response = await self.client.chat.complete_async(
+                model=self.model_name,
+                messages=messages,
+                **self.model_params,
+            )
+            content: str = ""
+            if response and response.choices:
+                possible_content = response.choices[0].message.content
+                if isinstance(possible_content, str):
+                    content = possible_content
+            return LLMResponse(content=content)
+        except SDKError as e:
+            raise LLMGenerationError(e)
+
+    # subsidiary methods
+    def get_messages(
+        self,
+        input: str,
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> list[Messages]:
+        """Constructs the message list for the Mistral chat completion model."""
+        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 cast(list[Messages], messages)
+
+    def get_brand_new_messages(
+        self,
+        input: list[LLMMessage],
+    ) -> list[Messages]:
+        """Constructs the message list for the Mistral chat completion model."""
+        messages: list[Messages] = []
+        for m in input:
+            if m["role"] == "system":
+                messages.append(SystemMessage(content=m["content"]))
+                continue
+            if m["role"] == "user":
+                messages.append(UserMessage(content=m["content"]))
+                continue
+            if m["role"] == "assistant":
+                messages.append(AssistantMessage(content=m["content"]))
+                continue
+            raise ValueError(f"Unknown role: {m['role']}")
+        return messages