Support for custom LLM clients

giskard/llm/client/__init__.py

@@ -3,7 +3,7 @@
 import logging
 import os
 
-from .base import LLMClient, LLMFunctionCall, LLMMessage, LLMToolCall
+from .base import ChatMessage, LLMClient
 from .logger import LLMLogger
 
 _default_client = None
@@ -58,7 +58,7 @@ def get_default_client() -> LLMClient:
         return _default_client
 
     # Setup the default client
-    from .openai import LegacyOpenAIClient, OpenAIClient
+    from .openai import OpenAIClient
 
     default_llm_api = get_default_llm_api()
 
@@ -68,21 +68,16 @@ def get_default_client() -> LLMClient:
 
         client = AzureOpenAI() if default_llm_api == "azure" else OpenAI()
 
-        _default_client = OpenAIClient(_default_llm_model, client)
+        _default_client = OpenAIClient(model=_default_llm_model, client=client)
     except ImportError:
-        # Fallback for openai<=0.28.1
-        if default_llm_api != "openai":
-            raise ValueError(f"LLM scan using {default_llm_api.name} require openai>=1.0.0")
-        _default_client = LegacyOpenAIClient(_default_llm_model)
+        raise ValueError(f"LLM scan using {default_llm_api.name} require openai>=1.0.0")
 
     return _default_client
 
 
 __all__ = [
     "LLMClient",
-    "LLMFunctionCall",
-    "LLMToolCall",
-    "LLMMessage",
+    "ChatMessage",
     "LLMLogger",
     "get_default_client",
     "set_llm_model",

giskard/llm/client/base.py

@@ -1,59 +1,33 @@
-from typing import Any, Dict, List, Optional, Sequence
+from typing import Optional, Sequence
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
-import numpy as np
-
 from .logger import LLMLogger
 
 
 @dataclass
-class LLMFunctionCall:
-    name: str
-    arguments: Any
-
-
-@dataclass
-class LLMToolCall:
-    id: str
-    type: str
-    function: LLMFunctionCall
-
-
-@dataclass
-class LLMMessage:
+class ChatMessage:
     role: str
     content: Optional[str] = None
-    function_call: Optional[LLMFunctionCall] = None
-    tool_calls: Optional[List[LLMToolCall]] = None
 
-    @staticmethod
-    def create_message(role: str, content: str):
-        return LLMMessage(role=role, content=content, function_call=None, tool_calls=None)
+
+_logger = LLMLogger()
 
 
 class LLMClient(ABC):
     @property
-    @abstractmethod
     def logger(self) -> LLMLogger:
-        ...
+        return _logger
 
     @abstractmethod
     def complete(
         self,
-        messages: Sequence[LLMMessage],
-        functions=None,
-        temperature=0.5,
-        max_tokens=None,
-        function_call: Optional[Dict] = None,
+        messages: Sequence[ChatMessage],
+        temperature: float = 1,
+        max_tokens: Optional[int] = None,
         caller_id: Optional[str] = None,
-        tools=None,
-        tool_choice=None,
         seed: Optional[int] = None,
-    ) -> LLMMessage:
-        ...
-
-    @abstractmethod
-    def embeddings(self, text) -> np.ndarray:
+        format=None,
+    ) -> ChatMessage:
         ...

giskard/llm/client/mistral.py

@@ -0,0 +1,66 @@
+from typing import Optional, Sequence
+
+from dataclasses import asdict
+from logging import warning
+
+from ..config import LLMConfigurationError
+from ..errors import LLMImportError
+from . import LLMClient
+from .base import ChatMessage
+
+try:
+    from mistralai.client import MistralClient as _MistralClient
+    from mistralai.models.chat_completion import ChatMessage as MistralChatMessage
+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):
+        self.model = model
+        self._client = client or _MistralClient()
+
+    def complete(
+        self,
+        messages: Sequence[ChatMessage],
+        temperature: float = 1.0,
+        max_tokens: Optional[int] = None,
+        caller_id: Optional[str] = None,
+        seed: Optional[int] = None,
+        format: str = None,
+    ) -> ChatMessage:
+        extra_params = dict()
+        if seed is not None:
+            extra_params["random_seed"] = seed
+
+        if format not in (None, "json", "json_object") and "large" not in self.model:
+            warning(f"Unsupported format '{format}', ignoring.")
+            format = None
+
+        if format == "json" or format == "json_object":
+            extra_params["response_format"] = {"type": "json_object"}
+
+        try:
+            completion = self._client.chat(
+                model=self.model,
+                messages=[MistralChatMessage(**asdict(m)) for m in messages],
+                temperature=temperature,
+                max_tokens=max_tokens,
+                **extra_params,
+            )
+        except RuntimeError as err:
+            raise LLMConfigurationError("Could not get response from Mistral API") from err
+
+        self.logger.log_call(
+            prompt_tokens=completion.usage.prompt_tokens,
+            sampled_tokens=completion.usage.completion_tokens,
+            model=self.model,
+            client_class=self.__class__.__name__,
+            caller_id=caller_id,
+        )
+
+        msg = completion.choices[0].message
+
+        return ChatMessage(role=msg.role, content=msg.content)