fix(threat-analyst): Async AI chat to handle slow LocalAI inference
- LocalAI inference takes 30-60s on ARM64 hardware - Changed RPCD chat handler to async pattern: - Returns poll_id immediately - Background process runs AI query (120s timeout) - Saves result to /var/lib/threat-analyst/chat_*.json - Client polls with poll_id to get result - Updated api.js with chatAsync() that polls automatically - Changed default LocalAI port from 8081 to 8091 - Frontend shows "Thinking..." message with spinner during inference - Uses curl instead of wget (BusyBox wget doesn't support --post-data=-) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
parent
223abb1114
commit
93b9636258
@ -36,10 +36,41 @@ var callGetPending = rpc.declare({
|
|||||||
var callChat = rpc.declare({
|
var callChat = rpc.declare({
|
||||||
object: 'luci.threat-analyst',
|
object: 'luci.threat-analyst',
|
||||||
method: 'chat',
|
method: 'chat',
|
||||||
params: ['message'],
|
params: ['message', 'poll_id'],
|
||||||
expect: { }
|
expect: { }
|
||||||
});
|
});
|
||||||
|
|
||||||
|
// Async chat with polling
|
||||||
|
function chatAsync(message, maxPolls, pollInterval) {
|
||||||
|
maxPolls = maxPolls || 20;
|
||||||
|
pollInterval = pollInterval || 3000;
|
||||||
|
|
||||||
|
return callChat(message, null).then(function(result) {
|
||||||
|
if (!result.pending) {
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
var pollId = result.poll_id;
|
||||||
|
var polls = 0;
|
||||||
|
|
||||||
|
return new Promise(function(resolve, reject) {
|
||||||
|
function poll() {
|
||||||
|
polls++;
|
||||||
|
callChat('poll', pollId).then(function(res) {
|
||||||
|
if (res.response || res.error) {
|
||||||
|
resolve(res);
|
||||||
|
} else if (polls < maxPolls) {
|
||||||
|
setTimeout(poll, pollInterval);
|
||||||
|
} else {
|
||||||
|
resolve({ error: 'AI response timeout (max polls reached)' });
|
||||||
|
}
|
||||||
|
}).catch(reject);
|
||||||
|
}
|
||||||
|
setTimeout(poll, pollInterval);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
var callGenerateRules = rpc.declare({
|
var callGenerateRules = rpc.declare({
|
||||||
object: 'luci.threat-analyst',
|
object: 'luci.threat-analyst',
|
||||||
method: 'generate_rules',
|
method: 'generate_rules',
|
||||||
@ -111,7 +142,8 @@ return baseclass.extend({
|
|||||||
getStatus: callStatus,
|
getStatus: callStatus,
|
||||||
getThreats: callGetThreats,
|
getThreats: callGetThreats,
|
||||||
getPending: callGetPending,
|
getPending: callGetPending,
|
||||||
chat: callChat,
|
chat: chatAsync,
|
||||||
|
chatSync: callChat,
|
||||||
generateRules: callGenerateRules,
|
generateRules: callGenerateRules,
|
||||||
approveRule: callApproveRule,
|
approveRule: callApproveRule,
|
||||||
rejectRule: callRejectRule,
|
rejectRule: callRejectRule,
|
||||||
|
|||||||
@ -250,9 +250,9 @@ return view.extend({
|
|||||||
input.value = '';
|
input.value = '';
|
||||||
messages.scrollTop = messages.scrollHeight;
|
messages.scrollTop = messages.scrollHeight;
|
||||||
|
|
||||||
// Add loading
|
// Add loading (AI can take 30-60s to respond)
|
||||||
var loading = E('div', { 'class': 'ta-message ai', 'id': 'ta-chat-loading' }, [
|
var loading = E('div', { 'class': 'ta-message ai', 'id': 'ta-chat-loading' }, [
|
||||||
E('div', { 'class': 'ta-message-bubble' }, 'Analyzing...')
|
E('div', { 'class': 'ta-message-bubble spinning' }, 'Thinking... (AI inference can take up to 60s)')
|
||||||
]);
|
]);
|
||||||
messages.appendChild(loading);
|
messages.appendChild(loading);
|
||||||
|
|
||||||
|
|||||||
@ -15,7 +15,7 @@ uci_get() { uci -q get "${CONFIG}.$1"; }
|
|||||||
[ -f "$LIB_DIR/analyzer.sh" ] && {
|
[ -f "$LIB_DIR/analyzer.sh" ] && {
|
||||||
localai_url=$(uci_get main.localai_url)
|
localai_url=$(uci_get main.localai_url)
|
||||||
localai_model=$(uci_get main.localai_model)
|
localai_model=$(uci_get main.localai_model)
|
||||||
[ -z "$localai_url" ] && localai_url="http://127.0.0.1:8081"
|
[ -z "$localai_url" ] && localai_url="http://127.0.0.1:8091"
|
||||||
[ -z "$localai_model" ] && localai_model="tinyllama-1.1b-chat-v1.0.Q4_K_M"
|
[ -z "$localai_model" ] && localai_model="tinyllama-1.1b-chat-v1.0.Q4_K_M"
|
||||||
. "$LIB_DIR/analyzer.sh"
|
. "$LIB_DIR/analyzer.sh"
|
||||||
. "$LIB_DIR/appliers.sh"
|
. "$LIB_DIR/appliers.sh"
|
||||||
@ -143,46 +143,54 @@ EOF
|
|||||||
exit 0
|
exit 0
|
||||||
fi
|
fi
|
||||||
|
|
||||||
# Check LocalAI
|
# Generate request ID
|
||||||
if ! wget -q -O /dev/null --timeout=2 "${localai_url}/v1/models" 2>/dev/null; then
|
req_id=$(head -c 8 /dev/urandom | md5sum | head -c 8)
|
||||||
|
req_file="$STATE_DIR/chat_${req_id}.json"
|
||||||
|
mkdir -p "$STATE_DIR"
|
||||||
|
|
||||||
|
# Check for pending response from previous request
|
||||||
|
poll_id=$(echo "$input" | jsonfilter -e '@.poll_id' 2>/dev/null)
|
||||||
|
if [ -n "$poll_id" ]; then
|
||||||
|
poll_file="$STATE_DIR/chat_${poll_id}.json"
|
||||||
|
if [ -f "$poll_file" ]; then
|
||||||
|
cat "$poll_file"
|
||||||
|
rm -f "$poll_file"
|
||||||
|
else
|
||||||
|
echo '{"pending":true,"poll_id":"'"$poll_id"'"}'
|
||||||
|
fi
|
||||||
|
exit 0
|
||||||
|
fi
|
||||||
|
|
||||||
|
# Check LocalAI quickly
|
||||||
|
if ! curl -s --max-time 2 "${localai_url}/v1/models" >/dev/null 2>&1; then
|
||||||
echo '{"error":"LocalAI not available","suggestion":"Start LocalAI: localaictl start"}'
|
echo '{"error":"LocalAI not available","suggestion":"Start LocalAI: localaictl start"}'
|
||||||
exit 0
|
exit 0
|
||||||
fi
|
fi
|
||||||
|
|
||||||
# Get threat context
|
# Start background AI query
|
||||||
threat_context=""
|
(
|
||||||
if command -v cscli >/dev/null 2>&1; then
|
escaped_message=$(printf '%s' "$message" | sed 's/\\/\\\\/g; s/"/\\"/g' | tr '\n' ' ')
|
||||||
threat_context=$(cscli alerts list -o json --limit 10 2>/dev/null | head -c 3000)
|
prompt="You are SecuBox AI security assistant. $escaped_message"
|
||||||
fi
|
|
||||||
|
|
||||||
# Build prompt
|
response=$(curl -s --max-time 120 -X POST "${localai_url}/v1/chat/completions" \
|
||||||
prompt="You are SecuBox Threat Analyst, an AI security assistant for an OpenWrt router.
|
-H "Content-Type: application/json" \
|
||||||
|
-d "{\"model\":\"$localai_model\",\"messages\":[{\"role\":\"user\",\"content\":\"$prompt\"}],\"max_tokens\":256,\"temperature\":0.7}" 2>/dev/null)
|
||||||
|
|
||||||
Current security context:
|
if [ -n "$response" ]; then
|
||||||
- Recent CrowdSec alerts: $threat_context
|
content=$(echo "$response" | jsonfilter -e '@.choices[0].message.content' 2>/dev/null)
|
||||||
|
if [ -n "$content" ]; then
|
||||||
|
content=$(printf '%s' "$content" | sed 's/\\/\\\\/g; s/"/\\"/g' | tr '\n' ' ')
|
||||||
|
printf '{"response":"%s"}' "$content" > "$req_file"
|
||||||
|
else
|
||||||
|
echo '{"error":"Empty AI response"}' > "$req_file"
|
||||||
|
fi
|
||||||
|
else
|
||||||
|
echo '{"error":"AI query failed"}' > "$req_file"
|
||||||
|
fi
|
||||||
|
) &
|
||||||
|
|
||||||
User message: $message
|
# Return poll ID immediately
|
||||||
|
echo '{"pending":true,"poll_id":"'"$req_id"'"}'
|
||||||
Provide helpful, actionable security advice. If asked about threats, analyze the context. If asked to generate rules, provide specific patterns for mitmproxy/CrowdSec/WAF."
|
|
||||||
|
|
||||||
# Call LocalAI
|
|
||||||
request=$(cat <<AIEOF
|
|
||||||
{"model":"$localai_model","messages":[{"role":"system","content":"You are SecuBox Threat Analyst AI."},{"role":"user","content":"$(echo "$prompt" | sed 's/"/\\"/g' | tr '\n' ' ')"}],"max_tokens":1024,"temperature":0.5}
|
|
||||||
AIEOF
|
|
||||||
)
|
|
||||||
|
|
||||||
response=$(echo "$request" | wget -q -O - --post-data=- \
|
|
||||||
--header="Content-Type: application/json" \
|
|
||||||
"${localai_url}/v1/chat/completions" 2>/dev/null)
|
|
||||||
|
|
||||||
if [ -n "$response" ]; then
|
|
||||||
content=$(echo "$response" | jsonfilter -e '@.choices[0].message.content' 2>/dev/null)
|
|
||||||
# Escape for JSON
|
|
||||||
content=$(echo "$content" | sed 's/\\/\\\\/g; s/"/\\"/g' | tr '\n' ' ')
|
|
||||||
printf '{"response":"%s"}' "$content"
|
|
||||||
else
|
|
||||||
echo '{"error":"AI query failed"}'
|
|
||||||
fi
|
|
||||||
;;
|
;;
|
||||||
|
|
||||||
analyze)
|
analyze)
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user