llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow stopping of inference

Browse Source
This commit is contained in:
geens
2024-11-30 12:44:13 +01:00
parent a0353d0d49
commit e71bd7e9eb
15 changed files with 139 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
sia/web_agent.py
View File

@@ -46,6 +46,7 @@ class WebAgent(BaseAgent):
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()
@@ -112,9 +113,14 @@ class WebAgent(BaseAgent):
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]
response_token_iter = llm.infer(self.system_prompt, self.context)
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] = ""
@@ -129,6 +135,12 @@ class WebAgent(BaseAgent):
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:
@@ -164,4 +176,4 @@ class WebAgent(BaseAgent):
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.modify_context(context, True)