180 lines
6.5 KiB
Python
180 lines
6.5 KiB
Python
from datetime import datetime, timezone
|
|
from enum import Enum, auto
|
|
from threading import Lock
|
|
from typing import Callable, Dict, List, Optional
|
|
from collections import defaultdict
|
|
|
|
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 not result.should_stop:
|
|
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)
|