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 .system_metrics import SystemMetrics
from .web.api import Api from .web.api import Api
from .web.static import Static from .web.static import Static
from .web.websockts import Websockets from .web.websockets import Websockets
from .web_agent import WebAgent from .web_agent import WebAgent
from .web_io_buffer import WebIOBuffer from .web_io_buffer import WebIOBuffer
from .working_memory import WorkingMemory from .working_memory import WorkingMemory

366
sia/auto_approver.py
View File

@@ -1,183 +1,183 @@
from threading import Thread, Event from threading import Thread, Event
from typing import Callable, TypeAlias, TypedDict from typing import Callable, TypeAlias, TypedDict
from .web_agent import WebAgent, LlmState from .web_agent import WebAgent, LlmState
class AutoApproverConfig(TypedDict): class AutoApproverConfig(TypedDict):
context_enabled: bool context_enabled: bool
response_enabled: bool response_enabled: bool
context_timeout: float context_timeout: float
response_timeout: float response_timeout: float
llm_name: str llm_name: str
ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None] ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]
class AutoApprover: class AutoApprover:
""" """
Handles automatic approvals for WebAgent states after configurable timeouts. Handles automatic approvals for WebAgent states after configurable timeouts.
""" """
def __init__(self, agent: WebAgent): def __init__(self, agent: WebAgent):
""" """
Initialize auto approver with a WebAgent instance. Initialize auto approver with a WebAgent instance.
""" """
self.agent = agent self.agent = agent
self._llm_name = next(iter(agent.llms.keys())) self._llm_name = next(iter(agent.llms.keys()))
self._context_timeout = 5.0 self._context_timeout = 5.0
self._response_timeout = 10.0 self._response_timeout = 10.0
self._context_enabled = False self._context_enabled = False
self._response_enabled = False self._response_enabled = False
self._stop_event = Event() self._stop_event = Event()
self._context_thread: Thread | None = None self._context_thread: Thread | None = None
self._response_thread: Thread | None = None self._response_thread: Thread | None = None
self._config_change_handlers: list[ConfigChangeHandler] = [] self._config_change_handlers: list[ConfigChangeHandler] = []
self.agent.add_llm_change_handler(self._handle_llm_state_change) self.agent.add_llm_change_handler(self._handle_llm_state_change)
self.agent.add_context_change_handler(self._handle_context_change) self.agent.add_context_change_handler(self._handle_context_change)
@property @property
def config(self) -> AutoApproverConfig: def config(self) -> AutoApproverConfig:
return AutoApproverConfig( return AutoApproverConfig(
context_enabled=self._context_enabled, context_enabled=self._context_enabled,
response_enabled=self._response_enabled, response_enabled=self._response_enabled,
context_timeout=self._context_timeout, context_timeout=self._context_timeout,
response_timeout=self._response_timeout, response_timeout=self._response_timeout,
llm_name=self._llm_name llm_name=self._llm_name
) )
def set_config(self, config: AutoApproverConfig) -> None: def set_config(self, config: AutoApproverConfig) -> None:
if config['llm_name'] not in self.agent.llms: if config['llm_name'] not in self.agent.llms:
raise ValueError(f"Unknown LLM: {config['llm_name']}") raise ValueError(f"Unknown LLM: {config['llm_name']}")
notify_config_change = self._notify_config_change notify_config_change = self._notify_config_change
self._notify_config_change = lambda: None self._notify_config_change = lambda: None
try: try:
self.context_enabled = False self.context_enabled = False
self.response_enabled = False self.response_enabled = False
self.context_timeout = config['context_timeout'] self.context_timeout = config['context_timeout']
self.response_timeout = config['response_timeout'] self.response_timeout = config['response_timeout']
self.llm_name = config['llm_name'] self.llm_name = config['llm_name']
self.context_enabled = config['context_enabled'] self.context_enabled = config['context_enabled']
self.response_enabled = config['response_enabled'] self.response_enabled = config['response_enabled']
finally: finally:
self._notify_config_change = notify_config_change self._notify_config_change = notify_config_change
self._notify_config_change() self._notify_config_change()
def add_config_change_handler(self, handler: ConfigChangeHandler) -> None: def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
self._config_change_handlers.append(handler) self._config_change_handlers.append(handler)
def _notify_config_change(self) -> None: def _notify_config_change(self) -> None:
current_config = self.config current_config = self.config
for handler in self._config_change_handlers: for handler in self._config_change_handlers:
handler(current_config) handler(current_config)
@property @property
def context_timeout(self) -> float: def context_timeout(self) -> float:
return self._context_timeout return self._context_timeout
@context_timeout.setter @context_timeout.setter
def context_timeout(self, timeout: float) -> None: def context_timeout(self, timeout: float) -> None:
if self._context_enabled: if self._context_enabled:
raise ValueError("Cannot change timeout while auto-approval is enabled") raise ValueError("Cannot change timeout while auto-approval is enabled")
self._context_timeout = timeout self._context_timeout = timeout
self._notify_config_change() self._notify_config_change()
@property @property
def response_timeout(self) -> float: def response_timeout(self) -> float:
return self._response_timeout return self._response_timeout
@response_timeout.setter @response_timeout.setter
def response_timeout(self, timeout: float) -> None: def response_timeout(self, timeout: float) -> None:
if self._response_enabled: if self._response_enabled:
raise ValueError("Cannot change timeout while auto-approval is enabled") raise ValueError("Cannot change timeout while auto-approval is enabled")
self._response_timeout = timeout self._response_timeout = timeout
self._notify_config_change() self._notify_config_change()
@property @property
def context_enabled(self) -> bool: def context_enabled(self) -> bool:
return self._context_enabled return self._context_enabled
@context_enabled.setter @context_enabled.setter
def context_enabled(self, enabled: bool) -> None: def context_enabled(self, enabled: bool) -> None:
if enabled == self._context_enabled: if enabled == self._context_enabled:
return return
self._context_enabled = enabled self._context_enabled = enabled
self._stop_context_thread() self._stop_context_thread()
if enabled and self.agent.llms[self._llm_name] == LlmState.NO_OUTPUT: if enabled and self.agent.llms[self._llm_name] == LlmState.IDLE:
self._start_context_thread() self._start_context_thread()
self._notify_config_change() self._notify_config_change()
@property @property
def response_enabled(self) -> bool: def response_enabled(self) -> bool:
return self._response_enabled return self._response_enabled
@response_enabled.setter @response_enabled.setter
def response_enabled(self, enabled: bool) -> None: def response_enabled(self, enabled: bool) -> None:
if enabled == self._response_enabled: if enabled == self._response_enabled:
return return
self._response_enabled = enabled self._response_enabled = enabled
self._stop_response_thread() self._stop_response_thread()
if enabled and self.agent.llms[self._llm_name] == LlmState.OUTPUT: if enabled and self.agent.llms[self._llm_name] == LlmState.IDLE:
self._start_response_thread() self._start_response_thread()
self._notify_config_change() self._notify_config_change()
@property @property
def llm_name(self) -> str: def llm_name(self) -> str:
return self._llm_name return self._llm_name
@llm_name.setter @llm_name.setter
def llm_name(self, name: str) -> None: def llm_name(self, name: str) -> None:
if name not in self.agent.llms: if name not in self.agent.llms:
raise ValueError(f"Unknown LLM: {name}") raise ValueError(f"Unknown LLM: {name}")
self._llm_name = name self._llm_name = name
self._notify_config_change() self._notify_config_change()
def _handle_llm_state_change(self, llm_name: str, state: LlmState) -> None: def _handle_llm_state_change(self, llm_name: str, state: LlmState) -> None:
if llm_name != self._llm_name: if llm_name != self._llm_name:
return return
if state == LlmState.OUTPUT and self._response_enabled: if state == LlmState.IDLE and self._response_enabled:
self._start_response_thread() self._start_response_thread()
else: else:
self._stop_response_thread() self._stop_response_thread()
def _handle_context_change(self, context: str, generated: bool) -> None: def _handle_context_change(self, context: str, generated: bool) -> None:
if generated and self._context_enabled: if generated and self._context_enabled:
self._start_context_thread() self._start_context_thread()
else: else:
self._stop_context_thread() self._stop_context_thread()
def _stop_context_thread(self) -> None: def _stop_context_thread(self) -> None:
if self._context_thread: if self._context_thread:
self._stop_event.set() self._stop_event.set()
self._context_thread = None self._context_thread = None
self._stop_event.clear() self._stop_event.clear()
def _stop_response_thread(self) -> None: def _stop_response_thread(self) -> None:
if self._response_thread: if self._response_thread:
self._stop_event.set() self._stop_event.set()
self._response_thread = None self._response_thread = None
self._stop_event.clear() self._stop_event.clear()
def _start_context_thread(self) -> None: def _start_context_thread(self) -> None:
self._context_thread = Thread(target=self._context_approval_thread) self._context_thread = Thread(target=self._context_approval_thread)
self._context_thread.start() self._context_thread.start()
def _start_response_thread(self) -> None: def _start_response_thread(self) -> None:
self._response_thread = Thread(target=self._response_approval_thread) self._response_thread = Thread(target=self._response_approval_thread)
self._response_thread.start() self._response_thread.start()
def _context_approval_thread(self) -> None: def _context_approval_thread(self) -> None:
if self._stop_event.wait(self._context_timeout): if self._stop_event.wait(self._context_timeout):
return return
if self._context_enabled: if self._context_enabled:
self.agent.run_inference(self._llm_name) self.agent.run_inference(self._llm_name)
def _response_approval_thread(self) -> None: def _response_approval_thread(self) -> None:
if self._stop_event.wait(self._response_timeout): if self._stop_event.wait(self._response_timeout):
return return
if (self._response_enabled and if (self._response_enabled and
self.agent.llms[self._llm_name] == LlmState.OUTPUT): self.agent.llms[self._llm_name] == LlmState.IDLE):
self.agent.approve_response(self._llm_name, self.agent.get_output(self._llm_name)) 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 typing import Callable, Iterator
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
class LlmEngine(ABC): class LlmEngine(ABC):
@abstractmethod @abstractmethod
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]:
pass pass
@abstractmethod @abstractmethod
def token_count(self, system_prompt: str, main_context: str) -> int: def token_count(self, system_prompt: str, main_context: str) -> int:
pass pass
@abstractmethod @abstractmethod
def token_limit(self) -> int: def token_limit(self) -> int:
pass pass

8
sia/llm_engine/local_llm_engine.py
View File

@@ -56,7 +56,7 @@ class LocalLlmEngine(LlmEngine):
else: else:
self._logits_processor = None 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. Run inference using the system prompt and main context.
@@ -70,10 +70,12 @@ class LocalLlmEngine(LlmEngine):
""" """
messages = [ messages = [
{"role": "system", "content": system_prompt}, {"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( prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True messages, tokenize=False,
add_generation_prompt=False,
) )
streamer = TextIteratorStreamer( streamer = TextIteratorStreamer(
self._tokenizer, self._tokenizer,

154
sia/llm_engine/mistral_llm_engine.py
View File

@@ -1,71 +1,83 @@
from typing import Iterator, Optional, Callable from typing import Iterator, Optional, Callable
from mistralai import Mistral from mistralai import Mistral
from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from . import LlmEngine from . import LlmEngine
from ..util import skip_prefix
class MistralLlmEngine(LlmEngine):
def __init__( class MistralLlmEngine(LlmEngine):
self, def __init__(
model: str, self,
temperature: float, model: str,
token_limit: int, temperature: float,
api_key: str, token_limit: int,
): api_key: str,
self._model = model ):
self._temperature = temperature self._model = model
self._token_limit = token_limit self._temperature = temperature
self._api_key = api_key self._token_limit = token_limit
self._client = Mistral(api_key=api_key) self._api_key = api_key
self._tokenizer = MistralTokenizer.v3() 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 = [ 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": "system",
}, "content": system_prompt,
{ },
"role": "user", {
"content": main_context, "role": "user",
}, "content": main_context,
{ },
"role": "assistant", {
"content": "<", "role": "assistant",
"prefix": True, "content": continuation_text,
}, "prefix": True,
] },
stream_response = self._client.chat.stream( ] if continuation_text else [
model=self._model, {
messages=messages, "role": "system",
temperature=self._temperature, "content": system_prompt,
) },
{
try: "role": "user",
for chunk in stream_response: "content": main_context,
if should_stop(): },
stream_response.response.close() ]
break stream_response = self._client.chat.stream(
if content := chunk.data.choices[0].delta.content: model=self._model,
yield content messages=messages,
finally: temperature=self._temperature,
stream_response.response.close() )
def token_count(self, system_prompt: str, main_context: str) -> int: try:
messages = [ def content_generator():
SystemMessage(content=system_prompt), for chunk in stream_response:
UserMessage(content=main_context), if should_stop():
] stream_response.response.close()
break
tokenized = self._tokenizer.encode_chat_completion( if content := chunk.data.choices[0].delta.content:
ChatCompletionRequest( yield content
messages=messages, yield from skip_prefix(content_generator(), continuation_text)
model=self._model finally:
) stream_response.response.close()
)
return len(tokenized.tokens) def token_count(self, system_prompt: str, main_context: str) -> int:
messages = [
def token_limit(self) -> int: SystemMessage(content=system_prompt),
return self._token_limit 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 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. Run inference using the system prompt and main context.
@@ -77,13 +77,14 @@ class QwQLlmEngine(LlmEngine):
""" """
messages = [ messages = [
{"role": "system", "content": system_prompt}, {"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( text = self._tokenizer.apply_chat_template(
messages, messages,
tokenize=False, tokenize=False,
add_generation_prompt=True, add_generation_prompt=False,
) )
streamer = TextIteratorStreamer( 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 datetime
import xml.dom.minidom import xml.dom.minidom
import xml.etree.ElementTree as ET import xml.etree.ElementTree as ET
from typing import Iterator, Optional from typing import Iterator, Optional
def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]: def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]:
""" """
Creates an iterator that yields values from the input iterator Creates an iterator that yields values from the input iterator
until it encounters the stop_value (exclusive). until it encounters the stop_value (exclusive).
Args: Args:
iterator: The source iterator iterator: The source iterator
stop_value: The value to stop before stop_value: The value to stop before
Yields: Yields:
Values from the iterator until stop_value is encountered Values from the iterator until stop_value is encountered
If stop_value is part of an item, yields the part before stop_value If stop_value is part of an item, yields the part before stop_value
""" """
for item in iterator: for item in iterator:
if stop_value in item: if stop_value in item:
split_point = item.index(stop_value) split_point = item.index(stop_value)
if split_point > 0: if split_point > 0:
yield item[:split_point] yield item[:split_point]
break break
yield item yield item
def pretty_print_element(elem: ET.Element, level: int = 0, max_line: int = 80) -> str: def skip_prefix(iterator: Iterator[str], prefix: str) -> Iterator[str]:
"""Convert ElementTree element to pretty-printed string with custom formatting.""" """
indent = ' ' * level Creates an iterator that skips a prefix from the input iterator
and yields only the content after the prefix.
# Handle empty elements
if not elem.text and not elem.tail and not elem.attrib and len(elem) == 0: Args:
return f'{indent}<{elem.tag}/>' iterator: The source iterator
prefix: The prefix to skip
# Build opening tag with attributes
tag = elem.tag Yields:
attrs = elem.attrib Values from the iterator after the prefix has been fully skipped
if attrs: """
attr_strings = [f'{k}="{v}"' for k, v in sorted(attrs.items())] if not prefix:
attr_str = ' ' + ' '.join(attr_strings) # If no prefix to skip, yield everything
yield from iterator
if len(indent) + len(tag) + len(attr_str) + 2 > max_line: return
attr_indent = indent + ' '
attr_str = '\n' + '\n'.join(f'{attr_indent}{a}' for a in attr_strings) prefix_remaining = prefix
else:
attr_str = '' for item in iterator:
if prefix_remaining:
parts = [f'{indent}<{tag}{attr_str}>'] # If the item starts with the remaining prefix
if prefix_remaining.startswith(item):
# Handle text content # Skip this item entirely
if elem.text and isinstance(elem.text, str) and elem.text.strip(): prefix_remaining = prefix_remaining[len(item):]
text = elem.text continue
if ']]>' in text: elif item.startswith(prefix_remaining):
escaped = text.replace('&', '&amp;') \ # Yield only the part after the prefix
.replace('<', '&lt;') \ yield item[len(prefix_remaining):]
.replace('>', '&gt;') \ prefix_remaining = ""
.replace('"', '&quot;') \ else:
.replace("'", '&apos;') # Item doesn't match prefix pattern, yield everything
parts.append(f'{indent} {escaped}') # This is unexpected but we handle it gracefully
else: yield item
parts.append(f'{indent} <![CDATA[{text}]]>') prefix_remaining = ""
else:
# Handle children # No prefix remaining, yield all content
for child in elem: yield item
parts.append(pretty_print_element(child, level + 1, max_line))
if child.tail and child.tail.strip(): def pretty_print_element(elem: ET.Element, level: int = 0, max_line: int = 80) -> str:
parts.append(f'{indent} {child.tail}') """Convert ElementTree element to pretty-printed string with custom formatting."""
indent = ' ' * level
parts.append(f'{indent}</{tag}>')
return '\n'.join(parts) # Handle empty elements
if not elem.text and not elem.tail and not elem.attrib and len(elem) == 0:
def format_timestamp(timestamp: datetime) -> str: 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] 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 pathlib import Path
from aiohttp import web from aiohttp import web
import json import json
import asyncio import asyncio
from ..auto_approver import AutoApprover from ..auto_approver import AutoApprover
from ..entry.entry_factory import EntryFactory from ..entry.entry_factory import EntryFactory
from ..iteration_parser import IterationParser from ..iteration_parser import IterationParser
from ..web_agent import WebAgent from ..web_agent import WebAgent
from ..web_io_buffer import WebIOBuffer from ..web_io_buffer import WebIOBuffer
from ..working_memory import WorkingMemory from ..working_memory import WorkingMemory
class Api: class Api:
def __init__( def __init__(
self, self,
work_dir: Path, work_dir: Path,
app: web.Application, app: web.Application,
agent: WebAgent, agent: WebAgent,
io_buffer: WebIOBuffer, io_buffer: WebIOBuffer,
working_memory: WorkingMemory, working_memory: WorkingMemory,
auto_approver: AutoApprover auto_approver: AutoApprover
): ):
self._work_dir = work_dir self._work_dir = work_dir
self._app = app self._app = app
self._agent = agent self._agent = agent
self._working_memory = working_memory self._working_memory = working_memory
self._io_buffer = io_buffer self._io_buffer = io_buffer
self._auto_approver = auto_approver self._auto_approver = auto_approver
self._init_routes() self._init_routes()
def _init_routes(self): def _init_routes(self):
"""Initialize REST API and WebSocket routes.""" """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}", self._run_inference)
self._app.router.add_post("/api/inference/{llm}/stop", self._stop_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/response", self._set_response)
self._app.router.add_post("/api/context", self._modify_context) self._app.router.add_post("/api/response/approve", self._approve_response)
self._app.router.add_post("/api/input", self._send_input) self._app.router.add_get("/api/response", self._get_response)
self._app.router.add_post("/api/clear", self._clear_output) self._app.router.add_post("/api/context", self._modify_context)
self._app.router.add_get("/api/output/{llm}", self._get_output) self._app.router.add_post("/api/input", self._send_input)
self._app.router.add_get("/api/llms", self._get_llms) self._app.router.add_post("/api/clear", self._clear_output)
self._app.router.add_get("/api/auto_approver/config", self._get_auto_approver_config) self._app.router.add_get("/api/output/{llm}", self._get_output)
self._app.router.add_post("/api/auto_approver/config", self._set_auto_approver_config) self._app.router.add_get("/api/llms", self._get_llms)
self._app.router.add_post("/api/auto_approver/context_enabled", self._set_context_enabled) self._app.router.add_get("/api/auto_approver/config", self._get_auto_approver_config)
self._app.router.add_post("/api/auto_approver/response_enabled", self._set_response_enabled) self._app.router.add_post("/api/auto_approver/config", self._set_auto_approver_config)
self._app.router.add_post("/api/auto_approver/context_timeout", self._set_context_timeout) self._app.router.add_post("/api/auto_approver/context_enabled", self._set_context_enabled)
self._app.router.add_post("/api/auto_approver/response_timeout", self._set_response_timeout) self._app.router.add_post("/api/auto_approver/response_enabled", self._set_response_enabled)
self._app.router.add_post("/api/auto_approver/llm", self._set_llm_name) self._app.router.add_post("/api/auto_approver/context_timeout", self._set_context_timeout)
self._app.router.add_get("/api/memory", self._get_memory) self._app.router.add_post("/api/auto_approver/response_timeout", self._set_response_timeout)
self._app.router.add_post("/api/memory/entry", self._create_entry) self._app.router.add_post("/api/auto_approver/llm", self._set_llm_name)
self._app.router.add_put("/api/memory/entry/{id}", self._save_entry) self._app.router.add_get("/api/memory", self._get_memory)
self._app.router.add_delete("/api/memory/entry/{id}", self._delete_entry) self._app.router.add_post("/api/memory/entry", self._create_entry)
self._app.router.add_post("/api/memory/entry/{id}/reset", self._reset_entry) self._app.router.add_put("/api/memory/entry/{id}", self._save_entry)
self._app.router.add_post("/api/memory/entry/{id}/update", self._update_entry) self._app.router.add_delete("/api/memory/entry/{id}", self._delete_entry)
self._app.router.add_post("/api/memory/load_iteration", self._load_iteration) 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)
async def _run_inference(self, request: web.Request) -> web.Response: self._app.router.add_post("/api/memory/load_iteration", self._load_iteration)
"""Start inference on specified LLM."""
llm_name = request.match_info["llm"] async def _run_inference(self, request: web.Request) -> web.Response:
try: """Start inference on specified LLM."""
await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name) llm_name = request.match_info["llm"]
return web.Response(status=200) try:
except (ValueError, RuntimeError) as e: await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name)
return web.Response(status=400, text=str(e)) return web.Response(status=200)
except (ValueError, RuntimeError) as e:
async def _stop_inference(self, request: web.Request) -> web.Response: return web.Response(status=400, text=str(e))
"""Stop inference on specified LLM."""
llm_name = request.match_info["llm"] async def _stop_inference(self, request: web.Request) -> web.Response:
try: """Stop inference on specified LLM."""
self._agent.stop_inference(llm_name) llm_name = request.match_info["llm"]
return web.Response(status=200) try:
except ValueError as e: self._agent.stop_inference(llm_name)
return web.Response(status=400, text=str(e)) return web.Response(status=200)
except ValueError as e:
async def _approve_response(self, request: web.Request) -> web.Response: return web.Response(status=400, text=str(e))
"""Approve response from specified LLM."""
llm_name = request.match_info["llm"] async def _set_response(self, request: web.Request) -> web.Response:
data = await request.json() """Edit the shared response buffer"""
response = data.get("response") try:
if not response: data = await request.json()
return web.Response(status=400, text="Missing response in request body") text = data.get("response")
try: if text is None:
self._agent.approve_response(llm_name, response) return web.Response(status=400, text="Missing response text in request body")
return web.Response(status=200) try:
except ValueError as e: self._agent._response_buffer.set_text(text)
return web.Response(status=400, text=str(e)) return web.Response(status=200)
except ValueError as e:
async def _modify_context(self, request: web.Request) -> web.Response: return web.Response(status=400, text=str(e))
"""Modify the current context.""" except json.JSONDecodeError:
data = await request.json() return web.Response(status=400, text="Invalid JSON in request body")
context = data.get("context")
if not context: async def _approve_response(self, request: web.Request) -> web.Response:
return web.Response(status=400, text="Missing context in request body") """Approve current buffer content"""
self._agent.modify_context(context) data = await request.json()
return web.Response(status=200) try:
self._agent.approve_response()
async def _send_input(self, request: web.Request) -> web.Response: return web.Response(status=200)
"""Send input to the IO buffer.""" except ValueError as e:
data = await request.json() return web.Response(status=400, text=str(e))
input_text = data.get("input")
if not input_text: async def _get_response(self, request: web.Request) -> web.Response:
return web.Response(status=400, text="Missing input in request body") """Get current buffer content"""
self._io_buffer.append_stdin(input_text) return web.Response(
return web.Response(status=200) text=json.dumps({
"response": self._agent.output_buffer.get_text(),
async def _clear_output(self, request: web.Request) -> web.Response: }),
"""Clear the stdout buffer.""" content_type="application/json"
self._io_buffer.clear_stdout() )
return web.Response(status=200)
async def _modify_context(self, request: web.Request) -> web.Response:
async def _get_output(self, request: web.Request) -> web.Response: """Modify the current context."""
"""Get complete output for specified LLM.""" data = await request.json()
llm_name = request.match_info["llm"] context = data.get("context")
try: if not context:
output = self._agent.get_output(llm_name) return web.Response(status=400, text="Missing context in request body")
return web.Response( self._agent.modify_context(context)
text=json.dumps({"output": output}), return web.Response(status=200)
content_type="application/json"
) async def _send_input(self, request: web.Request) -> web.Response:
except ValueError as e: """Send input to the IO buffer."""
return web.Response(status=400, text=str(e)) data = await request.json()
input_text = data.get("input")
async def _get_llms(self, request: web.Request) -> web.Response: if not input_text:
"""Get all LLMs and their current states.""" return web.Response(status=400, text="Missing input in request body")
states = self._agent.llms self._io_buffer.append_stdin(input_text)
return web.Response( return web.Response(status=200)
text=json.dumps({
name: state.name async def _clear_output(self, request: web.Request) -> web.Response:
for name, state in states.items() """Clear the stdout buffer."""
}), self._io_buffer.clear_stdout()
content_type="application/json" return web.Response(status=200)
)
async def _get_output(self, request: web.Request) -> web.Response:
async def _get_auto_approver_config(self, request: web.Request) -> web.Response: """Get complete output for specified LLM."""
"""Get current auto approver configuration.""" llm_name = request.match_info["llm"]
return web.Response( try:
text=json.dumps(self._auto_approver.config), output = self._agent.get_output(llm_name)
content_type="application/json" return web.Response(
) text=json.dumps({"output": output}),
content_type="application/json"
async def _set_auto_approver_config(self, request: web.Request) -> web.Response: )
"""Update auto approver configuration.""" except ValueError as e:
data = await request.json() return web.Response(status=400, text=str(e))
try:
self._auto_approver.set_config(data) async def _get_llms(self, request: web.Request) -> web.Response:
return web.Response(status=200) """Get all LLMs and their current states."""
except (ValueError, KeyError) as e: states = self._agent.llms
return web.Response(status=400, text=str(e)) return web.Response(
text=json.dumps({
async def _set_context_enabled(self, request: web.Request) -> web.Response: name: state.name
"""Set context auto-approval enabled state.""" for name, state in states.items()
data = await request.json() }),
enabled = data.get("enabled") content_type="application/json"
if enabled is None: )
return web.Response(status=400, text="Missing enabled parameter")
try: async def _get_auto_approver_config(self, request: web.Request) -> web.Response:
self._auto_approver.context_enabled = enabled """Get current auto approver configuration."""
return web.Response(status=200) return web.Response(
except ValueError as e: text=json.dumps(self._auto_approver.config),
return web.Response(status=400, text=str(e)) content_type="application/json"
)
async def _set_response_enabled(self, request: web.Request) -> web.Response:
"""Set response auto-approval enabled state.""" async def _set_auto_approver_config(self, request: web.Request) -> web.Response:
data = await request.json() """Update auto approver configuration."""
enabled = data.get("enabled") data = await request.json()
if enabled is None: try:
return web.Response(status=400, text="Missing enabled parameter") self._auto_approver.set_config(data)
try: return web.Response(status=200)
self._auto_approver.response_enabled = enabled except (ValueError, KeyError) as e:
return web.Response(status=200) return web.Response(status=400, text=str(e))
except ValueError 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."""
async def _set_context_timeout(self, request: web.Request) -> web.Response: data = await request.json()
"""Set context auto-approval timeout.""" enabled = data.get("enabled")
data = await request.json() if enabled is None:
timeout = data.get("timeout") return web.Response(status=400, text="Missing enabled parameter")
if timeout is None: try:
return web.Response(status=400, text="Missing timeout parameter") self._auto_approver.context_enabled = enabled
try: return web.Response(status=200)
self._auto_approver.context_timeout = float(timeout) except ValueError as e:
return web.Response(status=200) return web.Response(status=400, text=str(e))
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."""
async def _set_response_timeout(self, request: web.Request) -> web.Response: data = await request.json()
"""Set response auto-approval timeout.""" enabled = data.get("enabled")
data = await request.json() if enabled is None:
timeout = data.get("timeout") return web.Response(status=400, text="Missing enabled parameter")
if timeout is None: try:
return web.Response(status=400, text="Missing timeout parameter") self._auto_approver.response_enabled = enabled
try: return web.Response(status=200)
self._auto_approver.response_timeout = float(timeout) except ValueError as e:
return web.Response(status=200) return web.Response(status=400, text=str(e))
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."""
async def _set_llm_name(self, request: web.Request) -> web.Response: data = await request.json()
"""Set LLM name for auto-approval.""" timeout = data.get("timeout")
data = await request.json() if timeout is None:
name = data.get("name") return web.Response(status=400, text="Missing timeout parameter")
if name is None: try:
return web.Response(status=400, text="Missing name parameter") self._auto_approver.context_timeout = float(timeout)
try: return web.Response(status=200)
self._auto_approver.llm_name = name except ValueError as e:
return web.Response(status=200) return web.Response(status=400, text=str(e))
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."""
async def _get_memory(self, request: web.Request) -> web.Response: data = await request.json()
"""Get complete working memory state.""" timeout = data.get("timeout")
entries = self._working_memory.get_entries() if timeout is None:
return web.Response( return web.Response(status=400, text="Missing timeout parameter")
text=json.dumps([e.serialize() for e in entries]), try:
content_type="application/json" self._auto_approver.response_timeout = float(timeout)
) return web.Response(status=200)
except ValueError as e:
async def _create_entry(self, request: web.Request) -> web.Response: return web.Response(status=400, text=str(e))
"""Create a new entry in working memory."""
data = await request.json() async def _set_llm_name(self, request: web.Request) -> web.Response:
try: """Set LLM name for auto-approval."""
entry = EntryFactory.create_entry(data, self._work_dir, self._io_buffer) data = await request.json()
self._working_memory.add_entry(entry) name = data.get("name")
return web.Response( if name is None:
text=json.dumps({"id": entry.id}), return web.Response(status=400, text="Missing name parameter")
content_type="application/json" try:
) self._auto_approver.llm_name = name
except (ValueError, TypeError) as e: return web.Response(status=200)
return web.Response(status=400, text=str(e)) except ValueError 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.""" async def _get_memory(self, request: web.Request) -> web.Response:
entry_id = request.match_info["id"] """Get complete working memory state."""
data = await request.json() entries = self._working_memory.get_entries()
entry = self._working_memory.get_entry(entry_id) return web.Response(
if not entry: text=json.dumps([e.serialize() for e in entries]),
return web.Response(status=404, text="Entry not found") content_type="application/json"
try: )
EntryFactory.update_entry(entry, data)
entry.notify_change() async def _create_entry(self, request: web.Request) -> web.Response:
return web.Response(status=200) """Create a new entry in working memory."""
except ValueError as e: data = await request.json()
return web.Response(status=400, text=str(e)) try:
entry = EntryFactory.create_entry(data, self._work_dir, self._io_buffer)
async def _delete_entry(self, request: web.Request) -> web.Response: self._working_memory.add_entry(entry)
"""Delete an entry from working memory.""" return web.Response(
entry_id = request.match_info["id"] text=json.dumps({"id": entry.id}),
self._working_memory.remove_entry(entry_id) content_type="application/json"
return web.Response(status=200) )
except (ValueError, TypeError) as e:
async def _reset_entry(self, request: web.Request) -> web.Response: return web.Response(status=400, text=str(e))
"""Reset an entry's state."""
entry_id = request.match_info["id"] async def _save_entry(self, request: web.Request) -> web.Response:
entry = self._working_memory.get_entry(entry_id) """Update properties of an existing entry."""
if not entry: entry_id = request.match_info["id"]
return web.Response(status=404, text="Entry not found") data = await request.json()
entry.reset() entry = self._working_memory.get_entry(entry_id)
return web.Response(status=200) if not entry:
return web.Response(status=404, text="Entry not found")
async def _update_entry(self, request: web.Request) -> web.Response: try:
"""Update an entry's state.""" EntryFactory.update_entry(entry, data)
entry_id = request.match_info["id"] entry.notify_change()
entry = self._working_memory.get_entry(entry_id) return web.Response(status=200)
if not entry: except ValueError as e:
return web.Response(status=404, text="Entry not found") return web.Response(status=400, text=str(e))
try:
entry.update() async def _delete_entry(self, request: web.Request) -> web.Response:
entry.notify_change() """Delete an entry from working memory."""
return web.Response(status=200) entry_id = request.match_info["id"]
except ValueError as e: self._working_memory.remove_entry(entry_id)
return web.Response(status=400, text=str(e)) return web.Response(status=200)
async def _load_iteration(self, request: web.Request) -> web.Response: async def _reset_entry(self, request: web.Request) -> web.Response:
"""Load entries from iteration XML content into working memory""" """Reset an entry's state."""
data = await request.json() entry_id = request.match_info["id"]
content = data.get("content") entry = self._working_memory.get_entry(entry_id)
if not content: if not entry:
return web.Response(status=400, text="Missing content in request body") return web.Response(status=404, text="Entry not found")
entry.reset()
entries = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer) return web.Response(status=200)
for entry in entries: async def _update_entry(self, request: web.Request) -> web.Response:
self._working_memory.add_entry(entry) """Update an entry's state."""
entry_id = request.match_info["id"]
return web.Response(status=200) 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 aiohttp import web
from typing import Dict, Set
import asyncio
import json
from .util import wrap_async
from ..web_agent import WebAgent from ..web_agent import WebAgent
from .util import wrap_async
class TokenWebSocket:
""" class ResponseWebSocket:
WebSocket handler for LLM token streaming.
Broadcasts new tokens to all connected clients.
"""
def __init__(self, agent: WebAgent): def __init__(self, agent: WebAgent):
self._agent = agent self._agent = agent
self._clients: Set[web.WebSocketResponse] = set() self._clients = set()
self._agent.add_token_handler(wrap_async(self._handle_new_token)) self._agent._response_buffer.add_observer(wrap_async(self._handle_buffer_change))
async def _broadcast_message(self, message: Dict): async def _handle_buffer_change(self, buffer):
"""Broadcast message to all connected clients.""" await self._broadcast_message({
"buffer": buffer,
})
async def _broadcast_message(self, message):
disconnected = set() disconnected = set()
for ws in self._clients: for ws in self._clients:
try: try:
@@ -26,26 +23,22 @@ class TokenWebSocket:
except ConnectionResetError: except ConnectionResetError:
disconnected.add(ws) disconnected.add(ws)
self._clients -= disconnected 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: async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
"""Handle new WebSocket connections."""
ws = web.WebSocketResponse(heartbeat=30) ws = web.WebSocketResponse(heartbeat=30)
await ws.prepare(request) await ws.prepare(request)
self._clients.add(ws) self._clients.add(ws)
try: try:
await ws.send_json({
"buffer": self._agent.response_buffer.get_text(),
})
async for msg in ws: async for msg in ws:
if msg.type == WSMsgType.ERROR: if msg.type == WSMsgType.ERROR:
print(f"WebSocket connection closed with error: {ws.exception()}") print(f"WebSocket connection closed with error: {ws.exception()}")
finally: finally:
self._clients.remove(ws) 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 aiohttp import web
from ..auto_approver import AutoApprover from ..auto_approver import AutoApprover
from ..web_agent import WebAgent from ..web_agent import WebAgent
from ..web_io_buffer import WebIOBuffer from ..web_io_buffer import WebIOBuffer
from ..working_memory import WorkingMemory from ..working_memory import WorkingMemory
from .auto_approver_websocket import AutoApproverWebSocket from .auto_approver_websocket import AutoApproverWebSocket
from .context_websocket import ContextWebSocket from .context_websocket import ContextWebSocket
from .llm_websocket import LlmWebSocket from .llm_websocket import LlmWebSocket
from .memory_websocket import MemoryWebSocket from .memory_websocket import MemoryWebSocket
from .stdout_websocket import StdoutWebSocket from .response_websocket import ResponseWebSocket
from .token_websocket import TokenWebSocket from .stdout_websocket import StdoutWebSocket
class Websockets: class Websockets:
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory): def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory):
self._llm_ws = LlmWebSocket(agent) self._auto_approver_ws = AutoApproverWebSocket(auto_approver)
self._context_ws = ContextWebSocket(agent) self._context_ws = ContextWebSocket(agent)
self._token_ws = TokenWebSocket(agent) self._llm_ws = LlmWebSocket(agent)
self._stdout_ws = StdoutWebSocket(io_buffer) self._memory_ws = MemoryWebSocket(working_memory)
self._auto_approver_ws = AutoApproverWebSocket(auto_approver) self._response_ws = ResponseWebSocket(agent)
self._memory_ws = MemoryWebSocket(working_memory) self._stdout_ws = StdoutWebSocket(io_buffer)
app.router.add_get("/ws/llm", self._llm_ws.handle_connection) 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/context", self._context_ws.handle_connection)
app.router.add_get("/ws/token", self._token_ws.handle_connection) app.router.add_get("/ws/llm", self._llm_ws.handle_connection)
app.router.add_get("/ws/stdout", self._stdout_ws.handle_connection) app.router.add_get("/ws/memory", self._memory_ws.handle_connection)
app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection) app.router.add_get("/ws/response", self._response_ws.handle_connection)
app.router.add_get("/ws/memory", self._memory_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 collections import defaultdict
from datetime import datetime, timezone from datetime import datetime, timezone
from enum import Enum, auto from enum import Enum, auto
from sys import exit from sys import exit
from threading import Lock from threading import Lock
from typing import Callable, Dict, List, Optional from typing import Callable, Dict, List, Optional
from .base_agent import BaseAgent from .base_agent import BaseAgent
from .command import Command from .command import Command
from .command_result import CommandResult from .command_result import CommandResult
from .iteration_logger import IterationLogger from .iteration_logger import IterationLogger
from .llm_engine import LlmEngine from .llm_engine import LlmEngine
from .response_parser import ResponseParser from .response_buffer import ResponseBuffer
from .system_metrics import SystemMetrics from .response_parser import ResponseParser
from .working_memory import WorkingMemory from .system_metrics import SystemMetrics
from .xml_validator import XMLValidator from .working_memory import WorkingMemory
from .xml_validator import XMLValidator
class LlmState(Enum):
NO_OUTPUT = auto() class LlmState(Enum):
INFERENCE = auto() IDLE = auto()
OUTPUT = auto() INFERENCE = auto()
class WebAgent(BaseAgent): class WebAgent(BaseAgent):
def __init__( def __init__(
self, self,
system_prompt: str, system_prompt: str,
action_schema: str, action_schema: str,
working_memory: WorkingMemory, working_memory: WorkingMemory,
metrics: SystemMetrics, metrics: SystemMetrics,
llms: Dict[str, LlmEngine], llms: Dict[str, LlmEngine],
validator: XMLValidator, validator: XMLValidator,
parser: ResponseParser, parser: ResponseParser,
iteration_logger: IterationLogger, iteration_logger: IterationLogger,
): ):
super().__init__( super().__init__(
system_prompt, system_prompt,
action_schema, action_schema,
working_memory, working_memory,
metrics, metrics,
validator, validator,
parser parser
) )
self._llms = llms self._llms = llms
self._iteration_logger = iteration_logger self._iteration_logger = iteration_logger
self._llm_states: Dict[str, LlmState] = {name: LlmState.NO_OUTPUT for name in llms} self._llm_states: Dict[str, LlmState] = {name: LlmState.IDLE for name in llms}
self._llm_outputs: Dict[str, str] = defaultdict(str) self._response_buffer: ResponseBuffer = ResponseBuffer()
self._validation_error: Optional[str] = None self._validation_error: Optional[str] = None
self._command_result: Optional[CommandResult] = None self._command_result: Optional[CommandResult] = None
self._context = self._compile_context(next(iter(self._llms.values()))) self._context = self._compile_context(next(iter(self._llms.values())))
self._stop_flags: Dict[str, bool] = {name: False for name in llms} self._stop_flags: Dict[str, bool] = {name: False for name in llms}
# Locks # Locks
self._llm_lock = Lock() self._llm_lock = Lock()
self._output_lock = Lock() self._output_lock = Lock()
# Event handlers # Event handlers
self._llm_change_handlers: List[Callable[[str, LlmState], None]] = [] self._llm_change_handlers: List[Callable[[str, LlmState], None]] = []
self._token_handlers: List[Callable[[str, str], None]] = [] self._token_handlers: List[Callable[[str, str], None]] = []
self._context_change_handlers: List[Callable[[str, bool], None]] = [] self._context_change_handlers: List[Callable[[str, bool], None]] = []
# Working memory change handler # Change handlers
self._working_memory.add_change_handler(self._handle_memory_update) self._working_memory.add_change_handler(self._handle_memory_update)
@property @property
def llms(self) -> Dict[str, LlmState]: def llms(self) -> Dict[str, LlmState]:
"""Get current state of all LLMs""" """Get current state of all LLMs"""
with self._llm_lock: with self._llm_lock:
return self._llm_states.copy() return self._llm_states.copy()
@property @property
def context(self) -> str: def response_buffer(self) -> ResponseBuffer:
return self._context return self._response_buffer
@property @property
def command_result(self) -> Optional[CommandResult]: def context(self) -> str:
return self._command_result return self._context
@property @property
def validation_error(self) -> Optional[str]: def command_result(self) -> Optional[CommandResult]:
return self._validation_error return self._command_result
def add_llm_change_handler(self, handler: Callable[[str, LlmState], None]) -> None: @property
"""Add handler for LLM state changes""" def validation_error(self) -> Optional[str]:
if handler not in self._llm_change_handlers: return self._validation_error
self._llm_change_handlers.append(handler)
def add_llm_change_handler(self, handler: Callable[[str, LlmState], None]) -> None:
def add_token_handler(self, handler: Callable[[str, str], None]) -> None: """Add handler for LLM state changes"""
"""Add handler for new tokens""" if handler not in self._llm_change_handlers:
if handler not in self._token_handlers: self._llm_change_handlers.append(handler)
self._token_handlers.append(handler)
def add_context_change_handler(self, handler: Callable[[str, bool], None]) -> None:
def add_context_change_handler(self, handler: Callable[[str, bool], None]) -> None: """Add handler for context changes"""
"""Add handler for context changes""" if handler not in self._context_change_handlers:
if handler not in self._context_change_handlers: self._context_change_handlers.append(handler)
self._context_change_handlers.append(handler)
def modify_context(self, context: str, generated: bool = False) -> None:
def modify_context(self, context: str, generated: bool = False) -> None: """Update context and reset all LLM states"""
"""Update context and reset all LLM states""" with self._llm_lock:
with self._llm_lock: self._context = context
self._context = context self._response_buffer.clear()
self._llm_outputs.clear() for llm_name in self._llms:
for llm_name in self._llms: self._set_llm_state(llm_name, LlmState.IDLE)
self._set_llm_state(llm_name, LlmState.NO_OUTPUT)
for handler in self._context_change_handlers:
for handler in self._context_change_handlers: handler(context, generated)
handler(context, generated)
def run_inference(self, llm_name: str) -> None:
def run_inference(self, llm_name: str) -> None: """Start inference on specified LLM"""
"""Start inference on specified LLM""" if llm_name not in self._llms:
if llm_name not in self._llms: raise ValueError(f"Unknown LLM: {llm_name}")
raise ValueError(f"Unknown LLM: {llm_name}") with self._llm_lock:
if self._llm_states[llm_name] != LlmState.IDLE:
with self._llm_lock: raise RuntimeError(f"LLM {llm_name} is not ready for inference")
if self._llm_states[llm_name] != LlmState.NO_OUTPUT: self._set_llm_state(llm_name, LlmState.INFERENCE)
raise RuntimeError(f"LLM {llm_name} is not ready for inference") self._stop_flags[llm_name] = False
self._set_llm_state(llm_name, LlmState.INFERENCE) llm = self._llms[llm_name]
self._stop_flags[llm_name] = False def should_stop() -> bool:
return self._stop_flags[llm_name]
llm = self._llms[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:
def should_stop() -> bool: print(token, end='', flush=True)
return self._stop_flags[llm_name] with self._output_lock:
self._response_buffer.append_text(token)
response_token_iter = llm.infer(self.system_prompt, self.context, should_stop) with self._llm_lock:
self._set_llm_state(llm_name, LlmState.IDLE)
with self._output_lock:
self._llm_outputs[llm_name] = "" def stop_inference(self, llm_name: str) -> None:
"""Stop ongoing inference for specified LLM"""
for token in response_token_iter: if llm_name not in self._llms:
with self._output_lock: raise ValueError(f"Unknown LLM: {llm_name}")
self._llm_outputs[llm_name] += token self._stop_flags[llm_name] = True
for handler in self._token_handlers: def approve_response(self) -> None:
handler(llm_name, token) """Process approved response from specified LLM"""
timestamp = datetime.now(timezone.utc)
with self._llm_lock: self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text())
self._set_llm_state(llm_name, LlmState.OUTPUT) parse_result = self._parser.parse(timestamp, self._response_buffer.get_text())
if isinstance(parse_result, Command):
def stop_inference(self, llm_name: str) -> None: result = parse_result.execute(self._working_memory)
"""Stop ongoing inference for specified LLM""" self._command_result = result
if llm_name not in self._llms: if result.should_stop:
raise ValueError(f"Unknown LLM: {llm_name}") exit(42)
self._stop_flags[llm_name] = True else:
self._working_memory.update()
def get_output(self, llm_name: str) -> str: else:
"""Get complete output for specified LLM""" parse_result.update()
if llm_name not in self._llms: self._working_memory.update()
raise ValueError(f"Unknown LLM: {llm_name}") self._working_memory.add_entry(parse_result)
with self._output_lock: def _set_llm_state(self, llm_name: str, state: LlmState) -> None:
return self._llm_outputs[llm_name] """Update LLM state and notify handlers"""
self._llm_states[llm_name] = state
def approve_response(self, llm_name: str, response: str) -> None: for handler in self._llm_change_handlers:
"""Process approved response from specified LLM""" handler(llm_name, state)
if llm_name not in self._llms:
raise ValueError(f"Unknown LLM: {llm_name}") def _handle_memory_update(self) -> None:
"""Handle memory updates and update context"""
timestamp = datetime.now(timezone.utc) context = self._compile_context(next(iter(self._llms.values())))
self._iteration_logger.log_iteration(timestamp, self._context, response) self.modify_context(context, True)
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)