aboutsummaryrefslogtreecommitdiffstats
path: root/desktop/modules/ai/ai-state.sh
blob: 8b10368283d9d221ffa4a1e23650fd47da02b08d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
#!/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, "(sd-cli)" for a CLI run
#   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:--}"
}

# The basename of the model in a "pid args..." line, or "-". First match
# wins: sd-cli's prompt comes later and is free text that may contain "-m".
model_of() {
    local words i
    read -ra words <<<"$1"
    for ((i = 1; i < ${#words[@]} - 1; i++)); do
        case "${words[i]}" in
            --diffusion-model|-m) printf '%s' "${words[i + 1]##*/}"; return ;;
        esac
    done
    printf -
}

sd() {
    local line busy=idle mtime
    # A one-shot sd-cli (fanfictioner runs them) holds VRAM and generates for
    # as long as it lives, so it outranks whatever the web UI is doing.
    if line="$(pgrep -axo sd-cli)"; then
        row engine sd up busy "$(model_of "$line") (sd-cli)"
        return
    fi
    line="$(pgrep -axo sd-server)" || { row engine sd down idle -; return; }
    # 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_of "$line")"
}

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