Axon

Provider-independent model-routing research gateway for OpenAI-compatible chat requests.

Axon classifies a request into an abstract capability tier, reasoning effort, and required modalities, then uses a versioned registry to map those requirements to a concrete upstream model. Clients can request model: "auto"; provider names, model IDs, prices, context limits, and cache behaviour remain outside the classifier so they can change without retraining it.

Important

Axon v0 is a research candidate, not a production routing authority. The gateway, policy, ONNX inference path, schemas, and reproducible training pipeline are implemented and tested. The classifier has no execution-verified gold labels, the effort evaluation set is statistically inadequate, and no live-traffic pilot has been run.

Why Axon?

Most model routers either hard-code concrete model names into their labels or ask a large language model to choose another large language model. Axon instead separates two concerns:

1. Classifier: what capability does this task appear to require?
2. Gateway policy: which currently available model best satisfies that requirement under cost, latency, context, modality, cache, and availability constraints?

This makes the classifier provider-independent and keeps operational routing policy inspectable and replaceable.

OpenAI-compatible client
        |
        |  model: "auto"
        v
+------------------------- Axon --------------------------+
| request inspection                                        |
|        |                                                   |
|        v                                                   |
| MiniLM classifier -> tier + effort + modalities           |
|        |                                                   |
|        v                                                   |
| registry policy -> eligibility + cost + latency + cache   |
+-----------------------------------------------------------+
        |
        v
concrete upstream provider/model

What is implemented?

Area	Current state
OpenAI-compatible endpoint	POST /v1/chat/completions handler implemented
Automatic routing	model: "auto" classification and policy selection implemented
Explicit model routing	Registered concrete models can be requested directly
Classifier	Frozen MiniLM encoder plus trained ordinal/multilabel heads
Inference	Go request path using ONNX Runtime through cgo; no Python in the request path
Registry policy	Capability, modality, context, availability, cost, latency, cache, and stickiness
Streaming	Upstream event-stream responses are relayed and flushed
Telemetry	Privacy-safe routing decision events
Reproducibility	Versioned schemas, pinned sources, deterministic artefact digests
Standalone gateway executable	Yes — cmd/gateway (Apache-2.0)
Production model quality	Not established; there are no execution-verified gold labels

Quickstart

Prerequisites

Use only the prerequisites needed for the path you choose:

• Go 1.22+
• Python 3.12 and uv for training and Python tests
• ONNX Runtime 1.28 + pkg-config for the real Go classifier path (pkg-config --cflags --libs libonnxruntime must resolve; on macOS brew install onnxruntime pkg-config provides both)

1. Verify the gateway and policy

This path does not require Python or ONNX Runtime.

git clone https://github.com/mhingston/axon.git
cd axon

go test ./... -count=1

This validates the OpenAI-compatible handler, request inspection, registry validation, routing policy, fallback behaviour, stickiness, telemetry, and the portable classifier seams.

2. Verify the training package

cd training
uv sync --locked --all-groups
uv run --frozen pytest -q
cd ..

3. Run the real MiniLM + ONNX test path

The repository contains the v0 encoder, head, and vocabulary bundle under artifacts/bundles/. The Go classifier locates the ONNX Runtime with pkg-config — no platform-specific paths are baked into the source:

brew install onnxruntime pkg-config

pkg-config --cflags --libs libonnxruntime   # sanity check
go test -tags onnxruntime ./... -count=1

On Linux, install the ONNX Runtime shared library so that a libonnxruntime.pc is visible to pkg-config (set PKG_CONFIG_PATH if it lives outside the default search path).

This exercises the complete request path:

Go WordPiece tokeniser -> MiniLM encoder ONNX -> classifier head ONNX
-> ordinal reconstruction -> abstention contract -> gateway policy

4. Run the offline routing pilot

go run -tags onnxruntime ./cmd/pilot/

The pilot replays the challenge set through the real classifier and gateway, then writes:

• reports/pilot/phase7-static-shadow.v1.json
• reports/pilot/phase7-static-shadow.v1.md

It is deliberately offline. The transport is stubbed, the reference registry contains example endpoints, and no provider request is made.

5. Reproduce the classifier training run

cd training
uv sync --locked --all-groups
uv run --frozen --python 3.12 python -m axon_training.minilm_train \
    --config ../configs/experiments/phase6b-minilm-train.v1.json \
    --repository-root ..
cd ..

The run trains only the 4,620-parameter classifier head over the frozen MiniLM encoder and regenerates the content-addressed reports and ONNX bundle.

Running Axon as a gateway

cmd/gateway is the runnable OpenAI-compatible gateway server. It loads and validates the registry at startup, builds the MiniLM classifier from the ONNX bundle, configures server and upstream timeouts, emits JSON-lines decision telemetry, exposes /healthz, resolves per-upstream credentials, and shuts down gracefully on SIGINT/SIGTERM.

go build -tags onnxruntime -o axon-gateway ./cmd/gateway/

AXON_REGISTRY_PATH=configs/gateway/reference-registry.v1.json \
AXON_ENCODER_PATH=artifacts/bundles/minilm_classifier_v1_encoder.onnx \
AXON_HEAD_PATH=artifacts/bundles/minilm_classifier_v1_head.onnx \
AXON_VOCAB_PATH=artifacts/bundles/minilm_classifier_v1_vocab.json \
./axon-gateway   # serves on :8080 (override with AXON_LISTEN_ADDR)

Settings resolve from AXON_* environment variables with flag overrides (./axon-gateway --help lists them). The reference registry contains example endpoints; production deployments supply their own versioned registry.

Upstream credentials are resolved per host from AXON_UPSTREAM_KEY_<HOSTKEY> (uppercased hostname, non-alphanumerics replaced with _; api.openai.com → AXON_UPSTREAM_KEY_API_OPENAI_COM). The gateway never forwards a client's Authorization header upstream: it sets Authorization: Bearer <key> for configured hosts and strips the header for unconfigured ones (fail-closed).

A binary built without -tags onnxruntime is an intentional stub that prints a build-tag hint and exits 2 — it contains no classifier.

The module path is github.com/mhingston/axon, so the library surface (axon.NewMiniLMClassifier, axon.NewHandlerWithTelemetry) can also be imported directly by external Go modules.

Per-route reasoning-effort translation

The classifier selects an abstract reasoning effort (none/low/medium/high). How — or whether — that effort is communicated upstream is owned by the route, not the classifier:

• reasoning_mapping: { mode: "body_field", field: "reasoning_effort", values: { ... } } rewrites a JSON field on the forwarded request from a per-effort abstract→provider map. Abstract values not in the map are dropped to avoid leaking unsupported vocabulary upstream.
• reasoning_mapping: { mode: "none" } means effort is encoded in the model id; the gateway strips any caller-supplied abstract reasoning_effort from the forwarded body.
• No reasoning_mapping (backward-compatible default): only model is replaced.

ValidateConfig rejects mappings whose values reference abstract efforts not in supported_reasoning_efforts, and rejects mappings on routes that omit supported_reasoning_efforts. The reference registry ships one route per mode (economy-text → none, standard-general → body_field).

Request contract

Axon currently supports the OpenAI chat-completions shape at:

POST /v1/chat/completions

Automatic routing:

{
  "model": "auto",
  "messages": [
    {
      "role": "user",
      "content": "Diagnose why this distributed cache occasionally returns stale data."
    }
  ]
}

Explicit routing remains available by providing a concrete model registered in the active configuration. Axon validates its hard modality and context requirements before forwarding it.

Axon headers

Header	Purpose
X-Axon-Task-ID	Maintains route stickiness across turns in the same task
X-Axon-Task-Boundary: true	Clears the sticky route before classifying the request
X-Axon-Cache-Domain	Supplies cache-affinity information to routing policy

Axon replaces model: "auto" with the selected concrete model before forwarding the request. It preserves end-to-end request headers and relays the upstream response; deployment code remains responsible for authentication and provider-specific header policy.

Classifier contract

The classifier emits provider-independent structured data internally:

{
  "model_tier": "advanced",
  "reasoning_effort": "medium",
  "required_modalities": ["text"],
  "probabilities": {
    "model_tier": {
      "economy": 0.03,
      "standard": 0.18,
      "advanced": 0.71,
      "frontier": 0.08
    },
    "reasoning_effort": {
      "none": 0.02,
      "low": 0.19,
      "medium": 0.70,
      "high": 0.09
    },
    "required_modalities": {
      "text": 0.99,
      "image_input": 0.01,
      "audio_input": 0.01,
      "video_input": 0.01,
      "image_output": 0.01,
      "audio_output": 0.01
    }
  },
  "confidence": 0.70,
  "abstain": false
}

Labels are ordinal:

economy < standard < advanced < frontier
none < low < medium < high

The gateway converts low-confidence, invalid, or abstaining classifications into its configured safe fallback. Concrete provider and model names never appear in the classifier output.

Current input representation

Before classification the gateway packs the chat payload into a deterministic, bounded, role-preserving envelope:

[TASK]
<first user message>

[STATE]
turns=N user=.. assistant=.. tool=.. system=.. tool_calls=..

[RECENT]
<last 3 messages, chronological, "role: content">

[EVIDENCE]
<most recent tool/function message; the section is omitted for plain chats>

This keeps routing signals that live in the original task, tool errors, and verifier feedback visible to the classifier. Each message is truncated at 2,000 runes and identical payloads always produce identical envelopes. It remains a bounded task representation, not a full conversation-memory or trajectory encoder.

Post-processing and abstention contract

Confidence and abstention follow one versioned contract — configs/classifier/postprocessing.v1.json — shared by the Python evaluation harness, the Go gateway, and the model card. The classifier abstains when any of: tier confidence < 0.45, effort confidence < 0.45, top-two margin < 0.05, normalised entropy > 0.95, or input truncation ratio > 0.5. The Go gateway currently applies the contract to raw head probabilities; temperature scaling is fit inside the Python harness and is not yet exported into the v1 ONNX bundle (see docs/model-card-v0.md).

Routing policy

For model: "auto", Axon:

1. inspects request size and modalities;
2. obtains a tier, effort, modality set, probabilities, and confidence;
3. falls back safely on classifier error, invalid output, abstention, or low confidence;
4. filters routes by availability, modality, context, tier, and effort;
5. scores eligible routes using configured cost, latency, switch, and cache penalties;
6. applies task stickiness and hysteresis;
7. forwards the request with the selected concrete model;
8. emits a decision telemetry event.

The example registry is configs/gateway/reference-registry.v1.json. Its provider names, model IDs, prices, and URLs are illustrative and must not be used as production configuration.

Model architecture

The v0 classifier is deliberately small:

• Encoder: sentence-transformers/all-MiniLM-L6-v2, pinned revision, 22,713,216 parameters, frozen
• Embedding: 384-dimensional mean-pooled representation
• Tier head: three cumulative binary thresholds for four ordered tiers
• Effort head: three cumulative binary thresholds for four ordered effort levels
• Modality heads: six independent sigmoid outputs
• Trainable parameters: 4,620
• Total parameters: 22,717,836
• Export: separate encoder and head ONNX files, opset 17; the exported vocabulary pins max_length: 256 (the training sequence length, not the tokenizer's 512 default) so Go inference matches training

Go performs BERT WordPiece tokenisation and runs both ONNX sessions through the ONNX Runtime C API.

Evidence and limitations

The software pipeline is substantially more mature than the trained model. Keep those two claims separate.

What the current evidence supports

• The schemas, deterministic dataset materialisation, training run, ONNX export, Go inference path, routing policy, and telemetry can be reproduced.
• The Phase 6b test report contains 93 tier-labelled test rows, with tier accuracy 0.763 and macro-F1 0.381.
• The same report contains only three effort-labelled test rows, so its effort result is not statistically useful.
• The offline pilot demonstrates that requests can traverse the real classifier and policy without reaching a provider.

What the current evidence does not support

• There are no execution-verified gold routing labels.
• Model quality has not been validated on live SWE or agent trajectories.
• No claim can yet be made about production cost savings or success retention.
• The current static pilot contains 16 inputs and selects the safe fallback for 14 of them (the other 2 route to economy-text); it validates wiring rather than useful route discrimination. See reports/pilot/phase7-static-shadow.v1.md for the per-decision metrics, provenance, and per-stage timings.
• The frozen encoder has not been fine-tuned for routing.
• Platform-portable native packaging is not yet provided.

See docs/model-card-v0.md for the complete evidence, limitations, and promotion criteria.

Repository layout

.
├── proxy.go                         # OpenAI-compatible HTTP handler + envelope
├── gateway_policy.go                # eligibility, scoring, fallback, stickiness
├── postprocessing.go                # versioned abstention contract (Go)
├── minilm_classifier.go             # MiniLM + ONNX Runtime classifier
├── ordinal_bundle_classifier.go     # portable JSON classifier fallback
├── configs/
│   ├── gateway/                     # reference route registry
│   ├── classifier/                  # postprocessing.v1 abstention contract
│   ├── datasets/                    # pinned dataset releases
│   └── experiments/                 # versioned experiment configs
├── schemas/v1/                      # JSON Schema contracts
├── training/
│   ├── src/axon_training/           # materialisation, baselines, training, export
│   └── tests/                       # Python test suite
├── artifacts/                       # content-addressed runs and model bundles
├── reports/                         # dataset, experiment, and pilot evidence
├── docs/                            # design handoff and model card
├── testdata/                        # contracts, fixtures, and challenge inputs
├── cmd/gateway/                     # runnable gateway server
└── cmd/pilot/                       # offline static-shadow executable

Development checks

# Go formatting, tests, race detector, and vet
gofmt -l .
go test ./... -count=1
go test -race ./... -count=1
go vet ./...

bash scripts/check-doc-consistency.sh

# Python formatting, linting, and tests
cd training
uv sync --locked --all-groups
uv run --frozen ruff check .
uv run --frozen ruff format --check .
uv run --frozen pytest -q

The ONNX-tagged tests additionally require the native setup shown in the quickstart.

Recommended next milestones

The v0 release is intentionally frozen at a wiring-test state: every layer of the pipeline runs, but no claim is made about production routing quality. The substance of v1 is data, not code:

1. Execution-verified data campaign (top priority). Build the task × model tier × reasoning effort → success, quality, cost, latency, attempts execution matrix over 300–500 executable SWE/orchestration tasks. Current tier test support is 93 and effort support is 3 — no further model work is justified until this exists.
2. Create held-out repository and trajectory splits with meaningful support for tier, effort, and modality heads.
3. Export the fitted calibration temperatures into the ONNX bundle so gateway abstention decisions are bit-aligned with the offline harness.
4. Run dynamic shadow evaluation on real coding-agent trajectories.
5. Add sticky-route eviction/TTL, reproducible benchmark manifests, and container packaging for the ONNX runtime — only after the data campaign shows the classifier is worth the operational hardening.
6. Promote only after success-retention, under-routing, calibration, cost, and latency gates are met on the v1 data.

Design and evidence

• Normative implementation handoff
• v0 model card
• Source dataset audit
• Phase 6b training report
• Phase 7 offline pilot

Licence

Apache-2.0 — see LICENSE. The MiniLM backbone (sentence-transformers/all-MiniLM-L6-v2) is Apache-2.0; Axon's own code, routing policy, and trained heads are licensed under the same terms. The synthetic teacher data from SupraLabs/Prompt-Routing-Dataset is dataset-specific; see the source audit in reports/source-audit/.