Shared response buffer backend
This commit is contained in:
Niels Geens
341
sia/web_agent.py
341
sia/web_agent.py
@@ -1,182 +1,159 @@
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum, auto
|
||||
from sys import exit
|
||||
from threading import Lock
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
from .base_agent import BaseAgent
|
||||
from .command import Command
|
||||
from .command_result import CommandResult
|
||||
from .iteration_logger import IterationLogger
|
||||
from .llm_engine import LlmEngine
|
||||
from .response_parser import ResponseParser
|
||||
from .system_metrics import SystemMetrics
|
||||
from .working_memory import WorkingMemory
|
||||
from .xml_validator import XMLValidator
|
||||
|
||||
class LlmState(Enum):
|
||||
NO_OUTPUT = auto()
|
||||
INFERENCE = auto()
|
||||
OUTPUT = auto()
|
||||
|
||||
class WebAgent(BaseAgent):
|
||||
def __init__(
|
||||
self,
|
||||
system_prompt: str,
|
||||
action_schema: str,
|
||||
working_memory: WorkingMemory,
|
||||
metrics: SystemMetrics,
|
||||
llms: Dict[str, LlmEngine],
|
||||
validator: XMLValidator,
|
||||
parser: ResponseParser,
|
||||
iteration_logger: IterationLogger,
|
||||
):
|
||||
super().__init__(
|
||||
system_prompt,
|
||||
action_schema,
|
||||
working_memory,
|
||||
metrics,
|
||||
validator,
|
||||
parser
|
||||
)
|
||||
self._llms = llms
|
||||
self._iteration_logger = iteration_logger
|
||||
self._llm_states: Dict[str, LlmState] = {name: LlmState.NO_OUTPUT for name in llms}
|
||||
self._llm_outputs: Dict[str, str] = defaultdict(str)
|
||||
self._validation_error: Optional[str] = None
|
||||
self._command_result: Optional[CommandResult] = None
|
||||
self._context = self._compile_context(next(iter(self._llms.values())))
|
||||
self._stop_flags: Dict[str, bool] = {name: False for name in llms}
|
||||
|
||||
# Locks
|
||||
self._llm_lock = Lock()
|
||||
self._output_lock = Lock()
|
||||
|
||||
# Event handlers
|
||||
self._llm_change_handlers: List[Callable[[str, LlmState], None]] = []
|
||||
self._token_handlers: List[Callable[[str, str], None]] = []
|
||||
self._context_change_handlers: List[Callable[[str, bool], None]] = []
|
||||
|
||||
# Working memory change handler
|
||||
self._working_memory.add_change_handler(self._handle_memory_update)
|
||||
|
||||
@property
|
||||
def llms(self) -> Dict[str, LlmState]:
|
||||
"""Get current state of all LLMs"""
|
||||
with self._llm_lock:
|
||||
return self._llm_states.copy()
|
||||
|
||||
@property
|
||||
def context(self) -> str:
|
||||
return self._context
|
||||
|
||||
@property
|
||||
def command_result(self) -> Optional[CommandResult]:
|
||||
return self._command_result
|
||||
|
||||
@property
|
||||
def validation_error(self) -> Optional[str]:
|
||||
return self._validation_error
|
||||
|
||||
def add_llm_change_handler(self, handler: Callable[[str, LlmState], None]) -> None:
|
||||
"""Add handler for LLM state changes"""
|
||||
if handler not in self._llm_change_handlers:
|
||||
self._llm_change_handlers.append(handler)
|
||||
|
||||
def add_token_handler(self, handler: Callable[[str, str], None]) -> None:
|
||||
"""Add handler for new tokens"""
|
||||
if handler not in self._token_handlers:
|
||||
self._token_handlers.append(handler)
|
||||
|
||||
def add_context_change_handler(self, handler: Callable[[str, bool], None]) -> None:
|
||||
"""Add handler for context changes"""
|
||||
if handler not in self._context_change_handlers:
|
||||
self._context_change_handlers.append(handler)
|
||||
|
||||
def modify_context(self, context: str, generated: bool = False) -> None:
|
||||
"""Update context and reset all LLM states"""
|
||||
with self._llm_lock:
|
||||
self._context = context
|
||||
self._llm_outputs.clear()
|
||||
for llm_name in self._llms:
|
||||
self._set_llm_state(llm_name, LlmState.NO_OUTPUT)
|
||||
|
||||
for handler in self._context_change_handlers:
|
||||
handler(context, generated)
|
||||
|
||||
def run_inference(self, llm_name: str) -> None:
|
||||
"""Start inference on specified LLM"""
|
||||
if llm_name not in self._llms:
|
||||
raise ValueError(f"Unknown LLM: {llm_name}")
|
||||
|
||||
with self._llm_lock:
|
||||
if self._llm_states[llm_name] != LlmState.NO_OUTPUT:
|
||||
raise RuntimeError(f"LLM {llm_name} is not ready for inference")
|
||||
self._set_llm_state(llm_name, LlmState.INFERENCE)
|
||||
self._stop_flags[llm_name] = False
|
||||
|
||||
llm = self._llms[llm_name]
|
||||
|
||||
def should_stop() -> bool:
|
||||
return self._stop_flags[llm_name]
|
||||
|
||||
response_token_iter = llm.infer(self.system_prompt, self.context, should_stop)
|
||||
|
||||
with self._output_lock:
|
||||
self._llm_outputs[llm_name] = ""
|
||||
|
||||
for token in response_token_iter:
|
||||
with self._output_lock:
|
||||
self._llm_outputs[llm_name] += token
|
||||
|
||||
for handler in self._token_handlers:
|
||||
handler(llm_name, token)
|
||||
|
||||
with self._llm_lock:
|
||||
self._set_llm_state(llm_name, LlmState.OUTPUT)
|
||||
|
||||
def stop_inference(self, llm_name: str) -> None:
|
||||
"""Stop ongoing inference for specified LLM"""
|
||||
if llm_name not in self._llms:
|
||||
raise ValueError(f"Unknown LLM: {llm_name}")
|
||||
self._stop_flags[llm_name] = True
|
||||
|
||||
def get_output(self, llm_name: str) -> str:
|
||||
"""Get complete output for specified LLM"""
|
||||
if llm_name not in self._llms:
|
||||
raise ValueError(f"Unknown LLM: {llm_name}")
|
||||
|
||||
with self._output_lock:
|
||||
return self._llm_outputs[llm_name]
|
||||
|
||||
def approve_response(self, llm_name: str, response: str) -> None:
|
||||
"""Process approved response from specified LLM"""
|
||||
if llm_name not in self._llms:
|
||||
raise ValueError(f"Unknown LLM: {llm_name}")
|
||||
|
||||
timestamp = datetime.now(timezone.utc)
|
||||
self._iteration_logger.log_iteration(timestamp, self._context, response)
|
||||
parse_result = self._parser.parse(timestamp, response)
|
||||
if isinstance(parse_result, Command):
|
||||
result = parse_result.execute(self._working_memory)
|
||||
self._command_result = result
|
||||
if result.should_stop:
|
||||
exit(42)
|
||||
else:
|
||||
self._working_memory.update()
|
||||
else:
|
||||
parse_result.update()
|
||||
self._working_memory.update()
|
||||
self._working_memory.add_entry(parse_result)
|
||||
|
||||
def _set_llm_state(self, llm_name: str, state: LlmState) -> None:
|
||||
"""Update LLM state and notify handlers"""
|
||||
self._llm_states[llm_name] = state
|
||||
for handler in self._llm_change_handlers:
|
||||
handler(llm_name, state)
|
||||
|
||||
def _handle_memory_update(self) -> None:
|
||||
"""Handle memory updates and update context"""
|
||||
context = self._compile_context(next(iter(self._llms.values())))
|
||||
self.modify_context(context, True)
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum, auto
|
||||
from sys import exit
|
||||
from threading import Lock
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
from .base_agent import BaseAgent
|
||||
from .command import Command
|
||||
from .command_result import CommandResult
|
||||
from .iteration_logger import IterationLogger
|
||||
from .llm_engine import LlmEngine
|
||||
from .response_buffer import ResponseBuffer
|
||||
from .response_parser import ResponseParser
|
||||
from .system_metrics import SystemMetrics
|
||||
from .working_memory import WorkingMemory
|
||||
from .xml_validator import XMLValidator
|
||||
|
||||
class LlmState(Enum):
|
||||
IDLE = auto()
|
||||
INFERENCE = auto()
|
||||
|
||||
class WebAgent(BaseAgent):
|
||||
def __init__(
|
||||
self,
|
||||
system_prompt: str,
|
||||
action_schema: str,
|
||||
working_memory: WorkingMemory,
|
||||
metrics: SystemMetrics,
|
||||
llms: Dict[str, LlmEngine],
|
||||
validator: XMLValidator,
|
||||
parser: ResponseParser,
|
||||
iteration_logger: IterationLogger,
|
||||
):
|
||||
super().__init__(
|
||||
system_prompt,
|
||||
action_schema,
|
||||
working_memory,
|
||||
metrics,
|
||||
validator,
|
||||
parser
|
||||
)
|
||||
self._llms = llms
|
||||
self._iteration_logger = iteration_logger
|
||||
self._llm_states: Dict[str, LlmState] = {name: LlmState.IDLE for name in llms}
|
||||
self._response_buffer: ResponseBuffer = ResponseBuffer()
|
||||
self._validation_error: Optional[str] = None
|
||||
self._command_result: Optional[CommandResult] = None
|
||||
self._context = self._compile_context(next(iter(self._llms.values())))
|
||||
self._stop_flags: Dict[str, bool] = {name: False for name in llms}
|
||||
|
||||
# Locks
|
||||
self._llm_lock = Lock()
|
||||
self._output_lock = Lock()
|
||||
|
||||
# Event handlers
|
||||
self._llm_change_handlers: List[Callable[[str, LlmState], None]] = []
|
||||
self._token_handlers: List[Callable[[str, str], None]] = []
|
||||
self._context_change_handlers: List[Callable[[str, bool], None]] = []
|
||||
|
||||
# Change handlers
|
||||
self._working_memory.add_change_handler(self._handle_memory_update)
|
||||
|
||||
@property
|
||||
def llms(self) -> Dict[str, LlmState]:
|
||||
"""Get current state of all LLMs"""
|
||||
with self._llm_lock:
|
||||
return self._llm_states.copy()
|
||||
|
||||
@property
|
||||
def response_buffer(self) -> ResponseBuffer:
|
||||
return self._response_buffer
|
||||
|
||||
@property
|
||||
def context(self) -> str:
|
||||
return self._context
|
||||
|
||||
@property
|
||||
def command_result(self) -> Optional[CommandResult]:
|
||||
return self._command_result
|
||||
|
||||
@property
|
||||
def validation_error(self) -> Optional[str]:
|
||||
return self._validation_error
|
||||
|
||||
def add_llm_change_handler(self, handler: Callable[[str, LlmState], None]) -> None:
|
||||
"""Add handler for LLM state changes"""
|
||||
if handler not in self._llm_change_handlers:
|
||||
self._llm_change_handlers.append(handler)
|
||||
|
||||
def add_context_change_handler(self, handler: Callable[[str, bool], None]) -> None:
|
||||
"""Add handler for context changes"""
|
||||
if handler not in self._context_change_handlers:
|
||||
self._context_change_handlers.append(handler)
|
||||
|
||||
def modify_context(self, context: str, generated: bool = False) -> None:
|
||||
"""Update context and reset all LLM states"""
|
||||
with self._llm_lock:
|
||||
self._context = context
|
||||
self._response_buffer.clear()
|
||||
for llm_name in self._llms:
|
||||
self._set_llm_state(llm_name, LlmState.IDLE)
|
||||
|
||||
for handler in self._context_change_handlers:
|
||||
handler(context, generated)
|
||||
|
||||
def run_inference(self, llm_name: str) -> None:
|
||||
"""Start inference on specified LLM"""
|
||||
if llm_name not in self._llms:
|
||||
raise ValueError(f"Unknown LLM: {llm_name}")
|
||||
with self._llm_lock:
|
||||
if self._llm_states[llm_name] != LlmState.IDLE:
|
||||
raise RuntimeError(f"LLM {llm_name} is not ready for inference")
|
||||
self._set_llm_state(llm_name, LlmState.INFERENCE)
|
||||
self._stop_flags[llm_name] = False
|
||||
llm = self._llms[llm_name]
|
||||
def should_stop() -> bool:
|
||||
return self._stop_flags[llm_name]
|
||||
response_token_iter = llm.infer(self.system_prompt, self.context, self._response_buffer.get_text(), should_stop)
|
||||
for token in response_token_iter:
|
||||
print(token, end='', flush=True)
|
||||
with self._output_lock:
|
||||
self._response_buffer.append_text(token)
|
||||
with self._llm_lock:
|
||||
self._set_llm_state(llm_name, LlmState.IDLE)
|
||||
|
||||
def stop_inference(self, llm_name: str) -> None:
|
||||
"""Stop ongoing inference for specified LLM"""
|
||||
if llm_name not in self._llms:
|
||||
raise ValueError(f"Unknown LLM: {llm_name}")
|
||||
self._stop_flags[llm_name] = True
|
||||
|
||||
def approve_response(self) -> None:
|
||||
"""Process approved response from specified LLM"""
|
||||
timestamp = datetime.now(timezone.utc)
|
||||
self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text())
|
||||
parse_result = self._parser.parse(timestamp, self._response_buffer.get_text())
|
||||
if isinstance(parse_result, Command):
|
||||
result = parse_result.execute(self._working_memory)
|
||||
self._command_result = result
|
||||
if result.should_stop:
|
||||
exit(42)
|
||||
else:
|
||||
self._working_memory.update()
|
||||
else:
|
||||
parse_result.update()
|
||||
self._working_memory.update()
|
||||
self._working_memory.add_entry(parse_result)
|
||||
|
||||
def _set_llm_state(self, llm_name: str, state: LlmState) -> None:
|
||||
"""Update LLM state and notify handlers"""
|
||||
self._llm_states[llm_name] = state
|
||||
for handler in self._llm_change_handlers:
|
||||
handler(llm_name, state)
|
||||
|
||||
def _handle_memory_update(self) -> None:
|
||||
"""Handle memory updates and update context"""
|
||||
context = self._compile_context(next(iter(self._llms.values())))
|
||||
self.modify_context(context, True)
|
||||
Reference in New Issue
Block a user