New web interface, move llm engine to separate process
This commit is contained in:
145
sia/web_agent.py
145
sia/web_agent.py
@@ -1,13 +1,10 @@
|
||||
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 typing import Callable, Dict, List
|
||||
|
||||
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
|
||||
@@ -15,9 +12,10 @@ from .response_parser import ResponseParser
|
||||
from .system_metrics import SystemMetrics
|
||||
from .working_memory import WorkingMemory
|
||||
|
||||
class LlmState(Enum):
|
||||
class AgentState(Enum):
|
||||
IDLE = auto()
|
||||
INFERENCE = auto()
|
||||
PROCESSING_RESPONSE = auto()
|
||||
|
||||
class WebAgent(BaseAgent):
|
||||
def __init__(
|
||||
@@ -37,97 +35,81 @@ class WebAgent(BaseAgent):
|
||||
metrics,
|
||||
parser
|
||||
)
|
||||
|
||||
self._llms = llms
|
||||
self._selected_llm = list(llms.keys())[0]
|
||||
self._state: AgentState = AgentState.IDLE
|
||||
|
||||
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]] = []
|
||||
self._state_change_handlers: List[Callable[[AgentState], None]] = []
|
||||
self._selected_llm_change_handlers: List[Callable[[str], 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()
|
||||
self._update_compiled_context()
|
||||
self._working_memory.add_change_handler(self._update_compiled_context)
|
||||
|
||||
@property
|
||||
def response_buffer(self) -> ResponseBuffer:
|
||||
return self._response_buffer
|
||||
|
||||
@property
|
||||
def context(self) -> str:
|
||||
return self._context
|
||||
|
||||
def llms(self) -> List[str]:
|
||||
return self._llms.keys()
|
||||
|
||||
@property
|
||||
def command_result(self) -> Optional[CommandResult]:
|
||||
return self._command_result
|
||||
|
||||
def active_llm(self) -> str:
|
||||
return self._selected_llm
|
||||
|
||||
@active_llm.setter
|
||||
def active_llm(self, llm_name: str) -> None:
|
||||
if llm_name not in self._llms:
|
||||
raise ValueError(f"Invalid LLM name: {llm_name}")
|
||||
if self._selected_llm == llm_name:
|
||||
return
|
||||
self._selected_llm = llm_name
|
||||
self._update_compiled_context()
|
||||
self._notify_selected_llm_change()
|
||||
|
||||
def add_selected_llm_change_handler(self, handler: Callable[[str], None]) -> None:
|
||||
self._selected_llm_change_handlers.append(handler)
|
||||
|
||||
@property
|
||||
def validation_error(self) -> Optional[str]:
|
||||
return self._validation_error
|
||||
def state(self) -> AgentState:
|
||||
return self._state
|
||||
|
||||
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_state_change_handler(self, handler: Callable[[AgentState], None]) -> None:
|
||||
self._state_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
|
||||
for handler in self._context_change_handlers:
|
||||
handler(context, generated)
|
||||
|
||||
def run_inference(self, llm_name: str) -> None:
|
||||
def run_inference(self) -> 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)
|
||||
if self._state != AgentState.IDLE:
|
||||
raise RuntimeError(f"Not ready for inference")
|
||||
self._update_state(AgentState.INFERENCE)
|
||||
llm = self._llms[self.active_llm]
|
||||
response_token_iter = llm.infer(
|
||||
self.system_prompt,
|
||||
self._compiled_context,
|
||||
self._response_buffer.get_text()
|
||||
)
|
||||
for token in response_token_iter:
|
||||
with self._output_lock:
|
||||
self._response_buffer.append_text(token)
|
||||
with self._llm_lock:
|
||||
self._set_llm_state(llm_name, LlmState.IDLE)
|
||||
self._response_buffer.append_text(token)
|
||||
self._update_state(AgentState.IDLE)
|
||||
|
||||
def stop_inference(self, llm_name: str) -> None:
|
||||
def stop_inference(self) -> 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
|
||||
with self._llm_lock:
|
||||
self._set_llm_state(llm_name, LlmState.IDLE)
|
||||
self._llms[self.active_llm].restart()
|
||||
self._update_state(AgentState.IDLE)
|
||||
|
||||
def approve_response(self) -> None:
|
||||
"""Process approved response from specified LLM"""
|
||||
self._update_state(AgentState.PROCESSING_RESPONSE)
|
||||
timestamp = datetime.now(timezone.utc)
|
||||
self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text())
|
||||
self._iteration_logger.log_iteration(
|
||||
timestamp,
|
||||
self._compiled_context,
|
||||
self._response_buffer.get_text()
|
||||
)
|
||||
parse_result = self._parser.parse(timestamp, self._response_buffer.get_text())
|
||||
self._response_buffer.clear()
|
||||
if isinstance(parse_result, Command):
|
||||
@@ -141,14 +123,13 @@ class WebAgent(BaseAgent):
|
||||
parse_result.update()
|
||||
self._working_memory.update()
|
||||
self._working_memory.add_entry(parse_result)
|
||||
self._update_state(AgentState.IDLE)
|
||||
|
||||
def _set_llm_state(self, llm_name: str, state: LlmState) -> None:
|
||||
def _update_state(self, state: AgentState) -> 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)
|
||||
self._state = state
|
||||
for handler in self._state_change_handlers:
|
||||
handler(state)
|
||||
|
||||
def _update_compiled_context(self):
|
||||
self._compiled_context = self._compile_context(self._llms[self.active_llm])
|
||||
Reference in New Issue
Block a user