From dbc0068a824eb1f20b7a7dc9b6059fc31d04bfa8 Mon Sep 17 00:00:00 2001 From: geens Date: Sat, 23 Nov 2024 15:34:33 +0100 Subject: [PATCH] Server side implementation for interactive entry edits --- sia/__main__.py | 7 +- sia/base_agent.py | 11 +-- sia/{entry.py => entry/__init__.py} | 34 ++++++-- sia/{ => entry}/background_entry.py | 103 +++++++++++----------- sia/entry/entry_factory.py | 120 ++++++++++++++++++++++++++ sia/{ => entry}/parse_error_entry.py | 38 ++++---- sia/{ => entry}/read_entry.py | 50 ++++++----- sia/{ => entry}/reasoning_entry.py | 32 ++++--- sia/entry/repeat_entry.py | 105 +++++++++++++++++++++++ sia/entry/single_entry.py | 114 ++++++++++++++++++++++++ sia/{ => entry}/write_entry.py | 37 ++++---- sia/iteration_logger.py | 6 +- sia/repeat_entry.py | 122 -------------------------- sia/response_parser.py | 58 ++++++------- sia/single_entry.py | 124 --------------------------- sia/system_metrics.py | 86 ++----------------- sia/util.py | 8 +- sia/web/api.py | 68 ++++++++++++++- sia/web/memory_websocket.py | 59 +++++++++++++ sia/web/websockts.py | 16 ++-- sia/web_agent.py | 17 ++-- sia/working_memory.py | 63 +++++++++----- test/base_agent_test.py | 2 +- test/parse_error_entry_test.py | 2 +- test/response_parser_test.py | 4 +- test/web_agent_test.py | 2 +- test/working_memory_test.py | 4 +- test/write_entry_test.py | 10 +-- 28 files changed, 749 insertions(+), 553 deletions(-) rename sia/{entry.py => entry/__init__.py} (56%) rename sia/{ => entry}/background_entry.py (66%) create mode 100644 sia/entry/entry_factory.py rename sia/{ => entry}/parse_error_entry.py (61%) rename sia/{ => entry}/read_entry.py (55%) rename sia/{ => entry}/reasoning_entry.py (69%) create mode 100644 sia/entry/repeat_entry.py create mode 100644 sia/entry/single_entry.py rename sia/{ => entry}/write_entry.py (63%) delete mode 100644 sia/repeat_entry.py delete mode 100644 sia/single_entry.py create mode 100644 sia/web/memory_websocket.py diff --git a/sia/__main__.py b/sia/__main__.py index 60f4cf0..d24c17a 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -64,10 +64,11 @@ class Main: raise ValueError("No LLM engines enabled in configuration") self._io_buffer = WebIOBuffer() + self._working_memory = WorkingMemory() self._agent = WebAgent( system_prompt=self._system_prompt, action_schema=self._action_schema, - working_memory=WorkingMemory(), + working_memory=self._working_memory, metrics=SystemMetrics(), llms=self._llms, validator=XMLValidator(self._action_schema), @@ -77,8 +78,8 @@ class Main: self._auto_approver = AutoApprover(self._agent) self._app = web.Application() - self._api = Api(self._app, self._agent, self._io_buffer, self._auto_approver) - self._websockets = Websockets(self._app, self._agent, self._io_buffer, self._auto_approver) + self._api = Api(self._app, self._agent, self._io_buffer, self._working_memory, self._auto_approver) + self._websockets = Websockets(self._app, self._agent, self._io_buffer, self._auto_approver, self._working_memory) self._static = Static(self._app, self._config) return self diff --git a/sia/base_agent.py b/sia/base_agent.py index b115919..ebaf891 100644 --- a/sia/base_agent.py +++ b/sia/base_agent.py @@ -34,11 +34,6 @@ class BaseAgent(ABC): self._metrics = metrics self._validator = validator self._parser = parser - - def __del__(self): - """Clean up resources on deletion.""" - if hasattr(self, '_metrics'): - self._metrics.stop() @property def system_prompt(self) -> str: @@ -59,14 +54,12 @@ class BaseAgent(ABC): # Create context element context = ET.Element("context") context.set("time", metrics_data["timestamp"]) - context.set("cpu", str(metrics_data["cpu"])) - context.set("gpu", str(metrics_data["gpu"])) context.set("memory_used", str(metrics_data["memory_used"])) context.set("memory_total", str(metrics_data["memory_total"])) context.set("disk_used", str(metrics_data["disk_used"])) context.set("disk_total", str(metrics_data["disk_total"])) context.set("stdin", str(self._parser.io_buffer.buffer_length())) - context.set("context", "100") + context.set("context", "100%") for entry in memory_context: context.append(entry) @@ -79,6 +72,6 @@ class BaseAgent(ABC): context_usage = (float(token_count) / float(token_limit)) * 100.0 # Update context usage metric - context.set("context", str(round(context_usage, 2))) + context.set("context", f"{str(round(context_usage, 2))}%") return pretty_print_element(context) \ No newline at end of file diff --git a/sia/entry.py b/sia/entry/__init__.py similarity index 56% rename from sia/entry.py rename to sia/entry/__init__.py index beadf5e..ea374ec 100644 --- a/sia/entry.py +++ b/sia/entry/__init__.py @@ -1,32 +1,37 @@ from abc import ABC, abstractmethod -from datetime import datetime import xml.etree.ElementTree as ET +from typing import Callable, List class Entry(ABC): """ Abstract base class for all entry types in the working memory. + Provides observable pattern functionality. """ - def __init__(self, id: str, timestamp: datetime): + def __init__(self, id: str): """ Initialize a new entry with provided id and timestamp. Args: id: Unique identifier for this entry - timestamp: Creation timestamp for this entry """ self._id = id - self._timestamp = timestamp + self._change_handlers: List[Callable[['Entry'], None]] = [] @property def id(self) -> str: """Get entry's unique identifier.""" return self._id - @property - def timestamp(self) -> datetime: - """Get entry's creation timestamp.""" - return self._timestamp + def add_change_handler(self, handler: Callable[['Entry'], None]) -> None: + """Add a callback for entry changes.""" + if handler not in self._change_handlers: + self._change_handlers.append(handler) + + def notify_change(self) -> None: + """Notify all handlers of entry state change.""" + for handler in self._change_handlers: + handler(self) @abstractmethod def update(self) -> None: @@ -47,10 +52,23 @@ class Entry(ABC): """ pass + def serialize(self) -> dict[str, str]: + """ + Collect all public attributes of the entry into a dictionary. + """ + pass + def cleanup(self) -> None: """ Clean up any resources used by this entry. Should be overridden by classes that need cleanup. Default implementation does nothing. """ + pass + + def reset(self) -> None: + """ + Reset the entry as if update was never called. + Default implementation does nothing. + """ pass \ No newline at end of file diff --git a/sia/background_entry.py b/sia/entry/background_entry.py similarity index 66% rename from sia/background_entry.py rename to sia/entry/background_entry.py index bdb2132..cd98358 100644 --- a/sia/background_entry.py +++ b/sia/entry/background_entry.py @@ -1,26 +1,24 @@ -from datetime import datetime import subprocess import xml.etree.ElementTree as ET from typing import Optional -from .entry import Entry +from . import Entry class BackgroundEntry(Entry): """ Entry type for long-running background processes. - Private Attributes: - _script: The script/command to execute - _stdout: Captured standard output - _stderr: Captured standard error - _process: The running subprocess.Popen instance - _exit_code: Exit code when process completes + Attributes: + script: The script/command to execute + stdout: Captured standard output + stderr: Captured standard error + process: The running subprocess.Popen instance + exit_code: Exit code when process completes """ def __init__( self, id: str, - timestamp: datetime, script: str, ): """ @@ -28,35 +26,14 @@ class BackgroundEntry(Entry): Args: id: Unique identifier for this entry - timestamp: Creation timestamp for this entry script: The script/command to execute """ - super().__init__(id, timestamp) - self._script = script - self._stdout = "" - self._stderr = "" + super().__init__(id) + self.script = script self._process: Optional[subprocess.Popen] = None - self._exit_code: Optional[int] = None - - @property - def script(self) -> str: - """Get the script/command being executed.""" - return self._script - - @property - def stdout(self) -> str: - """Get the captured standard output.""" - return self._stdout - - @property - def stderr(self) -> str: - """Get the captured standard error.""" - return self._stderr - - @property - def exit_code(self) -> Optional[int]: - """Get the exit code of the process (None if still running).""" - return self._exit_code + self.stdout = "" + self.stderr = "" + self.exit_code = None @property def pid(self) -> Optional[int]: @@ -85,54 +62,71 @@ class BackgroundEntry(Entry): pass # Ignore cleanup errors finally: self._process = None + + def reset(self) -> None: + """ + Cleans up existing process and resets output buffers. + """ + self.cleanup() + self.stdout = "" + self.stderr = "" + self.exit_code = None + self.notify_change() def update(self) -> None: """ Start the process if not running and collect any new output. Updates stdout and stderr with any new output. """ - if self._process is None and self._exit_code is None: + if self._process is None and self.exit_code is None: self._process = subprocess.Popen( - self._script, + self.script, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1 # Line buffered ) + self.notify_change() return + if self._process is None: return + exit_code = self._process.poll() if exit_code is not None: try: remaining_out, remaining_err = self._process.communicate(timeout=0.1) - self._stdout += remaining_out - self._stderr += remaining_err + self.stdout += remaining_out + self.stderr += remaining_err except subprocess.TimeoutExpired: pass # Process didn't finish communicating, try again next update - self._exit_code = exit_code + self.exit_code = exit_code self.cleanup() + self.notify_change() return + if self._process.stdout: while True: try: line = self._process.stdout.readline() if not line: break - self._stdout += line + self.stdout += line except: break + if self._process.stderr: while True: try: line = self._process.stderr.readline() if not line: break - self._stderr += line + self.stderr += line except: break + self.notify_change() def generate_context(self) -> ET.Element: """ @@ -144,12 +138,25 @@ class BackgroundEntry(Entry): element = ET.Element("background", {"id": self._id}) if self._process is not None: element.set("pid", str(self._process.pid)) - elif self._exit_code is not None: - element.set("exit_code", str(self._exit_code)) - element.text = self._script + elif self.exit_code is not None: + element.set("exit_code", str(self.exit_code)) + element.text = self.script stdout_elem = ET.SubElement(element, "stdout") - stdout_elem.text = self._stdout + stdout_elem.text = self.stdout stderr_elem = ET.SubElement(element, "stderr") - stderr_elem.text = self._stderr + stderr_elem.text = self.stderr - return element \ No newline at end of file + return element + + def serialize(self): + """ + Collect all public attributes of the entry into a dictionary. + """ + return { + "type": "background", + "id": self._id, + "script": self.script, + "stdout": self.stdout, + "stderr": self.stderr, + "exit_code": self.exit_code, + } \ No newline at end of file diff --git a/sia/entry/entry_factory.py b/sia/entry/entry_factory.py new file mode 100644 index 0000000..d9a32a9 --- /dev/null +++ b/sia/entry/entry_factory.py @@ -0,0 +1,120 @@ +from typing import Optional + +from . import Entry +from ..io_buffer import IOBuffer +from .background_entry import BackgroundEntry +from .parse_error_entry import ParseErrorEntry +from .read_entry import ReadEntry +from .reasoning_entry import ReasoningEntry +from .repeat_entry import RepeatEntry +from .single_entry import SingleEntry +from .write_entry import WriteEntry + +class EntryFactory: + def create_entry(data: dict, io_buffer: IOBuffer) -> Entry: + entry_type = data.get("type") + entry_id = data.get("id") + + if entry_type == "background": + return BackgroundEntry(entry_id, data["script"]) + + elif entry_type == "read_stdin": + return ReadEntry(entry_id, io_buffer) + + elif entry_type == "reasoning": + return ReasoningEntry(entry_id, data["content"]) + + elif entry_type == "repeat": + return RepeatEntry( + entry_id, + data["script"], + data.get("timeout"), + data.get("limit") + ) + + elif entry_type == "single": + return SingleEntry( + entry_id, + data["script"], + data.get("timeout"), + data.get("limit") + ) + + elif entry_type == "write": + return WriteEntry(entry_id, data["content"], io_buffer) + + raise ValueError(f"Unknown entry type: {entry_type}") + + def update_entry(entry: Entry, data: dict): + if isinstance(entry, Entry): + if "id" in data: + entry.id = data["id"] + + if isinstance(entry, SingleEntry): + if "script" in data: + entry.script = data["script"] + if "timeout" in data: + entry.timeout = float(data["timeout"]) if data["timeout"] else None + if "limit" in data: + entry.limit = int(data["limit"]) if data["limit"] else None + if "stdout" in data: + entry.stdout = data["stdout"] + if "stderr" in data: + entry.stderr = data["stderr"] + if "exit_code" in data: + entry.exit_code = int(data["exit_code"]) if data["exit_code"] else None + if "executed" in data: + entry.executed = bool(data["executed"]) + if "timed_out" in data: + entry.timed_out = bool(data["timed_out"]) + + if isinstance(entry, RepeatEntry): + if "script" in data: + entry.script = data["script"] + if "timeout" in data: + entry.timeout = float(data["timeout"]) if data["timeout"] else None + if "limit" in data: + entry.limit = int(data["limit"]) if data["limit"] else None + if "stdout" in data: + entry.stdout = data["stdout"] + if "stderr" in data: + entry.stderr = data["stderr"] + if "exit_code" in data: + entry.exit_code = int(data["exit_code"]) if data["exit_code"] else None + if "executed" in data: + entry.executed = bool(data["executed"]) + if "timed_out" in data: + entry.timed_out = bool(data["timed_out"]) + + elif isinstance(entry, BackgroundEntry): + if "script" in data: + entry.script = data["script"] + if "stdout" in data: + entry.stdout = data["stdout"] + if "stderr" in data: + entry.stderr = data["stderr"] + if "exit_code" in data: + entry.exit_code = int(data["exit_code"]) if data["exit_code"] else None + + elif isinstance(entry, ParseErrorEntry): + if "content" in data: + entry.content = data["content"] + if "error" in data: + entry.error = data["error"] + + elif isinstance(entry, ReadEntry): + if "content" in data: + entry.content = data["content"] + if "read" in data: + entry.read = bool(data["read"]) + + elif isinstance(entry, ReasoningEntry): + if "content" in data: + entry.content = data["content"] + + elif isinstance(entry, WriteEntry): + if "content" in data: + entry.content = data["content"] + if "written" in data: + entry.written = bool(data["written"]) + entry.notify_change() \ No newline at end of file diff --git a/sia/parse_error_entry.py b/sia/entry/parse_error_entry.py similarity index 61% rename from sia/parse_error_entry.py rename to sia/entry/parse_error_entry.py index 19a3bc0..f2d4be9 100644 --- a/sia/parse_error_entry.py +++ b/sia/entry/parse_error_entry.py @@ -1,7 +1,6 @@ -from datetime import datetime import xml.etree.ElementTree as ET -from .entry import Entry +from . import Entry class ParseErrorEntry(Entry): """ @@ -11,7 +10,6 @@ class ParseErrorEntry(Entry): def __init__( self, id: str, - timestamp: datetime, content: str, error: str, ): @@ -20,23 +18,12 @@ class ParseErrorEntry(Entry): Args: id: Unique identifier for this entry - timestamp: Creation timestamp for this entry content: Original content that failed to parse error: Error message describing the failure """ - super().__init__(id, timestamp) - self._content = content - self._error = error - - @property - def content(self) -> str: - """Get the original content that failed to parse.""" - return self._content - - @property - def error(self) -> str: - """Get the error message describing the failure.""" - return self._error + super().__init__(id) + self.content = content + self.error = error def update(self) -> None: pass @@ -50,8 +37,19 @@ class ParseErrorEntry(Entry): """ element = ET.Element("parse_error", {"id": self._id}) error_elem = ET.SubElement(element, "error") - error_elem.text = self._error + error_elem.text = self.error content_elem = ET.SubElement(element, "content") - content_elem.text = self._content + content_elem.text = self.content - return element \ No newline at end of file + return element + + def serialize(self): + """ + Collect all public attributes of the entry into a dictionary. + """ + return { + "type": "parse_error", + "id": self._id, + "content": self.content, + "error": self.error, + } \ No newline at end of file diff --git a/sia/read_entry.py b/sia/entry/read_entry.py similarity index 55% rename from sia/read_entry.py rename to sia/entry/read_entry.py index 9998334..f841917 100644 --- a/sia/read_entry.py +++ b/sia/entry/read_entry.py @@ -1,52 +1,39 @@ -from datetime import datetime import xml.etree.ElementTree as ET -from .entry import Entry +from . import Entry +from ..io_buffer import IOBuffer class ReadEntry(Entry): """ Entry type for reading content from standard input. - - Attributes: - _content: Content read from stdin, empty until first update - _io_buffer: Buffer to use for IO operations - _read: Whether content has been read """ def __init__( self, id: str, - timestamp: datetime, - io_buffer, + io_buffer: IOBuffer, ): """ Initialize a new read entry. Args: id: Unique identifier for this entry - timestamp: Creation timestamp for this entry io_buffer: Buffer to use for IO operations """ - super().__init__(id, timestamp) - self._content = "" + super().__init__(id) + self.content: str = "", + self.read = False self._io_buffer = io_buffer - self._read = False - - @property - def content(self) -> str: - """ - Get the content read from stdin. - """ - return self._content def update(self) -> None: """ Read from stdin if not already read. Uses the provided IO buffer for the actual read operation. """ - if not self._read: - self._content = self._io_buffer.read() - self._read = True + if not self.read: + self.content = self._io_buffer.read() + self.read = True + self.notify_change() def generate_context(self) -> ET.Element: """ @@ -56,7 +43,18 @@ class ReadEntry(Entry): ET.Element: XML element containing the entry's data """ element = ET.Element("read_stdin", {"id": self._id}) - if self._read: - element.text = self._content + if self.read: + element.text = self.content - return element \ No newline at end of file + return element + + def serialize(self): + """ + Collect all public attributes of the entry into a dictionary. + """ + return { + "type": "read_stdin", + "id": self._id, + "content": self.content, + "read": self.read, + } \ No newline at end of file diff --git a/sia/reasoning_entry.py b/sia/entry/reasoning_entry.py similarity index 69% rename from sia/reasoning_entry.py rename to sia/entry/reasoning_entry.py index 2edef53..439e149 100644 --- a/sia/reasoning_entry.py +++ b/sia/entry/reasoning_entry.py @@ -1,20 +1,15 @@ -from datetime import datetime import xml.etree.ElementTree as ET -from .entry import Entry +from . import Entry class ReasoningEntry(Entry): """ Entry type for agent reasoning steps. - - Attributes: - _content: The reasoning text """ def __init__( self, id: str, - timestamp: datetime, content: str, ): """ @@ -25,15 +20,8 @@ class ReasoningEntry(Entry): timestamp: Creation timestamp for this entry content: The reasoning text """ - super().__init__(id, timestamp) - self._content = content - - @property - def content(self) -> str: - """ - Get the reasoning text. - """ - return self._content + super().__init__(id) + self.content = content def update(self) -> None: """No update needed for reasoning entries.""" @@ -47,6 +35,16 @@ class ReasoningEntry(Entry): ET.Element: XML element containing the entry's data """ element = ET.Element("reasoning", {"id": self._id}) - element.text = self._content + element.text = self.content - return element \ No newline at end of file + return element + + def serialize(self): + """ + Collect all public attributes of the entry into a dictionary. + """ + return { + "type": "reasoning", + "id": self._id, + "content": self.content, + } \ No newline at end of file diff --git a/sia/entry/repeat_entry.py b/sia/entry/repeat_entry.py new file mode 100644 index 0000000..1929a0d --- /dev/null +++ b/sia/entry/repeat_entry.py @@ -0,0 +1,105 @@ +import subprocess +import xml.etree.ElementTree as ET +from typing import Optional + +from . import Entry + +class RepeatEntry(Entry): + """ + Entry type for scripts that are executed on every update. + """ + + default_timeout = 1 + default_limit = 1024 + + def __init__( + self, + id: str, + script: str, + timeout: Optional[float] = None, + limit: Optional[int] = None, + ): + """ + Initialize a new repeat entry. + + Args: + id: Unique identifier for this entry + script: The script/command to execute + timeout: Maximum time to wait for script execution + limit: Maximum number of characters to capture from stdout/stderr + """ + super().__init__(id) + self.script = script + self.timeout = timeout + self.limit = limit + self.stdout: str = "", + self.stderr: str = "", + self.exit_code: Optional[int] = None, + self.timed_out: bool = False + + def update(self) -> None: + """ + Execute the script and update the output. + Captures stdout, stderr and exit code from each execution. + """ + try: + process = subprocess.run( + self.script, + timeout=(self.timeout or self.default_timeout), + shell=True, + capture_output=True, + text=True + ) + self.stdout = process.stdout + self.stderr = process.stderr + self.exit_code = process.returncode + self._timed_out = False + except subprocess.TimeoutExpired as e: + self._timed_out = True + self.notify_change() + + def generate_context(self) -> ET.Element: + """ + Generate an XML Element representing this repeat entry. + + Returns: + ET.Element: XML element containing the entry's data + """ + element = ET.Element("repeat", {"id": self.id}) + if self.timeout: + element.set("timeout", str(self.timeout)) + element.text = self.script + if self._timed_out: + element.set("timed_out", "true") + elif self.exit_code is not None: + element.set("exit_code", str(self.exit_code)) + if self.limit: + element.set("limit", str(self.limit)) + if len(self.stdout) > (self.limit or self.default_limit): + element.set("stdout_truncated", "true") + element.set("stdout_length", str(len(self.stdout))) + stdout_elem = ET.SubElement(element, "stdout") + stdout_elem.text = self.stdout[:(self.limit or self.default_limit)] + if len(self.stderr) > (self.limit or self.default_limit): + element.set("stderr_truncated", "true") + element.set("stderr_length", str(len(self.stderr))) + stderr_elem = ET.SubElement(element, "stderr") + stderr_elem.text = self.stderr[:(self.limit or self.default_limit)] + + return element + + def serialize(self): + """ + Collect all public attributes of the entry into a dictionary. + """ + return { + "type": "repeat", + "id": self.id, + "script": self.script, + "timeout": self.timeout, + "limit": self.limit, + "stdout": self.stdout, + "stderr": self.stderr, + "exit_code": self.exit_code, + "timed_out": self._timed_out + } diff --git a/sia/entry/single_entry.py b/sia/entry/single_entry.py new file mode 100644 index 0000000..b066a28 --- /dev/null +++ b/sia/entry/single_entry.py @@ -0,0 +1,114 @@ +import subprocess +import xml.etree.ElementTree as ET +from typing import Optional + +from . import Entry + +class SingleEntry(Entry): + """ + Entry type for one-time script executions. + """ + + default_timeout = 1 + default_limit = 1024 + + def __init__( + self, + id: str, + script: str, + timeout: Optional[float] = None, + limit: Optional[int] = None, + ): + """ + Initialize a new single shot entry. + + Args: + id: Unique identifier for this entry + script: The script/command to execute + timeout: Maximum time to wait for script execution + limit: Maximum number of characters to capture from stdout/stderr + """ + super().__init__(id) + self.script = script + self.timeout = timeout + self.limit = limit + self.stdout: str = "", + self.stderr: str = "", + self.exit_code: Optional[int] = None + self.executed: bool = False, + self.timed_out: bool = False + + def reset(self) -> None: + """Reset execution state to allow running again.""" + self.executed = False + self.timed_out = False + self.stdout = "" + self.stderr = "" + self.exit_code = None + + def update(self) -> None: + """ + Execute the script if not already executed. + Captures stdout, stderr and exit code. + """ + if self.executed: + return + self.executed = True + try: + process = subprocess.run( + self.script, + timeout=(self.timeout or self.default_timeout), + shell=True, + capture_output=True, + text=True + ) + self.stdout = process.stdout + self.stderr = process.stderr + self.exit_code = process.returncode + except subprocess.TimeoutExpired as e: + self.timed_out = True + self.notify_change() + + def generate_context(self) -> ET.Element: + """ + Generate an XML Element representing this single shot entry. + + Returns: + ET.Element: XML element containing the entry's data + """ + element = ET.Element("single", {"id": self.id}) + if self.timeout: + element.set("timeout", str(self.timeout)) + element.text = self.script + if self.timed_out: + element.set("timed_out", "true") + elif self.executed: + element.set("exit_code", str(self.exit_code)) + if self.limit: + element.set("limit", str(self.limit)) + if len(self.stdout) > (self.limit or self.default_limit): + element.set("stdout_truncated", "true") + element.set("stdout_length", str(len(self.stdout))) + stdout_elem = ET.SubElement(element, "stdout") + stdout_elem.text = self.stdout[:(self.limit or self.default_limit)] + if len(self.stderr) > (self.limit or self.default_limit): + element.set("stderr_truncated", "true") + element.set("stderr_length", str(len(self.stderr))) + stderr_elem = ET.SubElement(element, "stderr") + stderr_elem.text = self.stderr[:(self.limit or self.default_limit)] + return element + + def serialize(self): + """ + Collect all public attributes of the entry into a dictionary. + """ + return { + "type": "single", + "id": self.id, + "script": self.script, + "timeout": self.timeout, + "limit": self.limit, + "stdout": self.stdout, + "stderr": self.stderr, + "exit_code": self.exit_code, + } \ No newline at end of file diff --git a/sia/write_entry.py b/sia/entry/write_entry.py similarity index 63% rename from sia/write_entry.py rename to sia/entry/write_entry.py index 422a8fe..e4a92b4 100644 --- a/sia/write_entry.py +++ b/sia/entry/write_entry.py @@ -1,23 +1,16 @@ -from datetime import datetime import xml.etree.ElementTree as ET -from .entry import Entry -from .io_buffer import IOBuffer +from . import Entry +from ..io_buffer import IOBuffer class WriteEntry(Entry): """ Entry type for writing content to standard output. - - Attributes: - content: Text content to write to stdout - io_buffer: Buffer to use for IO operations - _written: Whether the content has been written """ def __init__( self, id: str, - timestamp: datetime, content: str, io_buffer: IOBuffer, ): @@ -26,23 +19,23 @@ class WriteEntry(Entry): Args: id: Unique identifier for this entry - timestamp: Creation timestamp for this entry content: Text to write to stdout io_buffer: Buffer to use for IO operations """ - super().__init__(id, timestamp) + super().__init__(id) self.content = content - self.io_buffer = io_buffer - self._written = False + self.written: bool = False + self._io_buffer = io_buffer def update(self) -> None: """ Write the content to stdout if not already written. Uses the provided IO buffer for the actual write operation. """ - if not self._written: - self.io_buffer.write(self.content) - self._written = True + if not self.written: + self._io_buffer.write(self.content) + self.written = True + self.notify_change() def generate_context(self) -> ET.Element: """ @@ -53,4 +46,14 @@ class WriteEntry(Entry): """ element = ET.Element("write_stdout", {"id": self.id}) element.text = self.content - return element \ No newline at end of file + return element + + def serialize(self): + """ + Collect all public attributes of the entry into a dictionary. + """ + return { + "type": "write", + "id": self._id, + "content": self.content, + } diff --git a/sia/iteration_logger.py b/sia/iteration_logger.py index 3e68389..6256ebc 100644 --- a/sia/iteration_logger.py +++ b/sia/iteration_logger.py @@ -3,6 +3,8 @@ from pathlib import Path import xml.etree.ElementTree as ET import hashlib +from .util import format_timestamp + class IterationLogger: """Logs agent iterations to XML files""" @@ -20,6 +22,7 @@ class IterationLogger: def log_iteration( self, + timestamp: datetime, context: str, response: str, ): @@ -30,8 +33,7 @@ class IterationLogger: context: The context as ElementTree response: Raw response from LLM """ - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3] - filename = f"iteration_{timestamp}.xml" + filename = f"iteration_{format_timestamp(timestamp)}.xml" filepath = self.iterations_dir / filename root = ET.Element("iteration") diff --git a/sia/repeat_entry.py b/sia/repeat_entry.py deleted file mode 100644 index 38979da..0000000 --- a/sia/repeat_entry.py +++ /dev/null @@ -1,122 +0,0 @@ -from datetime import datetime -import subprocess -import xml.etree.ElementTree as ET -from typing import Optional - -from .entry import Entry - -class RepeatEntry(Entry): - """ - Entry type for scripts that are executed on every update. - - Attributes: - script: The script/command to execute - timeout: Maximum time to wait for script execution - limit: Maximum number of characters to capture from stdout/stderr - stdout: Captured standard output from script execution - stderr: Captured standard error from script execution - exit_code: Process exit code after completion - _executed: Whether the script has been executed - """ - - default_timeout = 1 - default_limit = 1024 - - def __init__( - self, - id: str, - timestamp: datetime, - script: str, - timeout: Optional[float], - limit: Optional[int], - ): - """ - Initialize a new repeat entry. - - Args: - id: Unique identifier for this entry - timestamp: Creation timestamp for this entry - script: The script/command to execute - timeout: Maximum time to wait for script execution - limit: Maximum number of characters to capture from stdout/stderr - """ - super().__init__(id, timestamp) - self._script = script - self._timeout = timeout - self._limit = limit - self._stdout = "" - self._stderr = "" - self._exit_code: Optional[int] = None - self._timed_out = False - - @property - def script(self) -> str: - """Get the script/command being executed.""" - return self._script - - @property - def stdout(self) -> str: - """Get the captured standard output.""" - return self._stdout - - @property - def stderr(self) -> str: - """Get the captured standard error.""" - return self._stderr - - @property - def exit_code(self) -> Optional[int]: - """Get the exit code of the process (None if still running).""" - return self._exit_code - - def update(self) -> None: - """ - Execute the script and update the output. - Captures stdout, stderr and exit code from each execution. - """ - try: - process = subprocess.run( - self._script, - timeout=(self._timeout or self.default_timeout), - shell=True, - capture_output=True, - text=True - ) - self._stdout = process.stdout - self._stderr = process.stderr - self._exit_code = process.returncode - except subprocess.TimeoutExpired as e: - self._timed_out = True - - - def generate_context(self) -> ET.Element: - """ - Generate an XML Element representing this repeat entry. - - Returns: - ET.Element: XML element containing the entry's data - """ - element = ET.Element("repeat", { - "id": self.id, - }) - if self._timeout: - element.set("timeout", str(self._timeout)) - element.text = self.script - if self._timed_out: - element.set("timed_out", "true") - elif self._exit_code is not None: - element.set("exit_code", str(self._exit_code)) - if self._limit: - element.set("limit", str(self._limit)) - if len(self._stdout) > (self._limit or self.default_limit): - element.set("stdout_truncated", "true") - element.set("stdout_length", str(len(self._stdout))) - stdout_elem = ET.SubElement(element, "stdout") - stdout_elem.text = self._stdout[:(self._limit or self.default_limit)] - if len(self._stderr) > (self._limit or self.default_limit): - element.set("stderr_truncated", "true") - element.set("stderr_length", str(len(self._stderr))) - stderr_elem = ET.SubElement(element, "stderr") - stderr_elem.text = self._stderr[:(self._limit or self.default_limit)] - - return element \ No newline at end of file diff --git a/sia/response_parser.py b/sia/response_parser.py index a701e8e..017d7da 100644 --- a/sia/response_parser.py +++ b/sia/response_parser.py @@ -1,21 +1,22 @@ from datetime import datetime -import uuid import xml.etree.ElementTree as ET from typing import Union +from .util import format_timestamp + from .command import Command from .delete_command import DeleteCommand from .stop_command import StopCommand -from .background_entry import BackgroundEntry +from .entry.background_entry import BackgroundEntry from .entry import Entry from .io_buffer import IOBuffer -from .parse_error_entry import ParseErrorEntry -from .read_entry import ReadEntry -from .reasoning_entry import ReasoningEntry -from .repeat_entry import RepeatEntry -from .single_entry import SingleEntry -from .write_entry import WriteEntry +from .entry.parse_error_entry import ParseErrorEntry +from .entry.read_entry import ReadEntry +from .entry.reasoning_entry import ReasoningEntry +from .entry.repeat_entry import RepeatEntry +from .entry.single_entry import SingleEntry +from .entry.write_entry import WriteEntry class ResponseParser: """ @@ -48,7 +49,7 @@ class ResponseParser: """ return self._io_buffer - def parse(self, xml: str) -> Union[Command, Entry]: + def parse(self, timestamp: datetime, xml: str) -> Union[Command, Entry]: """ Parse XML response into a Command or Entry. @@ -59,8 +60,7 @@ class ResponseParser: Command or Entry based on the XML content ParseErrorEntry for any parsing or validation errors """ - entry_id = str(uuid.uuid4()) - timestamp = datetime.now() + entry_id = format_timestamp(timestamp) parser = ET.XMLPullParser(events=("start", "end")) parser.feed(xml) @@ -70,15 +70,15 @@ class ResponseParser: if event == "start": break except ET.ParseError as e: - return ParseErrorEntry(entry_id, timestamp, xml, f"Invalid XML: {str(e)}") + return ParseErrorEntry(entry_id, xml, f"Invalid XML: {str(e)}") try: if root.tag == 'delete': target_id = root.get('id') if not target_id: - return ParseErrorEntry(entry_id, timestamp, xml, "Delete command missing required 'id' attribute") + return ParseErrorEntry(entry_id, xml, "Delete command missing required 'id' attribute") if len(root.attrib) > 1: - return ParseErrorEntry(entry_id, timestamp, xml, "Delete command should only have 'id' attribute") + return ParseErrorEntry(entry_id, xml, "Delete command should only have 'id' attribute") return DeleteCommand(target_id) elif root.tag == 'stop': @@ -86,46 +86,46 @@ class ResponseParser: elif root.tag == 'background': if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '': - return ParseErrorEntry(entry_id, timestamp, xml, "Background entry requires (only) script content") - return BackgroundEntry(entry_id, timestamp, root.text) + return ParseErrorEntry(entry_id, xml, "Background entry requires (only) script content") + return BackgroundEntry(entry_id, root.text) elif root.tag == 'repeat': if len(root) != 0 or root.text is None or root.text.strip() == '': - return ParseErrorEntry(entry_id, timestamp, xml, "Repeat entry requires (only) script content") + return ParseErrorEntry(entry_id, xml, "Repeat entry requires (only) script content") if len(root.attrib) > 2 or any(k not in ('timeout', 'limit') for k in root.attrib): - return ParseErrorEntry(entry_id, timestamp, xml, "Repeat entry only accepts 'timeout' and 'limit' attributes") + return ParseErrorEntry(entry_id, xml, "Repeat entry only accepts 'timeout' and 'limit' attributes") timeout = root.get('timeout') limit = root.get('limit') timeout = float(timeout) if timeout is not None else None limit = int(limit) if limit is not None else None - return RepeatEntry(entry_id, timestamp, root.text, timeout, limit) + return RepeatEntry(entry_id, root.text, timeout, limit) elif root.tag == 'single': if len(root) != 0 or root.text is None or root.text.strip() == '': - return ParseErrorEntry(entry_id, timestamp, xml, "Single entry requires (only) script content") + return ParseErrorEntry(entry_id, xml, "Single entry requires (only) script content") if len(root.attrib) > 2 or any(k not in ('timeout', 'limit') for k in root.attrib): - return ParseErrorEntry(entry_id, timestamp, xml, "Single entry only accepts 'timeout' and 'limit' attributes") + return ParseErrorEntry(entry_id, xml, "Single entry only accepts 'timeout' and 'limit' attributes") timeout = root.get('timeout') limit = root.get('limit') timeout = float(timeout) if timeout is not None else None limit = int(limit) if limit is not None else None - return SingleEntry(entry_id, timestamp, root.text, timeout, limit) + return SingleEntry(entry_id, root.text, timeout, limit) elif root.tag == 'reasoning': if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '': - return ParseErrorEntry(entry_id, timestamp, xml, "Reasoning entry requires (only) text content") - return ReasoningEntry(entry_id, timestamp, root.text) + return ParseErrorEntry(entry_id, xml, "Reasoning entry requires (only) text content") + return ReasoningEntry(entry_id, root.text) elif root.tag == 'read_stdin': - return ReadEntry(entry_id, timestamp, self._io_buffer) + return ReadEntry(entry_id, self._io_buffer) elif root.tag == 'write_stdout': if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '': - return ParseErrorEntry(entry_id, timestamp, xml, "Write stdout entry requires (only) text content") - return WriteEntry(entry_id, timestamp, root.text, self._io_buffer) + return ParseErrorEntry(entry_id, xml, "Write stdout entry requires (only) text content") + return WriteEntry(entry_id, root.text, self._io_buffer) else: - return ParseErrorEntry(entry_id, timestamp, xml, f"Unknown root element: {root.tag}") + return ParseErrorEntry(entry_id, xml, f"Unknown root element: {root.tag}") except Exception as e: - return ParseErrorEntry(entry_id, timestamp, xml, f"Error parsing response: {str(e)}") \ No newline at end of file + return ParseErrorEntry(entry_id, xml, f"Error parsing response: {str(e)}") diff --git a/sia/single_entry.py b/sia/single_entry.py deleted file mode 100644 index 8650eea..0000000 --- a/sia/single_entry.py +++ /dev/null @@ -1,124 +0,0 @@ -from datetime import datetime -import subprocess -import xml.etree.ElementTree as ET -from typing import Optional - -from .entry import Entry - -class SingleEntry(Entry): - """ - Entry type for one-time script executions. - - Attributes: - script: The script/command to execute - timeout: Maximum time to wait for script execution - limit: Maximum number of characters to capture from stdout/stderr - stdout: Captured standard output from script execution - stderr: Captured standard error from script execution - exit_code: Process exit code after completion - _executed: Whether the script has been executed - """ - - default_timeout = 1 - default_limit = 1024 - - def __init__( - self, - id: str, - timestamp: datetime, - script: str, - timeout: Optional[float], - limit: Optional[int], - ): - """ - Initialize a new single shot entry. - - Args: - id: Unique identifier for this entry - timestamp: Creation timestamp for this entry - script: The script/command to execute - timeout: Maximum time to wait for script execution - limit: Maximum number of characters to capture from stdout/stderr - """ - super().__init__(id, timestamp) - self._script = script - self._timeout = timeout - self._limit = limit - self._stdout = "" - self._stderr = "" - self._exit_code: Optional[int] = None - self._executed = False - self._timed_out = False - - @property - def script(self) -> str: - """Get the script/command being executed.""" - return self._script - - @property - def stdout(self) -> str: - """Get the captured standard output.""" - return self._stdout - - @property - def stderr(self) -> str: - """Get the captured standard error.""" - return self._stderr - - @property - def exit_code(self) -> Optional[int]: - """Get the exit code of the process (None if still running).""" - return self._exit_code - - def update(self) -> None: - """ - Execute the script if not already executed. - Captures stdout, stderr and exit code. - """ - if self._executed: - return - self._executed = True - try: - process = subprocess.run( - self._script, - timeout=(self._timeout or self.default_timeout), - shell=True, - capture_output=True, - text=True - ) - self._stdout = process.stdout - self._stderr = process.stderr - self._exit_code = process.returncode - except subprocess.TimeoutExpired as e: - self._timed_out = True - - def generate_context(self) -> ET.Element: - """ - Generate an XML Element representing this single shot entry. - - Returns: - ET.Element: XML element containing the entry's data - """ - element = ET.Element("single", { - "id": self.id, - }) - if self._timeout: - element.set("timeout", str(self._timeout)) - element.text = self.script - if self._timed_out: - element.set("timed_out", "true") - elif self._executed: - element.set("exit_code", str(self._exit_code)) - if self._limit: - element.set("limit", str(self._limit)) - if len(self._stdout) > (self._limit or self.default_limit): - element.set("stdout_truncated", "true") - element.set("stdout_length", str(len(self._stdout))) - stdout_elem = ET.SubElement(element, "stdout") - stdout_elem.text = self._stdout[:(self._limit or self.default_limit)] - if len(self._stderr) > (self._limit or self.default_limit): - element.set("stderr_truncated", "true") - element.set("stderr_length", str(len(self._stderr))) - stderr_elem = ET.SubElement(element, "stderr") - stderr_elem.text = self._stderr[:(self._limit or self.default_limit)] - return element \ No newline at end of file diff --git a/sia/system_metrics.py b/sia/system_metrics.py index e0e778e..e0e02e6 100644 --- a/sia/system_metrics.py +++ b/sia/system_metrics.py @@ -1,72 +1,13 @@ -import threading import time -from typing import Optional, Dict +from typing import Dict import psutil +from .util import format_timestamp class SystemMetrics: """ - Tracks system metrics including CPU, GPU, memory and disk usage. - Maintains rolling averages for CPU and GPU. + Tracks system metrics including memory and disk usage. - Attributes: - _stop_event: Threading event to signal monitor thread to stop - _monitor_thread: Background thread that collects usage data - _cpu_samples: List of CPU usage samples since last reading - _gpu_samples: List of GPU usage samples since last reading - _metrics_lock: Lock for thread-safe access to samples """ - - def __init__(self, sample_interval: float = 0.1): - """ - Initialize system metrics monitoring. - - Args: - sample_interval: Seconds between usage samples - """ - self._stop_event = threading.Event() - self._cpu_samples = [] - self._gpu_samples = [] - self._metrics_lock = threading.Lock() - self._sample_interval = sample_interval - - # Start monitoring thread - self._monitor_thread = threading.Thread( - target=self._monitor_loop, - daemon=True - ) - self._monitor_thread.start() - - def _get_gpu_usage(self) -> Optional[float]: - """ - Get current GPU usage percentage. - Returns None if GPU monitoring not available. - """ - try: - import pynvml - pynvml.nvmlInit() - handle = pynvml.nvmlDeviceGetHandleByIndex(0) - info = pynvml.nvmlDeviceGetUtilizationRates(handle) - pynvml.nvmlShutdown() - return info.gpu - except: - return None - - def _monitor_loop(self): - """ - Background loop that collects system usage samples. - Runs until stop event is set. - """ - while not self._stop_event.is_set(): - with self._metrics_lock: - # Get CPU usage across all cores - self._cpu_samples.append(psutil.cpu_percent()) - - # Get GPU usage if available - gpu_usage = self._get_gpu_usage() - if gpu_usage is not None: - self._gpu_samples.append(gpu_usage) - - time.sleep(self._sample_interval) def get_metrics(self) -> Dict: """ @@ -75,29 +16,17 @@ class SystemMetrics: Returns: Dict containing: - - cpu: Average CPU usage percentage - - gpu: Average GPU usage percentage + - timestamp: Current timestamp - memory_used: Used memory in bytes - memory_total: Total memory in bytes - disk_used: Used disk space in bytes - disk_total: Total disk space in bytes - - timestamp: Current UTC timestamp in ISO format """ metrics = {} # Add timestamp metrics["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) - # Add usage metrics - with self._metrics_lock: - # CPU average (or 0 if no samples) - metrics["cpu"] = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0 - self._cpu_samples.clear() - - # GPU average (or 0 if no samples) - metrics["gpu"] = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0 - self._gpu_samples.clear() - # Memory usage in bytes memory = psutil.virtual_memory() metrics["memory_used"] = memory.used @@ -108,9 +37,4 @@ class SystemMetrics: metrics["disk_used"] = disk.used metrics["disk_total"] = disk.total - return metrics - - def stop(self): - """Stop the monitoring thread and clean up.""" - self._stop_event.set() - self._monitor_thread.join() + return metrics \ No newline at end of file diff --git a/sia/util.py b/sia/util.py index 3dbaa08..4c640f1 100644 --- a/sia/util.py +++ b/sia/util.py @@ -1,3 +1,4 @@ +import datetime import xml.dom.minidom import xml.etree.ElementTree as ET from typing import Iterator, Optional @@ -47,7 +48,7 @@ def pretty_print_element(elem: ET.Element, level: int = 0, max_line: int = 80) - parts = [f'{indent}<{tag}{attr_str}>'] # Handle text content - if elem.text and elem.text.strip(): + if elem.text and isinstance(elem.text, str) and elem.text.strip(): text = elem.text if ']]>' in text: escaped = text.replace('&', '&') \ @@ -66,4 +67,7 @@ def pretty_print_element(elem: ET.Element, level: int = 0, max_line: int = 80) - parts.append(f'{indent} {child.tail}') parts.append(f'{indent}') - return '\n'.join(parts) \ No newline at end of file + return '\n'.join(parts) + +def format_timestamp(timestamp: datetime) -> str: + return timestamp.strftime("%Y%m%d_%H%M%S_%f")[:-3] \ No newline at end of file diff --git a/sia/web/api.py b/sia/web/api.py index 966d668..405d531 100644 --- a/sia/web/api.py +++ b/sia/web/api.py @@ -5,11 +5,21 @@ import asyncio from ..auto_approver import AutoApprover from ..web_agent import WebAgent from ..web_io_buffer import WebIOBuffer +from ..working_memory import WorkingMemory +from ..entry.entry_factory import EntryFactory class Api: - def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover): + def __init__( + self, + app: web.Application, + agent: WebAgent, + io_buffer: WebIOBuffer, + working_memory: WorkingMemory, + auto_approver: AutoApprover + ): self._app = app self._agent = agent + self._working_memory = working_memory self._io_buffer = io_buffer self._auto_approver = auto_approver @@ -31,6 +41,11 @@ class Api: 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_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._update_entry) + self._app.router.add_delete("/api/memory/entry/{id}", self._delete_entry) + self._app.router.add_post("/api/memory/entry/{id}/reset", self._reset_entry) async def _run_inference(self, request: web.Request) -> web.Response: """Start inference on specified LLM.""" @@ -174,4 +189,53 @@ class Api: self._auto_approver.llm_name = name return web.Response(status=200) except ValueError as e: - return web.Response(status=400, text=str(e)) \ No newline at end of file + 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() + return web.Response( + text=json.dumps([e.serialize() for e in entries]), + content_type="application/json" + ) + + async def _create_entry(self, request: web.Request) -> web.Response: + """Create a new entry in working memory.""" + data = await request.json() + try: + entry = EntryFactory.create_entry(data, self._io_buffer) + self._working_memory.add_entry(entry) + return web.Response( + text=json.dumps({"id": entry.id}), + content_type="application/json" + ) + except (ValueError, TypeError) as e: + return web.Response(status=400, text=str(e)) + + async def _update_entry(self, request: web.Request) -> web.Response: + """Update properties of an existing entry.""" + entry_id = request.match_info["id"] + data = await request.json() + entry = self._working_memory.get_entry(entry_id) + if not entry: + return web.Response(status=404, text="Entry not found") + try: + EntryFactory.update_entry(entry, data) + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + + async def _delete_entry(self, request: web.Request) -> web.Response: + """Delete an entry from working memory.""" + entry_id = request.match_info["id"] + self._working_memory.remove_entry(entry_id) + return web.Response(status=200) + + async def _reset_entry(self, request: web.Request) -> web.Response: + """Reset an entry's state.""" + entry_id = request.match_info["id"] + entry = self._working_memory.get_entry(entry_id) + if not entry: + return web.Response(status=404, text="Entry not found") + entry.reset() + return web.Response(status=200) \ No newline at end of file diff --git a/sia/web/memory_websocket.py b/sia/web/memory_websocket.py new file mode 100644 index 0000000..82ee943 --- /dev/null +++ b/sia/web/memory_websocket.py @@ -0,0 +1,59 @@ +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 + +class MemoryWebSocket: + """ + WebSocket handler for working memory changes. + Broadcasts memory updates to all connected clients. + """ + + def __init__(self, working_memory: WorkingMemory): + self._working_memory = working_memory + self._clients: Set[web.WebSocketResponse] = set() + self._working_memory.add_change_handler(wrap_async(self._handle_memory_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_memory_change(self): + """Handle working memory changes.""" + entries = self._working_memory.get_entries() + await self._broadcast_message({ + "type": "memory_state", + "entries": [e.serialize() for e in entries] + }) + + 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 memory state + entries = self._working_memory.get_entries() + await ws.send_json({ + "type": "memory_state", + "entries": [e.serialize() for e in entries] + }) + + 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 diff --git a/sia/web/websockts.py b/sia/web/websockts.py index d756721..e5204dd 100644 --- a/sia/web/websockts.py +++ b/sia/web/websockts.py @@ -1,24 +1,28 @@ from aiohttp import web +from ..auto_approver import AutoApprover from ..web_agent import WebAgent from ..web_io_buffer import WebIOBuffer -from ..auto_approver import AutoApprover -from .llm_websocket import LlmWebSocket -from .context_websocket import ContextWebSocket -from .token_websocket import TokenWebSocket -from .stdout_websocket import StdoutWebSocket +from ..working_memory import WorkingMemory from .auto_approver_websocket import AutoApproverWebSocket +from .context_websocket import ContextWebSocket +from .llm_websocket import LlmWebSocket +from .memory_websocket import MemoryWebSocket +from .stdout_websocket import StdoutWebSocket +from .token_websocket import TokenWebSocket class Websockets: - def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover): + def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory): self._llm_ws = LlmWebSocket(agent) self._context_ws = ContextWebSocket(agent) self._token_ws = TokenWebSocket(agent) self._stdout_ws = StdoutWebSocket(io_buffer) self._auto_approver_ws = AutoApproverWebSocket(auto_approver) + self._memory_ws = MemoryWebSocket(working_memory) app.router.add_get("/ws/llm", self._llm_ws.handle_connection) app.router.add_get("/ws/context", self._context_ws.handle_connection) app.router.add_get("/ws/token", self._token_ws.handle_connection) app.router.add_get("/ws/stdout", self._stdout_ws.handle_connection) app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection) + app.router.add_get("/ws/memory", self._memory_ws.handle_connection) diff --git a/sia/web_agent.py b/sia/web_agent.py index dfb80ff..2120bed 100644 --- a/sia/web_agent.py +++ b/sia/web_agent.py @@ -1,3 +1,4 @@ +from datetime import datetime from enum import Enum, auto from threading import Lock from typing import Callable, Dict, List, Optional @@ -55,6 +56,9 @@ class WebAgent(BaseAgent): self._token_handlers: List[Callable[[str, str], None]] = [] self._context_change_handlers: List[Callable[[str, bool], None]] = [] + # Working memory change handler + self._working_memory.add_change_handler(self._handle_memory_update) + @property def llms(self) -> Dict[str, LlmState]: """Get current state of all LLMs""" @@ -138,9 +142,9 @@ class WebAgent(BaseAgent): if llm_name not in self._llms: raise ValueError(f"Unknown LLM: {llm_name}") - self._iteration_logger.log_iteration(self._context, response) - - parse_result = self._parser.parse(response) + timestamp = datetime.now() + self._iteration_logger.log_iteration(timestamp, self._context, response) + parse_result = self._parser.parse(timestamp, response) if isinstance(parse_result, Command): result = parse_result.execute(self._working_memory) self._command_result = result @@ -150,11 +154,14 @@ class WebAgent(BaseAgent): parse_result.update() self._working_memory.update() self._working_memory.add_entry(parse_result) - - self.modify_context(self._compile_context(self._llms[llm_name]), True) def _set_llm_state(self, llm_name: str, state: LlmState) -> 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 diff --git a/sia/working_memory.py b/sia/working_memory.py index ef7ee29..c34392a 100644 --- a/sia/working_memory.py +++ b/sia/working_memory.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import List, Optional, Callable, Set import xml.etree.ElementTree as ET from .entry import Entry @@ -10,11 +10,34 @@ class WorkingMemory: The working memory stores different types of entries (scripts, reasoning, errors, etc.) and provides methods to add, remove, and update them. It also generates XML context representing the current state. + + Notifies observers of entry additions, deletions and changes. During update(), + changes are batched and a single notification is sent after completion. """ def __init__(self): """Initialize an empty working memory.""" self._entries: List[Entry] = [] + self._change_handlers: List[Callable[[], None]] = [] + self._updating = False + self._changed_during_update: Set[Entry] = set() + + def add_change_handler(self, handler: Callable[[], None]) -> None: + """Add a callback for working memory changes.""" + if handler not in self._change_handlers: + self._change_handlers.append(handler) + + def _notify_change(self) -> None: + """Notify all handlers of working memory changes.""" + for handler in self._change_handlers: + handler() + + def _handle_entry_change(self, entry: Entry) -> None: + """Handle changes from individual entries.""" + if self._updating: + self._changed_during_update.add(entry) + else: + self._notify_change() def add_entry(self, entry: Entry) -> None: """ @@ -25,7 +48,9 @@ class WorkingMemory: """ if not isinstance(entry, Entry): raise TypeError("Entry must be an instance of Entry class") + entry.add_change_handler(self._handle_entry_change) self._entries.append(entry) + self._notify_change() def remove_entry(self, id: str) -> None: """ @@ -39,6 +64,7 @@ class WorkingMemory: if entry is not None: entry.cleanup() self._entries = [e for e in self._entries if e.id != id] + self._notify_change() def clear(self) -> None: """ @@ -48,12 +74,11 @@ class WorkingMemory: for entry in self._entries: entry.cleanup() self._entries.clear() + self._notify_change() def __del__(self): - """ - Clean up all entries when memory is deleted. - """ - if hasattr(self, '_entries'): # Check in case of partial initialization + """Clean up all entries when memory is deleted.""" + if hasattr(self, '_entries'): self.clear() def get_entry(self, id: str) -> Optional[Entry]: @@ -78,7 +103,7 @@ class WorkingMemory: Returns: List[Entry]: List of all entries """ - return self._entries.copy() # Return a copy to prevent external modification + return self._entries.copy() def get_entries_count(self) -> int: """ @@ -88,24 +113,22 @@ class WorkingMemory: int: Number of entries in working memory """ return len(self._entries) - - def get_entries_by_type(self, entry_type: type) -> List[Entry]: - """ - Get all entries of a specific type. - - Args: - entry_type: Type of entries to retrieve - - Returns: - List[Entry]: List of entries matching the specified type - """ - return [e for e in self._entries if isinstance(e, entry_type)] def update(self) -> None: - """Update all entries in working memory.""" + """ + Update all entries in working memory. + Batches change notifications and sends single update after completion. + """ + self._updating = True + self._changed_during_update.clear() + for entry in self._entries: entry.update() + self._updating = False + if self._changed_during_update: + self._notify_change() + def generate_context(self) -> List[ET.Element]: """ Generate XML Elements representing all entries. @@ -113,4 +136,4 @@ class WorkingMemory: Returns: List[ET.Element]: List of XML elements for each entry """ - return [entry.generate_context() for entry in self._entries] \ No newline at end of file + return [entry.generate_context() for entry in self._entries] diff --git a/test/base_agent_test.py b/test/base_agent_test.py index 80c01c4..3ee81d5 100644 --- a/test/base_agent_test.py +++ b/test/base_agent_test.py @@ -14,7 +14,7 @@ from sia.base_agent import BaseAgent from sia.command import Command from sia.command_result import CommandResult from sia.reasoning_entry import ReasoningEntry -from sia.parse_error_entry import ParseErrorEntry +from sia.entry.parse_error_entry import ParseErrorEntry from sia.delete_command import DeleteCommand from sia.entry import Entry diff --git a/test/parse_error_entry_test.py b/test/parse_error_entry_test.py index 86cdf5e..920ba0d 100644 --- a/test/parse_error_entry_test.py +++ b/test/parse_error_entry_test.py @@ -1,7 +1,7 @@ import unittest from datetime import datetime -from sia.parse_error_entry import ParseErrorEntry +from sia.entry.parse_error_entry import ParseErrorEntry class ParseErrorEntryTest(unittest.TestCase): def setUp(self): diff --git a/test/response_parser_test.py b/test/response_parser_test.py index 75fe9a4..7434f2d 100644 --- a/test/response_parser_test.py +++ b/test/response_parser_test.py @@ -7,12 +7,12 @@ from sia.delete_command import DeleteCommand from sia.stop_command import StopCommand from sia.entry import Entry from sia.background_entry import BackgroundEntry -from sia.parse_error_entry import ParseErrorEntry +from sia.entry.parse_error_entry import ParseErrorEntry from sia.read_entry import ReadEntry from sia.reasoning_entry import ReasoningEntry from sia.repeat_entry import RepeatEntry from sia.single_entry import SingleEntry -from sia.write_entry import WriteEntry +from sia.entry.write_entry import WriteEntry from sia.web_io_buffer import WebIOBuffer from sia.response_parser import ResponseParser diff --git a/test/web_agent_test.py b/test/web_agent_test.py index 2795dba..e03adbd 100644 --- a/test/web_agent_test.py +++ b/test/web_agent_test.py @@ -16,7 +16,7 @@ from sia.web_agent import WebAgent, WebAgentState from sia.command import Command from sia.command_result import CommandResult from sia.reasoning_entry import ReasoningEntry -from sia.parse_error_entry import ParseErrorEntry +from sia.entry.parse_error_entry import ParseErrorEntry from sia.entry import Entry class MockCommand(Command): diff --git a/test/working_memory_test.py b/test/working_memory_test.py index 20c62b4..115e426 100644 --- a/test/working_memory_test.py +++ b/test/working_memory_test.py @@ -6,12 +6,12 @@ import xml.etree.ElementTree as ET from sia.background_entry import BackgroundEntry from sia.entry import Entry -from sia.parse_error_entry import ParseErrorEntry +from sia.entry.parse_error_entry import ParseErrorEntry from sia.read_entry import ReadEntry from sia.reasoning_entry import ReasoningEntry from sia.web_io_buffer import WebIOBuffer from sia.working_memory import WorkingMemory -from sia.write_entry import WriteEntry +from sia.entry.write_entry import WriteEntry class MockEntry(Entry): diff --git a/test/write_entry_test.py b/test/write_entry_test.py index e7bbdc9..4d24fbf 100644 --- a/test/write_entry_test.py +++ b/test/write_entry_test.py @@ -2,7 +2,7 @@ import unittest from datetime import datetime from sia.web_io_buffer import WebIOBuffer -from sia.write_entry import WriteEntry +from sia.entry.write_entry import WriteEntry class WriteEntryTest(unittest.TestCase): def setUp(self): @@ -19,7 +19,7 @@ class WriteEntryTest(unittest.TestCase): self.assertEqual(entry.content, content) self.assertEqual(entry.id, self.test_id) self.assertEqual(entry.timestamp, self.test_timestamp) - self.assertFalse(entry._written) + self.assertFalse(entry.written) self.assertEqual(self.io_buffer.get_stdout(), "") def test_single_write(self): @@ -30,7 +30,7 @@ class WriteEntryTest(unittest.TestCase): # Perform update and verify content was written entry.update() self.assertEqual(self.io_buffer.get_stdout(), content) - self.assertTrue(entry._written) + self.assertTrue(entry.written) def test_multiple_updates(self): """Test that content is only written once even with multiple updates""" @@ -47,7 +47,7 @@ class WriteEntryTest(unittest.TestCase): # Verify no additional content was written self.assertEqual(self.io_buffer.get_stdout(), "") - self.assertTrue(entry._written) + self.assertTrue(entry.written) self.assertEqual(initial_output, content) def test_empty_content(self): @@ -56,7 +56,7 @@ class WriteEntryTest(unittest.TestCase): entry.update() self.assertEqual(self.io_buffer.get_stdout(), "") - self.assertTrue(entry._written) + self.assertTrue(entry.written) def test_special_characters(self): """Test writing content with special characters"""