From d4a4902b9487eaca4a5af3e0d49d9e3009ee5693 Mon Sep 17 00:00:00 2001 From: geens Date: Tue, 20 May 2025 09:43:17 +0200 Subject: [PATCH] New web interface, move llm engine to separate process --- .gitignore | 1 + Dockerfile | 12 +- Makefile | 20 +- README.md | 77 ++-- lib/llm_engine_utils/pyproject.toml | 13 + .../src/llm_engine_utils/__init__.py | 7 + .../src/llm_engine_utils}/dataset.py | 0 .../src/llm_engine_utils/iterators.py | 61 ++++ .../src/llm_engine_utils/llm_engine.py | 16 + .../src/llm_engine_utils/protocol.py | 87 +++++ lib/xml_schema_validator/pyproject.toml | 1 + .../llama_cpp_logits_processor.py | 3 +- llm_config.toml | 91 +++++ scripts/collect.sh | 2 +- scripts/install.sh | 46 ++- setup.py | 13 - sia/__main__.py | 52 +-- sia/auto_approver.py | 172 +++------ sia/base_agent.py | 14 +- sia/chat_io_buffer.py | 102 ++++++ sia/config.py | 262 +------------- sia/iteration_logger.py | 5 +- sia/llm_engine.py | 212 +++++++++++ sia/llm_engine/__init__.py | 15 - sia/util.py | 62 ---- sia/web/api.py | 182 +++++----- sia/web/auto_approver_websocket.py | 5 +- sia/web/chat_websocket.py | 64 ++++ sia/web/context_websocket.py | 55 --- sia/web/memory_websocket.py | 2 - .../{llm_websocket.py => state_websocket.py} | 33 +- sia/web/websockets.py | 19 +- sia/web_agent.py | 145 ++++---- sia/working_memory.py | 4 +- test/auto_approver_test.py | 24 +- test/base_agent_test.py | 29 +- test/llm_engine_test.py | 237 ++++++++++++ test/local_llm_engine_test.py | 53 --- test/util_test.py | 31 -- test/web_agent_test.py | 40 +-- test/xml_validator_test.py | 116 ------ tools/gemma_infer/pyproject.toml | 19 + tools/gemma_infer/src/gemma_infer/__main__.py | 42 +++ .../src/gemma_infer/gemma_llm_engine.py | 76 ++++ tools/mistral_infer/pyproject.toml | 18 + .../src/mistral_infer/__main__.py | 42 +++ .../src/mistral_infer}/mistral_llm_engine.py | 45 ++- tools/mistral_train/pyproject.toml | 19 + .../src/mistral_train/__main__.py} | 48 +-- .../mistral_train/src/mistral_train/config.py | 38 ++ .../old_llm_engines}/hf_llm_engine.py | 0 .../old_llm_engines}/local_llm_engine.py | 0 .../old_llm_engines}/openai_llm_engine.py | 0 .../old_llm_engines}/qwq_llm_engine.ipynb | 0 .../old_llm_engines}/qwq_llm_engine.py | 0 .../old_llm_engines}/qwq_vllm.ipynb | 0 tools/{train => qwq_train}/setup.py | 0 tools/{train => qwq_train}/train/__init__.py | 0 tools/{train => qwq_train}/train/qwen.ipynb | 0 tools/{train => qwq_train}/train/qwq.ipynb | 0 tools/{train => qwq_train}/train/qwq.py | 0 tools/train/bin/train | 15 - tools/train/readme.md | 68 ---- web/.dockerignore | 4 - web/.gitignore | 5 - web/index.html | 5 +- web/package.json | 44 +-- web/src/App.tsx | 77 ++++ web/src/components/App.jsx | 336 ------------------ web/src/components/Chat/ChatInput.tsx | 62 ++++ web/src/components/Chat/ChatMessage.tsx | 49 +++ web/src/components/Chat/ChatWindow.tsx | 172 +++++++++ web/src/components/Chat/index.ts | 4 + web/src/components/Controls/AgentSection.tsx | 67 ++++ .../Controls/AutoApproverSection.tsx | 102 ++++++ web/src/components/Controls/LLMSection.tsx | 91 +++++ web/src/components/Controls/MemorySection.tsx | 126 +++++++ web/src/components/Controls/index.ts | 4 + web/src/components/FloatingButton.jsx | 14 - web/src/components/Header.jsx | 58 --- web/src/components/Layout/MobileControls.tsx | 26 ++ web/src/components/Layout/MobileTaskbar.tsx | 53 +++ web/src/components/Layout/Ribbon.tsx | 30 ++ web/src/components/Layout/index.ts | 3 + web/src/components/Memory/MemoryEditor.tsx | 269 ++++++++++++++ .../Memory/entries/BackgroundEntry.tsx | 158 ++++++++ .../components/Memory/entries/BaseEntry.tsx | 157 ++++++++ web/src/components/Memory/entries/IOEntry.tsx | 87 +++++ .../Memory/entries/ParseErrorEntry.tsx | 75 ++++ .../components/Memory/entries/ReadEntry.tsx | 78 ++++ .../Memory/entries/ReasoningEntry.tsx | 53 +++ .../components/Memory/entries/RepeatEntry.tsx | 215 +++++++++++ .../components/Memory/entries/ScriptEntry.tsx | 127 +++++++ .../components/Memory/entries/SingleEntry.tsx | 228 ++++++++++++ .../components/Memory/entries/WriteEntry.tsx | 78 ++++ web/src/components/Memory/entries/index.ts | 8 + web/src/components/Memory/index.ts | 1 + .../components/Response/ResponseEditor.tsx | 77 ++++ web/src/components/Response/index.ts | 1 + web/src/components/Sidebar.jsx | 177 --------- .../editors/BackgroundEntryEditor.jsx | 124 ------- .../components/editors/BaseEntryEditor.jsx | 117 ------ web/src/components/editors/ContentEditor.jsx | 37 -- .../components/editors/DiffEditorWrapper.jsx | 60 ---- web/src/components/editors/MemoryEditor.jsx | 217 ----------- .../editors/ParseErrorEntryEditor.jsx | 53 --- .../components/editors/ReadEntryEditor.jsx | 54 --- .../editors/ReasoningEntryEditor.jsx | 40 --- .../components/editors/RepeatEntryEditor.jsx | 164 --------- .../components/editors/SingleEntryEditor.jsx | 176 --------- web/src/components/editors/StandardEditor.jsx | 61 ---- .../components/editors/WriteEntryEditor.jsx | 54 --- web/src/components/shared/Button.tsx | 58 +++ web/src/components/shared/index.ts | 1 + web/src/components/ui/alert.jsx | 29 -- web/src/components/ui/button.jsx | 36 -- web/src/components/ui/card.jsx | 35 -- web/src/components/ui/input.jsx | 15 - web/src/components/ui/scroll-area.jsx | 29 -- web/src/components/ui/select.jsx | 52 --- web/src/components/ui/switch.jsx | 21 -- web/src/components/ui/tabs.jsx | 33 -- web/src/components/ui/textarea.jsx | 14 - web/src/contexts/MonacoContext.tsx | 41 +++ web/src/contexts/WebSocketContext.tsx | 243 +++++++++++++ web/src/hooks/useScreenSize.ts | 24 ++ web/src/hooks/useWebSocket.jsx | 68 ---- web/src/index.css | 45 +-- web/src/{index.jsx => index.tsx} | 4 +- web/src/services/api.ts | 254 +++++++++++++ web/src/types/entries.ts | 54 +++ web/src/types/index.ts | 49 +++ web/tailwind.config.js | 70 +--- web/tsconfig.json | 25 ++ web/tsconfig.node.json | 10 + web/vite.config.js | 19 - web/vite.config.ts | 31 ++ 137 files changed, 4850 insertions(+), 3503 deletions(-) create mode 100644 lib/llm_engine_utils/pyproject.toml create mode 100644 lib/llm_engine_utils/src/llm_engine_utils/__init__.py rename {tools/train/train => lib/llm_engine_utils/src/llm_engine_utils}/dataset.py (100%) create mode 100644 lib/llm_engine_utils/src/llm_engine_utils/iterators.py create mode 100644 lib/llm_engine_utils/src/llm_engine_utils/llm_engine.py create mode 100644 lib/llm_engine_utils/src/llm_engine_utils/protocol.py create mode 100644 llm_config.toml create mode 100644 sia/chat_io_buffer.py create mode 100644 sia/llm_engine.py delete mode 100644 sia/llm_engine/__init__.py create mode 100644 sia/web/chat_websocket.py delete mode 100644 sia/web/context_websocket.py rename sia/web/{llm_websocket.py => state_websocket.py} (64%) create mode 100644 test/llm_engine_test.py delete mode 100644 test/local_llm_engine_test.py delete mode 100644 test/xml_validator_test.py create mode 100644 tools/gemma_infer/pyproject.toml create mode 100644 tools/gemma_infer/src/gemma_infer/__main__.py create mode 100644 tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py create mode 100644 tools/mistral_infer/pyproject.toml create mode 100644 tools/mistral_infer/src/mistral_infer/__main__.py rename {sia/llm_engine => tools/mistral_infer/src/mistral_infer}/mistral_llm_engine.py (60%) create mode 100644 tools/mistral_train/pyproject.toml rename tools/{train/train/mistral_api.py => mistral_train/src/mistral_train/__main__.py} (66%) create mode 100644 tools/mistral_train/src/mistral_train/config.py rename {sia/llm_engine => tools/old_llm_engines}/hf_llm_engine.py (100%) rename {sia/llm_engine => tools/old_llm_engines}/local_llm_engine.py (100%) rename {sia/llm_engine => tools/old_llm_engines}/openai_llm_engine.py (100%) rename {sia/llm_engine => tools/old_llm_engines}/qwq_llm_engine.ipynb (100%) rename {sia/llm_engine => tools/old_llm_engines}/qwq_llm_engine.py (100%) rename {sia/llm_engine => tools/old_llm_engines}/qwq_vllm.ipynb (100%) rename tools/{train => qwq_train}/setup.py (100%) rename tools/{train => qwq_train}/train/__init__.py (100%) rename tools/{train => qwq_train}/train/qwen.ipynb (100%) rename tools/{train => qwq_train}/train/qwq.ipynb (100%) rename tools/{train => qwq_train}/train/qwq.py (100%) delete mode 100644 tools/train/bin/train delete mode 100644 tools/train/readme.md delete mode 100644 web/.dockerignore delete mode 100644 web/.gitignore create mode 100644 web/src/App.tsx delete mode 100644 web/src/components/App.jsx create mode 100644 web/src/components/Chat/ChatInput.tsx create mode 100644 web/src/components/Chat/ChatMessage.tsx create mode 100644 web/src/components/Chat/ChatWindow.tsx create mode 100644 web/src/components/Chat/index.ts create mode 100644 web/src/components/Controls/AgentSection.tsx create mode 100644 web/src/components/Controls/AutoApproverSection.tsx create mode 100644 web/src/components/Controls/LLMSection.tsx create mode 100644 web/src/components/Controls/MemorySection.tsx create mode 100644 web/src/components/Controls/index.ts delete mode 100644 web/src/components/FloatingButton.jsx delete mode 100644 web/src/components/Header.jsx create mode 100644 web/src/components/Layout/MobileControls.tsx create mode 100644 web/src/components/Layout/MobileTaskbar.tsx create mode 100644 web/src/components/Layout/Ribbon.tsx create mode 100644 web/src/components/Layout/index.ts create mode 100644 web/src/components/Memory/MemoryEditor.tsx create mode 100644 web/src/components/Memory/entries/BackgroundEntry.tsx create mode 100644 web/src/components/Memory/entries/BaseEntry.tsx create mode 100644 web/src/components/Memory/entries/IOEntry.tsx create mode 100644 web/src/components/Memory/entries/ParseErrorEntry.tsx create mode 100644 web/src/components/Memory/entries/ReadEntry.tsx create mode 100644 web/src/components/Memory/entries/ReasoningEntry.tsx create mode 100644 web/src/components/Memory/entries/RepeatEntry.tsx create mode 100644 web/src/components/Memory/entries/ScriptEntry.tsx create mode 100644 web/src/components/Memory/entries/SingleEntry.tsx create mode 100644 web/src/components/Memory/entries/WriteEntry.tsx create mode 100644 web/src/components/Memory/entries/index.ts create mode 100644 web/src/components/Memory/index.ts create mode 100644 web/src/components/Response/ResponseEditor.tsx create mode 100644 web/src/components/Response/index.ts delete mode 100644 web/src/components/Sidebar.jsx delete mode 100644 web/src/components/editors/BackgroundEntryEditor.jsx delete mode 100644 web/src/components/editors/BaseEntryEditor.jsx delete mode 100644 web/src/components/editors/ContentEditor.jsx delete mode 100644 web/src/components/editors/DiffEditorWrapper.jsx delete mode 100644 web/src/components/editors/MemoryEditor.jsx delete mode 100644 web/src/components/editors/ParseErrorEntryEditor.jsx delete mode 100644 web/src/components/editors/ReadEntryEditor.jsx delete mode 100644 web/src/components/editors/ReasoningEntryEditor.jsx delete mode 100644 web/src/components/editors/RepeatEntryEditor.jsx delete mode 100644 web/src/components/editors/SingleEntryEditor.jsx delete mode 100644 web/src/components/editors/StandardEditor.jsx delete mode 100644 web/src/components/editors/WriteEntryEditor.jsx create mode 100644 web/src/components/shared/Button.tsx create mode 100644 web/src/components/shared/index.ts delete mode 100644 web/src/components/ui/alert.jsx delete mode 100644 web/src/components/ui/button.jsx delete mode 100644 web/src/components/ui/card.jsx delete mode 100644 web/src/components/ui/input.jsx delete mode 100644 web/src/components/ui/scroll-area.jsx delete mode 100644 web/src/components/ui/select.jsx delete mode 100644 web/src/components/ui/switch.jsx delete mode 100644 web/src/components/ui/tabs.jsx delete mode 100644 web/src/components/ui/textarea.jsx create mode 100644 web/src/contexts/MonacoContext.tsx create mode 100644 web/src/contexts/WebSocketContext.tsx create mode 100644 web/src/hooks/useScreenSize.ts delete mode 100644 web/src/hooks/useWebSocket.jsx rename web/src/{index.jsx => index.tsx} (60%) create mode 100644 web/src/services/api.ts create mode 100644 web/src/types/entries.ts create mode 100644 web/src/types/index.ts create mode 100644 web/tsconfig.json create mode 100644 web/tsconfig.node.json delete mode 100644 web/vite.config.js create mode 100644 web/vite.config.ts diff --git a/.gitignore b/.gitignore index 240f74e..665efe4 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ **/_unsloth_temporary_saved_buffers/ .env __pycache__/ +build/ cache/ collect.txt data/ diff --git a/Dockerfile b/Dockerfile index dad64a8..db8d221 100644 --- a/Dockerfile +++ b/Dockerfile @@ -49,7 +49,7 @@ RUN mkdir -p \ /root/desktop \ /root/venvs -# ITB tool setup +# Tool ITB setup FROM base AS itb-env RUN python3 -m venv /root/venvs/itb --system-site-packages COPY ./tools/itb/setup.py /root/sia/tools/itb/setup.py @@ -58,15 +58,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \ /root/venvs/itb/bin/pip install -r *.egg-info/requires.txt RUN rm -rf *.egg-info/ -# Train tool setup -FROM base AS train-env -RUN python3 -m venv /root/venvs/train --system-site-packages -COPY ./tools/train/setup.py /root/sia/tools/train/setup.py -RUN /root/venvs/train/bin/python /root/sia/tools/train/setup.py egg_info -RUN --mount=type=cache,target=/root/.cache/pip \ - /root/venvs/train/bin/pip install -r *.egg-info/requires.txt -RUN rm -rf *.egg-info/ - # SIA core setup FROM base AS sia-env RUN python3 -m venv /root/venvs/sia --system-site-packages @@ -110,7 +101,6 @@ RUN npm run build FROM base COPY --from=itb-env /root/venvs/itb /root/venvs/itb -COPY --from=train-env /root/venvs/train /root/venvs/train COPY --from=sia-env /root/venvs/sia /root/venvs/sia COPY --from=notebook-env /root/venvs/notebook /root/venvs/notebook COPY --from=web-build /app/dist /root/static/ diff --git a/Makefile b/Makefile index 2e51447..b121c5b 100644 --- a/Makefile +++ b/Makefile @@ -127,21 +127,17 @@ lint-rust: ## Lint Rust code # ================== .PHONY: test -test: test-python test-rust ## Run all tests +test: test-sia test-xml_schema_validator -.PHONY: test-python -test-python: ## Run Python tests +.PHONY: test-sia +test-sia: @echo "Running Python tests..." - python -m pytest -xvs + /root/venvs/sia/bin/python -m unittest discover -s test -p "*_test.py" -.PHONY: test-rust -test-rust: ## Run Rust tests +.PHONY: test-xml_schema_validator +test-xml_schema_validator: @echo "Running Rust tests..." - for dir in $(RUST_DIRS); do \ - if [ -f "$$dir/Cargo.toml" ]; then \ - (cd "$$dir" && cargo test); \ - fi; \ - done + cd lib/xml_schema_validator && cargo test # ================== # Build targets @@ -183,4 +179,4 @@ help: ## Show this help message .PHONY: collect collect: ## Run collect.sh script @echo "Running collect.sh script..." - ./scripts/collect.sh -s core -s lib -s webui -o collect.txt \ No newline at end of file + ./scripts/collect.sh -s core -s tests \ No newline at end of file diff --git a/README.md b/README.md index a8b8332..afd0fff 100644 --- a/README.md +++ b/README.md @@ -311,18 +311,51 @@ Or xml comments used for inline reasoning are not saved after parsing. LLM engine subprocesses receive input as XML documents containing paths to required files and the context. -**Example input to LLM engine subprocess:** +##### Example interaction with the LLM engine subprocess: + +**Core -> LLM engine** + ```xml - - /root/sia/system_prompt.md - /root/sia/action_schema.xsd - - - - - - - + +``` + +**LLM engine -> Core** + +``` +1024\u0004 +``` + +**Core -> LLM engine** + +```xml + + /root/sia/action_schema.xsd + + + +``` + +**LLM engine -> Core** + +``` +405\u0004 +``` + +**Core -> LLM engine** + +```xml + + /root/sia/action_schema.xsd + + + + +``` + +**LLM engine -> Core** + +```xml +... ``` Though the LLM can output any text, the goal is to output valid xml. @@ -336,13 +369,6 @@ This character was chosen because: - It's a single byte, making it efficient to process - It's standard across all platforms -**Example output from LLM engine subprocess:** -``` - -I should check the current state of the system and see if there are any pending tasks. -\u0004 -``` - The communication protocol between the SIA agent and LLM engine subprocesses has been designed with simplicity as the primary goal. Opting for a minimal approach that: @@ -453,16 +479,19 @@ classDiagram -monitor_loop() void } - class LLMEngine { - +LLMEngine(executable_path str) - +infer(system_prompt str, main_context str, prefix str) Iterator~str~ + class LlmEngine { + +LlmEngine(executable_path str, action_schema_path str) + +infer(system_prompt str, main_context ET, prefix str) Iterator~str~ + +token_count(system_prompt str, main_context ET, prefix str) int + +token_limit() int + +restart() } class BaseAgent { <> -working_memory: WorkingMemory -metrics: SystemMetrics - -llm: LLMEngine + -llm: LlmEngine -parser: ResponseParser -validator: XMLValidator -action_schema: str @@ -522,7 +551,7 @@ classDiagram } SystemMetrics "1" --* "1" BaseAgent - LLMEngine "1" --* "1" BaseAgent + LlmEngine "1" --* "1" BaseAgent XMLValidator "1" --* "1" BaseAgent BaseAgent "1" *-- "1" IOBuffer BaseAgent "1" *-- "1" WorkingMemory @@ -560,7 +589,7 @@ classDiagram <> -working_memory: WorkingMemory -metrics: SystemMetrics - -llm: LLMEngine + -llm: LlmEngine -parser: ResponseParser -validator: XMLValidator -action_schema: str diff --git a/lib/llm_engine_utils/pyproject.toml b/lib/llm_engine_utils/pyproject.toml new file mode 100644 index 0000000..cf5aebc --- /dev/null +++ b/lib/llm_engine_utils/pyproject.toml @@ -0,0 +1,13 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "llm_engine_utils" +version = "0.1.0" +requires-python = ">=3.8" + +[project.optional-dependencies] +dataset = [ + "torch>=4.0.0", +] \ No newline at end of file diff --git a/lib/llm_engine_utils/src/llm_engine_utils/__init__.py b/lib/llm_engine_utils/src/llm_engine_utils/__init__.py new file mode 100644 index 0000000..c734897 --- /dev/null +++ b/lib/llm_engine_utils/src/llm_engine_utils/__init__.py @@ -0,0 +1,7 @@ +try: + from . import dataset +except ImportError: + pass +from . import iterators +from . import protocol +from .llm_engine import LlmEngine \ No newline at end of file diff --git a/tools/train/train/dataset.py b/lib/llm_engine_utils/src/llm_engine_utils/dataset.py similarity index 100% rename from tools/train/train/dataset.py rename to lib/llm_engine_utils/src/llm_engine_utils/dataset.py diff --git a/lib/llm_engine_utils/src/llm_engine_utils/iterators.py b/lib/llm_engine_utils/src/llm_engine_utils/iterators.py new file mode 100644 index 0000000..f0c2f18 --- /dev/null +++ b/lib/llm_engine_utils/src/llm_engine_utils/iterators.py @@ -0,0 +1,61 @@ +from typing import Iterator + +def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]: + """ + Creates an iterator that yields values from the input iterator + until it encounters the stop_value (exclusive). + + Args: + iterator: The source iterator + stop_value: The value to stop before + + Yields: + Values from the iterator until stop_value is encountered + If stop_value is part of an item, yields the part before stop_value + """ + for item in iterator: + if stop_value in item: + split_point = item.index(stop_value) + if split_point > 0: + yield item[:split_point] + break + yield item + +def skip_prefix(iterator: Iterator[str], prefix: str) -> Iterator[str]: + """ + Creates an iterator that skips a prefix from the input iterator + and yields only the content after the prefix. + + Args: + iterator: The source iterator + prefix: The prefix to skip + + Yields: + Values from the iterator after the prefix has been fully skipped + """ + if not prefix: + # If no prefix to skip, yield everything + yield from iterator + return + + prefix_remaining = prefix + + for item in iterator: + if prefix_remaining: + # If the item starts with the remaining prefix + if prefix_remaining.startswith(item): + # Skip this item entirely + prefix_remaining = prefix_remaining[len(item):] + continue + elif item.startswith(prefix_remaining): + # Yield only the part after the prefix + yield item[len(prefix_remaining):] + prefix_remaining = "" + else: + # Item doesn't match prefix pattern, yield everything + # This is unexpected but we handle it gracefully + yield item + prefix_remaining = "" + else: + # No prefix remaining, yield all content + yield item \ No newline at end of file diff --git a/lib/llm_engine_utils/src/llm_engine_utils/llm_engine.py b/lib/llm_engine_utils/src/llm_engine_utils/llm_engine.py new file mode 100644 index 0000000..6fa8a65 --- /dev/null +++ b/lib/llm_engine_utils/src/llm_engine_utils/llm_engine.py @@ -0,0 +1,16 @@ +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Iterator + +class LlmEngine(ABC): + @abstractmethod + def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]: + pass + + @abstractmethod + def token_count(self, system: str, context: str) -> int: + pass + + @abstractmethod + def token_limit(self) -> int: + pass \ No newline at end of file diff --git a/lib/llm_engine_utils/src/llm_engine_utils/protocol.py b/lib/llm_engine_utils/src/llm_engine_utils/protocol.py new file mode 100644 index 0000000..a6133c2 --- /dev/null +++ b/lib/llm_engine_utils/src/llm_engine_utils/protocol.py @@ -0,0 +1,87 @@ +""" +Protocol handler for SIA LLM engine subprocess communication. + +This module provides tools to parse and handle the XML-based communication +protocol between the SIA core and LLM engine subprocesses as described +in the SIA README. +""" + +from .llm_engine import LlmEngine +from io import StringIO +from pathlib import Path +from typing import Optional +import sys +import xml.etree.ElementTree as ET + +# ASCII End of Transmission character used to signal the end of a response +EOT = '\u0004' + +def read_command() -> Optional[ET.Element]: + """ + Read an XML command from stdin. + + Returns: + ET.Element: The parsed XML element or None if EOF was reached. + """ + buffer = StringIO() + + while True: + char = sys.stdin.read(1) + if not char: # EOF + return None + + buffer.write(char) + content = buffer.getvalue() + + # Try to parse once we have a potentially complete XML document + if content.endswith('>'): + try: + return ET.fromstring(content) + except ET.ParseError: + # Not a complete XML document yet, continue reading + pass + +def process(engine: LlmEngine) -> None: + """ + Run the adapter loop, processing commands from stdin and writing responses to stdout. + """ + while True: + print("reading command", file=sys.stderr) + command = read_command() + if command is None: + print("EOF", file=sys.stderr) + # End of input stream, exit + return + + if command.tag == 'token_limit': + print("token_limit", file=sys.stderr) + limit = engine.token_limit() + sys.stdout.write(str(limit)) + sys.stdout.write(EOT) + sys.stdout.flush() + + elif command.tag == 'token_count': + print("token_count", file=sys.stderr) + system = command.find('system').text + context = command.find('context').text + count = engine.token_count(system, context) + sys.stdout.write(str(count)) + sys.stdout.write(EOT) + sys.stdout.flush() + + elif command.tag == 'infer_xml': + print("infer_xml", file=sys.stderr) + try: + schema = Path(command.find('schema').text) + system = command.find('system').text + context = command.find('context').text + prefix = command.find('prefix').text if command.find('prefix') is not None else None + for chunk in engine.infer_xml(schema, system, context, prefix): + sys.stdout.write(chunk) + sys.stdout.flush() + sys.stdout.write(EOT) + sys.stdout.flush() + except Exception as e: + print(f"Error: {e}", file=sys.stderr) + sys.stdout.write(EOT) + sys.stdout.flush() \ No newline at end of file diff --git a/lib/xml_schema_validator/pyproject.toml b/lib/xml_schema_validator/pyproject.toml index 8ba540b..e70daa3 100644 --- a/lib/xml_schema_validator/pyproject.toml +++ b/lib/xml_schema_validator/pyproject.toml @@ -5,6 +5,7 @@ description = "XML Schema validation library" requires-python = ">=3.10" dependencies = [ "maturin>=1.0,<2.0", + "torch>=2.0.0", ] optional-dependencies.dev = [ "pytest>=8", diff --git a/lib/xml_schema_validator/python/xml_schema_validator/llama_cpp_logits_processor.py b/lib/xml_schema_validator/python/xml_schema_validator/llama_cpp_logits_processor.py index b67022d..cf22b04 100644 --- a/lib/xml_schema_validator/python/xml_schema_validator/llama_cpp_logits_processor.py +++ b/lib/xml_schema_validator/python/xml_schema_validator/llama_cpp_logits_processor.py @@ -1,5 +1,6 @@ from transformers import AutoTokenizer from xml_schema_validator._rs import XmlLogitsProcessorCore +import sys class LlamaCppLogitsProcessor: """ @@ -57,7 +58,7 @@ class LlamaCppLogitsProcessor: ) marker_start = template_text.find("XML_LOGITS_PROCESSOR_MARKER") + len("XML_LOGITS_PROCESSOR_MARKER") start_marker = template_text[marker_start:] - print("schema validator start marker:", start_marker) + print("schema validator start marker:", start_marker, file=sys.stderr) return start_marker def get_processor(self): diff --git a/llm_config.toml b/llm_config.toml new file mode 100644 index 0000000..c488df4 --- /dev/null +++ b/llm_config.toml @@ -0,0 +1,91 @@ +gemma = """/root/venvs/gemma_infer/bin/gemma_infer --model /root/models/gemma3_1b.gguf --tokenizer /root/models/gemma3_1b_tokenizer""" +#mistral = """/root/venvs/mistral_infer/bin/mistral_infer""" +default = """#!/bin/bash + +# Mock LLM engine for testing +# This script simulates an LLM engine subprocess that responds to the three commands: +# , ..., and ... + +# Function to read XML input until a complete closing tag is found +read_xml_input() { + local input="" + local line + local char_count=0 + + # Read until we get a complete XML command + while IFS= read -r line; do + input="${input}${line}" + char_count=$((char_count + ${#line})) + + # Debug the actual content (first 30 chars) + + if [[ "$input" == *""* ]]; then + printf "1024" + printf "\\004" # EOT character (hex 04) + return + fi + + if [[ "$input" == *""*""* ]]; then + printf "405" + printf "\\004" # EOT character (hex 04) + return + fi + + if [[ "$input" == *""*""* ]]; then + generate_response + return + fi + + done +} + +# Function to generate a response token by token +generate_response() { + printf "<" + sleep 0.3 + + printf "reason" + sleep 0.3 + + printf "ing" + sleep 0.3 + + printf ">" + sleep 0.3 + + printf "This" + sleep 0.3 + + printf " is" + sleep 0.3 + + printf " a" + sleep 0.3 + + printf " test" + sleep 0.3 + + printf " response." + sleep 0.3 + + printf "" + sleep 0.3 + + printf "\\004" +} + +# Main loop - keep reading input and responding +iteration=0 +while true; do + iteration=$((iteration + 1)) + read_xml_input +done""" diff --git a/scripts/collect.sh b/scripts/collect.sh index 04cbdaf..9a65d40 100755 --- a/scripts/collect.sh +++ b/scripts/collect.sh @@ -5,7 +5,7 @@ set -euo pipefail declare -A FILTER_SETS=( ["py"]="-f .*(\\.py|requirements.txt)$" ["rust"]="-f .*\\.(rs|toml)$" - ["web"]="-f .*\\.(js|jsx|json|css|html)$" + ["web"]="-f .*\\.(js|jsx|json|css|html|ts|tsx|md)$" ["doc"]="-f .*\\.md$" ["deploy"]="-f .*(Dockerfile|akefile|\\.toml|\\.sh|\\.xsd|\\.yaml)$" diff --git a/scripts/install.sh b/scripts/install.sh index 87cd826..4a50dd2 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -1,28 +1,46 @@ #!/bin/bash +SIA_INSTALL_NO_NOTEBOOK=1 +#SIA_INSTALL_NO_CORE=1 +SIA_INSTALL_NO_ITB=1 +SIA_INSTALL_NO_MISTRAL_INFER=1 +#SIA_INSTALL_NO_GEMMA_INFER=1 +SIA_INSTALL_NO_TRAIN=1 + +if [ -z "${SIA_INSTALL_NO_NOTEBOOK}" ]; then + echo "Installing venv for running notebooks..." + python3 -m venv /root/venvs/notebook + /root/venvs/notebook/bin/pip install jupyter ipykernel ipywidgets + /root/venvs/notebook/bin/ipython kernel install --name=notebook +fi + +if [ -z "${SIA_INSTALL_NO_CORE}" ]; then + echo "Installing SIA core..." + python3 -m venv /root/venvs/sia + /root/venvs/sia/bin/pip install -e /root/sia +fi + if [ -z "${SIA_INSTALL_NO_ITB}" ]; then echo "Installing ITB tool..." python3 -m venv /root/venvs/itb /root/venvs/itb/bin/pip install -e /root/sia/tools/itb fi +if [ -z "${SIA_INSTALL_NO_MISTRAL_INFER}" ]; then + echo "Installing venv for mistral inference..." + python3 -m venv /root/venvs/mistral_infer + /root/venvs/mistral_infer/bin/pip install -e /root/sia/tools/mistral_infer +fi + +if [ -z "${SIA_INSTALL_NO_GEMMA_INFER}" ]; then + echo "Installing venv for gemma inference..." + python3 -m venv /root/venvs/gemma_infer + /root/venvs/gemma_infer/bin/pip install -e /root/sia/tools/gemma_infer +fi + if [ -z "${SIA_INSTALL_NO_TRAIN}" ]; then echo "Installing Train tool..." python3 -m venv /root/venvs/train /root/venvs/train/bin/pip install -e /root/sia/tools/train /root/venvs/train/bin/ipython kernel install --name=train -fi - -if [ -z "${SIA_INSTALL_NO_CORE}" ]; then - echo "Installing SIA core..." - python3 -m venv /root/venvs/sia - /root/venvs/sia/bin/pip install -e /root/sia - /root/venvs/sia/bin/ipython kernel install --name=sia -fi - -if [ -z "${SIA_INSTALL_NO_NOTEBOOK}"]; then - echo "Installing venv for running notebooks..." - python3 -m venv /root/venvs/notebook - /root/venvs/notebook/bin/pip install jupyter ipykernel ipywidgets - /root/venvs/notebook/bin/ipython kernel install --name=notebook fi \ No newline at end of file diff --git a/setup.py b/setup.py index 0fb6ac4..725699b 100644 --- a/setup.py +++ b/setup.py @@ -10,25 +10,12 @@ setup( ], }, install_requires=[ - 'accelerate>=0.26.0', 'aiohttp>=3.8.0', - 'bitsandbytes>=0.45', - 'dotenv-python>=0.0.1', - 'huggingface_hub>=0.16.0', 'ipykernel>=6.0.0', 'ipywidgets>=8.0.0', 'lxml>=4.9.0', - 'mistral-common>=1.0.0', - 'mistralai>=0.0.7', - 'openai>=1.0.0', 'psutil>=5.9.0', 'python-dotenv>=1.0.0', - 'tiktoken>=0.4.0', - 'torch>=2.0.0', - 'transformers>=4.30.0', - 'trl>=0.7.8', - 'unsloth>=2025.3', - 'vllm==0.8.2', 'xml_schema_validator @ file:///root/sia/lib/xml_schema_validator', ], python_requires='>=3.10', diff --git a/sia/__main__.py b/sia/__main__.py index 4a25d01..6e7f0ec 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -2,20 +2,16 @@ from aiohttp import web import asyncio from .auto_approver import AutoApprover +from .chat_io_buffer import ChatIOBuffer from .config import Config from .iteration_logger import IterationLogger -from .llm_engine.hf_llm_engine import HfLlmEngine -from .llm_engine.local_llm_engine import LocalLlmEngine -from .llm_engine.mistral_llm_engine import MistralLlmEngine -from .llm_engine.openai_llm_engine import OpenAILlmEngine -from .llm_engine.qwq_llm_engine import QwQLlmEngine +from .llm_engine import LlmEngine from .response_parser import ResponseParser from .system_metrics import SystemMetrics from .web.api import Api from .web.static import Static from .web.websockets import Websockets from .web_agent import WebAgent -from .web_io_buffer import WebIOBuffer from .working_memory import WorkingMemory class Main: @@ -30,49 +26,17 @@ class Main: # Initialize LLM engines based on config self._llms = {} - if config.local_enabled: - self._llms['local'] = LocalLlmEngine( - config.local_model, - config.local_temperature, - config.local_token_limit, - self._action_schema, - config.local_api_key, - ) - - if config.openai_enabled: - self._llms['openai'] = OpenAILlmEngine( - config.openai_model, - config.openai_temperature, - config.openai_token_limit, - config.openai_api_key, - ) - - if config.hf_enabled: - self._llms['hf'] = HfLlmEngine( - config.hf_model, - config.hf_temperature, - config.hf_api_key, - ) - - if config.mistral_enabled: - self._llms['mistral'] = MistralLlmEngine( - config.mistral_model, - config.mistral_temperature, - config.mistral_token_limit, - config.mistral_api_key, - ) - - if config.qwq_enabled: - self._llms['qwq'] = QwQLlmEngine( - config.qwq_model, - config.qwq_temperature, - self._action_schema, + # Use the config.llms property which returns only enabled LLMs + for llm_name, executable_path in config.llms.items(): + self._llms[llm_name] = LlmEngine( + executable_path=executable_path, + action_schema_path=str(config.action_schema) ) if not self._llms: raise ValueError("No LLM engines enabled in configuration") - self._io_buffer = WebIOBuffer() + self._io_buffer = ChatIOBuffer() self._working_memory = WorkingMemory() self._agent = WebAgent( system_prompt=self._system_prompt, diff --git a/sia/auto_approver.py b/sia/auto_approver.py index 3b6fdcf..320db2f 100644 --- a/sia/auto_approver.py +++ b/sia/auto_approver.py @@ -1,13 +1,14 @@ +from dataclasses import dataclass from threading import Thread, Event -from typing import Callable, TypeAlias, TypedDict -from .web_agent import WebAgent, LlmState +from typing import Callable, TypeAlias +from .web_agent import WebAgent, AgentState -class AutoApproverConfig(TypedDict): +@dataclass +class AutoApproverConfig: context_enabled: bool response_enabled: bool context_timeout: float response_timeout: float - llm_name: str ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None] @@ -21,48 +22,47 @@ class AutoApprover: Initialize auto approver with a WebAgent instance. """ self.agent = agent - self._llm_name = next(iter(agent.llms.keys())) - self._context_timeout = 5.0 - self._response_timeout = 10.0 - self._context_enabled = False - self._response_enabled = False + self._config = AutoApproverConfig( + context_enabled=False, + response_enabled=False, + context_timeout=1.0, + response_timeout=1.0, + ) self._stop_event = Event() self._context_thread: Thread | None = None self._response_thread: Thread | None = None self._config_change_handlers: list[ConfigChangeHandler] = [] - self.agent.add_llm_change_handler(self._handle_llm_state_change) - self.agent.add_context_change_handler(self._handle_context_change) + self._previous_state = agent.state + self.agent.add_state_change_handler(self._handle_state_change) @property def config(self) -> AutoApproverConfig: - return AutoApproverConfig( - context_enabled=self._context_enabled, - response_enabled=self._response_enabled, - context_timeout=self._context_timeout, - response_timeout=self._response_timeout, - llm_name=self._llm_name - ) + return self._config - def set_config(self, config: AutoApproverConfig) -> None: - if config['llm_name'] not in self.agent.llms: - raise ValueError(f"Unknown LLM: {config['llm_name']}") - - notify_config_change = self._notify_config_change - self._notify_config_change = lambda: None - - try: - self.context_enabled = False - self.response_enabled = False - - self.context_timeout = config['context_timeout'] - self.response_timeout = config['response_timeout'] - self.llm_name = config['llm_name'] - - self.context_enabled = config['context_enabled'] - self.response_enabled = config['response_enabled'] - finally: - self._notify_config_change = notify_config_change + @config.setter + def config(self, config: AutoApproverConfig) -> None: + if config == self._config: + return + if config.context_timeout != self._config.context_timeout: + if config.context_enabled and self._config.context_enabled: + raise ValueError("Cannot change context timeout while enabled") + if config.response_timeout != self._config.response_timeout: + if config.response_enabled and self._config.response_enabled: + raise ValueError("Cannot change response timeout while enabled") + self._stop_context_thread() + self._stop_response_thread() + # If both context and response are enabled, start response approval only. + # Better to run inference with a full prefix than approve an emmpty response. + if config.response_enabled and not self._config.response_enabled: + # response approval was just enabled + self._config = config + self._start_response_thread() + elif config.context_enabled and not self._config.context_enabled: + # context approval was just enabled + self._config = config + self._start_context_thread() + self._config = config self._notify_config_change() def add_config_change_handler(self, handler: ConfigChangeHandler) -> None: @@ -73,81 +73,17 @@ class AutoApprover: for handler in self._config_change_handlers: handler(current_config) - @property - def context_timeout(self) -> float: - return self._context_timeout - - @context_timeout.setter - def context_timeout(self, timeout: float) -> None: - if self._context_enabled: - raise ValueError("Cannot change timeout while auto-approval is enabled") - self._context_timeout = timeout - self._notify_config_change() - - @property - def response_timeout(self) -> float: - return self._response_timeout - - @response_timeout.setter - def response_timeout(self, timeout: float) -> None: - if self._response_enabled: - raise ValueError("Cannot change timeout while auto-approval is enabled") - self._response_timeout = timeout - self._notify_config_change() - - @property - def context_enabled(self) -> bool: - return self._context_enabled - - @context_enabled.setter - def context_enabled(self, enabled: bool) -> None: - if enabled == self._context_enabled: - return - self._context_enabled = enabled - self._stop_context_thread() - if enabled and self.agent.llms[self._llm_name] == LlmState.IDLE: - self._start_context_thread() - self._notify_config_change() - - @property - def response_enabled(self) -> bool: - return self._response_enabled - - @response_enabled.setter - def response_enabled(self, enabled: bool) -> None: - if enabled == self._response_enabled: - return - self._response_enabled = enabled - self._stop_response_thread() - if enabled and self.agent.llms[self._llm_name] == LlmState.IDLE: - self._start_response_thread() - self._notify_config_change() - - @property - def llm_name(self) -> str: - return self._llm_name - - @llm_name.setter - def llm_name(self, name: str) -> None: - if name not in self.agent.llms: - raise ValueError(f"Unknown LLM: {name}") - self._llm_name = name - self._notify_config_change() - - def _handle_llm_state_change(self, llm_name: str, state: LlmState) -> None: - if llm_name != self._llm_name: - return - - if state == LlmState.IDLE and self._response_enabled: - self._start_response_thread() - else: - self._stop_response_thread() - - def _handle_context_change(self, context: str, generated: bool) -> None: - if generated and self._context_enabled: - self._start_context_thread() - else: - self._stop_context_thread() + def _handle_state_change(self, state: AgentState) -> None: + if state == AgentState.IDLE: + if self._previous_state == AgentState.INFERENCE: + # finished inference + if self._config.response_enabled: + self._start_response_thread() + elif self._previous_state == AgentState.PROCESSING_RESPONSE: + # finished processing response + if self._config.context_enabled: + self._start_context_thread() + self._previous_state = state def _stop_context_thread(self) -> None: if self._context_thread: @@ -170,13 +106,13 @@ class AutoApprover: self._response_thread.start() def _context_approval_thread(self) -> None: - if self._stop_event.wait(self._context_timeout): + if self._stop_event.wait(self._config.context_timeout): return - if self._context_enabled: - self.agent.run_inference(self._llm_name) + if self._config.context_enabled: + self.agent.run_inference() def _response_approval_thread(self) -> None: - if self._stop_event.wait(self._response_timeout): + if self._stop_event.wait(self._config.response_timeout): return - if self._response_enabled: - self.agent.approve_response() + if self._config.response_enabled: + self.agent.approve_response() \ No newline at end of file diff --git a/sia/base_agent.py b/sia/base_agent.py index 48b34e0..703db0a 100644 --- a/sia/base_agent.py +++ b/sia/base_agent.py @@ -4,7 +4,6 @@ import xml.etree.ElementTree as ET from .llm_engine import LlmEngine from .response_parser import ResponseParser from .system_metrics import SystemMetrics -from .util import pretty_print_element from .working_memory import WorkingMemory class BaseAgent(ABC): @@ -36,8 +35,13 @@ class BaseAgent(ABC): def system_prompt(self) -> str: """Get the system prompt.""" return f"{self._system_prompt}\n{self._action_schema}" + + @property + def working_memory(self) -> WorkingMemory: + """Get the working memory.""" + return self._working_memory - def _compile_context(self, llmEngine: LlmEngine) -> str: + def _compile_context(self, llmEngine: LlmEngine) -> ET: """ Compile the current context for LLM inference. Includes system metrics and working memory entries. @@ -60,15 +64,13 @@ class BaseAgent(ABC): for entry in memory_context: context.append(entry) - - context_str = pretty_print_element(context) # Calculate token usage percentage - token_count = llmEngine.token_count(self.system_prompt, context_str) + token_count = llmEngine.token_count(self.system_prompt, context) token_limit = llmEngine.token_limit() context_usage = (float(token_count) / float(token_limit)) * 100.0 # Update context usage metric context.set("context", f"{str(round(context_usage, 2))}%") - return pretty_print_element(context) \ No newline at end of file + return context \ No newline at end of file diff --git a/sia/chat_io_buffer.py b/sia/chat_io_buffer.py new file mode 100644 index 0000000..d8d5602 --- /dev/null +++ b/sia/chat_io_buffer.py @@ -0,0 +1,102 @@ +from dataclasses import dataclass +from typing import List, Callable +from enum import Enum +import datetime + +from .io_buffer import IOBuffer +from .util import format_timestamp + +class MessageType(Enum): + USER = "user" + ASSISTANT = "assistant" + +@dataclass +class ChatMessage: + id: str # Formatted timestamp + content: str + message_type: MessageType + +class ChatIOBuffer(IOBuffer): + def __init__(self): + self._messages: List[ChatMessage] = [] + self._last_read_timestamp: str = format_timestamp(datetime.datetime.now()) + self._change_handlers: List[Callable] = [] + + @property + def messages(self) -> List[ChatMessage]: + return self._messages.copy() + + @property + def last_read_timestamp(self) -> str: + return self._last_read_timestamp + + def read(self) -> str: + unread_messages = self._unread_messages() + self._last_read_timestamp = self._format_now() + self._notify_handlers([], []) + return unread_messages + + def write(self, content: str) -> None: + message = ChatMessage( + id=self._format_now(), + content=content, + message_type=MessageType.ASSISTANT + ) + self._messages.append(message) + self._notify_handlers( + new_messages=[message], + deleted_message_ids=[] + ) + + def buffer_length(self) -> int: + return len(self._unread_messages()) + + def add_user_message(self, content: str): + message = ChatMessage( + id=self._format_now(), + content=content, + message_type=MessageType.USER + ) + self._messages.append(message) + self._notify_handlers( + new_messages=[message], + deleted_message_ids=[] + ) + + def delete_message(self, message_id: str): + for i, msg in enumerate(self._messages): + if msg.id == message_id: + self._messages.pop(i) + self._notify_handlers( + new_messages=[], + deleted_message_ids=[message_id] + ) + break + + def clear(self) -> None: + deleted_ids = [msg.id for msg in self._messages] + self._messages.clear() + self._last_read_timestamp = self._format_now() + self._notify_handlers( + new_messages=[], + deleted_message_ids=deleted_ids + ) + + def add_change_handler(self, handler: Callable[[ + List[ChatMessage], # new messages + List[str], # deleted message ids + str # last read timestamp + ], None]) -> None: + self._change_handlers.append(handler) + + def _notify_handlers(self, new_messages: List[ChatMessage], + deleted_message_ids: List[str]) -> None: + for handler in self._change_handlers: + handler(new_messages, deleted_message_ids, self._last_read_timestamp) + + def _unread_messages(self) -> List[ChatMessage]: + return "\n".join(msg.content for msg in self._messages + if msg.message_type == MessageType.USER and msg.id > self._last_read_timestamp) + + def _format_now(self) -> str: + return format_timestamp(datetime.datetime.now()) \ No newline at end of file diff --git a/sia/config.py b/sia/config.py index 011087b..093c476 100644 --- a/sia/config.py +++ b/sia/config.py @@ -1,9 +1,10 @@ from dataclasses import dataclass from dotenv import load_dotenv from pathlib import Path -from typing import Optional +from typing import Dict import argparse import os +import tomli as tomllib @dataclass class Config: @@ -59,170 +60,27 @@ class Config: parser.add_argument( '--static-files', type=Path, - default=self._parse_optional_path('SIA_STATIC_FILES', '/root/static/'), + default=os.getenv('SIA_STATIC_FILES', '/root/static/'), help='Path to static web files (default: /root/static/, env: SIA_STATIC_FILES)' ) - # Local LLM configuration + # LLM configuration parser.add_argument( - '--local-enable', - action='store_true', - default=self._parse_bool_env('SIA_LOCAL_ENABLED', False), - help='Enable local LLM engine (env: SIA_LOCAL_ENABLED)' - ) - parser.add_argument( - '--local-model', - type=str, - default=os.getenv('SIA_LOCAL_MODEL', '/root/models/current'), - help='Path to local model directory (default: /root/models/current, env: SIA_LOCAL_MODEL)' - ) - parser.add_argument( - '--local-temperature', - type=float, - default=float(os.getenv('SIA_LOCAL_TEMPERATURE', '0.1')), - help='Local LLM temperature (default: 0.1, env: SIA_LOCAL_TEMPERATURE)' - ) - parser.add_argument( - '--local-token-limit', - type=int, - 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( - '--openai-enable', - action='store_true', - default=self._parse_bool_env('SIA_OPENAI_ENABLED', False), - help='Enable OpenAI LLM engine (env: SIA_OPENAI_ENABLED)' - ) - parser.add_argument( - '--openai-model', - type=str, - default=os.getenv('SIA_OPENAI_MODEL', 'gpt-4o'), - help='OpenAI model name (default: gpt-4o, env: SIA_OPENAI_MODEL)' - ) - parser.add_argument( - '--openai-temperature', - type=float, - default=float(os.getenv('SIA_OPENAI_TEMPERATURE', '0.1')), - help='OpenAI temperature (default: 0.1, env: SIA_OPENAI_TEMPERATURE)' - ) - parser.add_argument( - '--openai-token-limit', - type=int, - default=int(os.getenv('SIA_OPENAI_TOKEN_LIMIT', '4096')), - help='OpenAI token limit (env: SIA_OPENAI_TOKEN_LIMIT)' - ) - parser.add_argument( - '--openai-api-key', - type=str, - default=os.getenv('SIA_OPENAI_API_KEY'), - help='OpenAI API key (env: SIA_OPENAI_API_KEY)' - ) - - # Hugging Face configuration - parser.add_argument( - '--hf-enable', - action='store_true', - default=self._parse_bool_env('SIA_HF_ENABLED', False), - help='Enable Hugging Face LLM engine (env: SIA_HF_ENABLED)' - ) - parser.add_argument( - '--hf-model', - type=str, - default=os.getenv('SIA_HF_MODEL'), - help='Hugging Face model name (env: SIA_HF_MODEL)' - ) - parser.add_argument( - '--hf-temperature', - type=float, - default=float(os.getenv('SIA_HF_TEMPERATURE', '0.1')), - help='Hugging Face temperature (default: 0.1, env: SIA_HF_TEMPERATURE)' - ) - parser.add_argument( - '--hf-api-key', - type=str, - default=os.getenv('SIA_HF_API_KEY'), - help='Hugging Face API key (env: SIA_HF_API_KEY)' - ) - - # Mistral configuration - parser.add_argument( - '--mistral-enable', - action='store_true', - default=self._parse_bool_env('SIA_MISTRAL_ENABLED', False), - help='Enable Mistral LLM engine (env: SIA_MISTRAL_ENABLED)' - ) - parser.add_argument( - '--mistral-model', - type=str, - default=os.getenv('SIA_MISTRAL_MODEL'), - help='Mistral model name (env: SIA_MISTRAL_MODEL)' - ) - parser.add_argument( - '--mistral-temperature', - type=float, - default=float(os.getenv('SIA_MISTRAL_TEMPERATURE', '0.1')), - help='Mistral temperature (default: 0.1, env: SIA_MISTRAL_TEMPERATURE)' - ) - parser.add_argument( - '--mistral-token-limit', - type=int, - default=int(os.getenv('SIA_MISTRAL_TOKEN_LIMIT', '4096')), - help='Mistral token limit (env: SIA_MISTRAL_TOKEN_LIMIT)' - ) - parser.add_argument( - '--mistral-api-key', - type=str, - default=os.getenv('SIA_MISTRAL_API_KEY'), - help='Mistral API key (env: SIA_MISTRAL_API_KEY)' - ) - # QwQ configuration - parser.add_argument( - '--qwq-enable', - action='store_true', - default=self._parse_bool_env('SIA_QWQ_ENABLED', False), - help='Enable QwQ LLM engine (env: SIA_QWQ_ENABLED)' - ) - parser.add_argument( - '--qwq-model', - type=str, - default=os.getenv('SIA_QWQ_MODEL', '/root/models/current'), - help='Path to QwQ model or HF model ID (default: /root/models/current, env: SIA_QWQ_MODEL)' - ) - parser.add_argument( - '--qwq-temperature', - type=float, - default=float(os.getenv('SIA_QWQ_TEMPERATURE', '0.6')), - help='QwQ temperature (default: 0.1, env: SIA_QWQ_TEMPERATURE)' - ) - parser.add_argument( - '--qwq-token-limit', - type=int, - default=int(os.getenv('SIA_QWQ_TOKEN_LIMIT', '4096')), - help='QwQ token limit (0 for model default, env: SIA_QWQ_TOKEN_LIMIT)' + '--llm-config', + type=Path, + default=os.getenv('SIA_LLM_CONFIG', '/root/sia/llm_config.toml'), + help='Path to TOML configuration file for LLMs (default: /root/sia/llm_config.toml, env: SIA_LLM_CONFIG)' ) self.args = parser.parse_args() + with open(self.llm_config, "rb") as f: + self._llm_configs = tomllib.load(f) def _parse_bool_env(self, env_var: str, default: bool) -> bool: val = os.getenv(env_var) if val is None: return default return val.lower() in ('true', '1', 'yes', 'on') - - def _parse_optional_path(self, env_var: str, default: Optional[Path]) -> Optional[Path]: - val = os.getenv(env_var) - if val is None: - return default - return Path(val) # Core properties @property @@ -258,100 +116,12 @@ class Config: def static_files(self) -> Path: return self.args.static_files - # Local LLM properties + # LLM properties @property - def local_enabled(self) -> bool: - return self.args.local_enable - - @property - def local_model(self) -> str: - return self.args.local_model - - @property - def local_temperature(self) -> float: - return self.args.local_temperature - - @property - def local_token_limit(self) -> int: - return self.args.local_token_limit + def llm_config(self) -> Path: + return self.args.llm_config @property - def local_api_key(self) -> Optional[str]: - return self.args.local_api_key - - # OpenAI properties - @property - def openai_enabled(self) -> bool: - return self.args.openai_enable - - @property - def openai_model(self) -> str: - return self.args.openai_model - - @property - def openai_temperature(self) -> float: - return self.args.openai_temperature - - @property - def openai_token_limit(self) -> int: - return self.args.openai_token_limit - - @property - def openai_api_key(self) -> Optional[str]: - return self.args.openai_api_key - - # Hugging Face properties - @property - def hf_enabled(self) -> bool: - return self.args.hf_enable - - @property - def hf_model(self) -> str: - return self.args.hf_model - - @property - def hf_temperature(self) -> float: - return self.args.hf_temperature - - @property - def hf_api_key(self) -> Optional[str]: - return self.args.hf_api_key - - # Mistral properties - @property - def mistral_enabled(self) -> bool: - return self.args.mistral_enable - - @property - def mistral_model(self) -> str: - return self.args.mistral_model - - @property - def mistral_temperature(self) -> float: - return self.args.mistral_temperature - - @property - def mistral_token_limit(self) -> int: - return self.args.mistral_token_limit - - @property - def mistral_api_key(self) -> Optional[str]: - return self.args.mistral_api_key - - # QwQ properties - @property - def qwq_enabled(self) -> bool: - return self.args.qwq_enable - - @property - def qwq_model(self) -> str: - return self.args.qwq_model - - @property - def qwq_temperature(self) -> float: - return self.args.qwq_temperature - - @property - def qwq_token_limit(self) -> Optional[int]: - # Return None if 0 to use model default - return self.args.qwq_token_limit if self.args.qwq_token_limit > 0 else None \ No newline at end of file + def llms(self) -> Dict[str, str]: + """Get only the enabled LLM configurations.""" + return self._llm_configs \ No newline at end of file diff --git a/sia/iteration_logger.py b/sia/iteration_logger.py index 6256ebc..b568707 100644 --- a/sia/iteration_logger.py +++ b/sia/iteration_logger.py @@ -23,7 +23,7 @@ class IterationLogger: def log_iteration( self, timestamp: datetime, - context: str, + context: ET, response: str, ): """ @@ -41,8 +41,7 @@ class IterationLogger: root.set("system_prompt_hash", self._system_prompt_hash) root.set("action_schema_hash", self._action_schema_hash) - context_elem = ET.SubElement(root, "context") - context_elem.text = context + root.append(context) response_elem = ET.SubElement(root, "response") response_elem.text = response diff --git a/sia/llm_engine.py b/sia/llm_engine.py new file mode 100644 index 0000000..58f758c --- /dev/null +++ b/sia/llm_engine.py @@ -0,0 +1,212 @@ +from typing import Iterator +from .util import pretty_print_element +import subprocess +import sys +import xml.etree.ElementTree as ET + +class LlmEngine: + """ + LlmEngine manages communication with LLM engine subprocesses. + Each LLM type runs in its own subprocess with a tailored environment. + """ + + EOT = '\x04' # EOT character (ASCII 4) as bytes + + def __init__(self, executable_path: str, action_schema_path: str): + """ + Initialize the LLM engine subprocess. + + Args: + executable_path (str): Path to the LLM engine executable + action_schema_path (str): Path to the XML action schema + """ + self.executable_path = executable_path + self.action_schema_path = action_schema_path + self.action_schema = open(action_schema_path, 'r').read() + self.process = None + self.restart() + + def _read_until_eot(self) -> str: + """ + Read from subprocess stdout until EOT character. + + Returns: + str: Complete response without the EOT character + """ + response = [] + + while True: + # Read available data + data = self.process.stdout.read(1024) # Read up to 1024 bytes at a time + + if not data: # process died + self.restart() + raise RuntimeError(f"LLM subprocess terminated unexpectedly") + + data = data.decode('utf-8') + + # Check if EOT is in the data + if self.EOT in data: + eot_index = data.index(self.EOT) + response.append(data[:eot_index]) # Add data before EOT + break + else: + response.append(data) + + return "".join(response) + + def token_limit(self) -> int: + """ + Get the maximum token limit of the LLM. + + Returns: + int: Maximum token limit + """ + self.process.stdin.write(b"\n") + self.process.stdin.flush() + response = self._read_until_eot() + return int(response.strip()) + + def token_count(self, system_prompt: str, main_context: ET.Element, prefix: str = "") -> int: + """ + Count the number of tokens in the prompt. + + Args: + system_prompt (str): System prompt text + main_context (ET.Element): Main context as ElementTree + prefix (str): Optional prefix for continuing generation + + Returns: + int: Token count + """ + # Create the XML document + root = ET.Element("token_count") + + # Add system prompt + system_prompt_elem = ET.SubElement(root, "system") + system_prompt_elem.text = self._append_action_schema(system_prompt) + + # Add context element + context_elem = ET.SubElement(root, "context") + context_elem.text = pretty_print_element(main_context) + + # Add prefix if provided + if prefix: + prefix_elem = ET.SubElement(root, "prefix") + prefix_elem.text = prefix + + # Send to subprocess - convert to bytes + xml_str = ET.tostring(root, encoding='utf-8') + self.process.stdin.write(xml_str + b"\n") + self.process.stdin.flush() + + # Read response + response = self._read_until_eot() + return int(response.strip()) + + def infer(self, system_prompt: str, main_context: ET.Element, prefix: str = "") -> Iterator[str]: + """ + Generate text from the LLM. + + Args: + system_prompt (str): System prompt text + main_context (ET.Element): Main context as ElementTree + prefix (str): Optional prefix for continuing generation + + Returns: + Iterator[str]: Generated text, yielded as it's produced + """ + # Create the XML document + root = ET.Element("infer_xml") + + # Add action schema path + schema_path_elem = ET.SubElement(root, "schema") + schema_path_elem.text = self.action_schema_path + + # Add system prompt in CDATA + system_prompt_elem = ET.SubElement(root, "system") + system_prompt_elem.text = self._append_action_schema(system_prompt) + + # Add context element + context_elem = ET.SubElement(root, "context") + context_elem.text = pretty_print_element(main_context) + + # Add prefix if provided + if prefix: + prefix_elem = ET.SubElement(root, "prefix") + prefix_elem.text = prefix + + # Send to subprocess - convert to bytes + xml_str = ET.tostring(root, encoding='utf-8') + self.process.stdin.write(xml_str + b"\n") + self.process.stdin.flush() + + while True: + # Read available data + data = self.process.stdout.read(1024) # Read up to 1024 bytes at a time + + if not data: # Process died + self.restart() + raise RuntimeError("LLM subprocess terminated unexpectedly") + + data = data.decode('utf-8') + + if self.EOT in data: + eot_index = data.index(self.EOT) + if eot_index > 0: + yield data[:eot_index] + break + else: + yield data + + def restart(self): + """Start the LLM engine subprocess.""" + # Ensure any existing process is terminated + if self.process: + self._terminate_process() + + # Start the subprocess with pipes for stdin/stdout and direct stderr + self.process = subprocess.Popen( + ["/bin/bash", "-c", self.executable_path], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=sys.stderr, + text=False, # Use binary mode to avoid buffering and encoding issues + bufsize=0, + ) + + # Check if the process started successfully + if self.process.poll() is not None: + raise RuntimeError(f"Failed to start LLM engine at {self.executable_path}") + + def _terminate_process(self): + """Terminate the LLM engine subprocess safely.""" + if self.process: + try: + self.process.stdin.close() + self.process.stdout.close() + self.process.terminate() + try: + self.process.wait(timeout=5) # Wait for process to terminate + except subprocess.TimeoutExpired: + # Force kill if termination takes too long + self.process.kill() + self.process.wait(timeout=2) + finally: + self.process = None + + def _append_action_schema(self, system_prompt: str) -> str: + """ + Append the action schema to the system prompt. + + Args: + system_prompt (str): Original system prompt + + Returns: + str: Updated system prompt with action schema + """ + return f"{system_prompt}\n\n--- Action Schema ---\n{self.action_schema}" + + def __del__(self): + """Cleanup when the object is destroyed.""" + self._terminate_process() \ No newline at end of file diff --git a/sia/llm_engine/__init__.py b/sia/llm_engine/__init__.py deleted file mode 100644 index 3ece6e2..0000000 --- a/sia/llm_engine/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -from typing import Callable, Iterator -from abc import ABC, abstractmethod - -class LlmEngine(ABC): - @abstractmethod - def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: - pass - - @abstractmethod - def token_count(self, system_prompt: str, main_context: str) -> int: - pass - - @abstractmethod - def token_limit(self) -> int: - pass \ No newline at end of file diff --git a/sia/util.py b/sia/util.py index 7cfc496..7377095 100644 --- a/sia/util.py +++ b/sia/util.py @@ -1,67 +1,5 @@ import datetime -import xml.dom.minidom import xml.etree.ElementTree as ET -from typing import Iterator, Optional - -def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]: - """ - Creates an iterator that yields values from the input iterator - until it encounters the stop_value (exclusive). - - Args: - iterator: The source iterator - stop_value: The value to stop before - - Yields: - Values from the iterator until stop_value is encountered - If stop_value is part of an item, yields the part before stop_value - """ - for item in iterator: - if stop_value in item: - split_point = item.index(stop_value) - if split_point > 0: - yield item[:split_point] - break - yield item - -def skip_prefix(iterator: Iterator[str], prefix: str) -> Iterator[str]: - """ - Creates an iterator that skips a prefix from the input iterator - and yields only the content after the prefix. - - Args: - iterator: The source iterator - prefix: The prefix to skip - - Yields: - Values from the iterator after the prefix has been fully skipped - """ - if not prefix: - # If no prefix to skip, yield everything - yield from iterator - return - - prefix_remaining = prefix - - for item in iterator: - if prefix_remaining: - # If the item starts with the remaining prefix - if prefix_remaining.startswith(item): - # Skip this item entirely - prefix_remaining = prefix_remaining[len(item):] - continue - elif item.startswith(prefix_remaining): - # Yield only the part after the prefix - yield item[len(prefix_remaining):] - prefix_remaining = "" - else: - # Item doesn't match prefix pattern, yield everything - # This is unexpected but we handle it gracefully - yield item - prefix_remaining = "" - else: - # No prefix remaining, yield all content - yield item def pretty_print_element(elem: ET.Element, level: int = 0, max_line: int = 80) -> str: """Convert ElementTree element to pretty-printed string with custom formatting.""" diff --git a/sia/web/api.py b/sia/web/api.py index 273c362..af8042e 100644 --- a/sia/web/api.py +++ b/sia/web/api.py @@ -1,14 +1,14 @@ -from pathlib import Path from aiohttp import web -import json +from pathlib import Path import asyncio +import json -from ..auto_approver import AutoApprover +from ..auto_approver import AutoApprover, AutoApproverConfig +from ..chat_io_buffer import ChatIOBuffer from ..entry.entry_factory import EntryFactory from ..iteration_parser import IterationParser from ..web_agent import WebAgent -from ..web_io_buffer import WebIOBuffer from ..working_memory import WorkingMemory class Api: @@ -17,7 +17,7 @@ class Api: work_dir: Path, app: web.Application, agent: WebAgent, - io_buffer: WebIOBuffer, + io_buffer: ChatIOBuffer, working_memory: WorkingMemory, auto_approver: AutoApprover ): @@ -32,23 +32,24 @@ 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_get("/api/llms", self._get_llms) + self._app.router.add_get("/api/llms/active", self._get_active_llm) + self._app.router.add_post("/api/llms/active", self._set_active_llm) + + self._app.router.add_post("/api/inference", self._run_inference) + self._app.router.add_post("/api/inference/stop", self._stop_inference) + + self._app.router.add_get("/api/response", self._get_response) self._app.router.add_post("/api/response", self._set_response) self._app.router.add_post("/api/response/approve", self._approve_response) - self._app.router.add_get("/api/response", self._get_response) - self._app.router.add_post("/api/context", self._modify_context) - self._app.router.add_post("/api/input", self._send_input) - self._app.router.add_post("/api/clear", self._clear_output) - self._app.router.add_get("/api/output/{llm}", self._get_output) - self._app.router.add_get("/api/llms", self._get_llms) - self._app.router.add_get("/api/auto_approver/config", self._get_auto_approver_config) - self._app.router.add_post("/api/auto_approver/config", self._set_auto_approver_config) - self._app.router.add_post("/api/auto_approver/context_enabled", self._set_context_enabled) - self._app.router.add_post("/api/auto_approver/response_enabled", self._set_response_enabled) - self._app.router.add_post("/api/auto_approver/context_timeout", self._set_context_timeout) - self._app.router.add_post("/api/auto_approver/response_timeout", self._set_response_timeout) - self._app.router.add_post("/api/auto_approver/llm", self._set_llm_name) + + self._app.router.add_post("/api/chat/message", self._add_chat_message) + self._app.router.add_delete("/api/chat/message/{id}", self._delete_chat_message) + self._app.router.add_delete("/api/chat", self._clear_chat) + + self._app.router.add_get("/api/auto_approver", self._get_auto_approver_config) + self._app.router.add_post("/api/auto_approver", self._set_auto_approver_config) + self._app.router.add_get("/api/memory", self._get_memory) self._app.router.add_post("/api/memory/entry", self._create_entry) self._app.router.add_put("/api/memory/entry/{id}", self._save_entry) @@ -57,16 +58,33 @@ class Api: self._app.router.add_post("/api/memory/entry/{id}/update", self._update_entry) self._app.router.add_post("/api/memory/load_iteration", self._load_iteration) - async def _run_inference(self, request: web.Request) -> web.Response: - """Start inference on specified LLM.""" + async def _get_llms(self, request: web.Request) -> web.Response: + return web.Response( + text=json.dumps( + [{"name": name} for name in self._agent.llms] + ), + content_type="application/json" + ) + + async def _get_active_llm(self, request: web.Request) -> web.Response: + return web.Response( + text=json.dumps( + {"active_llm": self._agent.active_llm} + ), + content_type="application/json" + ) + + async def _set_active_llm(self, request: web.Request) -> web.Response: 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) + self._agent.active_llm = data["active_llm"] + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) + + async def _run_inference(self, request: web.Request) -> web.Response: + try: + await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference) return web.Response(status=200) except Exception as e: return web.Response(status=400, text=str(e)) @@ -74,29 +92,7 @@ class Api: async def _stop_inference(self, request: web.Request) -> web.Response: """Stop inference on specified LLM.""" try: - llm_name = request.match_info["llm"] - self._agent.stop_inference(llm_name) - return web.Response(status=200) - except Exception as e: - return web.Response(status=400, text=str(e)) - - async def _set_response(self, request: web.Request) -> web.Response: - """Edit the shared response buffer""" - try: - data = await request.json() - text = data.get("response") - 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""" - data = await request.json() - try: - response = data.get("response") - self._agent.response_buffer.set_text(response) - self._agent.approve_response() + self._agent.stop_inference() return web.Response(status=200) except Exception as e: return web.Response(status=400, text=str(e)) @@ -110,53 +106,50 @@ class Api: content_type="application/json" ) - async def _modify_context(self, request: web.Request) -> web.Response: - """Modify the current context.""" + async def _set_response(self, request: web.Request) -> web.Response: + """Edit the shared response buffer""" try: data = await request.json() - context = data.get("context") - self._agent.modify_context(context) + text = data.get("response") + 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 _send_input(self, request: web.Request) -> web.Response: - """Send input to the IO buffer.""" + async def _approve_response(self, request: web.Request) -> web.Response: + """Approve current buffer content""" try: - data = await request.json() - input_text = data.get("input") - self._io_buffer.append_stdin(input_text) + self._agent.approve_response() 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.""" - self._io_buffer.clear_stdout() - return web.Response(status=200) - - async def _get_output(self, request: web.Request) -> web.Response: - """Get complete output for specified LLM.""" + + async def _add_chat_message(self, request: web.Request) -> web.Response: + """Add a chat message""" 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" - ) + data = await request.json() + message = data.get("message") + self._io_buffer.add_user_message(message) + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) + + async def _delete_chat_message(self, request: web.Request) -> web.Response: + """Delete a chat message""" + try: + message_id = request.match_info["id"] + self._io_buffer.delete_message(message_id) + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) + + async def _clear_chat(self, request: web.Request) -> web.Response: + """Clear chat messages""" + try: + self._io_buffer.clear() + return web.Response(status=200) except Exception as e: return web.Response(status=400, text=str(e)) - - async def _get_llms(self, request: web.Request) -> web.Response: - """Get all LLMs and their current states.""" - states = self._agent.llms - return web.Response( - text=json.dumps( - [{"name": name, "state": state.name} - for name, state in states.items()] - ), - content_type="application/json" - ) async def _get_auto_approver_config(self, request: web.Request) -> web.Response: """Get current auto approver configuration.""" @@ -169,7 +162,7 @@ class Api: """Update auto approver configuration.""" try: data = await request.json() - self._auto_approver.set_config(data) + self._auto_approver.config = AutoApproverConfig(**data) return web.Response(status=200) except Exception as e: return web.Response(status=400, text=str(e)) @@ -219,18 +212,6 @@ class Api: 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.""" - 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 Exception as e: - return web.Response(status=400, text=str(e)) - async def _get_memory(self, request: web.Request) -> web.Response: """Get complete working memory state.""" entries = self._working_memory.get_entries() @@ -304,13 +285,12 @@ class Api: if not content: return web.Response(status=400, text="Missing content in request body") - (context, response, entries) = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer) + (_context, response, entries) = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer) for entry in entries: self._working_memory.add_entry(entry) - self._agent.modify_context(context, True) self._agent.response_buffer.set_text(response) return web.Response(status=200) except Exception as e: - return web.Response(status=400, text=str(e)) + return web.Response(status=400, text=str(e)) \ No newline at end of file diff --git a/sia/web/auto_approver_websocket.py b/sia/web/auto_approver_websocket.py index 27d9bc2..ecaf8ac 100644 --- a/sia/web/auto_approver_websocket.py +++ b/sia/web/auto_approver_websocket.py @@ -1,4 +1,5 @@ from aiohttp import web, WSMsgType +from dataclasses import asdict from typing import Dict, Set from ..auto_approver import AutoApprover, AutoApproverConfig @@ -28,7 +29,7 @@ class AutoApproverWebSocket: async def _handle_config_change(self, config: AutoApproverConfig): """Handle config changes from the AutoApprover.""" await self._broadcast_message({ - "config": config + "config": asdict(config) }) async def handle_connection(self, request: web.Request) -> web.WebSocketResponse: @@ -41,7 +42,7 @@ class AutoApproverWebSocket: try: # Send initial config await ws.send_json({ - "config": self._auto_approver.config + "config": asdict(self._auto_approver.config) }) async for msg in ws: diff --git a/sia/web/chat_websocket.py b/sia/web/chat_websocket.py new file mode 100644 index 0000000..6b537ed --- /dev/null +++ b/sia/web/chat_websocket.py @@ -0,0 +1,64 @@ +from aiohttp import web, WSMsgType +from typing import Dict, Set + +from .util import wrap_async +from ..chat_io_buffer import ChatIOBuffer + +class ChatWebSocket: + + def __init__(self, io_buffer: ChatIOBuffer): + self._io_buffer = io_buffer + self._clients: Set[web.WebSocketResponse] = set() + self._io_buffer.add_change_handler(wrap_async(self._handle_io_buffer_change)) + + async def _broadcast_message(self, message: Dict): + """Broadcast message to all connected clients.""" + disconnected = set() + for ws in self._clients: + try: + await ws.send_json(message) + except ConnectionResetError: + disconnected.add(ws) + self._clients -= disconnected + + async def _handle_io_buffer_change(self, new_messages, deleted_message_ids, last_read_timestamp): + """Handle changes to the ChatIOBuffer.""" + await self._broadcast_message({ + "messages": [ + { + "id": message.id, + "content": message.content, + "message_type": message.message_type.value + } + for message in new_messages + ], + "deleted_message_ids": deleted_message_ids, + "last_read_timestamp": last_read_timestamp + }) + + async def handle_connection(self, request: web.Request) -> web.WebSocketResponse: + """Handle new WebSocket connections.""" + ws = web.WebSocketResponse(heartbeat=30) + await ws.prepare(request) + + try: + # Send initial state + await ws.send_json({ + "messages": [ + { + "id": message.id, + "content": message.content, + "message_type": message.message_type.value + } + for message in self._io_buffer.messages + ], + "last_read_timestamp": self._io_buffer.last_read_timestamp + }) + self._clients.add(ws) + async for msg in ws: + if msg.type == WSMsgType.ERROR: + print(f"WebSocket connection closed with error: {ws.exception()}") + finally: + self._clients.remove(ws) + + return ws \ No newline at end of file diff --git a/sia/web/context_websocket.py b/sia/web/context_websocket.py deleted file mode 100644 index f2ba962..0000000 --- a/sia/web/context_websocket.py +++ /dev/null @@ -1,55 +0,0 @@ -from aiohttp import web, WSMsgType -from typing import Dict, Set - -from .util import wrap_async -from ..web_agent import WebAgent - -class ContextWebSocket: - """ - WebSocket handler for context changes. - Broadcasts context updates to all connected clients. - """ - - def __init__(self, agent: WebAgent): - self._agent = agent - self._clients: Set[web.WebSocketResponse] = set() - self._agent.add_context_change_handler(wrap_async(self._handle_context_change)) - - async def _broadcast_message(self, message: Dict): - """Broadcast message to all connected clients.""" - disconnected = set() - for ws in self._clients: - try: - await ws.send_json(message) - except ConnectionResetError: - disconnected.add(ws) - self._clients -= disconnected - - async def _handle_context_change(self, context: str, generated: bool): - """Handle context changes from the WebAgent.""" - await self._broadcast_message({ - "context": context, - "generated": generated - }) - - async def handle_connection(self, request: web.Request) -> web.WebSocketResponse: - """Handle new WebSocket connections.""" - ws = web.WebSocketResponse(heartbeat=30) - await ws.prepare(request) - - self._clients.add(ws) - - try: - # Send initial context - await ws.send_json({ - "context": self._agent.context, - "generated": True - }) - - async for msg in ws: - if msg.type == WSMsgType.ERROR: - print(f"WebSocket connection closed with error: {ws.exception()}") - finally: - self._clients.remove(ws) - - return ws diff --git a/sia/web/memory_websocket.py b/sia/web/memory_websocket.py index 82ee943..a91b8c5 100644 --- a/sia/web/memory_websocket.py +++ b/sia/web/memory_websocket.py @@ -1,8 +1,6 @@ from aiohttp import web, WSMsgType from typing import Dict, Set -from ..entry import Entry -from ..web_agent import WebAgent from ..working_memory import WorkingMemory from .util import wrap_async diff --git a/sia/web/llm_websocket.py b/sia/web/state_websocket.py similarity index 64% rename from sia/web/llm_websocket.py rename to sia/web/state_websocket.py index ccd9167..4e9447f 100644 --- a/sia/web/llm_websocket.py +++ b/sia/web/state_websocket.py @@ -2,18 +2,19 @@ from aiohttp import web, WSMsgType from typing import Dict, Set from .util import wrap_async -from ..web_agent import WebAgent, LlmState +from ..web_agent import WebAgent, AgentState -class LlmWebSocket: +class StateWebSocket: """ - WebSocket handler for LLM state changes. + WebSocket handler for agent state changes. Broadcasts state updates to all connected clients. """ def __init__(self, agent: WebAgent): self._agent = agent self._clients: Set[web.WebSocketResponse] = set() - self._agent.add_llm_change_handler(wrap_async(self._handle_state_change)) + self._agent.add_state_change_handler(wrap_async(self._handle_change)) + self._agent.add_selected_llm_change_handler(wrap_async(self._handle_change)) async def _broadcast_message(self, message: Dict): """Broadcast message to all connected clients.""" @@ -24,12 +25,12 @@ class LlmWebSocket: except ConnectionResetError: disconnected.add(ws) self._clients -= disconnected - - async def _handle_state_change(self, llm_name: str, new_state: LlmState): - """Handle state changes from the WebAgent.""" + + async def _handle_change(self, arg): + """Handle changes to the active LLM.""" await self._broadcast_message({ - "llm": llm_name, - "state": new_state.name + "state": self._agent.state.name, + "active_llm": self._agent.active_llm }) async def handle_connection(self, request: web.Request) -> web.WebSocketResponse: @@ -38,16 +39,12 @@ class LlmWebSocket: await ws.prepare(request) try: - # Send initial states for all LLMs - states = self._agent.llms - for llm_name, state in states.items(): - await ws.send_json({ - "llm": llm_name, - "state": state.name - }) - + # Send initial state + await ws.send_json({ + "state": self._agent.state.name, + "active_llm": self._agent.active_llm + }) self._clients.add(ws) - async for msg in ws: if msg.type == WSMsgType.ERROR: print(f"WebSocket connection closed with error: {ws.exception()}") diff --git a/sia/web/websockets.py b/sia/web/websockets.py index 1fd6ede..6e79d02 100644 --- a/sia/web/websockets.py +++ b/sia/web/websockets.py @@ -1,28 +1,25 @@ from aiohttp import web from ..auto_approver import AutoApprover +from ..chat_io_buffer import ChatIOBuffer from ..web_agent import WebAgent -from ..web_io_buffer import WebIOBuffer from ..working_memory import WorkingMemory from .auto_approver_websocket import AutoApproverWebSocket -from .context_websocket import ContextWebSocket -from .llm_websocket import LlmWebSocket +from .chat_websocket import ChatWebSocket from .memory_websocket import MemoryWebSocket from .response_websocket import ResponseWebSocket -from .stdout_websocket import StdoutWebSocket +from .state_websocket import StateWebSocket class Websockets: - def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory): + def __init__(self, app: web.Application, agent: WebAgent, io_buffer: ChatIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory): self._auto_approver_ws = AutoApproverWebSocket(auto_approver) - self._context_ws = ContextWebSocket(agent) - self._llm_ws = LlmWebSocket(agent) + self._chat_ws = ChatWebSocket(io_buffer) self._memory_ws = MemoryWebSocket(working_memory) self._response_ws = ResponseWebSocket(agent) - self._stdout_ws = StdoutWebSocket(io_buffer) + self._state_ws = StateWebSocket(agent) app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection) - app.router.add_get("/ws/context", self._context_ws.handle_connection) - app.router.add_get("/ws/llms", self._llm_ws.handle_connection) + app.router.add_get("/ws/chat", self._chat_ws.handle_connection) app.router.add_get("/ws/memory", self._memory_ws.handle_connection) app.router.add_get("/ws/response", self._response_ws.handle_connection) - app.router.add_get("/ws/stdout", self._stdout_ws.handle_connection) + app.router.add_get("/ws/state", self._state_ws.handle_connection) \ No newline at end of file diff --git a/sia/web_agent.py b/sia/web_agent.py index 0003399..0421318 100644 --- a/sia/web_agent.py +++ b/sia/web_agent.py @@ -1,13 +1,10 @@ -from collections import defaultdict from datetime import datetime, timezone from enum import Enum, auto from sys import exit -from threading import Lock -from typing import Callable, Dict, List, Optional +from typing import Callable, Dict, List from .base_agent import BaseAgent from .command import Command -from .command_result import CommandResult from .iteration_logger import IterationLogger from .llm_engine import LlmEngine from .response_buffer import ResponseBuffer @@ -15,9 +12,10 @@ from .response_parser import ResponseParser from .system_metrics import SystemMetrics from .working_memory import WorkingMemory -class LlmState(Enum): +class AgentState(Enum): IDLE = auto() INFERENCE = auto() + PROCESSING_RESPONSE = auto() class WebAgent(BaseAgent): def __init__( @@ -37,97 +35,81 @@ class WebAgent(BaseAgent): metrics, parser ) + self._llms = llms + self._selected_llm = list(llms.keys())[0] + self._state: AgentState = AgentState.IDLE + self._iteration_logger = iteration_logger - self._llm_states: Dict[str, LlmState] = {name: LlmState.IDLE for name in llms} self._response_buffer: ResponseBuffer = ResponseBuffer() - 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() - self._output_lock = Lock() - - # Event handlers - self._llm_change_handlers: List[Callable[[str, LlmState], None]] = [] - self._token_handlers: List[Callable[[str, str], None]] = [] - self._context_change_handlers: List[Callable[[str, bool], None]] = [] + self._state_change_handlers: List[Callable[[AgentState], None]] = [] + self._selected_llm_change_handlers: List[Callable[[str], None]] = [] - # Change handlers - self._working_memory.add_change_handler(self._handle_memory_update) - - @property - def llms(self) -> Dict[str, LlmState]: - """Get current state of all LLMs""" - with self._llm_lock: - return self._llm_states.copy() + self._update_compiled_context() + self._working_memory.add_change_handler(self._update_compiled_context) @property def response_buffer(self) -> ResponseBuffer: return self._response_buffer @property - def context(self) -> str: - return self._context - + def llms(self) -> List[str]: + return self._llms.keys() + @property - def command_result(self) -> Optional[CommandResult]: - return self._command_result - + def active_llm(self) -> str: + return self._selected_llm + + @active_llm.setter + def active_llm(self, llm_name: str) -> None: + if llm_name not in self._llms: + raise ValueError(f"Invalid LLM name: {llm_name}") + if self._selected_llm == llm_name: + return + self._selected_llm = llm_name + self._update_compiled_context() + self._notify_selected_llm_change() + + def add_selected_llm_change_handler(self, handler: Callable[[str], None]) -> None: + self._selected_llm_change_handlers.append(handler) + @property - def validation_error(self) -> Optional[str]: - return self._validation_error + def state(self) -> AgentState: + return self._state - def add_llm_change_handler(self, handler: Callable[[str, LlmState], None]) -> None: - """Add handler for LLM state changes""" - if handler not in self._llm_change_handlers: - self._llm_change_handlers.append(handler) + def add_state_change_handler(self, handler: Callable[[AgentState], None]) -> None: + self._state_change_handlers.append(handler) - def add_context_change_handler(self, handler: Callable[[str, bool], None]) -> None: - """Add handler for context changes""" - if handler not in self._context_change_handlers: - self._context_change_handlers.append(handler) - - def modify_context(self, context: str, generated: bool = False) -> None: - """Update context and reset all LLM states""" - with self._llm_lock: - self._context = context - for handler in self._context_change_handlers: - handler(context, generated) - - def run_inference(self, llm_name: str) -> None: + def run_inference(self) -> None: """Start inference on specified LLM""" - if llm_name not in self._llms: - raise ValueError(f"Unknown LLM: {llm_name}") - with self._llm_lock: - if self._llm_states[llm_name] != LlmState.IDLE: - 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] - def should_stop() -> bool: - return self._stop_flags[llm_name] - response_token_iter = llm.infer(self.system_prompt, self.context, self._response_buffer.get_text(), should_stop) + if self._state != AgentState.IDLE: + raise RuntimeError(f"Not ready for inference") + self._update_state(AgentState.INFERENCE) + llm = self._llms[self.active_llm] + response_token_iter = llm.infer( + self.system_prompt, + self._compiled_context, + self._response_buffer.get_text() + ) for token in response_token_iter: - with self._output_lock: - self._response_buffer.append_text(token) - with self._llm_lock: - self._set_llm_state(llm_name, LlmState.IDLE) + self._response_buffer.append_text(token) + self._update_state(AgentState.IDLE) - def stop_inference(self, llm_name: str) -> None: + def stop_inference(self) -> 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 - with self._llm_lock: - self._set_llm_state(llm_name, LlmState.IDLE) + self._llms[self.active_llm].restart() + self._update_state(AgentState.IDLE) def approve_response(self) -> None: """Process approved response from specified LLM""" + self._update_state(AgentState.PROCESSING_RESPONSE) timestamp = datetime.now(timezone.utc) - self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text()) + self._iteration_logger.log_iteration( + timestamp, + self._compiled_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): @@ -141,14 +123,13 @@ class WebAgent(BaseAgent): parse_result.update() self._working_memory.update() self._working_memory.add_entry(parse_result) + self._update_state(AgentState.IDLE) - def _set_llm_state(self, llm_name: str, state: LlmState) -> None: + def _update_state(self, state: AgentState) -> None: """Update LLM state and notify handlers""" - self._llm_states[llm_name] = state - for handler in self._llm_change_handlers: - handler(llm_name, state) - - 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._state = state + for handler in self._state_change_handlers: + handler(state) + + def _update_compiled_context(self): + self._compiled_context = self._compile_context(self._llms[self.active_llm]) \ No newline at end of file diff --git a/sia/working_memory.py b/sia/working_memory.py index a87a8ff..7fc5171 100644 --- a/sia/working_memory.py +++ b/sia/working_memory.py @@ -85,8 +85,8 @@ class WorkingMemory: def __del__(self): """Clean up all entries when memory is deleted.""" - if hasattr(self, '_entries'): - self.clear() + self._change_handlers.clear() + self.clear() def get_entry(self, id: str) -> Optional[Entry]: """ diff --git a/test/auto_approver_test.py b/test/auto_approver_test.py index dd9e8a7..3e55824 100644 --- a/test/auto_approver_test.py +++ b/test/auto_approver_test.py @@ -4,14 +4,14 @@ from threading import Event import time from sia.auto_approver import AutoApprover -from sia.web_agent import WebAgent, LlmState +from sia.web_agent import WebAgent, AgentState class AutoApproverTest(unittest.TestCase): def setUp(self): """Create mock agent and approver for each test""" self.mock_agent = Mock(spec=WebAgent) # Set up llms attribute as a dictionary - self.mock_agent.llms = {"default": LlmState.IDLE} + self.mock_agent.llms = {"default": AgentState.IDLE} self.approver = AutoApprover(self.mock_agent) def test_timeout_setters(self): @@ -30,7 +30,7 @@ class AutoApproverTest(unittest.TestCase): def test_enable_starts_thread_in_correct_state(self): """Test enabling only starts thread in matching state""" # Mock the agent's llm state - self.mock_agent.llms["default"] = LlmState.IDLE + self.mock_agent.llms["default"] = AgentState.IDLE self.approver.context_enabled = True time.sleep(0.1) # Allow thread to start self.assertIsNotNone(self.approver._context_thread) @@ -40,7 +40,7 @@ class AutoApproverTest(unittest.TestCase): def test_state_change_stops_threads(self): """Test state changes stop irrelevant threads""" # Start in idle state with context thread - self.mock_agent.llms["default"] = LlmState.IDLE + self.mock_agent.llms["default"] = AgentState.IDLE # Mock _stop_context_thread to verify it's called with patch.object(self.approver, '_stop_context_thread') as mock_stop: @@ -48,31 +48,31 @@ class AutoApproverTest(unittest.TestCase): time.sleep(0.1) # Change state and verify stop method was called - self.mock_agent.llms["default"] = LlmState.INFERENCE - self.approver._handle_llm_state_change("default", LlmState.INFERENCE) + self.mock_agent.llms["default"] = AgentState.INFERENCE + self.approver._handle_state_change(AgentState.INFERENCE) mock_stop.assert_called_once() def test_quick_state_cycle(self): """Test thread stops when state changes before timeout""" self.approver.context_timeout = 5.0 - self.mock_agent.llms["default"] = LlmState.IDLE + self.mock_agent.llms["default"] = AgentState.IDLE self.approver.context_enabled = True # Change state quickly time.sleep(0.1) - self.mock_agent.llms["default"] = LlmState.INFERENCE - self.approver._handle_llm_state_change("default", LlmState.INFERENCE) + self.mock_agent.llms["default"] = AgentState.INFERENCE + self.approver._handle_state_change(AgentState.INFERENCE) # Verify no approval happened self.mock_agent.run_inference.assert_not_called() def test_disable_stops_thread(self): """Test disabling stops active thread""" - self.mock_agent.llms["default"] = LlmState.IDLE + self.mock_agent.llms["default"] = AgentState.IDLE # Create a new instance with a fresh mock for this test mock_agent = Mock(spec=WebAgent) - mock_agent.llms = {"default": LlmState.IDLE} + mock_agent.llms = {"default": AgentState.IDLE} test_approver = AutoApprover(mock_agent) # Start with a mocked stop method @@ -91,7 +91,7 @@ class AutoApproverTest(unittest.TestCase): def test_successful_auto_approval(self): """Test thread approves after timeout when conditions met""" self.approver.context_timeout = 0.1 - self.mock_agent.llms["default"] = LlmState.IDLE + self.mock_agent.llms["default"] = AgentState.IDLE self.approver.llm_name = "default" # Use threading Event to properly synchronize the test diff --git a/test/base_agent_test.py b/test/base_agent_test.py index ecac014..5108d9b 100644 --- a/test/base_agent_test.py +++ b/test/base_agent_test.py @@ -8,7 +8,6 @@ from sia.web_io_buffer import WebIOBuffer from sia.working_memory import WorkingMemory from sia.system_metrics import SystemMetrics from sia.llm_engine import LlmEngine -from sia.xml_validator import XMLValidator from sia.response_parser import ResponseParser from sia.base_agent import BaseAgent from sia.command import Command @@ -57,7 +56,6 @@ class BaseAgentTest(unittest.TestCase): self.mock_llm = Mock(spec=LlmEngine) self.mock_llm.token_count.return_value = 100 self.mock_llm.token_limit.return_value = 1000 - self.mock_validator = Mock(spec=XMLValidator) self.io_buffer = WebIOBuffer() self.working_memory = WorkingMemory() self.work_dir = Path("/tmp") # Use a temp directory for tests @@ -79,7 +77,6 @@ class BaseAgentTest(unittest.TestCase): action_schema="test schema", working_memory=self.working_memory, metrics=self.mock_metrics, - validator=self.mock_validator, parser=self.parser ) @@ -100,7 +97,6 @@ class BaseAgentTest(unittest.TestCase): """Test agent initialization and component setup.""" self.assertEqual(self.agent._working_memory, self.working_memory) self.assertEqual(self.agent._metrics, self.mock_metrics) - self.assertEqual(self.agent._validator, self.mock_validator) self.assertEqual(self.agent._parser, self.parser) self.assertEqual(self.agent._action_schema, "test schema") @@ -114,20 +110,19 @@ class BaseAgentTest(unittest.TestCase): context = self.agent._compile_context(self.mock_llm) # Parse context and verify structure - root = ET.fromstring(context) - self.assertEqual(root.tag, "context") + self.assertEqual(context.tag, "context") # Check metrics - self.assertEqual(root.get("memory_used"), "1000") - self.assertEqual(root.get("memory_total"), "2000") - self.assertEqual(root.get("disk_used"), "5000") - self.assertEqual(root.get("disk_total"), "10000") + self.assertEqual(context.get("memory_used"), "1000") + self.assertEqual(context.get("memory_total"), "2000") + self.assertEqual(context.get("disk_used"), "5000") + self.assertEqual(context.get("disk_total"), "10000") # Check context size - self.assertEqual(root.get("context"), "10.0%") + self.assertEqual(context.get("context"), "10.0%") # Check no memory entries - self.assertEqual(len(list(root)), 0) + self.assertEqual(len(list(context)), 0) def test_compile_context_with_entries(self): """Test context compilation with a single entry.""" @@ -135,13 +130,12 @@ class BaseAgentTest(unittest.TestCase): self.agent._working_memory.add_entry(entry) context = self.agent._compile_context(self.mock_llm) - root = ET.fromstring(context) # Check context size reflects one entry - self.assertEqual(root.get("context"), "10.0%") + self.assertEqual(context.get("context"), "10.0%") # Verify entry content - reasoning_elem = root.find("reasoning") + reasoning_elem = context.find("reasoning") self.assertIsNotNone(reasoning_elem) self.assertEqual(reasoning_elem.get("id"), "test-id") self.assertEqual(reasoning_elem.text.strip(), "test reasoning") @@ -157,13 +151,12 @@ class BaseAgentTest(unittest.TestCase): self.agent._working_memory.add_entry(entry) context = self.agent._compile_context(self.mock_llm) - root = ET.fromstring(context) # Check context size reflects three entries - self.assertEqual(root.get("context"), "10.0%") + self.assertEqual(context.get("context"), "10.0%") # Verify entry order maintained - reasoning_elems = root.findall("reasoning") + reasoning_elems = context.findall("reasoning") self.assertEqual(len(reasoning_elems), 3) for i, elem in enumerate(reasoning_elems): self.assertEqual(elem.get("id"), f"id-{i}") diff --git a/test/llm_engine_test.py b/test/llm_engine_test.py new file mode 100644 index 0000000..c059177 --- /dev/null +++ b/test/llm_engine_test.py @@ -0,0 +1,237 @@ +import unittest +import os +import tempfile +import time +from xml.etree import ElementTree as ET + +from sia.llm_engine import LlmEngine + +class TestLlmEngine(unittest.TestCase): + def setUp(self): + """Set up test environment before each test.""" + # Create a temporary bash script to simulate the LLM engine + self.temp_script = self._create_mock_llm_script() + os.chmod(self.temp_script, 0o755) # Make it executable + + # Initialize the LlmEngine with our mock script + self.engine = LlmEngine(self.temp_script, "/root/sia/action_schema.xsd") + + def tearDown(self): + """Clean up after the test.""" + if hasattr(self, 'engine') and self.engine: + self.engine._terminate_process() + + def _create_mock_llm_script(self): + """Create a bash script that simulates an LLM engine.""" + fd, path = tempfile.mkstemp(suffix='.sh') + with os.fdopen(fd, 'w') as f: + # Paste the content from the mock_llm_script_debug artifact + # This is a placeholder - you'll replace this with the actual script content + f.write('''#!/bin/bash + +# Mock LLM engine for testing +# This script simulates an LLM engine subprocess that responds to the three commands: +# , ..., and ... + +# Function to read XML input until a complete closing tag is found +read_xml_input() { + local input="" + local line + local char_count=0 + + # Read until we get a complete XML command + while IFS= read -r line; do + input="${input}${line}" + char_count=$((char_count + ${#line})) + + # Debug the actual content (first 30 chars) + + if [[ "$input" == *""* ]]; then + printf "1024" + printf "\004" # EOT character (hex 04) + return + fi + + if [[ "$input" == *""*""* ]]; then + printf "405" + printf "\004" # EOT character (hex 04) + return + fi + + if [[ "$input" == *""*""* ]]; then + generate_response + return + fi + + done +} + +# Function to generate a response token by token +generate_response() { + printf "<" + sleep 0.01 + + printf "reason" + sleep 0.01 + + printf "ing" + sleep 0.01 + + printf ">" + sleep 0.01 + + printf "This" + sleep 0.01 + + printf " is" + sleep 0.01 + + printf " a" + sleep 0.01 + + printf " test" + sleep 0.01 + + printf " response." + sleep 0.01 + + printf "" + sleep 0.01 + + printf "\004" +} + +# Main loop - keep reading input and responding +iteration=0 +while true; do + iteration=$((iteration + 1)) + read_xml_input +done''') + return path + + def test_token_limit(self): + """Test retrieving the token limit from the LLM engine.""" + limit = self.engine.token_limit() + self.assertEqual(limit, 1024, "Token limit should be 1024") + + def test_token_count(self): + """Test counting tokens in a prompt.""" + system_prompt = "You are a helpful AI assistant." + + # Create a simple context ElementTree + context_et = ET.Element("context") + context_et.set("time", "2024-10-18T12:00:00Z") + context_et.text = "Test context" + + prefix = "" + + count = self.engine.token_count(system_prompt, context_et, prefix) + self.assertEqual(count, 405, "Token count should be 405") + + def test_inference_token_by_token(self): + """Test inference with the LLM engine, ensuring tokens come one by one.""" + system_prompt = "You are a helpful AI assistant." + + # Create a simple context ElementTree + context_et = ET.Element("context") + context_et.set("time", "2024-10-18T12:00:00Z") + context_et.text = "Test context" + + prefix = "" + + # Get the generator for token-by-token output + token_generator = self.engine.infer(system_prompt, context_et, prefix) + + # Collect tokens + response = "" + tokens_received = 0 + + for token in token_generator: + response += token + tokens_received += 1 + # Each "token" should be small (in the mock script, it's up to 10 characters) + self.assertLessEqual(len(token), 10, "Each token should be small") + + # Verify we received multiple tokens + self.assertGreater(tokens_received, 10, "Should receive multiple tokens") + + expected = "This is a test response." + self.assertEqual(response, expected, "Inference response should match expected output") + + def test_multiple_inferences(self): + """Test running multiple inference cycles one after another.""" + system_prompt = "You are a helpful AI assistant." + + # Create a simple context ElementTree + context_et = ET.Element("context") + context_et.set("time", "2024-10-18T12:00:00Z") + context_et.text = "Test context" + + prefix = "" + + # Run first inference + token_generator = self.engine.infer(system_prompt, context_et, prefix) + response = "" + for token in token_generator: + response += token + + expected = "This is a test response." + self.assertEqual(response, expected, "First inference response should match expected output") + + # Add a small delay between inferences to see if it helps + time.sleep(0.1) + + # Run second inference + token_generator = self.engine.infer(system_prompt, context_et, prefix) + response = "" + for token in token_generator: + response += token + + self.assertEqual(response, expected, "Second inference response should match expected output") + + # We should be able to get the token limit again after inference + limit = self.engine.token_limit() + self.assertEqual(limit, 1024, "Token limit should still be 1024 after multiple inferences") + + def test_restart(self): + """Test restarting the LLM engine.""" + # First get a token limit + limit = self.engine.token_limit() + self.assertEqual(limit, 1024, "Token limit should be 1024") + + # Restart the engine + self.engine.restart() + + # Get token limit again + limit = self.engine.token_limit() + self.assertEqual(limit, 1024, "Token limit should still be 1024 after restart") + + # Engine should still work for inference after restart + system_prompt = "You are a helpful AI assistant." + + # Create a simple context ElementTree + context_et = ET.Element("context") + context_et.set("time", "2024-10-18T12:00:00Z") + context_et.text = "Test context" + + prefix = "" + + token_generator = self.engine.infer(system_prompt, context_et, prefix) + response = "" + for token in token_generator: + response += token + + expected = "This is a test response." + self.assertEqual(response, expected, "Inference should work after restart") + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/test/local_llm_engine_test.py b/test/local_llm_engine_test.py deleted file mode 100644 index 6d6c526..0000000 --- a/test/local_llm_engine_test.py +++ /dev/null @@ -1,53 +0,0 @@ -import unittest - -from itertools import tee - -from . import test_data - -from sia.llm_engine import LlmEngine -from sia.llm_engine.local_llm_engine import LocalLlmEngine - -class LocalLlmEngineTest(unittest.TestCase): - def setUp(self): - self.model_path = "/root/models/current" - - @unittest.skip("Takes too long") - def test_initialization(self): - llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024) - self.assertIsInstance(llm_engine, LlmEngine) - - @unittest.skip("Takes too long") - def test_infer(self): - main_context = "This is a test" - llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024) - tokens = llm_engine.infer(test_data.echo_system_prompt, main_context, "") - print_tokens, result_tokens = tee(tokens) - for token in print_tokens: - print(token, end="", flush=True) - result = ''.join(result_tokens) - self.assertEqual(result, f"{main_context}{main_context}") - - @unittest.skip("Takes too long") - def test_token_count(self): - """Test token counting returns reasonable numbers""" - llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024) - - # Test with short inputs - count = llm_engine.token_count("Short prompt", "Brief context") - self.assertGreater(count, 5) - self.assertLess(count, 50) - - # Test with longer inputs - long_prompt = "A detailed system prompt with multiple sentences. " * 5 - long_context = "An extensive context containing various details. " * 10 - count = llm_engine.token_count(long_prompt, long_context) - self.assertGreater(count, 100) - self.assertLess(count, 1000) - - @unittest.skip("Takes too long") - def test_token_limit(self): - """Test token limit is within expected range for modern LLMs""" - llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024) - limit = llm_engine.token_limit() - self.assertGreaterEqual(limit, 2048) - self.assertLessEqual(limit, 1048576) \ No newline at end of file diff --git a/test/util_test.py b/test/util_test.py index d21f388..00648d9 100644 --- a/test/util_test.py +++ b/test/util_test.py @@ -5,37 +5,6 @@ import xml.etree.ElementTree as ET from sia import util class UtilTest(unittest.TestCase): - def test_stop_before_value(self): - # Helper function to create iterator from list - def create_iterator(items: list) -> Iterator[str]: - for item in items: - yield item - - # Test case 1: Stop value in middle of sequence - input_sequence = ['hello', 'world', 'STOP', 'ignored'] - result = list(util.stop_before_value(create_iterator(input_sequence), 'STOP')) - self.assertEqual(result, ['hello', 'world']) - - # Test case 2: Stop value not in sequence - input_sequence = ['hello', 'world'] - result = list(util.stop_before_value(create_iterator(input_sequence), 'STOP')) - self.assertEqual(result, ['hello', 'world']) - - # Test case 3: Stop value at start of sequence - input_sequence = ['STOP', 'ignored'] - result = list(util.stop_before_value(create_iterator(input_sequence), 'STOP')) - self.assertEqual(result, []) - - # Test case 4: Stop value as part of an item - input_sequence = ['hello', 'woSTOPrld', 'ignored'] - result = list(util.stop_before_value(create_iterator(input_sequence), 'STOP')) - self.assertEqual(result, ['hello', 'wo']) - - # Test case 5: Empty sequence - input_sequence = [] - result = list(util.stop_before_value(create_iterator(input_sequence), 'STOP')) - self.assertEqual(result, []) - def test_pretty_print_element_basic(self): """Test basic element printing""" elem = ET.Element('root') diff --git a/test/web_agent_test.py b/test/web_agent_test.py index 63e815a..6562d1b 100644 --- a/test/web_agent_test.py +++ b/test/web_agent_test.py @@ -10,10 +10,9 @@ from sia.web_io_buffer import WebIOBuffer from sia.working_memory import WorkingMemory from sia.system_metrics import SystemMetrics from sia.llm_engine import LlmEngine -from sia.xml_validator import XMLValidator from sia.response_parser import ResponseParser from sia.iteration_logger import IterationLogger -from sia.web_agent import WebAgent, LlmState +from sia.web_agent import WebAgent, AgentState from sia.command import Command from sia.command_result import CommandResult from sia.entry.reasoning_entry import ReasoningEntry @@ -22,7 +21,7 @@ from sia.entry import Entry from sia.response_buffer import ResponseBuffer from pathlib import Path -class MockLlmEngine(LlmEngine): +class MockLlmEngine: """Mock LLM engine for testing.""" def __init__(self, output: str = "test reasoning"): @@ -67,14 +66,10 @@ class WebAgentTest(unittest.TestCase): 'default': MockLlmEngine(), 'alternative': MockLlmEngine("alternative reasoning") } - self.mock_validator = Mock(spec=XMLValidator) self.io_buffer = WebIOBuffer() self.working_memory = WorkingMemory() self.work_dir = Path("/tmp") # Use temp directory for tests - # Mock validator to pass by default - self.mock_validator.validate.return_value = None - # Create parser with IO buffer self.parser = ResponseParser(self.work_dir, self.io_buffer) @@ -98,7 +93,6 @@ class WebAgentTest(unittest.TestCase): working_memory=self.working_memory, metrics=self.mock_metrics, llms=self.mock_llms, - validator=self.mock_validator, parser=self.parser, iteration_logger=self.mock_iteration_logger ) @@ -107,7 +101,7 @@ class WebAgentTest(unittest.TestCase): self.llm_state_changes = [] self.context_changes = [] - def llm_change_handler(self, llm_name: str, state: LlmState): + def llm_change_handler(self, llm_name: str, state: AgentState): """Track LLM state changes for verification.""" self.llm_state_changes.append((llm_name, state)) @@ -124,23 +118,11 @@ class WebAgentTest(unittest.TestCase): # Check LLM states llm_states = self.agent.llms self.assertEqual(len(llm_states), 2) - self.assertEqual(llm_states['default'], LlmState.IDLE) - self.assertEqual(llm_states['alternative'], LlmState.IDLE) + self.assertEqual(llm_states['default'], AgentState.IDLE) + self.assertEqual(llm_states['alternative'], AgentState.IDLE) # Check response buffer self.assertIsInstance(self.agent.response_buffer, ResponseBuffer) - - def test_handler_registration(self): - """Test adding state and context change handlers.""" - self.agent.add_llm_change_handler(self.llm_change_handler) - self.agent.add_context_change_handler(self.context_change_handler) - - # Adding same handler twice should not duplicate - self.agent.add_llm_change_handler(self.llm_change_handler) - self.agent.add_context_change_handler(self.context_change_handler) - - self.assertEqual(len(self.agent._llm_change_handlers), 1) - self.assertEqual(len(self.agent._context_change_handlers), 1) def test_run_inference(self): """Test running inference.""" @@ -153,8 +135,8 @@ class WebAgentTest(unittest.TestCase): # Check state changes self.assertEqual(len(self.llm_state_changes), 2) # IDLE -> INFERENCE -> IDLE - self.assertEqual(self.llm_state_changes[0], ('default', LlmState.INFERENCE)) - self.assertEqual(self.llm_state_changes[1], ('default', LlmState.IDLE)) + self.assertEqual(self.llm_state_changes[0], ('default', AgentState.INFERENCE)) + self.assertEqual(self.llm_state_changes[1], ('default', AgentState.IDLE)) # Verify LLM was called self.assertTrue(self.mock_llms['default'].infer_called) @@ -171,7 +153,7 @@ class WebAgentTest(unittest.TestCase): def test_run_inference_busy_llm(self): """Test running inference on a busy LLM.""" # Set LLM state to INFERENCE - with patch.object(self.agent, '_llm_states', {'default': LlmState.INFERENCE}): + with patch.object(self.agent, '_llm_states', {'default': AgentState.INFERENCE}): with self.assertRaises(RuntimeError): self.agent.run_inference('default') @@ -187,7 +169,6 @@ class WebAgentTest(unittest.TestCase): working_memory=WorkingMemory(), metrics=self.mock_metrics, llms={'test': mock_engine}, - validator=self.mock_validator, parser=self.parser, iteration_logger=self.mock_iteration_logger ) @@ -224,7 +205,6 @@ class WebAgentTest(unittest.TestCase): working_memory=WorkingMemory(), metrics=self.mock_metrics, llms=self.mock_llms, - validator=self.mock_validator, parser=self.parser, iteration_logger=self.mock_iteration_logger ) @@ -254,7 +234,6 @@ class WebAgentTest(unittest.TestCase): working_memory=WorkingMemory(), metrics=self.mock_metrics, llms=self.mock_llms, - validator=self.mock_validator, parser=self.parser, iteration_logger=self.mock_iteration_logger ) @@ -290,7 +269,7 @@ class WebAgentTest(unittest.TestCase): self.assertTrue(mock_command.executed) # Check command result was stored - self.assertEqual(self.agent.command_result, command_result) + self.assertEqual(self.agent._command_result, command_result) # Verify iteration was logged self.mock_iteration_logger.log_iteration.assert_called_once() @@ -329,7 +308,6 @@ class WebAgentTest(unittest.TestCase): working_memory=test_memory, metrics=self.mock_metrics, llms=self.mock_llms, - validator=self.mock_validator, parser=self.parser, iteration_logger=self.mock_iteration_logger ) diff --git a/test/xml_validator_test.py b/test/xml_validator_test.py deleted file mode 100644 index faf5198..0000000 --- a/test/xml_validator_test.py +++ /dev/null @@ -1,116 +0,0 @@ -import unittest - -from sia.xml_validator import XMLValidator - -class XMLValidatorTest(unittest.TestCase): - def setUp(self): - """Set up test schema""" - self.test_schema = """ - - - - - - - - - - - - - - - - - - - """ - - self.validator = XMLValidator(self.test_schema) - - def test_initialization(self): - """Test validator initialization and root element extraction""" - valid_elements = self.validator.get_valid_root_elements() - self.assertEqual(valid_elements, {'test_element', 'simple_element'}) - - def test_invalid_schema(self): - """Test handling invalid schema""" - with self.assertRaises(ValueError): - XMLValidator("invalid schema") - - def test_valid_xml(self): - """Test validation of valid XML""" - xml = 'test content' - result = self.validator.validate(xml) - self.assertIsNone(result) - - # Test without optional attribute - xml = 'test content' - result = self.validator.validate(xml) - self.assertIsNone(result) - - # Test simple element - xml = 'test content' - result = self.validator.validate(xml) - self.assertIsNone(result) - - def test_invalid_root_element(self): - """Test validation with invalid root element""" - xml = 'test content' - result = self.validator.validate(xml) - self.assertIsNotNone(result) - self.assertIn("Invalid root element", result) - - def test_missing_required_attribute(self): - """Test validation with missing required attribute""" - xml = 'test content' - result = self.validator.validate(xml) - self.assertIsNotNone(result) - self.assertIn("Missing required attribute", result) - - def test_invalid_integer_attribute(self): - """Test validation with invalid integer attribute""" - xml = 'test content' - result = self.validator.validate(xml) - self.assertIsNotNone(result) - self.assertIn("Invalid integer value", result) - - def test_unexpected_attribute(self): - """Test validation with unexpected attribute""" - xml = 'test content' - result = self.validator.validate(xml) - self.assertIsNotNone(result) - self.assertIn("Unexpected attribute", result) - - def test_invalid_xml_syntax(self): - """Test validation with invalid XML syntax""" - xml = 'unclosed tag' - result = self.validator.validate(xml) - self.assertIsNotNone(result) - self.assertIn("Invalid XML", result) - - def test_whitespace_handling(self): - """Test validation with whitespace in XML""" - xml = """ - - test content - - """ - result = self.validator.validate(xml) - self.assertIsNone(result) - - def test_empty_content(self): - """Test validation with empty element content""" - xml = '' - result = self.validator.validate(xml) - self.assertIsNone(result) - - xml = '' - result = self.validator.validate(xml) - self.assertIsNone(result) - - def test_special_characters(self): - """Test validation with special characters""" - xml = ' " \' chars]]>' - result = self.validator.validate(xml) - self.assertIsNone(result) \ No newline at end of file diff --git a/tools/gemma_infer/pyproject.toml b/tools/gemma_infer/pyproject.toml new file mode 100644 index 0000000..6754d7b --- /dev/null +++ b/tools/gemma_infer/pyproject.toml @@ -0,0 +1,19 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "gemma_infer" +version = "0.1.0" +requires-python = ">=3.8" + +dependencies = [ + "llama-cpp-python @ git+https://github.com/abetlen/llama-cpp-python.git#egg=llama-cpp-python&env=CMAKE_ARGS=-DLLAMA_BUILD=OFF", + "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils", + "python-dotenv>=1.0.0", + "transformers>=4.0.0", + "xml_schema_validator @ file:///root/sia/lib/xml_schema_validator", +] + +[project.scripts] +gemma_infer = "gemma_infer.__main__:main" \ No newline at end of file diff --git a/tools/gemma_infer/src/gemma_infer/__main__.py b/tools/gemma_infer/src/gemma_infer/__main__.py new file mode 100644 index 0000000..477c2a3 --- /dev/null +++ b/tools/gemma_infer/src/gemma_infer/__main__.py @@ -0,0 +1,42 @@ +from .gemma_llm_engine import GemmaLlmEngine +from dotenv import load_dotenv +from llm_engine_utils.protocol import process +import argparse +import os + +def main(): + load_dotenv() + parser = argparse.ArgumentParser(description='Gemma Inference') + + parser.add_argument( + '--model', + type=str, + default=os.getenv('SIA_GEMMA_MODEL', '/root/models/current/model.gguf'), + help='Model name (default: /root/models/current/model.gguf, env: SIA_GEMMA_MODEL)' + ) + + parser.add_argument( + '--tokenizer', + type=str, + default=os.getenv('SIA_GEMMA_TOKENIZER', '/root/models/current/tokenizer'), + help='Model name (default: /root/models/current/tokenizer, env: SIA_GEMMA_TOKENIZER)' + ) + + parser.add_argument( + '--token-limit', + type=int, + default=os.getenv('SIA_GEMMA_TOKEN_LIMIT', 4096), + help='Token limit (default: 4096, env: SIA_GEMMA_TOKEN_LIMIT)' + ) + + args = parser.parse_args() + + gemma_llm_engine = GemmaLlmEngine( + model=args.model, + tokenizer=args.tokenizer, + token_limit=args.token_limit, + ) + process(gemma_llm_engine) + +if __name__ == "__main__": + main() diff --git a/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py b/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py new file mode 100644 index 0000000..34f077f --- /dev/null +++ b/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py @@ -0,0 +1,76 @@ +import os +os.environ["LLAMA_CPP_LIB_PATH"] = "/usr/local/lib" +os.environ["LD_LIBRARY_PATH"] += ":/usr/local/lib" +os.chdir("/usr/local/lib") + +from llama_cpp import Llama, LogitsProcessorList +from llm_engine_utils import LlmEngine +from pathlib import Path +from transformers import AutoTokenizer +from typing import Iterator +from xml_schema_validator import LlamaCppLogitsProcessor + +class GemmaLlmEngine(LlmEngine): + def __init__( + self, + model: str, + tokenizer: str, + token_limit: int, + ): + self._model = model + self._tokenizer = AutoTokenizer.from_pretrained(tokenizer) + self._token_limit = token_limit + self._llm = Llama( + model_path=model, + n_gpu_layers=100, + n_ctx=token_limit, + #verbose=False, # Disable most logging + ) + + def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]: + xml_schema_text = Path(schema).read_text() + logits_processor = LlamaCppLogitsProcessor(self._tokenizer, xml_schema_text).get_processor() + logits_processor_list = LogitsProcessorList([logits_processor]) + prompt = self._format_messages(system, context, prefix) + stream = self._llm.create_completion( + prompt=prompt, + max_tokens=self._token_limit, + stream=True, + logits_processor=logits_processor_list + ) + for output in stream: + if 'text' in output["choices"][0]: + yield output["choices"][0]['text'] + + def token_count(self, system: str, context: str) -> int: + return len(self._format_messages(system, context, None)) + + def token_limit(self) -> int: + return self._token_limit + + def _format_messages(self, system: str, context: str, prefix: str) -> str: + messages = [ + { + "role": "user", + "content": f"{system}\n\n--- Context ---\n{context}", + }, + { + "role": "assistant", + "content": prefix, + "prefix": True, + }, + ] if prefix else [ + { + "role": "system", + "content": system, + }, + { + "role": "user", + "content": context, + }, + ] + return self._tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True + ).removeprefix("") \ No newline at end of file diff --git a/tools/mistral_infer/pyproject.toml b/tools/mistral_infer/pyproject.toml new file mode 100644 index 0000000..18b1c62 --- /dev/null +++ b/tools/mistral_infer/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "mistral_infer" +version = "0.1.0" +requires-python = ">=3.8" + +dependencies = [ + "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils", + "mistral-common>=1.0.0", + "mistralai>=0.0.7", + "python-dotenv>=1.0.0", +] + +[project.scripts] +mistral_infer = "mistral_infer.__main__:main" \ No newline at end of file diff --git a/tools/mistral_infer/src/mistral_infer/__main__.py b/tools/mistral_infer/src/mistral_infer/__main__.py new file mode 100644 index 0000000..1b02275 --- /dev/null +++ b/tools/mistral_infer/src/mistral_infer/__main__.py @@ -0,0 +1,42 @@ +from .mistral_llm_engine import MistralLlmEngine +from dotenv import load_dotenv +from llm_engine_utils.protocol import process +import argparse +import os + +def main(): + load_dotenv() + parser = argparse.ArgumentParser(description='Mistral API Inference') + + parser.add_argument( + '--model', + type=str, + default=os.getenv('SIA_MISTRAL_MODEL', 'mistral-large-latest'), + help='Model name (default: mistral-large-latest, env: SIA_MISTRAL_MODEL)' + ) + + parser.add_argument( + '--token-limit', + type=int, + default=int(os.getenv('SIA_MISTRAL_TOKEN_LIMIT', 128000)), + help='Token limit for the model (default: 128000, env: SIA_MISTRAL_TOKEN_LIMIT)' + ) + + parser.add_argument( + '--api-key', + type=str, + default=os.getenv('SIA_MISTRAL_API_KEY'), + help='API key for the model (required, env: SIA_MISTRAL_API_KEY)' + ) + + args = parser.parse_args() + + mistral_llm_engine = MistralLlmEngine( + model=args.model, + token_limit=args.token_limit, + api_key=args.api_key, + ) + process(mistral_llm_engine) + +if __name__ == "__main__": + main() diff --git a/sia/llm_engine/mistral_llm_engine.py b/tools/mistral_infer/src/mistral_infer/mistral_llm_engine.py similarity index 60% rename from sia/llm_engine/mistral_llm_engine.py rename to tools/mistral_infer/src/mistral_infer/mistral_llm_engine.py index 984b778..e068fbd 100644 --- a/sia/llm_engine/mistral_llm_engine.py +++ b/tools/mistral_infer/src/mistral_infer/mistral_llm_engine.py @@ -1,74 +1,69 @@ -from typing import Iterator, Optional, Callable -from mistralai import Mistral +from llm_engine_utils import LlmEngine +from llm_engine_utils.iterators import skip_prefix from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage from mistral_common.protocol.instruct.request import ChatCompletionRequest from mistral_common.tokens.tokenizers.mistral import MistralTokenizer - -from . import LlmEngine -from ..util import skip_prefix +from mistralai import Mistral +from pathlib import Path +from typing import Iterator class MistralLlmEngine(LlmEngine): def __init__( self, model: str, - temperature: float, - token_limit: int, api_key: str, + token_limit: int, ): self._model = model - self._temperature = temperature - self._token_limit = token_limit self._api_key = api_key + self._token_limit = token_limit self._client = Mistral(api_key=api_key) self._tokenizer = MistralTokenizer.v3() - def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: + def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]: messages = [ { "role": "system", - "content": system_prompt, + "content": system, }, { "role": "user", - "content": main_context, + "content": context, }, { "role": "assistant", - "content": continuation_text, + "content": prefix, "prefix": True, }, - ] if continuation_text else [ + ] if prefix else [ { "role": "system", - "content": system_prompt, + "content": system, }, { "role": "user", - "content": main_context, + "content": context, }, ] stream_response = self._client.chat.stream( model=self._model, messages=messages, - temperature=self._temperature, + #temperature=self._temperature, ) try: def content_generator(): for chunk in stream_response: - if should_stop(): - stream_response.response.close() - break if content := chunk.data.choices[0].delta.content: yield content - yield from skip_prefix(content_generator(), continuation_text) + yield from skip_prefix(content_generator(), prefix) finally: stream_response.response.close() - def token_count(self, system_prompt: str, main_context: str) -> int: + def token_count(self, system: str, context: str) -> int: messages = [ - SystemMessage(content=system_prompt), - UserMessage(content=main_context), + SystemMessage(content=system), + UserMessage(content=context), ] tokenized = self._tokenizer.encode_chat_completion( @@ -80,4 +75,4 @@ class MistralLlmEngine(LlmEngine): return len(tokenized.tokens) def token_limit(self) -> int: - return self._token_limit + return self._token_limit \ No newline at end of file diff --git a/tools/mistral_train/pyproject.toml b/tools/mistral_train/pyproject.toml new file mode 100644 index 0000000..26d1d58 --- /dev/null +++ b/tools/mistral_train/pyproject.toml @@ -0,0 +1,19 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "mistral_train" +version = "0.1.0" +requires-python = ">=3.8" + +dependencies = [ + "llm_engine_utils[dataset] @ file:///root/sia/lib/llm_engine_utils", + "mistral-common>=1.0.0", + "mistralai>=0.0.7", + "python-dotenv>=1.0.0", + "requests>=2.28.0", +] + +[project.scripts] +mistral_train = "mistral_train.__main__:main" \ No newline at end of file diff --git a/tools/train/train/mistral_api.py b/tools/mistral_train/src/mistral_train/__main__.py similarity index 66% rename from tools/train/train/mistral_api.py rename to tools/mistral_train/src/mistral_train/__main__.py index b7df173..218e0ba 100644 --- a/tools/train/train/mistral_api.py +++ b/tools/mistral_train/src/mistral_train/__main__.py @@ -1,54 +1,11 @@ -#!/root/venvs/train/bin/python -""" -Script for fine-tuning Mistral models for SIA using the Mistral API. -""" -from dataclasses import dataclass from pathlib import Path -import argparse import json import os import sys import tempfile import requests -# Import from our shared library -from .util import TrainingParams, DatasetCreator - -@dataclass -class Config: - def __init__(self): - parser = argparse.ArgumentParser(description='Train SIA model using Mistral API') - parser.add_argument( - '--config', - type=Path, - default=Path('/root/sia/training/config.yaml'), - help='Path to config file' - ) - parser.add_argument( - '--model', - type=str, - default='mistral-large-latest', - help='Base model for fine-tuning' - ) - parser.add_argument( - '--api-key', - type=str, - default=os.environ.get('SIA_MISTRAL_API_KEY'), - help='Mistral API key' - ) - self.args = parser.parse_args() - - @property - def config_path(self) -> Path: - return self.args.config - - @property - def model(self) -> str: - return self.args.model - - @property - def api_key(self) -> str: - return self.args.api_key +from .config import Config def upload_file(api_key: str, file_path: Path) -> str: """Upload a file to the Mistral API and return the file ID""" @@ -97,9 +54,6 @@ def start_finetune_job(api_key: str, model: str, file_id: str, params: sia_train def main(): config = Config() - if not config.api_key: - print("Error: Mistral API key not found. Set SIA_MISTRAL_API_KEY environment variable.") - return 1 training_data, train_params, commit_hash = sia_train_lib.prepare_training_data(config.config_path) diff --git a/tools/mistral_train/src/mistral_train/config.py b/tools/mistral_train/src/mistral_train/config.py new file mode 100644 index 0000000..0f1020b --- /dev/null +++ b/tools/mistral_train/src/mistral_train/config.py @@ -0,0 +1,38 @@ +from pathlib import Path +import argparse +import os + +class Config: + def __init__(self): + parser = argparse.ArgumentParser(description='Train SIA model using Mistral API') + parser.add_argument( + '--config', + type=Path, + default=Path('/root/sia/training/config.yaml'), + help='Path to config file' + ) + parser.add_argument( + '--model', + type=str, + default='mistral-large-latest', + help='Base model for fine-tuning' + ) + parser.add_argument( + '--api-key', + type=str, + default=os.environ.get('SIA_MISTRAL_API_KEY'), + help='Mistral API key' + ) + self.args = parser.parse_args() + + @property + def config_path(self) -> Path: + return self.args.config + + @property + def model(self) -> str: + return self.args.model + + @property + def api_key(self) -> str: + return self.args.api_key \ No newline at end of file diff --git a/sia/llm_engine/hf_llm_engine.py b/tools/old_llm_engines/hf_llm_engine.py similarity index 100% rename from sia/llm_engine/hf_llm_engine.py rename to tools/old_llm_engines/hf_llm_engine.py diff --git a/sia/llm_engine/local_llm_engine.py b/tools/old_llm_engines/local_llm_engine.py similarity index 100% rename from sia/llm_engine/local_llm_engine.py rename to tools/old_llm_engines/local_llm_engine.py diff --git a/sia/llm_engine/openai_llm_engine.py b/tools/old_llm_engines/openai_llm_engine.py similarity index 100% rename from sia/llm_engine/openai_llm_engine.py rename to tools/old_llm_engines/openai_llm_engine.py diff --git a/sia/llm_engine/qwq_llm_engine.ipynb b/tools/old_llm_engines/qwq_llm_engine.ipynb similarity index 100% rename from sia/llm_engine/qwq_llm_engine.ipynb rename to tools/old_llm_engines/qwq_llm_engine.ipynb diff --git a/sia/llm_engine/qwq_llm_engine.py b/tools/old_llm_engines/qwq_llm_engine.py similarity index 100% rename from sia/llm_engine/qwq_llm_engine.py rename to tools/old_llm_engines/qwq_llm_engine.py diff --git a/sia/llm_engine/qwq_vllm.ipynb b/tools/old_llm_engines/qwq_vllm.ipynb similarity index 100% rename from sia/llm_engine/qwq_vllm.ipynb rename to tools/old_llm_engines/qwq_vllm.ipynb diff --git a/tools/train/setup.py b/tools/qwq_train/setup.py similarity index 100% rename from tools/train/setup.py rename to tools/qwq_train/setup.py diff --git a/tools/train/train/__init__.py b/tools/qwq_train/train/__init__.py similarity index 100% rename from tools/train/train/__init__.py rename to tools/qwq_train/train/__init__.py diff --git a/tools/train/train/qwen.ipynb b/tools/qwq_train/train/qwen.ipynb similarity index 100% rename from tools/train/train/qwen.ipynb rename to tools/qwq_train/train/qwen.ipynb diff --git a/tools/train/train/qwq.ipynb b/tools/qwq_train/train/qwq.ipynb similarity index 100% rename from tools/train/train/qwq.ipynb rename to tools/qwq_train/train/qwq.ipynb diff --git a/tools/train/train/qwq.py b/tools/qwq_train/train/qwq.py similarity index 100% rename from tools/train/train/qwq.py rename to tools/qwq_train/train/qwq.py diff --git a/tools/train/bin/train b/tools/train/bin/train deleted file mode 100644 index 868254b..0000000 --- a/tools/train/bin/train +++ /dev/null @@ -1,15 +0,0 @@ -#!/bin/bash - -set -e - -SIA_DIR="/root/sia" -OUTPUT_DIR="${1:-/root/models/$(cd "$SIA_DIR" && git rev-parse HEAD)}" - -if [ -n "$(cd "$SIA_DIR" && git status --porcelain)" ]; then - echo "Uncommitted changes in SIA directory" - #exit 1 -fi - -mkdir -p "$OUTPUT_DIR" - -/root/venvs/train/bin/python -m train.qwq --output-dir "$OUTPUT_DIR" \ No newline at end of file diff --git a/tools/train/readme.md b/tools/train/readme.md deleted file mode 100644 index 965e2e5..0000000 --- a/tools/train/readme.md +++ /dev/null @@ -1,68 +0,0 @@ -# SIA Training Tool - -This tool provides command-line utilities for fine-tuning SIA's language models. - -## Supported Models - -- DeepSeek R1 models (including distilled versions) -- Mistral models - -## Commands - -### train_deepseek - -Fine-tune DeepSeek models using Unsloth optimization. - -```bash -train_deepseek --base-model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --output-dir /root/models/DeepSeek-R1-Distill-Qwen-1.5B -``` - -Options: -- `--config`: Path to training configuration file (default: /root/sia/training/config.yaml) -- `--base-model`: HuggingFace model ID for the base model (required) -- `--output-dir`: Directory to save model (required) -- `--api-key`: HuggingFace API key (optional, will use SIA_HF_API_KEY) - -### train_mistral - -Fine-tune Mistral models using Mistral's API. - -```bash -train_mistral --model mistral-large-latest -``` - -Options: -- `--config`: Path to training configuration file (default: /root/sia/training/config.yaml) -- `--model`: Base model name (default: mistral-large-latest) -- `--api-key`: Mistral API key (optional, will use SIA_MISTRAL_API_KEY) - -## Configuration Format - -The training configuration file (YAML) should include: - -```yaml -model: - system_prompt_path: "/root/sia/system_prompt.md" - action_schema: "/root/sia/action_schema.xsd" -params: - learning_rate: 1e-5 - epochs: 3 -data: - - "/root/sia/training/data_dir1/" - - "/root/sia/training/data_dir2/" -``` - -## Data Format - -Training data should be XML files in the following format: - -```xml - - - - - - - - -``` \ No newline at end of file diff --git a/web/.dockerignore b/web/.dockerignore deleted file mode 100644 index b440869..0000000 --- a/web/.dockerignore +++ /dev/null @@ -1,4 +0,0 @@ -node_modules -dist -.git -*.log \ No newline at end of file diff --git a/web/.gitignore b/web/.gitignore deleted file mode 100644 index bfe46fb..0000000 --- a/web/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -.DS_Store -coverage -dist -node_modules -package-lock.json \ No newline at end of file diff --git a/web/index.html b/web/index.html index 1e04a06..b96c5a2 100644 --- a/web/index.html +++ b/web/index.html @@ -1,12 +1,13 @@ - + + SIA Web Interface
- + \ No newline at end of file diff --git a/web/package.json b/web/package.json index 6ec8d88..8948ad4 100644 --- a/web/package.json +++ b/web/package.json @@ -1,43 +1,27 @@ { - "name": "sia-web", - "private": true, + "name": "sia-web-interface", "version": "0.1.0", "type": "module", "scripts": { "dev": "vite", - "build": "vite build", - "preview": "vite preview", - "test": "jest", - "test:watch": "jest --watch" + "build": "tsc && vite build", + "preview": "vite preview" }, "dependencies": { "@monaco-editor/react": "^4.6.0", - "@radix-ui/react-alert-dialog": "^1.0.5", - "@radix-ui/react-scroll-area": "^1.0.5", - "@radix-ui/react-select": "^2.0.0", - "@radix-ui/react-tabs": "^1.0.4", - "class-variance-authority": "^0.7.0", - "clsx": "^2.1.0", + "monaco-editor": "^0.52.0", "lucide-react": "^0.263.1", "react": "^18.2.0", - "react-dom": "^18.2.0", - "tailwind-merge": "^2.2.0", - "tailwindcss-animate": "^1.0.7" + "react-dom": "^18.2.0" }, "devDependencies": { - "@babel/preset-env": "^7.23.9", - "@babel/preset-react": "^7.23.9", - "@testing-library/jest-dom": "^6.4.2", - "@testing-library/react": "^14.2.1", - "@types/react": "^18.2.55", - "@types/react-dom": "^18.2.19", - "@vitejs/plugin-react": "^4.2.1", - "autoprefixer": "^10.4.17", - "babel-jest": "^29.7.0", - "jest": "^29.7.0", - "jest-environment-jsdom": "^29.7.0", - "postcss": "^8.4.35", - "tailwindcss": "^3.4.1", - "vite": "^5.1.0" + "@types/react": "^18.2.0", + "@types/react-dom": "^18.2.0", + "@vitejs/plugin-react": "^4.2.0", + "autoprefixer": "^10.4.14", + "postcss": "^8.4.31", + "tailwindcss": "^3.3.5", + "typescript": "^5.2.0", + "vite": "^6.3.5" } -} +} \ No newline at end of file diff --git a/web/src/App.tsx b/web/src/App.tsx new file mode 100644 index 0000000..6c815c7 --- /dev/null +++ b/web/src/App.tsx @@ -0,0 +1,77 @@ +import React, { useState } from 'react'; +import useScreenSize from './hooks/useScreenSize'; +import { MobileControls, Ribbon, MobileTaskbar } from './components/Layout'; +import { ResponseEditor } from './components/Response'; +import { MemoryEditor } from './components/Memory'; +import { ChatWindow } from './components/Chat'; +import { WebSocketProvider } from './contexts/WebSocketContext'; +import { MonacoProvider } from './contexts/MonacoContext'; + +const App: React.FC = () => { + const [mobileView, setMobileView] = useState<'controls' | 'memory' | 'response' | 'chat'>('controls'); + const screenSize = useScreenSize(); + + return ( + + +
+ +
+
+ +
+ +
+ +
+ +
+ +
+ +
+ +
+ +
+
+ +
+
+ +
+
+ +
+ +
+
+ +
+
+ {screenSize === 'mobile' && } + + {screenSize !== 'mobile' && } +
+ + + ); +}; + +export default App; \ No newline at end of file diff --git a/web/src/components/App.jsx b/web/src/components/App.jsx deleted file mode 100644 index d233307..0000000 --- a/web/src/components/App.jsx +++ /dev/null @@ -1,336 +0,0 @@ -import React, { useState, useEffect } from 'react'; -import { ContentEditor } from '@/components/editors/ContentEditor'; -import { FloatingButton } from '@/components/FloatingButton'; -import { Header, Tabs } from '@/components/Header'; -import { MemoryEditor } from '@/components/editors/MemoryEditor'; -import { Sidebar } from '@/components/Sidebar'; -import { StandardEditor } from '@/components/editors/StandardEditor'; -import { useWebSocket, WebSocketState } from '@/hooks/useWebSocket'; - -export const LlmState = { - IDLE: 'IDLE', - INFERENCE: 'INFERENCE' -}; - -const App = () => { - // Editor content state - const [generatedContext, setGeneratedContext] = useState(''); - const [modifiedContext, setModifiedContext] = useState(''); - const [generatedResponse, setGeneratedResponse] = useState(''); - const [modifiedResponse, setModifiedResponse] = useState(''); - const [input, setInput] = useState(''); - const [output, setOutput] = useState(''); - - // LLM state - const [llms, setLlms] = useState({}); - const [activeLlm, setActiveLlm] = useState(null); - - // Auto approver state - const [autoApproverConfig, setAutoApproverConfig] = useState({ - context_enabled: false, - response_enabled: false, - context_timeout: 5.0, - response_timeout: 10.0, - llm_name: null - }); - - // UI state - const [activeTab, setActiveTab] = useState(Tabs.CONTEXT); - const [showDiff, setShowDiff] = useState(false); - const [showSidebar, setShowSidebar] = useState(false); - const [autoScroll, setAutoScroll] = useState(false); - - // WebSocket connections - const wsRoot = `${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/ws` - const llmWs = useWebSocket(`${wsRoot}/llms`); - const contextWs = useWebSocket(`${wsRoot}/context`); - const responseWs = useWebSocket(`${wsRoot}/response`); - const stdoutWs = useWebSocket(`${wsRoot}/stdout`); - const autoApproverWs = useWebSocket(`${wsRoot}/auto_approver`); - const memoryWs = useWebSocket(`${wsRoot}/memory`); - const [entries, setEntries] = useState([]); - const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED); - - // Handle llm state changes - useEffect(() => { - llmWs.addMessageHandler((data) => { - setLlms(prevLlms => ({ - ...prevLlms, - [data.llm]: data.state - })); - }); - }, []); - - // Handle context changes - useEffect(() => { - contextWs.addMessageHandler((data) => { - setModifiedContext(data.context); - if (data.generated) { - setGeneratedContext(data.context); - } - }); - }, []); - - // Handle response changes - useEffect(() => { - responseWs.addMessageHandler((data) => { - setModifiedResponse(data.response); - setGeneratedResponse(data.response); - }); - }, []); - - // Handle stdout changes - useEffect(() => { - stdoutWs.addMessageHandler((data) => { - setOutput(data.output); - }); - }, []); - - // Handle auto approver config updates - useEffect(() => { - autoApproverWs.addMessageHandler((data) => { - setAutoApproverConfig(data.config); - }); - }, []); - - // Handle memory state changes - useEffect(() => { - memoryWs.addMessageHandler((data) => { - if (data.type === 'memory_state') { - setEntries(data.entries); - } - }); - }, []); - - // Initialize active llm - useEffect(() => { - if (activeLlm === null && Object.keys(llms).length > 0) { - setActiveLlm(Object.keys(llms)[0]); - } - }, [llms, activeLlm]); - - const handleNextLlm = () => { - const llmNames = Object.keys(llms); - const currentIndex = llmNames.indexOf(activeLlm); - const nextIndex = (currentIndex + 1) % llmNames.length; - setActiveLlm(llmNames[nextIndex]); - }; - - const handleApprove = () => { - fetch('/api/response/approve', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ response: modifiedResponse }) - }); - }; - - const handleInference = () => { - fetch(`/api/inference/${activeLlm}`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - response: modifiedResponse, - context: modifiedContext, - }) - }); - }; - - const handleStop = () => { - // We still specify which LLM to stop inference for - fetch(`/api/inference/${activeLlm}/stop`, { - method: 'POST' - }); - }; - - const handleSendInput = () => { - if (input.trim()) { - fetch('/api/input', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ input: input }) - }); - setInput(''); - } - }; - - const handleClearOutput = () => { - fetch('/api/clear', { method: 'POST' }); - }; - - const handleContextEdit = (context) => { - setContextDirty(true); - // When context is edited, we need to reset all LLM states to IDLE - // since responses are no longer valid - const resetLlms = {}; - for (const llm of Object.keys(llms)) { - resetLlms[llm] = LlmState.IDLE; - } - setLlms(resetLlms); - setModifiedContext(context); - }; - - const handleResponseEdit = (response) => { - setModifiedResponse(response); - }; - - const handleAutoApproverConfigChange = async (config) => { - try { - const response = await fetch('/api/auto_approver/config', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(config) - }); - if (!response.ok) { - throw new Error('Failed to update auto approver config'); - } - } catch (error) { - console.error('Error updating auto approver config:', error); - } - }; - - const handleCreateEntry = async (entry) => { - const response = await fetch('/api/memory/entry', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(entry) - }); - if (!response.ok) throw new Error('Failed to create entry'); - }; - - const handleSaveEntry = async (id, entry) => { - const response = await fetch(`/api/memory/entry/${id}`, { - method: 'PUT', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(entry) - }); - if (!response.ok) throw new Error('Failed to update entry'); - }; - - const handleDeleteEntry = async (id) => { - const response = await fetch(`/api/memory/entry/${id}`, { - method: 'DELETE' - }); - if (!response.ok) throw new Error('Failed to delete entry'); - }; - - const handleResetEntry = async (id) => { - const response = await fetch(`/api/memory/entry/${id}/reset`, { - method: 'POST' - }); - if (!response.ok) throw new Error('Failed to reset entry'); - }; - - const handleUpdateEntry = async (id) => { - const response = await fetch(`/api/memory/entry/${id}/update`, { - method: 'POST' - }); - if (!response.ok) throw new Error('Failed to update entry'); - }; - - // Determine active LLM's current state for UI feedback - const activeLlmState = activeLlm ? llms[activeLlm] : null; - - // Track overall WebSocket connection state - useEffect(() => { - if (llmWs.wsState === WebSocketState.CONNECTED && - contextWs.wsState === WebSocketState.CONNECTED && - responseWs.wsState === WebSocketState.CONNECTED && - stdoutWs.wsState === WebSocketState.CONNECTED && - autoApproverWs.wsState === WebSocketState.CONNECTED && - memoryWs.wsState === WebSocketState.CONNECTED) { - setWsState(WebSocketState.CONNECTED); - } else if (llmWs.wsState === WebSocketState.CONNECTING || - contextWs.wsState === WebSocketState.CONNECTING || - responseWs.wsState === WebSocketState.CONNECTING || - stdoutWs.wsState === WebSocketState.CONNECTING || - autoApproverWs.wsState === WebSocketState.CONNECTING || - memoryWs.wsState === WebSocketState.CONNECTING) { - setWsState(WebSocketState.CONNECTING); - } else { - setWsState(WebSocketState.DISCONNECTED); - } - }, [llmWs.wsState, contextWs.wsState, responseWs.wsState, stdoutWs.wsState, autoApproverWs.wsState, memoryWs.wsState]); - - return ( -
-
- -
-
- {activeTab === Tabs.CONTEXT && ( - - )} - {activeTab === Tabs.RESPONSE && ( - - )} - {activeTab === Tabs.INPUT && ( - setInput(value)} - language="plaintext" - /> - )} - {activeTab === Tabs.OUTPUT && ( - - )} - {activeTab === Tabs.MEMORY && ( - - )} -
- - setShowDiff(!showDiff)} - onSendInput={handleSendInput} - onClearOutput={handleClearOutput} - onLlmChange={setActiveLlm} - onNextLlm={handleNextLlm} - onAutoApproverConfigChange={handleAutoApproverConfigChange} - /> - - setShowSidebar(!showSidebar)} /> -
-
- ); -}; - -export default App; \ No newline at end of file diff --git a/web/src/components/Chat/ChatInput.tsx b/web/src/components/Chat/ChatInput.tsx new file mode 100644 index 0000000..54ebbca --- /dev/null +++ b/web/src/components/Chat/ChatInput.tsx @@ -0,0 +1,62 @@ +import React, { useState, useRef, KeyboardEvent } from 'react'; +import { Send } from 'lucide-react'; + +interface ChatInputProps { + onSendMessage: (message: string) => void; + disabled?: boolean; +} + +const ChatInput: React.FC = ({ onSendMessage, disabled = false }) => { + const [message, setMessage] = useState(''); + const inputRef = useRef(null); + + const handleSend = () => { + if (message.trim() && !disabled) { + onSendMessage(message.trim()); + setMessage(''); + setTimeout(() => { + if (inputRef.current) { + inputRef.current.focus(); + } + }, 0); + } + }; + + const handleKeyDown = (e: KeyboardEvent) => { + if (e.key === 'Enter') { + e.preventDefault(); + handleSend(); + } + }; + + return ( +
+
+ setMessage(e.target.value)} + onKeyDown={handleKeyDown} + placeholder="Type a message..." + className="flex-1 px-3 py-2 border rounded-lg text-sm" + disabled={disabled} + autoFocus + /> + +
+
+ ); +}; + +export default ChatInput; \ No newline at end of file diff --git a/web/src/components/Chat/ChatMessage.tsx b/web/src/components/Chat/ChatMessage.tsx new file mode 100644 index 0000000..b2ae5cd --- /dev/null +++ b/web/src/components/Chat/ChatMessage.tsx @@ -0,0 +1,49 @@ +import React from 'react'; +import { Trash, CheckCheck } from 'lucide-react'; +import { ChatMessage as ChatMessageType } from '../../types'; +import { useWebSocketContext } from '../../contexts/WebSocketContext'; + +type ChatMessageProps = ChatMessageType & { + onDelete: (id: string) => void; +}; + +const ChatMessage: React.FC = ({ + id, + content, + message_type, + onDelete +}) => { + const { lastReadTimestamp } = useWebSocketContext(); + const isUser = message_type === 'user'; + const isRead = lastReadTimestamp !== null && id <= lastReadTimestamp; + + return ( +
+
+
+ {id} +
+ {isUser && ( + + + + )} + +
+
+

{content}

+
+
+ ); +}; + +export default ChatMessage; \ No newline at end of file diff --git a/web/src/components/Chat/ChatWindow.tsx b/web/src/components/Chat/ChatWindow.tsx new file mode 100644 index 0000000..8caab70 --- /dev/null +++ b/web/src/components/Chat/ChatWindow.tsx @@ -0,0 +1,172 @@ +import React, { useState, useRef, useEffect } from 'react'; +import { MessageCircle, X, GripHorizontal } from 'lucide-react'; +import ChatMessage from './ChatMessage'; +import ChatInput from './ChatInput'; +import useScreenSize from '../../hooks/useScreenSize'; +import { useWebSocketContext } from '../../contexts/WebSocketContext'; +import { api } from '../../services/api'; + +interface ChatWindowProps { + className?: string; +} + +const ChatWindow: React.FC = ({ className = '' }) => { + const screenSize = useScreenSize(); + const [chatOpen, setChatOpen] = useState(false); + const [chatHeight, setChatHeight] = useState(400); + const [chatWidth, setChatWidth] = useState(360); + const [isDragging, setIsDragging] = useState(false); + const [isLoading, setIsLoading] = useState(false); + const messagesEndRef = useRef(null); + const startHeightRef = useRef(0); + const startWidthRef = useRef(0); + const startPosRef = useRef<{ x: number, y: number }>({ x: 0, y: 0 }); + + const { messages = [], isConnected } = useWebSocketContext(); + + useEffect(() => { + messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' }); + }, [messages]); + + const handleResizeStart = (e: React.MouseEvent) => { + e.preventDefault(); + setIsDragging(true); + startHeightRef.current = chatHeight; + startWidthRef.current = chatWidth; + startPosRef.current = { x: e.clientX, y: e.clientY }; + }; + + useEffect(() => { + const handleMouseMove = (e: MouseEvent) => { + if (!isDragging) return; + + const deltaY = startPosRef.current.y - e.clientY; + const deltaX = startPosRef.current.x - e.clientX; + + const newHeight = Math.max(300, startHeightRef.current + deltaY); + const newWidth = Math.max(300, startWidthRef.current + deltaX); + + setChatHeight(newHeight); + setChatWidth(newWidth); + }; + + const handleMouseUp = () => { + setIsDragging(false); + }; + + if (isDragging) { + document.addEventListener('mousemove', handleMouseMove); + document.addEventListener('mouseup', handleMouseUp); + } + + return () => { + document.removeEventListener('mousemove', handleMouseMove); + document.removeEventListener('mouseup', handleMouseUp); + }; + }, [isDragging]); + + const handleSendMessage = async (content: string) => { + try { + setIsLoading(true); + await api.addChatMessage({ message: content }); + } catch (error) { + console.error('Failed to send message:', error); + } finally { + setIsLoading(false); + } + }; + + const handleDeleteMessage = async (id: string) => { + try { + await api.deleteChatMessage(id); + } catch (error) { + console.error('Failed to delete message:', error); + } + }; + + if (screenSize === 'mobile') { + return ( +
+
+

Chat

+
+
+ {messages.length === 0 ? ( +
+

No messages yet

+
+ ) : ( + messages.map((message) => ( + + )) + )} +
+
+ +
+ ); + } + + return ( + <> + {!chatOpen && ( + + )} + + {chatOpen && ( +
+
+
+ +

Chat

+
+ +
+ +
+ {messages.length === 0 ? ( +
+

No messages yet

+
+ ) : ( + messages.map((message) => ( + + )) + )} +
+
+ + +
+ )} + + ); +}; + +export default ChatWindow; \ No newline at end of file diff --git a/web/src/components/Chat/index.ts b/web/src/components/Chat/index.ts new file mode 100644 index 0000000..7170533 --- /dev/null +++ b/web/src/components/Chat/index.ts @@ -0,0 +1,4 @@ +// src/components/Chat/index.ts +export { default as ChatWindow } from './ChatWindow'; +export { default as ChatMessage } from './ChatMessage'; +export { default as ChatInput } from './ChatInput'; \ No newline at end of file diff --git a/web/src/components/Controls/AgentSection.tsx b/web/src/components/Controls/AgentSection.tsx new file mode 100644 index 0000000..af69f7a --- /dev/null +++ b/web/src/components/Controls/AgentSection.tsx @@ -0,0 +1,67 @@ +import React, { useState } from 'react'; +import { useWebSocketContext } from '../../contexts/WebSocketContext'; + +interface AgentSectionProps { + className?: string; +} + +const AgentSection: React.FC = ({ className = '' }) => { + const { state, isConnected, llms, setActiveLLM } = useWebSocketContext(); + const [isLoading, setIsLoading] = useState(false); + + const handleLLMChange = async (event: React.ChangeEvent) => { + const newLLM = event.target.value; + try { + setIsLoading(true); + await setActiveLLM(newLLM); + } catch (error) { + console.error('Failed to set active LLM:', error); + } finally { + setIsLoading(false); + } + }; + + const stateColors = { + 'IDLE': 'bg-green-100 text-green-800', + 'INFERENCE': 'bg-yellow-100 text-yellow-800', + 'PROCESSING_RESPONSE': 'bg-blue-100 text-blue-800' + }; + + return ( +
+

Agent

+
+ Connection: + + {isConnected ? 'Connected' : 'Disconnected'} + + + State: + + {state.agentState} + + + Engine: + +
+
+ ); +}; + +export default AgentSection; \ No newline at end of file diff --git a/web/src/components/Controls/AutoApproverSection.tsx b/web/src/components/Controls/AutoApproverSection.tsx new file mode 100644 index 0000000..60d25c0 --- /dev/null +++ b/web/src/components/Controls/AutoApproverSection.tsx @@ -0,0 +1,102 @@ +import React, { useState } from 'react'; +import { useWebSocketContext } from '../../contexts/WebSocketContext'; + +interface AutoApproverSectionProps { + className?: string; +} + +const AutoApproverSection: React.FC = ({ className = '' }) => { + const { autoApprover, setAutoApproverConfig } = useWebSocketContext(); + const [isLoading, setIsLoading] = useState(false); + + const handleCheckboxChange = async (field: 'context_enabled' | 'response_enabled') => { + try { + setIsLoading(true); + await setAutoApproverConfig({ [field]: !autoApprover[field] }); + } catch (error) { + console.error(`Failed to update ${field}:`, error); + } finally { + setIsLoading(false); + } + }; + + const handleTimeoutChange = async (field: 'context_timeout' | 'response_timeout', value: string) => { + const numValue = parseFloat(value); + if (!isNaN(numValue) && numValue >= 0) { + try { + setIsLoading(true); + await setAutoApproverConfig({ [field]: numValue }); + } catch (error) { + console.error(`Failed to update ${field}:`, error); + } finally { + setIsLoading(false); + } + } + }; + + return ( +
+

Auto Approver

+
+
+ + +
+ +
+
+ Context: +
+ handleTimeoutChange('context_timeout', e.target.value)} + step="0.1" + min="0" + className="w-28 sm:w-16 px-1 py-0.5 border rounded text-xs text-right" + disabled={isLoading || autoApprover.context_enabled} + /> + s +
+
+ +
+ Response: +
+ handleTimeoutChange('response_timeout', e.target.value)} + step="0.1" + min="0" + className="w-28 sm:w-16 px-1 py-0.5 border rounded text-xs text-right" + disabled={isLoading || autoApprover.response_enabled} + /> + s +
+
+
+
+
+ ); +}; + +export default AutoApproverSection; \ No newline at end of file diff --git a/web/src/components/Controls/LLMSection.tsx b/web/src/components/Controls/LLMSection.tsx new file mode 100644 index 0000000..55efef0 --- /dev/null +++ b/web/src/components/Controls/LLMSection.tsx @@ -0,0 +1,91 @@ +import React, { useState } from 'react'; +import { Play, StopCircle, Check } from 'lucide-react'; +import { Button } from '../shared'; +import { api } from '../../services/api'; +import { useWebSocketContext } from '../../contexts/WebSocketContext'; + +interface LLMSectionProps { + className?: string; +} + +const LLMSection: React.FC = ({ className = '' }) => { + const { state } = useWebSocketContext(); + const [isInferring, setIsInferring] = useState(false); + const [isStopping, setIsStopping] = useState(false); + const [isApproving, setIsApproving] = useState(false); + + const isInferenceBlocked = state.agentState === 'INFERENCE' || state.agentState === 'PROCESSING_RESPONSE'; + + const handleInfer = async () => { + try { + setIsInferring(true); + await api.startInference(); + } catch (error) { + console.error('Failed to start inference:', error); + } finally { + setIsInferring(false); + } + }; + + const handleStop = async () => { + try { + setIsStopping(true); + await api.stopInference(); + } catch (error) { + console.error('Failed to stop inference:', error); + } finally { + setIsStopping(false); + } + }; + + const handleApprove = async () => { + try { + setIsApproving(true); + await api.approveResponse(); + } catch (error) { + console.error('Failed to approve response:', error); + } finally { + setIsApproving(false); + } + }; + + return ( +
+

LLM

+
+ + + +
+
+ ); +}; + +export default LLMSection; \ No newline at end of file diff --git a/web/src/components/Controls/MemorySection.tsx b/web/src/components/Controls/MemorySection.tsx new file mode 100644 index 0000000..a2ee7f8 --- /dev/null +++ b/web/src/components/Controls/MemorySection.tsx @@ -0,0 +1,126 @@ +import React, { useState, useRef } from 'react'; +import { FileText, Trash, Save } from 'lucide-react'; +import { Button } from '../shared'; +import { api } from '../../services/api'; + +interface MemorySectionProps { + className?: string; +} + +const MemorySection: React.FC = ({ className = '' }) => { + const [isLoading, setIsLoading] = useState(false); + const fileInputRef = useRef(null); + + const handleLoad = async () => { + fileInputRef.current?.click(); + }; + + const handleFileLoad = async (event: React.ChangeEvent) => { + const file = event.target.files?.[0]; + if (!file) return; + + try { + setIsLoading(true); + const content = await file.text(); + await api.loadIteration(content); + console.log('Iteration loaded successfully'); + } catch (error) { + console.error('Failed to load iteration:', error); + alert('Failed to load iteration file'); + } finally { + setIsLoading(false); + if (fileInputRef.current) { + fileInputRef.current.value = ''; + } + } + }; + + const handleClear = async () => { + if (!window.confirm('Are you sure you want to clear all memory entries?')) { + return; + } + + try { + setIsLoading(true); + const memoryState = await api.getMemory(); + const entries = memoryState.entries || []; + + for (const entry of entries) { + await api.deleteMemoryEntry(entry.id); + } + + console.log('Memory cleared successfully'); + } catch (error) { + console.error('Failed to clear memory:', error); + alert('Failed to clear memory'); + } finally { + setIsLoading(false); + } + }; + + const handleUpdate = async () => { + try { + setIsLoading(true); + const memoryState = await api.getMemory(); + const entries = memoryState.entries || []; + + for (const entry of entries) { + await api.updateMemoryEntry(entry.id); + } + + console.log('Memory updated successfully'); + } catch (error) { + console.error('Failed to update memory:', error); + alert('Failed to update memory'); + } finally { + setIsLoading(false); + } + }; + + return ( +
+

Memory

+
+ + + + +
+
+ ); +}; + +export default MemorySection; \ No newline at end of file diff --git a/web/src/components/Controls/index.ts b/web/src/components/Controls/index.ts new file mode 100644 index 0000000..751c18e --- /dev/null +++ b/web/src/components/Controls/index.ts @@ -0,0 +1,4 @@ +export { default as AgentSection } from './AgentSection'; +export { default as AutoApproverSection } from './AutoApproverSection'; +export { default as LLMSection } from './LLMSection'; +export { default as MemorySection } from './MemorySection'; \ No newline at end of file diff --git a/web/src/components/FloatingButton.jsx b/web/src/components/FloatingButton.jsx deleted file mode 100644 index 41a2d69..0000000 --- a/web/src/components/FloatingButton.jsx +++ /dev/null @@ -1,14 +0,0 @@ -import React from 'react'; -import { Button } from '@/components/ui/button'; -import { Menu } from 'lucide-react'; - -export const FloatingButton = ({ onClick }) => ( - -); \ No newline at end of file diff --git a/web/src/components/Header.jsx b/web/src/components/Header.jsx deleted file mode 100644 index 7a99b29..0000000 --- a/web/src/components/Header.jsx +++ /dev/null @@ -1,58 +0,0 @@ -import React from 'react'; -import { WebSocketState } from '../hooks/useWebSocket'; - -export const Tabs = { - CONTEXT: 'CONTEXT', - RESPONSE: 'RESPONSE', - INPUT: 'INPUT', - OUTPUT: 'OUTPUT', - MEMORY: 'MEMORY' -}; - -export const Header = ({ - wsState, - agentState, - activeTab, - setActiveTab, -}) => ( -
-
-
-

SIA Control Interface

-
-
-
- {wsState} -
- {agentState && ( -
- {agentState} -
- )} -
-
-
-
- {Object.values(Tabs).map((tab) => ( - - ))} -
-
-
-); \ No newline at end of file diff --git a/web/src/components/Layout/MobileControls.tsx b/web/src/components/Layout/MobileControls.tsx new file mode 100644 index 0000000..005e995 --- /dev/null +++ b/web/src/components/Layout/MobileControls.tsx @@ -0,0 +1,26 @@ +import React from 'react'; +import { AgentSection, LLMSection, AutoApproverSection, MemorySection } from '../Controls'; + +const MobileControls: React.FC = () => { + return ( +
+

Agent Controls

+ + + +
+ +
+ +
+ +
+ +
+ +
+
+ ); +}; + +export default MobileControls; \ No newline at end of file diff --git a/web/src/components/Layout/MobileTaskbar.tsx b/web/src/components/Layout/MobileTaskbar.tsx new file mode 100644 index 0000000..0fa35b2 --- /dev/null +++ b/web/src/components/Layout/MobileTaskbar.tsx @@ -0,0 +1,53 @@ +import React from 'react'; + +interface MobileTaskbarProps { + mobileView: 'controls' | 'memory' | 'response' | 'chat'; + setMobileView: (view: 'controls' | 'memory' | 'response' | 'chat') => void; +} + +const MobileTaskbar: React.FC = ({ mobileView, setMobileView }) => { + return ( +
+
+ + + + +
+
+ ); +}; + +export default MobileTaskbar; \ No newline at end of file diff --git a/web/src/components/Layout/Ribbon.tsx b/web/src/components/Layout/Ribbon.tsx new file mode 100644 index 0000000..6743b83 --- /dev/null +++ b/web/src/components/Layout/Ribbon.tsx @@ -0,0 +1,30 @@ +import React from 'react'; +import useScreenSize from '../../hooks/useScreenSize'; +import { AgentSection, LLMSection, AutoApproverSection, MemorySection } from '../Controls'; + +const Ribbon: React.FC = () => { + const screenSize = useScreenSize(); + + if (screenSize === 'mobile') return null; + + return ( +
+
+
+ +
+
+ +
+
+ +
+
+ +
+
+
+ ); +}; + +export default Ribbon; \ No newline at end of file diff --git a/web/src/components/Layout/index.ts b/web/src/components/Layout/index.ts new file mode 100644 index 0000000..11b1a02 --- /dev/null +++ b/web/src/components/Layout/index.ts @@ -0,0 +1,3 @@ +export { default as MobileControls } from './MobileControls'; +export { default as Ribbon } from './Ribbon'; +export { default as MobileTaskbar } from './MobileTaskbar'; \ No newline at end of file diff --git a/web/src/components/Memory/MemoryEditor.tsx b/web/src/components/Memory/MemoryEditor.tsx new file mode 100644 index 0000000..eb7377a --- /dev/null +++ b/web/src/components/Memory/MemoryEditor.tsx @@ -0,0 +1,269 @@ +import React, { useState } from 'react'; +import { useWebSocketContext } from '../../contexts/WebSocketContext'; +import { api, EntryTypes, getInitialEntryData } from '../../services/api'; +import { + ReasoningEntry, + SingleEntry, + RepeatEntry, + ReadEntry, + WriteEntry, + ParseErrorEntry, + BackgroundEntry +} from './entries'; +import { + EntryData, + ReasoningEntryData, + ScriptEntryData, + IOEntryData, + ParseErrorEntryData, + BackgroundEntryData +} from '../../types/entries'; +import { Plus } from 'lucide-react'; + +interface MemoryEditorProps { + className?: string; +} + +const MemoryEditor: React.FC = ({ className = '' }) => { + const { memory } = useWebSocketContext(); + const [isLoading, setIsLoading] = useState(false); + const [showNewEntryMenu, setShowNewEntryMenu] = useState(false); + + const handleDelete = async (id: string) => { + if (!window.confirm('Are you sure you want to delete this entry?')) { + return; + } + + try { + setIsLoading(true); + await api.deleteMemoryEntry(id); + } catch (error) { + console.error('Failed to delete entry:', error); + alert('Failed to delete entry'); + } finally { + setIsLoading(false); + } + }; + + const handleUpdate = async (id: string) => { + try { + setIsLoading(true); + await api.updateMemoryEntry(id); + } catch (error) { + console.error('Failed to update entry:', error); + alert('Failed to update entry'); + } finally { + setIsLoading(false); + } + }; + + const handleReset = async (id: string) => { + try { + setIsLoading(true); + await api.resetMemoryEntry(id); + } catch (error) { + console.error('Failed to reset entry:', error); + alert('Failed to reset entry'); + } finally { + setIsLoading(false); + } + }; + + const handleSaveEntry = async (id: string, data: T) => { + try { + setIsLoading(true); + await api.saveMemoryEntry(id, data); + } catch (error) { + console.error('Failed to save entry:', error); + alert('Failed to save entry'); + } finally { + setIsLoading(false); + } + }; + + const handleCreateEntry = async (type: string) => { + try { + setIsLoading(true); + setShowNewEntryMenu(false); + + const entryType = type as typeof EntryTypes[keyof typeof EntryTypes]; + const newEntryData = getInitialEntryData(entryType); + + await api.createMemoryEntry(newEntryData); + } catch (error) { + console.error('Failed to create entry:', error); + alert('Failed to create entry'); + } finally { + setIsLoading(false); + } + }; + + const renderEntry = (entry: EntryData) => { + switch (entry.type) { + case EntryTypes.REASONING: + return ( + + ); + case EntryTypes.SINGLE: + return ( + + ); + case EntryTypes.REPEAT: + return ( + + ); + case EntryTypes.READ_STDIN: + return ( + + ); + case EntryTypes.WRITE: + return ( + + ); + case EntryTypes.PARSE_ERROR: + return ( + + ); + case EntryTypes.BACKGROUND: + return ( + + ); + default: + return null; + } + }; + + const renderNewEntryMenu = () => ( +
+
Create New Entry
+
+ + + + + + + +
+
+ ); + + return ( +
+
+

Memory / Context

+ + {showNewEntryMenu && renderNewEntryMenu()} +
+ +
+ {isLoading && ( +
+
+
+ )} + + {memory.length === 0 ? ( +
+

No memory entries

+
+ ) : ( + memory.map((entry) => renderEntry(entry)) + )} +
+
+ ); +}; + +export default MemoryEditor; \ No newline at end of file diff --git a/web/src/components/Memory/entries/BackgroundEntry.tsx b/web/src/components/Memory/entries/BackgroundEntry.tsx new file mode 100644 index 0000000..84e61b7 --- /dev/null +++ b/web/src/components/Memory/entries/BackgroundEntry.tsx @@ -0,0 +1,158 @@ +import React from 'react'; +import { BaseEntry } from './BaseEntry'; +import { BackgroundEntryData } from '../../../types/entries'; + +interface BackgroundEntryProps { + entry: BackgroundEntryData; + onDelete: (id: string) => void; + onUpdate: (id: string) => void; + onReset: (id: string) => void; + onSave: (id: string, data: BackgroundEntryData) => Promise; +} + +const BackgroundEntry: React.FC = ({ + entry, + onDelete, + onUpdate, + onReset, + onSave +}) => { + const renderContent = (data: BackgroundEntryData) => ( +
+
+ Script: + 
{data.script || ''}
+
+ + {data.pid !== null && data.pid !== undefined && ( +
+ Process ID: + {data.pid} +
+ )} + + {data.stdout && ( +
+ Standard Output: + 
{data.stdout}
+
+ )} + + {data.stderr && ( +
+ Standard Error: + 
{data.stderr}
+
+ )} + + {data.exit_code !== null && data.exit_code !== undefined && ( +
+ Exit Code: + {data.exit_code} +
+ )} +
+ ); + + const renderEditForm = ( + data: BackgroundEntryData, + onChange: (data: BackgroundEntryData) => void, + isLoading: boolean + ) => ( +
+
+ +