From e71bd7e9eb6f7085d6de50744b7d7a79664392d3 Mon Sep 17 00:00:00 2001 From: geens Date: Sat, 30 Nov 2024 12:44:13 +0100 Subject: [PATCH] Allow stopping of inference --- run.sh | 13 +++++++-- sia/__main__.py | 3 ++- sia/config.py | 10 +++++++ sia/hf_llm_engine.py | 29 ++++++++++++--------- sia/llm_engine.py | 4 +-- sia/local_llm_engine.py | 24 ++++++++++------- sia/mistral_llm_engine.py | 25 +++++++++--------- sia/openai_llm_engine.py | 26 ++++++++---------- sia/web/api.py | 10 +++++++ sia/web_agent.py | 16 ++++++++++-- tasks/fix_unit_tests.txt | 4 +++ tasks/improve_sia.txt | 5 ++++ web/src/components/App.jsx | 7 +++++ web/src/components/Sidebar.jsx | 9 +++---- web/src/components/editors/MemoryEditor.jsx | 23 ++++++++++------ 15 files changed, 139 insertions(+), 69 deletions(-) create mode 100644 tasks/fix_unit_tests.txt create mode 100644 tasks/improve_sia.txt diff --git a/run.sh b/run.sh index b4140ac..d0028f8 100755 --- a/run.sh +++ b/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 "$@" \ No newline at end of file diff --git a/sia/__main__.py b/sia/__main__.py index d24c17a..84dfb50 100644 --- a/sia/__main__.py +++ b/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: diff --git a/sia/config.py b/sia/config.py index 812cb0f..b70bc6d 100644 --- a/sia/config.py +++ b/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( @@ -233,6 +239,10 @@ class Config: @property 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 diff --git a/sia/hf_llm_engine.py b/sia/hf_llm_engine.py index 98a1207..734a410 100644 --- a/sia/hf_llm_engine.py +++ b/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,36 +30,39 @@ 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, - temperature=self._temperature, - stream=True - ) - + 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 - return stream_wrapper() + finally: + stream.close() def token_count(self, system_prompt: str, main_context: str) -> int: messages = [ diff --git a/sia/llm_engine.py b/sia/llm_engine.py index ffee779..7b9bff2 100644 --- 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 diff --git a/sia/local_llm_engine.py b/sia/local_llm_engine.py index 28e7469..203fc16 100644 --- a/sia/local_llm_engine.py +++ b/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) - thread.start() - return util.stop_before_value(streamer, '<|eot_id|>') + )) + generation_thread.start() + + 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: """ @@ -112,4 +118,4 @@ class LocalLlmEngine(LlmEngine): if self._token_limit is not None: return self._token_limit else: - return self._pipeline.model.config.max_position_embeddings \ No newline at end of file + return self._pipeline.model.config.max_position_embeddings diff --git a/sia/mistral_llm_engine.py b/sia/mistral_llm_engine.py index 9b4aeef..68504be 100644 --- a/sia/mistral_llm_engine.py +++ b/sia/mistral_llm_engine.py @@ -1,6 +1,4 @@ -from typing import Iterator, Optional -from abc import ABC, abstractmethod -import os +from typing import Iterator, Optional, Callable 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]: - messages = [ - SystemMessage(content=system_prompt), - UserMessage(content=main_context), - ] + def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: messages = [ { "role": "system", @@ -48,9 +42,16 @@ class MistralLlmEngine(LlmEngine): messages=messages, temperature=self._temperature, ) - for chunk in stream_response: - if content := chunk.data.choices[0].delta.content: - yield chunk.data.choices[0].delta.content + + try: + for chunk in stream_response: + if should_stop(): + stream_response.close() + break + if content := chunk.data.choices[0].delta.content: + yield content + finally: + stream_response.close() def token_count(self, system_prompt: str, main_context: str) -> int: messages = [ @@ -67,4 +68,4 @@ class MistralLlmEngine(LlmEngine): return len(tokenized.tokens) def token_limit(self) -> int: - return self._token_limit \ No newline at end of file + return self._token_limit diff --git a/sia/openai_llm_engine.py b/sia/openai_llm_engine.py index ad161a2..d4e9d7b 100644 --- a/sia/openai_llm_engine.py +++ b/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: """ diff --git a/sia/web/api.py b/sia/web/api.py index 7ea1033..e0523ce 100644 --- a/sia/web/api.py +++ b/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"] diff --git a/sia/web_agent.py b/sia/web_agent.py index 2120bed..e4d62c7 100644 --- a/sia/web_agent.py +++ b/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: @@ -164,4 +176,4 @@ class WebAgent(BaseAgent): def _handle_memory_update(self) -> None: """Handle memory updates and update context""" context = self._compile_context(next(iter(self._llms.values()))) - self.modify_context(context, True) \ No newline at end of file + self.modify_context(context, True) diff --git a/tasks/fix_unit_tests.txt b/tasks/fix_unit_tests.txt new file mode 100644 index 0000000..5fc12f6 --- /dev/null +++ b/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. \ No newline at end of file diff --git a/tasks/improve_sia.txt b/tasks/improve_sia.txt new file mode 100644 index 0000000..e1d7337 --- /dev/null +++ b/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. \ No newline at end of file diff --git a/web/src/components/App.jsx b/web/src/components/App.jsx index a9f57de..3572f73 100644 --- a/web/src/components/App.jsx +++ b/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} diff --git a/web/src/components/Sidebar.jsx b/web/src/components/Sidebar.jsx index 34e1de7..185fbbf 100644 --- a/web/src/components/Sidebar.jsx +++ b/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'} `}>
- {/* Existing controls */}