LLMRouter

A tiny neural router that classifies the first user message of a coding-assistant session into one of five task types and picks a reasoning-effort level. It runs as a sidecar next to an AI gateway so the gateway can route cheap models to simple tasks and powerful models to hard ones.

openCode/Codex → gateway :4100 → router sidecar :5560 → Qwen3-0.6B → head → {task,reasoning}
                                  │
                                  └─ (fallback) gateway default_model

The router is intentionally small:

• Encoder: Qwen/Qwen3-0.6B (runs locally on MPS/CPU).
• Head: a single 8 × 1024 weight matrix — no bias, no activation.
• Output: 5 task logits (explore, plan, build, quick, vision) + 3 reasoning logits (low, medium, high).

The router only returns a task type and reasoning effort. The gateway maps that to a combo/tier (gateway-side wiring is a separate step):

pip install -e .
python -m uvicorn llmrouter.service:app --host 0.0.0.0 --port 5560

The first call downloads Qwen/Qwen3-0.6B (~1.5 GB) if it is not cached.

docker compose up --build

The container exposes :5560 and uses the committed data/head_weights.pt.

Set via environment variables or an optional YAML file (ROUTER_CONFIG):

• GET /healthz — process is alive.
• GET /readyz — encoder + head have finished warmup; returns "ready": true.
• POST /route — classify a transcript.

Example:

curl -s http://localhost:5560/route \
  -H 'Content-Type: application/json' \
  -d '{"transcript": "Search the codebase for User model references"}'

Response:

{
  "task_type": "explore",
  "reasoning_effort": "low",
  "task_probs": {"explore": 0.2, "plan": 0.2, "build": 0.2, "quick": 0.2, "vision": 0.2},
  "reason_probs": {"low": 0.5, "medium": 0.25, "high": 0.25},
  "latency_ms": 0.1
}

The placeholder weights return uniform-ish probabilities; replace them with real trained weights for useful predictions.

The committed data/head_weights.pt is an untrained placeholder. To regenerate real weights on a machine with the 0.6B encoder (MPS is strongly recommended; CPU is very slow):

python -m llmrouter.train --synthesize --epochs 200

This command:

1. Extracts real labeled sessions from ~/.local/share/opencode/opencode.db.
2. Adds synthetic samples for class balance (templates + boundary cases).
3. Encodes every sample through Qwen/Qwen3-0.6B with mean pooling.
4. Trains the 8-logit head and overwrites data/head_weights.pt.

The artifact now includes metadata (encoding_mode, head_version), and service.py will refuse to load a head trained with a different encoding mode or logit count.

ai-router/
├── src/llmrouter/
│   ├── service.py      # FastAPI sidecar
│   ├── encoder.py      # Qwen3-0.6B hidden-state extractor
│   ├── head.py         # 8-logit linear coordinator head
│   ├── task_types.py   # TaskType / ReasoningEffort enums
│   ├── train.py        # Training script
│   ├── data.py         # opencode.db extraction
│   ├── synthesize.py   # Synthetic data generation
│   └── config.py       # RouterConfig (env/yaml)
├── tests/              # pytest suite (uses stub encoder, no model download)
├── data/
│   ├── head_weights.pt # Placeholder weights (replace via retrain)
│   └── README.md       # Artifact documentation
├── Dockerfile
├── docker-compose.yml
├── pyproject.toml
└── requirements.txt

pip install -e ".[dev]"
ruff check src tests
mypy src
PYTHONPATH=src pytest tests -q

• The router is designed as a localhost/private-network sidecar. Do not expose /route to the public internet without adding authentication.
• If the router fails or is unreachable, the gateway should fall back to its configured default_model (the blog describes a circuit-breaker for this).
• Vision requests are classified from text only; the gateway should still short-circuit to the vision combo when the request actually contains image attachments.