Skip to content

raulmn00/agent-router

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

114 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

agent-router

Multi-agent system whose routing layer is a DistilBERT classifier I fine-tuned, instead of paying tokens to a frontier LLM to route each request.

Live demo:

URL
Frontend (React + TS, Vercel) https://agent-router-five.vercel.app
Backend API (FastAPI, Cloud Run) https://agent-router-909428365094.us-central1.run.app
Model release tarball https://github.com/raulmn00/agent-router/releases/tag/v0.1.0

The frontend has two modes: Roteamento (calls POST /route and renders the dispatched intent + agent trace) and Comparação (calls POST /compare and renders DistilBERT vs LLM zero-shot vs Embeddings+LogReg side by side with measured latency and cost). The DistilBERT model is fetched on cold start by the Cloud Run container from the GitHub Release above.

Showcases, in one repo:

  • Own fine-tuning of distilbert-base-uncased on a synthetic dataset for 4 intent classes (simple_qa, complex_task, document_qa, chitchat).
  • Multi-agent orchestration (Planner / Executor / Critic) built by hand on top of a pluggable LLMProvider (OpenAI / Anthropic / Fake).
  • FastAPI backend with two endpoints: POST /route (one input → dispatched answer + agent trace) and POST /compare (one input → all three routers side-by-side with measured latency, confidence, and cost). Rate-limited via slowapi.
  • React + TypeScript frontend that consumes both endpoints with a side-by-side comparison view whose column reveal is timed by the measured latency.
  • Comparative evaluation of three router strategies: the fine-tuned model, an LLM zero-shot baseline, and an embeddings + LogisticRegression baseline — measured for accuracy, F1, latency, and cost on a held-out testset.

Architecture

                       ┌─────────────────────────────┐
                       │      POST /route            │
                       │   { "input": "user text" }  │
                       └──────────────┬──────────────┘
                                      │
                                      ▼
                       ┌─────────────────────────────┐
                       │  DistilBERT intent router   │   <— fine-tuned
                       │   classify() -> intent      │       (this repo)
                       └──────────────┬──────────────┘
                                      │
        ┌─────────────────┬───────────┴──────────┬───────────────────┐
        ▼                 ▼                      ▼                   ▼
  simple_qa          complex_task           document_qa         chitchat
        │                 │                      │                   │
        ▼                 ▼                      ▼                   ▼
  1 direct call    Orchestrator              RAG stub            1 direct call
  to OpenAI        ┌─────────────┐           (integration        to OpenAI
  (gpt-4o-mini)    │  Planner    │            point for          (short max_tokens,
                   │     ↓       │            retriever)          warm/social tone)
                   │  Executors  │
                   │     ↓       │
                   │   Critic    │ — re-plans once if rejected
                   └─────────────┘
        │                 │                      │                   │
        └─────────────────┴──────────┬───────────┴───────────────────┘
                                     ▼
                       ┌─────────────────────────────┐
                       │      RouteResponse          │
                       │  {intent, confidence,       │
                       │   answer, path_taken,       │
                       │   trace}                    │
                       └─────────────────────────────┘

Why these choices

Choice Why Trade-off accepted
DistilBERT local instead of LLM-as-router Routing is on the hot path of every request; paying an LLM call just to pick a branch is wasteful. Local inference removes that hop entirely once the model is trained. Need to train + version + ship a 256 MB binary.
Hand-rolled multi-agent loop (no LangChain / CrewAI / etc.) The whole point of the demo is to show I understand what an agent loop is and what Planner → Executor → Critic → retry looks like in plain code. Frameworks hide exactly the part the reader needs to see. Less abstraction reuse if I later want fan-out / async tools / state machines.
LLMProvider ABC with OpenAI / Anthropic / Fake impls Tests run with FakeProvider (no network, no spend). Production can pick the provider via env var. Both real providers go through the same call surface. One extra layer to maintain.
Synthetic dataset of 600 templated examples for fine-tuning A held-out, handcrafted testset of 40 lives in eval/ and is what gets reported. The synthetic 600 is just enough signal for DistilBERT to separate the 4 classes; bigger gains would come from real labelled data. F1 on the template-distribution test split is artificially high (the test split shares templates with training). The honest number is the one in the evaluation table below, against the handcrafted held-out set.
Pure-function metrics (accuracy, f1_macro, cost_per_1k) Testable without network or training. Pricing is a documented constant — comparison stays auditable. Cost is a function of declared token assumptions, not measured token counts. Replace with measured counts if you want a precise number.

How to run

Prerequisites

  • Python 3.12+ (tested on 3.13 too — the pyproject pin is conservative)
  • OPENAI_API_KEY in .env (copy from .env.example) — required for the backend and for the eval baselines
  • ML stack: pip install -r backend/requirements.txt (or install the per-package extras you need)

1. Tests (no network, no spend)

python -m pytest

2. Train the DistilBERT router

cd router
python -m router.train

Saves model + tokenizer to router/model/. Device auto-detected (CUDA → MPS → CPU). The training dataset is committed at router/data/intents.jsonl (regenerate with python -m router.dataset if you change the templates).

3. Fit the embeddings LogReg (one-time)

cd eval
python -m eval.fit_embed_router

Writes eval/models/embed_router.joblib (~25 KB, committed in this repo). Required by POST /compare's embed-router row; if missing, that row just carries an error and the other two routers still respond.

4. Run the backend

cd backend
uvicorn app.api:app --reload --port 8000

Endpoints: GET / (health), POST /route (single dispatch, 30 req/min/IP), POST /compare (three-way comparison, 10 req/min/IP — burns real tokens, hence the tighter limit). CORS is open by default; set CORS_ALLOW_ORIGINS env var to restrict.

5. Frontend (optional)

cd frontend
cp .env.example .env       # default VITE_API_URL=http://localhost:8000
npm install
npm run dev                # http://localhost:5173

6. Compare the three routers offline (no backend)

cd eval
python -m eval.compare_routers --runs 3

Generates eval/results/comparison.md and comparison.csv. The --runs N flag aggregates N independent runs, clearing the embedding cache between them so the embed router stays cold-cache (production-like) for every measurement. Reuses the same code that powers POST /compare.

Evaluation & honest analysis

Two views into the model: the held-out test split (a perfect score that tells you less than it looks) and a generalization probe against inputs from outside the training distribution (much more informative).

1. Held-out test split — 120 examples, 30 per class

Trained on Google Colab (CUDA, fp16, distilbert-base-uncased, 4 epochs, batch 16, lr 5e-5, weight_decay 0.01, max_length 64) over the 480/120 train/test split from router/data/intents.jsonl. Every row of the confusion matrix lands on the diagonal — zero misclassifications.

Confusion matrix

              precision    recall  f1-score   support

   simple_qa     1.0000    1.0000    1.0000        30
complex_task     1.0000    1.0000    1.0000        30
 document_qa     1.0000    1.0000    1.0000        30
    chitchat     1.0000    1.0000    1.0000        30

    accuracy                         1.0000       120
   macro avg     1.0000    1.0000    1.0000       120
weighted avg     1.0000    1.0000    1.0000       120

Versioned evidence in router/results/: confusion_matrix.png · confusion_matrix.csv · classification_report.json · classification_report.txt · training_meta.json · generalization_test.txt.

Reading the 100% honestly. This is not a flex. The training set is synthetic — every example was generated from the template bank in router/router/dataset.py, and the test split is sampled from the same template distribution as training (disjoint sentences, but shared per-class markers). With four classes and consistent surface cues, a fine-tuned DistilBERT separates them trivially. The perfect score says the task is easily separable on this dataset, not that the model is exceptional. The number that actually matters is what happens off-distribution — section 2 below.

2. Generalization probe — inputs from outside the training distribution

I wrote 20 fresh sentences (different topics and phrasings, never seen at training time) and ran IntentClassifier.classify() on each. Full output committed at router/results/generalization_test.txt.

16 sentences with obvious class markers (different from the training templates, but a human would still know the intent):

simple_qa     0.857  What's the capital of Australia?
simple_qa     0.847  How many milliliters are in a cup?
simple_qa     0.849  Who wrote the novel Dom Casmurro?
simple_qa     0.851  What year did the Berlin Wall fall?
complex_task  0.939  Design a scalable architecture for a food delivery app...
complex_task  0.825  Plan a 7-day trip to Japan...
complex_task  0.910  Help me migrate a monolith to microservices step by step.
complex_task  0.930  Create a go-to-market strategy for a B2B SaaS...
document_qa   0.871  According to the attached contract...
document_qa   0.879  In the PDF I uploaded, which section covers the refund policy?
document_qa   0.813  Summarize the key findings from this research paper.
document_qa   0.880  Based on the document, what are the eligibility requirements?
chitchat      0.542  Hey, how's it going today?
chitchat      0.560  Good morning! Hope you're having a nice day.
chitchat      0.518  Haha that's pretty funny.
chitchat      0.562  Thanks so much, you've been really helpful!

All 16 classified correctly. Confidence sits around 0.81–0.94 for the first three intents — and notably 0.52–0.56 for chitchat, meaningfully lower even when the answer is right.

4 ambiguous sentences crafted to remove obvious markers:

chitchat      0.385  Tell me about the requirements
chitchat      0.510  Build me something cool
simple_qa     0.588  What does it say about pricing?
chitchat      0.417  Can you explain how this works?

These are deliberately underspecified — without surrounding context a human reader could argue for two or three classes for each. Confidence drops to 0.39–0.59.

What this actually tells me:

  • Calibration is decent. The model doesn't crank every prediction up to 0.99. Confidence visibly drops on inputs it should be unsure about. That's the property that matters for a routing layer — a route picked at 0.42 should be treated very differently from one picked at 0.92.

  • chitchat is doing double duty. It's both a legitimate intent ("hey, thanks!") AND the soft fallback the model lands on when nothing else fits — 3 out of 4 ambiguous probes ended up there with low confidence. The boundary between real chitchat and I don't really know is the fuzziest one in the model, which is also why even legitimate chitchat tops out around 0.56. This isn't a model bug — it's a property of the class distribution and the reason the threshold is per-class (below), not global.

  • Per-class confidence thresholds are empirically justified — and now implemented. The first cut used a single global threshold of 0.65, and the production smoke test (committed at scripts/results/prod_routing_report.md) caught the predictable consequence: a legitimate "Good morning! Hope you're having a nice day." scored 0.560 and got demoted to the fallback — a false positive driven by chitchat's naturally-lower confidence band. The dispatcher now uses a different threshold per intent:

    Class Threshold Why
    simple_qa 0.65 Clear inputs sit at 0.81–0.86; 0.65 catches genuinely ambiguous ones (0.39–0.59).
    complex_task 0.65 Clear inputs 0.82–0.94.
    document_qa 0.65 Clear inputs 0.81–0.88.
    chitchat 0.45 Legitimate chitchat tops out around 0.52–0.56; a 0.65 floor was demoting real conversation. 0.45 still catches ambiguous-classified-as-chitchat (which sit at 0.39–0.51).

    Resolution order (later overrides earlier):

    1. hardcoded defaults (above)
    2. legacy CONFIDENCE_THRESHOLD scalar env var → applied uniformly to every class
    3. CONFIDENCE_THRESHOLDS env var → JSON object overriding specific classes, e.g.
      CONFIDENCE_THRESHOLDS='{"chitchat": 0.50, "simple_qa": 0.70}'
      

    Malformed JSON, unknown intent keys, non-numeric values → logged at WARNING level and ignored; the service never crashes on a misconfigured env var. Predictions below the threshold for their own class skip the four dispatch arms and land on path_taken = "low_confidence_fallback" with no LLM call. See backend/app/dispatch.py (Dispatcher._low_confidence_fallback and get_thresholds()).

3. Latency: parallel orchestrator (v0.2.3)

The interesting latency story is on the complex_task path, where the orchestrator fans out N independent subtasks generated by the Planner. v0.2.2 ran them sequentially — and the production smoke test (committed at scripts/results/prod_routing_report.md) caught the cost: a 6-subtask migration prompt took 31.7 s warm and a 5-subtask architecture prompt timed out at 60 s (cold start + sequential 5×). Each Executor is stateless, so the subtasks can fan out.

v0.2.3 parallelizes within an attempt via asyncio.gather bounded by asyncio.Semaphore(MAX_CONCURRENT_EXECUTORS) (default 5, env-configurable). The wrapping is done at Executor.execute_async, which delegates to asyncio.to_thread(self.execute, …) — see the comments in agents/agents/executor.py for why asyncio.to_thread over AsyncOpenAI / AsyncAnthropic (short version: keep the sync LLMProvider ABC; the GIL releases on the SDK's underlying socket I/O so threads actually overlap). Failure isolation via gather(return_exceptions=True): one subtask raising leaves an inline [execution failed: <Type>: <msg>] marker in the answer and an EXEC[i] FAILED … line in the trace — the orchestrator never tears down a whole attempt because of a single failure.

Same prompts measured against production on v0.2.3 (full report in scripts/results/prod_routing_report.md):

Prompt Subtasks v0.2.2 (sequential) v0.2.3 (parallel)
"Plan and execute a database migration from MySQL to Postgres for a 2TB user table with zero downtime." 6 — (new measurement) 11.9 s warm
"Help me migrate a monolith to microservices step by step." 6 31.7 s warm 9.2 s warm (~3.4× faster)
"Design a scalable architecture for a food delivery app and outline the main trade-offs." 5 timed out @ 60 s 26.4 s cold-start (~7.9 s warm in earlier measurements)
"What year did the Berlin Wall fall?" (simple_qa, baseline) n/a 627 ms 667 ms (no regression)
"According to the attached contract…" (document_qa, RAG stub) n/a 210 ms 186 ms (no regression)

The trace includes a single summary line that makes the parallelism legible:

executed 6 subtask(s) concurrently (max_concurrent=5, failures=0)

The same smoke-test script that flagged the v0.2.2 regression on chitchat now reports 7/7 clear PASS (was 5/6 FAIL — the "Good morning! …" case is fixed by the per-class threshold above) and 4/4 ambiguous PASS (fallback or conf < 0.65).

The Orchestrator.run_task(task) public API is preserved (sync signature) — it internally calls asyncio.run(run_task_async(task)), so the existing FastAPI dispatcher (running in Starlette's threadpool) keeps working without any changes. An async-native caller can await orchestrator.run_task_async(task) directly to skip the wrapper.

4. Comparison against baselines — measured

Ran python -m eval.compare_routers --runs 3 against the 40-example held-out testset in eval/data/routing_testset.jsonl. Three cold-cache runs (the script drops eval/results/.embed_cache.json between runs so the embed router pays a fresh API call each time). Generated artifacts committed: eval/results/comparison.md, eval/results/comparison.csv.

Approach Accuracy F1 (macro) Mean latency (ms) ± stdev Cost / 1k (USD) N
DistilBERT (fine-tuned) 0.975 0.975 25.0 ± 4.4 $0.0000 40
LLM zero-shot (gpt-4o-mini) 1.000 1.000 795.7 ± 68.6 $0.0120 40
Embeddings + LogReg 0.975 0.975 780.7 ± 574.9 $0.0003 40

Cost values come from the documented assumptions in eval/eval/metrics.py (PRICING_USD_PER_M_TOKENS, TOKEN_ASSUMPTIONS), not from measured per-call token counts.

Reading the table honestly. Accuracies sit within 2.5 percentage points of each other — the LLM is the only one to reach 40/40, DistilBERT and the embed router each miss 1/40 (likely the same hard case). On this dataset the accuracy column doesn't separate the approaches. The decision is latency and cost:

  • DistilBERT is ~32× faster than either API-based option (25 ms vs ~780-796 ms) and pays no per-request cost. No network in the path.
  • The LLM zero-shot wins accuracy by one case at 40× the cost of the embed router and ~33,000× the cost of DistilBERT.
  • The embed router's latency standard deviation (574.9 ms over 3 runs) is notably high — embedding-API tail latency was much noisier than the chat-completions API in this run. The mean of 780.7 ms isn't representative of a "typical" call.

The 1-of-40 DistilBERT miss is not a model failure — it's "Build me something cool", an input deliberately included in the testset as ambiguous (no class markers). The model lands it on chitchat at confidence 0.510 (see router/results/generalization_test.txt); under the per-class threshold of 0.45 for chitchat, the live dispatcher routes it as chitchat. The eval harness measures raw classifier accuracy without the threshold gate, so it counts as a miss.

When each approach makes sense

The numbers above show three routers within a couple of percentage points of each other on accuracy. That's a property of this dataset, not a property of the approaches in general — and it's worth taking seriously when reading the conclusion:

  • DistilBERT (fine-tuned) wins on this dataset primarily on operational metrics: ~32× lower latency than the API options, zero per-call cost. Accuracy is in a statistical tie with the embed router (39/40 each). It wins when the set of intents is closed and stable, you have (or can synthesize) training data, and the router is on the hot path. The trade-off is that you carry a trained model artifact and have to retrain when intents change.

  • LLM zero-shot is the only approach that scored 40/40, but on a dataset this easily separable that's a one-case difference. It wins when intents change every week, you don't yet have labelled data, or traffic is low enough that token cost doesn't matter (it's still 40× more expensive than the embed router and ~33,000× more than DistilBERT per request). The accuracy advantage of a frontier LLM should be more visible on a harder distribution — see the generalization probe in section 2: ambiguous inputs sit at 0.39-0.59 confidence, which is exactly the regime where a stronger zero-shot model could pull ahead.

  • Embeddings + LogReg sits between the two operationally (cheaper than the LLM, slower than DistilBERT, with high latency variance from the embedding API itself). Adding a class is fast — re-fit the LogReg in seconds. The accuracy is tied with DistilBERT here; whether it stays competitive in a harder setting depends on the embedding model's discriminative power for that specific intent set.

Honest summary for this dataset: the synthetic training distribution is highly separable, so the accuracy column doesn't decide. The latency and cost columns do — and they point clearly at the fine-tuned local model. The accuracy advantage a frontier LLM could offer would only become legible on a dataset with more genuinely ambiguous inputs.

Project layout

agent-router/
├── router/         # DistilBERT intent classifier — train + classify
├── agents/         # Planner / Executor / Critic / Orchestrator + LLMProvider
├── backend/        # FastAPI: /route + /compare + slowapi rate limiting
├── eval/           # Comparative evaluation (DistilBERT vs LLM vs embed)
├── frontend/       # React + TS demo UI (Vite, hand-written CSS)
├── pyproject.toml  # Workspace root (pytest config)
└── .env.example    # API keys + optional overrides

Each subpackage has its own pyproject.toml (or package.json) and README.md:

Deployment (Google Cloud Run)

The model artifact is 256 MB and is not committed to git. The Docker image fetches it at container start from a configured URL (e.g. a GitHub Release tarball). See backend/Dockerfile and backend/entrypoint.sh for the wiring, and backend/README.md for the exact gcloud run deploy command.

Key flags:

  • --memory 2Gi (model load + working set + the parallel orchestrator's thread pool fit comfortably here)
  • --cpu 2 (DistilBERT inference is single-threaded; 2 vCPUs leaves headroom for the LLM dispatch path's I/O)
  • --min-instances 0 is fine; cold start pays one model load (the lazy classifier in backend/app/api.py defers that until the first /route request).

Runtime env vars used in the live deploy:

Var Value Notes
MODEL_RELEASE_URL GitHub Release tarball entrypoint.sh extracts it to /app/router/model on boot.
CORS_ALLOW_ORIGINS Vercel preview + prod origins, localhost:5173 No wildcard in production.
ENABLE_API_DOCS false Hides /docs, /redoc, /openapi.json.
MAX_BODY_BYTES 10000 Hard cap before Pydantic validation.
CONFIDENCE_THRESHOLDS (unset → defaults) JSON map; chitchat=0.45, others 0.65.
MAX_CONCURRENT_EXECUTORS 5 Semaphore cap on parallel Executor calls in the orchestrator.
MODEL_VERSION_TAG colab-cuda-2026-06-06 Pure annotation; bumps trigger a new revision so the new model is re-fetched.
OPENAI_API_KEY Secret Manager Mounted via --set-secrets.

Testing

router/tests/    schema, balance, stratified split, classifier mocks,
                 train.py artifact generators                                          (20)
agents/tests/    orchestrator happy path + critic-reproves-and-retries
                 + parallel execution + semaphore cap + failure isolation              (21)
backend/tests/   /route + /compare + 413/422/429/503/500 + security headers +
                 sanitization + per-class confidence thresholds                        (40)
eval/tests/      accuracy, f1_macro, cost_per_1k pinning + sklearn cross-check         (17)

98 tests (python -m pytest); runs without GPU and without spending API credits. Heavy ML deps are skipif-gated so schema and orchestration tests still run in environments that don't have torch/transformers installed. The parallel-orchestrator tests use a thread-safe FakeProvider with a content-dispatched responder so they're deterministic regardless of which asyncio worker thread fires first. Frontend's npm run build runs tsc -b first — type-check failures break the build.

Security posture (v0.2.1)

Layer Mitigation
Cloud Run runtime Dedicated SA agent-router-runtime@… with roles/secretmanager.secretAccessor on openai-api-key only — no project-level Editor.
Container Runs as non-root app (UID 1000); /app is the only writable path.
Transport Cloud Run + Vercel enforce TLS. Backend sets HSTS (max-age=31536000), nosniff, X-Frame-Options: DENY, no-referrer, Cross-Origin-Resource-Policy: same-site.
CORS CORS_ALLOW_ORIGINS allowlist (no wildcard in production). Combination * + allow_credentials is auto-corrected (credentials disabled) instead of silently broken.
Rate limiting slowapi: 30 req/min/IP on /route, 10 req/min/IP on /compare.
Request size Hard cap at MAX_BODY_BYTES (default 10 000) before Pydantic. 413 returned on oversize.
Error surface 503 generic ("upstream LLM provider is not configured"). 500 opaque. /compare per-row errors drawn from a fixed safe allowlist — never echoes raw exception text.
Secrets .env gitignored; production reads OPENAI_API_KEY from Secret Manager via Cloud Run's secret mount.
API surface /docs, /redoc, /openapi.json return 404 when ENABLE_API_DOCS=false.
Dependencies npm audit clean (vite 8 + plugin-react 6). Python deps pinned in backend/requirements.txt.

About

Multi-agent system with a fine-tuned DistilBERT router (own training, in place of LLM-as-router). FastAPI backend, Planner/Executor/Critic loop, comparative evaluation.

Resources

Stars

1 star

Watchers

0 watching

Forks

Packages

 
 
 

Contributors