Commit graph

1 commit

Author SHA1 Message Date
Krishna Chaitanya
57e8dcb425
feat(proxy): add opt-in cost-aware model router (#1706) (#2205)
## Description

Adds an optional, configuration driven model router (closes #1706). With
`HEADROOM_MODEL_ROUTER_ENABLED` set, ordered rules in
`HEADROOM_MODEL_ROUTES` rewrite the upstream model by estimated input
size and tool presence, complementary to content compression, for
example sending small, tool-free requests to a cheaper model. First
matching rule wins and every decision is logged with a reason. Off by
default so behavior is unchanged, skipped under
`x-headroom-bypass`/passthrough, and wired on the Anthropic
`/v1/messages` path. Malformed rules fail open, so a bad rule is skipped
rather than silently widened.

Closes #1706

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [x] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [x] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/proxy/model_router.py`: new `ModelRouter` component (ordered
rules, first-match decision with reason, fail-open env parsing,
tokenizer-free input estimate).
- `headroom/proxy/models.py` + `headroom/proxy/server.py`:
`ProxyConfig.model_router` field, env loader
(`HEADROOM_MODEL_ROUTER_ENABLED` / `HEADROOM_MODEL_ROUTES`), and proxy
wiring.
- `headroom/proxy/handlers/anthropic.py`: apply routing on
`/v1/messages` after the bypass gate, tracked as a body mutation.
- Tests, docs (`configuration.mdx`), and a CHANGELOG entry.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ pytest -q tests/test_proxy/test_model_router.py tests/test_proxy/test_model_router_wiring.py
36 passed, 1 warning

$ ruff check .
All checks passed!

$ mypy headroom --ignore-missing-imports
Success: no issues found in 477 source files
```

## Real Behavior Proof

- Environment: local, macOS, Python 3.12, headroom `.venv`, upstream
mocked (no live provider call).
- Exact command / steps: enable the router via
`ProxyConfig(model_router=...)`, POST `/v1/messages` through
`TestClient` with a rule routing low-risk requests to a cheaper model;
repeat with header `x-headroom-bypass: true`.
- Observed result: the forwarded upstream body model is rewritten from
`claude-sonnet-4-6` to `claude-haiku-4-5` when the router is enabled,
and is left unchanged under bypass (see
`tests/test_proxy/test_model_router_wiring.py`).
- Not tested: the OpenAI and Gemini handler paths (this PR wires the
Anthropic path only); no live provider request (upstream is mocked).

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

Happy to adjust the interface or scope (for example OpenAI and Gemini
parity) if you'd prefer a different shape.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: reneleonhardt <65483435+reneleonhardt@users.noreply.github.com>
2026-07-15 19:58:17 +00:00