feat: add litellm.compress() — BM25-based prompt compression with ret…

[krrish-berri-2]

…rieval tool (#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.

---

Relevant issues

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Screenshots / Proof of Fix

Type

🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test

Changes

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Apr 14, 2026 5:56pm

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing litellm_dev_04_13_2026_p1 (6d2b942) with main (4a71583)}

[greptile-apps]

Greptile Summary

This PR introduces a new litellm.compress() utility that reduces conversation token counts using BM25 relevance scoring (with optional embedding-based semantic scoring). The change is entirely additive — a new litellm/compression/ package plus one __init__.py import — so there is no risk of regression to existing functionality. Several style issues from previous review threads (inline imports, stale docstring, zip truncation on mismatched score lists) remain unaddressed.

Confidence Score: 5/5

• Safe to merge — purely additive new module with no changes to existing request paths
• All findings are P2 (style/approximation quality). The token-budget overshoot from the char/token approximation is a minor accuracy concern on the target, not a correctness or data-integrity failure. No auth, proxy, or database code is touched.
• litellm/compression/compress.py (token budget overshoot), tests/test_litellm/test_compression.py (open items from previous threads)

Important Files Changed

Filename	Overview
litellm/compression/compress.py	Core orchestration for BM25/embedding scoring and message stubbing; token budget can slightly overshoot compression_target due to character-to-token approximation in truncate_message
litellm/compression/scoring/bm25.py	Pure-Python BM25 implementation with prefix-expansion for inflected tokens; _extract_content is duplicated verbatim from embedding_scorer.py
litellm/compression/scoring/embedding_scorer.py	Cosine-similarity scoring via litellm.embedding(); lazy-imports litellm inside the function (pre-existing thread) and duplicates _extract_content with bm25.py
litellm/compression/message_stubbing.py	Message stubbing and line-based truncation logic; clean, no major issues
litellm/compression/content_detection.py	Lightweight code/JSON/text classifier using regex and indentation heuristics; straightforward and well-contained
litellm/compression/retrieval_tool.py	Builds OpenAI-format tool definition for on-demand content retrieval; simple and correct
litellm/types/compression.py	TypedDict for CompressedResult; clean and complete
tests/test_litellm/test_compression.py	Good unit test coverage; test_embedding_scorer still makes real API calls when OPENAI_API_KEY is set (flagged in previous thread); test_compress_default_target docstring still documents 50% target despite 70% default (flagged in previous thread)
litellm/init.py	Single-line addition exposing litellm.compress(); correct and minimal

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[litellm.compress called] --> B{total tokens\n> compression_trigger?}
    B -- No --> C[Return messages unchanged\ncompression_ratio=0.0]
    B -- Yes --> D[Extract last user message\nas BM25 query]
    D --> E[bm25_score_messages]
    E --> F{embedding_model\nprovided?}
    F -- Yes --> G[embedding_score_messages\nvia litellm.embedding]
    G --> H[_combine_scores\nweighted average]
    F -- No --> I[Use BM25 scores only]
    H --> J[Sort indices by score desc]
    I --> J
    J --> K[_get_protected_indices\nsystem + last user + last assistant]
    K --> L[Fill token budget\nfrom highest-scoring messages]
    L --> M{Message fits\nin remaining budget?}
    M -- Yes --> N[Keep as-is]
    M -- No, remaining>=100 --> O[truncate_message\nfirst 70% + last 30% of lines]
    M -- No, budget full --> P[stub_message\npointer stub + cache original]
    N --> Q[Build compressed_messages list]
    O --> Q
    P --> Q
    Q --> R[build_retrieval_tool\nif cache non-empty]
    R --> S[Return CompressedResult\nmessages, cache, tools, ratios]

Loading

_{Reviews (2): Last reviewed commit: "fix(mypy): resolve type errors in compre..." | Re-trigger Greptile}

[greptile-apps]

Real network call in tests/test_litellm/

This test calls litellm.compress() with embedding_model="text-embedding-3-small" which issues a real litellm.embedding() call to the OpenAI API when OPENAI_API_KEY is in the environment. All tests in tests/test_litellm/ must be mock-only — this will make real API calls in any CI environment where the secret is available and can add latency/cost.

Replace with a fully-mocked version (as done in test_compress_forwards_embedding_model_params and test_embedding_scorer_forwards_embedding_model_params) or move this to an integration test directory.

Rule Used: What: prevent any tests from being added here that... (source)

[greptile-apps]

Docstring says // 2 but code computes * 7 // 10

The docstring for compression_target states "Defaults to compression_trigger // 2" (50%), but the actual default is compression_trigger * 7 // 10 (70%). Callers relying on the docstring to reason about expected output sizes will be surprised by the real behaviour.

Suggested change

	embedding_model: If provided, use BM25 + embeddings for scoring.
	If ``None``, BM25 only.
	embedding_model_params: Optional kwargs forwarded to
	compression_target: Target token count after compression.
	Defaults to ``70% of compression_trigger``
	(i.e. ``compression_trigger * 7 // 10``).

[greptile-apps]

Stale comment contradicts actual default

The comment says compression_trigger // 2 (50%) but the current default is compression_trigger * 7 // 10 (70%). The test itself doesn't pin the compression_target argument, so the assertion is testing 70%, not 50%.

Suggested change

	{"role": "user", "content": "query"},
	"""compression_target defaults to 70% of compression_trigger."""

[greptile-apps]

Inline import inside function body

The from litellm.compression.scoring.embedding_scorer import embedding_score_messages import is placed inside the if embedding_model: branch rather than at module level. Per the project's CLAUDE.md style guide, imports belong at the top of the file; the only accepted exception is where a module-level import would create a circular dependency. No circular dependency exists here — embedding_scorer.py already guards the litellm import itself — so this can safely be moved to the top of the file.

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!

[greptile-apps]

Inline import litellm inside function body

import litellm is placed inside embedding_score_messages rather than at module level. If this is intentional to avoid a circular import (since litellm/__init__.py imports from litellm.compression), a brief comment should explain it; otherwise the import should be promoted to module level per the project style guide.

Suggested change

	import litellm
	# `import litellm` is intentionally deferred here to avoid a circular
	# import: litellm/__init__.py -> litellm.compression -> embedding_scorer.
	import litellm

Context Used: CLAUDE.md (source)

[greptile-apps]

zip silently truncates when score lists differ in length

If embedding_score_messages ever returns a shorter list than bm25_score_messages (e.g. because the embedding API batches and drops an item), the zip in the return expression will silently produce a combined list that is shorter than messages. Downstream code that iterates ranked_indices by index into messages would then mis-sort or entirely omit some messages without any error.

A length assertion before the zip would make this detectable:

assert len(bm25_scores) == len(emb_scores) == len(messages), (
    f"Score list length mismatch: bm25={len(bm25_scores)}, "
    f"emb={len(emb_scores)}, messages={len(messages)}"
)

[qdrddr]

Would be great if LiteLLM AI Gateway replace PGVector extention with VectorChord.
PGVector supports only 2k dimensions and does not support BM25 & HybridSearch.
While VectorChord PG extension supports up to 64k dimensions and has a build-in BM25 and Hybrid search

BerriAI/litellm-pgvector#9

FYI @krrish-berri-2