From ef94d0d1a0496ac02b451a340906d8ebd2e19175 Mon Sep 17 00:00:00 2001
From: geens <niels.geens@gmail.com>
Date: Sun, 17 Nov 2024 16:53:44 +0100
Subject: [PATCH] Added mistral llm engine

---
 install.sh                |  2 +-
 requirements.txt          |  2 ++
 sia/__main__.py           | 10 +++++-
 sia/local_llm_engine.py   |  2 +-
 sia/mistral_llm_engine.py | 70 +++++++++++++++++++++++++++++++++++++++
 sia/openai_llm_engine.py  |  2 +-
 6 files changed, 84 insertions(+), 4 deletions(-)
 create mode 100644 sia/mistral_llm_engine.py

diff --git a/install.sh b/install.sh
index 4c198ce..1abdab7 100644
--- a/install.sh
+++ b/install.sh
@@ -20,4 +20,4 @@ vim .env
 cd /root/SIA
 python3 -m sia
 
-#The SIA source is located in /SIA. Not all features are implemented yet. Look at the readme and code to find what is missing. Make sure to unit test your work.
\ No newline at end of file
+#The SIA source is located in /root/sia. Not all features are implemented yet. Look at the readme and code to find what is missing. Make sure to unit test your work.
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index c39ba5f..ca87b53 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,8 @@
 accelerate
 aiohttp
 bs4
+mistral_common
+mistralai
 openai
 python-dotenv
 tiktoken
diff --git a/sia/__main__.py b/sia/__main__.py
index b667b2b..6db2e3a 100644
--- a/sia/__main__.py
+++ b/sia/__main__.py
@@ -9,6 +9,7 @@ from .config import Config
 from .hf_llm_engine import HfLlmEngine
 from .llm_engine import LlmEngine
 from .local_llm_engine import LocalLlmEngine
+from .mistral_llm_engine import MistralLlmEngine
 from .openai_llm_engine import OpenAILlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
@@ -39,8 +40,8 @@ class Main:
                 self._llm = LocalLlmEngine(
                     self._config.model,
                     self._config.temperature,
+                    self._config.token_limit,
                     self._config.api_token,
-                    self._config.token_limit
                 )
             case "hf":
                 self._llm = HfLlmEngine(
@@ -50,6 +51,13 @@ class Main:
                 )
             case "openai":
                 self._llm = OpenAILlmEngine(
+                    self._config.model,
+                    self._config.temperature,
+                    self._config.token_limit,
+                    self._config.api_token,
+                )
+            case "mistral":
+                self._llm = MistralLlmEngine(
                     self._config.model,
                     self._config.temperature,
                     self._config.api_token,
diff --git a/sia/local_llm_engine.py b/sia/local_llm_engine.py
index 8aa9d9c..28e7469 100644
--- a/sia/local_llm_engine.py
+++ b/sia/local_llm_engine.py
@@ -26,7 +26,7 @@ class LocalLlmEngine(LlmEngine):
         """
         self._temperature = temperature
         self._token_limit = token_limit
-        self._tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
         model = AutoModelForCausalLM.from_pretrained(
             model_path,
             return_dict=True,
diff --git a/sia/mistral_llm_engine.py b/sia/mistral_llm_engine.py
new file mode 100644
index 0000000..d054f8d
--- /dev/null
+++ b/sia/mistral_llm_engine.py
@@ -0,0 +1,70 @@
+from typing import Iterator, Optional
+from abc import ABC, abstractmethod
+import os
+from mistralai import Mistral
+from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+
+from .llm_engine import LlmEngine
+
+class MistralLlmEngine(LlmEngine):
+    def __init__(
+            self,
+            model: str,
+            temperature: float,
+            api_key: str,
+            token_limit: int
+        ):
+        self._model = model
+        self._temperature = temperature
+        self._api_key = api_key
+        self._token_limit = token_limit
+        self._client = Mistral(api_key=api_key)
+        self._tokenizer = MistralTokenizer.v3()
+
+    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
+        messages = [
+            SystemMessage(content=system_prompt),
+            UserMessage(content=main_context),
+        ]
+        messages = [
+            {
+                "role": "system",
+                "content": system_prompt,
+            },
+            {
+                "role": "user",
+                "content": main_context,
+            },
+            {
+                "role": "assistant",
+                "content": "<",
+                "prefix": True,
+            },
+        ]
+        stream_response = self._client.chat.stream(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+        )
+        for chunk in stream_response:
+            if content := chunk.data.choices[0].delta.content:
+                yield chunk.data.choices[0].delta.content
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        messages = [
+            SystemMessage(content=system_prompt),
+            UserMessage(content=main_context),
+        ]
+
+        tokenized = self._tokenizer.encode_chat_completion(
+            ChatCompletionRequest(
+                messages=messages,
+                model=self._model
+            )
+        )
+        return len(tokenized.tokens)
+
+    def token_limit(self) -> int:
+        return self._token_limit
\ No newline at end of file
diff --git a/sia/openai_llm_engine.py b/sia/openai_llm_engine.py
index 43464a7..ad161a2 100644
--- a/sia/openai_llm_engine.py
+++ b/sia/openai_llm_engine.py
@@ -14,8 +14,8 @@ class OpenAILlmEngine(LlmEngine):
         self,
         model: str,
         temperature: float,
+        token_limit: int,
         api_key: str,
-        token_limit: int = 0,
     ):
         """
         Initialize the OpenAI LLM Engine.