Fixed auto approver and inference continuation

sia/llm_engine/hf_llm_engine.py

@@ -1,84 +1,87 @@
-from huggingface_hub import InferenceClient
-from transformers import AutoTokenizer, AutoConfig
-from typing import Iterator, Optional, Callable
-
-from . import LlmEngine
-
-class HfLlmEngine(LlmEngine):
-    """
-    LLM Engine implementation using HuggingFace's InferenceClient.
-    """
-    
-    def __init__(
-        self,
-        model: str,
-        temperature: float,
-        api_token: Optional[str],
-    ):
-        """
-        Initialize the HuggingFace Inference API LLM Engine.
-
-        Args:
-            model: HuggingFace model ID to use
-            temperature: Sampling temperature
-            api_token: HuggingFace API token
-        """
-        self._model = model
-        self._temperature = temperature
-
-        self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
-        self._config = AutoConfig.from_pretrained(model, token=api_token)
-        self._client = InferenceClient(token=api_token)
-                        
-    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
-        """
-        Run inference using the system prompt and main context.
-
-        Args:
-            system_prompt: The system prompt string
-            main_context: The main context string after templating
-            should_stop: Callback that returns True when inference should stop
-
-        Returns:
-            Iterator[str]: An iterator that yields the generated text.
-        """
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": main_context}
-        ]
-        
-        stream = self._client.chat_completion(
-            model=self._model,
-            messages=messages,
-            temperature=self._temperature,
-            stream=True
-        )
-        
-        try:
-            for response in stream:
-                if should_stop():
-                    stream.close()
-                    break
-                if content := response.choices[0].delta.content:
-                    yield content
-        finally:
-            stream.close()
-
-    def token_count(self, system_prompt: str, main_context: str) -> int:
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": main_context}
-        ]
-        prompt = self._tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        return len(self._tokenizer.encode(prompt))
-
-    def token_limit(self) -> int:
-        """
-        Get the model's context window size.
-        
-        Returns:
-            int: Maximum number of tokens the model can process
-        """
-        return self._config.max_position_embeddings
+from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoConfig
+from typing import Iterator, Optional, Callable
+
+from . import LlmEngine
+
+class HfLlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using HuggingFace's InferenceClient.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        temperature: float,
+        api_token: Optional[str],
+    ):
+        """
+        Initialize the HuggingFace Inference API LLM Engine.
+
+        Args:
+            model: HuggingFace model ID to use
+            temperature: Sampling temperature
+            api_token: HuggingFace API token
+        """
+        self._model = model
+        self._temperature = temperature
+
+        self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
+        self._config = AutoConfig.from_pretrained(model, token=api_token)
+        self._client = InferenceClient(token=api_token)
+                        
+    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.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+            continuation_text: Part of the response that is already generated
+            should_stop: Callback that returns True when inference should stop
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
+        ]
+        
+        stream = self._client.chat_completion(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+            add_generation_prompt=False,
+            stream=True
+        )
+        
+        try:
+            for response in stream:
+                if should_stop():
+                    stream.close()
+                    break
+                if content := response.choices[0].delta.content:
+                    yield content
+        finally:
+            stream.close()
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        return len(self._tokenizer.encode(prompt))
+
+    def token_limit(self) -> int:
+        """
+        Get the model's context window size.
+        
+        Returns:
+            int: Maximum number of tokens the model can process
+        """
+        return self._config.max_position_embeddings

sia/llm_engine/local_llm_engine.py

@@ -63,6 +63,7 @@ class LocalLlmEngine(LlmEngine):
         Args:
             system_prompt: The system prompt string
             main_context: The main context string after templating
+            continuation_text: Part of the response that is already generated
             should_stop: Callback that returns True when inference should stop
 
         Returns:

sia/llm_engine/openai_llm_engine.py

@@ -1,75 +1,77 @@
-from typing import Callable, Iterator
-import openai
-import tiktoken
-
-from . import LlmEngine
-
-class OpenAILlmEngine(LlmEngine):
-    """
-    LLM Engine implementation using OpenAI's API.
-    Supports streaming responses from chat completion models.
-    """
-    
-    def __init__(
-        self,
-        model: str,
-        temperature: float,
-        token_limit: int,
-        api_key: str,
-    ):
-        """
-        Initialize the OpenAI LLM Engine.
-
-        Args:
-            model: OpenAI model to use
-            temperature: Temperature for sampling
-            api_key: OpenAI API key
-            token_limit: Maximum number of tokens to generate
-        """
-        self._model = model
-        self._temperature = temperature
-        self._token_limit = token_limit
-
-        self._client = openai.Client(
-            api_key=api_key,
-        )
-                        
-    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}
-        ]
-        
-        stream = self._client.chat.completions.create(
-            model=self._model,
-            messages=messages,
-            temperature=self._temperature,
-            stream=True,
-        )
-
-        try:
-            for chunk in stream:
-                if should_stop():
-                    break
-                if content := chunk.choices[0].delta.content:
-                    yield content
-        finally:
-            stream.close()
-            #stream.response.close()
-    
-    def token_count(self, system_prompt: str, main_context: str) -> int:
-        """
-        Calculate the total token count for the system prompt and context.
-        
-        Args:
-            system_prompt: The system prompt string
-            main_context: The main context string
-            
-        Returns:
-            int: Total number of tokens
-        """
-        encoding = tiktoken.encoding_for_model(self._model)
-        return len(encoding.encode(system_prompt)) + len(encoding.encode(main_context))
-    
-    def token_limit(self) -> int:
+from typing import Callable, Iterator
+import openai
+import tiktoken
+
+from . import LlmEngine
+
+class OpenAILlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using OpenAI's API.
+    Supports streaming responses from chat completion models.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        temperature: float,
+        token_limit: int,
+        api_key: str,
+    ):
+        """
+        Initialize the OpenAI LLM Engine.
+
+        Args:
+            model: OpenAI model to use
+            temperature: Temperature for sampling
+            api_key: OpenAI API key
+            token_limit: Maximum number of tokens to generate
+        """
+        self._model = model
+        self._temperature = temperature
+        self._token_limit = token_limit
+
+        self._client = openai.Client(
+            api_key=api_key,
+        )
+                        
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        if continuation_text:
+            print("OpenAI LLM Engine: continuation_text is not supported")
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        
+        stream = self._client.chat.completions.create(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+            stream=True,
+        )
+
+        try:
+            for chunk in stream:
+                if should_stop():
+                    break
+                if content := chunk.choices[0].delta.content:
+                    yield content
+        finally:
+            stream.close()
+            #stream.response.close()
+    
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        """
+        Calculate the total token count for the system prompt and context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        encoding = tiktoken.encoding_for_model(self._model)
+        return len(encoding.encode(system_prompt)) + len(encoding.encode(main_context))
+    
+    def token_limit(self) -> int:
         return self._token_limit

sia/llm_engine/qwq_llm_engine.py

@@ -70,6 +70,7 @@ class QwQLlmEngine(LlmEngine):
         Args:
             system_prompt: The system prompt string
             main_context: The main context string after templating
+            continuation_text: Part of the response that is already generated
             should_stop: Callback that returns True when inference should stop
 
         Returns: