diff options
Diffstat (limited to 'desktop/modules/ai/ai-state.sh')
| -rwxr-xr-x | desktop/modules/ai/ai-state.sh | 106 |
1 files changed, 106 insertions, 0 deletions
diff --git a/desktop/modules/ai/ai-state.sh b/desktop/modules/ai/ai-state.sh new file mode 100755 index 0000000..b3b9642 --- /dev/null +++ b/desktop/modules/ai/ai-state.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# +# Copyright (C) 2026 Danilo M. <danix@danix.xyz> +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License version 2 as +# published by the Free Software Foundation. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# One sweep of the local AI stack for the drawer's AI page. Always prints +# these six lines, tab separated, and exits 0: +# +# engine<TAB>llama<TAB>up|down<TAB>idle|busy<TAB>model +# engine<TAB>sd<TAB>up|down<TAB>idle|busy<TAB>model +# app<TAB>assistant<TAB>running|stopped<TAB>- +# app<TAB>chat<TAB>running|stopped<TAB>- +# app<TAB>imggen<TAB>running|stopped<TAB>model +# app<TAB>fanfic<TAB>running|stopped<TAB>pid +# +# An empty value is "-". A probe that fails reads as down or stopped, not as +# an error: every row is something that may simply not be running. +# +# The test stubs curl, pgrep, llamachat and imggen on PATH and points +# AI_ASSISTANT and AI_SD_LOG at fixtures; see test-ai-state.sh. + +set -u + +LLAMA="${AI_LLAMA_URL:-http://127.0.0.1:8181}" +SD_LOG="${AI_SD_LOG:-${XDG_CACHE_HOME:-$HOME/.cache}/sd-server.log}" +ASSISTANT="${AI_ASSISTANT:-$HOME/Programming/GIT/desktop-assistant/assistant.sh}" + +row() { local IFS=$'\t'; printf '%s\n' "$*"; } + +llama() { + local models id busy=idle + models="$(curl -sf -m 2 "$LLAMA/models")" || { row engine llama down idle -; return; } + id="$(jq -r 'first(.data[] | select(.status.value == "loaded") | .id) // empty' \ + <<<"$models" 2>/dev/null)" + # autoload=false is not optional: without it this probe loads an unloaded + # model into VRAM, and the model can unload between the two requests. + if [[ -n "$id" ]] && + curl -sf -m 2 -G --data-urlencode "model=$id" -d autoload=false "$LLAMA/slots" | + jq -e 'any(.[]; .is_processing)' >/dev/null 2>&1; then + busy=busy + fi + row engine llama up "$busy" "${id:--}" +} + +sd() { + local line words i model=- busy=idle mtime + line="$(pgrep -axo sd-server)" || { row engine sd down idle -; return; } + read -ra words <<<"$line" + for ((i = 1; i < ${#words[@]} - 1; i++)); do + case "${words[i]}" in + --diffusion-model|-m) model="${words[i + 1]##*/}" ;; + esac + done + # ponytail: sd-server has no progress endpoint, but its log streams + # progress bars while generating, so a log written in the last 5s means + # busy. A step slower than 5s reads as idle. + mtime="$(stat -c %Y "$SD_LOG" 2>/dev/null)" || mtime=0 + (( $(date +%s) - mtime < 5 )) && busy=busy + row engine sd up "$busy" "$model" +} + +apps() { + local out pid + + if "$ASSISTANT" status >/dev/null 2>&1; then + row app assistant running - + else + row app assistant stopped - + fi + + if llamachat --ping >/dev/null 2>&1; then + row app chat running - + else + row app chat stopped - + fi + + # imggen status exits 0 either way: the daemon's JSON when it is up, + # "daemon down" when it is not. + out="$(imggen status 2>/dev/null)" + if [[ "$out" == "{"* ]]; then + row app imggen running "$(jq -r '.model // "-"' <<<"$out" 2>/dev/null || echo -)" + else + row app imggen stopped - + fi + + # It runs through env, so the command line is python3 then the script, + # by path from ~/bin or by ./fanfictioner from the repo. + if pid="$(pgrep -u "$USER" -fo '^[^ ]*python3[^ ]* [^ ]*fanfictioner( |$)')"; then + row app fanfic running "$pid" + else + row app fanfic stopped - + fi +} + +llama +sd +apps +exit 0 |