laya-mlx 0.1.0__py3-none-any.whl

laya_mlx/__init__.py

@@ -0,0 +1,34 @@
+"""Laya typed decisions on Apple silicon with MLX."""
+
+from .agent import Agent, RLAgent, load
+from .email import clean_email_body, email_state
+from .lang import analyse as detect_language
+from .lang import detect_script, is_english
+from .presets import (
+    email_questions,
+    guard_questions,
+    moderation_questions,
+    router_questions,
+    triage_questions,
+)
+from .router import DEFAULT_MODELS, RouteDecision, Router
+
+__version__ = "0.1.0"
+__all__ = [
+    "Agent",
+    "RLAgent",
+    "load",
+    "Router",
+    "RouteDecision",
+    "DEFAULT_MODELS",
+    "detect_language",
+    "detect_script",
+    "is_english",
+    "clean_email_body",
+    "email_state",
+    "email_questions",
+    "guard_questions",
+    "moderation_questions",
+    "router_questions",
+    "triage_questions",
+]

laya_mlx/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+main()

laya_mlx/agent.py

@@ -0,0 +1,262 @@
+"""Public MLX inference runtime; prompt and result formats follow upstream Laya."""
+
+import json
+import math
+from pathlib import Path, PurePosixPath
+
+import mlx.core as mx
+import numpy as np
+from huggingface_hub import snapshot_download
+
+from .common import QTYPES, build_sequence, confidence_from_probs, render_options, temp_bucket
+from .model import DecisionModel, EncoderConfig, sanitize_weights
+from .prepared import PrefixCache
+from .tokenizer import Tokenizer
+
+DTYPES = {"float32": mx.float32, "float16": mx.float16, "bfloat16": mx.bfloat16}
+
+
+def resolve_model(model_id_or_path, *, token=None, subfolder=None, revision=None):
+    if subfolder:
+        part = PurePosixPath(subfolder)
+        if part.is_absolute() or ".." in part.parts:
+            raise ValueError("subfolder must be a relative path inside the model repository")
+    path = Path(model_id_or_path).expanduser()
+    if not path.exists():
+        value = str(model_id_or_path)
+        if value.startswith(("/", "./", "../", "~")) or isinstance(model_id_or_path, Path):
+            raise FileNotFoundError(f"Local model directory does not exist: {value}")
+        prefix = subfolder.rstrip("/") + "/" if subfolder else ""
+        patterns = [
+            prefix + name
+            for name in (
+                "model.safetensors",
+                "rl_agent_config.json",
+                "encoder/config.json",
+                "tokenizer/*",
+                "mlx_config.json",
+            )
+        ]
+        path = Path(
+            snapshot_download(value, token=token, revision=revision, allow_patterns=patterns)
+        )
+    if subfolder:
+        path /= subfolder
+    for name in ("model.safetensors", "rl_agent_config.json", "encoder/config.json"):
+        if not (path / name).is_file():
+            raise FileNotFoundError(f"Not a complete Laya checkpoint: {path / name} is missing")
+    return path
+
+
+def collate_items(items, pad_id, *, pad_to_multiple=None, max_length=None):
+    if not items:
+        raise ValueError("Cannot collate an empty batch")
+    n, length = len(items), max(len(item["ids"]) for item in items)
+    if pad_to_multiple:
+        length = ((length + pad_to_multiple - 1) // pad_to_multiple) * pad_to_multiple
+        if max_length is not None:
+            length = min(length, max_length)
+    count = max(2, max(len(item["markers"]) for item in items))
+    batch = {
+        "input_ids": np.full((n, length), pad_id, dtype=np.int32),
+        "attention_mask": np.zeros((n, length), dtype=np.bool_),
+        "marker_pos": np.zeros((n, count), dtype=np.int32),
+        "marker_mask": np.zeros((n, count), dtype=np.bool_),
+        "qtype": np.array([item["qtype"] for item in items], dtype=np.int32),
+    }
+    for i, item in enumerate(items):
+        length, count = len(item["ids"]), len(item["markers"])
+        batch["input_ids"][i, :length] = item["ids"]
+        batch["attention_mask"][i, :length] = True
+        batch["marker_pos"][i, :count] = item["markers"]
+        batch["marker_mask"][i, :count] = True
+    return batch
+
+
+class Agent:
+    def __init__(
+        self,
+        model_id_or_path="convaiinnovations/laya",
+        device=None,
+        token=None,
+        subfolder=None,
+        *,
+        dtype="float16",
+        revision=None,
+        batch_size=16,
+        compile=False,
+        pad_to_multiple=None,
+        cache_prompts=False,
+    ):
+        if dtype not in DTYPES:
+            raise ValueError(f"dtype must be one of {list(DTYPES)}")
+        if device not in (None, "gpu", "metal", "cpu"):
+            raise ValueError("MLX device must be 'gpu', 'metal', or 'cpu'")
+        if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size < 1:
+            raise ValueError("batch_size must be a positive integer")
+        self.device = (
+            mx.default_device() if device is None else (mx.cpu if device == "cpu" else mx.gpu)
+        )
+        self.dtype = DTYPES[dtype]
+        self.batch_size = batch_size
+        if pad_to_multiple is not None and (
+            not isinstance(pad_to_multiple, int)
+            or isinstance(pad_to_multiple, bool)
+            or pad_to_multiple < 1
+        ):
+            raise ValueError("pad_to_multiple must be a positive integer or None")
+        self.pad_to_multiple = pad_to_multiple
+        self._prefix_cache = PrefixCache() if cache_prompts else None
+        self.model_id = str(model_id_or_path)
+        self.revision = revision
+        self.model_dir = resolve_model(
+            model_id_or_path, token=token, subfolder=subfolder, revision=revision
+        )
+        self.cfg = json.loads((self.model_dir / "rl_agent_config.json").read_text())
+        self.encoder_cfg = json.loads((self.model_dir / "encoder/config.json").read_text())
+        if "encoder" not in self.cfg or "head_layers" not in self.cfg:
+            raise ValueError("Laya config must specify encoder and head_layers")
+        enc_cfg = EncoderConfig.from_dict(self.encoder_cfg)
+        max_len = self.cfg.get("max_len", 512)
+        head_max_len = self.cfg.get("head_max_len", 192)
+        if not 4 < head_max_len < max_len <= enc_cfg.max_position_embeddings:
+            raise ValueError("Expected 4 < head_max_len < max_len <= max_position_embeddings")
+        self.temperature = self.cfg.get("temperature", [1.0, 1.0, 1.0])
+        self.temperature_by_options = self.cfg.get("temperature_by_options", {})
+        if len(self.temperature) != 3 or any(
+            not math.isfinite(float(t)) or float(t) <= 0
+            for t in [*self.temperature, *self.temperature_by_options.values()]
+        ):
+            raise ValueError("Calibration temperatures must be finite and positive")
+        self.tok = Tokenizer(self.model_dir / "tokenizer")
+        with mx.stream(self.device):
+            self.model = DecisionModel(enc_cfg, self.cfg)
+            weights = sanitize_weights(mx.load(str(self.model_dir / "model.safetensors")))
+            weights = {k: v.astype(self.dtype) for k, v in weights.items()}
+            self.model.load_weights(list(weights.items()), strict=True)
+            self.model.eval()
+            mx.eval(self.model.parameters())
+        # Frozen inference instance: changing weights or module structure requires a new Agent.
+        self._inference = mx.compile(self.model) if compile else self.model
+
+    @staticmethod
+    def _to_internal(qdef):
+        if not isinstance(qdef, dict):
+            raise ValueError("Each question must be a dictionary")
+        kind = qdef.get("type")
+        if kind not in QTYPES:
+            raise ValueError(f"Unknown question type {kind!r}; expected choice, score, or noul")
+        if "instructions" not in qdef:
+            raise ValueError("Question is missing instructions")
+        criteria = qdef.get("criteria")
+        if kind == "choice":
+            if isinstance(criteria, list):
+                if not all(isinstance(c, str) for c in criteria):
+                    raise ValueError("Choice labels must be strings")
+                if len(set(criteria)) != len(criteria):
+                    raise ValueError("Choice labels must be unique")
+                criteria = dict.fromkeys(criteria)
+            if not isinstance(criteria, dict) or not criteria:
+                raise ValueError("Choice criteria must be a nonempty dictionary or list")
+            if not all(isinstance(k, str) for k in criteria):
+                raise ValueError("Choice labels must be strings")
+        elif kind == "score":
+            if not isinstance(criteria, list) or not criteria:
+                raise ValueError("Score criteria must be a nonempty list")
+        elif criteria is not None and not isinstance(criteria, dict):
+            raise ValueError("Noul criteria must be a dictionary with false/true descriptions")
+        instructions = qdef["instructions"]
+        if not isinstance(instructions, str):
+            instructions = json.dumps(instructions)
+        return {"t": kind, "ins": instructions, "crit": criteria}
+
+    def prepare(self, state, questions):
+        """Construct upstream-compatible CPU inputs, useful for parity and profiling."""
+        if self._prefix_cache is not None:
+            return self._prefix_cache.prepare(self, state, questions)
+        if not isinstance(questions, dict):
+            raise ValueError("questions must be a dictionary keyed by question id")
+        items, internal = [], []
+        for qid, definition in questions.items():
+            q = self._to_internal(definition)
+            ids, markers = build_sequence(
+                self.tok, state, q, self.cfg.get("max_len", 512), self.cfg.get("head_max_len", 192)
+            )
+            if len(markers) != len(render_options(q)):
+                raise ValueError(f"Question {qid!r} has too many options for the token budget")
+            items.append({"ids": ids, "markers": markers, "qtype": QTYPES[q["t"]]})
+            internal.append(q)
+        return items, internal
+
+    def forward(self, batch):
+        """Run one prepared batch and return evaluated MLX logits on this agent's device."""
+        with mx.stream(self.device):
+            tensors = {k: mx.array(v) for k, v in batch.items()}
+            result = self._inference(**tensors)
+            mx.eval(result)
+        return result
+
+    def system_one(self, state, questions):
+        items, internal = self.prepare(state, questions)
+        answers = {}
+        question_ids = list(questions)
+        for start in range(0, len(items), self.batch_size):
+            chunk = items[start : start + self.batch_size]
+            batch = collate_items(
+                chunk,
+                self.tok.pad_token_id,
+                pad_to_multiple=self.pad_to_multiple,
+                max_length=self.cfg.get("max_len", 512),
+            )
+            logits, act = self.forward(batch)
+            logits, act = np.asarray(logits), np.asarray(act)
+            if not np.isfinite(logits).all() or not np.isfinite(act).all():
+                raise FloatingPointError("Non-finite model outputs; retry with dtype='float32'")
+            act = np.exp(act - act.max(axis=-1, keepdims=True))
+            act /= act.sum(axis=-1, keepdims=True)
+            for row, item in enumerate(chunk):
+                qid, q = question_ids[start + row], internal[start + row]
+                k, qt = len(item["markers"]), item["qtype"]
+                scale = self.temperature_by_options.get(temp_bucket(qt, k), self.temperature[qt])
+                z = logits[row, :k] / max(1e-3, float(scale))
+                p = np.exp(z - z.max())
+                p /= p.sum()
+                answer = {
+                    "type": q["t"],
+                    "confidence": round(confidence_from_probs(p, k), 4),
+                    "action": {"act_probability": round(float(act[row, 0]), 4)},
+                }
+                if q["t"] == "choice":
+                    labels = list(q["crit"])
+                    answer.update(
+                        choice=labels[int(p.argmax())],
+                        probabilities={label: round(float(v), 4) for label, v in zip(labels, p)},
+                    )
+                elif q["t"] == "score":
+                    answer.update(
+                        score=round(float((np.arange(k) * p).sum()), 4),
+                        legend={str(i): value for i, value in enumerate(q["crit"])},
+                        probabilities={str(i): round(float(v), 4) for i, v in enumerate(p)},
+                    )
+                else:
+                    answer.update(
+                        noul=round(float(p[1]), 4),
+                        confidence=round(max(float(p[1]), 1.0 - float(p[1])), 4),
+                    )
+                answers[qid] = answer
+        return {
+            "model": "laya-rl-agent",
+            "answers": answers,
+            "usage": {"input_tokens": sum(len(item["ids"]) for item in items), "output_tokens": 0},
+        }
+
+    predict = system_one
+
+
+RLAgent = Agent
+
+
+def load(
+    model_id_or_path="convaiinnovations/laya", device=None, token=None, subfolder=None, **kwargs
+):
+    return Agent(model_id_or_path, device=device, token=token, subfolder=subfolder, **kwargs)

laya_mlx/cli.py

@@ -0,0 +1,55 @@
+"""Command-line prediction and checkpoint conversion."""
+
+import argparse
+import json
+from pathlib import Path
+
+from . import __version__
+from .agent import DTYPES, Agent
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="laya-mlx", description=__doc__)
+    parser.add_argument("--version", action="version", version=__version__)
+    commands = parser.add_subparsers(dest="command", required=True)
+    for command in ("predict", "convert"):
+        sub = commands.add_parser(command)
+        sub.add_argument("--model", default="convaiinnovations/laya")
+        sub.add_argument("--subfolder")
+        sub.add_argument("--revision")
+        sub.add_argument("--dtype", choices=DTYPES, default="float16")
+        if command == "predict":
+            source = sub.add_mutually_exclusive_group(required=True)
+            source.add_argument("--state", help="Plain text input")
+            source.add_argument("--state-file", type=Path, help="JSON state file")
+            sub.add_argument(
+                "--questions", required=True, type=Path, help="JSON question definitions"
+            )
+            sub.add_argument("--device", choices=("gpu", "cpu"), default="gpu")
+            sub.add_argument("--batch-size", type=int, default=16)
+        else:
+            sub.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args(argv)
+    if args.command == "convert":
+        from .convert import convert
+
+        result = convert(
+            args.model,
+            args.output,
+            dtype=args.dtype,
+            revision=args.revision,
+            subfolder=args.subfolder,
+        )
+        print(json.dumps({"output": str(result), "dtype": args.dtype}))
+    else:
+        state = args.state if args.state is not None else json.loads(args.state_file.read_text())
+        questions = json.loads(args.questions.read_text())
+        agent = Agent(
+            args.model,
+            device=args.device,
+            dtype=args.dtype,
+            revision=args.revision,
+            subfolder=args.subfolder,
+            batch_size=args.batch_size,
+        )
+        print(json.dumps(agent.predict(state, questions), ensure_ascii=False, indent=2))

laya_mlx/common.py

@@ -0,0 +1,119 @@
+"""Laya prompt construction and calibration, adapted from upstream (see NOTICE)."""
+
+import json
+import math
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+QTYPES = {"choice": 0, "score": 1, "noul": 2}
+QTYPE_NAMES = {v: k for k, v in QTYPES.items()}
+
+
+def serialize_state(state: Union[str, dict, list]) -> str:
+    if isinstance(state, str):
+        return state
+    return json.dumps(state, ensure_ascii=False)
+
+
+def render_criterion(value) -> str:
+    """Render one criterion value as text.
+
+    Strings pass through; anything structured (dict, list, number) becomes compact JSON, so a
+    rubric reads as JSON rather than a Python repr. Without this a dict-valued criterion
+    crashed `noul` outright and leaked `{'desc': ...}` into `choice` and `score` prompts.
+    """
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False, separators=(", ", ": "), default=str)
+
+
+def render_options(q: Dict) -> List[str]:
+    """Render option texts in label-index order. Noul is always [false, true]."""
+    t, crit = q["t"], q.get("crit")
+    if t == "choice":
+        # only None/"" mean "no description"; 0 and False are legitimate criterion values
+        return [
+            k if v is None or v == "" else "%s: %s" % (k, render_criterion(v))
+            for k, v in crit.items()
+        ]
+    if t == "score":
+        return ["level %d: %s" % (i, render_criterion(c)) for i, c in enumerate(crit)]
+    crit = crit or {}
+    false_crit, true_crit = crit.get("false"), crit.get("true")
+    return [
+        "false: "
+        + (
+            render_criterion(false_crit)
+            if false_crit not in (None, "")
+            else "no, the statement does not hold"
+        ),
+        "true: "
+        + (
+            render_criterion(true_crit)
+            if true_crit not in (None, "")
+            else "yes, the statement holds"
+        ),
+    ]
+
+
+def build_prefix(tok, q: Dict, head_max_len: int = 192, option_order=None):
+    """Build the question-only prefix, before state tokens and final truncation."""
+    mask_tok = tok.mask_token
+    opts = render_options(q)
+    order = option_order if option_order is not None else list(range(len(opts)))
+    ins = str(q["ins"]).replace(mask_tok, " ")
+    head_ids = tok("%s question: %s" % (q["t"], ins), add_special_tokens=False)["input_ids"]
+    opt_ids = []
+    for i in order:
+        opt_ids.append(
+            [tok.mask_token_id]
+            + tok(" " + opts[i].replace(mask_tok, " "), add_special_tokens=False)["input_ids"][:48]
+        )
+    opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+    if opt_budget < 16:
+        per = max(4, (head_max_len - 16) // max(1, len(opt_ids)))
+        opt_ids = [o[:per] for o in opt_ids]
+        opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+    head_ids = head_ids[: max(8, opt_budget)]
+    ids = [tok.cls_token_id] + head_ids + [tok.sep_token_id]
+    markers = []
+    for o in opt_ids:
+        markers.append(len(ids))
+        ids.extend(o)
+    ids.append(tok.sep_token_id)
+    return ids, markers
+
+
+def build_sequence(
+    tok,
+    state: Union[str, dict, list],
+    q: Dict,
+    max_len: int = 512,
+    head_max_len: int = 192,
+    option_order: Optional[List[int]] = None,
+    truncate_left: bool = False,
+):
+    """Format: [CLS] <type> instructions [SEP] [MASK] opt0 [MASK] opt1 ... [SEP] state [SEP]."""
+    ids, markers = build_prefix(tok, q, head_max_len, option_order)
+    room = max(0, max_len - len(ids) - 1)
+    st = tok(serialize_state(state).replace(tok.mask_token, " "), add_special_tokens=False)[
+        "input_ids"
+    ]
+    st = st[-room:] if truncate_left else st[:room]
+    ids = ids + st + [tok.sep_token_id]
+    return ids[:max_len], [m for m in markers if m < max_len]
+
+
+def confidence_from_probs(p: np.ndarray, k: int) -> float:
+    """Normalized Shannon entropy confidence: 1 - H(p) / log(k)."""
+    if k < 2:
+        return 1.0
+    p = p[:k]
+    ent = -(p * np.log(np.clip(p, 1e-12, 1.0))).sum()
+    return float(np.clip(1.0 - ent / math.log(k), 0.0, 1.0))
+
+
+def temp_bucket(qtype: int, k: int) -> str:
+    size = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+"
+    return "%s:%s" % (QTYPE_NAMES[int(qtype)], size)

laya_mlx/convert.py

@@ -0,0 +1,48 @@
+"""Export a standalone MLX checkpoint. Weight downloads stay outside Git."""
+
+import json
+import shutil
+from pathlib import Path
+
+import mlx.core as mx
+from mlx.utils import tree_flatten
+
+from .agent import Agent
+
+
+def convert(
+    model_id_or_path, output, *, dtype="float16", revision=None, subfolder=None, token=None
+):
+    """Convert an upstream checkpoint to MLX parameter names and the requested dtype.
+
+    The destination must not exist. If export fails, the newly created partial
+    directory is removed; existing checkpoints are never overwritten.
+    """
+    output = Path(output).expanduser()
+    if output.exists():
+        raise FileExistsError(f"Output already exists: {output}")
+    agent = Agent(
+        model_id_or_path, dtype=dtype, revision=revision, subfolder=subfolder, token=token
+    )
+    output.mkdir(parents=True)
+    try:
+        shutil.copytree(agent.model_dir / "tokenizer", output / "tokenizer")
+        (output / "encoder").mkdir()
+        (output / "encoder/config.json").write_text(json.dumps(agent.encoder_cfg, indent=2) + "\n")
+        (output / "rl_agent_config.json").write_text(json.dumps(agent.cfg, indent=2) + "\n")
+        mx.save_safetensors(
+            str(output / "model.safetensors"), dict(tree_flatten(agent.model.parameters()))
+        )
+        metadata = {
+            "format": "laya-mlx",
+            "format_version": 1,
+            "dtype": dtype,
+            "source": str(model_id_or_path),
+            "revision": revision,
+            "subfolder": subfolder,
+        }
+        (output / "mlx_config.json").write_text(json.dumps(metadata, indent=2) + "\n")
+    except BaseException:
+        shutil.rmtree(output)
+        raise
+    return output

laya_mlx/email.py

@@ -0,0 +1,101 @@
+# Derived from Laya (Apache-2.0); see NOTICE. Modified for laya-mlx.
+"""Email utilities for cleaning and structuring email inputs in laya."""
+
+import re
+from typing import Dict, Optional
+
+_QUOTE_HEADERS = [
+    re.compile(r"^\s*On .{0,300}wrote:\s*$", re.I),
+    re.compile(r"^\s*-{2,}\s*(Original|Forwarded) Message\s*-{2,}", re.I),
+    re.compile(r"^\s*_{8,}\s*$"),
+    re.compile(r"^\s*From:\s.+$", re.I),
+]
+_SIGNATURE_MARKERS = [
+    re.compile(r"^\s*--\s*$"),
+    re.compile(
+        r"^\s*(best|kind|warm|many thanks|thanks|thank you|regards|cheers|sincerely)[\w ,!.]*$",
+        re.I,
+    ),
+    re.compile(r"^\s*sent from my (iphone|android|mobile|ipad)", re.I),
+]
+_DISCLAIMER = re.compile(
+    r"(confidential|intended (solely )?for the (use of the )?(named )?(addressee|recipient)|"
+    r"if you (have )?received this (e-?mail|message) in error)",
+    re.I,
+)
+
+
+def clean_email_body(body: str, max_chars: int = 3000) -> str:
+    """Remove quoted email history, signatures and disclaimers to keep input focused."""
+    text = (body or "").replace("\r\n", "\n").replace("\r", "\n").replace("\\n", "\n")
+    lines = []
+    for line in text.split("\n"):
+        if any(p.match(line) for p in _QUOTE_HEADERS) and lines:
+            break
+        if line.lstrip().startswith(">"):
+            continue
+        lines.append(line.rstrip())
+    cut = len(lines)
+    for i in range(max(1, min(int(len(lines) * 0.6), len(lines) - 8)), len(lines)):
+        if len(lines[i].strip()) <= 40 and any(p.match(lines[i]) for p in _SIGNATURE_MARKERS):
+            cut = i
+            break
+    lines = lines[:cut]
+    paragraphs = [p for p in re.split(r"\n\s*\n", "\n".join(lines)) if not _DISCLAIMER.search(p)]
+    text = re.sub(r"[ \t]+", " ", "\n\n".join(p.strip() for p in paragraphs if p.strip()))
+    return text[:max_chars]
+
+
+def email_state(
+    subject: str, body: str, sender: Optional[str] = None, clean: bool = True, **extra
+) -> Dict:
+    """Construct a clean state dictionary for email classification."""
+    state = {
+        "subject": (subject or "").strip(),
+        "body": clean_email_body(body) if clean else (body or ""),
+    }
+    if sender:
+        state["from"] = sender
+    state.update({k: v for k, v in extra.items() if v is not None})
+    return state
+
+
+def email_questions(categories: Optional[Dict[str, str]] = None) -> Dict:
+    """Standard pre-built questions for email triage."""
+    categories = categories or {
+        "billing": "invoices, payments, refunds",
+        "technical": "bugs, outages, integrations",
+        "sales": "pricing, demos, new purchases",
+        "security": "phishing, scams, account compromise",
+        "hr": "hiring, leave, payroll",
+        "other": "none of the above",
+    }
+    return {
+        "category": {
+            "type": "choice",
+            "instructions": "Which team should handle the email in `body`?",
+            "criteria": categories,
+        },
+        "is_spam": {
+            "type": "noul",
+            "instructions": "Is this email unsolicited spam or bulk marketing?",
+        },
+        "is_phishing": {
+            "type": "noul",
+            "instructions": "Is this email a phishing or scam attempt to steal money, credentials, or personal data?",
+            "criteria": {"true": "phishing, scam, or fraud", "false": "a legitimate email"},
+        },
+        "urgency": {
+            "type": "score",
+            "instructions": "How urgent is the request in `body`?",
+            "criteria": [
+                "no time pressure",
+                "needs attention soon",
+                "blocking issue or hard deadline",
+            ],
+        },
+        "needs_reply": {
+            "type": "noul",
+            "instructions": "Does the sender expect a reply?",
+        },
+    }