llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Shared response buffer backend

Browse Source
This commit is contained in:
Niels Geens
2025-04-17 18:32:23 +02:00
parent 9477575421
commit 34fb5d814f
12 changed files with 941 additions and 869 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sia/__main__.py
View File

@@ -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

366
sia/auto_approver.py
View File

@@ -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))

28
sia/llm_engine/__init__.py
View File

@@ -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

8
sia/llm_engine/local_llm_engine.py
View File

@@ -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,

154
sia/llm_engine/mistral_llm_engine.py
View File

@@ -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

7
sia/llm_engine/qwq_llm_engine.py
View File

@@ -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(

26
sia/response_buffer.py Normal file
View File

@@ -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)

183
sia/util.py
View File

@@ -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('&', '&amp;') \
.replace('<', '&lt;') \
.replace('>', '&gt;') \
.replace('"', '&quot;') \
.replace("'", '&apos;')
parts.append(f'{indent} {escaped}')
else:
parts.append(f'{indent} <![CDATA[{text}]]>')
# 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}</{tag}>')
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('&', '&amp;') \
.replace('<', '&lt;') \
.replace('>', '&gt;') \
.replace('"', '&quot;') \
.replace("'", '&apos;')
parts.append(f'{indent} {escaped}')
else:
parts.append(f'{indent} <![CDATA[{text}]]>')
# 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}</{tag}>')
return '\n'.join(parts)
def format_timestamp(timestamp: datetime) -> str:
return timestamp.strftime("%Y%m%d_%H%M%S_%f")[:-3]

594
sia/web/api.py
View File

@@ -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)

45
sia/web/token_websocket.py → sia/web/response_websocket.py
View File

@@ -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

56
sia/web/websockts.py → sia/web/websockets.py
View File

@@ -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)

341
sia/web_agent.py
View File

@@ -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)