feat(guardrails): add Compresr guardrail for query-aware context compression - #33019
feat(guardrails): add Compresr guardrail for query-aware context compression#33019charafkamel wants to merge 7 commits into
Conversation
|
|
charafkamel
changed the title
Greptile SummaryThis PR introduces a new Compresr guardrail that runs on
Confidence Score: 4/5The change is on the critical request path and mutates messages before the LLM sees them; the recovery store is per-process and silently disabled when virtual-key auth is absent, meaning compressed content can become unrecoverable in some deployment configurations. The core logic is carefully implemented with good security controls (SSRF mitigations, upstream-error redaction, cross-tenant store partitioning). The outstanding concern is that recovery is silently disabled without per-key auth scope while compression still runs, affecting operators who deploy the guardrail without configuring virtual keys. litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py - specifically the recovery-enabled/disabled logic around _scoped_store_key and the _warned_no_scope_recovery one-shot warning.
|
| Filename | Overview |
|---|---|
| litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py | Core guardrail: ~1180 lines implementing compression, recovery store, and agentic-loop hooks. One minor inconsistency in null-guard pattern for model_call_details in async_build_agentic_loop_plan. |
| litellm/proxy/guardrails/guardrail_hooks/compresr/init.py | Wiring module: registers initializer and class in the guardrail registry; correctly defers litellm import into the function body to avoid circular imports. |
| litellm/types/proxy/guardrails/guardrail_hooks/compresr.py | Pydantic config model with well-documented fields; correctly follows the existing GuardrailConfigModel[OptionalParams] pattern seen in other integrations. |
| litellm/types/guardrails.py | Adds COMPRESR to SupportedGuardrailIntegrations enum and mixes CompresrGuardrailConfigModel into LitellmParams; follows the exact same pattern as all other guardrail integrations. |
| litellm/llms/openai/responses/guardrail_translation/handler.py | Adds result.extend(remapped[j:]) to keep guardrail-appended tools that exceed the original_tools length; fix is minimal and has a direct regression test. |
| tests/test_litellm/proxy/guardrails/guardrail_hooks/test_compresr.py | 2002-line test suite covering compression, recovery, agentic-loop hooks, fail-open/closed, multimodal content, bypass header, and batch mismatches - all via mocks with no real network calls. |
| tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py | Adds two tests for the _merge_tools_after_guardrail fix: a unit test for the merge logic and an end-to-end async test verifying an injected tool survives the Responses API write-back path. |
Reviews (4): Last reviewed commit: "refactor(guardrails): extract _existing_..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # Compresr Guardrail — query-aware, recoverable context compression | ||
|
|
||
| [Compresr](https://compresr.ai) compresses bulky message content (tool outputs, | ||
| RAG chunks, search results) before the request reaches the LLM, cutting prompt | ||
| tokens without losing the information the model actually needs for the current |
There was a problem hiding this comment.
Per the team's rule, documentation must not be added here — it should live in the litellm-docs repo instead. Please remove this README from the main codebase and add it there.
Rule Used: Prevent documentation from being added - needs to ... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 3 · PR risk: 0/10 |
…ression Adds a first-class guardrail that compresses bulky message content (tool outputs, RAG chunks, search results) through the Compresr API before the request reaches the LLM, via the apply_guardrail / structured_messages hook so it covers /chat/completions, /v1/messages, and /v1/responses (the latter through the texts channel, mirrored only when the replacement is unambiguous; anything ambiguous is left uncompressed). Distinct from whole-conversation compressors: - Query-aware: each message is compressed against the intent that produced it (a tool output against its originating tool call's name + arguments, resolved via tool_call_id; otherwise the last user message). - Recoverable: each compressed message carries a hash marker and the request gains a compresr_retrieve tool, so the model can pull the original content back through the agentic loop when the compressed version is not enough. Originals are cached in-process, scoped to the caller's virtual-key hash plus the request's litellm_call_id, with a TTL and a per-call byte cap; recovery is skipped when no caller scope is available so one caller can never read another's originals. The store is per-process, so multi-worker deployments need sticky routing (or enable_retrieval=false). Fail-closed by default (fail_open configurable), SSRF-validated api_base (alternate IP-literal encodings included), cross-tenant-isolated recovery store, and upstream errors redacted from client-facing responses. The outbound client follows redirects and re-resolves DNS per request, so the api_base host/IP checks are defense-in-depth, not a full SSRF guarantee; this is documented as a known limitation. Requests where nothing was actually compressed are returned untouched (same object identity) so handlers skip the write-back. Auto-discovered via the guardrail_hooks registry.
431dd46 to
29411b6
Compare
The recovery store bounded bytes per call and entry count, but had no aggregate cap: 256 tracked call ids at the 10 MiB per-call default could retain ~2.5 GiB per worker. A flood of requests with distinct x-litellm-call-id values and large compressible tool outputs could exhaust a shared proxy worker. Add a global byte budget (_MAX_TOTAL_STORE_BYTES, 256 MiB) across all entries. A running total is maintained on every insert/eviction so the cap is enforced without re-encoding the whole store on the request path; oldest entries are evicted once the budget is exceeded, always keeping the most-recent entry so recovery still works for the request populating the store. +2 regression tests.
92934c5 to
672969a
Compare
Two hardening fixes to the compresr_retrieve agentic loop: 1. Only run the loop when a retrieve call resolves to recovery state this guardrail actually created for the request. Previously the gate checked only that the caller-supplied tool list contained a compresr_retrieve function and that the model emitted a call, so a caller could define their own same-named tool and force an extra provider round-trip with nothing to recover. The plan now returns run_agentic_loop=False when no requested hash resolves. 2. Bound the follow-up against retrieval amplification: each distinct hash is expanded at most once (repeats get a short marker) and at most _MAX_RETRIEVALS_PER_LOOP calls are honored, so prompting the model to call compresr_retrieve many times with the same marker cannot balloon the follow-up. _retrieve_original now returns None on miss. +3 regression tests; two existing security tests updated to assert the stronger veto behavior (forged/cross-tenant hashes now stop the loop entirely instead of returning a not-found follow-up).
df5f912 to
69f0662
Compare
…scope When enable_retrieval is on (the default) but the proxy has no per-key auth, the request has no caller scope, so recovery is silently disabled: content is compressed but the compresr_retrieve tool is never injected and the originals are dropped, with no runtime indication. Emit a one-shot call-time warning so operators can see recovery is being suppressed and configure virtual-key auth. +1 regression test.
Condense the verbose multi-line inline comments and the api_base docstring to concise form. No behavior change.
|
bugbot run |
|
@greptileai review |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit e109f51. Configure here.
| else: | ||
| merged_tools = list(existing_tools) if isinstance(existing_tools, list) else [retrieve_tool] | ||
|
|
||
| compressed_inputs["tools"] = merged_tools |
There was a problem hiding this comment.
Responses drops recovery tool
High Severity
The CompresrGuardrail adds the compresr_retrieve tool to requests. However, for the Responses API, if a request already has tools, the API's tool write-back mechanism drops the injected compresr_retrieve tool. This prevents the model from recovering compressed content, even when compression markers are present.
Reviewed by Cursor Bugbot for commit e109f51. Configure here.
| else: | ||
| merged_tools = list(existing_tools) if isinstance(existing_tools, list) else [retrieve_tool] | ||
|
|
||
| compressed_inputs["tools"] = merged_tools |
There was a problem hiding this comment.
Markers outlive byte-cap eviction
Medium Severity
Recovery markers and the compresr_retrieve tool are decided from applied.originals before _store_originals runs _bound_call_bytes. If the per-call byte cap evicts some or all hashes, those markers still ship and the tool may still be injected even though nothing retrievable remains for those hashes.
Additional Locations (2)
Reviewed by Cursor Bugbot for commit e109f51. Configure here.
There was a problem hiding this comment.
@charafkamel Can you take a look at this issue? Thanks
… markers by byte cap Two fixes for reviewer-flagged defects in the Compresr guardrail: - Responses API: _merge_tools_after_guardrail iterated only over the request's original tools, dropping any tool a guardrail appended (the compresr_retrieve recovery tool) whenever the request already had tools. Keep the appended tools so recovery works on /v1/responses. - Recovery markers: markers + originals were built for every compressed target before the per-call byte cap trimmed the store, so an evicted original left a marker the model could never retrieve. Attach recovery only while the store (existing entries under the same key + this call's originals) stays within the cap, so a shipped marker is always retrievable -- including on a later turn that reuses the store key. Adds regression tests for both paths.
…rail under the complexity gate The byte-cap fix added a branch to apply_guardrail, tipping it past the C901 complexity ceiling. Move the store lookup into a small helper; no behavior change.
|
@greptileai review |
|
@yucheng-berri can you please let me know what the status is here? is there any issue I have to resolve? |
|
merged in #33295 |
3 participants
Changes
A new Compresr guardrail that shrinks bulky context (tool outputs, RAG chunks, search results) before the request hits the model, so you spend fewer tokens with zero application changes. Runs on
apply_guardrail, so it covers/chat/completions,/v1/messages, and/v1/responses.Two things make it different from whole-conversation compressors:
compresr_retrievetool; if the model needs the original, it asks and the agentic loop feeds it back. Originals live in memory briefly, scoped per caller (no cross-tenant reads), with a TTL and byte cap.Fails closed by default (configurable), validates
api_base, redacts upstream errors, and sanitizes model-supplied input before logging. Nothing compressed -> request passed through untouched. Auto-discovered via theguardrail_hooksregistry. Defaults:latte_v2model, adaptive compression on, recovery on.6 files:
.../compresr/compresr.py.../compresr/__init__.py.../compresr/README.mdlitellm/types/proxy/guardrails/guardrail_hooks/compresr.pylitellm/types/guardrails.pytests/.../test_compresr.pySame shape as the existing guardrail integrations.
Quickstart
Note
High Risk
Sits on the critical request path, mutates messages/tools, makes outbound calls with API keys, and holds tenant-scoped originals in per-process memory (multi-worker recovery limitations); security controls are thoughtful but the blast radius is large if misconfigured.
Overview
Adds a new Compresr proxy guardrail (
guardrail: compresr) that runs on pre-call requests and calls the Compresr API to shrink large message text (tool outputs by default) before the LLM sees it. Compression is query-aware (tool results use the originating tool call intent; otherwise the last user message), with optional flags for system/history/last-user targets, adaptivelatte_v2settings, and passthroughcompression_params.When recovery is enabled (default), compressed chunks get hash markers, originals are stored in a bounded in-process cache keyed by virtual-key hash + framework call id, and a
compresr_retrievetool is injected. The guardrail hooks the agentic loop to satisfy retrieve calls and rebuild follow-up turns for chat, Anthropic, and Responses API shapes—without echoing unrelated parallel tool calls.Also wires initializer/class registries, a Pydantic config model, and extends shared guardrail types (
SupportedGuardrailIntegrations.COMPRESR,LitellmParams,unreachable_fallbackdocs). Behavior includesfail_closed/fail_openon Compresr outages, optionalx-compresr-bypass,api_basescheme/metadata checks, safe error redaction, and mirroring compression into the Responsestextschannel when unambiguous. Covered by a largetest_compresr.pysuite.Reviewed by Cursor Bugbot for commit e109f51. Bugbot is set up for automated code reviews on this repo. Configure here.