Enable multiple llms
This commit is contained in:
Niels Geens
257
sia/web_agent.py
257
sia/web_agent.py
@@ -1,6 +1,7 @@
|
||||
from enum import Enum, auto
|
||||
from threading import Thread, Lock
|
||||
from typing import Callable, List, Optional
|
||||
from threading import Lock
|
||||
from typing import Callable, Dict, List, Optional
|
||||
from collections import defaultdict
|
||||
|
||||
from .base_agent import BaseAgent
|
||||
from .command import Command
|
||||
@@ -12,200 +13,148 @@ from .system_metrics import SystemMetrics
|
||||
from .working_memory import WorkingMemory
|
||||
from .xml_validator import XMLValidator
|
||||
|
||||
class WebAgentState(Enum):
|
||||
"""
|
||||
States for the web agent state machine.
|
||||
"""
|
||||
UPDATE = auto()
|
||||
"""Updating system metrics and working memory entries"""
|
||||
CONTEXT_APPROVAL = auto()
|
||||
"""Waiting for human approval of context"""
|
||||
class LlmState(Enum):
|
||||
NO_OUTPUT = auto()
|
||||
INFERENCE = auto()
|
||||
"""Processing context through LLM"""
|
||||
RESPONSE_APPROVAL = auto()
|
||||
"""Waiting for human approval of LLM response"""
|
||||
STOPPED = auto()
|
||||
OUTPUT = auto()
|
||||
|
||||
class WebAgent(BaseAgent):
|
||||
"""
|
||||
Agent implementation for interactive web interface.
|
||||
|
||||
Uses a state machine to allow human intervention between steps.
|
||||
Broadcasts state changes to registered handlers.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
system_prompt: str,
|
||||
action_schema: str,
|
||||
working_memory: WorkingMemory,
|
||||
metrics: SystemMetrics,
|
||||
llm: LlmEngine,
|
||||
llms: Dict[str, LlmEngine],
|
||||
validator: XMLValidator,
|
||||
parser: ResponseParser,
|
||||
iteration_logger: IterationLogger,
|
||||
):
|
||||
"""
|
||||
Initialize web agent with required components.
|
||||
"""
|
||||
super().__init__(
|
||||
system_prompt,
|
||||
action_schema,
|
||||
working_memory,
|
||||
metrics, llm,
|
||||
metrics,
|
||||
validator,
|
||||
parser
|
||||
)
|
||||
self._response = ""
|
||||
self._state = WebAgentState.CONTEXT_APPROVAL
|
||||
self._validation_error: Optional[str] = None
|
||||
self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
|
||||
self._context_change_handlers: List[Callable[[str], None]] = []
|
||||
self._response_change_handlers: List[Callable[[str, str], None]] = []
|
||||
self._command_result: Optional[CommandResult] = None
|
||||
self._context = self._compile_context()
|
||||
self._state_lock = Lock()
|
||||
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())))
|
||||
|
||||
# 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]] = []
|
||||
|
||||
@property
|
||||
def state(self) -> WebAgentState:
|
||||
"""Get the current state of the agent."""
|
||||
return self._state
|
||||
|
||||
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]:
|
||||
"""Get the result of the last command execution."""
|
||||
return self._command_result
|
||||
|
||||
@property
|
||||
def validation_error(self) -> Optional[str]:
|
||||
"""Get the validation error, if any."""
|
||||
return self._validation_error
|
||||
|
||||
@property
|
||||
def response(self) -> str:
|
||||
"""Get the current response."""
|
||||
return self._response
|
||||
|
||||
@property
|
||||
def context(self) -> str:
|
||||
"""Get the current context."""
|
||||
return self._context
|
||||
|
||||
def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
|
||||
"""
|
||||
Add a callback for state changes.
|
||||
|
||||
Args:
|
||||
handler: Function to call with new state
|
||||
"""
|
||||
if handler not in self._state_change_handlers:
|
||||
self._state_change_handlers.append(handler)
|
||||
|
||||
def add_context_change_handler(self, handler: Callable[[str], None]) -> None:
|
||||
"""
|
||||
Add a callback for context changes.
|
||||
|
||||
Args:
|
||||
handler: Function to call with new context
|
||||
"""
|
||||
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 add_response_change_handler(self, handler: Callable[[str, str], None]) -> None:
|
||||
"""
|
||||
Add a callback for response changes.
|
||||
|
||||
Args:
|
||||
handler: Function to call with new response
|
||||
"""
|
||||
if handler not in self._response_change_handlers:
|
||||
self._response_change_handlers.append(handler)
|
||||
|
||||
def modify_context(self, context: str) -> None:
|
||||
with self._state_lock:
|
||||
if self._state != WebAgentState.CONTEXT_APPROVAL:
|
||||
error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
|
||||
raise Exception(error_msg)
|
||||
Thread(target=self._set_context, args=(context,)).start()
|
||||
|
||||
def modify_response(self, response: str) -> None:
|
||||
started = False
|
||||
try:
|
||||
self._state_lock.acquire()
|
||||
if self._state != WebAgentState.RESPONSE_APPROVAL:
|
||||
error_msg = f"Not in RESPONSE_APPROVAL state (current: {self._state})"
|
||||
raise Exception(error_msg)
|
||||
Thread(target=self._set_response, args=(response, self._state_lock)).start()
|
||||
started = True
|
||||
finally:
|
||||
if not started:
|
||||
self._state_lock.release()
|
||||
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)
|
||||
|
||||
def approve_context(self) -> None:
|
||||
with self._state_lock:
|
||||
if self._state != WebAgentState.CONTEXT_APPROVAL:
|
||||
error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
|
||||
raise Exception(error_msg)
|
||||
self._set_state(WebAgentState.INFERENCE)
|
||||
Thread(target=self._approve_context_thread).start()
|
||||
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)
|
||||
|
||||
llm = self._llms[llm_name]
|
||||
response_token_iter = llm.infer(self.system_prompt, self.context)
|
||||
|
||||
with self._output_lock:
|
||||
self._llm_outputs[llm_name] = ""
|
||||
|
||||
def approve_response(self) -> None:
|
||||
with self._state_lock:
|
||||
if self._state != WebAgentState.RESPONSE_APPROVAL:
|
||||
raise Exception("Not in RESPONSE_APPROVAL state")
|
||||
self._set_state(WebAgentState.UPDATE)
|
||||
Thread(target=self._approve_response_thread).start()
|
||||
|
||||
def _approve_context_thread(self) -> None:
|
||||
self._set_response("")
|
||||
response_token_iter = self._llm.infer(self.system_prompt, self.context)
|
||||
response = ""
|
||||
for token in response_token_iter:
|
||||
response += token
|
||||
self._set_response(response)
|
||||
print(f"{token}", end='')
|
||||
self._set_state(WebAgentState.RESPONSE_APPROVAL)
|
||||
print()
|
||||
|
||||
def _approve_response_thread(self) -> None:
|
||||
self._iteration_logger.log_iteration(self._context, self._response)
|
||||
parse_result = self._parser.parse(self._response)
|
||||
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 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}")
|
||||
|
||||
self._iteration_logger.log_iteration(self._context, response)
|
||||
|
||||
parse_result = self._parser.parse(response)
|
||||
if isinstance(parse_result, Command):
|
||||
result = parse_result.execute(self._working_memory)
|
||||
self._command_result = result
|
||||
if result.should_stop:
|
||||
self._set_state(WebAgentState.STOPPED)
|
||||
return
|
||||
self._working_memory.update()
|
||||
if not result.should_stop:
|
||||
self._working_memory.update()
|
||||
else:
|
||||
parse_result.update()
|
||||
self._working_memory.update()
|
||||
self._working_memory.add_entry(parse_result)
|
||||
|
||||
self._set_context(self._compile_context())
|
||||
self._set_state(WebAgentState.CONTEXT_APPROVAL)
|
||||
|
||||
def _set_state(self, state) -> None:
|
||||
"""Notify all handlers of state change."""
|
||||
self._state = state
|
||||
for handler in self._state_change_handlers:
|
||||
handler(self._state)
|
||||
|
||||
def _set_context(self, context) -> None:
|
||||
"""Notify all handlers of context change."""
|
||||
self._context = context
|
||||
for handler in self._context_change_handlers:
|
||||
handler(context)
|
||||
|
||||
def _set_response(self, response: str, lock: Optional[Lock] = None) -> None:
|
||||
"""Set response and notify handlers."""
|
||||
try:
|
||||
self._response = response
|
||||
validation_error = self._validator.validate(response)
|
||||
self._validation_error = validation_error
|
||||
for handler in self._response_change_handlers:
|
||||
handler(response, validation_error)
|
||||
finally:
|
||||
if lock is not None:
|
||||
lock.release()
|
||||
self.modify_context(self._compile_context(self._llms[llm_name]), True)
|
||||
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user