477 lines
22 KiB
Python
477 lines
22 KiB
Python
"""Motore interno: carica un modello MLX (mlx-lm) dentro l'app, senza server esterni.
|
|
|
|
Modelli da mlx-community (Hugging Face), scaricati nella cache locale al primo uso. Gli strumenti (memoria,
|
|
plugin, ricerca web) passano dal template di chat del modello nel formato Hermes/Qwen:
|
|
<tool_call>{"name": "...", "arguments": {...}}</tool_call>
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import json
|
|
import os
|
|
import re
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from avatar.memory_tools import memory_prompt, openai_tools, parse_args, run_tool
|
|
from avatar.plugins import registry
|
|
from avatar.websearch import brave_search
|
|
from .base import Emit, History, compact_history, persona_text, summary_block, today_label, user_block
|
|
from .openai_compat import SEARCH_TOOL, Aborted
|
|
|
|
MAX_ROUNDS = 8
|
|
MAX_HISTORY = 30
|
|
MAX_TOKENS = 3000
|
|
SAMPLING = {"temp": 0.6, "top_p": 0.8, "top_k": 20} # poco sotto i valori consigliati da Qwen: più costante nell'uso degli strumenti
|
|
DEBUG = bool(os.environ.get("MLX_DEBUG")) # stampa il testo grezzo generato (tag compresi)
|
|
DEFAULT_MODEL = "mlx-community/Qwen3-30B-A3B-Instruct-2507-4bit"
|
|
HUB = Path.home() / ".cache" / "huggingface" / "hub"
|
|
_EXCLUDE = ("whisper", "tts", "embed", "rerank", "clip", "vl-", "-vl", "vision", "diffusion", "parakeet")
|
|
|
|
_lock = threading.Lock()
|
|
_loaded: dict = {"name": None, "model": None, "tokenizer": None, "cache": None, "tokens": [], "snap": None}
|
|
|
|
|
|
def cached_models() -> list[str]:
|
|
"""Modelli MLX già scaricati nella cache di Hugging Face (esclusi voce/visione)."""
|
|
out = []
|
|
if HUB.exists():
|
|
for d in sorted(HUB.glob("models--*")):
|
|
name = d.name[len("models--"):].replace("--", "/")
|
|
low = name.lower()
|
|
if "mlx" in low and not any(x in low for x in _EXCLUDE) and (d / "snapshots").exists():
|
|
out.append(name)
|
|
return out
|
|
|
|
|
|
def load_model(name: str):
|
|
"""Carica (o riusa) il modello. Un solo modello in memoria per volta."""
|
|
with _lock:
|
|
if _loaded["name"] == name and _loaded["model"] is not None:
|
|
return _loaded["model"], _loaded["tokenizer"]
|
|
unload()
|
|
try:
|
|
import optiq # noqa: F401 registra architetture extra (es. spark2_5) in mlx-lm, se installato
|
|
except Exception:
|
|
pass
|
|
from mlx_lm import load
|
|
model, tokenizer = load(name)
|
|
_loaded.update(name=name, model=model, tokenizer=tokenizer, cache=None, tokens=[], snap=None)
|
|
return model, tokenizer
|
|
|
|
|
|
def unload() -> None:
|
|
if _loaded["model"] is not None:
|
|
_loaded.update(name=None, model=None, tokenizer=None, cache=None, tokens=[], snap=None)
|
|
try:
|
|
import gc
|
|
import mlx.core as mx
|
|
gc.collect()
|
|
mx.clear_cache()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class TagFilter:
|
|
"""Separa il testo visibile dai blocchi racchiusi tra `start` e `end` (pensieri o chiamate di strumenti)."""
|
|
|
|
def __init__(self, start: str, end: str, keep: bool) -> None:
|
|
self.start, self.end, self.keep = start, end, keep
|
|
self.inside = False
|
|
self.strip_next = False # dopo un blocco chiuso, via gli spazi iniziali del testo che segue
|
|
self.carry = ""
|
|
self.buf = ""
|
|
self.blocks: list[str] = []
|
|
|
|
def _emit(self, piece: str) -> str:
|
|
if self.strip_next:
|
|
piece = piece.lstrip().lstrip("-").lstrip()
|
|
if piece:
|
|
self.strip_next = False
|
|
return piece
|
|
|
|
def _partial(self, text: str, tag: str) -> int:
|
|
"""Lunghezza del pezzo finale di `text` che potrebbe essere l'inizio di `tag`."""
|
|
for n in range(min(len(tag) - 1, len(text)), 0, -1):
|
|
if text.endswith(tag[:n]):
|
|
return n
|
|
return 0
|
|
|
|
def push(self, delta: str) -> str:
|
|
text, self.carry, out = self.carry + delta, "", ""
|
|
while text:
|
|
if self.inside:
|
|
e = text.find(self.end)
|
|
if e < 0:
|
|
n = self._partial(text, self.end)
|
|
self.buf += text[: len(text) - n] if n else text
|
|
self.carry = text[len(text) - n:] if n else ""
|
|
return out
|
|
self.buf += text[:e]
|
|
if self.keep:
|
|
self.blocks.append(self.buf.strip())
|
|
self.buf, self.inside, self.strip_next = "", False, True
|
|
text = text[e + len(self.end):]
|
|
else:
|
|
st = text.find(self.start)
|
|
if st < 0:
|
|
n = self._partial(text, self.start)
|
|
out += self._emit(text[: len(text) - n] if n else text)
|
|
self.carry = text[len(text) - n:] if n else ""
|
|
return out
|
|
out += self._emit(text[:st])
|
|
text = text[st + len(self.start):]
|
|
self.inside = True
|
|
return out
|
|
|
|
def flush(self) -> str:
|
|
if self.inside:
|
|
if self.keep and (self.buf + self.carry).strip():
|
|
self.blocks.append((self.buf + self.carry).strip())
|
|
self.buf = self.carry = ""
|
|
self.inside = False
|
|
return ""
|
|
c, self.carry = self.carry, ""
|
|
return self._emit(c)
|
|
|
|
|
|
# Formati di chiamata: Hermes/Qwen (JSON in <tool_call>), Gemma 4 (call:nome{chiave:<|"|>testo<|"|>,...}),
|
|
# argkey/Spark-GLM (<tool_call>nome<arg_key>k</arg_key><arg_value>v</arg_value>...</tool_call>).
|
|
# "thinking": alcuni modelli (Spark) usano gli strumenti in modo affidabile solo col ragionamento acceso.
|
|
# "raw_args": il template riproduce gli argomenti come stringa grezza (prompt identico a ciò che il modello ha scritto).
|
|
FORMATS = {
|
|
"hermes": {"think": ("<think>", "</think>"), "call": ("<tool_call>", "</tool_call>"), "thinking": False, "raw_args": False},
|
|
"gemma4": {"think": ("<|channel>", "<channel|>"), "call": ("<|tool_call>", "<tool_call|>"), "thinking": False, "raw_args": True},
|
|
"argkey": {"think": ("<think>", "</think>"), "call": ("<tool_call>", "</tool_call>"), "thinking": True, "raw_args": False},
|
|
}
|
|
GEMMA_ESC = '<|"|>'
|
|
|
|
|
|
def detect_format(tokenizer) -> str:
|
|
try:
|
|
vocab = tokenizer.get_vocab()
|
|
except Exception:
|
|
vocab = {}
|
|
if "<|tool_call>" in vocab:
|
|
return "gemma4"
|
|
if "<arg_key>" in str(getattr(tokenizer, "chat_template", "") or ""):
|
|
return "argkey"
|
|
return "hermes"
|
|
|
|
|
|
def _gemma_to_json(body: str) -> str:
|
|
"""Converte gli argomenti nel formato Gemma in JSON: stringhe tra <|"|>, chiavi nude."""
|
|
parts = body.split(GEMMA_ESC)
|
|
out = []
|
|
for i, part in enumerate(parts):
|
|
if i % 2 == 1:
|
|
out.append(json.dumps(part, ensure_ascii=False))
|
|
else:
|
|
out.append(re.sub(r'([{,\[]\s*)([A-Za-z_][\w.-]*)\s*:', r'\1"\2":', part))
|
|
return "".join(out)
|
|
|
|
|
|
def _parse_call(raw: str, fmt: str) -> tuple[str, dict, str] | None:
|
|
"""(nome, argomenti, argomenti grezzi come li ha scritti il modello)."""
|
|
raw = raw.strip()
|
|
if fmt == "gemma4" or raw.startswith("call:"):
|
|
m = re.match(r"call:\s*([\w.-]+)\s*(\{.*\})?\s*$", raw, re.S)
|
|
if not m:
|
|
return None
|
|
body = m.group(2) or "{}"
|
|
try:
|
|
args = json.loads(_gemma_to_json(body))
|
|
except Exception:
|
|
return None
|
|
return m.group(1), dict(args or {}), body[1:-1]
|
|
if fmt == "argkey" or "<arg_key>" in raw:
|
|
m = re.match(r"\s*([\w.-]+)\s*(.*)$", raw, re.S)
|
|
if not m:
|
|
return None
|
|
args = {}
|
|
for k, v in re.findall(r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>", m.group(2), re.S):
|
|
v = v.strip()
|
|
try:
|
|
args[k.strip()] = json.loads(v)
|
|
except Exception:
|
|
args[k.strip()] = v
|
|
return m.group(1), args, json.dumps(args, ensure_ascii=False)
|
|
try:
|
|
d = json.loads(raw)
|
|
except Exception:
|
|
m = re.search(r"\{.*\}", raw, re.S)
|
|
if not m:
|
|
return None
|
|
try:
|
|
d = json.loads(m.group(0))
|
|
except Exception:
|
|
return None
|
|
if not isinstance(d, dict) or not d.get("name"):
|
|
return None
|
|
args = d.get("arguments", d.get("parameters", {}))
|
|
if isinstance(args, str):
|
|
args = parse_args(args)
|
|
return str(d["name"]), dict(args or {}), json.dumps(args or {}, ensure_ascii=False)
|
|
|
|
|
|
class MLXEngine:
|
|
name = "mlx"
|
|
|
|
def __init__(self, model_name: str, search_api_key: str, assistant_name: str, user_name: str, thinking: str = "auto") -> None:
|
|
self.model_name = model_name or DEFAULT_MODEL
|
|
self.thinking = thinking # auto | on | off
|
|
self.search_api_key = search_api_key
|
|
self.assistant_name, self.user_name = assistant_name, user_name
|
|
self.history = History("mlx")
|
|
self.model = self.tokenizer = None
|
|
self.fmt = "hermes"
|
|
self.load_error: str | None = None
|
|
threading.Thread(target=self._preload, daemon=True).start()
|
|
|
|
def _preload(self) -> None:
|
|
try:
|
|
self.model, self.tokenizer = load_model(self.model_name)
|
|
self.fmt = detect_format(self.tokenizer)
|
|
self._warm()
|
|
except Exception as err:
|
|
self.load_error = str(err)
|
|
print(f"[mlx] caricamento fallito: {err}")
|
|
|
|
def _warm(self) -> None:
|
|
"""Pre-calcola la parte stabile del prompt (persona, memoria, strumenti, cronologia) e la fotografa."""
|
|
if self.model is None:
|
|
return
|
|
with _lock:
|
|
if _loaded["snap"] is not None:
|
|
return
|
|
saved = list(self.history.messages)
|
|
try:
|
|
renders = []
|
|
for dummy in ("a", "b"):
|
|
self.history.messages = saved + [{"role": "user", "content": dummy}]
|
|
renders.append(list(self.tokenizer.encode(self._prompt())))
|
|
finally:
|
|
self.history.messages = saved
|
|
n = 0
|
|
while n < min(len(renders[0]), len(renders[1])) and renders[0][n] == renders[1][n]:
|
|
n += 1
|
|
self._prefill(renders[0][:n], n)
|
|
|
|
def _prefill(self, tokens: list[int], cut: int):
|
|
"""Cache che contiene esattamente `tokens`, riusando la fotografia precedente se ne è un'estensione.
|
|
La nuova fotografia viene scattata a `cut` (fine dell'ultimo messaggio): il turno successivo riparte da lì."""
|
|
from mlx_lm.models.cache import make_prompt_cache, trim_prompt_cache, can_trim_prompt_cache
|
|
snap, cache, fed = _loaded["snap"], None, 0
|
|
if snap and len(tokens) >= len(snap["tokens"]) and tokens[: len(snap["tokens"])] == snap["tokens"]:
|
|
cache, fed = copy.deepcopy(snap["cache"]), len(snap["tokens"])
|
|
elif _loaded["cache"] is not None and _loaded["tokens"]:
|
|
# cache viva del turno precedente: ritaglia al prefisso comune (solo se il modello lo consente)
|
|
old = _loaded["tokens"]
|
|
n = 0
|
|
while n < min(len(old), len(tokens)) and old[n] == tokens[n]:
|
|
n += 1
|
|
if n > 0 and can_trim_prompt_cache(_loaded["cache"]):
|
|
trim_prompt_cache(_loaded["cache"], len(old) - n)
|
|
cache, fed = _loaded["cache"], n
|
|
_loaded.update(cache=None, tokens=[])
|
|
if cache is None:
|
|
cache = make_prompt_cache(self.model)
|
|
if cut > fed:
|
|
self._feed(cache, tokens[fed:cut])
|
|
_loaded["snap"] = {"tokens": list(tokens[:cut]), "cache": copy.deepcopy(cache)}
|
|
fed = cut
|
|
if len(tokens) > fed:
|
|
self._feed(cache, tokens[fed:])
|
|
return cache
|
|
|
|
def _feed(self, cache, tokens: list[int], step: int = 2048) -> None:
|
|
"""Inserisce esattamente `tokens` nella cache, senza campionare (stream_generate aggiungerebbe un token in più)."""
|
|
import mlx.core as mx
|
|
for i in range(0, len(tokens), step):
|
|
self.model(mx.array(tokens[i:i + step])[None], cache=cache)
|
|
mx.eval([c.state for c in cache])
|
|
mx.clear_cache()
|
|
|
|
def close(self) -> None:
|
|
self.model = self.tokenizer = None
|
|
unload()
|
|
|
|
def reset(self) -> None:
|
|
self.history.clear()
|
|
|
|
def _ensure_loaded(self, emit: Emit) -> None:
|
|
if self.model is not None:
|
|
return
|
|
emit({"type": "status", "status": "working", "detail": "carico il modello"})
|
|
self.load_error = None
|
|
self.model, self.tokenizer = load_model(self.model_name)
|
|
self.fmt = detect_format(self.tokenizer)
|
|
|
|
def _thinking(self) -> bool:
|
|
if self.thinking == "on":
|
|
return True
|
|
if self.thinking == "off":
|
|
return False
|
|
return FORMATS[self.fmt]["thinking"]
|
|
|
|
def _system(self) -> str:
|
|
note = ("\n\nRegole sugli strumenti (obbligatorie):\n"
|
|
"- Quando l'utente ti chiede di ricordare qualcosa, o ti dice un fatto importante su di sé, DEVI chiamare salva_memoria prima di rispondere. Non dire mai di aver salvato senza averlo chiamato davvero.\n"
|
|
"- Per cancellare una memoria chiama dimentica_memoria; per cercarne una non presente nel prompt chiama cerca_memoria.\n"
|
|
"- Per azioni sul Mac (per esempio il calendario) usa gli strumenti dedicati. Se uno risponde con [CONFIRMATION_PENDING], chiedi all'utente di confermare sul pannello e non dire che è fatto.\n"
|
|
"- Non inventare mai dati reali: meteo, calendario, mail, messaggi, file, ora, server e memoria li ottieni SOLO chiamando lo strumento corrispondente, anche a metà conversazione. Rispondere senza averlo chiamato è un errore grave.\n"
|
|
"- Chiama gli strumenti solo nel formato previsto dal tuo template e mai descrivendoli a parole.")
|
|
note += ("\n- Per informazioni aggiornate chiama cerca_web e rispondi in base ai risultati." if self.search_api_key
|
|
else "\n- Non hai accesso al web: se ti chiedono informazioni aggiornate, dillo chiaramente.")
|
|
return f"{persona_text(self.assistant_name)}\n\n{user_block(self.user_name, memory_prompt())}{summary_block(self.history)}{note}"
|
|
|
|
def _recent(self) -> list:
|
|
msgs = self.history.messages
|
|
if len(msgs) <= MAX_HISTORY:
|
|
return msgs
|
|
start = len(msgs) - MAX_HISTORY
|
|
while start < len(msgs) and msgs[start].get("role") != "user":
|
|
start += 1
|
|
return msgs[start:]
|
|
|
|
def _tools(self) -> list:
|
|
return openai_tools() + registry.openai_tools() + ([SEARCH_TOOL] if self.search_api_key else [])
|
|
|
|
def _prompt(self, generation: bool = True) -> str:
|
|
"""Prompt completo; con generation=False si ferma alla fine dell'ultimo messaggio (parte stabile)."""
|
|
messages = [{"role": "system", "content": self._system()}, *self._recent()]
|
|
try:
|
|
return self.tokenizer.apply_chat_template(messages, tools=self._tools(), add_generation_prompt=generation, tokenize=False, enable_thinking=self._thinking())
|
|
except TypeError:
|
|
# Template senza supporto strumenti: li descrive nel sistema.
|
|
messages[0]["content"] += "\n\nStrumenti disponibili (JSON):\n" + json.dumps(self._tools(), ensure_ascii=False) + \
|
|
"\nPer usarne uno scrivi solo: <tool_call>{\"name\": \"…\", \"arguments\": {…}}</tool_call>"
|
|
return self.tokenizer.apply_chat_template(messages, add_generation_prompt=generation, tokenize=False, enable_thinking=self._thinking())
|
|
|
|
def _prompt_tokens(self) -> tuple[list[int], int, str]:
|
|
"""Token del prompt, lunghezza della parte stabile (comune al turno successivo) e testo del prompt."""
|
|
prompt = self._prompt()
|
|
tokens = list(self.tokenizer.encode(prompt))
|
|
stable = list(self.tokenizer.encode(self._prompt(generation=False)))
|
|
cut = 0
|
|
while cut < min(len(tokens), len(stable)) and tokens[cut] == stable[cut]:
|
|
cut += 1
|
|
return tokens, cut, prompt
|
|
|
|
def _run_tool(self, name: str, args: dict, emit: Emit, sources: list) -> str:
|
|
if name == "cerca_web":
|
|
emit({"type": "status", "status": "searching"})
|
|
try:
|
|
hits = brave_search(str(args.get("query", "")), self.search_api_key)
|
|
except Exception as err:
|
|
return f"Errore nella ricerca: {err}"
|
|
for h in hits[:5]:
|
|
if all(s["url"] != h["url"] for s in sources):
|
|
sources.append({"title": h["title"], "url": h["url"]})
|
|
return "\n".join(f"{i+1}. {h['title']}\n {h['url']}\n {h['description']}" for i, h in enumerate(hits)) or "Nessun risultato."
|
|
if registry.has(name):
|
|
emit({"type": "status", "status": "working", "detail": name})
|
|
return registry.run(name, args)
|
|
emit({"type": "status", "status": "memory"})
|
|
res, ev = run_tool(name, args)
|
|
if ev:
|
|
emit(ev)
|
|
return res
|
|
|
|
def _generate(self, emit: Emit, abort: threading.Event, announced: list) -> tuple[str, list[str], bool]:
|
|
from mlx_lm import stream_generate
|
|
from mlx_lm.sample_utils import make_sampler
|
|
f = FORMATS[self.fmt]
|
|
think, tools, text, truncated = TagFilter(*f["think"], keep=False), TagFilter(*f["call"], keep=True), "", False
|
|
sampler = make_sampler(**SAMPLING)
|
|
# Cache del prompt: la parte già vista (fotografia del turno precedente) non viene ricalcolata.
|
|
tokens, cut, prompt = self._prompt_tokens()
|
|
think.inside = prompt.rstrip().endswith(f["think"][0]) # il template ha già aperto il blocco di pensiero
|
|
cache = self._prefill(tokens[:-1], min(cut, len(tokens) - 1))
|
|
generated: list[int] = []
|
|
try:
|
|
for r in stream_generate(self.model, self.tokenizer, tokens[-1:], max_tokens=MAX_TOKENS, sampler=sampler, prompt_cache=cache):
|
|
if abort.is_set():
|
|
raise Aborted()
|
|
generated.append(r.token)
|
|
if DEBUG:
|
|
print(r.text, end="", flush=True)
|
|
visible = tools.push(think.push(r.text))
|
|
if visible:
|
|
if not announced:
|
|
announced.append(True)
|
|
emit({"type": "status", "status": "responding"})
|
|
text += visible
|
|
emit({"type": "text", "delta": visible})
|
|
if r.finish_reason == "length":
|
|
truncated = True
|
|
except Aborted:
|
|
raise
|
|
else:
|
|
# generate_step inserisce nella cache ogni token emesso (anche l'ultimo, anche quando si ferma per limite).
|
|
_loaded.update(cache=cache, tokens=tokens + generated)
|
|
tail = tools.flush()
|
|
if tail:
|
|
text += tail
|
|
emit({"type": "text", "delta": tail})
|
|
return text, tools.blocks, truncated
|
|
|
|
def send(self, user_text: str, emit: Emit, abort: threading.Event, **_) -> None:
|
|
self.history.messages.append({"role": "user", "content": f"{user_text}\n\n(Adesso è {today_label()}.)"})
|
|
full, sources, announced = "", [], []
|
|
try:
|
|
self._ensure_loaded(emit)
|
|
emit({"type": "status", "status": "thinking"})
|
|
for rnd in range(MAX_ROUNDS):
|
|
with _lock:
|
|
round_text, raw_calls, truncated = self._generate(emit, abort, announced)
|
|
full += round_text
|
|
calls = [c for c in (_parse_call(r, self.fmt) for r in raw_calls) if c]
|
|
if not calls:
|
|
self.history.messages.append({"role": "assistant", "content": round_text.strip()})
|
|
break
|
|
self.history.messages.append({
|
|
"role": "assistant", "content": round_text.strip(),
|
|
"tool_calls": [{"id": f"call_{rnd}_{i}", "type": "function", "function": {"name": n, "arguments": raw if FORMATS[self.fmt]["raw_args"] else a}} for i, (n, a, raw) in enumerate(calls)],
|
|
})
|
|
for i, (n, a, raw) in enumerate(calls):
|
|
result = self._run_tool(n, a, emit, sources)
|
|
self.history.messages.append({"role": "tool", "tool_call_id": f"call_{rnd}_{i}", "name": n, "content": str(result)})
|
|
emit({"type": "status", "status": "thinking"})
|
|
if truncated:
|
|
break
|
|
full = full.strip() or "(nessuna risposta)"
|
|
self.history.save()
|
|
emit({"type": "done", "text": full, "sources": sources})
|
|
with _lock:
|
|
if compact_history(self.history, self.summarize):
|
|
_loaded["snap"] = None # il prompt cambia: la fotografia non vale più
|
|
except Aborted:
|
|
self._rollback()
|
|
emit({"type": "error", "message": "Risposta interrotta.", "aborted": True})
|
|
except Exception as err:
|
|
self._rollback()
|
|
emit({"type": "error", "message": f"Modello interno: {err}"})
|
|
|
|
def summarize(self, prompt: str) -> str:
|
|
"""Genera un testo breve col modello caricato, senza strumenti né cache condivisa (usato per compattare la cronologia)."""
|
|
from mlx_lm import stream_generate
|
|
from mlx_lm.sample_utils import make_sampler
|
|
from mlx_lm.models.cache import make_prompt_cache
|
|
msgs = [{"role": "system", "content": "Sei un assistente che riassume conversazioni in italiano, in modo fedele e compatto."},
|
|
{"role": "user", "content": prompt}]
|
|
text = self.tokenizer.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False)
|
|
f = FORMATS[self.fmt]
|
|
think = TagFilter(*f["think"], keep=False)
|
|
think.inside = text.rstrip().endswith(f["think"][0])
|
|
out = ""
|
|
for r in stream_generate(self.model, self.tokenizer, text, max_tokens=400, sampler=make_sampler(temp=0.3, top_p=0.9), prompt_cache=make_prompt_cache(self.model)):
|
|
out += think.push(r.text)
|
|
return (out + think.flush()).strip()
|
|
|
|
def _rollback(self) -> None:
|
|
msgs = self.history.messages
|
|
if msgs and msgs[-1].get("role") == "user":
|
|
msgs.pop()
|
|
self.history.save()
|