diff --git a/sia/__main__.py b/sia/__main__.py index 7ff3842..7ab0562 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -13,7 +13,7 @@ from .response_parser import ResponseParser from .system_metrics import SystemMetrics from .web.api import Api from .web.static import Static -from .web.websockts import Websockets +from .web.websockets import Websockets from .web_agent import WebAgent from .web_io_buffer import WebIOBuffer from .working_memory import WorkingMemory diff --git a/sia/auto_approver.py b/sia/auto_approver.py index 78e44fc..0844257 100644 --- a/sia/auto_approver.py +++ b/sia/auto_approver.py @@ -1,183 +1,183 @@ -from threading import Thread, Event -from typing import Callable, TypeAlias, TypedDict -from .web_agent import WebAgent, LlmState - -class AutoApproverConfig(TypedDict): - context_enabled: bool - response_enabled: bool - context_timeout: float - response_timeout: float - llm_name: str - -ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None] - -class AutoApprover: - """ - Handles automatic approvals for WebAgent states after configurable timeouts. - """ - - def __init__(self, agent: WebAgent): - """ - 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._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) - - @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 - ) - - 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 - self._notify_config_change() - - def add_config_change_handler(self, handler: ConfigChangeHandler) -> None: - self._config_change_handlers.append(handler) - - def _notify_config_change(self) -> None: - current_config = self.config - 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.NO_OUTPUT: - 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.OUTPUT: - 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.OUTPUT 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 _stop_context_thread(self) -> None: - if self._context_thread: - self._stop_event.set() - self._context_thread = None - self._stop_event.clear() - - def _stop_response_thread(self) -> None: - if self._response_thread: - self._stop_event.set() - self._response_thread = None - self._stop_event.clear() - - def _start_context_thread(self) -> None: - self._context_thread = Thread(target=self._context_approval_thread) - self._context_thread.start() - - def _start_response_thread(self) -> None: - self._response_thread = Thread(target=self._response_approval_thread) - self._response_thread.start() - - def _context_approval_thread(self) -> None: - if self._stop_event.wait(self._context_timeout): - return - if self._context_enabled: - self.agent.run_inference(self._llm_name) - - def _response_approval_thread(self) -> None: - if self._stop_event.wait(self._response_timeout): - return - if (self._response_enabled and - self.agent.llms[self._llm_name] == LlmState.OUTPUT): - self.agent.approve_response(self._llm_name, self.agent.get_output(self._llm_name)) +from threading import Thread, Event +from typing import Callable, TypeAlias, TypedDict +from .web_agent import WebAgent, LlmState + +class AutoApproverConfig(TypedDict): + context_enabled: bool + response_enabled: bool + context_timeout: float + response_timeout: float + llm_name: str + +ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None] + +class AutoApprover: + """ + Handles automatic approvals for WebAgent states after configurable timeouts. + """ + + def __init__(self, agent: WebAgent): + """ + 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._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) + + @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 + ) + + 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 + self._notify_config_change() + + def add_config_change_handler(self, handler: ConfigChangeHandler) -> None: + self._config_change_handlers.append(handler) + + def _notify_config_change(self) -> None: + current_config = self.config + 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 _stop_context_thread(self) -> None: + if self._context_thread: + self._stop_event.set() + self._context_thread = None + self._stop_event.clear() + + def _stop_response_thread(self) -> None: + if self._response_thread: + self._stop_event.set() + self._response_thread = None + self._stop_event.clear() + + def _start_context_thread(self) -> None: + self._context_thread = Thread(target=self._context_approval_thread) + self._context_thread.start() + + def _start_response_thread(self) -> None: + self._response_thread = Thread(target=self._response_approval_thread) + self._response_thread.start() + + def _context_approval_thread(self) -> None: + if self._stop_event.wait(self._context_timeout): + return + if self._context_enabled: + self.agent.run_inference(self._llm_name) + + def _response_approval_thread(self) -> None: + if self._stop_event.wait(self._response_timeout): + return + if (self._response_enabled and + self.agent.llms[self._llm_name] == LlmState.IDLE): + self.agent.approve_response(self._llm_name, self.agent.get_output(self._llm_name)) diff --git a/sia/llm_engine/__init__.py b/sia/llm_engine/__init__.py index 7b9bff2..3ece6e2 100644 --- a/sia/llm_engine/__init__.py +++ b/sia/llm_engine/__init__.py @@ -1,15 +1,15 @@ -from typing import Callable, Iterator -from abc import ABC, abstractmethod - -class LlmEngine(ABC): - @abstractmethod - def infer(self, system_prompt: str, main_context: 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: +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/llm_engine/local_llm_engine.py b/sia/llm_engine/local_llm_engine.py index 439f747..e29a3d6 100644 --- a/sia/llm_engine/local_llm_engine.py +++ b/sia/llm_engine/local_llm_engine.py @@ -56,7 +56,7 @@ class LocalLlmEngine(LlmEngine): else: self._logits_processor = None - def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: + def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: """ Run inference using the system prompt and main context. @@ -70,10 +70,12 @@ class LocalLlmEngine(LlmEngine): """ messages = [ {"role": "system", "content": system_prompt}, - {"role": "user", "content": main_context} + {"role": "user", "content": main_context}, + {"role": "assistant", "content": continuation_text}, ] prompt = self._tokenizer.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True + messages, tokenize=False, + add_generation_prompt=False, ) streamer = TextIteratorStreamer( self._tokenizer, diff --git a/sia/llm_engine/mistral_llm_engine.py b/sia/llm_engine/mistral_llm_engine.py index a0d0f80..984b778 100644 --- a/sia/llm_engine/mistral_llm_engine.py +++ b/sia/llm_engine/mistral_llm_engine.py @@ -1,71 +1,83 @@ -from typing import Iterator, Optional, Callable -from mistralai import Mistral -from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage -from mistral_common.protocol.instruct.request import ChatCompletionRequest -from mistral_common.tokens.tokenizers.mistral import MistralTokenizer - -from . import LlmEngine - -class MistralLlmEngine(LlmEngine): - def __init__( - self, - model: str, - temperature: float, - token_limit: int, - api_key: str, - ): - self._model = model - self._temperature = temperature - self._token_limit = token_limit - self._api_key = api_key - self._client = Mistral(api_key=api_key) - self._tokenizer = MistralTokenizer.v3() - - def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: - messages = [ - { - "role": "system", - "content": system_prompt, - }, - { - "role": "user", - "content": main_context, - }, - { - "role": "assistant", - "content": "<", - "prefix": True, - }, - ] - stream_response = self._client.chat.stream( - model=self._model, - messages=messages, - temperature=self._temperature, - ) - - try: - for chunk in stream_response: - if should_stop(): - stream_response.response.close() - break - if content := chunk.data.choices[0].delta.content: - yield content - finally: - stream_response.response.close() - - def token_count(self, system_prompt: str, main_context: str) -> int: - messages = [ - SystemMessage(content=system_prompt), - UserMessage(content=main_context), - ] - - tokenized = self._tokenizer.encode_chat_completion( - ChatCompletionRequest( - messages=messages, - model=self._model - ) - ) - return len(tokenized.tokens) - - def token_limit(self) -> int: - return self._token_limit +from typing import Iterator, Optional, Callable +from mistralai import Mistral +from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage +from mistral_common.protocol.instruct.request import ChatCompletionRequest +from mistral_common.tokens.tokenizers.mistral import MistralTokenizer + +from . import LlmEngine +from ..util import skip_prefix + +class MistralLlmEngine(LlmEngine): + def __init__( + self, + model: str, + temperature: float, + token_limit: int, + api_key: str, + ): + self._model = model + self._temperature = temperature + self._token_limit = token_limit + self._api_key = api_key + 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]: + messages = [ + { + "role": "system", + "content": system_prompt, + }, + { + "role": "user", + "content": main_context, + }, + { + "role": "assistant", + "content": continuation_text, + "prefix": True, + }, + ] if continuation_text else [ + { + "role": "system", + "content": system_prompt, + }, + { + "role": "user", + "content": main_context, + }, + ] + stream_response = self._client.chat.stream( + model=self._model, + messages=messages, + 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) + finally: + stream_response.response.close() + + def token_count(self, system_prompt: str, main_context: str) -> int: + messages = [ + SystemMessage(content=system_prompt), + UserMessage(content=main_context), + ] + + tokenized = self._tokenizer.encode_chat_completion( + ChatCompletionRequest( + messages=messages, + model=self._model + ) + ) + return len(tokenized.tokens) + + def token_limit(self) -> int: + return self._token_limit diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py index d8770ad..afba73f 100644 --- a/sia/llm_engine/qwq_llm_engine.py +++ b/sia/llm_engine/qwq_llm_engine.py @@ -63,7 +63,7 @@ class QwQLlmEngine(LlmEngine): self._logits_processor = None - def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: + def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: """ Run inference using the system prompt and main context. @@ -77,13 +77,14 @@ class QwQLlmEngine(LlmEngine): """ messages = [ {"role": "system", "content": system_prompt}, - {"role": "user", "content": main_context} + {"role": "user", "content": main_context}, + {"role": "assistant", "content": continuation_text}, ] text = self._tokenizer.apply_chat_template( messages, tokenize=False, - add_generation_prompt=True, + add_generation_prompt=False, ) streamer = TextIteratorStreamer( diff --git a/sia/response_buffer.py b/sia/response_buffer.py new file mode 100644 index 0000000..2345a0b --- /dev/null +++ b/sia/response_buffer.py @@ -0,0 +1,26 @@ +class ResponseBuffer: + def __init__(self): + self._text = "" + self._observers = [] + + def add_observer(self, callback): + self._observers.append(callback) + + def get_text(self): + return self._text + + def append_text(self, text): + self._text += text + self._notify_observers() + + def set_text(self, text): + self._text = text + self._notify_observers() + + def clear(self): + self._text = "" + self._notify_observers() + + def _notify_observers(self): + for observer in self._observers: + observer(self._text) \ No newline at end of file diff --git a/sia/util.py b/sia/util.py index 4c640f1..7cfc496 100644 --- a/sia/util.py +++ b/sia/util.py @@ -1,73 +1,112 @@ -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 pretty_print_element(elem: ET.Element, level: int = 0, max_line: int = 80) -> str: - """Convert ElementTree element to pretty-printed string with custom formatting.""" - indent = ' ' * level - - # Handle empty elements - if not elem.text and not elem.tail and not elem.attrib and len(elem) == 0: - return f'{indent}<{elem.tag}/>' - - # Build opening tag with attributes - tag = elem.tag - attrs = elem.attrib - if attrs: - attr_strings = [f'{k}="{v}"' for k, v in sorted(attrs.items())] - attr_str = ' ' + ' '.join(attr_strings) - - if len(indent) + len(tag) + len(attr_str) + 2 > max_line: - attr_indent = indent + ' ' - attr_str = '\n' + '\n'.join(f'{attr_indent}{a}' for a in attr_strings) - else: - attr_str = '' - - parts = [f'{indent}<{tag}{attr_str}>'] - - # Handle text content - if elem.text and isinstance(elem.text, str) and elem.text.strip(): - text = elem.text - if ']]>' in text: - escaped = text.replace('&', '&') \ - .replace('<', '<') \ - .replace('>', '>') \ - .replace('"', '"') \ - .replace("'", ''') - parts.append(f'{indent} {escaped}') - else: - parts.append(f'{indent} ') - - # Handle children - for child in elem: - parts.append(pretty_print_element(child, level + 1, max_line)) - if child.tail and child.tail.strip(): - parts.append(f'{indent} {child.tail}') - - parts.append(f'{indent}') - return '\n'.join(parts) - -def format_timestamp(timestamp: datetime) -> str: +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.""" + indent = ' ' * level + + # Handle empty elements + if not elem.text and not elem.tail and not elem.attrib and len(elem) == 0: + return f'{indent}<{elem.tag}/>' + + # Build opening tag with attributes + tag = elem.tag + attrs = elem.attrib + if attrs: + attr_strings = [f'{k}="{v}"' for k, v in sorted(attrs.items())] + attr_str = ' ' + ' '.join(attr_strings) + + if len(indent) + len(tag) + len(attr_str) + 2 > max_line: + attr_indent = indent + ' ' + attr_str = '\n' + '\n'.join(f'{attr_indent}{a}' for a in attr_strings) + else: + attr_str = '' + + parts = [f'{indent}<{tag}{attr_str}>'] + + # Handle text content + if elem.text and isinstance(elem.text, str) and elem.text.strip(): + text = elem.text + if ']]>' in text: + escaped = text.replace('&', '&') \ + .replace('<', '<') \ + .replace('>', '>') \ + .replace('"', '"') \ + .replace("'", ''') + parts.append(f'{indent} {escaped}') + else: + parts.append(f'{indent} ') + + # Handle children + for child in elem: + parts.append(pretty_print_element(child, level + 1, max_line)) + if child.tail and child.tail.strip(): + parts.append(f'{indent} {child.tail}') + + parts.append(f'{indent}') + 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 82d4fd9..9e1d212 100644 --- a/sia/web/api.py +++ b/sia/web/api.py @@ -1,286 +1,308 @@ -from pathlib import Path -from aiohttp import web -import json -import asyncio - - -from ..auto_approver import AutoApprover -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: - def __init__( - self, - work_dir: Path, - app: web.Application, - agent: WebAgent, - io_buffer: WebIOBuffer, - working_memory: WorkingMemory, - auto_approver: AutoApprover - ): - self._work_dir = work_dir - self._app = app - self._agent = agent - self._working_memory = working_memory - self._io_buffer = io_buffer - self._auto_approver = auto_approver - - self._init_routes() - - def _init_routes(self): - """Initialize REST API and WebSocket routes.""" - self._app.router.add_post("/api/inference/{llm}", self._run_inference) - self._app.router.add_post("/api/inference/{llm}/stop", self._stop_inference) - self._app.router.add_post("/api/approve/{llm}", self._approve_response) - self._app.router.add_post("/api/context", self._modify_context) - self._app.router.add_post("/api/input", self._send_input) - 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_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) - 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) - 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.""" - llm_name = request.match_info["llm"] - try: - await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name) - return web.Response(status=200) - except (ValueError, RuntimeError) as e: - return web.Response(status=400, text=str(e)) - - async def _stop_inference(self, request: web.Request) -> web.Response: - """Stop inference on specified LLM.""" - llm_name = request.match_info["llm"] - try: - self._agent.stop_inference(llm_name) - return web.Response(status=200) - except ValueError as e: - return web.Response(status=400, text=str(e)) - - async def _approve_response(self, request: web.Request) -> web.Response: - """Approve response from specified LLM.""" - llm_name = request.match_info["llm"] - data = await request.json() - response = data.get("response") - if not response: - return web.Response(status=400, text="Missing response in request body") - try: - self._agent.approve_response(llm_name, response) - return web.Response(status=200) - except ValueError as e: - return web.Response(status=400, text=str(e)) - - async def _modify_context(self, request: web.Request) -> web.Response: - """Modify the current context.""" - data = await request.json() - context = data.get("context") - if not context: - return web.Response(status=400, text="Missing context in request body") - self._agent.modify_context(context) - return web.Response(status=200) - - async def _send_input(self, request: web.Request) -> web.Response: - """Send input to the IO buffer.""" - data = await request.json() - input_text = data.get("input") - if not input_text: - return web.Response(status=400, text="Missing input in request body") - self._io_buffer.append_stdin(input_text) - return web.Response(status=200) - - 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.""" - llm_name = request.match_info["llm"] - try: - output = self._agent.get_output(llm_name) - return web.Response( - text=json.dumps({"output": output}), - content_type="application/json" - ) - except ValueError 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: 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.""" - return web.Response( - text=json.dumps(self._auto_approver.config), - content_type="application/json" - ) - - async def _set_auto_approver_config(self, request: web.Request) -> web.Response: - """Update auto approver configuration.""" - data = await request.json() - try: - self._auto_approver.set_config(data) - return web.Response(status=200) - except (ValueError, KeyError) as e: - return web.Response(status=400, text=str(e)) - - async def _set_context_enabled(self, request: web.Request) -> web.Response: - """Set context auto-approval enabled state.""" - data = await request.json() - enabled = data.get("enabled") - if enabled is None: - return web.Response(status=400, text="Missing enabled parameter") - try: - self._auto_approver.context_enabled = enabled - return web.Response(status=200) - except ValueError as e: - return web.Response(status=400, text=str(e)) - - async def _set_response_enabled(self, request: web.Request) -> web.Response: - """Set response auto-approval enabled state.""" - data = await request.json() - enabled = data.get("enabled") - if enabled is None: - return web.Response(status=400, text="Missing enabled parameter") - try: - self._auto_approver.response_enabled = enabled - return web.Response(status=200) - except ValueError as e: - return web.Response(status=400, text=str(e)) - - async def _set_context_timeout(self, request: web.Request) -> web.Response: - """Set context auto-approval timeout.""" - data = await request.json() - timeout = data.get("timeout") - if timeout is None: - return web.Response(status=400, text="Missing timeout parameter") - try: - self._auto_approver.context_timeout = float(timeout) - return web.Response(status=200) - except ValueError as e: - return web.Response(status=400, text=str(e)) - - async def _set_response_timeout(self, request: web.Request) -> web.Response: - """Set response auto-approval timeout.""" - data = await request.json() - timeout = data.get("timeout") - if timeout is None: - return web.Response(status=400, text="Missing timeout parameter") - try: - self._auto_approver.response_timeout = float(timeout) - return web.Response(status=200) - except ValueError as e: - return web.Response(status=400, text=str(e)) - - async def _set_llm_name(self, request: web.Request) -> web.Response: - """Set LLM name for auto-approval.""" - data = await request.json() - name = data.get("name") - if name is None: - return web.Response(status=400, text="Missing name parameter") - try: - self._auto_approver.llm_name = name - return web.Response(status=200) - except ValueError 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() - 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._work_dir, self._io_buffer) - self._working_memory.add_entry(entry) - return web.Response( - text=json.dumps({"id": entry.id}), - content_type="application/json" - ) - except (ValueError, TypeError) as e: - return web.Response(status=400, text=str(e)) - - async def _save_entry(self, request: web.Request) -> web.Response: - """Update properties of an existing entry.""" - entry_id = request.match_info["id"] - data = await request.json() - entry = self._working_memory.get_entry(entry_id) - if not entry: - return web.Response(status=404, text="Entry not found") - try: - EntryFactory.update_entry(entry, data) - entry.notify_change() - 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) - - async def _update_entry(self, request: web.Request) -> web.Response: - """Update an entry's state.""" - entry_id = request.match_info["id"] - entry = self._working_memory.get_entry(entry_id) - if not entry: - return web.Response(status=404, text="Entry not found") - try: - entry.update() - entry.notify_change() - return web.Response(status=200) - except ValueError as e: - return web.Response(status=400, text=str(e)) - - async def _load_iteration(self, request: web.Request) -> web.Response: - """Load entries from iteration XML content into working memory""" - data = await request.json() - content = data.get("content") - if not content: - return web.Response(status=400, text="Missing content in request body") - - entries = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer) - - for entry in entries: - self._working_memory.add_entry(entry) - - return web.Response(status=200) +from pathlib import Path +from aiohttp import web +import json +import asyncio + + +from ..auto_approver import AutoApprover +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: + def __init__( + self, + work_dir: Path, + app: web.Application, + agent: WebAgent, + io_buffer: WebIOBuffer, + working_memory: WorkingMemory, + auto_approver: AutoApprover + ): + self._work_dir = work_dir + self._app = app + self._agent = agent + self._working_memory = working_memory + self._io_buffer = io_buffer + self._auto_approver = auto_approver + + self._init_routes() + + def _init_routes(self): + """Initialize REST API and WebSocket routes.""" + self._app.router.add_post("/api/inference/{llm}", self._run_inference) + self._app.router.add_post("/api/inference/{llm}/stop", self._stop_inference) + self._app.router.add_post("/api/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_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) + 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) + 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.""" + llm_name = request.match_info["llm"] + try: + await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name) + return web.Response(status=200) + except (ValueError, RuntimeError) as e: + return web.Response(status=400, text=str(e)) + + async def _stop_inference(self, request: web.Request) -> web.Response: + """Stop inference on specified LLM.""" + llm_name = request.match_info["llm"] + try: + self._agent.stop_inference(llm_name) + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + + async def _set_response(self, request: web.Request) -> web.Response: + """Edit the shared response buffer""" + try: + data = await request.json() + text = data.get("response") + if text is None: + return web.Response(status=400, text="Missing response text in request body") + try: + self._agent._response_buffer.set_text(text) + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + except json.JSONDecodeError: + return web.Response(status=400, text="Invalid JSON in request body") + + async def _approve_response(self, request: web.Request) -> web.Response: + """Approve current buffer content""" + data = await request.json() + try: + self._agent.approve_response() + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + + async def _get_response(self, request: web.Request) -> web.Response: + """Get current buffer content""" + return web.Response( + text=json.dumps({ + "response": self._agent.output_buffer.get_text(), + }), + content_type="application/json" + ) + + async def _modify_context(self, request: web.Request) -> web.Response: + """Modify the current context.""" + data = await request.json() + context = data.get("context") + if not context: + return web.Response(status=400, text="Missing context in request body") + self._agent.modify_context(context) + return web.Response(status=200) + + async def _send_input(self, request: web.Request) -> web.Response: + """Send input to the IO buffer.""" + data = await request.json() + input_text = data.get("input") + if not input_text: + return web.Response(status=400, text="Missing input in request body") + self._io_buffer.append_stdin(input_text) + return web.Response(status=200) + + 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.""" + llm_name = request.match_info["llm"] + try: + output = self._agent.get_output(llm_name) + return web.Response( + text=json.dumps({"output": output}), + content_type="application/json" + ) + except ValueError 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: 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.""" + return web.Response( + text=json.dumps(self._auto_approver.config), + content_type="application/json" + ) + + async def _set_auto_approver_config(self, request: web.Request) -> web.Response: + """Update auto approver configuration.""" + data = await request.json() + try: + self._auto_approver.set_config(data) + return web.Response(status=200) + except (ValueError, KeyError) as e: + return web.Response(status=400, text=str(e)) + + async def _set_context_enabled(self, request: web.Request) -> web.Response: + """Set context auto-approval enabled state.""" + data = await request.json() + enabled = data.get("enabled") + if enabled is None: + return web.Response(status=400, text="Missing enabled parameter") + try: + self._auto_approver.context_enabled = enabled + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + + async def _set_response_enabled(self, request: web.Request) -> web.Response: + """Set response auto-approval enabled state.""" + data = await request.json() + enabled = data.get("enabled") + if enabled is None: + return web.Response(status=400, text="Missing enabled parameter") + try: + self._auto_approver.response_enabled = enabled + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + + async def _set_context_timeout(self, request: web.Request) -> web.Response: + """Set context auto-approval timeout.""" + data = await request.json() + timeout = data.get("timeout") + if timeout is None: + return web.Response(status=400, text="Missing timeout parameter") + try: + self._auto_approver.context_timeout = float(timeout) + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + + async def _set_response_timeout(self, request: web.Request) -> web.Response: + """Set response auto-approval timeout.""" + data = await request.json() + timeout = data.get("timeout") + if timeout is None: + return web.Response(status=400, text="Missing timeout parameter") + try: + self._auto_approver.response_timeout = float(timeout) + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + + async def _set_llm_name(self, request: web.Request) -> web.Response: + """Set LLM name for auto-approval.""" + data = await request.json() + name = data.get("name") + if name is None: + return web.Response(status=400, text="Missing name parameter") + try: + self._auto_approver.llm_name = name + return web.Response(status=200) + except ValueError 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() + 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._work_dir, self._io_buffer) + self._working_memory.add_entry(entry) + return web.Response( + text=json.dumps({"id": entry.id}), + content_type="application/json" + ) + except (ValueError, TypeError) as e: + return web.Response(status=400, text=str(e)) + + async def _save_entry(self, request: web.Request) -> web.Response: + """Update properties of an existing entry.""" + entry_id = request.match_info["id"] + data = await request.json() + entry = self._working_memory.get_entry(entry_id) + if not entry: + return web.Response(status=404, text="Entry not found") + try: + EntryFactory.update_entry(entry, data) + entry.notify_change() + 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) + + async def _update_entry(self, request: web.Request) -> web.Response: + """Update an entry's state.""" + entry_id = request.match_info["id"] + entry = self._working_memory.get_entry(entry_id) + if not entry: + return web.Response(status=404, text="Entry not found") + try: + entry.update() + entry.notify_change() + return web.Response(status=200) + except ValueError as e: + return web.Response(status=400, text=str(e)) + + async def _load_iteration(self, request: web.Request) -> web.Response: + """Load entries from iteration XML content into working memory""" + data = await request.json() + content = data.get("content") + if not content: + return web.Response(status=400, text="Missing content in request body") + + entries = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer) + + for entry in entries: + self._working_memory.add_entry(entry) + + return web.Response(status=200) diff --git a/sia/web/token_websocket.py b/sia/web/response_websocket.py similarity index 55% rename from sia/web/token_websocket.py rename to sia/web/response_websocket.py index 6dccd14..e77934f 100644 --- a/sia/web/token_websocket.py +++ b/sia/web/response_websocket.py @@ -1,24 +1,21 @@ -from aiohttp import web, WSMsgType -from typing import Dict, Set -import asyncio -import json +from aiohttp import web -from .util import wrap_async from ..web_agent import WebAgent +from .util import wrap_async -class TokenWebSocket: - """ - WebSocket handler for LLM token streaming. - Broadcasts new tokens to all connected clients. - """ - + +class ResponseWebSocket: def __init__(self, agent: WebAgent): self._agent = agent - self._clients: Set[web.WebSocketResponse] = set() - self._agent.add_token_handler(wrap_async(self._handle_new_token)) - - async def _broadcast_message(self, message: Dict): - """Broadcast message to all connected clients.""" + self._clients = set() + self._agent._response_buffer.add_observer(wrap_async(self._handle_buffer_change)) + + async def _handle_buffer_change(self, buffer): + await self._broadcast_message({ + "buffer": buffer, + }) + + async def _broadcast_message(self, message): disconnected = set() for ws in self._clients: try: @@ -26,26 +23,22 @@ class TokenWebSocket: except ConnectionResetError: disconnected.add(ws) self._clients -= disconnected - - async def _handle_new_token(self, llm_name: str, token: str): - """Handle new tokens from the WebAgent.""" - await self._broadcast_message({ - "llm": llm_name, - "token": token - }) - + 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: + await ws.send_json({ + "buffer": self._agent.response_buffer.get_text(), + }) + 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 + return ws \ No newline at end of file diff --git a/sia/web/websockts.py b/sia/web/websockets.py similarity index 85% rename from sia/web/websockts.py rename to sia/web/websockets.py index e5204dd..fc623d4 100644 --- a/sia/web/websockts.py +++ b/sia/web/websockets.py @@ -1,28 +1,28 @@ -from aiohttp import web - -from ..auto_approver import AutoApprover -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 .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, 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) +from aiohttp import web + +from ..auto_approver import AutoApprover +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 .memory_websocket import MemoryWebSocket +from .response_websocket import ResponseWebSocket +from .stdout_websocket import StdoutWebSocket + +class Websockets: + def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory): + self._auto_approver_ws = AutoApproverWebSocket(auto_approver) + self._context_ws = ContextWebSocket(agent) + self._llm_ws = LlmWebSocket(agent) + self._memory_ws = MemoryWebSocket(working_memory) + self._response_ws = ResponseWebSocket(agent) + self._stdout_ws = StdoutWebSocket(io_buffer) + + 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/llm", self._llm_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) diff --git a/sia/web_agent.py b/sia/web_agent.py index 8e713b0..37208fd 100644 --- a/sia/web_agent.py +++ b/sia/web_agent.py @@ -1,182 +1,159 @@ -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 .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_parser import ResponseParser -from .system_metrics import SystemMetrics -from .working_memory import WorkingMemory -from .xml_validator import XMLValidator - -class LlmState(Enum): - NO_OUTPUT = auto() - INFERENCE = auto() - OUTPUT = auto() - -class WebAgent(BaseAgent): - def __init__( - self, - system_prompt: str, - action_schema: str, - working_memory: WorkingMemory, - metrics: SystemMetrics, - llms: Dict[str, LlmEngine], - validator: XMLValidator, - parser: ResponseParser, - iteration_logger: IterationLogger, - ): - super().__init__( - system_prompt, - action_schema, - working_memory, - metrics, - validator, - parser - ) - self._llms = llms - self._iteration_logger = iteration_logger - self._llm_states: Dict[str, LlmState] = {name: LlmState.NO_OUTPUT for name in llms} - self._llm_outputs: Dict[str, str] = defaultdict(str) - 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]] = [] - - # 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""" - with self._llm_lock: - return self._llm_states.copy() - - @property - def context(self) -> str: - return self._context - - @property - def command_result(self) -> Optional[CommandResult]: - return self._command_result - - @property - def validation_error(self) -> Optional[str]: - return self._validation_error - - 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_token_handler(self, handler: Callable[[str, str], None]) -> None: - """Add handler for new tokens""" - if handler not in self._token_handlers: - self._token_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 - self._llm_outputs.clear() - for llm_name in self._llms: - self._set_llm_state(llm_name, LlmState.NO_OUTPUT) - - for handler in self._context_change_handlers: - handler(context, generated) - - def run_inference(self, llm_name: str) -> 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.NO_OUTPUT: - raise RuntimeError(f"LLM {llm_name} is not ready for inference") - self._set_llm_state(llm_name, LlmState.INFERENCE) - self._stop_flags[llm_name] = False - - llm = self._llms[llm_name] - - def should_stop() -> bool: - return self._stop_flags[llm_name] - - response_token_iter = llm.infer(self.system_prompt, self.context, should_stop) - - with self._output_lock: - self._llm_outputs[llm_name] = "" - - for token in response_token_iter: - with self._output_lock: - self._llm_outputs[llm_name] += token - - for handler in self._token_handlers: - handler(llm_name, token) - - with self._llm_lock: - self._set_llm_state(llm_name, LlmState.OUTPUT) - - def stop_inference(self, llm_name: str) -> None: - """Stop ongoing inference for specified LLM""" - if llm_name not in self._llms: - raise ValueError(f"Unknown LLM: {llm_name}") - self._stop_flags[llm_name] = True - - def get_output(self, llm_name: str) -> str: - """Get complete output for specified LLM""" - if llm_name not in self._llms: - raise ValueError(f"Unknown LLM: {llm_name}") - - with self._output_lock: - return self._llm_outputs[llm_name] - - def approve_response(self, llm_name: str, response: str) -> None: - """Process approved response from specified LLM""" - if llm_name not in self._llms: - raise ValueError(f"Unknown LLM: {llm_name}") - - timestamp = datetime.now(timezone.utc) - 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 - if result.should_stop: - exit(42) - else: - self._working_memory.update() - else: - parse_result.update() - self._working_memory.update() - self._working_memory.add_entry(parse_result) - - 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) +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 .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 +from .response_parser import ResponseParser +from .system_metrics import SystemMetrics +from .working_memory import WorkingMemory +from .xml_validator import XMLValidator + +class LlmState(Enum): + IDLE = auto() + INFERENCE = auto() + +class WebAgent(BaseAgent): + def __init__( + self, + system_prompt: str, + action_schema: str, + working_memory: WorkingMemory, + metrics: SystemMetrics, + llms: Dict[str, LlmEngine], + validator: XMLValidator, + parser: ResponseParser, + iteration_logger: IterationLogger, + ): + super().__init__( + system_prompt, + action_schema, + working_memory, + metrics, + validator, + parser + ) + self._llms = llms + 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]] = [] + + # 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() + + @property + def response_buffer(self) -> ResponseBuffer: + return self._response_buffer + + @property + def context(self) -> str: + return self._context + + @property + def command_result(self) -> Optional[CommandResult]: + return self._command_result + + @property + def validation_error(self) -> Optional[str]: + return self._validation_error + + 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_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 + self._response_buffer.clear() + for llm_name in self._llms: + self._set_llm_state(llm_name, LlmState.IDLE) + + for handler in self._context_change_handlers: + handler(context, generated) + + def run_inference(self, llm_name: str) -> 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) + for token in response_token_iter: + print(token, end='', flush=True) + with self._output_lock: + self._response_buffer.append_text(token) + with self._llm_lock: + self._set_llm_state(llm_name, LlmState.IDLE) + + def stop_inference(self, llm_name: str) -> None: + """Stop ongoing inference for specified LLM""" + if llm_name not in self._llms: + raise ValueError(f"Unknown LLM: {llm_name}") + self._stop_flags[llm_name] = True + + def approve_response(self) -> None: + """Process approved response from specified LLM""" + timestamp = datetime.now(timezone.utc) + self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text()) + parse_result = self._parser.parse(timestamp, self._response_buffer.get_text()) + if isinstance(parse_result, Command): + result = parse_result.execute(self._working_memory) + self._command_result = result + if result.should_stop: + exit(42) + else: + self._working_memory.update() + else: + parse_result.update() + self._working_memory.update() + self._working_memory.add_entry(parse_result) + + 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