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

11
run.sh
View File

@@ -2,6 +2,15 @@
set -e
function chown_iterations() {
if [ -d "./iterations" ] && [ "$(find ./iterations/ ! -user $USER -o ! -group $USER 2>/dev/null)" ]; then
echo "Chowning iterations directory"
sudo chown -R $USER:$USER ./iterations/
fi
}
trap chown_iterations EXIT
docker build \
--tag sia \
.
@@ -13,6 +22,6 @@ docker run \
-p 8080:8080 \
-v /$(pwd)/model/:/root/model/ \
-v /$(pwd)/.env:/root/.env \
-v /$(pwd)/iterations/:/root/sia/iterations/ \
-v /$(pwd)/iterations/:/root/iterations/ \
-v /$(pwd)/.git/:/root/sia/.git/ \
sia "$@"

3
sia/__main__.py
View File

@@ -34,7 +34,8 @@ class Main:
self._llms['local'] = LocalLlmEngine(
config.local_model,
config.local_temperature,
config.local_token_limit
config.local_token_limit,
config.local_api_key,
)
if config.openai_enabled:

10
sia/config.py
View File

@@ -82,6 +82,12 @@ class Config:
default=int(os.getenv('SIA_LOCAL_TOKEN_LIMIT', '2048')),
help='Local LLM token limit (env: SIA_LOCAL_TOKEN_LIMIT)'
)
parser.add_argument(
'--local-api-key',
type=str,
default=os.getenv('SIA_LOCAL_API_KEY'),
help='API key for local models (env: SIA_LOCAL_API_KEY)'
)
# OpenAI configuration
parser.add_argument(
@@ -234,6 +240,10 @@ class Config:
def local_token_limit(self) -> int:
return self.args.local_token_limit
@property
def local_api_key(self) -> Optional[str]:
return self.args.local_api_key
# OpenAI properties
@property
def openai_enabled(self) -> bool:

15
sia/hf_llm_engine.py
View File

@@ -1,6 +1,6 @@
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoConfig
from typing import Iterator, Optional
from typing import Iterator, Optional, Callable
from .llm_engine import LlmEngine
@@ -30,25 +30,23 @@ class HfLlmEngine(LlmEngine):
self._config = AutoConfig.from_pretrained(model, token=api_token)
self._client = InferenceClient(token=api_token)
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
"""
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
token_count=self.token_count(system_prompt, main_context)
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
def stream_wrapper():
stream = self._client.chat_completion(
model=self._model,
messages=messages,
@@ -56,10 +54,15 @@ class HfLlmEngine(LlmEngine):
stream=True
)
try:
for response in stream:
if should_stop():
stream.close()
break
if content := response.choices[0].delta.content:
yield content
return stream_wrapper()
finally:
stream.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
messages = [

4
sia/llm_engine.py
View File

@@ -1,9 +1,9 @@
from typing import Iterator
from typing import Callable, Iterator
from abc import ABC, abstractmethod
class LlmEngine(ABC):
@abstractmethod
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
pass
@abstractmethod

22
sia/local_llm_engine.py
View File

@@ -1,5 +1,5 @@
from threading import Thread
from typing import Iterator, Optional
from typing import Iterator, Optional, Callable
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
import torch
@@ -49,13 +49,14 @@ class LocalLlmEngine(LlmEngine):
return_full_text=False,
)
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
"""
Run inference using the system prompt and main context, while validating actions against the provided XML schema.
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
@@ -71,16 +72,21 @@ class LocalLlmEngine(LlmEngine):
self._tokenizer,
skip_prompt=True
)
pipeline_kwargs = dict(
generation_thread = Thread(target=self._pipeline, kwargs=dict(
text_inputs=prompt,
do_sample=True,
temperature=self._temperature,
max_new_tokens=self.token_limit(),
streamer=streamer
)
thread = Thread(target=self._pipeline, kwargs=pipeline_kwargs)
thread.start()
return util.stop_before_value(streamer, '<|eot_id|>')
))
generation_thread.start()
for text in util.stop_before_value(streamer, '<|eot_id|>'):
yield text
if should_stop():
break
generation_thread.join()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""

19
sia/mistral_llm_engine.py
View File

@@ -1,6 +1,4 @@
from typing import Iterator, Optional
from abc import ABC, abstractmethod
import os
from typing import Iterator, Optional, Callable
from mistralai import Mistral
from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest
@@ -23,11 +21,7 @@ class MistralLlmEngine(LlmEngine):
self._client = Mistral(api_key=api_key)
self._tokenizer = MistralTokenizer.v3()
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
messages = [
SystemMessage(content=system_prompt),
UserMessage(content=main_context),
]
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
messages = [
{
"role": "system",
@@ -48,9 +42,16 @@ class MistralLlmEngine(LlmEngine):
messages=messages,
temperature=self._temperature,
)
try:
for chunk in stream_response:
if should_stop():
stream_response.close()
break
if content := chunk.data.choices[0].delta.content:
yield chunk.data.choices[0].delta.content
yield content
finally:
stream_response.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
messages = [

20
sia/openai_llm_engine.py
View File

@@ -1,4 +1,4 @@
from typing import Iterator
from typing import Callable, Iterator
import openai
import tiktoken
@@ -34,17 +34,7 @@ class OpenAILlmEngine(LlmEngine):
api_key=api_key,
)
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
"""
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
Returns:
Iterator[str]: An iterator that yields the generated text in chunks.
"""
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
@@ -57,9 +47,15 @@ class OpenAILlmEngine(LlmEngine):
stream=True,
)
try:
for chunk in stream:
if should_stop():
break
if content := chunk.choices[0].delta.content:
yield content
finally:
stream.close()
#stream.response.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""

10
sia/web/api.py
View File

@@ -30,6 +30,7 @@ class Api:
def _init_routes(self):
"""Initialize REST API and WebSocket routes."""
self._app.router.add_post("/api/inference/{llm}", self._run_inference)
self._app.router.add_post("/api/inference/{llm}/stop", self._stop_inference)
self._app.router.add_post("/api/approve/{llm}", self._approve_response)
self._app.router.add_post("/api/context", self._modify_context)
self._app.router.add_post("/api/input", self._send_input)
@@ -60,6 +61,15 @@ class Api:
except (ValueError, RuntimeError) as e:
return web.Response(status=400, text=str(e))
async def _stop_inference(self, request: web.Request) -> web.Response:
"""Stop inference on specified LLM."""
llm_name = request.match_info["llm"]
try:
self._agent.stop_inference(llm_name)
return web.Response(status=200)
except ValueError as e:
return web.Response(status=400, text=str(e))
async def _approve_response(self, request: web.Request) -> web.Response:
"""Approve response from specified LLM."""
llm_name = request.match_info["llm"]

14
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:

4
tasks/fix_unit_tests.txt Normal file
View File

@@ -0,0 +1,4 @@
The unit tests for the sia project are no longer up-to-date.
Find failing unit tests and fix them.
Find missing unit tests and write them.
Make regular commits to the git repo in the /root/sia directory but don't push them.

5
tasks/improve_sia.txt Normal file
View File

@@ -0,0 +1,5 @@
The SIA source is located in /root/sia.
Not all features are implemented yet.
Look at the readme and code to find what is missing.
Make sure to unit test your work.
Make commits to the git repo when a features are done but don't push them.

7
web/src/components/App.jsx
View File

@@ -159,6 +159,12 @@ const App = () => {
}
};
const handleStop = () => {
fetch(`/api/inference/${activeLlm}/stop`, {
method: 'POST'
});
};
const handleSendInput = () => {
if (input.trim()) {
fetch('/api/input', {
@@ -347,6 +353,7 @@ const App = () => {
output={output}
autoApproverConfig={autoApproverConfig}
onApprove={handleApprove}
onStop={handleStop}
onToggleDiff={() => setShowDiff(!showDiff)}
onSendInput={handleSendInput}
onClearOutput={handleClearOutput}

9
web/src/components/Sidebar.jsx
View File

@@ -16,6 +16,7 @@ export const Sidebar = ({
output,
autoApproverConfig,
onApprove,
onStop,
onToggleDiff,
onSendInput,
onClearOutput,
@@ -31,7 +32,6 @@ export const Sidebar = ({
${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
`}>
<div className="p-4 space-y-4 overflow-y-auto h-full">
{/* Existing controls */}
<Select value={activeLlm} onValueChange={onLlmChange}>
<SelectTrigger>
<SelectValue placeholder="Select LLM" />
@@ -62,14 +62,13 @@ export const Sidebar = ({
</div>
<Button
onClick={onApprove}
disabled={llmState === LlmState.INFERENCE}
onClick={llmState === LlmState.INFERENCE ? onStop : onApprove}
className="w-full"
>
{
llmState === LlmState.NO_OUTPUT ? 'Inference' :
llmState === LlmState.OUTPUT ? 'Approve' :
'In Progress'
llmState === LlmState.INFERENCE ? 'Stop' :
'Approve'
}
</Button>

9
web/src/components/editors/MemoryEditor.jsx
View File

@@ -46,6 +46,7 @@ export const MemoryEditor = ({
const handleLoadIteration = async (event) => {
const file = event.target.files[0];
if (file) {
try {
const content = await file.text();
const response = await fetch('/api/memory/load_iteration', {
method: 'POST',
@@ -53,10 +54,16 @@ export const MemoryEditor = ({
body: JSON.stringify({ content: content })
});
if (!response.ok) throw new Error('Failed to load iteration');
} catch (error) {
if (error instanceof DOMException) {
alert('Failed to read the file. Check permissions. chown -R $USER:$USER iteration');
} else {
throw error;
}
}
}
};
const handleLoadIterationClick = () => {
const input = document.createElement('input');
input.type = 'file';