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

8
sia/auto_approver.py
View File

@@ -105,7 +105,7 @@ class AutoApprover:
return
self._context_enabled = enabled
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._notify_config_change()
@@ -119,7 +119,7 @@ class AutoApprover:
return
self._response_enabled = enabled
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._notify_config_change()
@@ -138,7 +138,7 @@ class AutoApprover:
if llm_name != self._llm_name:
return
if state == LlmState.OUTPUT and self._response_enabled:
if state == LlmState.IDLE and self._response_enabled:
self._start_response_thread()
else:
self._stop_response_thread()
@@ -179,5 +179,5 @@ class AutoApprover:
if self._stop_event.wait(self._response_timeout):
return
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))

2
sia/llm_engine/__init__.py
View File

@@ -3,7 +3,7 @@ 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]:
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
pass
@abstractmethod

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,

16
sia/llm_engine/mistral_llm_engine.py
View File

@@ -5,6 +5,7 @@ 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__(
@@ -21,7 +22,7 @@ class MistralLlmEngine(LlmEngine):
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]:
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
messages = [
{
"role": "system",
@@ -33,9 +34,18 @@ class MistralLlmEngine(LlmEngine):
},
{
"role": "assistant",
"content": "<",
"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,
@@ -44,12 +54,14 @@ class MistralLlmEngine(LlmEngine):
)
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()

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)

39
sia/util.py
View File

@@ -24,6 +24,45 @@ def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]
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

40
sia/web/api.py
View File

@@ -34,7 +34,9 @@ class Api:
"""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/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)
@@ -73,18 +75,38 @@ class Api:
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")
async def _set_response(self, request: web.Request) -> web.Response:
"""Edit the shared response buffer"""
try:
self._agent.approve_response(llm_name, response)
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."""

37
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))
self._clients = set()
self._agent._response_buffer.add_observer(wrap_async(self._handle_buffer_change))
async def _broadcast_message(self, message: Dict):
"""Broadcast message to all connected clients."""
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:
@@ -27,21 +24,17 @@ class TokenWebSocket:
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()}")

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

@@ -8,21 +8,21 @@ 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
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._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/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/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)

61
sia/web_agent.py
View File

@@ -10,15 +10,15 @@ 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):
NO_OUTPUT = auto()
IDLE = auto()
INFERENCE = auto()
OUTPUT = auto()
class WebAgent(BaseAgent):
def __init__(
@@ -42,8 +42,8 @@ class WebAgent(BaseAgent):
)
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._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())))
@@ -58,7 +58,7 @@ class WebAgent(BaseAgent):
self._token_handlers: List[Callable[[str, str], 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)
@property
@@ -67,6 +67,10 @@ class WebAgent(BaseAgent):
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
@@ -84,11 +88,6 @@ class WebAgent(BaseAgent):
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:
@@ -98,9 +97,9 @@ class WebAgent(BaseAgent):
"""Update context and reset all LLM states"""
with self._llm_lock:
self._context = context
self._llm_outputs.clear()
self._response_buffer.clear()
for llm_name in self._llms:
self._set_llm_state(llm_name, LlmState.NO_OUTPUT)
self._set_llm_state(llm_name, LlmState.IDLE)
for handler in self._context_change_handlers:
handler(context, generated)
@@ -109,32 +108,21 @@ class WebAgent(BaseAgent):
"""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:
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, should_stop)
with self._output_lock:
self._llm_outputs[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._llm_outputs[llm_name] += token
for handler in self._token_handlers:
handler(llm_name, token)
self._response_buffer.append_text(token)
with self._llm_lock:
self._set_llm_state(llm_name, LlmState.OUTPUT)
self._set_llm_state(llm_name, LlmState.IDLE)
def stop_inference(self, llm_name: str) -> None:
"""Stop ongoing inference for specified LLM"""
@@ -142,22 +130,11 @@ class WebAgent(BaseAgent):
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:
def approve_response(self) -> 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)
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