Mistral

Author: leila-messallem

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -16,20 +16,15 @@
 
 import os
 from typing import Any, Optional, Union
+from pydantic import ValidationError
 
-from ..exceptions import LLMGenerationError
-from .base import LLMInterface
-from .types import LLMResponse
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm.base import LLMInterface
+from neo4j_graphrag.llm.types import LLMResponse, MessageList, SystemMessage, UserMessage
 
 try:
-    from mistralai import Mistral
-    from mistralai.models.assistantmessage import AssistantMessage
+    from mistralai import Mistral, Messages
     from mistralai.models.sdkerror import SDKError
-    from mistralai.models.systemmessage import SystemMessage
-    from mistralai.models.toolmessage import ToolMessage
-    from mistralai.models.usermessage import UserMessage
-
-    MessageType = Union[AssistantMessage, SystemMessage, ToolMessage, UserMessage]
 except ImportError:
     Mistral = None  # type: ignore
     SDKError = None  # type: ignore
@@ -63,9 +58,19 @@ def __init__(
         if api_key is None:
             api_key = os.getenv("MISTRAL_API_KEY", "")
         self.client = Mistral(api_key=api_key, **kwargs)
-
-    def get_messages(self, input: str) -> list[MessageType]:
-        return [UserMessage(content=input)]
+    
+    def get_messages(self, input: str, chat_history: list) -> list[Messages]:
+        messages = []
+        if self.system_instruction:
+            messages.append(SystemMessage(content=self.system_instruction).model_dump())
+        if chat_history:
+            try:
+                MessageList(messages=chat_history)
+            except ValidationError as e:
+                raise LLMGenerationError(e.errors()) from e
+            messages.extend(chat_history)
+        messages.append(UserMessage(content=input).model_dump())
+        return messages
 
     def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None) -> LLMResponse:
         """Sends a text input to the Mistral chat completion model
@@ -82,9 +87,10 @@ def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None
             LLMGenerationError: If anything goes wrong.
         """
         try:
+            messages = self.get_messages(input, chat_history)
             response = self.client.chat.complete(
                 model=self.model_name,
-                messages=self.get_messages(input),
+                messages=messages,
                 **self.model_params,
             )
             if response is None or response.choices is None or not response.choices:
@@ -112,9 +118,10 @@ async def ainvoke(
             LLMGenerationError: If anything goes wrong.
         """
         try:
+            messages = self.get_messages(input, chat_history)
             response = await self.client.chat.complete_async(
                 model=self.model_name,
-                messages=self.get_messages(input),
+                messages=messages,
                 **self.model_params,
             )
             if response is None or response.choices is None or not response.choices:

tests/unit/llm/test_mistralaillm.py tests/unit/llm/test_mistralai_llm.pytests/unit/llm/test_mistralaillm.py renamed to tests/unit/llm/test_mistralai_llm.py

@@ -22,13 +22,13 @@
 
 
 @patch("neo4j_graphrag.llm.mistralai_llm.Mistral", None)
-def test_mistral_ai_llm_missing_dependency() -> None:
+def test_mistralai_llm_missing_dependency() -> None:
     with pytest.raises(ImportError):
         MistralAILLM(model_name="mistral-model")
 
 
 @patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
-def test_mistral_ai_llm_invoke(mock_mistral: Mock) -> None:
+def test_mistralai_llm_invoke(mock_mistral: Mock) -> None:
     mock_mistral_instance = mock_mistral.return_value
 
     chat_response_mock = MagicMock()
@@ -46,9 +46,64 @@ def test_mistral_ai_llm_invoke(mock_mistral: Mock) -> None:
     assert res.content == "mistral response"
 
 
+@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
+def test_mistralai_llm_invoke_with_chat_history(mock_mistral: Mock) -> None:
+    mock_mistral_instance = mock_mistral.return_value
+    chat_response_mock = MagicMock()
+    chat_response_mock.choices = [
+        MagicMock(message=MagicMock(content="mistral response"))
+    ]
+    mock_mistral_instance.chat.complete.return_value = chat_response_mock
+    model = "mistral-model"
+    system_instruction = "You are a helpful assistant."
+
+    llm = MistralAILLM(model_name=model, system_instruction=system_instruction)
+
+    chat_history = [
+        {"role": "user", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "content": "Usually around 6am."},
+    ]
+    question = "What about next season?"
+    res = llm.invoke(question, chat_history)
+
+    assert isinstance(res, LLMResponse)
+    assert res.content == "mistral response"
+    messages = [{"role": "system", "content": system_instruction}]
+    messages.extend(chat_history)
+    messages.append({"role": "user", "content": question})
+    llm.client.chat.complete.assert_called_once_with(
+        messages=messages,
+        model=model,
+    )
+
+
+@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
+def test_mistralai_llm_invoke_with_chat_history_validation_error(mock_mistral: Mock) -> None:
+    mock_mistral_instance = mock_mistral.return_value
+    chat_response_mock = MagicMock()
+    chat_response_mock.choices = [
+        MagicMock(message=MagicMock(content="mistral response"))
+    ]
+    mock_mistral_instance.chat.complete.return_value = chat_response_mock
+    model = "mistral-model"
+    system_instruction = "You are a helpful assistant."
+
+    llm = MistralAILLM(model_name=model, system_instruction=system_instruction)
+
+    chat_history = [
+        {"role": "user", "content": "When does the sun come up in the summer?"},
+        {"role": "monkey", "content": "Usually around 6am."},
+    ]
+    question = "What about next season?"
+
+    with pytest.raises(LLMGenerationError) as exc_info:
+        llm.invoke(question, chat_history)
+    assert "Input should be 'user' or 'assistant'" in str(exc_info.value)
+
+
 @pytest.mark.asyncio
 @patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
-async def test_mistral_ai_llm_ainvoke(mock_mistral: Mock) -> None:
+async def test_mistralai_llm_ainvoke(mock_mistral: Mock) -> None:
     mock_mistral_instance = mock_mistral.return_value
 
     async def mock_complete_async(*args: Any, **kwargs: Any) -> MagicMock:
@@ -69,7 +124,7 @@ async def mock_complete_async(*args: Any, **kwargs: Any) -> MagicMock:
 
 
 @patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
-def test_mistral_ai_llm_invoke_sdkerror(mock_mistral: Mock) -> None:
+def test_mistralai_llm_invoke_sdkerror(mock_mistral: Mock) -> None:
     mock_mistral_instance = mock_mistral.return_value
     mock_mistral_instance.chat.complete.side_effect = SDKError("Some error")
 
@@ -81,7 +136,7 @@ def test_mistral_ai_llm_invoke_sdkerror(mock_mistral: Mock) -> None:
 
 @pytest.mark.asyncio
 @patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
-async def test_mistral_ai_llm_ainvoke_sdkerror(mock_mistral: Mock) -> None:
+async def test_mistralai_llm_ainvoke_sdkerror(mock_mistral: Mock) -> None:
     mock_mistral_instance = mock_mistral.return_value
 
     async def mock_complete_async(*args: Any, **kwargs: Any) -> None: