mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(rust): wire ICM compressor into Rust proxy on /v1/messages
Adds an opt-in compression interceptor that buffers Anthropic
/v1/messages requests, runs IntelligentContextManager over the
messages array, and forwards the (possibly trimmed) body upstream.
All other paths, methods, and content-types stay on the original
streaming passthrough — so existing operators see zero change.
Behaviour gates ALL must be true to buffer + compress:
- --compression flag (or HEADROOM_PROXY_COMPRESSION=1)
- method == POST
- path == /v1/messages
- Content-Type: application/json
- ICM constructed successfully at startup
Falls through to streaming on any failure: parse, missing fields,
unknown model, body-too-large. Compression must never break a
request — that's the safety contract.
Model context windows come from a vendored LiteLLM snapshot at
crates/headroom-proxy/data/model_prices_and_context_window.json
parsed once into an OnceLock<HashMap>. Refresh via
scripts/refresh_model_limits.sh. Rationale documented inline:
hardcoded tables silently rot; LiteLLM is the canonical source
the entire LLM-tooling ecosystem relies on.
New tests:
- 16 unit tests across compression::{anthropic, icm, model_limits}
- 5 integration tests: off-passthrough, on-short-passthrough,
on-oversized-trim, on-non-json-skip, on-non-llm-path-skip
Verification:
- cargo test --workspace -> 884 passed, 0 failed
- cargo clippy --workspace -- -D warnings -> clean
- cargo fmt --check -> clean
This commit is contained in:
parent
f726d0e280
commit
fa5fbfabf4
17 changed files with 40945 additions and 20 deletions
48
scripts/refresh_model_limits.sh
Executable file
48
scripts/refresh_model_limits.sh
Executable file
|
|
@ -0,0 +1,48 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# Refresh the vendored LiteLLM model_prices_and_context_window.json
|
||||
# used by `crates/headroom-proxy/src/compression/model_limits.rs`.
|
||||
#
|
||||
# We vendor the snapshot rather than fetching at build/runtime so the
|
||||
# proxy binary ships with no network dependency at startup. Operators
|
||||
# tracking new model releases run this script and commit the diff.
|
||||
#
|
||||
# Validation:
|
||||
# 1. JSON parses
|
||||
# 2. Contains a known-stable Claude model entry
|
||||
# 3. Contains a known-stable GPT model entry
|
||||
# These guard against accidentally vendoring an empty / malformed file.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
DEST="$REPO_ROOT/crates/headroom-proxy/data/model_prices_and_context_window.json"
|
||||
URL="https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
|
||||
|
||||
echo "Fetching $URL"
|
||||
TMP="$(mktemp)"
|
||||
trap 'rm -f "$TMP"' EXIT
|
||||
curl -fsSL "$URL" -o "$TMP"
|
||||
|
||||
# Validate the snapshot before swapping it in.
|
||||
python3 -c "
|
||||
import json, sys
|
||||
with open('$TMP') as f:
|
||||
data = json.load(f)
|
||||
if not isinstance(data, dict):
|
||||
sys.exit('top-level not an object')
|
||||
if 'sample_spec' not in data:
|
||||
sys.exit('missing sample_spec entry — schema may have changed')
|
||||
# Spot-check stable entries.
|
||||
required = ['claude-sonnet-4-5-20250929', 'gpt-4o-mini', 'gpt-4-turbo']
|
||||
missing = [k for k in required if k not in data]
|
||||
if missing:
|
||||
sys.exit(f'missing required entries: {missing!r}')
|
||||
print(f'OK: {len(data)} entries, including {required}')
|
||||
"
|
||||
|
||||
mv "$TMP" "$DEST"
|
||||
trap - EXIT
|
||||
|
||||
echo "Updated $DEST"
|
||||
echo "Run 'cargo test -p headroom-proxy --lib compression::model_limits' to verify."
|
||||
Loading…
Add table
Add a link
Reference in a new issue