feat(router): per-tenant kNN model router with preference collection - #188
Open
njbrake wants to merge 11 commits into
Open
feat(router): per-tenant kNN model router with preference collection#188njbrake wants to merge 11 commits into
njbrake wants to merge 11 commits into
Conversation
njbrake
temporarily deployed
to
integration-tests
June 19, 2026 17:29 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
|
Warning Review limit reached
More reviews will be available in 12 minutes and 28 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan refill rate. For paid Pro and Pro+ PR reviews, CodeRabbit uses rolling per-developer review limits. Reviews become available again as older review attempts age out of the rolling limit window. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (7)
WalkthroughAdds a cost-optimizing kNN model router to the Otari gateway. The backend selects cheaper candidate models using per-tenant embedding-based nearest-neighbor voting, with preference collection endpoints, chat integration, trace-sticky conversation pinning, and task partitioning. Includes a standalone offline React demo app and full docs. ChangesGateway kNN Router Feature
React Router Demo App
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
dpoulopoulos
self-requested a review
June 22, 2026 14:54
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 18:36 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 18:52 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 18:58 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 19:04 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 20:38 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 21:04 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 21:25 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 21:25 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 22:15 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 23:01 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 22, 2026 23:36 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 23, 2026 01:45 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 23, 2026 10:47 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
temporarily deployed
to
integration-tests
June 23, 2026 10:53 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Optional, learn-from-your-data model router for standalone mode, so easy prompts can be served by a cheaper model while hard prompts stay on the strong one. Off by default (OTARI_ROUTER_BACKEND); when disabled the chat path is unchanged. - Pluggable RouterBackend seam (RoutingContext / RoutingDecision) plus a no-op backend, shaped like the sandbox / web-search backends. - KnnRoutingMemory: per-tenant cosine kNN over (embedding, model, quality, cost) records, scoring mean_quality - alpha * normalized_cost. Cold start, sparse neighborhoods, sub-floor confidence, and tool requests fall back to the requested model, so routing is never worse than off. - trace_sticky granularity reuses a conversation's first decision, keyed on an Otari-Conversation-Id header; hard per-task partitions are selectable per request via Otari-Router-Task and warm independently; Otari-Router opts a single request out. - Preference API (standalone only): /v1/router/preferences/compare and /rank, and a per-pool /v1/router/status overview. Storage is RoutingMemory plus a RouterPreference audit row with an Alembic migration; the new config fields are validated at load. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Unit tests for the scoring, confidence, trace-key, and granularity logic, the default-off seam, and the request headers. Integration tests drive the FastAPI app over SQLite/Postgres: the preference endpoints, end-to-end routing through the chat seam, per-tenant and per-task isolation, and the status overview. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add docs/routing.md (setup, the Otari-Router / Otari-Conversation-Id / Otari-Router-Task headers, task partitions, tuning) and docs/routing-scaling.md (how trace stickiness behaves across replicas and the target shared-store design). Update the configuration and API references and regenerate the OpenAPI spec. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A self-contained, offline demo (demo/router/) with two guided walkthroughs: "How it learns" steps through the preference API (the per-tenant/per-task model, compare, scoring, rank, and per-partition warm-up), and "See it route" animates the kNN decision (gates, embedding, neighbors, quality-vs-cost scoring, and the resolved pick). No gateway or provider calls at runtime; the prompts, answers, and embeddings are bundled. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The demo and guide threaded task_id through the canonical rank and chat requests, which made it read as required. Drop it from the demo's default rank payload and 200 OK, and introduce it only as an opt-in: the isolation callout now shows how to add task_id plus the Otari-Router-Task header, and both that header and Otari-Router are marked optional in "See it route". Reword the warm step and the routing guide so per-task partitioning is clearly the optional layer over the per-tenant default pool, and fix a stale GET /v1/router/status?task_id reference (status now returns the full overview). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The "See it route" request carries Otari-Router-Task: <task>, but the demo's kNN sim searched every task, so the neighbor list and graph mixed partitions, inconsistent with the filter the request advertises. Route the sim within the query's task partition (only same-task prompts are candidate neighbors), draw just the searched partition in the graph (other tasks dimmed, no edges), and drop the caveat that said the demo searches one shared store. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
record_preference wrote one row per scored model, all sharing the prompt's
embedding, so the kNN treated N rows of one example as N independent neighbors.
That conflated examples with rows: the seed gate counted rows (~N x examples),
the top-k neighborhood was dominated by the single nearest example's N rows, and
the confidence metric was structurally capped near 1/N. /rank also reported
recorded = number of models (e.g. 4 for one ranked prompt), which is what
surfaced in the demo as "one example, four samples".
Store one record per example: the prompt embedding plus a {model: quality} map
(RoutingMemory.qualities replaces model/quality/cost; migration updated). The
kNN now votes over distinct prompts, the seed gate and /status counts are in
examples, /rank returns recorded = 1, and confidence is the genuine share of the
k neighbor prompts whose own best-scoring model is the chosen one.
Demo: the rank step shows one record with the per-model score map and
recorded: 1; the warm step counts examples (shared SEED_COUNT lowered to 3).
Docs and tests updated.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Resolve conflicts:
- chat.py: keep both the router header imports (CONVERSATION_HEADER,
ROUTER_HEADER, ROUTER_TASK_HEADER) and main's new GatewayUsage import.
- config.py: keep the router_* fields and adopt main's updated mode/platform
descriptions ("standalone or hybrid", "otari.ai connection settings").
- Alembic: main added a migration that reused revision id d4e6f8a0b2c4, the same
id as the router tables migration. Re-chain the router migration to a new id
(f0a1b2c3d4e5) descending from main's d4e6f8a0b2c4, giving a single linear head.
- Fix a warn_return_any mypy error in test_router_seam_e2e._post_chat surfaced by
main's updated test deps.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
njbrake
temporarily deployed
to
integration-tests
June 23, 2026 15:20 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
The guide and API reference drifted from the code after later branch changes:
- routing.md step 3 used a "ranking" list and showed {recorded: 2, tenant_warm}.
The rank API takes a required `scores` map ({model: 0..1}); there is no
`ranking` field, so the example 422s. The response is {recorded, task_id,
warm} with recorded = 1 (one record per example). Reword the score-don't-rank
flow to match.
- "How it works" described records as (embedding, model, quality, cost); they
are now one per example: (embedding, {model: quality}).
- api-reference.md: rank writes one record of per-model scores (not a
rank-to-scalar ranking); /status returns the default pool plus each task
partition, not a single record count.
- Note the low-confidence safety net only applies when a confidence floor is set
(off by default).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
njbrake
temporarily deployed
to
integration-tests
June 23, 2026 15:30 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
…ndidate pool Two router correctness fixes (both in knn_router.py, so kept together): - Require pricing for every OTARI_ROUTER_CANDIDATES model. The router scores by cost, so an unpriced candidate has nothing to weigh. Drop the _UNKNOWN_PRICE fallback (now dead): validate at startup (fail fast with RouterPricingError) and raise again at request time if a candidate is unpriced. - Confidence counts only candidate models. A neighbor's favorite model that is not in the routing pool no longer collapses confidence to 0, which, with a confidence floor set, would have vetoed a correct cheaper pick back to the pricier requested model. Adds a startup-validation integration test (the harness now configures the provider + pricing for its candidates) and a confidence unit test. Docs note that candidates must be priced. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
njbrake
temporarily deployed
to
integration-tests
June 23, 2026 17:02 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
njbrake
marked this pull request as ready for review
June 23, 2026 17:27
njbrake
commented
Jun 23, 2026
Member
Author
There was a problem hiding this comment.
I built this demo as a nice way to walkthrough how the knn routing is working but I'm also open to deleting this if it feel unnecessary to commit to git
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 13
🧹 Nitpick comments (2)
docs/routing.md (1)
138-138: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd language markers to fenced code blocks showing example headers.
Lines 138 and 168 have code fences without a language specifier. To satisfy markdown linting (and to make the intent clear to readers), add a language marker. Since these are example HTTP headers, use
`text`or`bash`:# Line 138: - `Otari-Router: off` serves... Instead of: +Otari-Conversation-Id: 4f9c2b10-... Use: + ``` text +Otari-Conversation-Id: 4f9c2b10-... +```This keeps the visual clarity of code blocks while satisfying the linter and readers.
Also applies to: 168-168
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/routing.md` at line 138, The fenced code blocks at lines 138 and 168 in the routing.md file are missing language specifiers on their opening code fence markers. Add a language marker to each opening fence by changing the bare triple backticks (```) to include a language identifier such as `text` or `bash` (for example, ```text or ```bash). This satisfies markdown linting requirements and provides clarity to readers about the content type of the code blocks containing example HTTP headers.Source: Linters/SAST tools
demo/router/generate_demo_dataset.py (1)
251-269: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the repository logger instead of
These progress/status lines should use the project logger pattern so output is structured and consistent with repo logging conventions.
As per coding guidelines,
**/*.py: use module logger fromgateway.log_configwith structured/contextual log messages using%sformatting placeholders.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@demo/router/generate_demo_dataset.py` around lines 251 - 269, Replace all print statements in the main function with structured logging using the project's logger from gateway.log_config. The three print statements that display progress output (showing candidate index and spread information, divergence summary, and file write confirmation) should be converted to logger calls using percent-style formatting placeholders (%s) instead of f-strings. Import the logger module at the top of the file and use appropriate log levels (info or debug) for these status messages to maintain consistency with repository logging conventions.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@demo/router/src/components/StepWalkthrough.tsx`:
- Around line 33-40: The global keyboard event listener in the useEffect hook
unconditionally triggers step navigation for ArrowLeft and ArrowRight keys,
which interferes with interactive controls like sliders that also use arrow
keys. Inside the onKey handler function, add a check to verify that the event
target is not an interactive element (such as an input, slider, or other form
control) before calling next() or prev(). You can check the event target's type
or role to determine if it is an interactive control, and only proceed with
navigation if the event originated from a non-interactive element.
In `@demo/router/src/components/TeachWalkthrough.tsx`:
- Around line 24-38: The code assumes demo.items is non-empty, causing crashes
when it has zero items (accessing demo.items[0] on line 24 and reducing an empty
array on line 38). Add an early empty-state guard clause at the beginning of the
TeachWalkthrough component that checks if demo.items.length is zero and handles
this case before deriving the item variable and initializing the useState hooks
that depend on having data. This ensures the component gracefully handles empty
datasets before attempting to access array indices or perform reduce operations.
In `@demo/router/src/main.tsx`:
- Around line 2-7: The import statement in main.tsx is using a default import
for ReactDOM from react-dom/client, but React 19 exports createRoot as a named
export only. Change the import statement from `import ReactDOM from
"react-dom/client"` to `import { createRoot } from "react-dom/client"`, then
update the ReactDOM.createRoot call to use createRoot directly, replacing
`ReactDOM.createRoot(document.getElementById("root") as
HTMLElement).render(...)` with `createRoot(document.getElementById("root") as
HTMLElement).render(...)`.
In `@demo/router/src/router-sim.ts`:
- Around line 63-91: When no eligible neighbors exist in a task partition, the
neighborItems array becomes empty, causing the meanQuality calculation to
default to 0 for all candidates, which then selects the cheapest model instead
of preserving the fallback behavior. After the neighborItems assignment and
before the pool.map call that creates the candidates array, add an explicit
check: if neighborItems.length is 0, set winner to the fallback value
(strongest) and either skip the candidate scoring loop or handle it as a special
case to preserve the documented fallback behavior.
In `@src/gateway/api/routes/router.py`:
- Around line 8-10: The module docstring for the POST
/v1/router/preferences/rank endpoint currently describes the old behavior of
writing one routing-memory record per scored model, but the implementation now
writes a single routing-memory example with a score map instead. Update the
docstring (lines 8-10) to accurately reflect the current storage model where one
submission creates one routing-memory example containing a score map for all
models, plus one audit row, rather than describing multiple records per model.
- Around line 51-52: The `models` field in the schema has no maximum length
constraint, allowing users to submit arbitrarily large arrays that trigger
unbounded concurrent provider calls at line 145, creating excessive upstream
load. Add a `max_length` parameter to the Field definition for the `models` list
to enforce a schema-level hard cap on the number of models allowed.
Additionally, implement a semaphore or similar concurrency limiting mechanism
around line 145 where the provider calls are launched asynchronously to ensure
that only a small, bounded number of concurrent calls execute regardless of the
input array size.
- Around line 141-143: The ModelResponse being returned in the exception handler
is exposing raw exception details via str(exc) directly to clients, which can
leak sensitive upstream provider information. Replace the str(exc) value passed
to the error parameter in the ModelResponse with a generic, user-friendly error
message that does not reveal internal details. The actual error details are
already being captured server-side in the logger.warning call, so keep those
specifics there and only return a safe generic message to the client.
In `@src/gateway/models/entities.py`:
- Around line 310-313: The docstring for the RouterPreference class incorrectly
describes the storage model by stating that rank submissions write one
RoutingMemory row per scored model. Update this docstring to accurately reflect
the new storage model where one row is written per example with per-model
qualities stored together, rather than separate rows per model. This ensures the
documentation matches the actual implementation and prevents confusion during
future maintenance.
In `@src/gateway/services/knn_router.py`:
- Around line 235-247: The database write operation in the RoutingMemory
creation block does not have proper error handling for commit failures. Wrap the
await db.commit() call in a try/except block that catches SQLAlchemyError, calls
db.rollback() in the except clause on any database error, and re-raises an
appropriate mapped API or domain error. This same transaction handling pattern
must also be applied to the other database write operation mentioned at lines
374-399 to ensure consistent error handling across all router DB writes and
prevent sessions from remaining in failed transaction states.
- Around line 216-248: The record_preference() method opens and commits its own
database session independently, causing a race condition where the
RouterPreference row can be persisted before this method is called, but if the
embedding or insert fails here, the audit row exists without the corresponding
routing-memory example. Modify record_preference() to accept an optional
AsyncSession parameter and use that session directly instead of creating a new
one with async with create_session(). This allows the calling code in the route
handler to manage both the RouterPreference and RoutingMemory writes as a single
atomic transaction, ensuring both succeed or fail together.
- Around line 193-195: In the exception handler block where embedding fails (in
the except Exception clause with logger.warning), change the logging to avoid
logging the raw exception text which may contain user-derived content or
sensitive data from provider exceptions. Instead, log only the exception class
name or type using type(exc).__name__ rather than logging the full exc object
with the %s placeholder. This prevents privacy leaks from user payloads or
provider-specific error messages in logs while still providing debugging
context.
In `@tests/integration/test_config_env_loading.py`:
- Around line 150-166: The test_router_knob_defaults function only clears
OTARI_ROUTER_* environment variables but GatewayConfig also accepts legacy
GATEWAY_ROUTER_* fallback environment variables, making the assertions
environment-dependent. Add the legacy GATEWAY_ROUTER_* environment variables
(such as GATEWAY_ROUTER_ALPHA, GATEWAY_ROUTER_K, GATEWAY_ROUTER_SEED_COUNT,
GATEWAY_ROUTER_GRANULARITY, GATEWAY_ROUTER_CANDIDATES, and
GATEWAY_ROUTER_EMBEDDING_MODEL) to the loop in monkeypatch.delenv calls so both
namespaces are cleared before instantiating GatewayConfig().
In `@tests/integration/test_router_preferences.py`:
- Around line 87-105: The _build_client function hardcodes "test-key" as the
OpenAI API key in the providers configuration, which causes issues when running
in live mode where the OPENAI_API_KEY environment variable is available. Modify
the _build_client function to check if the OPENAI_API_KEY environment variable
exists and use it instead of the hardcoded "test-key" value in the providers
dict. Apply this same fix to all other occurrences of this pattern in the file
(at lines 131-135 and 463-467 as indicated) to ensure consistency across all
test client builders.
---
Nitpick comments:
In `@demo/router/generate_demo_dataset.py`:
- Around line 251-269: Replace all print statements in the main function with
structured logging using the project's logger from gateway.log_config. The three
print statements that display progress output (showing candidate index and
spread information, divergence summary, and file write confirmation) should be
converted to logger calls using percent-style formatting placeholders (%s)
instead of f-strings. Import the logger module at the top of the file and use
appropriate log levels (info or debug) for these status messages to maintain
consistency with repository logging conventions.
In `@docs/routing.md`:
- Line 138: The fenced code blocks at lines 138 and 168 in the routing.md file
are missing language specifiers on their opening code fence markers. Add a
language marker to each opening fence by changing the bare triple backticks
(```) to include a language identifier such as `text` or `bash` (for example,
```text or ```bash). This satisfies markdown linting requirements and provides
clarity to readers about the content type of the code blocks containing example
HTTP headers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: a240f8a6-be5d-4911-8049-619de44f244b
⛔ Files ignored due to path filters (1)
docs/public/openapi.jsonis excluded by!docs/public/openapi.json
📒 Files selected for processing (45)
alembic/versions/f0a1b2c3d4e5_add_router_tables.pydemo/router/.gitignoredemo/router/README.mddemo/router/generate_demo_dataset.pydemo/router/index.htmldemo/router/package.jsondemo/router/src/App.tsxdemo/router/src/components/CodeJson.tsxdemo/router/src/components/KnnGraph.tsxdemo/router/src/components/StepWalkthrough.tsxdemo/router/src/components/TeachWalkthrough.tsxdemo/router/src/components/Walkthrough.tsxdemo/router/src/components/ui/button.tsxdemo/router/src/components/ui/card.tsxdemo/router/src/components/ui/index.tsdemo/router/src/demo-config.tsdemo/router/src/demo_prompts.jsondemo/router/src/format.tsdemo/router/src/globals.cssdemo/router/src/main.tsxdemo/router/src/router-sim.tsdemo/router/src/styles/brand-tokens.cssdemo/router/src/types.tsdemo/router/tsconfig.jsondemo/router/vite.config.tsdocs/api-reference.mddocs/configuration.mddocs/index.mddocs/routing-scaling.mddocs/routing.mdsrc/gateway/api/deps.pysrc/gateway/api/main.pysrc/gateway/api/routes/chat.pysrc/gateway/api/routes/router.pysrc/gateway/core/config.pysrc/gateway/main.pysrc/gateway/models/entities.pysrc/gateway/services/knn_router.pysrc/gateway/services/router_backend.pytests/integration/test_config_env_loading.pytests/integration/test_router_preferences.pytests/integration/test_router_seam_e2e.pytests/unit/test_knn_router.pytests/unit/test_router_backend.pytests/unit/test_router_header.py
- /preferences/compare: cap the model fan-out (max 8) so one request cannot launch unbounded concurrent provider calls, and return a generic per-model error instead of the raw provider exception text. - knn_router: log the embedding-failure exception class only (not raw text, which can carry request content), and wrap the routing-memory and eviction commits in try / SQLAlchemyError rollback per the project's DB-write convention. - /preferences/rank: write the routing-memory record before the audit row, so a failure never leaves an orphan RouterPreference row without its RoutingMemory. - Refresh stale "one record per scored model" docstrings (one record per example now) on the router module and RouterPreference. - Tests: use the real OPENAI_API_KEY in the live integration client (the pricing fix had hardcoded a placeholder), and clear the legacy GATEWAY_ROUTER_* env namespace in the router-defaults test. - Demo: arrow-key navigation ignores interactive controls (score sliders), and main.tsx uses the named createRoot import. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
njbrake
temporarily deployed
to
integration-tests
June 23, 2026 18:14 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
There was a problem hiding this comment.
Pull request overview
This PR adds an optional, standalone-only model routing system to the gateway, including a pluggable router backend seam and a kNN-based routing-memory implementation that can be trained via preference collection endpoints.
Changes:
- Introduces a
RouterBackendprotocol withnone(default),noop, andknnbackends, plus request headers to control routing and conversation stickiness. - Implements kNN routing memory (embedding, per-tenant/task partitions, confidence gating, trace-sticky caching) and persists training/audit data via new DB tables and migration.
- Adds standalone router preference and status endpoints, documentation, OpenAPI updates, and an offline demo app.
Reviewed changes
Copilot reviewed 45 out of 46 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/gateway/core/config.py |
Adds router config knobs and header constants, plus validation for backend and granularity. |
src/gateway/services/router_backend.py |
Defines the router backend seam, no-op backend, and cached knn backend factory. |
src/gateway/services/knn_router.py |
Implements the kNN routing-memory algorithm, persistence, trace stickiness, and pricing validation. |
src/gateway/api/routes/chat.py |
Wires routing into the standalone chat path with opt-out and conversation/task headers. |
src/gateway/api/routes/router.py |
Adds standalone preference collection endpoints (compare, rank) and routing-memory status. |
src/gateway/api/main.py |
Registers the router endpoints in standalone mode only. |
src/gateway/api/deps.py |
Clears the router backend cache during test config resets. |
src/gateway/models/entities.py |
Adds RoutingMemory and RouterPreference ORM entities. |
src/gateway/main.py |
Validates router candidate pricing at startup when router_backend=knn. |
alembic/versions/f0a1b2c3d4e5_add_router_tables.py |
Creates the routing_memory and router_preferences tables and indexes. |
tests/unit/test_router_backend.py |
Unit tests for router backend selection and noop behavior. |
tests/unit/test_router_header.py |
Unit tests for Otari-Router, Otari-Conversation-Id, and Otari-Router-Task parsing. |
tests/unit/test_knn_router.py |
Unit tests for kNN scoring, confidence, tool gating, and stickiness logic. |
tests/integration/test_router_seam_e2e.py |
End-to-end tests proving routing is inert by default and applied on standalone (including streaming). |
tests/integration/test_router_preferences.py |
Integration tests for preference endpoints, warming, task partitions, and startup pricing enforcement. |
tests/integration/test_config_env_loading.py |
Tests config defaults and OTARI env overrides for router settings. |
docs/routing.md |
User documentation for enabling, training, and operating routing. |
docs/routing-scaling.md |
Design note for stickiness behavior under scale and planned durable design. |
docs/configuration.md |
Adds configuration reference for router env vars. |
docs/api-reference.md |
Documents new router endpoints in the API reference table. |
docs/index.md |
Links to routing docs from the docs index. |
docs/public/openapi.json |
Updates generated OpenAPI spec with new schemas and endpoints. |
demo/router/package.json |
Adds an offline demo app package manifest and deps. |
demo/router/vite.config.ts |
Vite config for the router demo. |
demo/router/tsconfig.json |
TypeScript config for the demo app. |
demo/router/index.html |
Demo app HTML entry point. |
demo/router/.gitignore |
Demo-specific ignore rules. |
demo/router/README.md |
Demo instructions and explanation. |
demo/router/src/main.tsx |
Demo React entry. |
demo/router/src/App.tsx |
Demo app view switch and top-level layout. |
demo/router/src/types.ts |
Demo data types. |
demo/router/src/router-sim.ts |
Client-side reproduction of kNN routing for the offline demo. |
demo/router/src/demo-config.ts |
Shared demo seed gate constant. |
demo/router/src/format.ts |
Shared score formatting helper. |
demo/router/src/globals.css |
Demo global styles and imports. |
demo/router/src/styles/brand-tokens.css |
Brand token CSS variables for theming. |
demo/router/src/components/StepWalkthrough.tsx |
Shared step walkthrough component with keyboard nav. |
demo/router/src/components/TeachWalkthrough.tsx |
Demo walkthrough for preference collection and warming. |
demo/router/src/components/Walkthrough.tsx |
Demo walkthrough for explaining routing decisions. |
demo/router/src/components/KnnGraph.tsx |
Demo visualization of prompt neighborhood graph. |
demo/router/src/components/CodeJson.tsx |
Demo JSON pretty-print block with simple highlighting. |
demo/router/src/components/ui/button.tsx |
Demo UI button wrapper. |
demo/router/src/components/ui/card.tsx |
Demo UI panel wrapper. |
demo/router/src/components/ui/index.ts |
Demo UI exports. |
demo/router/generate_demo_dataset.py |
Script to generate the bundled demo dataset using live models. |
Comment on lines
+208
to
+209
| prices = await self._candidate_prices(pool) | ||
| decision = self._score(ctx, pool, neighbors, prices) |
Comment on lines
+128
to
+155
| @router.post("/preferences/compare", response_model=CompareResponse) | ||
| async def compare_models( | ||
| request: CompareRequest, | ||
| api_key: Annotated[APIKey, Depends(verify_api_key)], | ||
| config: Annotated[GatewayConfig, Depends(get_config)], | ||
| ) -> CompareResponse: | ||
| """Fan the prompt out to every candidate model concurrently.""" | ||
| messages: list[dict[str, Any]] = [] | ||
| if request.system: | ||
| messages.append({"role": "system", "content": request.system}) | ||
| messages.append({"role": "user", "content": request.prompt}) | ||
|
|
||
| async def _one(model: str) -> ModelResponse: | ||
| try: | ||
| provider, bare = AnyLLM.split_model_provider(model) | ||
| kwargs = get_provider_kwargs(config, provider) | ||
| if request.max_tokens is not None: | ||
| kwargs["max_tokens"] = request.max_tokens | ||
| result = await acompletion(model=bare, provider=provider, messages=cast(Any, messages), **kwargs) | ||
| choices = getattr(result, "choices", None) | ||
| content = choices[0].message.content if choices else None | ||
| return ModelResponse(model=model, content=content) | ||
| except Exception as exc: # one model failing must not sink the comparison | ||
| logger.warning("compare: model %s failed (%s)", model, type(exc).__name__) | ||
| return ModelResponse(model=model, content=None, error="model invocation failed") | ||
|
|
||
| responses = await asyncio.gather(*[_one(m) for m in request.models]) | ||
| return CompareResponse(prompt=request.prompt, responses=list(responses)) |
Comment on lines
+98
to
+102
| | Method | Path | Description | Auth | | ||
| |--------|------|-------------|------| | ||
| | `POST` | `/v1/router/preferences/compare` | Fan a prompt out to several candidate models and return their responses side by side. | API key or master key | | ||
| | `POST` | `/v1/router/preferences/rank` | Record per-model quality scores (0–1) for a prompt; writes one routing-memory record. Requires the `knn` backend. | API key or master key | | ||
| | `GET` | `/v1/router/status` | Report routing-memory progress: the default pool plus each task partition, with record counts and warmth. | API key or master key | |
khaledosman
reviewed
Jul 16, 2026
khaledosman
left a comment
khaledosman
left a comment
Member
There was a problem hiding this comment.
Reviewed the full diff (gateway source, migration, tests, docs, demo) against AGENTS.md and .github/instructions/security-review.instructions.md. Solid, well-layered work: the seam is genuinely inert by default, the kNN logic is clean and well-tested, tenant isolation is done right, and billing correctly attributes the routed model (tested). Four inline comments below, two of them budget/settlement issues per the repo's own security checklist.
Two notes that don't fit a line:
- The branch is 56 commits behind main, including #275 (completion routes consolidated onto the shared pipeline), so the
chat.pychanges will conflict on rebase. Main also now requiresscripts/generate_postman.pyregeneration (make postman-check) for API changes; that script doesn't exist on this branch, so the new/v1/router/*endpoints will need it post-rebase. - +1 to Copilot's still-unaddressed point that
docs/api-reference.mdlists the router endpoints as "API key or master key" whileverify_api_keyrejects the master key (it isn't in theapi_keystable). Docs should say "API key", anddocs/routing.mdcould note the master key can't be used to teach (it has no tenant).
This review was created by Claude Code.
|
|
||
| provider, model = AnyLLM.split_model_provider(request.model) | ||
| call_kwargs = {**get_provider_kwargs(config, provider), **request_fields} | ||
| routed_model = await _resolve_standalone_model( |
Member
There was a problem hiding this comment.
Budget-reservation leak: _resolve_standalone_model runs after resolve_request_context has taken the reservation but outside run_standalone_non_stream's settlement, and nothing here releases it on error (same at the streaming call site, line 485). backend.route() can raise: RouterPricingError when the requested model has no pricing (require_pricing=false, or a pricing row deleted after startup; startup validation only covers configured candidates while the requested model joins the pool per request), or SQLAlchemyError from _load_records/_candidate_prices. Result: a 500 plus a permanently leaked reservation (spend + reserved creeps toward max_budget), which is the exact §0.3 pattern the security instructions call a finding, and it contradicts the knn_router module docstring's promise that an unpriced candidate pool falls back to the requested model.
Fix: extend the best-effort guard that already covers _embed to the whole decision, e.g. wrap the backend.route(...) call (in _resolve_standalone_model or in KnnRoutingMemory.route) in try/except Exception that logs the exception type and returns the requested model. That also makes the docstring true and adds a unit test slot next to test_embedding_failure_passes_through.
| logger.warning("compare: model %s failed (%s)", model, type(exc).__name__) | ||
| return ModelResponse(model=model, content=None, error="model invocation failed") | ||
|
|
||
| responses = await asyncio.gather(*[_one(m) for m in request.models]) |
Member
There was a problem hiding this comment.
Building on Copilot's budget-bypass point with the concrete exposure: any active user API key, including one whose budget is exhausted, can fan a request out to 8 arbitrary provider:model pairs per call, not just the configured router_candidates, with no budget reservation, no usage log, no check_rate_limit, and no require_pricing gate. Per the repo's security instructions (§0.3/0.4) a billable route without the reserve→settle lifecycle is a finding. The mitigation docs/routing.md suggests ("gate who holds keys that can reach it") has no mechanism: every chat-capable key can reach this endpoint, and the master key actually cannot (verify_api_key 401s it).
Suggested hardening, in order of value: (1) restrict request.models to the configured OTARI_ROUTER_CANDIDATES pool (they're the only models worth scoring anyway), (2) apply check_rate_limit, (3) meter the fan-out through the normal reserve/reconcile + usage-log path, or at minimum write usage rows so the spend is visible.
| trace_key=conversation_id, | ||
| ) | ||
| decision = await backend.route(routing_ctx) | ||
| return decision.ordered_models[0] |
Member
There was a problem hiding this comment.
Routing decisions are invisible in production: decision.confidence and decision.rationale are computed and well-tested but never logged or exposed, so an operator can only detect a reroute by diffing usage logs against requests. A logger.debug here with tenant, requested → chosen model, confidence, and rationale (all static strings, no user content) would make the first rollout debuggable and give the fast-follow eval work something to grep.
|
|
||
| ## Notes | ||
|
|
||
| - **Blind ranking** (Settings, on by default) hides which model wrote each answer |
Member
There was a problem hiding this comment.
Stale: the shipped demo has no Settings, no blind-ranking mode, and no localStorage use (nothing in src/ references any of these; it became a read-only walkthrough). Drop this bullet and the "Settings persist in localStorage" line.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note: this PR description was drafted by Claude via back-and-forth with @njbrake. The reasoning and decisions are his; the prose and code are Claude's.
Description
Adds an optional, learn-from-your-own-data model router to Otari (standalone mode) so easy prompts can be served by a cheaper model while hard prompts stay on the strong one. Off by default; when disabled the chat path is byte-for-byte unchanged.
The work lands in layers:
RouterBackendprotocol (RoutingContext/RoutingDecision) shaped likeSandboxBackend/WebSearchBackend, plus a no-op backend and theOTARI_ROUTER_BACKENDflag. Withnone(default) the chat handler behaves as it did before.OTARI_ROUTER_BACKEND=knn).KnnRoutingMemoryembeds the task signal viaany_llmand runs per-tenant cosine kNN over one record per scored example (the prompt embedding plus a{model: quality}map), scoring each candidate asmean_quality(m|neighbors) - alpha * normalized_cost(m). Cold start, a sparse neighborhood, sub-floor confidence, tool requests, and an unpriced candidate pool all fall back to the requested model, so the safe default is never worse than no routing.trace_stickygranularity (default) decides once per conversation and reuses that decision on later turns. The trace identity is a client-suppliedOtari-Conversation-Idheader, or a per-tenant hash of the conversation's opening messages when absent. A client can skip routing for one request withOtari-Router: off.Otari-Router-Taskvotes only over that task's records and warms independently of other tasks; preferences are filed into a partition with an optionaltask_idon/rank. Without a task, everything shares one default pool.POST /v1/router/preferences/comparefans a prompt out to N models;/rankrecords per-model quality scores (0 to 1) as one routing-memory record per example;GET /v1/router/statusreports the per-pool overview (the default pool plus each task partition, with record counts and warmth).OTARI_ROUTER_CANDIDATESmust have configured pricing or the gateway fails fast at startup.RoutingMemory(one record per example,embedding_modelinvalidation tag, per-tenant cap, oldest-first eviction) andRouterPreference(preference audit) plus an Alembic migration.Config:
OTARI_ROUTER_BACKEND,OTARI_ROUTER_CANDIDATES,OTARI_ROUTER_ALPHA,OTARI_ROUTER_K,OTARI_ROUTER_EMBEDDING_MODEL,OTARI_ROUTER_CONFIDENCE_FLOOR,OTARI_ROUTER_SEED_COUNT,OTARI_ROUTER_GRANULARITY,OTARI_ROUTER_MAX_VECTORS_PER_TENANT.Try it
A self-contained, fully offline demo lives in
demo/router/(no gateway, no API key, nothing called at runtime):/preferences/compare, scoring each answer (the gpt-5.4 judge's 0 to 1 scores, editable),/preferences/rankand the routing-memory record it writes, and watching each task partition warm independently.The bundled data (
demo/router/src/demo_prompts.json) is a four-tier ladder (gpt-5.4 / gpt-4o / gpt-4o-mini / gpt-3.5-turbo) over general-knowledge, hard-code, and hard-math prompts, built bydemo/router/generate_demo_dataset.py. Seedemo/router/README.md.The kNN scoring design was validated offline against the RouterBench dataset during development: the router beat a cost-matched blend of cheap and strong models, and a clear gap to the oracle confirmed the signal is region-level rather than per-prompt, which motivates the task-metadata fast-follow. That exploratory harness is not part of this PR.
Deviations from the v1 plan in #187
This PR ships a v1 of #187 but does not match the original v1 checklist exactly:
effective_costterm and switch penalty (withOTARI_ROUTER_SWITCH_PENALTY), and theRoutingOutcome/record_outcomepassive-learning hook. Scoring is currently on normalized list price, and only the preference flow writes memory.Otari-Router-Task), theOtari-Conversation-Ididentity for trace stickiness, the per-pool/statusoverview, require-pricing for candidates, and a routing-at-scale design note (docs/routing-scaling.md).What remains for #187 (fast-follow)
record_outcomeseam).fasteststrategy and for latency reporting.cheapest/fastest).UsageLog("a cheaper model would have sufficed on X%").PR Type
Relevant issues
Implements a v1 of #187, with the deviations noted above. Fast-follow items remain, so this does not close the issue.
Checklist
tests/unit,tests/integration).make lint,make typecheck,make test).uv run python scripts/generate_openapi.py).Note on test scope: ruff and mypy are clean (187 files); the full unit suite (532) passes; the router unit and integration suites pass. The Docker-based testcontainers path was not available in my environment, so the router integration suite was run against SQLite via
TEST_DATABASE_URLrather than testcontainers PostgreSQL.AI Usage
AI Model/Tool used: Claude Code (Opus 4.8)
Any additional AI details you'd like to share: Code, tests, docs, and the demo were generated by Claude through iterative back-and-forth with @njbrake, who directed the design decisions (standalone-first, kNN default, preference-collection flow) and reviewed the output.
What Changed
This PR adds a model routing system that intelligently selects which AI model to use for each request, optimizing for cost while maintaining quality. The router learns from user feedback and automatically routes straightforward requests to cheaper models while keeping complex ones on more capable (expensive) models.
Core Components Added
Backend Service: A new kNN-based routing engine that:
API Endpoints (standalone mode only):
POST /v1/router/preferences/compare— collect responses from multiple candidate modelsPOST /v1/router/preferences/rank— record quality scores for model responsesGET /v1/router/status— check router warmup progress per tenantConfiguration: Nine environment variables control backend selection, candidate models, routing aggressiveness, and safety thresholds
Database: Two new tables (
routing_memory,router_preferences) with Alembic migration to store embeddings, quality scores, and preference audit trailsRequest Control:
Otari-Conversation-Idheader enables "sticky" routing (same model within a conversation)Otari-Router: offheader opts a single request out of routingOtari-Router-Taskheader partitions routing by task type, maintaining independent learning per partitionDemo Application
A fully offline React/Vite application (
demo/router/) with:Safety Guards
Routing only occurs when conditions are met:
Benefits
Documentation