Allow stopping of inference

sia/mistral_llm_engine.py

@@ -1,6 +1,4 @@
-from typing import Iterator, Optional
-from abc import ABC, abstractmethod
-import os
+from typing import Iterator, Optional, Callable
 from mistralai import Mistral
 from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
@@ -23,11 +21,7 @@ class MistralLlmEngine(LlmEngine):
         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),
-        ]
+    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
         messages = [
             {
                 "role": "system",
@@ -48,9 +42,16 @@ class MistralLlmEngine(LlmEngine):
             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
+        
+        try:
+            for chunk in stream_response:
+                if should_stop():
+                    stream_response.close()
+                    break
+                if content := chunk.data.choices[0].delta.content:
+                    yield content
+        finally:
+            stream_response.close()
 
     def token_count(self, system_prompt: str, main_context: str) -> int:
         messages = [
@@ -67,4 +68,4 @@ class MistralLlmEngine(LlmEngine):
         return len(tokenized.tokens)
 
     def token_limit(self) -> int:
-        return self._token_limit
+        return self._token_limit