"""Write moth training chats with a teacher behind an OpenAI-compatible server (vLLM).
python -m moth.gen_dialogs --pool --n-per-folder 250 # teacher-written home files -> out/file_pool.json
python -m moth.gen_dialogs --dialogs 12000 --start 0 # chats -> out/dialogs.jsonl (resumable)
One chat: sample a home and a (persona, scenario, length); the teacher plays the human (USER_SIM_PROMPT)
and moth (TEACHER_BIBLE + real bash/read/write tools run in the sandbox). Records use tiny_agent's
canonical message format; filtering and lowercasing happen later in build_sft.py.
"""
from __future__ import annotations
import argparse
import asyncio
import collections
import json
import os
import random
import re
import sys
import time
from openai import AsyncOpenAI
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from moth import _tools as tools # noqa: E402 (vendored copy of tiny_agent/tools.py)
from moth.agent import (AGENT_CHAT_REMIND, AGENT_LENGTHS, AGENT_LENSES, AGENT_PROMPT, AGENT_REMIND, # noqa: E402
AGENT_SCENARIOS, AGENT_TAG, AGENT_TASKS, AGENT_USER_PERSONAS, HARNESS_STYLES,
MOTH_AGENT_RULES, OUTCOMES, agent_turn)
from moth.home import load_pool, sample_home # noqa: E402
from moth.persona import (FILE_POOL_PROMPT, FOLDER_HINTS, FOLDERS, HOME_SCENARIOS, SCENARIOS, # noqa: E402
TALK_SCENARIOS, TEACHER_BIBLE, TOPICS, USER_PERSONAS, USER_SIM_PROMPT)
if os.environ.get("MOTH_NO_BWRAP") == "1": # box without user namespaces: plain bash in the workspace
def _plain(workdir, command, timeout):
return ["timeout", "-k", "1", str(timeout), "env", "-i", f"HOME={workdir}", "PATH=/usr/bin:/bin",
"LANG=C.UTF-8", "bash", "-c", f"cd {workdir} && ulimit -t {timeout}; {command}"]
tools.bwrap_cmd = _plain
OAI_TOOLS = [
{"type": "function", "function": {"name": "bash", "description": "Run a shell command in your home folder ~ (ls, cat, wc, grep...).",
"parameters": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}}},
{"type": "function", "function": {"name": "read", "description": "Read a file with line numbers.",
"parameters": {"type": "object", "properties": {"path": {"type": "string"}, "offset": {"type": "integer"},
"limit": {"type": "integer"}}, "required": ["path"]}}},
{"type": "function", "function": {"name": "write", "description": "Create or overwrite a file.",
"parameters": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}},
"required": ["path", "content"]}}},
]
SIM_SWITCH = ("[DIRECTOR NOTE: for this message only, you are NOT moth. You write the HUMAN side of the chat. "
"Output only the human's next message.]\n\n")
SIM_SYSTEM = ("You write one side of a chat for a story: the HUMAN who is chatting with moth, a tiny cute chatbot moth. "
"You are never moth. Type like a real person in a chat app, in the voice of the person you play: a grandparent "
"or formal person writes proper sentences with capitals and punctuation, a kid or teen types lowercase with slang, "
"a sloppy typer makes typos. Short (often 1-10 words), no greeting after the first message, rarely emojis, never "
"the same message twice. Reply with only the human's next chat message: no preface, no quotes, no notes.")
SIM_SHORT = ("\n\nIMPORTANT: this time write a tiny lazy reply of 1 to 5 words to moth's last message, typed your way, "
"that sounds {mood}. if it names something, use a word moth really said.") # r3: literal examples got copied
SHORT_MOODS = ["amused", "doubtful", "impressed", "teasing", "agreeing", "curious about one word moth used", "bored",
"surprised", "touched", "a bit confused", "playfully annoyed", "sleepy", "delighted", "unconvinced"]
# r3: "what lamp" x1346, mostly about words moth never said: a share of the confusion turns is built in code from
# moth's real last reply (no teacher call)
WHAT_X = ["what {x}?", "what {x}", "wait, {x}?", "{x}?", "what do you mean {x}?", "huh, {x}?", "What {x}?", "a {x}?",
"{x}? what", "why {x}?", "{x}?? lol", "what's {x}"]
WHAT_GENERIC = ["what?", "huh?", "what do you mean?", "wdym", "What do you mean?", "Pardon?", "wait what", "??",
"i don't get it", "what does that mean", "hm?", "Sorry, what?"]
WORD_STOP = set("that this with have from they them then than there their what when where which would could should "
"about maybe think really very just like little tiny small some much many more here know your yours "
"mine been being sure okay also still only once ever never always good nice thing things something "
"sorry thank thanks wait well kind sort feel feels want need make made said says look tell told it's "
"i'm don't that's what's you're can't didn't isn't doesn't i'll i've let's every other same into over "
"back down come came goes going gone does done want these those because while again even most "
"moth friend idea".split())
DET = set("the a an my your his her our some this that these those its their any no".split())
ADJ = set("little big tiny small soft warm old new nice good long short whole same other only first last best very cosy "
"cozy fuzzy dusty bright dark sleepy happy sweet great real funny tall round".split())
WHATX_RE = re.compile(r"\b(?:what|which|why)(?: is| a| an| the)? ([a-z']+)\W*$", re.I)
RECALL_FACTS = [("your name (make one up)", "your name"), ("your pet's name and what animal it is (make it up)", "your pet"),
("your favourite food", "your favourite food"), ("one small thing that happened to you today",
"what happened to you today"), ("your favourite colour", "your favourite colour"),
("where you're going this weekend", "your weekend plan"), ("what your job is", "your job"),
("your best friend's name (make one up)", "your best friend")]
FACT_WORDS = {"name", "favourite", "favorite", "food", "colour", "color", "pet", "today", "friend", "best", "job", "weekend",
"work", "love", "really", "go", "call", "happen", "moth", "thing", "day"}
RECALL_SHARE = {"p": 0.15} # --recall-share: talk chats where the human shares a fact early and asks for it 5+ turns later
LONG_SHARE = {"p": 0.15} # --long-share: talk chats with 16-24 user turns
SIM_NO_REPEAT = "\n\nYou already sent: {prev}. Write something new, not any of these."
NO_THINK = {"chat_template_kwargs": {"enable_thinking": False}}
MOTH_REMIND = ("\n\n(stay moth: one short lowercase line, usually under 15 words. don't say what you are (a moth, small, "
"on a card) unless they ask what you are. files only if they ask about your home or want something saved: "
"one tiny command.)") # --think: LFM drifts long; teacher view only, not in msgs
QBACK = ("\n(this time hand the turn back, on the same line: end with one small question about them, or bring back "
"something they told you earlier.)")
NO_QBACK = "\n(this time just answer: no question at the end, unless a game needs one.)" # r4 test: 61% asked back
RECALL_REMIND = "\n(they told you this earlier in this chat: answer from the chat.)"
LEAKS = ("", "", "<|tool_call", "<|im_") # reasoning or raw tool-call markup in the reply = broken call
TURNS = [3, 4, 5, 6, 6, 7, 8, 8, 9, 10, 12, 14] # back-and-forth: the thing round 2 couldn't do
HOME_SHARE = {"p": 0.15} # --home-share: fraction of chats about moth's files
THINK = {"on": False, "budget": 256} # --think: teacher always reasons first (LFM2.5), capped at --think-budget tokens
TURN_CHOICES: list[int] = [] # --turns: override the chat-length distribution (e.g. long chats)
SCENARIO_IDS: list[int] = [] # --scenario-ids: restrict a run to some scenarios (e.g. small talk only)
def _sft_keep():
"""build_sft.keep (the build-time filter), so moth_turn resamples replies it would cut. build_sft imports torch,
tokenizers and tiny_agent at the top for the token rows only: stub whichever this env lacks."""
import importlib.util
import types
try:
for name, attrs in (("torch", {}), ("tokenizers", {"Tokenizer": object})):
if name not in sys.modules and importlib.util.find_spec(name) is None:
sys.modules[name] = types.ModuleType(name)
sys.modules[name].__dict__.update(attrs)
if "tiny_agent.chat" not in sys.modules:
sys.modules.setdefault("tiny_agent", types.ModuleType("tiny_agent"))
sys.modules["tiny_agent.chat"] = types.ModuleType("tiny_agent.chat")
sys.modules["tiny_agent.chat"].render = None
from moth.build_sft import keep
return keep
except Exception as e: # noqa: BLE001 build_sft changed shape: generate without the check rather than crash
print(f"WARNING: build_sft.keep not importable ({type(e).__name__}: {e}); no build-time checks", flush=True)
return None
SFT_KEEP = _sft_keep()
def sft_problem(meta: dict, msgs: list[dict], text: str) -> str | None:
"""Why build_sft would cut the chat at this new reply (None if it keeps it, or the chat is already cut earlier)."""
if SFT_KEEP is None:
return None
def cut(t):
why = collections.Counter()
m = msgs + [{"role": "assistant", "think": None, "content": t, "tool_calls": []}]
out = SFT_KEEP({**meta, "home_after": {}, "messages": m}, why)
return None if out is not None and len(out) == len(m) else (next(iter(why), "cut").split(":")[0])
try:
r = cut(text)
return r if r and not cut("ok.") else None # "ok." cut too: the problem is earlier (a user turn), not here
except Exception: # noqa: BLE001
return None
def coerce(name: str, args: dict) -> dict:
if name == "read":
for k in ("offset", "limit"):
if k in args:
try:
args[k] = int(args[k])
except (TypeError, ValueError):
args.pop(k)
return args
def transcript(msgs: list[dict]) -> str:
out = []
for m in msgs:
if m["role"] == "user" and m["content"].startswith(AGENT_TAG):
body = m["content"][len(AGENT_TAG):].strip()
out.append("[your AI agent's latest message to you, which moth can also see]\n" + body[:600])
elif m["role"] == "user":
out.append(f"you: {m['content']}")
elif m["role"] == "assistant":
for c in m.get("tool_calls") or []:
a = c["arguments"]
out.append(f"[moth used {c['name']}: {a.get('command') or a.get('path')}]")
if m.get("content"):
out.append(f"moth: {m['content']}")
elif m["role"] == "tool":
for r in m["results"]:
out.append("[output]\n" + (r if len(r) < 400 else r[:400] + " ..."))
return "\n".join(out) or "(nothing yet: write your first message)"
async def chat(client, model, messages, budget=None, **kw):
# --think: LFM2.5's template always opens . Never send enable_thinking=False then: vLLM's qwen3 reasoning parser
# would leave the whole reasoning in .content. vLLM's thinking_token_budget forces after `budget` tokens.
extra = {"thinking_token_budget": THINK["budget"] if budget is None else budget} if THINK["on"] else NO_THINK
for attempt in range(4):
try:
return await client.chat.completions.create(model=model, messages=messages, extra_body=extra, **kw)
except Exception as e: # noqa: BLE001 server restarts / timeouts: back off and retry
if attempt == 3:
raise
await asyncio.sleep(5 * (attempt + 1))
last = e # noqa: F841
META = re.compile(r"\bthe user\b|\bin character\b|\bmoth (should|would|is asked)\b|\blet me (think|check|re)|"
r"\bno markdown\b|\bkeep it (short|simple)\b|\b\d+ words\b|\blet me try\b|\bmisunderstand", re.I)
COMPLEX_CMD = re.compile(r"\bfor\b|\bwhile\b|\||&&|;|\*|-la\b|2>|\$\(|\becho\b|<<|\s-\w*R") # = build_sft: those
_P = r"(?:~/)?(?:[\w.-]+/)?[\w.-]+" # chats get dropped
TINY_CMD = re.compile(rf"\s*(?:ls(?: -[1a])?(?: (?:~/?|(?:~/)?[\w-]+/?))*|cat {_P}(?: {_P})*|"
rf"(?:wc|head)(?: -[a-z]*\d*| -n \d+)? {_P})\s*")
# agent chats: build_sft's TECHY (copied, it imports torch) + over 25 words + the "i'm just a moth / can't understand" dodge
AGENT_BAD = re.compile(r"\b[a-z]\w*_\w+\b|\w\(\)|\b\w+\.(?:py|js|ts|tsx|json|csv|md|ya?ml|toml|sh|rs|go|html|css|sql|cfg|ini)\b|"
r"\b(?:ruff|mypy|pytest|npm|pip|git|github|docker|eslint|webpack|cargo|kubectl|ci|lint|linter|"
r"endpoints?|commits?|regex|api|json|csv|yaml|http|https|redis|schema|parser|tuple|repo)\b|"
r"\btests? (?:all )?(?:pass|passed|passes|passing)\b|\b(?:all )?\d+ tests\b|^(?:\S+\s+){25}\S|"
r"\bi(?:['’]?m| am) (?:just |only )?a (?:little |tiny |small )?moth\b|"
r"\bcan(?:['’]?t|not) (?:understand|read|follow)\b|\bdon['’]?t (?:understand|get it)\b")
BUG = re.compile(r"\bbugs?\b|\bgrand ?ma\b", re.I)
MOTIF = re.compile(r"\bfans?\b", re.I)
def grams(t: str, n=5) -> set:
w = re.sub(r"[^\w\s']", " ", t.lower()).split()
return {tuple(w[i:i + n]) for i in range(len(w) - n + 1)}
EXAMPLE_GRAMS = set().union(*map(grams, re.findall(r'-> "([^"]+)"', MOTH_AGENT_RULES) +
["ohhh it fixed a bug! i'm a bug too. i hope it was nice to it."])) # r3 bible line
STOP = set("the and for that this with you your are was were has have had not but all can its into from what when then "
"than they them there their our out just one two too very like little tiny she her his who how why yes now "
"new agent said".split())
VIBES = [(r"\bsorry|\bapolog", lambda m: re.search(r"sorry|apolog", m)), (r"\bask|\bquestion", lambda m: "?" in m),
(r"\bwords\b|\blong\b|\bbusy\b|\btired\b", lambda m: len(m.split()) > 80),
(r"\bproud\b|\bexcited\b|\bhappy\b|\bpolite\b", lambda m: True)] # vibe it really has, not a made-up "sorry"
def stems(t: str) -> set:
return {re.sub(r"(es|s|ed|ing)$", "", w) for w in re.findall(r"[a-z]{3,}", t.lower()) if w not in STOP}
def agent_stale(text: str, msgs: list[dict]) -> bool:
"""r3 agent chats: LFM copied the few-shot lines, repeated itself, said "bug" to 1 in 3 agent messages that never
mentioned one (the bible's grandma lore), and riffed on words that aren't in the message at all."""
mine = set().union(*(grams(m["content"]) for m in msgs if m["role"] == "assistant" and m.get("content")))
if grams(text) & (EXAMPLE_GRAMS | mine):
return True
if MOTIF.search(text) and any(MOTIF.search(m["content"] or "") for m in msgs if m["role"] == "assistant"):
return True # build_sft drops chats with the fan twice
last = msgs[-1]["content"]
if not last.startswith(AGENT_TAG):
return False
if BUG.search(text) and not re.search(r"\bbugs?\b", last, re.I):
return True
body, t = last[len(AGENT_TAG):].lower(), text.lower()
return not (stems(t) & stems(body) or any(re.search(c, t) and ok(body) for c, ok in VIBES)) # a word or its vibe
def complex_cmd(cmd: str) -> bool:
return bool(cmd) and bool(COMPLEX_CMD.search(cmd) or not TINY_CMD.fullmatch(cmd)) # "" = read/write call
def bash_cmd(tc) -> str:
try:
return str(json.loads(tc.function.arguments or "{}").get("command", "")) if tc.function.name == "bash" else ""
except (json.JSONDecodeError, AttributeError):
return ""
def teacher_view(u: str, remind: str | None = None, extra: str = "") -> str:
if remind: # agent chats: their own reminder, always
return u + remind
return u + (MOTH_REMIND if THINK["on"] else "") + extra
def reply_text(r, cut_ok=False) -> str | None:
"""The reply, or None when the call broke: cut off by max_tokens, or reasoning/tool markup leaked into the text."""
c = r.choices[0]
t = (c.message.content or "").strip()
if (c.finish_reason == "length" and not cut_ok) or any(x in t for x in LEAKS):
return None
return t
def last_moth(msgs: list[dict]) -> str:
return next((m["content"] for m in reversed(msgs) if m["role"] == "assistant" and m.get("content")), "")
def confused(last: str, mine: list[str]) -> str:
"""A "what X?" turn about a word moth really said (or a plain "what?"), built in code."""
toks, words = [w.strip("'") for w in re.findall(r"[a-z']+", last.lower())], []
for j, w in enumerate(toks): # nouns, roughly: the first real word after "the/a/my/..." (skipping adjectives)
if w in DET:
x = next((v for v in toks[j + 1:j + 4] if v not in ADJ), "")
if len(x) >= 3 and x not in WORD_STOP and x not in DET:
words.append(x)
for _ in range(4):
t = random.choice(WHAT_X).format(x=random.choice(words)) if words and random.random() < 0.75 \
else random.choice(WHAT_GENERIC)
if t not in mine:
return t
return t
async def user_turn(client, model, persona, scenario, turns, i, msgs, topic="anything", directive="") -> str | None:
p = USER_SIM_PROMPT.format(persona=persona, scenario=scenario, topic=topic, turns=turns, i=i, transcript=transcript(msgs))
mine = [m["content"] for m in msgs if m["role"] == "user" and not m["content"].startswith(AGENT_TAG)]
last = last_moth(msgs)
short = not directive and msgs and random.random() < 0.4 # without reasoning LFM ignores "one in four very short"
if short and last and random.random() < 0.3:
return confused(last, mine)
if mine: # r3: lazy short lines ("ok and?") came back word for word
p += SIM_NO_REPEAT.format(prev=" / ".join(f'"{x[:100]}"' for x in mine[-8:]))
p += directive
if THINK["on"]: # LFM2.5: under the character sheet it plays moth (even emits tool calls); its prefix cache
# works in 16-token blocks, so sharing the sheet buys nothing. No thinking: capped reasoning leaks into the message
if short:
p += SIM_SHORT.format(mood=random.choice(SHORT_MOODS))
r = await chat(client, model, [{"role": "system", "content": SIM_SYSTEM}, {"role": "user", "content": p}],
budget=0, temperature=1.0, top_p=0.95, max_tokens=100)
else: # same system + tools as moth's calls: Qwen3.x's hybrid prefix cache reuses only whole ~1.5k-token blocks,
# so both roles share one identical leading block (the character sheet) and only the tail is prefilled
r = await chat(client, model, [{"role": "system", "content": TEACHER_BIBLE}, {"role": "user", "content": SIM_SWITCH + p}],
tools=OAI_TOOLS, tool_choice="none", temperature=1.0, top_p=0.95, max_tokens=80)
t = reply_text(r, cut_ok=True)
if not t or "[END]" in t:
return None
lines = [x.strip() for x in t.split("\n") if x.strip()]
if len(lines) > 1 and lines[0].endswith(":"): # "Sure! Here's the human's next message:"
lines = lines[1:]
t = lines[0].strip('"\u201c\u201d').strip()
t = re.sub(r"^(you|me|user|friend|human)\s*:\s*", "", t, flags=re.I).strip()
if not t or re.match(r"moth\s*:", t, re.I):
return None
w = WHATX_RE.search(t) if len(t.split()) <= 5 and last else None
if w and w.group(1).lower() not in WORD_STOP | {"do", "does", "is", "are", "you", "u", "mean", "happened", "now", "else",
"about", "it", "that", "this", "then", "not", "tho", "so"} \
and stems(w.group(1)).isdisjoint(stems(last)) and w.group(1).lower() not in last.lower():
return confused(last, mine) # asked about a word moth never said
return t[:300]
async def moth_turn(client, model, oai: list[dict], msgs: list[dict], ws, stats, talk_only=False, must=None,
meta=None) -> bool | str:
"""Extend oai (teacher view) and msgs (canonical) with moth's reply to the last user message."""
avoid: list[str] = [] # agent chats: rejected lines, shown back to the teacher (else it resamples the same one)
for k in range(8 if talk_only else 5): # up to 3 looks, broken replies are resampled; late rounds may only talk
# agent chats get no tools at all (LFM ran `echo hi` / wrote notes after agent messages, and with
# tool_choice="none" it still writes tool-call markup into the text)
tk = {} if talk_only else {"tools": OAI_TOOLS, "tool_choice": "auto" if k < 3 else "none"}
view = oai if not avoid else oai[:-1] + [{**oai[-1], "content": oai[-1]["content"] + "\n(not these, they are "
"stale, dodgy or too techy: " + " / ".join(avoid[-3:]) + ")"}]
r = await chat(client, model, view, temperature=1.0 if avoid else 0.85, top_p=0.95,
max_tokens=THINK["budget"] + 300 if THINK["on"] else 400, **tk)
m = r.choices[0].message
if r.choices[0].finish_reason == "length":
continue
if m.tool_calls and THINK["on"] and any(complex_cmd(bash_cmd(tc)) for tc in m.tool_calls):
continue # LFM loves `ls -la` and pipes: resample rather than keep a chat build_sft will drop
if m.tool_calls:
calls, results, oai_calls = [], [], []
for tc in m.tool_calls[:3]:
try:
args = coerce(tc.function.name, json.loads(tc.function.arguments or "{}"))
except json.JSONDecodeError:
return False
if tc.function.name not in ("bash", "read", "write"):
return False
res = await asyncio.to_thread(ws.call, tc.function.name, args)
calls.append({"name": tc.function.name, "arguments": args})
results.append(res)
oai_calls.append({"id": tc.id, "type": "function",
"function": {"name": tc.function.name, "arguments": json.dumps(args)}})
oai.append({"role": "assistant", "content": None, "tool_calls": oai_calls})
for tc, res in zip(oai_calls, results):
oai.append({"role": "tool", "tool_call_id": tc["id"], "content": res})
msgs.append({"role": "assistant", "think": None, "content": "", "tool_calls": calls})
msgs.append({"role": "tool", "results": results})
stats["tool_calls"] += len(calls)
continue
text = re.sub(r"^moth\s*:\s*", "", reply_text(r) or "", flags=re.I)
if not text or META.search(text): # empty, cut off, or the teacher's reasoning spilled into the reply
continue
if must and k < 4 and stems(text).isdisjoint(must):
stats["recall_resample"] += 1 # recall turn: the answer must use a word of the fact the friend shared
continue
why = sft_problem(meta, msgs, text) if meta else None # identity, motif cap, explains, bug, style...
if why and (talk_only or k < 4):
stats[f"sft_{why}"] += 1
avoid.append(text)
continue
if talk_only and AGENT_BAD.search(text.lower()):
stats["agent_resample"] += 1 # dodged or explained the code (build_sft would drop the chat): try again
avoid.append(text)
continue
if talk_only and agent_stale(text, msgs):
stats["agent_stale"] += 1
avoid.append(text)
continue
oai.append({"role": "assistant", "content": text})
msgs.append({"role": "assistant", "think": None, "content": text, "tool_calls": []})
if why: # last try still bad: build_sft cuts the chat here, so stop writing it
stats[f"sft_stop_{why}"] += 1
return "stop"
return True
return False
AGENT_POOL: list[dict] = [] # real agent messages from public traces (moth/agent_pool.json)
AGENT_SHARE = {"p": 0.25} # --agent-share: fraction of chats where moth watches a big agent's messages
MID_RUN = re.compile(r":\s*\Z|\b(?:Now )?[Ll]et me (?!know)|\bPhase \d|\bI'll now\b|\bLet's\b|^(?:Now|Next),? |"
r"^I (?:see|notice)\b|^(?:Excellent|Perfect|Great)! Now\b", re.M)
async def agent_message(client, model, rng) -> tuple[str, str]:
"""(raw message, source): a real one from the pool most of the time, else one written by the teacher."""
if AGENT_POOL and rng.random() < 0.75:
m = rng.choice(AGENT_POOL)
return m["text"], m["source"]
p = AGENT_PROMPT.format(style=rng.choice(HARNESS_STYLES), task=rng.choice(AGENT_TASKS), outcome=rng.choice(OUTCOMES),
length=rng.choice(AGENT_LENGTHS))
r = await chat(client, model, [{"role": "user", "content": p}], temperature=0.9, top_p=0.95,
max_tokens=THINK["budget"] + 900 if THINK["on"] else 700)
return reply_text(r) or "", "teacher"
async def agent_dialog(client, model, seed, rng, files, stats) -> dict | None:
"""moth watches a big coding agent: agent message -> moth's silly take -> the human vents or chats -> repeat."""
persona, scenario = rng.choice(AGENT_USER_PERSONAS), rng.choice(AGENT_SCENARIOS)
msgs: list[dict] = []
oai = [{"role": "system", "content": TEACHER_BIBLE}]
sources = []
async def say(u: str, remind: str) -> bool:
msgs.append({"role": "user", "content": u})
oai.append({"role": "user", "content": teacher_view(u, remind)})
if await moth_turn(client, model, oai, msgs, ws, stats, talk_only=True,
meta={"type": "agent", "home": files, "topic": ""}):
return True
msgs.pop() # no fresh line after all retries: end the chat before this message
return False
with tools.Workspace(files) as ws:
ok = True
for k in range(rng.randint(2, 4)):
if k == 0 and rng.random() < 0.4: # sometimes the human says something first
u = await user_turn(client, model, persona, scenario, 6, 1, msgs, "the AI agent")
if u and not (ok := await say(u, AGENT_CHAT_REMIND)):
break
raw, src = await agent_message(client, model, rng)
if not raw:
break
if not (ok := await say(agent_turn(raw), AGENT_REMIND.format(lens=rng.choice(AGENT_LENSES)))):
break
sources.append(src)
for j in range(rng.choice([0, 1, 1, 2, 3])):
u = await user_turn(client, model, persona, scenario, 6, j + 2, msgs, "the AI agent")
if u is None or not (ok := await say(u, AGENT_CHAT_REMIND)):
break
if not ok:
break
if not sources: # moth never reacted to an agent message
stats["bad"] += 1
return None
stats["agent_cut"] += not ok
return {"seed": seed, "type": "agent", "persona": persona, "scenario": scenario, "agent_sources": sources,
"turns": len(msgs), "home": files, "home_after": files, "messages": msgs}
async def one_dialog(client, model, seed, pool, stats) -> dict | None:
rng = random.Random(seed)
files = sample_home(rng, pool)
if not SCENARIO_IDS and rng.random() < AGENT_SHARE["p"]:
return await agent_dialog(client, model, seed, rng, files, stats)
if SCENARIO_IDS:
scen = [SCENARIOS[i] for i in SCENARIO_IDS]
else: # the product is conversation: most chats are pure talk, a few use moth's home folder
scen = HOME_SCENARIOS if rng.random() < HOME_SHARE["p"] else TALK_SCENARIOS
persona, scenario, turns = rng.choice(USER_PERSONAS), rng.choice(scen), rng.choice(TURN_CHOICES or TURNS)
topic = rng.choice(TOPICS)
talk = scen is TALK_SCENARIOS
if talk and not TURN_CHOICES and rng.random() < LONG_SHARE["p"]:
turns = rng.randint(16, 24) # r3 had no chat over 14 turns
recall = None
if talk and rng.random() < RECALL_SHARE["p"]: # r3: ~50 of 21k chats asked for something back
fact, short = rng.choice(RECALL_FACTS)
fi = rng.choice([1, 2])
recall = {"fact": short, "share_turn": fi, "ask_turn": fi + rng.randint(5, 8)}
turns = max(turns, recall["ask_turn"] + rng.randint(1, 3))
msgs: list[dict] = []
oai = [{"role": "system", "content": TEACHER_BIBLE}]
with tools.Workspace(files) as ws:
for i in range(1, turns + 1):
directive, extra = "", QBACK if rng.random() < 0.2 else NO_QBACK
if recall and i == recall["share_turn"]:
directive = f"\n\nIn THIS message, casually tell moth {fact}."
elif recall and i == recall["ask_turn"]:
directive = (f"\n\nIn THIS message, ask moth if it remembers {short} (you told it earlier). Ask in your "
"own words and don't give the answer.")
extra = RECALL_REMIND
if (recall and i < recall["ask_turn"]) or (turns >= 16 and i < turns - 3):
directive += "\n\nThe chat is not finished: do not write [END] yet."
u = await user_turn(client, model, persona, scenario, turns, i, msgs, topic, directive)
for _ in range(3 if "[END]" in directive else 0): # long / recall chat: broken sim calls (leaked
if u is not None: # reasoning) ended 13 of 20 long chats early
break
u = await user_turn(client, model, persona, scenario, turns, i, msgs, topic, directive)
if u is None:
break
if recall and i == recall["share_turn"]:
recall["said"] = u
must = None
if recall and i == recall["ask_turn"] and recall.get("said"):
recall["asked"] = u
must = stems(recall["said"]) - FACT_WORDS or None
msgs.append({"role": "user", "content": u})
oai.append({"role": "user", "content": teacher_view(u, extra=extra)})
r = await moth_turn(client, model, oai, msgs, ws, stats, must=must,
meta={"type": "home" if scen is HOME_SCENARIOS else "talk", "home": files, "topic": topic})
if not r:
stats["bad"] += 1
return None
if r == "stop":
break
if not any(m["role"] == "assistant" for m in msgs):
return None
final = {}
for root, _, names in os.walk(ws.root):
for n in names:
p = os.path.join(root, n)
final[os.path.relpath(p, ws.root)] = open(p, errors="replace").read()[:2000]
return {"seed": seed, "type": "home" if scen is HOME_SCENARIOS else "talk", "persona": persona,
"scenario": scenario, "topic": topic, "turns": turns, "recall": recall, "home": files,
"home_after": final, "messages": msgs}
async def run_dialogs(a, client, model):
pool = load_pool(a.pool_file)
path = os.path.join(a.out, "dialogs.jsonl")
done = set()
if os.path.exists(path):
for line in open(path):
try:
done.add(json.loads(line)["seed"])
except (ValueError, KeyError):
pass
todo = [s for s in range(a.start, a.start + a.dialogs) if s not in done]
stats = collections.Counter(ok=0, bad=0, err=0, tool_calls=0, agent_resample=0, agent_stale=0, agent_cut=0,
recall_resample=0)
sem = asyncio.Semaphore(a.concurrency)
f = open(path, "a")
t0 = time.time()
async def go(seed):
async with sem:
try:
d = await one_dialog(client, model, seed, pool, stats)
except Exception as e: # noqa: BLE001
stats["err"] += 1
print(f"seed {seed}: {type(e).__name__}: {e}", flush=True)
return
if d:
f.write(json.dumps(d, ensure_ascii=False) + "\n")
f.flush()
stats["ok"] += 1
if stats["ok"] % 100 == 0:
el = time.time() - t0
print(f"{stats} {stats['ok'] / el * 60:.0f} dialogs/min, {len(todo) - stats['ok']} left", flush=True)
await asyncio.gather(*(go(s) for s in todo))
print(f"done {stats} in {(time.time() - t0) / 60:.1f} min", flush=True)
async def run_pool(a, client, model):
path = a.pool_file
pool = json.load(open(path)) if os.path.exists(path) else {}
sem = asyncio.Semaphore(a.concurrency)
async def batch(folder, k):
async with sem:
have = list(pool.get(folder, {}))
avoid = ", ".join(random.Random(k).sample(have, min(40, len(have)))) or "none"
p = FILE_POOL_PROMPT.format(n=10, folder=folder, hint=FOLDER_HINTS[folder], avoid=avoid, bible=TEACHER_BIBLE)
try:
r = await chat(client, model, [{"role": "user", "content": p}], temperature=1.0, top_p=0.95,
max_tokens=6000 if THINK["on"] else 1500)
txt = r.choices[0].message.content or ""
items = json.loads(txt[txt.index("["): txt.rindex("]") + 1])
except Exception: # noqa: BLE001
return
for it in items:
name = re.sub(r"[^a-z0-9_.]", "_", str(it.get("name", "")).lower().strip())
if not name.endswith(".txt") or len(name) > 48 or not it.get("content"):
continue
pool.setdefault(folder, {}).setdefault(name, str(it["content"]).strip()[:400])
await asyncio.gather(*(batch(fo, k) for fo in FOLDERS for k in range(a.n_per_folder // 10)))
json.dump(pool, open(path, "w"), indent=1, ensure_ascii=False)
print("pool sizes:", {k: len(v) for k, v in pool.items()}, flush=True)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--base-url", default="http://127.0.0.1:8000/v1")
ap.add_argument("--model", default=None)
ap.add_argument("--out", default="out")
ap.add_argument("--pool", action="store_true")
ap.add_argument("--n-per-folder", type=int, default=250)
ap.add_argument("--pool-file", default=None)
ap.add_argument("--dialogs", type=int, default=0)
ap.add_argument("--start", type=int, default=0)
ap.add_argument("--concurrency", type=int, default=96)
ap.add_argument("--agent-share", type=float, default=0.25, help="fraction of chats watching a big agent")
ap.add_argument("--agent-pool", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "agent_pool.json"))
ap.add_argument("--eval-agent-msgs", default=os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"data", "eval_agent_msgs.json"),
help="held-out eval agent messages: dropped from the agent pool")
ap.add_argument("--home-share", type=float, default=0.15, help="fraction of chats about moth's home/files")
ap.add_argument("--recall-share", type=float, default=0.15, help="talk chats: share a fact early, ask 5+ turns later")
ap.add_argument("--long-share", type=float, default=0.15, help="talk chats with 16-24 user turns")
ap.add_argument("--think", action="store_true", help="teacher always reasons first (LFM2.5): cap it, split it off")
ap.add_argument("--think-budget", type=int, default=256, help="--think: max reasoning tokens per moth/agent call")
ap.add_argument("--turns", default="", help="comma-separated user-turn counts to sample from")
ap.add_argument("--scenario-ids", default="", help="comma-separated indexes into persona.SCENARIOS")
a = ap.parse_args()
SCENARIO_IDS.extend(int(x) for x in a.scenario_ids.split(",") if x)
TURN_CHOICES.extend(int(x) for x in a.turns.split(",") if x)
THINK["on"], THINK["budget"] = a.think, a.think_budget
HOME_SHARE["p"] = a.home_share
RECALL_SHARE["p"], LONG_SHARE["p"] = a.recall_share, a.long_share
AGENT_SHARE["p"] = a.agent_share
if os.path.exists(a.agent_pool): # final messages only: drop mid-run narration ("Now let me...", "Phase 2", "...:")
AGENT_POOL.extend(m for m in json.load(open(a.agent_pool)) if not MID_RUN.search(m["text"]))
n = len(AGENT_POOL)
if os.path.exists(a.eval_agent_msgs): # held-out eval messages never go into training chats
norm = lambda t: " ".join(str(t).split()) # noqa: E731
ev = [m["text"] if isinstance(m, dict) else m for m in json.load(open(a.eval_agent_msgs))]
exact, heads = {norm(t) for t in ev}, {norm(t)[:200] for t in ev}
AGENT_POOL[:] = [m for m in AGENT_POOL if norm(m["text"]) not in exact and norm(m["text"])[:200] not in heads]
print(f"agent pool: removed {n - len(AGENT_POOL)} eval messages ({len(ev)} held out)", flush=True)
else:
print(f"agent pool: WARNING no eval file at {a.eval_agent_msgs}, nothing held out", flush=True)
print(f"agent pool: {len(AGENT_POOL)} final messages", flush=True)
os.makedirs(a.out, exist_ok=True)
a.pool_file = a.pool_file or os.path.join(a.out, "file_pool.json")
client = AsyncOpenAI(base_url=a.base_url, api_key="x", timeout=300)
async def run():
model = a.model or (await client.models.list()).data[0].id
if a.pool:
await run_pool(a, client, model)
if a.dialogs:
await run_dialogs(a, client, model)
asyncio.run(run())
if __name__ == "__main__":
main()