moth-112m / code /moth /serve.py
darioooooo0o's picture
moth 0.3.1
b7c2357 verified
Raw History Blame Contribute Delete
61.5 kB
"""Chat with moth in the browser: python -m moth.serve -> http://127.0.0.1:4747 (or just: python -m moth)
One conversation at a time. moth's home (~/.moth/home in the release layout, live/home/ here) survives restarts,
so what it writes stays. Its bash tool is moth/shell.py: a tiny pure-Python shell that knows only the commands
moth was trained on (ls, cat, head, wc, echo), stays inside the home and answers exactly like the bash it trained with.
No subprocess, no sandbox needed, so it runs the same on Linux, macOS and Windows.
The KV cache is kept across turns; when it fills up, the oldest turns are dropped and the rest is re-read.
Runs on CPU by default (112M params: fast enough).
Local server hardening: only Host 127.0.0.1/localhost:<port> (DNS rebinding), only local pages over CORS, and a
per-install token (~/.moth/token, made on first run) for everything that reads or drives the chat. The page moth
serves gets the token in its HTML; adapters and displays read the file. GET /api/says is the token-free view
(moth's own lines only) for desktop pets that can't read files.
Several coding-agent sessions at once (any mix of harnesses): each adapter sends a session id and a label (the project
folder), shown on the cards as "claude · moth-repo said". Reactions are collected over a short window and moth reacts
to the most interesting message (Moth.hear_agent), and its history keeps one agent turn per session, so it doesn't
mix them up (Moth._one_turn_per_session). The model's input stays exactly agent_turn(text): no session names.
"""
from __future__ import annotations
import argparse
import hmac
import json
import os
import random
import re
import secrets
import shutil
import socketserver
import sys
import threading
import time
from collections import deque
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlsplit
import torch
TA = os.environ.get("TA_REPO") or os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, TA)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tiny_agent.chat import parse_assistant, render # noqa: E402
from tiny_agent.generate import KVCache, append # noqa: E402
from tiny_agent.tools import BASH_TIMEOUT_S, Workspace, truncate # noqa: E402
from moth.agent import AGENT_TAG, agent_turn, clean_agent_message # noqa: E402
from moth.home import load_pool, sample_home # noqa: E402
from moth.norepeat import PhraseMemory, vocab_classes # noqa: E402
from moth.persona import SYSTEM_STUDENT # noqa: E402
from moth.shell import MARKER, Shell, _windows_unsafe # noqa: E402
HERE = os.path.dirname(os.path.abspath(__file__))
MAX_LEN = 4096 # KV cache size (hard limit)
CTX_BUDGET = 3000 # trained on chats + stitched sessions up to 4k tokens (quality flat to ~4k, tests/ctx_probe.py)
CTX_KEEP = 1800 # when over budget, keep the system prompt + the newest turns that fit in this
REPLY_ROOM = 200 # tokens kept free for moth's next reply
CALL_ROOM = 60 # ... and for each later tool call in the same turn
COMMON = ("i i'm im the a an and or to of it it's its is in on at my me you your we be was are not no yes "
"just so that this but like with for do don't can can't oh . , ! ? ' \" - : ... s t m")
FALLBACK = "oh. my head went all fuzzy for a second. what were we saying?"
SPECIALS = ("<|im_start|>", "<|im_end|>", "<tool_call>", "</tool_call>", "<think>", "</think>",
"<tool_response>", "</tool_response>", "<|endoftext|>", "<|pad|>") # + whatever the tokenizer adds
# Pages allowed to read moth over CORS. Full log, chat, agent, reset: local dev pages only (http(s)://localhost or
# 127.0.0.1, any port). /api/says (moth's own lines, nothing else) also: packaged desktop shells (Hermes Desktop
# = file:// page = Origin "null"; app://). No wildcard anywhere.
LOCAL_ORIGIN = re.compile(r"https?://(localhost|127\.0\.0\.1)(:\d+)?")
PET_ORIGIN = re.compile(r"app://[^/\s]*|file://|null")
TOKEN_HEADER = "X-Moth-Token"
MAX_BODY = 64 * 1024
LOG_SIZE = 500
# Reacting to the coding agent: at most once per REACT_EVERY s, unless it sounds bad or is very long; never while
# the human is mid-vent (chatted in the last QUIET_AFTER_CHAT s): the reaction waits.
REACT_EVERY = 120.0
QUIET_AFTER_CHAT = 60.0
LONG_AGENT_WORDS = 250
URGENT = re.compile(r"\b(fail\w*|errors?|couldn'?t|could not|can'?t|cannot|sorry|apologi\w+|broke\w*|crash\w*)\b",
re.I)
COMPACT_WORDS = 25 # what moth keeps of an agent message in its history once it has reacted
# Several agent sessions at once: when moth may react, it waits BATCH_WINDOW s for more finished turns, then reacts
# to the most interesting one (interest()) and shows the others as skipped cards.
BATCH_WINDOW = 4.0
QUESTION = re.compile(r"\?|\b(should i|shall i|do you want|would you like|want me to|let me know|which one|"
r"which option|your call|okay to|ok to)\b", re.I)
ACTIVE_AFTER = 30 * 60.0 # a session that posted in the last 30 min counts as active (the page's status line)
SOURCE = re.compile(r"[a-z0-9][a-z0-9_-]{0,31}") # "claude-code", "codex", "hermes", "pi", "opencode", "dsh"
SESSION_JUNK = re.compile(r"[^A-Za-z0-9._:-]") # session ids: uuids, "ses_...", "20261018_120000_ab12"
ANSI = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]|\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)?|\x1b[@-_]|[\x00-\x08\x0b-\x1f\x7f-\x9f]")
def strip_ansi(text: str) -> str:
"""Terminal colour/cursor codes and other control characters out; newlines and tabs stay."""
return ANSI.sub("", text)
def clean_session(v) -> str | None:
"""An adapter's session id: a short id charset, at most 80 chars, else None."""
if not isinstance(v, str):
return None
return SESSION_JUNK.sub("", v)[:80] or None
def clean_label(v) -> str | None:
"""An adapter's session label (the project folder's name): one plain line of at most 40 chars, else None."""
if not isinstance(v, str):
return None
return " ".join(strip_ansi(v).split())[:40] or None
def interest(text: str) -> int:
"""Which of several agent turns moth reacts to: failure/apology/error 3 > a question to the user 2 > long 1 >
routine "done" 0. The question is looked for at the end, where agents ask it."""
if URGENT.search(text):
return 3
words = text.split()
if QUESTION.search(" ".join(words[-60:])):
return 2
return 1 if len(words) > LONG_AGENT_WORDS else 0
def default_paths() -> dict:
"""Weights next to code/ (the Hugging Face repo layout) or this project's runs/ and data/."""
root = os.path.dirname(HERE)
repo = os.path.dirname(root)
from moth import weights # stdlib only
if weights.mode() == "hf": # also a download still missing its weights: main() says how to finish it
return {"model": repo, "tokenizer": os.path.join(repo, "tokenizer.json"),
"home": os.path.join(os.path.expanduser("~"), ".moth", "home"),
"pool": os.path.join(HERE, "file_pool.json")}
if weights.mode() == "pip": # pip install: weights in ~/.cache/moth/<repo>/, fetched in main() (not on --help)
cache = weights.cache_dir()
return {"model": cache, "tokenizer": os.path.join(cache, "tokenizer.json"),
"home": os.path.join(os.path.expanduser("~"), ".moth", "home"),
"pool": os.path.join(HERE, "file_pool.json")}
return {"model": os.path.join(root, "runs/sft1/model"),
"tokenizer": os.path.join(root, "data/tiny-agent-112m/tokenizer.json"),
"home": os.path.join(root, "live/home"),
"pool": os.path.join(HERE, "file_pool.json")}
INSTALL_ROOT = os.path.dirname(os.path.dirname(HERE)) # the download folder (moth-112m/), for moth update
def version() -> str:
try:
with open(os.path.join(HERE, "VERSION")) as f:
return f.read().strip()
except OSError:
return "unknown"
def pidfile_path(port: int) -> str:
return os.path.join(os.path.expanduser("~"), ".moth", f"serve-{port}.pid")
def write_pidfile(port: int) -> str | None:
"""~/.moth/serve-<port>.pid, so `python -m moth update` can tell a moth from this install is running."""
path = pidfile_path(port)
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump({"pid": os.getpid(), "port": port, "install_root": INSTALL_ROOT}, f)
return path
except OSError:
return None # moth runs fine without it; update then just can't see this server
def remove_pidfile(path: str | None):
try:
with open(path) as f:
mine = json.load(f).get("pid") == os.getpid() # not one a newer moth on this port wrote
if mine:
os.remove(path) # after the file is closed: Windows can't delete an open file
except (OSError, TypeError, ValueError, AttributeError):
pass
def token_path() -> str:
return os.environ.get("MOTH_TOKEN_FILE") or os.path.join(os.path.expanduser("~"), ".moth", "token")
def load_token(path: str | None = None) -> str:
"""This install's secret, made on first run (owner-only file). Adapters send it as X-Moth-Token."""
path = path or token_path()
try:
with open(path) as f:
tok = f.read().strip()
if tok:
return tok
except FileNotFoundError:
pass
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
tok = secrets.token_urlsafe(24)
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
with os.fdopen(fd, "w") as f:
f.write(tok + "\n")
return tok
def _opt(a, name: str, default):
"""An option eval.py's namespace may not have."""
v = getattr(a, name, None)
return default if v is None else v
def _is_marker_path(path: str) -> bool:
"""Some part of path is moth's home marker as a case-insensitive file system or Windows would see it."""
return any(p.lower().rstrip(". ") == MARKER for p in path.replace("\\", "/").split("/"))
class Home(Workspace):
"""A Workspace rooted at a fixed, persistent directory. bash goes to moth.shell, read/write as in training."""
def __init__(self, root: str, seed: int, pool_file: str | None):
self.files, self.read_paths, self.submitted = {}, set(), None
self.root = os.path.realpath(root)
if not os.path.isdir(self.root):
os.makedirs(self.root)
# "new home" only ever deletes folders with this
with open(os.path.join(self.root, MARKER), "w", encoding="utf-8", newline="\n") as f:
f.write("moth made this folder. its \"new home\" button may delete it.\n")
for rel, content in sample_home(random.Random(seed), load_pool(pool_file)).items():
p = os.path.join(self.root, rel)
os.makedirs(os.path.dirname(p), exist_ok=True)
with open(p, "w", encoding="utf-8", newline="\n") as f:
f.write(content)
self.shell = Shell(self.root)
@property
def owned(self) -> bool:
"""moth created this folder itself (so it may delete it)."""
return os.path.isfile(os.path.join(self.root, MARKER)) and not os.path.islink(os.path.join(self.root, MARKER))
def close(self):
pass
def _resolve(self, path: str) -> str | None:
"""As Workspace's, but moth's marker is off limits: a write must never make a folder look like moth's. On
case-insensitive file systems ".MOTH-HOME" is the marker, and on Windows ".moth-home." / ".moth-home " too,
so any path part that folds to it is refused, on every OS (and Windows-only path tricks, as in moth.shell)."""
if isinstance(path, str) and (_is_marker_path(path) or (os.name == "nt" and _windows_unsafe(path))):
return None
full = super()._resolve(path)
if full is None or _is_marker_path(os.path.relpath(full, self.root).replace(os.sep, "/")):
return None
return full
def tool_bash(self, command: str, timeout: int = BASH_TIMEOUT_S) -> str:
return self.shell.run(str(command))
def tree(self) -> list[dict]:
out = []
for d in sorted(os.listdir(self.root)):
full = os.path.join(self.root, d)
if os.path.isdir(full):
out.append({"dir": d, "files": sorted(f for f in os.listdir(full) if not f.startswith("."))})
return out
class Log:
"""Everything shown in the chat window, polled by displays: the last `size` events, with absolute indexes.
The epoch changes on every restart and reset, so a display knows to start over instead of guessing."""
def __init__(self, size: int = LOG_SIZE):
self.lock = threading.Lock()
self.size = size
self.clear()
def clear(self):
self.events: deque = deque(maxlen=self.size)
self.start = 0 # absolute index of events[0]
self.epoch = secrets.token_hex(4)
def add(self, *evs: dict):
with self.lock:
for ev in evs:
if len(self.events) == self.size:
self.start += 1
self.events.append(ev)
def since(self, n: int, epoch: str | None = None, says_only: bool = False) -> dict:
"""Events from absolute index n (clamped to what is kept). A different epoch: everything kept."""
with self.lock:
end = self.start + len(self.events)
if epoch is not None and epoch != self.epoch:
n = self.start
n = min(max(n, self.start), end)
evs = list(self.events)[n - self.start:]
if says_only:
evs = [{"say": e["say"]} for e in evs if "say" in e]
return {"events": evs, "next": end, "epoch": self.epoch}
def __len__(self):
return self.start + len(self.events)
class Moth:
"""moth's engine on the tiny-agent backend (moth/backends/tiny_agent.py loads the model): the KV cache is kept
across turns and only new text is fed (_feed / _prefill_all / _turn). Other backends: moth.backends.engine's
GenericMoth, the same engine re-rendering the chat each turn (make_engine picks)."""
backend = None # the moth.backends object (None: tests and benchmarks that set tok/model themselves)
manifest = None # its moth.json (moth/backends/manifest.py); None: moth-112m's defaults
def __init__(self, a, manifest=None):
from moth.backends import load_backend, resolve
from moth.backends.manifest import apply_args
m = manifest or getattr(a, "manifest", None) or resolve(a.model, network=False)
if m.backend != "tiny-agent":
raise SystemExit(f"moth: {m.backend} models run on moth.serve.make_engine(a), not Moth(a)")
if not m.is_default:
self.manifest = m
a = apply_args(a, m)
b = self.backend = load_backend(m, a) # tokenizer, weights (+ int8 on the CPU), the session's KV cache
self.tok, self.model, self.quant, self.kv_dtype, self._kv = b.tok, b.model, b.quant, b.kv_dtype, b.kv
self.setup(a)
if a.device != "cpu": # capture the GPU graphs (decode + short-chunk feeds) now, not on the first turns
self.cache = self._kv
self._decoder().warmup(a.max_turn_tokens, self.im_end, a.temperature, a.top_p,
getattr(a, "rep_penalty", 1.0))
self.cache = None
def setup(self, a):
"""Everything but the weights (tests build a Moth with a fake model and call this)."""
self.a = a
# what the manifest changes above the model (no manifest: moth-112m as trained)
man = self.manifest
self.system = man.system_prompt(SYSTEM_STUDENT) if man else SYSTEM_STUDENT
self.character = man.character if man else "moth"
self.card_pre, self.card_post = man.card_parts() if man else (AGENT_TAG + "\n", "")
self.lowercase = man.lowercase if man else True
self.im_end = self.tok.token_to_id("<|im_end|>")
self.ban = [self.tok.token_to_id(t) for t in ("<tool_call>", "<think>")] # for turns where moth must talk
self.lock = threading.RLock()
self.home = Home(a.home, a.seed, a.pool_file)
self.log = Log()
self.budget = getattr(a, "ctx_budget", None) or CTX_BUDGET
self.keep = self.budget * CTX_KEEP // CTX_BUDGET
self.common = {i for w in COMMON.split() for v in (w, " " + w) for i in self.enc(v)}
# moth's last reactions, for the phrase ban (moth/norepeat.py); kept across reset() and "new chat" on purpose,
# wiped by forget(everything=True) (moth forget) and by stopping moth
self.phrases = PhraseMemory(_opt(a, "phrase_memory", 16), _opt(a, "phrase_n", 3))
self.ban_chat = bool(getattr(a, "phrase_ban_chat", False))
self.word_ids, self.boring_ids = vocab_classes(self.tok, self.common)
self._hook = None # this turn's logits_hook (set in _chat)
self._carry = set() # this turn's extra repetition-penalty ids: the last reactions' (set in _chat)
try:
added = [t.content for t in self.tok.get_added_tokens_decoder().values()]
except AttributeError: # older tokenizers
added = []
specials = sorted(set(SPECIALS) | set(added), key=len, reverse=True)
self.specials = re.compile("|".join(re.escape(s) for s in specials))
self.react_every = _opt(a, "react_every", REACT_EVERY)
self.quiet_after = _opt(a, "quiet_after", QUIET_AFTER_CHAT)
self.clock = time.monotonic
self.batch_window = _opt(a, "batch_window", BATCH_WINDOW)
self.last_react = self.last_user = float("-inf")
self.pending = None # the agent message (an item(), see hear_agent) moth will react to next
self.batch, self.batch_until = None, 0.0 # agent messages collected in the open batch window, and its end
self.sessions: dict[str, dict] = {} # session key -> {"source", "session", "label", "last", "skips"}
self.wake = threading.Condition(threading.RLock())
self.reset()
def enc(self, text: str) -> list[int]:
return self.tok.encode(text, add_special_tokens=False).ids
def defuse(self, text: str) -> str:
"""Literal special-token strings in outside text ("<|im_end|>" in a file or a message) would be read as the
real tokens and could end a turn or fake a tool call: "<|im_end|>" -> "< |im_end|>"."""
return self.specials.sub(lambda m: "< " + m.group(0)[1:], text)
def reset(self):
self.messages = [{"role": "system", "content": self.system}]
self.cache = None
self.said = deque(maxlen=3) # token ids of moth's last replies (repetition penalty)
self.protect: set[int] = set() # ids from this turn's user message and tool results: free to copy
def forget(self, everything: bool = False):
"""New chat (the page's button / POST /api/reset): history, log (new epoch), any pending reaction or open
batch, and the list of active agent sessions. Kept on purpose: self.phrases, the token ids of moth's own last
reactions (RAM only, never shown anywhere), so a new chat doesn't bring back the same joke.
everything=True (POST /api/reset {"everything": true}, what `moth forget` sends a running moth): those too,
so nothing moth remembers about this session is left in RAM. Stopping moth wipes all of it as well."""
with self.lock:
self.reset()
self.log.clear()
# what the last turn left behind (rebuilt at the next turn anyway): the phrase-ban hook (reaction n-grams,
# words of user messages) and the reactions' penalty ids; in the decoder (kept: it serves the same reused
# KV cache, and a new one would re-capture its graphs) only its per-turn buffers, cleared in place
self._hook, self._carry = None, set()
dec = getattr(self, "_dec", None)
if dec is not None:
dec.clear_turn_state()
if everything:
self.phrases.clear()
with self.wake:
self.pending, self.batch = None, None
self.sessions.clear()
# --- the model ------------------------------------------------------------------------------------------
def _prefill_all(self):
"""(Re)build the cache from self.messages, dropping the oldest turns if they don't fit."""
while True:
ids = self.enc(render(self.messages, add_generation_prompt=True))
if len(ids) < self.keep or len(self.messages) <= 2:
break
# drop the oldest user turn and everything up to the next user turn
i = next((k for k in range(2, len(self.messages)) if self.messages[k]["role"] == "user"), None)
if i is None: # only the current turn is left: keep it even if it's long
break
self.messages = self.messages[:1] + self.messages[i:]
self.cache = self._new_cache()
logits = self._feed(ids)
if logits is None: # the current turn alone is over budget (tool results are cut to prevent this)
raise RuntimeError(f"turn too long for moth's context ({len(ids)} tokens)")
return logits
def _new_cache(self):
kv = getattr(self, "_kv", None)
if kv is None: # tests build a Moth without weights
return KVCache(self.model, 1, MAX_LEN, self.a.device, dtype=next(self.model.parameters()).dtype)
kv.reset([0]) # same as a fresh cache: everything past the length is masked
return kv
def _feed(self, ids):
"""Append ids to the cache in one pass (tiny_agent.generate.append: the whole prompt in one forward)."""
if int(self.cache.lens_cpu[0]) + len(ids) + REPLY_ROOM > self.budget:
return None # over budget: the caller re-reads a trimmed conversation
x = torch.tensor([ids], device=self.a.device)
return append(self.model, self.cache, x, [len(ids)])
def _pen_ids(self) -> set[int]:
"""Token ids of moth's last replies that the repetition penalty hits this turn (not common words, not ids from
this turn's user message or tool results): fixed for a whole turn."""
if getattr(self.a, "rep_penalty", 1.0) == 1.0 or not (self.said or self._carry):
return set()
return ({i for r in self.said for i in r} | self._carry) - self.common - self.protect
def _sample(self, logits, talk_only=False):
"""One token from logits (1, V): ban (talk_only), repetition penalty, temperature, top-p. On the logits'
device, vectorised (masks, no per-id indexing); the same values and draws as the old per-token code."""
from tiny_agent.generate import sample_top_p
V = logits.size(-1)
pm = torch.zeros(V, dtype=torch.bool)
ids = self._pen_ids()
if ids:
pm[list(ids)] = True
bm = torch.zeros(V, dtype=torch.bool)
if talk_only:
bm[self.ban] = True
return int(sample_top_p(logits, self.a.temperature, self.a.top_p, getattr(self.a, "rep_penalty", 1.0),
pm.to(logits.device) if ids else None, bm.to(logits.device) if talk_only else None))
def _decoder(self):
from tiny_agent.generate import Decoder
dec = getattr(self, "_dec", None)
if dec is None or dec.cache is not self.cache:
graph = False if getattr(self.a, "no_graph", False) else None # None: device graphs on CUDA/XPU
dec = self._dec = Decoder(self.model, self.cache, graph=graph, compile=getattr(self.a, "compile", False))
return dec
def _turn(self, logits, talk_only=False) -> str:
"""Sample moth's turn until <|im_end|> (appended to the cache too) or max_turn_tokens (then <|im_end|> is
appended). CPU: an eager per-token loop, token-for-token the same as the old one for a given seed; CUDA/XPU:
one captured graph per token, sampling on the device, tokens read back every few steps (generate.Decoder)."""
out = self._decoder().generate(
logits, self.a.max_turn_tokens, self.im_end, temperature=self.a.temperature, top_p=self.a.top_p,
pen=getattr(self.a, "rep_penalty", 1.0), pen_ids=self._pen_ids(), ban_ids=self.ban if talk_only else (),
logits_hook=self._hook)
return self.tok.decode(out, skip_special_tokens=False)
def _cut(self, text: str, max_tokens: int) -> str:
"""Shorten one tool result to max_tokens, with the truncation marker moth saw in training."""
ids = self.enc(text)
if len(ids) <= max_tokens:
return text
keep = int(len(text) * max_tokens / len(ids)) - 80 # room for the marker line
text = truncate(text, keep) if keep >= 30 else "[... truncated, too long ...]"
ids = self.enc(text)
return text if len(ids) <= max_tokens else self.tok.decode(ids[:max(1, max_tokens)])
def _fit_results(self, results: list[str], turn_msg: dict, rounds_left: int) -> list[str]:
"""Cut tool results so the current turn fits the context budget even after every older turn is dropped:
a turn that can't be trimmed must never overflow (that used to leave the cache empty and crash).
Each round's results get at most half of what is left, after room for the later rounds' calls."""
start = next((i for i, m in enumerate(self.messages) if m is turn_msg), 1)
fixed = len(self.enc(render(self.messages[:1] + self.messages[start:], add_generation_prompt=True)))
room = self.budget - REPLY_ROOM - fixed - 16 * len(results) - 8 - CALL_ROOM * rounds_left
per = max(8, room // (2 * max(1, len(results))) if rounds_left else room // max(1, len(results)))
return [self._cut(r, per) for r in results]
# --- one turn -------------------------------------------------------------------------------------------
@torch.no_grad()
def chat(self, user: str, talk_only: bool = False) -> list[dict]:
"""Returns the events of moth's turn: {"tool": name, "args":…, "result":…} and {"say": text}.
talk_only: no tools in any round (reactions to the coding agent). Any failure rolls the turn back."""
with self.lock:
snap = (list(self.messages), list(self.said), set(self.protect))
try:
return self._chat(user, talk_only)
except Exception as e: # noqa: BLE001 one bad turn must never wedge the chat
self.messages, self.said, self.protect = snap[0], deque(snap[1], maxlen=3), snap[2]
self.cache = None # appends mutate the cache in place: rebuild it from the restored messages
print(f"moth: turn rolled back ({type(e).__name__}: {e})", file=sys.stderr, flush=True)
return [{"say": FALLBACK}]
def _chat(self, user: str, talk_only: bool = False) -> list[dict]:
user = self.defuse(user)
msg = {"role": "user", "content": user}
self.protect = set(self.enc(" " + user)) | set(self.enc(user))
# the phrase ban: reactions always; chat turns only with --phrase-ban-chat, sparing every user message's words
quote = self.enc(" " + user) if talk_only else self.enc(" ".join(
[m["content"] for m in self.messages if m["role"] == "user"] + [user]))
self._hook = self.phrases.hook(self.word_ids, self.boring_ids, quote) if talk_only or self.ban_chat else None
self._carry = self.phrases.recent() if talk_only else set() # reactions: the soft penalty spans "new chat"
logits = None
if self.cache is not None:
logits = self._feed(self.enc("\n" + render([msg], add_generation_prompt=True)))
self.messages.append(msg)
if logits is None:
logits = self._prefill_all()
events, seen_calls = [], set()
for rnd in range(self.a.max_tool_rounds + 1):
last = talk_only or rnd == self.a.max_tool_rounds # this turn must talk
m = parse_assistant(self._turn(logits, talk_only=last))
if last:
m["tool_calls"] = [] # banned tokens can still be spelled out in plain text
keys = {json.dumps(c, sort_keys=True) for c in m["tool_calls"]}
cut = "<tool_call>" in (m["content"] or "") # a call cut off by the token limit
if (m["tool_calls"] and keys <= seen_calls) or cut:
# the same call again (a loop) or a broken call: rewind this turn and make it talk instead
self.cache = None
logits = self._prefill_all()
m = parse_assistant(self._turn(logits, talk_only=True))
m["tool_calls"], keys = [], set()
m["content"] = m["content"] or ""
if self.lowercase: # moth only speaks lowercase (the odd "I" slips through)
m["content"] = m["content"].lower()
self.messages.append(m)
if not m["tool_calls"]:
events.append({"say": m["content"] or "..."})
self.said.append(self.enc(" " + (m["content"] or "")))
break
if m["content"]:
events.append({"say": m["content"]})
seen_calls |= keys
results = []
for c in m["tool_calls"]:
r = c["error"] if "error" in c else self.home.call(c["name"], c["arguments"])
results.append(r)
events.append({"tool": c.get("name", "?"), "args": c.get("arguments", {}), "result": r})
results = self._fit_results([self.defuse(r) for r in results], msg, self.a.max_tool_rounds - rnd - 1)
for r in results:
self.protect |= set(self.enc(" " + r))
tool_msg = {"role": "tool", "results": results}
self.messages.append(tool_msg)
logits = self._feed(self.enc("\n" + render([tool_msg], add_generation_prompt=True)))
if logits is None:
logits = self._prefill_all()
return events
# --- the human ------------------------------------------------------------------------------------------
def say_to(self, text: str) -> list[dict]:
"""The human talks to moth."""
self.last_user = self.clock()
with self.lock:
self.log.add({"you": text})
evs = self.chat(text)
self.log.add(*evs)
self.last_user = self.clock() # the quiet window starts after moth answered
with self.wake:
self.wake.notify_all()
return evs
# --- the coding agent -----------------------------------------------------------------------------------
@staticmethod
def card(text: str, source: str | None = None, **extra) -> dict:
ev = {"agent": agent_turn(strip_ansi(text))[len(AGENT_TAG):].strip()}
if source:
ev["source"] = source
ev.update({k: v for k, v in extra.items() if v is not None})
return ev
@staticmethod
def item(text: str, source: str | None = None, session: str | None = None, label: str | None = None) -> dict:
"""One agent message as hear_agent, the batch and the reactor carry it. key = the session it came from; an
adapter that sends no session id counts as one session per harness."""
return {"text": text, "source": source, "session": session, "label": label,
"key": f"{source or 'agent'}:{session}" if session else f"src:{source or 'agent'}",
"score": interest(text),
"urgent": bool(URGENT.search(text)) or len(text.split()) > LONG_AGENT_WORDS}
def _seen(self, it: dict):
"""Mark the message's session active (under self.wake). No label: a short session id. A label another active
session already shows (two agents in one folder) gets the end of the session id after it."""
now = self.clock()
for k in [k for k, s in self.sessions.items() if now - s["last"] >= ACTIVE_AFTER]:
del self.sessions[k]
label, sid = it["label"], it["session"]
if not label and sid:
label = sid[:8]
if label and sid and any(k != it["key"] and s["label"] == label for k, s in self.sessions.items()):
label = f"{label} ({sid[-4:]})"
s = self.sessions.setdefault(it["key"], {"source": it["source"], "session": sid, "skips": deque(maxlen=200)})
s["label"] = it["label"] = label
s["last"] = now
def _log_card(self, it: dict, skipped: bool):
self.log.add(self.card(it["text"], it["source"], session=it["session"], label=it["label"],
skipped=True if skipped else None))
if skipped and it["key"] in self.sessions:
self.sessions[it["key"]]["skips"].append(self.clock())
def agents(self) -> dict:
"""The page's status line and session filter: the sessions that posted in the last ACTIVE_AFTER s, and how
many messages moth let pass in that time. Only in the token-guarded /api/log (labels are folder names)."""
with self.wake:
now = self.clock()
act = [s for s in self.sessions.values() if now - s["last"] < ACTIVE_AFTER]
return {"active": len(act), "skipped": sum(now - t < ACTIVE_AFTER for s in act for t in s["skips"]),
"sessions": [{"source": s["source"], "session": s["session"], "label": s["label"]} for s in act]}
def agent_said(self, text: str, source: str | None = None, session: str | None = None,
label: str | None = None) -> list[dict]:
"""The big agent (Claude Code, Hermes, ...) finished a reply: moth sees it and reacts now (no throttle)."""
it = self.item(text, source, session, label)
with self.wake:
self._seen(it)
self._log_card(it, skipped=False)
return self.react(text, item=it)
def react(self, text: str, hold_for_human: bool = False, item: dict | None = None) -> list[dict]:
"""moth's talk-only reaction to an agent message; afterwards only ~25 words of it stay in moth's history,
so agent turns don't push the human's conversation out of context. (The cache still holds the full text
until the next re-read; trimming decides on the short version.) The model sees exactly agent_card(text)
(= agent_turn(text) for moth-112m): session names never go into moth's text. item: the message's session (see _one_turn_per_session)."""
msg = self.agent_card(text)
key = item["key"] if item else None
with self.lock:
if hold_for_human and self.clock() - self.last_user < self.quiet_after: # the human got in first
with self.wake:
if self.pending is None:
self.pending = item or self.item(text)
return []
self.last_react = self.clock()
self._one_turn_per_session(key)
evs = self.chat(msg, talk_only=True)
for e in evs: # the phrase ban's memory (not the fallback line: that one isn't moth's own)
if e.get("say") and e["say"] != FALLBACK:
self.phrases.add(self.enc(e["say"]))
want = self.defuse(msg)
pre, post = len(self.defuse(self.card_pre)), len(self.defuse(self.card_post))
for i in range(len(self.messages) - 1, 0, -1):
m = self.messages[i]
if m["role"] == "user":
if m["content"] == want:
words = want[pre:len(want) - post].split()
short = " ".join(words[:COMPACT_WORDS]) + (" ..." if len(words) > COMPACT_WORDS else "")
self.messages[i] = {**m, "content": f"{self.card_pre}{short}{self.card_post}", "session": key}
break
# the session this reaction answers, for /api/says?source=&session= (the token-free view strips it)
tag = {k: item[k] for k in ("source", "session") if item and item.get(k)}
self.log.add(*({**e, "re": tag} if tag and "say" in e else e for e in evs))
return evs
def agent_card(self, text: str) -> str:
"""The user-role message the model gets for an agent's final message: agent_turn(text) for moth-112m
("[the agent said]\n" + the cleaned text), else the manifest's agent_card with the cleaned text in it."""
return f"{self.card_pre}{clean_agent_message(strip_ansi(text))}{self.card_post}"
def _one_turn_per_session(self, key: str | None):
"""Sessions don't mix: before moth reacts to a message from session `key`, its history keeps only the newest
earlier agent turn of every OTHER session (already cut to COMPACT_WORDS) and none of `key`'s own (the new
message, in full, replaces it). Agent turns carry a "session" field in self.messages, the marker; render()
never shows it to the model. A dropped agent turn goes together with moth's reply to it, so the chat stays
the user/assistant turns moth trained on. When anything was dropped the cache is re-read from the trimmed
history, so the dropped text is really gone."""
seen, drop = {key}, set()
for i in range(len(self.messages) - 1, 0, -1):
m = self.messages[i]
if m["role"] == "user" and "session" in m:
if m["session"] in seen:
drop.add(i)
seen.add(m["session"])
if not drop:
return
kept, dropping = [], False
for i, m in enumerate(self.messages):
if m["role"] == "user":
dropping = i in drop
if not dropping:
kept.append(m)
self.messages, self.cache = kept, None
def hear_agent(self, text: str, source: str | None = None, session: str | None = None,
label: str | None = None) -> str:
"""POST /api/agent: react later on the reactor thread, at most once per react_every s unless it sounds bad
(sorry, failed, error...) or is very long, and only once the human has been quiet for quiet_after s.
Several agent sessions at once: a message moth may react to opens a batch_window s window, and every message
that arrives in it joins. When the window closes, moth takes the most interesting one (interest(): failure >
question > long > routine, the newest on a tie) and the rest are skipped. Batched cards are logged when the
window closes, with their final skipped flag (displays draw each event once), the chosen one last so moth's
reply sits right under it. A message already waiting for the human to go quiet stays unless the batch has
one at least as interesting. batch_window 0 or react_every 0: no window, the card is logged at once and
(react_every 0) every message is queued (the judge tooling's "react to everything right away").
Returns "queued", "batched" (in a window: decided when it closes) or "skipped"."""
it = self.item(text, source, session, label)
with self.wake:
self._seen(it)
if self.batch is not None:
self.batch.append(it)
return "batched"
p = self.pending
if p and p["urgent"] and not it["urgent"]:
status = "skipped"
elif not it["urgent"] and not p and self.clock() - self.last_react < self.react_every:
status = "skipped"
elif self.batch_window > 0 and self.react_every > 0:
self.batch, self.batch_until = [it], self.clock() + self.batch_window
self.wake.notify_all()
return "batched"
else:
status = "queued"
self._log_card(it, skipped=status == "skipped")
if status == "queued": # after the card is logged, so moth's reply always comes after it
self.pending = it
self.wake.notify_all()
return status
def _close_batch(self):
"""The window is over (under self.wake): log the cards, the winner last, and queue it."""
batch, self.batch = self.batch, None
win = max(reversed(batch), key=lambda it: it["score"]) # max keeps the first best: reversed = the newest
if self.pending is not None and self.pending["score"] > win["score"]:
win = None # the message already waiting is more interesting: the whole batch passes
for it in batch:
if it is not win:
self._log_card(it, skipped=True)
if win is not None:
self._log_card(win, skipped=False)
self.pending = win
def due(self) -> tuple[dict | None, float]:
"""(the agent message (an item()) to react to now, or None; seconds to wait before asking again)."""
with self.wake:
if self.batch is not None:
wait = self.batch_until - self.clock()
if wait > 0:
return None, wait
self._close_batch()
if self.pending is None:
return None, 3600.0
wait = self.last_user + self.quiet_after - self.clock()
if wait > 0:
return None, wait
it, self.pending = self.pending, None
return it, 0.0
def start_reactor(self):
def loop():
while True:
with self.wake: # check and wait under one hold of the condition: no lost wakeups
it, wait = self.due()
if it is None:
self.wake.wait(min(wait, 3600.0))
continue
try:
self.react(it["text"], hold_for_human=True, item=it)
except Exception as e: # noqa: BLE001 keep the reactor alive
print(f"moth: reaction failed ({type(e).__name__}: {e})", file=sys.stderr, flush=True)
threading.Thread(target=loop, daemon=True, name="moth-reactor").start()
# --- http ---------------------------------------------------------------------------------------------------
class Server(ThreadingHTTPServer):
"""ThreadingHTTPServer minus the reverse-DNS lookup HTTPServer.server_bind does (socket.getfqdn of the bind
address, used only for CGI's SERVER_NAME). Where reverse DNS is slow that lookup sits between bind() and listen(),
so the port refuses connections for that long (35 s on GitHub's macOS runners) and the pet gives up waiting."""
daemon_threads = True
def server_bind(self):
socketserver.TCPServer.server_bind(self)
self.server_name, self.server_port = str(self.server_address[0]), self.server_address[1]
def make_handler(moth: Moth, token: str, hosts: tuple[str, ...] = ("127.0.0.1", "localhost", "[::1]")):
page = os.path.join(HERE, "static", "index.html")
class H(BaseHTTPRequestHandler):
server_version = "moth"
sys_version = ""
def log_message(self, *a):
pass
def _send(self, code: int, body: bytes, ctype: str, cors: bool = True):
try:
self.send_response(code)
self.send_header("Content-Type", ctype)
self.send_header("Content-Length", str(len(body)))
self.send_header("X-Frame-Options", "DENY")
self.send_header("X-Content-Type-Options", "nosniff")
self.send_header("Cache-Control", "no-store")
self.send_header("Referrer-Policy", "no-referrer")
if cors:
self._cors()
self.end_headers()
self.wfile.write(body)
except (BrokenPipeError, ConnectionResetError):
pass # agent adapters hang up after 1 s on purpose; moth still keeps the message
def _json(self, obj, code=200):
self._send(code, json.dumps(obj).encode(), "application/json")
def _cors(self) -> bool:
"""Reflect the Origin of a local page; desktop-shell origins only for /api/says; others get nothing."""
origin = self.headers.get("Origin", "")
pet = urlsplit(self.path).path == "/api/says"
if not (LOCAL_ORIGIN.fullmatch(origin) or (pet and PET_ORIGIN.fullmatch(origin))):
return False
self.send_header("Access-Control-Allow-Origin", origin)
self.send_header("Vary", "Origin")
return True
def _host_ok(self) -> bool:
"""DNS rebinding: a page on evil.example resolving to 127.0.0.1 still sends Host: evil.example."""
port = self.server.server_address[1]
return (self.headers.get("Host") or "").lower() in {f"{h}:{port}" for h in hosts}
def _authed(self) -> bool:
return hmac.compare_digest(self.headers.get(TOKEN_HEADER, "").encode(), token.encode())
def do_OPTIONS(self): # CORS preflight for a JSON POST / a token header from a local page
if not self._host_ok():
return self._json({"error": "bad host"}, 403)
self.send_response(204)
if self._cors():
self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
self.send_header("Access-Control-Allow-Headers", f"Content-Type, {TOKEN_HEADER}")
self.send_header("Access-Control-Max-Age", "600")
self.send_header("Content-Length", "0")
self.end_headers()
def do_GET(self):
if not self._host_ok():
return self._json({"error": "bad host"}, 403)
url = urlsplit(self.path)
q = parse_qs(url.query)
if url.path == "/api/health": # no token (Host still checked): which install runs here, for moth update
health = {"ok": True, "pid": os.getpid(), "install_root": INSTALL_ROOT, "version": version()}
if getattr(moth, "manifest", None) is not None: # another model (moth.json): who the pet shows
health.update(character=moth.character, backend=moth.manifest.backend)
return self._send(200, json.dumps(health).encode(), "application/json", cors=False)
if url.path == "/":
with open(page, encoding="utf-8") as f:
html = f.read().replace("__MOTH_TOKEN__", token)
return self._send(200, html.encode(), "text/html; charset=utf-8", cors=False) # holds the token
if url.path not in ("/api/says", "/api/log", "/api/home"):
return self._json({"error": "not found"}, 404)
if url.path != "/api/says" and not self._authed():
return self._json({"error": "missing or wrong token (see ~/.moth/token)"}, 401)
if url.path == "/api/home":
return self._json({"tree": moth.home.tree()})
since = (q.get("since") or ["0"])[0]
if not re.fullmatch(r"\d{1,18}", since):
return self._json({"error": "since must be a number"}, 400)
epoch = (q.get("epoch") or [None])[0]
# /api/says?source=&session=: only moth's reactions to that harness / session (an in-harness display's
# own), never chat or another session's; filtered on the tagged log, then stripped to {"say"}
want = {} # cleaned the way /api/agent cleans them; an id that cleans to nothing matches nothing
if q.get("source"):
want["source"] = q["source"][0].lower()
if q.get("session"):
want["session"] = clean_session(q["session"][0]) or "-"
if url.path != "/api/says" or not want:
d = moth.log.since(int(since), epoch, says_only=url.path == "/api/says")
else:
d = moth.log.since(int(since), epoch)
d["events"] = [{"say": e["say"]} for e in d["events"] if "say" in e
and all((e.get("re") or {}).get(k) == v for k, v in want.items())]
if url.path == "/api/log":
d["tree"] = moth.home.tree()
d["agents"] = moth.agents() # session labels are folder names: never in the token-free /api/says
self._json(d)
def _body(self) -> dict | None:
"""The JSON object of a POST, or None after answering 4xx."""
if self.headers.get_content_type() != "application/json":
return self._json({"error": "Content-Type must be application/json"}, 415)
n = self.headers.get("Content-Length", "")
if not (n.isascii() and n.isdigit()):
return self._json({"error": "Content-Length required"}, 411)
if int(n) > MAX_BODY:
return self._json({"error": "too big"}, 413)
try:
body = json.loads(self.rfile.read(int(n)) or b"{}")
except (ValueError, UnicodeDecodeError):
return self._json({"error": "bad JSON"}, 400)
if not isinstance(body, dict):
return self._json({"error": "expected a JSON object"}, 400)
return body
def do_POST(self):
if not self._host_ok():
return self._json({"error": "bad host"}, 403)
origin = self.headers.get("Origin")
if origin is not None and not LOCAL_ORIGIN.fullmatch(origin):
return self._json({"error": "foreign origin"}, 403)
path = urlsplit(self.path).path
if path not in ("/api/chat", "/api/agent", "/api/reset"):
return self._json({"error": "not found"}, 404)
if not self._authed():
return self._json({"error": "missing or wrong token (see ~/.moth/token)"}, 401)
body = self._body()
if body is None:
return
if path == "/api/chat":
text = body.get("text")
text = text.strip()[:500] if isinstance(text, str) else ""
if not text:
return self._json({"error": "empty"}, 400)
self._json({"events": moth.say_to(text), "tree": moth.home.tree()})
elif path == "/api/agent": # from moth/integrations/* : one-way, never replies to the agent
text = body.get("text")
text = text.strip()[:20000] if isinstance(text, str) else ""
if not text:
return self._json({"error": "empty"}, 400)
source = body.get("source")
source = source.lower() if isinstance(source, str) and SOURCE.fullmatch(source.lower()) else None
status = moth.hear_agent(text, source, session=clean_session(body.get("session")),
label=clean_label(body.get("label")))
self._json({"ok": True, "status": status})
else: # /api/reset
with moth.lock:
if body.get("new_home") is True and not moth.home.owned:
return self._json({"error": f"moth didn't make {moth.home.root} itself, so it won't delete "
f"it. move it away by hand for a fresh home.",
"tree": moth.home.tree()}, 409)
if body.get("everything") is True: # moth forget: also the phrase memory
moth.forget(everything=True)
else:
moth.forget()
if body.get("new_home") is True:
shutil.rmtree(moth.home.root, ignore_errors=True)
moth.home = Home(moth.a.home, random.randrange(10**6), moth.a.pool_file)
self._json({"ok": True, "tree": moth.home.tree(), "epoch": moth.log.epoch})
return H
CPU_QUANT = "int8" # --quant auto on the CPU (tiny_agent.generate.quantize_cpu; see release/README.md "Speed")
def resolve_quant(a) -> str:
"""--quant auto -> CPU_QUANT on the CPU, none (unquantized) on a GPU. A namespace without quant: none."""
q = getattr(a, "quant", None) or "none"
if q == "auto":
return CPU_QUANT if getattr(a, "device", "cpu") == "cpu" else "none"
return q
def default_threads() -> int:
"""At most 8. On a 16-core 5950X with free cores, 16 threads are 10-30% faster for both int8 decode and long
prefill (so no separate prefill count), but more threads than free cores (a busy desktop, a CPU quota) cost a lot:
under a 4-CPU quota 8 threads halve int8 decode vs 4. 8 is the safe middle (tests/bench_speed.py --threads)."""
return max(1, min(8, os.cpu_count() or 1))
def build_parser(manifest=None) -> argparse.ArgumentParser:
"""manifest (a moth.json, moth/backends/manifest.py): its sampling defaults and tokenizer become the defaults;
flags given on the command line still win. None: moth-112m's defaults."""
d = default_paths()
s = manifest.sampling if manifest is not None else {}
tok = d["tokenizer"]
if manifest is not None and manifest.backend == "tiny-agent" and manifest.folder:
own = os.path.join(manifest.folder, manifest.tokenizer or "tokenizer.json")
tok = own if (manifest.tokenizer or manifest.repo or os.path.isfile(own)) else tok # repo: comes with it
ap = argparse.ArgumentParser(description="moth's local chat server")
ap.add_argument("--model", default=d["model"],
help="a weights folder, an HF repo id or openai:<base_url>#<model> (default: `moth config model`, "
"else this install's moth; see moth/docs/MODEL_BACKENDS.md)")
ap.add_argument("--tokenizer", default=tok)
ap.add_argument("--home", default=d["home"], help="moth's home folder (kept between runs)")
ap.add_argument("--pool-file", default=d["pool"])
ap.add_argument("--seed", type=int, default=3, help="which home moth starts with")
ap.add_argument("--device", default="cpu")
ap.add_argument("--temperature", type=float, default=s.get("temperature", 0.8))
ap.add_argument("--top-p", type=float, default=s.get("top_p", 0.95))
ap.add_argument("--rep-penalty", type=float, default=s.get("rep_penalty", 1.3),
help="on words from moth's own recent replies")
ap.add_argument("--phrase-memory", type=int, default=16,
help="a new reaction can't repeat a 3-token phrase from moth's last N reactions, and the "
"repetition penalty covers the last 3 of them, also after a new chat (0: off)")
ap.add_argument("--phrase-n", type=int, default=3, help=argparse.SUPPRESS) # the phrase ban's n-gram length
ap.add_argument("--phrase-ban-chat", action="store_true",
help="use the phrase ban in chat replies too (off: chat recall is never affected)")
ap.add_argument("--ctx-budget", type=int, default=None, help="max tokens of chat moth keeps in mind")
ap.add_argument("--max-turn-tokens", type=int, default=s.get("max_tokens", 200))
ap.add_argument("--max-tool-rounds", type=int, default=3)
ap.add_argument("--quant", default="auto", choices=["auto", "none", "int8", "bf16"],
help=f"Linear weights on the CPU: auto (default) = {CPU_QUANT} on the CPU, none on a GPU; "
"int8 (weight-only, per-channel) and bf16 are the fast modes; none = exact fp32 (slowest)")
ap.add_argument("--fast", action="store_const", const=CPU_QUANT, dest="quant",
help=f"same as --quant {CPU_QUANT} (already the CPU default)")
ap.add_argument("--threads", type=int, default=None, help="CPU threads (default: min(8, CPUs))")
ap.add_argument("--no-graph", action="store_true", help="GPU: decode without device graphs (debugging)")
ap.add_argument("--compile", action="store_true",
help="GPU: torch.compile the decode step (~1.4x faster replies on the B70; ~80 s at startup)")
ap.add_argument("--react-every", type=float, default=REACT_EVERY,
help="seconds between reactions to the coding agent (sorry/failed/long messages always count)")
ap.add_argument("--quiet-after", type=float, default=QUIET_AFTER_CHAT,
help="hold reactions until you haven't chatted for this many seconds")
ap.add_argument("--batch-window", type=float, default=BATCH_WINDOW,
help="when moth may react, wait this many seconds for other agent sessions to finish too, then "
"react to the most interesting message (0: no waiting; --react-every 0 never waits either)")
ap.add_argument("--token-file", default=None, help="default ~/.moth/token (or $MOTH_TOKEN_FILE)")
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--port", type=int, default=4747)
ap.add_argument("--exit-with-parent", action="store_true",
help="for launchers (the desktop pet): stop when stdin closes, i.e. when the launcher that holds "
"the other end of the pipe dies")
return ap
def trim_heap():
"""Give the memory freed while loading (the fp32 weights --quant replaces, load buffers) back to the system: glibc
keeps freed heap otherwise (~15 MB less RSS for an all-day server). Not glibc: nothing."""
if sys.platform.startswith("linux"):
try:
import ctypes
ctypes.CDLL(None).malloc_trim(0)
except (OSError, AttributeError):
pass
def _given(argv: list, flag: str) -> bool:
return any(x == flag or x.startswith(flag + "=") for x in argv)
def parse_args(argv=None) -> argparse.Namespace:
"""build_parser()'s options plus the model: --model, else `moth config model`, else this install's moth. A model
with a moth.json brings its own sampling defaults (flags still win). a.manifest: the resolved manifest; a.model:
the weights folder for tiny-agent models (the spec as given for the others)."""
argv = list(sys.argv[1:] if argv is None else argv)
a = build_parser().parse_args(argv)
from moth import backends
spec = a.model if _given(argv, "--model") else (backends.configured() or a.model)
kind, val = backends.parse_spec(spec or "")
if kind == "folder": # a model newer than this code: say how to upgrade, not a traceback from the model code
from moth.backends import compat
compat.require(val)
try:
m = backends.resolve(spec)
except backends.ManifestError as e:
raise SystemExit(f"moth: {e}")
if not m.is_default:
a = build_parser(m).parse_args(argv)
if m.backend == "tiny-agent" and m.folder: # (an HF repo id: the folder it downloads to)
a.model = m.folder
a.manifest = m
return a
def make_engine(a) -> Moth:
"""Moth for tiny-agent models (the KV-cache path), moth.backends.engine.GenericMoth for the others."""
m = getattr(a, "manifest", None)
if m is None or m.backend == "tiny-agent":
return Moth(a)
from moth.backends.engine import GenericMoth
return GenericMoth(a)
def main(argv=None, on_ready=None):
a = parse_args(argv)
from moth import backends, weights
live = {}
if a.exit_with_parent: # first, so a pet killed while the model loads takes this server too (moth/launch.py)
import _thread
from moth.launch import exit_with_parent
exit_with_parent(lambda: live["srv"].shutdown() if "srv" in live else _thread.interrupt_main())
if a.manifest.backend == "tiny-agent":
weights.ensure_default(a.model) # pip install, first run: download the model files
backends.prepare(a.manifest) # a tiny-agent model from another Hub repo: the same first-run download
from moth.backends import compat
compat.require(a.model)
if not a.manifest.is_default:
print(f"moth: model {a.manifest.describe()}", flush=True)
for n in a.manifest.notes:
print(f"moth: moth.json: {n}", file=sys.stderr, flush=True)
note = backends.privacy_note(a.manifest)
if note:
print(f"moth: WARNING: {note}. moth is meant to stay local; point --model / `moth config model` at "
"localhost to keep it so.", file=sys.stderr, flush=True)
token = load_token(a.token_file)
torch.set_num_threads(a.threads or default_threads())
moth = make_engine(a)
trim_heap()
hosts = ("127.0.0.1", "localhost", "[::1]") + ((a.host,) if a.host not in ("0.0.0.0", "::", "") else ())
srv = Server((a.host, a.port), make_handler(moth, token, hosts))
moth.start_reactor()
url = f"http://127.0.0.1:{a.port}" if a.host in ("0.0.0.0", "::", "") else f"http://{a.host}:{a.port}"
print(f"moth is awake at {url} (home: {moth.home.root})", flush=True)
if on_ready:
on_ready(url)
live["srv"] = srv # from here the lifeline stops serve_forever (and the finally below removes the pidfile)
pidfile = write_pidfile(srv.server_address[1])
try:
srv.serve_forever()
finally:
remove_pidfile(pidfile)
if __name__ == "__main__":
main()