feat: add litellm.compress() — BM25-based prompt compression with retrieval tool - #25637
Conversation
Adds a compress() utility that reduces context size for LLM calls using BM25 relevance scoring (with optional semantic embeddings via litellm.embedding()). Messages below a token threshold pass through unchanged; messages above are scored, ranked, and the lowest-relevance ones replaced with stubs. Originals are cached and a retrieval tool is injected so the model can recover dropped content on demand. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…ng them When a relevant message was too large to fit in the token budget it was replaced with a stub, leaving the LLM with no real content to work with. Now the highest-scoring overflow message is truncated (first 70% + last 30% of words) to fill the remaining budget, so the LLM always receives actual content rather than just a retrieval pointer. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…kens "cook" now matches "cooking", "auth" matches "authentication", etc. Without this, short query terms scored 0 against longer inflected forms in documents, causing the wrong message to be kept. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…ess()
- test_simple_compression: parametrized test verifying BM25 routes the
right message based on query ("How to cook?" keeps cooking, "Fix auth"
keeps auth content)
- eval_compression.py: end-to-end eval harness comparing baseline vs
compressed model performance on HumanEval-style coding problems
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
|
Greptile SummaryThis PR introduces Confidence Score: 4/5Safe to merge for BM25-only usage; embedding path has an unresolved P1 (dict-style access on Pydantic v2 objects) from prior rounds that still needs fixing. Prior review rounds flagged three P1 issues (Pydantic v2 dict access in embedding_scorer, token budget/final count method mismatch, inline import) that remain unaddressed. The embedding path will raise TypeError in production. BM25-only compression is solid and well-tested. litellm/compression/scoring/embedding_scorer.py (dict-style access bug), litellm/compression/compress.py (token counting inconsistency + inline import)
|
| Filename | Overview |
|---|---|
| litellm/compression/compress.py | Main orchestrator for compress(); uses text= for budget tracking but messages= for final count — prior comment flagged this discrepancy; also has inline import and docstring mismatch noted above. |
| litellm/compression/scoring/bm25.py | Pure-Python BM25 scorer with prefix expansion; logic is correct; if term not in idf guard is unreachable dead code (all query terms are guaranteed present in idf), but benign. |
| litellm/compression/scoring/embedding_scorer.py | Uses dict-style item["embedding"] on Pydantic v2 model objects (TypeError at runtime); inline import litellm is justified to avoid circular import; both issues previously flagged. |
| tests/test_litellm/test_compression.py | 23 unit tests covering BM25, stubbing, detection, and end-to-end flow; _MockResponse uses plain dicts so the dict-style-access bug in embedding_scorer is not exercised, and test_compress_default_target has a stale comment. |
| tests/eval_swe_bench.py | SWE-bench eval harness placed in tests/; makes live litellm.completion() calls requiring real API keys — breaks CI for contributors without credentials (previously flagged). |
| litellm/compression/message_stubbing.py | Stub and truncation helpers; line-based truncation logic is clear and correct; duplicate-key handling terminates safely. |
| litellm/compression/retrieval_tool.py | Builds the OpenAI-format litellm_content_retrieve tool definition; straightforward, no issues. |
| litellm/compression/content_detection.py | Heuristic code/JSON/text classifier; lightweight and self-contained with no external dependencies. |
| litellm/types/compression.py | CompressedResult TypedDict; clean, minimal, correct. |
| litellm/init.py | Adds from .compression import compress — clean, no circular import issues at module load time. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[litellm.compress called] --> B{tokens <= trigger?}
B -- yes --> C[Return messages unchanged]
B -- no --> D[Extract last user message as query]
D --> E[bm25_score_messages]
E --> F{embedding_model set?}
F -- yes --> G[embedding_score_messages]
G --> H[_combine_scores weighted avg]
F -- no --> I[Use BM25 scores]
H --> J[Rank indices by score desc]
I --> J
J --> K[_get_protected_indices]
K --> L[Fill token budget from ranked list]
L --> M{Fits in budget?}
M -- yes --> N[Keep message as-is]
M -- no, remaining >= 100 --> O[truncate_message]
M -- no, remaining < 100 --> P[Skip candidate]
N --> Q[Build compressed_messages and cache]
O --> Q
P --> L
Q --> R{cache non-empty?}
R -- yes --> S[build_retrieval_tool]
R -- no --> T[tools is empty list]
S --> U[Return CompressedResult]
T --> U
Reviews (4): Last reviewed commit: "Revert "feat: add compression_intercepti..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| """ | ||
| Unit tests for litellm.compress(). | ||
| """ | ||
|
|
||
| import os | ||
|
|
||
| import pytest | ||
|
|
||
| import litellm | ||
| from litellm.compression.scoring.bm25 import bm25_score_messages |
There was a problem hiding this comment.
tests/test_litellm/
These tests live at tests/test_compression.py, but make test-unit (the CI gate) only discovers tests under tests/test_litellm/. Every test in this file will be invisible to the required CI check, so the checklist item "My PR passes all unit tests on make test-unit" is vacuously true here. Per CLAUDE.md, unit tests belong in tests/test_litellm/.
Context Used: CLAUDE.md (source)
| def call_llm(model: str, messages: list[dict]) -> dict: | ||
| """Call model via litellm. Returns dict with response text and usage.""" | ||
| t0 = time.time() | ||
| resp = litellm.completion( | ||
| model=model, messages=messages, temperature=0.0, max_tokens=2048 | ||
| ) | ||
| latency_ms = (time.time() - t0) * 1000 | ||
|
|
||
| text = resp.choices[0].message.content or "" | ||
| usage = resp.usage | ||
|
|
||
| return { | ||
| "text": text, | ||
| "prompt_tokens": usage.prompt_tokens, | ||
| "completion_tokens": usage.completion_tokens, | ||
| "total_tokens": usage.total_tokens, | ||
| "latency_ms": latency_ms, | ||
| } |
There was a problem hiding this comment.
tests/ directory
call_llm makes a live litellm.completion() call that requires actual API credentials and hits external endpoints. Only mock tests are permitted in the tests/ directory — real network calls break CI for contributors without credentials. Per the project rule, this file should either be moved out of tests/ (e.g., a top-level eval/ directory) or all LLM calls must be mocked.
Rule Used: What: prevent any tests from being added here that... (source)
| if embedding_model: | ||
| from litellm.compression.scoring.embedding_scorer import ( | ||
| embedding_score_messages, | ||
| ) | ||
|
|
||
| emb_scores = embedding_score_messages( | ||
| query, messages, model=embedding_model, cache=compression_cache | ||
| ) |
There was a problem hiding this comment.
Per CLAUDE.md, imports should be at module level, not inside functions. The import of embedding_score_messages is conditional on embedding_model being truthy, but this isn't a circular-import situation — litellm.compression.scoring.embedding_scorer doesn't import back into compress.py. Moving this to the top of the file is straightforward.
| if embedding_model: | |
| from litellm.compression.scoring.embedding_scorer import ( | |
| embedding_score_messages, | |
| ) | |
| emb_scores = embedding_score_messages( | |
| query, messages, model=embedding_model, cache=compression_cache | |
| ) | |
| from litellm.compression.scoring.embedding_scorer import embedding_score_messages |
(Add to the module-level import block at the top of the file and remove the inline import here.)
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!
| # Count tokens for protected messages | ||
| current_tokens = 0 | ||
| for i in kept_indices: | ||
| current_tokens += token_counter( | ||
| model=model, text=messages[i].get("content", "") or "" | ||
| ) | ||
|
|
||
| # Fill token budget from highest-scoring messages. | ||
| # For each candidate (ranked by relevance): | ||
| # - If it fits entirely → keep it as-is. | ||
| # - If it doesn't fit but there's meaningful remaining budget → truncate it | ||
| # to fill that budget (so the LLM always has real content to work with). | ||
| # - Otherwise → stub it (pointer only, content goes to cache). | ||
| # We only truncate one message (the highest-scoring one that overflows) so | ||
| # the budget is consumed and the rest are stubbed cleanly. | ||
| truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict | ||
|
|
||
| for idx in ranked_indices: | ||
| if idx in kept_indices: | ||
| continue | ||
| msg_content = messages[idx].get("content", "") or "" | ||
| msg_tokens = token_counter(model=model, text=msg_content) | ||
| remaining = compression_target - current_tokens | ||
|
|
||
| if current_tokens + msg_tokens <= compression_target: | ||
| # Fits entirely | ||
| kept_indices.add(idx) | ||
| current_tokens += msg_tokens | ||
| elif remaining >= 100 and idx not in truncated_overrides: | ||
| # Too large to fit whole, but we have budget — truncate it. | ||
| # Only do this once (the highest-scoring overflow message). | ||
| truncated = truncate_message(messages[idx], remaining) | ||
| truncated_overrides[idx] = truncated | ||
| kept_indices.add(idx) | ||
| current_tokens = compression_target # budget consumed |
There was a problem hiding this comment.
text= but final count uses messages=
The budget accumulation calls token_counter(model=model, text=...) which counts only raw content tokens, while the final compressed count uses token_counter(model=model, messages=...) which adds per-message overhead (~4 tokens/message for OpenAI models). For small compression_target values (e.g. < 1000), the accumulated overhead across many messages can cause the actual compressed_tokens to meaningfully exceed compression_target. Both loops should use the same counting method for consistency.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Uses princeton-nlp/SWE-bench_Lite_bm25_27K which bundles ~27k tokens of BM25-retrieved repo context per problem — large enough to meaningfully stress litellm.compress() without Docker or GitHub API calls. Proxy eval metrics (no test runner needed): - has_diff: model produced a valid unified diff - file_overlap: fraction of gold-patch files in generated patch - exact_file_match: generated patch touches exactly the right files Run: python tests/eval_swe_bench.py --model gpt-4o --problems 10 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Add HuggingFace API fallback so the SWE-bench loader doesn't need the `datasets` library (avoids pyarrow/numpy binary compat issues) - Insert repo root into sys.path so compression module resolves from worktrees - Use direct import of litellm_compress to avoid __getattr__ issues Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…get, 70% default target - Switch truncate_message from word-based to line-based splitting to preserve code structure (function boundaries, indentation) - Allow multiple messages to be truncated instead of burning entire budget on one overflow message - Raise default compression target from 50% to 70% of trigger for better quality/cost tradeoff - Add --compression-target CLI arg to SWE-bench eval harness - Move tests to canonical locations (tests/test_litellm/, scripts/) - Add docs page and sidebar entries for compress() Eval results (5 problems, Opus, trigger=10k): Hunk overlap delta improved from -0.417 to -0.221 Content similarity now matches baseline (+0.006) Cost savings: 72% Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Include benchmark table from Opus eval (5 problems, trigger=10k) showing 72% cost savings with file-level quality fully preserved. Add metric explanations and eval runner examples. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
krrish-berri-2
changed the title
The exact line-number matching was too brittle — LLM-generated patches often target the right code region but with slightly offset line numbers. Switch to hunk-level overlap with a 10-line tolerance window so nearby edits count as matches. This better reflects actual patch quality. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Add a proxy callback that automatically compresses incoming /v1/messages payloads above a configurable token threshold, runs the retrieval tool loop server-side, and returns the final response. This brings compress() support to proxy deployments (e.g. Claude Code via /v1/messages). - New callback: litellm/integrations/compression_interception/ - Proxy config: compression_interception_params in litellm_settings - Support for input_type param in compress() (openai vs anthropic) - Docs: proxy setup instructions with YAML config example - Tests: 139-line unit test suite for the interception handler Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| response = litellm.embedding(**kwargs) | ||
|
|
||
| # Extract embedding vectors | ||
| embeddings = [item["embedding"] for item in response.data] |
There was a problem hiding this comment.
response.data returns a list of EmbeddingObject instances (Pydantic v2 models). Pydantic v2 models do not expose __getitem__ by default, so item["embedding"] raises TypeError. Use attribute access instead:
| embeddings = [item["embedding"] for item in response.data] | |
| embeddings = [item.embedding for item in response.data] |
| kwargs_for_followup = { | ||
| k: v | ||
| for k, v in kwargs.items() | ||
| if not k.startswith("_compression_interception") | ||
| and not k.startswith("_websearch_interception") | ||
| and k != "litellm_logging_obj" | ||
| } | ||
| optional_params_without_max_tokens = { | ||
| k: v | ||
| for k, v in anthropic_messages_optional_request_params.items() | ||
| if k != "max_tokens" | ||
| } | ||
|
|
||
| max_tokens = anthropic_messages_optional_request_params.get( | ||
| "max_tokens", kwargs.get("max_tokens", 1024) | ||
| ) | ||
| full_model_name = model | ||
| if logging_obj is not None: | ||
| agentic_params = logging_obj.model_call_details.get( | ||
| "agentic_loop_params", {} | ||
| ) | ||
| full_model_name = agentic_params.get("model", model) | ||
|
|
||
| return await anthropic_messages.acreate( | ||
| max_tokens=max_tokens, | ||
| messages=follow_up_messages, | ||
| model=full_model_name, | ||
| **optional_params_without_max_tokens, | ||
| **kwargs_for_followup, | ||
| ) |
There was a problem hiding this comment.
TypeError from duplicate keyword arguments in acreate
kwargs_for_followup is constructed from kwargs with only _compression_interception*, _websearch_interception*, and litellm_logging_obj stripped out. In typical LiteLLM proxy usage kwargs also carries model, messages, and max_tokens. Those three are also passed as explicit keyword arguments to acreate, so Python raises TypeError: acreate() got multiple values for keyword argument 'messages' (and similarly for model and max_tokens).
The existing unit test avoids this by passing kwargs={}, so the path is untested with a realistic payload. At minimum, filter those keys out of kwargs_for_followup:
kwargs_for_followup = {
k: v
for k, v in kwargs.items()
if not k.startswith("_compression_interception")
and not k.startswith("_websearch_interception")
and k not in ("litellm_logging_obj", "model", "messages", "max_tokens", "tools", "stream")
}This reverts commit 72bd5cb.
26c7412
into
litellm_dev_04_13_2026_p1
…rieval tool (BerriAI#25637) * feat: add litellm.compress() for BM25-based context compression Adds a compress() utility that reduces context size for LLM calls using BM25 relevance scoring (with optional semantic embeddings via litellm.embedding()). Messages below a token threshold pass through unchanged; messages above are scored, ranked, and the lowest-relevance ones replaced with stubs. Originals are cached and a retrieval tool is injected so the model can recover dropped content on demand. * fix(compress): truncate high-scoring messages instead of fully stubbing them When a relevant message was too large to fit in the token budget it was replaced with a stub, leaving the LLM with no real content to work with. Now the highest-scoring overflow message is truncated (first 70% + last 30% of words) to fill the remaining budget, so the LLM always receives actual content rather than just a retrieval pointer. * fix(bm25): add prefix expansion so query terms match inflected doc tokens "cook" now matches "cooking", "auth" matches "authentication", etc. Without this, short query terms scored 0 against longer inflected forms in documents, causing the wrong message to be kept. * test: add routing correctness test and eval harness for litellm.compress() - test_simple_compression: parametrized test verifying BM25 routes the right message based on query ("How to cook?" keeps cooking, "Fix auth" keeps auth content) - eval_compression.py: end-to-end eval harness comparing baseline vs compressed model performance on HumanEval-style coding problems * feat(eval): add SWE-bench Lite compression eval harness Uses princeton-nlp/SWE-bench_Lite_bm25_27K which bundles ~27k tokens of BM25-retrieved repo context per problem — large enough to meaningfully stress litellm.compress() without Docker or GitHub API calls. Proxy eval metrics (no test runner needed): - has_diff: model produced a valid unified diff - file_overlap: fraction of gold-patch files in generated patch - exact_file_match: generated patch touches exactly the right files Run: python tests/eval_swe_bench.py --model gpt-4o --problems 10 * fix(eval): robust dataset loading + sys.path fix for worktree imports - Add HuggingFace API fallback so the SWE-bench loader doesn't need the `datasets` library (avoids pyarrow/numpy binary compat issues) - Insert repo root into sys.path so compression module resolves from worktrees - Use direct import of litellm_compress to avoid __getattr__ issues * improve compression quality: line-based truncation, multi-message budget, 70% default target - Switch truncate_message from word-based to line-based splitting to preserve code structure (function boundaries, indentation) - Allow multiple messages to be truncated instead of burning entire budget on one overflow message - Raise default compression target from 50% to 70% of trigger for better quality/cost tradeoff - Add --compression-target CLI arg to SWE-bench eval harness - Move tests to canonical locations (tests/test_litellm/, scripts/) - Add docs page and sidebar entries for compress() Eval results (5 problems, Opus, trigger=10k): Hunk overlap delta improved from -0.417 to -0.221 Content similarity now matches baseline (+0.006) Cost savings: 72% * docs: add SWE-bench performance results to compress() docs Include benchmark table from Opus eval (5 problems, trigger=10k) showing 72% cost savings with file-level quality fully preserved. Add metric explanations and eval runner examples. * fix(eval): use tolerance-based hunk overlap metric The exact line-number matching was too brittle — LLM-generated patches often target the right code region but with slightly offset line numbers. Switch to hunk-level overlap with a 10-line tolerance window so nearby edits count as matches. This better reflects actual patch quality. * feat: add compression_interception callback for LiteLLM Proxy Add a proxy callback that automatically compresses incoming /v1/messages payloads above a configurable token threshold, runs the retrieval tool loop server-side, and returns the final response. This brings compress() support to proxy deployments (e.g. Claude Code via /v1/messages). - New callback: litellm/integrations/compression_interception/ - Proxy config: compression_interception_params in litellm_settings - Support for input_type param in compress() (openai vs anthropic) - Docs: proxy setup instructions with YAML config example - Tests: 139-line unit test suite for the interception handler * Revert "feat: add compression_interception callback for LiteLLM Proxy" This reverts commit 72bd5cb. ---------
3 participants
Summary
Adds
litellm.compress()— a BM25-based prompt compression utility that reduces context size for LLM calls while preserving quality through intelligent message ranking and a retrieval tool for on-demand content recovery.How it works
litellm_content_retrievetool so the model can recover any stubbed content on demandKey design decisions
SWE-bench performance (Claude Opus, 5 problems, trigger=10k)
File-level targeting fully preserved. Content similarity matches baseline. 72% cost savings.
New files
litellm/compression/— core compression module (compress, BM25, embedding scorer, stubbing, retrieval tool, content detection)litellm/types/compression.py—CompressedResultTypedDicttests/test_litellm/test_compression.py— 23+ unit teststests/eval_swe_bench.py— SWE-bench eval harness with proxy metrics (file overlap, hunk overlap, content similarity)scripts/eval_compression.py— HumanEval-style eval harnessdocs/my-website/docs/completion/prompt_compression.md— docs with performance resultsTest plan
pytest tests/test_litellm/test_compression.py