Shared response buffer backend

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,