154 lines
5.4 KiB
Python
154 lines
5.4 KiB
Python
from datetime import datetime, timezone
|
|
from enum import Enum, auto
|
|
from sys import exit
|
|
from typing import Callable, Dict, List
|
|
|
|
from .base_agent import BaseAgent
|
|
from .command import Command
|
|
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
|
|
|
|
class AgentState(Enum):
|
|
IDLE = auto()
|
|
INFERENCE = auto()
|
|
PROCESSING_RESPONSE = auto()
|
|
|
|
class WebAgent(BaseAgent):
|
|
def __init__(
|
|
self,
|
|
system_prompt: str,
|
|
action_schema: str,
|
|
working_memory: WorkingMemory,
|
|
metrics: SystemMetrics,
|
|
llms: Dict[str, LlmEngine],
|
|
parser: ResponseParser,
|
|
iteration_logger: IterationLogger,
|
|
):
|
|
super().__init__(
|
|
system_prompt,
|
|
action_schema,
|
|
working_memory,
|
|
metrics,
|
|
parser
|
|
)
|
|
|
|
self._llms = llms
|
|
self._selected_llm = list(llms.keys())[0]
|
|
self._state: AgentState = AgentState.IDLE
|
|
|
|
self._iteration_logger = iteration_logger
|
|
self._response_buffer: ResponseBuffer = ResponseBuffer()
|
|
|
|
self._state_change_handlers: List[Callable[[AgentState], None]] = []
|
|
self._selected_llm_change_handlers: List[Callable[[str], None]] = []
|
|
self._context_change_handlers: List[Callable[[Dict], None]] = []
|
|
|
|
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 llms(self) -> List[str]:
|
|
return self._llms.keys()
|
|
|
|
@property
|
|
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 state(self) -> AgentState:
|
|
return self._state
|
|
|
|
def add_state_change_handler(self, handler: Callable[[AgentState], None]) -> None:
|
|
self._state_change_handlers.append(handler)
|
|
|
|
@property
|
|
def context_info(self) -> Dict:
|
|
return dict(self._compiled_context.attrib)
|
|
|
|
def add_context_change_handler(self, handler: Callable[[Dict], None]) -> None:
|
|
self._context_change_handlers.append(handler)
|
|
|
|
def run_inference(self) -> None:
|
|
"""Start inference on specified LLM"""
|
|
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:
|
|
self._response_buffer.append_text(token)
|
|
self._update_state(AgentState.IDLE)
|
|
|
|
def stop_inference(self) -> None:
|
|
"""Stop ongoing inference for specified LLM"""
|
|
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._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):
|
|
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)
|
|
self._update_state(AgentState.IDLE)
|
|
|
|
def _update_state(self, state: AgentState) -> None:
|
|
"""Update LLM state and notify handlers"""
|
|
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])
|
|
self._notify_context_change()
|
|
|
|
def _notify_context_change(self):
|
|
"""Notify all context change handlers."""
|
|
for handler in self._context_change_handlers:
|
|
handler(self.context_info)
|
|
|
|
def _notify_selected_llm_change(self):
|
|
"""Notify all selected LLM change handlers."""
|
|
for handler in self._selected_llm_change_handlers:
|
|
handler(self._selected_llm) |