From c09f0766c1414e66d7ff9b27513bcd6e8b33018c Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Fri, 18 Apr 2025 11:36:17 +0200 Subject: [PATCH] Fixed auto approver and inference continuation --- sia/auto_approver.py | 4 +- sia/llm_engine/hf_llm_engine.py | 171 +++++++++--------- sia/llm_engine/local_llm_engine.py | 1 + sia/llm_engine/openai_llm_engine.py | 150 +++++++-------- sia/llm_engine/qwq_llm_engine.py | 1 + sia/web/api.py | 164 ++++++++--------- sia/web_agent.py | 8 +- web/src/components/App.jsx | 23 +-- web/src/components/editors/StandardEditor.jsx | 98 +++++----- 9 files changed, 304 insertions(+), 316 deletions(-) diff --git a/sia/auto_approver.py b/sia/auto_approver.py index aecaef6..3b6fdcf 100644 --- a/sia/auto_approver.py +++ b/sia/auto_approver.py @@ -178,7 +178,5 @@ class AutoApprover: def _response_approval_thread(self) -> None: if self._stop_event.wait(self._response_timeout): return - if (self._response_enabled and - self.agent.llms[self._llm_name] == LlmState.IDLE): - response = self.agent.response_buffer.get_text() + if self._response_enabled: self.agent.approve_response() diff --git a/sia/llm_engine/hf_llm_engine.py b/sia/llm_engine/hf_llm_engine.py index 1f0193f..6a9bce3 100644 --- a/sia/llm_engine/hf_llm_engine.py +++ b/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 diff --git a/sia/llm_engine/local_llm_engine.py b/sia/llm_engine/local_llm_engine.py index e29a3d6..f4ecb57 100644 --- a/sia/llm_engine/local_llm_engine.py +++ b/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: diff --git a/sia/llm_engine/openai_llm_engine.py b/sia/llm_engine/openai_llm_engine.py index 656768e..c061bb4 100644 --- a/sia/llm_engine/openai_llm_engine.py +++ b/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 \ No newline at end of file diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py index afba73f..d628f6e 100644 --- a/sia/llm_engine/qwq_llm_engine.py +++ b/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: diff --git a/sia/web/api.py b/sia/web/api.py index bb8edef..510ae6e 100644 --- a/sia/web/api.py +++ b/sia/web/api.py @@ -59,20 +59,25 @@ class Api: async def _run_inference(self, request: web.Request) -> web.Response: """Start inference on specified LLM.""" - llm_name = request.match_info["llm"] try: + llm_name = request.match_info["llm"] + data = await request.json() + text = data.get("response") + context = data.get("context") + self._agent._response_buffer.set_text(text) + self._agent.modify_context(context) await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name) return web.Response(status=200) - except (ValueError, RuntimeError) as e: + except Exception 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: + llm_name = request.match_info["llm"] self._agent.stop_inference(llm_name) return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _set_response(self, request: web.Request) -> web.Response: @@ -80,15 +85,10 @@ class Api: try: data = await request.json() text = data.get("response") - if text is None: - return web.Response(status=400, text="Missing response text in request body") - try: - self._agent._response_buffer.set_text(text) - return web.Response(status=200) - except ValueError as e: - return web.Response(status=400, text=str(e)) - except json.JSONDecodeError: - return web.Response(status=400, text="Invalid JSON in request body") + self._agent._response_buffer.set_text(text) + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) async def _approve_response(self, request: web.Request) -> web.Response: """Approve current buffer content""" @@ -98,35 +98,37 @@ class Api: self._agent.response_buffer.set_text(response) self._agent.approve_response() return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _get_response(self, request: web.Request) -> web.Response: """Get current buffer content""" return web.Response( text=json.dumps({ - "response": self._agent.output_buffer.get_text(), + "response": self._agent.response_buffer.get_text(), }), content_type="application/json" ) async def _modify_context(self, request: web.Request) -> web.Response: """Modify the current context.""" - data = await request.json() - context = data.get("context") - if not context: - return web.Response(status=400, text="Missing context in request body") - self._agent.modify_context(context) - return web.Response(status=200) + try: + data = await request.json() + context = data.get("context") + self._agent.modify_context(context) + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) async def _send_input(self, request: web.Request) -> web.Response: """Send input to the IO buffer.""" - data = await request.json() - input_text = data.get("input") - if not input_text: - return web.Response(status=400, text="Missing input in request body") - self._io_buffer.append_stdin(input_text) - return web.Response(status=200) + try: + data = await request.json() + input_text = data.get("input") + self._io_buffer.append_stdin(input_text) + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) async def _clear_output(self, request: web.Request) -> web.Response: """Clear the stdout buffer.""" @@ -135,14 +137,14 @@ class Api: async def _get_output(self, request: web.Request) -> web.Response: """Get complete output for specified LLM.""" - llm_name = request.match_info["llm"] try: + llm_name = request.match_info["llm"] output = self._agent.get_output(llm_name) return web.Response( text=json.dumps({"output": output}), content_type="application/json" ) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _get_llms(self, request: web.Request) -> web.Response: @@ -165,71 +167,68 @@ class Api: async def _set_auto_approver_config(self, request: web.Request) -> web.Response: """Update auto approver configuration.""" - data = await request.json() try: + data = await request.json() self._auto_approver.set_config(data) return web.Response(status=200) - except (ValueError, KeyError) as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _set_context_enabled(self, request: web.Request) -> web.Response: """Set context auto-approval enabled state.""" - data = await request.json() - enabled = data.get("enabled") - if enabled is None: - return web.Response(status=400, text="Missing enabled parameter") try: - self._auto_approver.context_enabled = enabled + data = await request.json() + enabled = data.get("enabled") + self._auto_approver.context_enabled = enabled or False return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _set_response_enabled(self, request: web.Request) -> web.Response: """Set response auto-approval enabled state.""" - data = await request.json() - enabled = data.get("enabled") - if enabled is None: - return web.Response(status=400, text="Missing enabled parameter") try: - self._auto_approver.response_enabled = enabled + data = await request.json() + enabled = data.get("enabled") + self._auto_approver.response_enabled = enabled or False return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _set_context_timeout(self, request: web.Request) -> web.Response: """Set context auto-approval timeout.""" - data = await request.json() - timeout = data.get("timeout") - if timeout is None: - return web.Response(status=400, text="Missing timeout parameter") try: + data = await request.json() + timeout = data.get("timeout") + if timeout is None: + return web.Response(status=400, text="Missing timeout parameter") self._auto_approver.context_timeout = float(timeout) return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) + async def _set_response_timeout(self, request: web.Request) -> web.Response: """Set response auto-approval timeout.""" - data = await request.json() - timeout = data.get("timeout") - if timeout is None: - return web.Response(status=400, text="Missing timeout parameter") try: + data = await request.json() + timeout = data.get("timeout") + if timeout is None: + return web.Response(status=400, text="Missing timeout parameter") self._auto_approver.response_timeout = float(timeout) return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _set_llm_name(self, request: web.Request) -> web.Response: """Set LLM name for auto-approval.""" - data = await request.json() - name = data.get("name") - if name is None: - return web.Response(status=400, text="Missing name parameter") try: + data = await request.json() + name = data.get("name") + if name is None: + return web.Response(status=400, text="Missing name parameter") self._auto_approver.llm_name = name return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _get_memory(self, request: web.Request) -> web.Response: @@ -244,29 +243,29 @@ class Api: ) async def _create_entry(self, request: web.Request) -> web.Response: """Create a new entry in working memory.""" - data = await request.json() try: + data = await request.json() entry = EntryFactory.create_entry(data, self._work_dir, self._io_buffer) self._working_memory.add_entry(entry) return web.Response( text=json.dumps({"id": entry.id}), content_type="application/json" ) - except (ValueError, TypeError) as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _save_entry(self, request: web.Request) -> web.Response: """Update properties of an existing entry.""" - entry_id = request.match_info["id"] - data = await request.json() - entry = self._working_memory.get_entry(entry_id) - if not entry: - return web.Response(status=404, text="Entry not found") try: + entry_id = request.match_info["id"] + data = await request.json() + entry = self._working_memory.get_entry(entry_id) + if not entry: + return web.Response(status=404, text="Entry not found") EntryFactory.update_entry(entry, data) entry.notify_change() return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _delete_entry(self, request: web.Request) -> web.Response: @@ -286,27 +285,30 @@ class Api: async def _update_entry(self, request: web.Request) -> web.Response: """Update an entry's state.""" - entry_id = request.match_info["id"] - entry = self._working_memory.get_entry(entry_id) - if not entry: - return web.Response(status=404, text="Entry not found") try: + entry_id = request.match_info["id"] + entry = self._working_memory.get_entry(entry_id) + if not entry: + return web.Response(status=404, text="Entry not found") entry.update() entry.notify_change() return web.Response(status=200) - except ValueError as e: + except Exception as e: return web.Response(status=400, text=str(e)) async def _load_iteration(self, request: web.Request) -> web.Response: """Load entries from iteration XML content into working memory""" - data = await request.json() - content = data.get("content") - if not content: - return web.Response(status=400, text="Missing content in request body") + try: + data = await request.json() + content = data.get("content") + if not content: + return web.Response(status=400, text="Missing content in request body") + + entries = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer) - entries = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer) - - for entry in entries: - self._working_memory.add_entry(entry) - - return web.Response(status=200) + for entry in entries: + self._working_memory.add_entry(entry) + + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) diff --git a/sia/web_agent.py b/sia/web_agent.py index 70d0d43..133ba21 100644 --- a/sia/web_agent.py +++ b/sia/web_agent.py @@ -97,10 +97,6 @@ class WebAgent(BaseAgent): """Update context and reset all LLM states""" with self._llm_lock: self._context = context - self._response_buffer.clear() - for llm_name in self._llms: - self._set_llm_state(llm_name, LlmState.IDLE) - for handler in self._context_change_handlers: handler(context, generated) @@ -118,7 +114,6 @@ class WebAgent(BaseAgent): return self._stop_flags[llm_name] response_token_iter = llm.infer(self.system_prompt, self.context, self._response_buffer.get_text(), should_stop) for token in response_token_iter: - print(token, end='', flush=True) with self._output_lock: self._response_buffer.append_text(token) with self._llm_lock: @@ -132,11 +127,10 @@ class WebAgent(BaseAgent): def approve_response(self) -> None: """Process approved response from specified LLM""" - if self.llms.get(llm_name) != LlmState.IDLE: - return timestamp = datetime.now(timezone.utc) self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text()) parse_result = self._parser.parse(timestamp, self._response_buffer.get_text()) + self._response_buffer.clear() if isinstance(parse_result, Command): result = parse_result.execute(self._working_memory) self._command_result = result diff --git a/web/src/components/App.jsx b/web/src/components/App.jsx index 9436c9b..d233307 100644 --- a/web/src/components/App.jsx +++ b/web/src/components/App.jsx @@ -16,7 +16,6 @@ const App = () => { // Editor content state const [generatedContext, setGeneratedContext] = useState(''); const [modifiedContext, setModifiedContext] = useState(''); - const [contextDirty, setContextDirty] = useState(false); const [generatedResponse, setGeneratedResponse] = useState(''); const [modifiedResponse, setModifiedResponse] = useState(''); const [input, setInput] = useState(''); @@ -65,7 +64,6 @@ const App = () => { // Handle context changes useEffect(() => { contextWs.addMessageHandler((data) => { - setContextDirty(false); setModifiedContext(data.context); if (data.generated) { setGeneratedContext(data.context); @@ -127,20 +125,13 @@ const App = () => { }; const handleInference = () => { - fetch('/api/context', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ context: modifiedContext }) - }); - - fetch('/api/response', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ response: modifiedResponse }) - }); - fetch(`/api/inference/${activeLlm}`, { method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + response: modifiedResponse, + context: modifiedContext, + }) }); }; @@ -176,10 +167,6 @@ const App = () => { } setLlms(resetLlms); setModifiedContext(context); - - // Reset response buffers - setGeneratedResponse(''); - setModifiedResponse(''); }; const handleResponseEdit = (response) => { diff --git a/web/src/components/editors/StandardEditor.jsx b/web/src/components/editors/StandardEditor.jsx index 4da0a31..0c91eff 100644 --- a/web/src/components/editors/StandardEditor.jsx +++ b/web/src/components/editors/StandardEditor.jsx @@ -4,58 +4,58 @@ import { Card, CardContent } from '@/components/ui/card'; import { Alert, AlertDescription } from '@/components/ui/alert'; export const StandardEditor = ({ - content, - onChange, - readOnly = false, - language = 'xml', - validationError = null, - autoScroll = false, + content, + onChange, + readOnly = false, + language = 'xml', + validationError = null, + autoScroll = false, }) => { - const editorRef = React.useRef(null); + const editorRef = React.useRef(null); - const scrollToBottom = () => { - if (editorRef.current && autoScroll) { - const model = editorRef.current.getModel(); - const lineCount = model.getLineCount(); - editorRef.current.revealLine(lineCount, monaco.editor.ScrollType.Smooth); - } - }; + const scrollToBottom = () => { + if (editorRef.current && autoScroll) { + const model = editorRef.current.getModel(); + const lineCount = model.getLineCount(); + editorRef.current.revealLine(lineCount, monaco.editor.ScrollType.Smooth); + } + }; - const handleEditorDidMount = (editor) => { - editorRef.current = editor; - setTimeout(() => { - scrollToBottom(); - }, 10); - }; + const handleEditorDidMount = (editor) => { + editorRef.current = editor; + setTimeout(() => { + scrollToBottom(); + }, 10); + }; - React.useEffect(() => { - scrollToBottom(); - }, [content, autoScroll]); + React.useEffect(() => { + scrollToBottom(); + }, [content, autoScroll]); - return ( - - - {validationError && ( - - {validationError} - - )} - - - - ); + return ( + + + {validationError && ( + + {validationError} + + )} + + + + ); }; \ No newline at end of file