Merge pull request #2025 from Giskard-AI/dep/upgrade-mistralai

Upgraded `mistralai` dep to >= 1

Author: kevinmessiaen

giskard/llm/client/mistral.py

@@ -1,5 +1,6 @@
 from typing import Optional, Sequence
 
+import os
 from dataclasses import asdict
 from logging import warning
 
@@ -9,18 +10,17 @@
 from .base import ChatMessage
 
 try:
-    from mistralai.client import MistralClient as _MistralClient
-    from mistralai.models.chat_completion import ChatMessage as MistralChatMessage
+    from mistralai import Mistral
 except ImportError as err:
     raise LLMImportError(
         flavor="llm", msg="To use Mistral models, please install the `mistralai` package with `pip install mistralai`"
     ) from err
 
 
 class MistralClient(LLMClient):
-    def __init__(self, model: str = "mistral-large-latest", client: _MistralClient = None):
+    def __init__(self, model: str = "mistral-large-latest", client: Mistral = None):
         self.model = model
-        self._client = client or _MistralClient()
+        self._client = client or Mistral(api_key=os.getenv("MISTRAL_API_KEY", ""))
 
     def complete(
         self,
@@ -43,9 +43,9 @@ def complete(
             extra_params["response_format"] = {"type": "json_object"}
 
         try:
-            completion = self._client.chat(
+            completion = self._client.chat.complete(
                 model=self.model,
-                messages=[MistralChatMessage(**asdict(m)) for m in messages],
+                messages=[asdict(m) for m in messages],
                 temperature=temperature,
                 max_tokens=max_tokens,
                 **extra_params,

pdm.lock

@@ -56,7 +56,7 @@ dev = [
     "pytest-asyncio>=0.21.1",
     "pydantic>=2",
     "avidtools",
-    "mistralai>=0.1.8, <1",
+    "mistralai>=1",
     "boto3>=1.34.88",
     "scikit-learn==1.4.2",
 ]

pyproject.toml

@@ -4,9 +4,6 @@
 import pydantic
 import pytest
 from google.generativeai.types import ContentDict
-from mistralai.models.chat_completion import ChatCompletionResponse, ChatCompletionResponseChoice
-from mistralai.models.chat_completion import ChatMessage as MistralChatMessage
-from mistralai.models.chat_completion import FinishReason, UsageInfo
 from openai.types import CompletionUsage
 from openai.types.chat import ChatCompletion, ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
@@ -35,22 +32,6 @@
 )
 
 
-DEMO_MISTRAL_RESPONSE = ChatCompletionResponse(
-    id="2d62260a7a354e02922a4f6ad36930d3",
-    object="chat.completion",
-    created=1630000000,
-    model="mistral-large",
-    choices=[
-        ChatCompletionResponseChoice(
-            index=0,
-            message=MistralChatMessage(role="assistant", content="This is a test!", name=None, tool_calls=None),
-            finish_reason=FinishReason.stop,
-        )
-    ],
-    usage=UsageInfo(prompt_tokens=9, total_tokens=89, completion_tokens=80),
-)
-
-
 def test_llm_complete_message():
     client = Mock()
     client.chat.completions.create.return_value = DEMO_OPENAI_RESPONSE
@@ -69,19 +50,36 @@ def test_llm_complete_message():
 
 @pytest.mark.skipif(not PYDANTIC_V2, reason="Mistral raise an error with pydantic < 2")
 def test_mistral_client():
+    from mistralai.models import ChatCompletionChoice, ChatCompletionResponse, UsageInfo
+
+    demo_response = ChatCompletionResponse(
+        id="2d62260a7a354e02922a4f6ad36930d3",
+        object="chat.completion",
+        created=1630000000,
+        model="mistral-large",
+        choices=[
+            ChatCompletionChoice(
+                index=0,
+                message={"role": "assistant", "content": "This is a test!"},
+                finish_reason="stop",
+            )
+        ],
+        usage=UsageInfo(prompt_tokens=9, total_tokens=89, completion_tokens=80),
+    )
+
     client = Mock()
-    client.chat.return_value = DEMO_MISTRAL_RESPONSE
+    client.chat.complete.return_value = demo_response
 
     from giskard.llm.client.mistral import MistralClient
 
     res = MistralClient(model="mistral-large", client=client).complete(
         [ChatMessage(role="user", content="Hello")], temperature=0.11, max_tokens=12
     )
 
-    client.chat.assert_called_once()
-    assert client.chat.call_args[1]["messages"] == [MistralChatMessage(role="user", content="Hello")]
-    assert client.chat.call_args[1]["temperature"] == 0.11
-    assert client.chat.call_args[1]["max_tokens"] == 12
+    client.chat.complete.assert_called_once()
+    assert client.chat.complete.call_args[1]["messages"] == [{"role": "user", "content": "Hello"}]
+    assert client.chat.complete.call_args[1]["temperature"] == 0.11
+    assert client.chat.complete.call_args[1]["max_tokens"] == 12
 
     assert isinstance(res, ChatMessage)
     assert res.content == "This is a test!"