Allow stopping of inference

2024-11-30 12:44:13 +01:00
parent a0353d0d49
commit e71bd7e9eb
15 changed files with 139 additions and 69 deletions
--- a/sia/llm_engine.py
+++ b/sia/llm_engine.py
@@ -1,9 +1,9 @@
-from typing import Iterator
+from typing import Callable, Iterator
 from abc import ABC, abstractmethod

 class LlmEngine(ABC):
    @abstractmethod
-    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
+    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
        pass

    @abstractmethod