Prompt Compression - add it to the proxy - #25729
Conversation
simplifies how to create logic for tool based multi llm calls
ensures claude code messages can run through proxy easily
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
|
krrish-berri-2
changed the title
Greptile SummaryThis PR adds server-side prompt compression as a proxy callback (
Confidence Score: 3/5Not safe to merge: the streaming path exposes internal tool calls to clients, contradicting the transparent-compression contract. Two P1 findings remain: the streaming iterator forwards Phase 1 tool-use blocks to clients (breaking transparency for any standard Anthropic SDK consumer), and the truncation token-count may silently undercount for Anthropic structured-content messages. Previous P0/P1 concerns (empty-tools injection, safety-guard propagation, print leaks, monkeypatch target) appear resolved, which raises the score from what would otherwise be lower. litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py (streaming transparency), litellm/compression/compress.py (truncation token-count for Anthropic structured content)
|
| Filename | Overview |
|---|---|
| litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py | New two-phase streaming iterator that yields Phase 1 SSE bytes live and chains Phase 2 on exhaustion; exposes internal litellm_content_retrieve tool calls to standard SSE clients before Phase 2 is sent, and contains a dead list comprehension (line 218–223) that computes but discards a debug summary string. |
| litellm/compression/compress.py | Adds Anthropic-Messages input support: content-block normalisation, atomic tool-exchange span detection, and call_type dispatch; the truncation token-count branch passes a list to token_counter for structured Anthropic content, which may cause silent budget miscalculation. |
| litellm/integrations/compression_interception/handler.py | New CustomLogger implementing pre-call compression, retrieval-tool injection, and agentic-loop plan building; previous issues (empty-tools injection, monkeypatch target) are fixed, in-memory cache is per-process as documented. |
| litellm/llms/custom_httpx/llm_http_handler.py | Refactors agentic-loop dispatch to typed AgenticLoopPlan; safety guards now run outside per-callback try/except (previous issue fixed); streaming path delegates to AgenticAnthropicStreamingIterator; print-statement leaks removed. |
| litellm/proxy/litellm_pre_call_utils.py | Adds generic x-<vendor>-session-id header extraction for chain-ID propagation; broad regex could collide across unrelated tenants sharing a short session ID value. |
| litellm/integrations/websearch_interception/handler.py | Ported to new AgenticLoopPlan pattern; retains async_run_agentic_loop for backward compatibility with the old hook path; no regressions found. |
| litellm/types/integrations/custom_logger.py | Adds AgenticLoopRequestPatch and AgenticLoopPlan typed models for the new hook abstraction; clean, well-typed addition. |
| litellm/proxy/common_utils/callback_utils.py | Wires up compression_interception string callback to CompressionInterceptionLogger.initialize_from_proxy_config; straightforward and well-guarded. |
Sequence Diagram
sequenceDiagram
participant C as Client
participant P as LiteLLM Proxy
participant CI as CompressionInterceptionLogger
participant M as Upstream Model
participant H as AgenticAnthropicStreamingIterator
C->>P: POST /v1/messages (large context, stream=true)
P->>CI: async_pre_call_deployment_hook
CI->>CI: compress() → stub low-score messages
CI-->>P: kwargs[messages]=compressed, kwargs[tools]=[litellm_content_retrieve]
P->>M: Compressed request (~20k tokens)
M-->>H: SSE stream Phase 1
Note over H,C: Phase 1 bytes yielded in real-time
H-->>C: SSE: message_start, content (tool_use blocks), message_delta(stop_reason=tool_use), message_stop
Note over C: Client sees litellm_content_retrieve tool_use — standard SDK stops here
H->>H: _process_agentic_hooks (rebuild response)
H->>CI: async_should_run_agentic_loop → True
H->>CI: async_build_agentic_loop_plan → AgenticLoopPlan
H->>M: Follow-up request with tool_results (Phase 2)
M-->>H: SSE stream Phase 2
H-->>C: SSE: Phase 2 events (final answer)
Note over C: Standard clients already closed stream after Phase 1 message_stop
Reviews (12): Last reviewed commit: "Merge branch 'litellm_internal_staging' ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if input_type == "anthropic_messages": | ||
| # Lazy import to avoid introducing provider transformation imports | ||
| # during module import for non-Anthropic call paths. | ||
| from litellm.llms.anthropic.chat.transformation import AnthropicConfig | ||
|
|
||
| anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools) | ||
| return cast(List[dict], anthropic_tools) | ||
|
|
||
| return openai_tools |
There was a problem hiding this comment.
llms/
_build_retrieval_tools imports and calls AnthropicConfig()._map_tools() directly, embedding Anthropic-specific transformation logic inside litellm/compression/ — outside the llms/ directory. Per the project's style guide, provider-specific code should live in llms/ so it can evolve independently. If AnthropicConfig._map_tools is renamed or its return type changes, this will silently break Anthropic-format tool injection with no indication in the compression module itself.
A cleaner approach is to return the retrieval tool in OpenAI format unconditionally and let the existing Anthropic transformation layer (which already converts tools during request building) handle the format conversion.
Rule Used: What: Avoid writing provider-specific code outside... (source)
| @@ -150,7 +395,7 @@ def compress( | |||
|
|
|||
| emb_scores = embedding_score_messages( | |||
There was a problem hiding this comment.
embedding_score_messages is imported inside the compress() function body with no circular-import justification. CLAUDE.md requires all imports to be at module level — inline imports inside functions make dependencies harder to trace. The comment in _build_retrieval_tools explains the circular-import rationale for that lazy import, but no such comment or rationale exists here.
| from litellm.compression.scoring.embedding_scorer import embedding_score_messages |
Move this to the top of the file along with the other litellm.compression.* imports.
Context Used: CLAUDE.md (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!
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Error
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
There was a problem hiding this comment.
Medium: Agentic loop depth limits silently bypassed by broad exception handler
This PR adds prompt compression with an agentic loop mechanism that reruns LLM calls server-side. The new depth and fingerprint safety checks (depth >= max_loops, repeated fingerprint detection) raise ValueError, but these are caught by the surrounding except Exception handler that only logs and continues. While the loop does terminate in practice (no agentic response is returned), the safety mechanism is degraded — callers receive the last LLM response with no indication that a safety limit was hit.
- medium: agentic loop safety checks swallowed — litellm/llms/custom_httpx/llm_http_handler.py
- low: unbounded in-memory cache — litellm/integrations/compression_interception/handler.py
| fingerprint=fingerprint, | ||
| ) | ||
|
|
||
| except Exception as e: |
There was a problem hiding this comment.
Medium: Agentic loop safety checks silently swallowed
The ValueError raised at lines 4631 and 4635 for repeated fingerprints and exceeded max_agentic_loops is caught by this broad except Exception. Instead of propagating the safety error to the caller, it's logged and discarded. The next callback in the loop (or no callback at all) then runs, and the original response is returned silently.
Consider either re-raising specific safety-check exceptions before this catch-all, or narrowing the catch to exclude ValueError:
except (ValueError) as e:
raise
except Exception as e:Alternatively, move the depth/fingerprint checks outside the per-callback try/except so they propagate correctly.
| # First hook that runs agentic loop wins | ||
| return agentic_response | ||
|
|
||
| except Exception as e: |
There was a problem hiding this comment.
Medium: Same swallowed safety check in chat completion path
Same issue as the Anthropic Messages path — the ValueError from depth/fingerprint checks at lines 4797-4803 is caught here and silently discarded. Should be handled consistently with the fix above.
GitHub Actions
Inactive
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
|
| GitGuardian id | GitGuardian status | Secret | Commit | Filename | |
|---|---|---|---|---|---|
| 29203053 | Triggered | Generic Password | 6dfad8b | .circleci/config.yml | View secret |
| 29203065 | Triggered | JSON Web Token | f123e55 | tests/test_litellm/proxy/test_litellm_pre_call_utils.py | View secret |
🛠 Guidelines to remediate hardcoded secrets
- Understand the implications of revoking this secret by investigating where it is used in your code.
- Replace and store your secrets safely. Learn here the best practices.
- Revoke and rotate these secrets.
- If possible, rewrite git history. Rewriting git history is not a trivial act. You might completely break other contributing developers' workflow and you risk accidentally deleting legitimate data.
To avoid such incidents in the future consider
- following these best practices for managing and storing secrets including API keys and other credentials
- install secret detection on pre-commit to catch secret before it leaves your machine and ease remediation.
🦉 GitGuardian detects secrets in your source code to help developers and security teams secure the modern development process. You are seeing this because you or someone else with access to this repository has authorized GitGuardian to scan your pull request.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
|
||
| import litellm | ||
| from litellm._logging import verbose_logger | ||
| from litellm.integrations.custom_logger import CustomLogger |
| AgenticLoopPlan, | ||
| AgenticLoopRequestPatch, | ||
| ) | ||
| from litellm.types.utils import CallTypes |
| from litellm.integrations.compression_interception.handler import ( | ||
| CompressionInterceptionLogger, | ||
| ) |
| if fingerprint in fingerprints: | ||
| raise ValueError( | ||
| "Agentic loop detected repeated tool-call fingerprint; aborting rerun" | ||
| ) | ||
| if depth >= max_loops: | ||
| raise ValueError( | ||
| f"Exceeded max_agentic_loops={max_loops} for model={model}" | ||
| ) |
There was a problem hiding this comment.
try/except
The fingerprint-repeat and max-loops checks raise ValueError inside the try/except Exception block that wraps each callback iteration (line 4697). This means the guard is caught, logged, and the for loop simply continues to the next callback — it never hard-stops the agentic chain. When no subsequent callback handles the tool calls (the typical single-callback case), _call_agentic_completion_hooks returns None, and the caller falls through to returning the original model response — which still contains the raw litellm_content_retrieve tool_use blocks — directly to the client.
The same pattern repeats in the chat-completion path at lines 4797–4804.
Move these checks before the try block, or raise after logging so the error propagates to the caller as a deliberate abort rather than an opaque exception log:
# Outside the per-callback try/except
if depth >= max_loops:
raise ValueError(
f"Exceeded max_agentic_loops={max_loops} for model={model}"
)
if fingerprint in fingerprints:
raise ValueError(
"Agentic loop detected repeated tool-call fingerprint; aborting rerun"
)
ishaan-berri
left a comment
ishaan-berri
left a comment
There was a problem hiding this comment.
why not use existing call types ? vs introducing a new var to maintain input_types
| compressed = litellm.compress( | ||
| messages=messages, | ||
| model="gpt-4o", | ||
| input_type="openai_chat_completions", |
There was a problem hiding this comment.
why introduce a new input_type ? We already have call types completion, messages etc
Restores the 34 HTML files under _experimental/out/ to their pre-PR paths (X/index.html -> X.html). All renames are R100 (content unchanged); no other files are touched.
GitHub Actions
Inactive
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
- Skip ``kwargs["tools"] = []`` injection when compression is a no-op — Anthropic Messages rejects empty tool arrays on requests that did not originally declare tools. - Move agentic-loop safety guards (fingerprint cycle / max depth) out of the per-callback try/except so they propagate instead of being swallowed by the generic exception handler. Extracted _check_agentic_loop_safety. - Gate generic ``x-<vendor>-session-id`` capture behind the LITELLM_CAPTURE_VENDOR_SESSION_HEADERS env var (off by default) to preserve backwards compatibility; explicit x-litellm-* headers are unaffected. - Fix monkeypatch target in pre-call-hook test to patch the actual module-level binding (litellm.integrations.compression_interception.handler.compress). - Add regression tests for empty-tools skip and opt-in session capture. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
Generic x-<vendor>-session-id header capture is a new feature and only runs *after* the explicit x-litellm-trace-id / x-litellm-session-id checks, so it does not change behavior for any existing caller that was already using the LiteLLM headers — no backwards-incompatibility to gate. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
GitHub Actions
Inactive
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
Drop the bespoke ``CompressionInputType`` literal and use the existing ``litellm.types.utils.CallTypes`` enum instead. ``litellm.compress()`` now takes ``call_type: Union[CallTypes, str]`` (default ``CallTypes.completion``) — no new concept to learn, and the enum is already the way the rest of the codebase talks about request shapes. Supported values: ``completion`` / ``acompletion`` (OpenAI chat-completions shape) and ``anthropic_messages`` (Anthropic structured content blocks). Updated: compress(), the compression_interception handler, tests, docs, and the two eval scripts. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
| [ | ||
| f"{b.get('type')}({b.get('name', '')})" | ||
| if b.get("type") == "tool_use" | ||
| else b.get("type") | ||
| for b in rebuilt.get("content", []) | ||
| ] |
There was a problem hiding this comment.
This list comprehension computes a list of content-block type strings but its result is never assigned to a variable or passed to a logger. On every agentic hook invocation the work is done and immediately garbage-collected, and the intended debug information is never surfaced.
The shape of the expression (building a human-readable block summary) strongly suggests it was meant to be passed to verbose_logger.debug(...) but the call was accidentally omitted.
| [ | |
| f"{b.get('type')}({b.get('name', '')})" | |
| if b.get("type") == "tool_use" | |
| else b.get("type") | |
| for b in rebuilt.get("content", []) | |
| ] | |
| block_summary = [ | |
| f"{b.get('type')}({b.get('name', '')})" | |
| if b.get("type") == "tool_use" | |
| else b.get("type") | |
| for b in rebuilt.get("content", []) | |
| ] | |
| verbose_logger.debug( | |
| "AgenticStreamingIterator: rebuilt response content blocks: %s", | |
| block_summary, | |
| ) |
GitHub Actions
Inactive
GitHub Actions
Inactive
| block = content_blocks.get(idx) | ||
| if block and block.get("type") == "tool_use": | ||
| partial = block.pop("_partial_json", "") | ||
| if partial: | ||
| try: | ||
| block["input"] = json.loads(partial) | ||
| except (json.JSONDecodeError, ValueError): | ||
| block["input"] = {"_raw": partial} | ||
|
|
||
|
|
||
| def _handle_message_delta(data: Dict, response: Dict) -> None: | ||
| delta = data.get("delta", {}) | ||
| if "stop_reason" in delta: | ||
| response["stop_reason"] = delta["stop_reason"] | ||
| if "stop_sequence" in delta: | ||
| response["stop_sequence"] = delta["stop_sequence"] | ||
| usage = data.get("usage", {}) | ||
| if usage.get("output_tokens") is not None: | ||
| response["usage"]["output_tokens"] = usage["output_tokens"] |
There was a problem hiding this comment.
litellm_content_retrieve tool calls to clients
AgenticAnthropicStreamingIterator unconditionally yields Phase 1 bytes — including the stop_reason: "tool_use" event with litellm_content_retrieve content blocks — to the caller before Phase 2 even begins. Standard Anthropic SDK streaming clients (e.g. Claude Code with anthropic-sdk-python) mark the stream as complete on the first message_stop event and call get_final_message(). They never read Phase 2 bytes and receive a tool_use stop response for an unknown tool, which breaks the interaction.
The PR architecture diagram explicitly states "Final answer (no tool calls visible to client)" but the streaming implementation contradicts this: the internal retrieval tool calls are streamed verbatim to the caller in Phase 1.
A transparent streaming implementation would need to buffer Phase 1 entirely (don't forward it) until it is known whether Phase 2 is needed, and only then stream the final answer to the client — at the cost of initial latency. The current approach is only viable if every downstream client is aware of and can ignore litellm_content_retrieve tool calls.
5 participants
Problem
Long-context workloads (Claude Code, RAG pipelines, document processing) frequently send requests with 50k–200k+ token inputs. This drives up cost and latency even when most of that context isn't relevant to the model's actual reasoning for the current turn.
LiteLLM already has a
compress()utility, but it was OpenAI-only and ran client-side. There was no way for a proxy operator to transparently apply compression to all traffic without clients opting in.Solution
This PR adds server-side prompt compression as a first-class proxy callback. The proxy intercepts inbound Anthropic Messages requests, compresses them using BM25/embedding scoring, and fulfills content retrieval calls server-side — all transparently, with no client changes required.
Architecture
Streaming is fully supported. The new
AgenticAnthropicStreamingIteratorwraps the SSE byte stream, yields every chunk to the client in real time, reconstructs the full response on stream exhaustion, and chains a Phase 2 stream if an agentic hook fires.What Changed
1.
compression_interceptioncallbackNew file:
litellm/integrations/compression_interception/handler.pyCompressionInterceptionLoggerimplements the full server-side flow as aCustomLogger:async_pre_call_deployment_hooklitellm_call_idasync_should_run_agentic_loopTruewhen the model response containslitellm_content_retrievetool callsasync_build_agentic_loop_plantool_resultblocks, returns typed rerun specGuard rails built in:
_agentic_loop_depth > 0(no double-compression on reruns)Enable via proxy config:
2. Typed
AgenticLoopPlanevent hookModified:
litellm/types/integrations/custom_logger.py,litellm/integrations/custom_logger.pyPreviously, the agentic loop logic for web-search interception was baked into
llm_http_handler.pywith no clean extension point. This PR abstracts it into two typed hooks onCustomLogger:AgenticLoopPlancarries arequest_patch: AgenticLoopRequestPatchthat specifies what to override on the next call (messages, tools, optional params, max_tokens).llm_http_handler.pyexecutes the plan without knowing what triggered it.Benefit: Any
CustomLoggercan now implement tool-driven multi-LLM loops without touching core handler code. Web-search interception has been ported to this pattern.3. Anthropic Messages support in
litellm.compress()Modified:
litellm/compression/compress.py,litellm/types/compression.pylitellm.compress()now acceptsinput_type: Literal["anthropic_messages", "openai_chat_completions"].Key changes for Anthropic input:
Content normalization — Anthropic messages use structured content blocks (
text,tool_use,tool_result,thinking). A new_content_to_text()helper extracts only text-bearing fields for BM25/embedding scoring while preserving the original block structure for the compressed output.Atomic tool exchange spans — An assistant
tool_useblock and the following usertool_resultblock are treated as an indivisible pair. If the scorer decides to stub the assistant message, the subsequent tool_result is stubbed too (and vice versa). This prevents the proxy from emitting malformed Anthropic message sequences that the API would reject.Tool schema remapping — The retrieval tool is built in OpenAI function-tool schema and then remapped to Anthropic's custom tool schema via
AnthropicConfig._map_tools()wheninput_type="anthropic_messages".New field:
compression_skipped_reason: Optional[str]onCompressedResult— logged at DEBUG level to explain why compression was a no-op (e.g. below token threshold, no messages eligible).4. Streaming agentic loop
New file:
litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.pyAgenticAnthropicStreamingIteratoris anAsyncIteratorthat:AnthropicMessagesResponsefrom the SSE events on exhaustion_call_agentic_completion_hooksPreviously, streaming requests bypassed the agentic loop entirely. This unblocks streaming support for both compression interception and web-search interception.
Tests
tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.pytests/test_litellm/test_compression.pyinput_typecases; atomic tool exchange span handling; budget allocation;compression_skipped_reasontests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.pytests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.pytests/test_litellm/llms/custom_httpx/test_llm_http_handler.pyllm_http_handlertests/test_litellm/proxy/common_utils/test_callback_utils.pycompression_interceptionwired up via callback utilsType
compression_interceptioncallback; streaming agentic loopAgenticLoopPlanhook abstraction; web-search ported to new patterncompress()Limitations / Follow-ups
compression_triggeris measured in tokens usingtoken_counter; the BM25 scorer operates on whitespace-split terms — semantic compression quality improves significantly withembedding_modelset.max_agentic_loops(default 3) with fingerprint-based cycle detection to prevent infinite loops on malformed tool responses.