Allow stopping of inference

sia/openai_llm_engine.py

@@ -1,4 +1,4 @@
-from typing import Iterator
+from typing import Callable, Iterator
 import openai
 import tiktoken
 
@@ -34,17 +34,7 @@ class OpenAILlmEngine(LlmEngine):
             api_key=api_key,
         )
                         
-    def infer(self, system_prompt: str, main_context: str) -> 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
-
-        Returns:
-            Iterator[str]: An iterator that yields the generated text in chunks.
-        """
+    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}
@@ -57,9 +47,15 @@ class OpenAILlmEngine(LlmEngine):
             stream=True,
         )
 
-        for chunk in stream:
-            if content := chunk.choices[0].delta.content:
-                yield content
+        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:
         """