Shared response buffer backend

sia/llm_engine/__init__.py

@@ -1,15 +1,15 @@
-from typing import Callable, Iterator
-from abc import ABC, abstractmethod
-
-class LlmEngine(ABC):
-    @abstractmethod
-    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
-        pass
-
-    @abstractmethod
-    def token_count(self, system_prompt: str, main_context: str) -> int:
-        pass
-
-    @abstractmethod
-    def token_limit(self) -> int:
+from typing import Callable, Iterator
+from abc import ABC, abstractmethod
+
+class LlmEngine(ABC):
+    @abstractmethod
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        pass
+
+    @abstractmethod
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        pass
+
+    @abstractmethod
+    def token_limit(self) -> int:
         pass

sia/llm_engine/local_llm_engine.py

@@ -56,7 +56,7 @@ class LocalLlmEngine(LlmEngine):
         else:
             self._logits_processor = None
 
-    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
         """
         Run inference using the system prompt and main context.
 
@@ -70,10 +70,12 @@ class LocalLlmEngine(LlmEngine):
         """
         messages = [
             {"role": "system", "content": system_prompt},
-            {"role": "user", "content": main_context}
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
         ]
         prompt = self._tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
+            messages, tokenize=False,
+            add_generation_prompt=False,
         )
         streamer = TextIteratorStreamer(
             self._tokenizer,

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

sia/llm_engine/qwq_llm_engine.py

@@ -63,7 +63,7 @@ class QwQLlmEngine(LlmEngine):
             self._logits_processor = None
 
                         
-    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
         """
         Run inference using the system prompt and main context.
 
@@ -77,13 +77,14 @@ class QwQLlmEngine(LlmEngine):
         """
         messages = [
             {"role": "system", "content": system_prompt},
-            {"role": "user", "content": main_context}
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
         ]
 
         text = self._tokenizer.apply_chat_template(
             messages,
             tokenize=False,
-            add_generation_prompt=True,
+            add_generation_prompt=False,
         )
             
         streamer = TextIteratorStreamer(