Allow stopping of inference

run.sh

@@ -2,6 +2,15 @@
 
 set -e
 
+function chown_iterations() {
+    if [ -d "./iterations" ] && [ "$(find ./iterations/ ! -user $USER -o ! -group $USER 2>/dev/null)" ]; then
+        echo "Chowning iterations directory"
+        sudo chown -R $USER:$USER ./iterations/
+    fi
+}
+
+trap chown_iterations EXIT
+
 docker build \
     --tag sia \
     .
@@ -13,6 +22,6 @@ docker run \
     -p 8080:8080 \
     -v /$(pwd)/model/:/root/model/ \
     -v /$(pwd)/.env:/root/.env \
-    -v /$(pwd)/iterations/:/root/sia/iterations/ \
+    -v /$(pwd)/iterations/:/root/iterations/ \
     -v /$(pwd)/.git/:/root/sia/.git/ \
     sia "$@"

sia/__main__.py

@@ -34,7 +34,8 @@ class Main:
             self._llms['local'] = LocalLlmEngine(
                 config.local_model,
                 config.local_temperature,
-                config.local_token_limit
+                config.local_token_limit,
+                config.local_api_key,
             )
             
         if config.openai_enabled:

sia/config.py

@@ -82,6 +82,12 @@ class Config:
             default=int(os.getenv('SIA_LOCAL_TOKEN_LIMIT', '2048')),
             help='Local LLM token limit (env: SIA_LOCAL_TOKEN_LIMIT)'
         )
+        parser.add_argument(
+            '--local-api-key',
+            type=str,
+            default=os.getenv('SIA_LOCAL_API_KEY'),
+            help='API key for local models (env: SIA_LOCAL_API_KEY)'
+        )
 
         # OpenAI configuration
         parser.add_argument(
@@ -234,6 +240,10 @@ class Config:
     def local_token_limit(self) -> int:
         return self.args.local_token_limit
     
+    @property
+    def local_api_key(self) -> Optional[str]:
+        return self.args.local_api_key
+
     # OpenAI properties
     @property
     def openai_enabled(self) -> bool:

sia/hf_llm_engine.py

@@ -1,6 +1,6 @@
 from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer, AutoConfig
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Callable
 
 from .llm_engine import LlmEngine
 
@@ -30,25 +30,23 @@ class HfLlmEngine(LlmEngine):
         self._config = AutoConfig.from_pretrained(model, token=api_token)
         self._client = InferenceClient(token=api_token)
                         
-    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]:
         """
         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.
         """
-        token_count=self.token_count(system_prompt, main_context)
-
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": main_context}
         ]
         
-        def stream_wrapper():
         stream = self._client.chat_completion(
             model=self._model,
             messages=messages,
@@ -56,10 +54,15 @@ class HfLlmEngine(LlmEngine):
             stream=True
         )
         
+        try:
             for response in stream:
+                if should_stop():
+                    stream.close()
+                    break
                 if content := response.choices[0].delta.content:
                     yield content
-        return stream_wrapper()
+        finally:
+            stream.close()
 
     def token_count(self, system_prompt: str, main_context: str) -> int:
         messages = [

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

sia/local_llm_engine.py

@@ -1,5 +1,5 @@
 from threading import Thread
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Callable
 
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
 import torch
@@ -49,13 +49,14 @@ class LocalLlmEngine(LlmEngine):
             return_full_text=False,
         )
 
-    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]:
         """
-        Run inference using the system prompt and main context, while validating actions against the provided XML schema.
+        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.
@@ -71,16 +72,21 @@ class LocalLlmEngine(LlmEngine):
             self._tokenizer,
             skip_prompt=True
         )
-        pipeline_kwargs = dict(
+        generation_thread = Thread(target=self._pipeline, kwargs=dict(
             text_inputs=prompt,
             do_sample=True,
             temperature=self._temperature,
             max_new_tokens=self.token_limit(),
             streamer=streamer
-        )
+        ))
-        thread = Thread(target=self._pipeline, kwargs=pipeline_kwargs)
+        generation_thread.start()
-        thread.start()
+
-        return util.stop_before_value(streamer, '<|eot_id|>')
+        for text in util.stop_before_value(streamer, '<|eot_id|>'):
+            yield text
+            if should_stop():
+                break
+
+        generation_thread.join()
 
     def token_count(self, system_prompt: str, main_context: str) -> int:
         """

sia/mistral_llm_engine.py

@@ -1,6 +1,4 @@
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Callable
-from abc import ABC, abstractmethod
-import os
 from mistralai import Mistral
 from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
@@ -23,11 +21,7 @@ class MistralLlmEngine(LlmEngine):
         self._client = Mistral(api_key=api_key)
         self._tokenizer = MistralTokenizer.v3()
 
-    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]:
-        messages = [
-            SystemMessage(content=system_prompt),
-            UserMessage(content=main_context),
-        ]
         messages = [
             {
                 "role": "system",
@@ -48,9 +42,16 @@ class MistralLlmEngine(LlmEngine):
             messages=messages,
             temperature=self._temperature,
         )
+        
+        try:
             for chunk in stream_response:
+                if should_stop():
+                    stream_response.close()
+                    break
                 if content := chunk.data.choices[0].delta.content:
-                yield chunk.data.choices[0].delta.content
+                    yield content
+        finally:
+            stream_response.close()
 
     def token_count(self, system_prompt: str, main_context: str) -> int:
         messages = [

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]:
+    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
-
-        Returns:
-            Iterator[str]: An iterator that yields the generated text in chunks.
-        """
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": main_context}
@@ -57,9 +47,15 @@ class OpenAILlmEngine(LlmEngine):
             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:
         """

sia/web/api.py

@@ -30,6 +30,7 @@ class Api:
     def _init_routes(self):
         """Initialize REST API and WebSocket routes."""
         self._app.router.add_post("/api/inference/{llm}", self._run_inference)
+        self._app.router.add_post("/api/inference/{llm}/stop", self._stop_inference)
         self._app.router.add_post("/api/approve/{llm}", self._approve_response)
         self._app.router.add_post("/api/context", self._modify_context)
         self._app.router.add_post("/api/input", self._send_input)
@@ -60,6 +61,15 @@ class Api:
         except (ValueError, RuntimeError) as e:
             return web.Response(status=400, text=str(e))
 
+    async def _stop_inference(self, request: web.Request) -> web.Response:
+        """Stop inference on specified LLM."""
+        llm_name = request.match_info["llm"]
+        try:
+            self._agent.stop_inference(llm_name)
+            return web.Response(status=200)
+        except ValueError as e:
+            return web.Response(status=400, text=str(e))
+
     async def _approve_response(self, request: web.Request) -> web.Response:
         """Approve response from specified LLM."""
         llm_name = request.match_info["llm"]

sia/web_agent.py

@@ -46,6 +46,7 @@ class WebAgent(BaseAgent):
         self._validation_error: Optional[str] = None
         self._command_result: Optional[CommandResult] = None
         self._context = self._compile_context(next(iter(self._llms.values())))
+        self._stop_flags: Dict[str, bool] = {name: False for name in llms}
         
         # Locks
         self._llm_lock = Lock()
@@ -112,9 +113,14 @@ class WebAgent(BaseAgent):
             if self._llm_states[llm_name] != LlmState.NO_OUTPUT:
                 raise RuntimeError(f"LLM {llm_name} is not ready for inference")
             self._set_llm_state(llm_name, LlmState.INFERENCE)
+            self._stop_flags[llm_name] = False
 
         llm = self._llms[llm_name]
-        response_token_iter = llm.infer(self.system_prompt, self.context)
+        
+        def should_stop() -> bool:
+            return self._stop_flags[llm_name]
+        
+        response_token_iter = llm.infer(self.system_prompt, self.context, should_stop)
         
         with self._output_lock:
             self._llm_outputs[llm_name] = ""
@@ -129,6 +135,12 @@ class WebAgent(BaseAgent):
         with self._llm_lock:
             self._set_llm_state(llm_name, LlmState.OUTPUT)
 
+    def stop_inference(self, llm_name: str) -> None:
+        """Stop ongoing inference for specified LLM"""
+        if llm_name not in self._llms:
+            raise ValueError(f"Unknown LLM: {llm_name}")
+        self._stop_flags[llm_name] = True
+
     def get_output(self, llm_name: str) -> str:
         """Get complete output for specified LLM"""
         if llm_name not in self._llms:

tasks/fix_unit_tests.txt

@@ -0,0 +1,4 @@
+The unit tests for the sia project are no longer up-to-date.
+Find failing unit tests and fix them.
+Find missing unit tests and write them.
+Make regular commits to the git repo in the /root/sia directory but don't push them.

tasks/improve_sia.txt

@@ -0,0 +1,5 @@
+The SIA source is located in /root/sia.
+Not all features are implemented yet.
+Look at the readme and code to find what is missing.
+Make sure to unit test your work.
+Make commits to the git repo when a features are done but don't push them.

web/src/components/App.jsx

@@ -159,6 +159,12 @@ const App = () => {
     }
   };
 
+  const handleStop = () => {
+    fetch(`/api/inference/${activeLlm}/stop`, {
+      method: 'POST'
+    });
+  };
+
   const handleSendInput = () => {
     if (input.trim()) {
       fetch('/api/input', {
@@ -347,6 +353,7 @@ const App = () => {
           output={output}
           autoApproverConfig={autoApproverConfig}
           onApprove={handleApprove}
+          onStop={handleStop}
           onToggleDiff={() => setShowDiff(!showDiff)}
           onSendInput={handleSendInput}
           onClearOutput={handleClearOutput}

web/src/components/Sidebar.jsx

@@ -16,6 +16,7 @@ export const Sidebar = ({
   output,
   autoApproverConfig,
   onApprove,
+  onStop,
   onToggleDiff,
   onSendInput,
   onClearOutput,
@@ -31,7 +32,6 @@ export const Sidebar = ({
     ${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
   `}>
     <div className="p-4 space-y-4 overflow-y-auto h-full">
-      {/* Existing controls */}
       <Select value={activeLlm} onValueChange={onLlmChange}>
         <SelectTrigger>
           <SelectValue placeholder="Select LLM" />
@@ -62,14 +62,13 @@ export const Sidebar = ({
       </div>
 
       <Button 
-        onClick={onApprove}
+        onClick={llmState === LlmState.INFERENCE ? onStop : onApprove}
-        disabled={llmState === LlmState.INFERENCE}
         className="w-full"
       >
         {
           llmState === LlmState.NO_OUTPUT ? 'Inference' :
-          llmState === LlmState.OUTPUT ? 'Approve' :
+          llmState === LlmState.INFERENCE ? 'Stop' :
-          'In Progress'
+          'Approve'
         }
       </Button>
 

web/src/components/editors/MemoryEditor.jsx

@@ -46,6 +46,7 @@ export const MemoryEditor = ({
   const handleLoadIteration = async (event) => {
     const file = event.target.files[0];
     if (file) {
+      try {
         const content = await file.text();
         const response = await fetch('/api/memory/load_iteration', {
           method: 'POST',
@@ -53,10 +54,16 @@ export const MemoryEditor = ({
           body: JSON.stringify({ content: content })
         });
         if (!response.ok) throw new Error('Failed to load iteration');
+      } catch (error) {
+        if (error instanceof DOMException) {
+          alert('Failed to read the file. Check permissions. chown -R $USER:$USER iteration');
+        } else {
+          throw error;
+        }
+      }
     }
   };
 
-
   const handleLoadIterationClick = () => {
     const input = document.createElement('input');
     input.type = 'file';