path: root/assistant.py

#!/usr/bin/env python3
# desktop-assistant: push-to-talk local voice assistant.
# 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.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
"""Mic button on: record. Mic button off: transcribe (whisper-server),
ask the loaded llama-server model, speak the reply (Kokoro).

Run: python3 assistant.py        (--test runs the self-check)
"""

import ctypes
import glob
import json
import os
import queue
import random
import re
import signal
import subprocess
import sys
import threading
import time
import tomllib
import unicodedata
import urllib.request
import wave

CONFIG_DIR = os.path.join(os.environ.get("XDG_CONFIG_HOME") or os.path.expanduser("~/.config"),
                          "desktop-assistant")
RAW = os.path.join(os.environ.get("XDG_RUNTIME_DIR", "/tmp"), "desktop-assistant.raw")
WAV = RAW[:-3] + "wav"
DEFAULTS = {
    "card": "Microphone",                 # ALSA card id of the USB mic
    "llm": "http://localhost:8181/v1",
    "system": ("You are a voice assistant. Your replies are spoken aloud, so answer "
               "briefly in plain sentences: no markdown, no lists, no emoji. "
               "Reply in the language the user speaks, Italian or English. "
               "Facts about the user's computer, devices, accounts, mail, calendar "
               "or money come only from a tool: call one, and if none can tell, "
               "say you do not know. Never guess them."),
    "tool_rounds": 3,
    "whisper": {"model": "/data/voice-models/whisper/ggml-large-v3-turbo-q5_0.bin",
                "device": 0, "port": 8182},
    "kokoro": {"model": "/data/voice-models/kokoro/kokoro-v1.0.onnx",
               "voices": "/data/voice-models/kokoro/voices-v1.0.bin"},
    "voice": {"it": {"lang": "it", "blend": {"if_sara": 0.8, "af_bella": 0.2}},
              "en": {"lang": "en-us", "blend": {"af_heart": 1.0}}},
    "env": {},
}
OPEN = {"env", "blend"}                   # free-form tables: replaced whole, not merged


def merge(base, over, where=""):
    # A typo must fail loudly, not fall back to the default in silence.
    out = dict(base)
    for k, v in over.items():
        if k not in base:
            raise SystemExit(f"config: unknown key {where}{k}")
        if isinstance(base[k], dict) and k not in OPEN:
            if not isinstance(v, dict):
                raise SystemExit(f"config: {where}{k} must be a table")
            v = merge(base[k], v, f"{where}{k}.")
        out[k] = v
    return out


def load_config(path):
    try:
        with open(path, "rb") as f:
            return merge(DEFAULTS, tomllib.load(f))
    except FileNotFoundError:
        return DEFAULTS


def load_tools(d):
    tools = {}
    for path in sorted(glob.glob(os.path.join(d, "*.toml"))):
        with open(path, "rb") as f:
            t = tomllib.load(f)
        if not isinstance(t.get("description"), str) or not t.get("command"):
            raise SystemExit(f"{path}: needs description and command")
        tools[os.path.basename(path)[:-5]] = t
    return tools


def schema(tools):
    return [{"type": "function", "function": {
        "name": name, "description": t["description"],
        "parameters": {"type": "object", "properties": {
            p: {"type": s.get("type", "string"), "description": s.get("description", "")}
            for p, s in t.get("params", {}).items()}}}}
        for name, t in tools.items()]


CFG = load_config(os.path.join(CONFIG_DIR, "config.toml"))
WHISPER = f"http://127.0.0.1:{CFG['whisper']['port']}"
KLANG = {l: v["lang"] for l, v in CFG["voice"].items()}
TOOLS_DIR = os.path.join(CONFIG_DIR, "tools")
TOOLS = load_tools(TOOLS_DIR)
TOOL_SCHEMA = schema(TOOLS) or None


def log(*a):
    print(time.strftime("%H:%M:%S"), *a, flush=True)


def mic_on():
    out = subprocess.run(["amixer", "-c", CFG["card"], "cget", "numid=2"],
                         capture_output=True, text=True).stdout
    return out.rstrip().endswith("values=on")


def presenting(d=None):
    # During calls the user turns on presentation mode; the button then
    # belongs to the call, not to us. Presentation is the OR of the manual
    # flag and the game detector's (quickshell status registry: the files are
    # its interface); `statusctl presentation get` reads only the manual half.
    d = d or os.environ.get("XDG_RUNTIME_DIR", "/tmp")

    def flag(name):
        try:
            with open(os.path.join(d, f"status.{name}")) as f:
                return f.read().strip() == "1"
        except OSError:                   # a missing file means off
            return False
    return flag("presentation") or flag("gaming")


def latin(text):
    return all(unicodedata.name(c, "").startswith("LATIN") for c in text if c.isalpha())


def transcribe(lang="auto"):
    # Always send language: whisper-server's own default is en, not auto.
    out = subprocess.run(["curl", "-s", "-F", f"file=@{WAV}", "-F", f"language={lang}",
                          "-F", "response_format=verbose_json", WHISPER + "/inference"],
                         capture_output=True, text=True).stdout
    d = json.loads(out)
    return d.get("language_probabilities", {}), " ".join(d["text"].split())


def stt():
    # large-v3-turbo's detector is reliable for this speaker (small's was
    # not). Only it and en are spoken: if detection lands elsewhere, or the
    # text is non-Latin (a hallucination, both are Latin script), redo the
    # pass forced to the likelier of the two.
    p, text = transcribe()
    lang = max(KLANG, key=lambda l: p.get(l, 0))
    if max(p, key=p.get, default=lang) != lang or not latin(text):
        text = transcribe(lang)[1]
    return lang, text


IT_WORDS = set("il lo la gli le di che è non per una sono ho hai con mi ti ci del della "
               "come questo questa anche ma più se nel alla puoi posso sei".split())
EN_WORDS = set("the is are you i to of and it that for with not this what be have "
               "can your my do there".split())


def lang_of(sentence, fallback):
    # The voice must follow the reply's language, not the user's: the model
    # may answer in the other one. ponytail: stopword count, a real language
    # identifier if this misfires.
    words = re.findall(r"[a-zàèéìòù]+", sentence.lower())
    it = sum(w in IT_WORDS for w in words)
    en = sum(w in EN_WORDS for w in words)
    return "it" if it > en else "en" if en > it else fallback


def split_sentences(buf):
    parts = re.split(r"(?<=[.!?…])\s+|\n+", buf)
    return [p.strip() for p in parts[:-1] if p.strip()], parts[-1]


def run_tool(tools, name, arguments, d=TOOLS_DIR):
    # Whatever goes wrong becomes the tool's answer, so the model can say it
    # could not find out; a tool never takes the assistant down.
    t = tools.get(name)
    if t is None:
        return f"error: no tool named {name}"
    try:
        args = json.loads(arguments or "{}")
    except json.JSONDecodeError:
        args = None
    if not isinstance(args, dict):
        return "error: arguments are not valid JSON"
    vals = {p: "" if args.get(p) is None else str(args[p]) for p in t.get("params", {})}
    # No shell: each {name} fills in inside one argv element.
    argv = [re.sub(r"\{(\w+)\}", lambda m: vals.get(m.group(1), ""), a) for a in t["command"]]
    if "/" not in argv[0] and os.path.exists(os.path.join(d, argv[0])):
        argv[0] = os.path.join(d, argv[0])
    env = {**os.environ, **{k: str(v) for k, v in CFG["env"].items()}}
    try:
        # Own process group: a timeout kills the script's children too (curl),
        # which would otherwise keep the pipes open and block us.
        p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
                             env=env, start_new_session=True)
    except OSError as e:
        return f"error: cannot run {argv[0]}: {e.strerror}"
    try:
        out, err = p.communicate(timeout=t.get("timeout", 10))
    except subprocess.TimeoutExpired:
        os.killpg(p.pid, signal.SIGKILL)
        p.communicate()
        return "error: timed out"
    if p.returncode:
        last = err.strip().splitlines()[-1:]
        return ("error: " + (last[0] if last else f"exit status {p.returncode}"))[:2000]
    return out.strip()[:2000]


def collect_calls(calls, fragments):
    # A streamed tool call arrives in pieces keyed by index: id and name in
    # the first, the JSON arguments spread over the rest.
    for f in fragments:
        c = calls.setdefault(f["index"], {"id": "", "name": "", "arguments": ""})
        fn = f.get("function") or {}
        c["id"] = f.get("id") or c["id"]
        c["name"] = fn.get("name") or c["name"]
        c["arguments"] += fn.get("arguments") or ""


def announcement(tools, calls, lang):
    # Something to hear while a slow tool runs: there is no screen.
    for c in calls.values():
        phrases = tools.get(c["name"], {}).get("announce", {}).get(lang)
        if phrases:
            return random.choice(phrases)
    return None


def llm_model():
    # The router keeps one model resident; asking for another would swap it
    # out under other clients, so use whatever is not unloaded.
    with urllib.request.urlopen(CFG["llm"] + "/models") as r:
        models = json.load(r)["data"]
    return next((m["id"] for m in models if m["status"]["value"] != "unloaded"),
                models[0]["id"])


def ask(messages, calls, tools=None):
    # Yields reply text as it streams; tool-call fragments go into calls.
    body = {"model": llm_model(), "messages": messages, "stream": True,
            "stream_options": {"include_usage": True},
            "chat_template_kwargs": {"enable_thinking": False}}
    if tools:
        body["tools"] = tools
    req = urllib.request.Request(CFG["llm"] + "/chat/completions", json.dumps(body).encode(),
                                 {"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as r:
        for line in r:
            if not line.startswith(b"data: ") or line.strip() == b"data: [DONE]":
                continue
            chunk = json.loads(line[6:])
            if chunk.get("usage"):
                # How full the model's context is: history, tool results and
                # tool definitions all count.
                log(f"context: {chunk['usage']['prompt_tokens']} prompt tokens")
            choices = chunk["choices"]
            if not choices:
                continue
            delta = choices[0]["delta"]
            collect_calls(calls, delta.get("tool_calls") or [])
            if delta.get("content"):
                yield delta["content"]


turn = 0          # bumped by every button press: older replies and audio are stale
player = None     # the aplay speaking now, killed by a new press


def speaker(q):
    global player
    from kokoro_onnx import Kokoro
    k = Kokoro(CFG["kokoro"]["model"], CFG["kokoro"]["voices"])
    # Each voice is a weighted blend of Kokoro style vectors.
    voices = {l: sum(w * k.get_voice_style(n) for n, w in v["blend"].items())
              for l, v in CFG["voice"].items()}
    play, playing = None, None            # aplay and the turn it speaks for
    while True:
        t, item = q.get()
        if t != turn:                     # from a reply the user cut off
            continue
        if play and playing != t:         # never reuse a cut-off stream
            play.kill()
            play = None
        try:
            if item is None:              # end of reply
                if play:
                    play.stdin.close()
                    play.wait()
                    play = None
                continue
            if isinstance(item, float):   # seconds of silence
                if play:
                    play.stdin.write(bytes(int(sr * item) * 2))
                continue
            text, lang = item
            audio, sr = k.create(text, voice=voices[lang], lang=KLANG[lang])
            if t != turn:                 # cut off while synthesising
                continue
            if play is None:
                play = subprocess.Popen(["aplay", "-q", "-r", str(sr), "-f", "S16_LE", "-c", "1"],
                                        stdin=subprocess.PIPE)
                playing, player = t, play
                # A new stream wakes the suspended sink and the first ~0.2 s is
                # lost: lead with silence so it eats that, not the first word.
                play.stdin.write(bytes(int(sr * 0.3) * 2))
            # The pipe buffer lets synthesis of the next sentence overlap playback.
            play.stdin.write((audio.clip(-1, 1) * 32767).astype("<i2").tobytes())
        except BrokenPipeError:           # killed by a new press mid-write
            play = None


def die_with_parent():
    # Runs in the child before exec: the kernel sends SIGTERM when we die,
    # however we die, so no whisper-server or amixer is ever orphaned.
    ctypes.CDLL("libc.so.6").prctl(1, signal.SIGTERM)   # PR_SET_PDEATHSIG


def start_whisper():
    def up():
        try:
            urllib.request.urlopen(WHISPER, timeout=1)
            return True
        except OSError:
            return False
    if up():
        return
    p = subprocess.Popen(["whisper-server", "-m", CFG["whisper"]["model"],
                          "-dev", str(CFG["whisper"]["device"]),
                          "--host", "127.0.0.1", "--port", str(CFG["whisper"]["port"])],
                         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                         preexec_fn=die_with_parent)
    for _ in range(60):
        if up():
            return
        time.sleep(0.5)
    p.kill()
    sys.exit("whisper-server did not come up")


def reply(history, q, mine):
    # Runs on its own thread; returns as soon as a newer press makes it stale.
    # ponytail: a press during this STT can overwrite WAV mid-upload, one
    # garbled question at worst; per-turn file names if that ever shows.
    pcm = open(RAW, "rb").read()
    if len(pcm) < 16000 * 2 * 0.3:               # under 0.3 s: a stray press
        return
    # Trailing silence: whisper drops the last word when audio ends mid-word,
    # which it does when the button is pressed while still talking.
    with wave.open(WAV, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(16000)
        w.writeframes(pcm + bytes(16000 * 2 // 2))
    t = time.perf_counter()
    lang, text = stt()
    if not text or text.startswith("[") or turn != mine:   # [BLANK_AUDIO] and friends
        return
    log(f"you ({lang}, {time.perf_counter() - t:.2f}s): {text}")
    question = {"role": "user", "content": text}
    # ponytail: last 20 messages, summarise if long sessions need memory
    # The model has no clock: without the date it makes one up for tools.
    now = time.strftime("It is %A %d %B %Y, %H:%M.")
    msgs = [{"role": "system", "content": f"{CFG['system']} {now}"}] + history[-20:] + [question]
    t, said, voice, announced = time.perf_counter(), [], lang, False

    def say(s):
        nonlocal voice
        if not said:
            log(f"first sentence after {time.perf_counter() - t:.2f}s")
        said.append(s)
        voice = lang_of(s, voice)
        q.put((mine, (re.sub(r"[*_#`]", "", s), voice)))

    for rnd in range(CFG["tool_rounds"] + 1):
        calls, buf = {}, ""
        # The last round offers no tools, which forces a spoken answer.
        for delta in ask(msgs, calls, TOOL_SCHEMA if rnd < CFG["tool_rounds"] else None):
            if turn != mine:
                log("interrupted")
                return
            buf += delta
            done, buf = split_sentences(buf)
            for s in done:
                say(s)
        if buf.strip():
            say(buf.strip())
        if not calls or turn != mine:
            break
        msgs.append({"role": "assistant", "content": "", "tool_calls": [
            {"id": c["id"], "type": "function",
             "function": {"name": c["name"], "arguments": c["arguments"]}}
            for c in calls.values()]})
        phrase = None if announced else announcement(TOOLS, calls, lang)
        if phrase:
            # The question's language: the reply has no text yet to follow.
            q.put((mine, (phrase, lang)))
            # A beat after it: straight into the answer sounds rushed.
            q.put((mine, 0.6))
            announced = True
        # ponytail: a press does not kill a running tool, its result is dropped
        for c in calls.values():
            result = run_tool(TOOLS, c["name"], c["arguments"])
            log(f"tool {c['name']}({c['arguments']}): {result[:120]!r}")
            msgs.append({"role": "tool", "tool_call_id": c["id"], "content": result})
    if turn != mine:
        log("interrupted")
        return
    q.put((mine, None))
    # Tool calls stay inside the turn: history keeps only what was said. A cut
    # off exchange is dropped whole, so user and assistant still alternate.
    history += [question, {"role": "assistant", "content": " ".join(said)}]
    log("assistant:", " ".join(said))


def main():
    global turn
    # SIGTERM must run the cleanup below too, or the recording is left behind.
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
    start_whisper()
    q = queue.Queue()
    threading.Thread(target=speaker, args=(q,), daemon=True).start()
    history, rec = [], None
    events = subprocess.Popen(["stdbuf", "-oL", "amixer", "-c", CFG["card"], "events"],
                              stdout=subprocess.PIPE, text=True, preexec_fn=die_with_parent)
    prev = mic_on()
    log(f"ready, mic is {'on' if prev else 'off'}")
    try:
        for line in events.stdout:
            if "Mic Capture Switch" not in line:
                continue
            on = mic_on()
            if on == prev:
                continue
            prev = on
            if on and presenting():
                log("presentation mode: ignored")
            elif on:
                # Barge-in: a press silences the reply now and makes it stale.
                turn += 1
                if player:
                    player.kill()
                rec = subprocess.Popen(["arecord", "-q", "-t", "raw", "-f", "S16_LE", "-r", "16000",
                                        "-c", "1", RAW], stderr=subprocess.DEVNULL)
            elif rec:
                rec.send_signal(signal.SIGINT)
                rec.wait()
                rec = None
                threading.Thread(target=reply, args=(history, q, turn), daemon=True).start()
    except KeyboardInterrupt:
        pass
    finally:
        for f in (RAW, WAV):                    # the user's voice: never keep it
            if os.path.exists(f):
                os.remove(f)


def test():
    assert split_sentences("Ciao! Come stai? Io") == (["Ciao!", "Come stai?"], "Io")
    assert split_sentences("Pi is 3.14 today") == ([], "Pi is 3.14 today")
    assert split_sentences("one\n\ntwo") == (["one"], "two")
    assert latin("Perché no? Hi!") and not latin("早, 早!")
    assert lang_of("Non posso fornire il meteo in tempo reale.", "en") == "it"
    assert lang_of("That is a very exciting place to be.", "it") == "en"
    assert lang_of("Tokyo!", "it") == "it"
    c = merge(DEFAULTS, {"card": "X", "whisper": {"port": 9},
                         "voice": {"it": {"blend": {"a": 1.0}}}, "env": {"K": "v"}})
    assert c["card"] == "X" and c["whisper"] == {**DEFAULTS["whisper"], "port": 9}
    assert c["voice"]["it"] == {"lang": "it", "blend": {"a": 1.0}}      # blend replaced whole
    assert c["voice"]["en"] == DEFAULTS["voice"]["en"] and c["env"] == {"K": "v"}
    assert DEFAULTS["card"] == "Microphone"                             # defaults untouched
    for bad in ({"whisper": {"prot": 1}}, {"cards": "X"}, {"whisper": 3}):
        try:
            merge(DEFAULTS, bad)
            assert False, bad
        except SystemExit as e:
            assert "config:" in str(e)
    import tempfile
    with tempfile.TemporaryDirectory() as d:
        with open(os.path.join(d, "echo.sh"), "w") as f:
            f.write('#!/bin/sh\nprintf "%s|" "$@"\n')
        os.chmod(os.path.join(d, "echo.sh"), 0o755)
        with open(os.path.join(d, "now.toml"), "w") as f:
            f.write('description = "d"\ncommand = ["date"]\n'
                    '[params.x]\ntype = "string"\ndescription = "px"\n')
        assert load_tools(d) == {"now": {"description": "d", "command": ["date"],
                                         "params": {"x": {"type": "string", "description": "px"}}}}
        assert schema(load_tools(d)) == [{"type": "function", "function": {
            "name": "now", "description": "d", "parameters": {"type": "object", "properties": {
                "x": {"type": "string", "description": "px"}}}}}]
        tools = {"echo": {"command": ["echo.sh", "{a}", "x{b}y"], "params": {"a": {}, "b": {}}},
                 "slow": {"command": ["sh", "-c", "sleep 5; true"], "timeout": 0.3},
                 "fail": {"command": ["sh", "-c", "echo oops >&2; exit 3"]},
                 "gone": {"command": ["no-such-binary-here"]}}
        # A value is one argv element, never a shell command or a second argument.
        assert run_tool(tools, "echo", '{"a": "x; rm -rf ~", "b": 2}', d) == "x; rm -rf ~|x2y|"
        assert run_tool(tools, "echo", "{}", d) == "|xy|"
        t0 = time.time()
        assert run_tool(tools, "slow", "{}", d) == "error: timed out"
        assert time.time() - t0 < 2                 # the grandchild sleep was killed too
        assert run_tool(tools, "fail", "{}", d) == "error: oops"
        assert run_tool(tools, "gone", "{}", d).startswith("error: cannot run")
        assert run_tool(tools, "nope", "{}", d) == "error: no tool named nope"
        assert run_tool(tools, "echo", "{bad", d) == "error: arguments are not valid JSON"
        assert run_tool(tools, "echo", "[1]", d) == "error: arguments are not valid JSON"
    with tempfile.TemporaryDirectory() as d:
        assert not presenting(d)
        for name, value, want in (("gaming", "1", True), ("gaming", "0", False),
                                  ("presentation", "1\n", True)):
            with open(os.path.join(d, f"status.{name}"), "w") as f:
                f.write(value)
            assert presenting(d) == want, (name, value)
    calls = {}
    collect_calls(calls, [{"index": 0, "id": "c1", "type": "function",
                           "function": {"name": "weather", "arguments": "{"}}])
    for piece in ['"place', '":', '"', "Rome", '"', "}"]:
        collect_calls(calls, [{"index": 0, "function": {"arguments": piece}}])
    collect_calls(calls, [{"index": 1, "id": "c2", "function": {"name": "now", "arguments": "{}"}}])
    assert calls == {0: {"id": "c1", "name": "weather", "arguments": '{"place":"Rome"}'},
                     1: {"id": "c2", "name": "now", "arguments": "{}"}}
    tools = {"now": {}, "search": {"announce": {"it": ["Cerco."], "en": ["Searching."]}}}
    assert announcement(tools, {0: {"name": "now"}}, "it") is None
    assert announcement(tools, {0: {"name": "now"}, 1: {"name": "search"}}, "en") == "Searching."
    assert announcement(tools, {0: {"name": "unknown"}}, "it") is None
    print("ok")


if __name__ == "__main__":
    test() if sys.argv[1:] == ["--test"] else main()