Shared response buffer backend

sia/__main__.py

@@ -13,7 +13,7 @@ from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
 from .web.api import Api
 from .web.static import Static
-from .web.websockts import Websockets
+from .web.websockets import Websockets
 from .web_agent import WebAgent
 from .web_io_buffer import WebIOBuffer
 from .working_memory import WorkingMemory

sia/auto_approver.py

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

sia/llm_engine/__init__.py

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

sia/llm_engine/local_llm_engine.py

@@ -56,7 +56,7 @@ class LocalLlmEngine(LlmEngine):
         else:
             self._logits_processor = None
 
-    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
         """
         Run inference using the system prompt and main context.
 
@@ -70,10 +70,12 @@ class LocalLlmEngine(LlmEngine):
         """
         messages = [
             {"role": "system", "content": system_prompt},
-            {"role": "user", "content": main_context}
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
         ]
         prompt = self._tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
+            messages, tokenize=False,
+            add_generation_prompt=False,
         )
         streamer = TextIteratorStreamer(
             self._tokenizer,

sia/llm_engine/mistral_llm_engine.py

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

sia/llm_engine/qwq_llm_engine.py

@@ -63,7 +63,7 @@ class QwQLlmEngine(LlmEngine):
             self._logits_processor = None
 
                         
-    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
         """
         Run inference using the system prompt and main context.
 
@@ -77,13 +77,14 @@ class QwQLlmEngine(LlmEngine):
         """
         messages = [
             {"role": "system", "content": system_prompt},
-            {"role": "user", "content": main_context}
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
         ]
 
         text = self._tokenizer.apply_chat_template(
             messages,
             tokenize=False,
-            add_generation_prompt=True,
+            add_generation_prompt=False,
         )
             
         streamer = TextIteratorStreamer(

sia/response_buffer.py

@@ -0,0 +1,26 @@
+class ResponseBuffer:
+    def __init__(self):
+        self._text = ""
+        self._observers = []
+    
+    def add_observer(self, callback):
+        self._observers.append(callback)
+    
+    def get_text(self):
+        return self._text
+    
+    def append_text(self, text):
+        self._text += text
+        self._notify_observers()
+    
+    def set_text(self, text):
+        self._text = text
+        self._notify_observers()
+    
+    def clear(self):
+        self._text = ""
+        self._notify_observers()
+    
+    def _notify_observers(self):
+        for observer in self._observers:
+            observer(self._text)

sia/util.py

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

sia/web/api.py

@@ -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:
+    async def _set_response(self, request: web.Request) -> web.Response:
-        """Approve response from specified LLM."""
+        """Edit the shared response buffer"""
-        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)
+            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."""

sia/web/response_websocket.py

@@ -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 .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._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()
         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()}")

sia/web/websockets.py

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

sia/web_agent.py

@@ -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_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._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, self._response_buffer.get_text(), should_stop)
-        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:
+            print(token, end='', flush=True)
             with self._output_lock:
-                self._llm_outputs[llm_name] += token
+                self._response_buffer.append_text(token)
-            
-            for handler in self._token_handlers:
-                handler(llm_name, 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:
+    def approve_response(self) -> None:
-        """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)
+        self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text())
-        parse_result = self._parser.parse(timestamp, response)
+        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