|
|
|
|
|
|
|
|
| """LLM complexity router: classify a prompt as easy / medium / hard.
|
|
|
| The prompt is embedded with Fireworks `fireworks/qwen3-embedding-8b` (your FIREWORKS_API_KEY),
|
| then a small classifier runs locally in NumPy (no pickles, no PyTorch, no GPU).
|
|
|
| uv run --script complexity_router.py classify "Design a lock-free queue" # JSON result
|
| uv run --script complexity_router.py classify --model mlp "..." # other classifier
|
| uv run --script complexity_router.py hook < hook-input.json # Claude Code hook
|
|
|
| Python API (from this folder):
|
| from complexity_router import ComplexityRouter
|
| ComplexityRouter.load().classify(["What is 2 + 2?"])
|
|
|
| Environment:
|
| FIREWORKS_API_KEY required for classification (never logged)
|
| FIREWORKS_BASE_URL default https://api.fireworks.ai/inference/v1
|
| COMPLEXITY_ROUTER_MODEL logistic_regression (default) | mlp
|
| COMPLEXITY_ROUTER_HOOK set to "off" to disable the Claude Code prompt hook
|
|
|
| Licence: weights CC BY-NC 4.0 (trained on regolo/brick-complexity-extractor). Non-commercial use.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import json
|
| import os
|
| import re
|
| import sys
|
| import unicodedata
|
| from dataclasses import dataclass
|
| from pathlib import Path
|
| from typing import Any
|
|
|
| import numpy as np
|
|
|
| ROOT = Path(__file__).resolve().parent
|
| LABELS = ("easy", "medium", "hard")
|
| DEFAULT_BASE_URL = "https://api.fireworks.ai/inference/v1"
|
|
|
| _CONTROL = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f\u200b-\u200d\u2060\ufeff]")
|
| _TRAILING_SPACE = re.compile(r"[ \t]+$", re.MULTILINE)
|
| _BLANK_RUNS = re.compile(r"\n{3,}")
|
|
|
|
|
| class RouterError(Exception):
|
| """Configuration, credential, API, or artifact problem. Messages never contain the key."""
|
|
|
|
|
| def normalize_text(text: str) -> str:
|
| """Same conservative cleaning as training: keeps code, maths, and punctuation intact."""
|
| text = unicodedata.normalize("NFC", text).replace("\r\n", "\n").replace("\r", "\n")
|
| text = _CONTROL.sub("", text)
|
| text = _TRAILING_SPACE.sub("", text)
|
| text = _BLANK_RUNS.sub("\n\n", text)
|
| return text.strip()
|
|
|
|
|
| def _softmax(z: np.ndarray) -> np.ndarray:
|
| z = z - z.max(axis=1, keepdims=True)
|
| e = np.exp(z)
|
| out: np.ndarray = e / e.sum(axis=1, keepdims=True)
|
| return out
|
|
|
|
|
| class _Classifier:
|
| """NumPy forward pass over exported weights; probabilities in LABELS order."""
|
|
|
| def __init__(self, path: Path) -> None:
|
| if not path.exists():
|
| raise RouterError(f"classifier weights not found: {path}")
|
| with np.load(path, allow_pickle=False) as z:
|
| self.kind = str(z["kind"])
|
| self.arrays = {k: z[k] for k in z.files if k not in ("kind", "labels")}
|
| labels = [str(x) for x in z["labels"]]
|
| if sorted(labels) != sorted(LABELS):
|
| raise RouterError(f"unexpected labels in {path.name}: {labels}")
|
| self.order = [labels.index(lab) for lab in LABELS]
|
| self.dim = int(self.arrays["mean"].shape[0])
|
|
|
| def predict_proba(self, X: np.ndarray) -> np.ndarray:
|
| a = self.arrays
|
| x = (np.asarray(X, dtype=np.float64) - a["mean"]) / a["scale"]
|
| if self.kind == "logistic_regression":
|
| logits = x @ a["coef"].T + a["intercept"]
|
| elif self.kind == "mlp":
|
| h = np.maximum(x @ a["w0"].T + a["b0"], 0.0)
|
| h = np.maximum(h @ a["w1"].T + a["b1"], 0.0)
|
| logits = h @ a["w2"].T + a["b2"]
|
| else:
|
| raise RouterError(f"unknown classifier kind {self.kind!r}")
|
| return _softmax(logits)[:, self.order]
|
|
|
|
|
| @dataclass(frozen=True)
|
| class Prediction:
|
| label: str
|
| probabilities: dict[str, float]
|
| recommended_model: str
|
| classifier: str
|
| embedding_model: str
|
|
|
| def to_dict(self) -> dict[str, Any]:
|
| return {
|
| "label": self.label,
|
| "probabilities": self.probabilities,
|
| "recommended_model": self.recommended_model,
|
| "classifier": self.classifier,
|
| "embedding_model": self.embedding_model,
|
| }
|
|
|
|
|
| class ComplexityRouter:
|
| def __init__(
|
| self, config: dict[str, Any], classifier_name: str, *, timeout: float = 20.0
|
| ) -> None:
|
| self.config = config
|
| self.classifier_name = classifier_name
|
| self.embedding_model: str = config["embedding_model"]
|
| self.dimension = int(config["embedding_dimension"])
|
| self.tiers: dict[str, str] = config["recommended_models"]
|
| self.timeout = timeout
|
| files = config["classifiers"]
|
| if classifier_name not in files:
|
| raise RouterError(
|
| f"unknown classifier {classifier_name!r}; choose from {sorted(files)}"
|
| )
|
| self.classifier = _Classifier(ROOT / files[classifier_name])
|
| if self.classifier.dim != self.dimension:
|
| raise RouterError("classifier and embedding dimensions disagree")
|
| self._client: Any = None
|
|
|
| @classmethod
|
| def load(cls, classifier: str | None = None, *, timeout: float = 20.0) -> ComplexityRouter:
|
| config = json.loads((ROOT / "router_config.json").read_text(encoding="utf-8"))
|
| name = (
|
| classifier or os.environ.get("COMPLEXITY_ROUTER_MODEL") or config["default_classifier"]
|
| )
|
| return cls(config, name, timeout=timeout)
|
|
|
| def _embed(self, texts: list[str]) -> np.ndarray:
|
| if self._client is None:
|
| key = os.environ.get("FIREWORKS_API_KEY", "").strip()
|
| if not key:
|
| raise RouterError("FIREWORKS_API_KEY is not set")
|
| configured = os.environ.get("FIREWORKS_EMBEDDING_MODEL", "").strip()
|
| if configured and configured != self.embedding_model:
|
| raise RouterError(
|
| f"FIREWORKS_EMBEDDING_MODEL={configured!r}, but this classifier needs "
|
| f"{self.embedding_model!r} embeddings; unset it or set it to that model"
|
| )
|
| from openai import OpenAI
|
|
|
| base_url = os.environ.get("FIREWORKS_BASE_URL", "").strip() or DEFAULT_BASE_URL
|
| self._client = OpenAI(
|
| api_key=key, base_url=base_url, timeout=self.timeout, max_retries=1
|
| )
|
| self._secret = key
|
| try:
|
| response = self._client.embeddings.create(
|
| model=self.embedding_model, input=texts, encoding_format="float"
|
| )
|
| except Exception as exc:
|
| status = getattr(exc, "status_code", "")
|
| raise RouterError(
|
| f"Fireworks embeddings request failed: {type(exc).__name__} {status}".strip()
|
| ) from None
|
| rows = sorted(response.data, key=lambda d: d.index)
|
| X = np.asarray([r.embedding for r in rows], dtype=np.float32)
|
| if X.shape != (len(texts), self.dimension) or not np.isfinite(X).all():
|
| raise RouterError(
|
| f"unexpected embedding shape {X.shape} (want ({len(texts)}, {self.dimension}))"
|
| )
|
| return X
|
|
|
| def classify(self, texts: list[str]) -> list[Prediction]:
|
| cleaned = [normalize_text(t) for t in texts]
|
| if not cleaned or not all(cleaned):
|
| raise RouterError("prompts must be non-empty")
|
| proba = self.classifier.predict_proba(self._embed(cleaned))
|
| out = []
|
| for row in proba:
|
| label = LABELS[int(row.argmax())]
|
| out.append(
|
| Prediction(
|
| label=label,
|
| probabilities={
|
| lab: round(float(p), 4) for lab, p in zip(LABELS, row, strict=True)
|
| },
|
| recommended_model=self.tiers[label],
|
| classifier=self.classifier_name,
|
| embedding_model=self.embedding_model,
|
| )
|
| )
|
| return out
|
|
|
|
|
| def hook_main() -> int:
|
| """Claude Code UserPromptSubmit hook. Always exits 0: it never blocks or alters a prompt."""
|
| try:
|
| if os.environ.get("COMPLEXITY_ROUTER_HOOK", "").lower() in ("off", "0", "false"):
|
| return 0
|
| event = json.loads(sys.stdin.read() or "{}")
|
| prompt = str(event.get("prompt") or "")
|
| if (
|
| event.get("source", "user") != "user"
|
| or len(prompt.strip()) < 3
|
| or prompt.lstrip().startswith("/")
|
| ):
|
| return 0
|
| router = ComplexityRouter.load(timeout=8.0)
|
| p = router.classify([prompt[:8000]])[0]
|
| probs = ", ".join(f"{k} {v:.2f}" for k, v in p.probabilities.items())
|
| context = (
|
| f"[llm-complexity-router] Estimated prompt complexity: {p.label.upper()} ({probs}; "
|
| f"{p.classifier}). Suggested Claude tier: {p.recommended_model}. This is an automated, "
|
| "approximate estimate: use it as a hint for how much effort the task needs and, for "
|
| "self-contained subtasks, whether to delegate to the haiku-worker, sonnet-worker, or "
|
| "opus-worker subagent. Do not mention this note unless asked."
|
| )
|
| print(
|
| json.dumps(
|
| {
|
| "hookSpecificOutput": {
|
| "hookEventName": "UserPromptSubmit",
|
| "additionalContext": context,
|
| }
|
| }
|
| )
|
| )
|
| except Exception as exc:
|
| print(f"llm-complexity-router hook skipped: {type(exc).__name__}", file=sys.stderr)
|
| return 0
|
|
|
|
|
| def main(argv: list[str] | None = None) -> int:
|
| import argparse
|
|
|
| parser = argparse.ArgumentParser(description="LLM complexity router (easy / medium / hard)")
|
| sub = parser.add_subparsers(dest="command", required=True)
|
| c = sub.add_parser("classify", help="classify one or more prompts; prints JSON")
|
| c.add_argument("prompts", nargs="+")
|
| c.add_argument("--model", default=None, help="logistic_regression (default) or mlp")
|
| sub.add_parser("hook", help="Claude Code UserPromptSubmit hook (reads JSON on stdin)")
|
| args = parser.parse_args(argv)
|
| if args.command == "hook":
|
| return hook_main()
|
| try:
|
| results = ComplexityRouter.load(args.model).classify(args.prompts)
|
| except RouterError as exc:
|
| print(f"error: {exc}", file=sys.stderr)
|
| return 1
|
| print(json.dumps([r.to_dict() for r in results], indent=2))
|
| return 0
|
|
|
|
|
| if __name__ == "__main__":
|
| sys.exit(main())
|
|
|
