Shared response buffer backend

2025-04-17 18:32:23 +02:00
parent 9477575421
commit 34fb5d814f
12 changed files with 941 additions and 869 deletions
--- a/sia/llm_engine/mistral_llm_engine.py
+++ b/sia/llm_engine/mistral_llm_engine.py
@@ -1,71 +1,83 @@
-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
-from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
-
-from . import LlmEngine
-
-class MistralLlmEngine(LlmEngine):
-    def __init__(
-            self,
-            model: str,
-            temperature: float,
-            token_limit: int,
-            api_key: str,
-        ):
-        self._model = model
-        self._temperature = temperature
-        self._token_limit = token_limit
-        self._api_key = api_key
-        self._client = Mistral(api_key=api_key)
-        self._tokenizer = MistralTokenizer.v3()
-
-    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
-        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,
-        )
-        
-        try:
-            for chunk in stream_response:
-                if should_stop():
-                    stream_response.response.close()
-                    break
-                if content := chunk.data.choices[0].delta.content:
-                    yield content
-        finally:
-            stream_response.response.close()
-
-    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
+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
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+
+from . import LlmEngine
+from ..util import skip_prefix
+
+class MistralLlmEngine(LlmEngine):
+    def __init__(
+            self,
+            model: str,
+            temperature: float,
+            token_limit: int,
+            api_key: str,
+        ):
+        self._model = model
+        self._temperature = temperature
+        self._token_limit = token_limit
+        self._api_key = api_key
+        self._client = Mistral(api_key=api_key)
+        self._tokenizer = MistralTokenizer.v3()
+
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        messages = [
+            {
+                "role": "system",
+                "content": system_prompt,
+            },
+            {
+                "role": "user",
+                "content": main_context,
+            },
+            {
+                "role": "assistant",
+                "content": continuation_text,
+                "prefix": True,
+            },
+        ] if continuation_text else [
+            {
+                "role": "system",
+                "content": system_prompt,
+            },
+            {
+                "role": "user",
+                "content": main_context,
+            },
+        ]
+        stream_response = self._client.chat.stream(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+        )
+        
+        try:
+            def content_generator():
+                for chunk in stream_response:
+                    if should_stop():
+                        stream_response.response.close()
+                        break
+                    if content := chunk.data.choices[0].delta.content:
+                        yield content
+            yield from skip_prefix(content_generator(), continuation_text)
+        finally:
+            stream_response.response.close()
+
+    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