ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #25650 from BerriAI/litellm_dev_04_13_2026_p1

Browse Source 
feat: add litellm.compress() — BM25-based prompt compression with ret…
This commit is contained in:
ishaan-berri 2026-04-14 12:24:47 -07:00 committed by GitHub
commit 0e43050a01
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 3065 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

123
docs/my-website/docs/completion/prompt_compression.md Normal file
View File

@ -0,0 +1,123 @@
# Prompt Compression (`compress()`)
Use `litellm.compress()` to shrink long conversation history before calling `completion()`.
The function keeps high-relevance and recent context, replaces low-relevance content with lightweight stubs, and returns a retrieval tool so the model can request full content only when needed.
## Quickstart
```python
import litellm
messages = [
{"role": "system", "content": "You are a coding assistant."},
{"role": "user", "content": "# auth.py\n" + "def authenticate():\n pass\n" * 2000},
{"role": "user", "content": "# utils.py\n" + "def helper():\n pass\n" * 2000},
{"role": "user", "content": "Fix the bug in auth.py"},
]
compressed = litellm.compress(
messages=messages,
model="gpt-4o",
compression_trigger=1000,
compression_target=500,
)
response = litellm.completion(
model="gpt-4o",
messages=compressed["messages"],
tools=compressed["tools"],
)
```
## What It Returns
`compress()` returns a dictionary with:
- `messages`: compressed conversation messages
- `original_tokens`: token count before compression
- `compressed_tokens`: token count after compression
- `compression_ratio`: fraction of tokens removed
- `cache`: key-value mapping of stub key -> original full content
- `tools`: retrieval tool definition (`litellm_content_retrieve`) for on-demand restoration
## Parameters
- `messages` (`List[dict]`, required): input conversation messages
- `model` (`str`, required): model name used for token counting
- `compression_trigger` (`int`, default `200000`): compress only if input token count exceeds this
- `compression_target` (`Optional[int]`, default `70% of compression_trigger`): desired post-compression token budget
- `embedding_model` (`Optional[str]`): if set, combines BM25 + embedding relevance scoring
- `embedding_model_params` (`Optional[dict]`): additional kwargs passed to `litellm.embedding()`
- `compression_cache` (`Optional[DualCache]`): optional cache used by embedding scoring
## Behavior Notes
- Messages below `compression_trigger` are passed through unchanged.
- System messages, the last user message, and the last assistant message are always preserved.
- If a relevant message does not fully fit the remaining budget, `compress()` may keep a truncated version of it.
- Compressed-out content is never lost; it is stored in `cache` and addressable by `litellm_content_retrieve`.
## Handling Retrieval Tool Calls
If the model calls `litellm_content_retrieve`, look up the requested key in `compressed["cache"]` and return that value as tool output.
```python
import json
tool_call = response.choices[0].message.tool_calls[0]
args = json.loads(tool_call.function.arguments)
full_content = compressed["cache"][args["key"]]
```
## Performance
Benchmarked on [SWE-bench Lite](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite_bm25_27K) (real GitHub issues with ~27k tokens of BM25-retrieved repo context per problem).
### Claude Opus — 5 problems, trigger=10k
| Metric | Baseline | Compressed | Delta |
|---|---|---|---|
| File overlap | 1.000 | 1.000 | +0.000 |
| Exact file match | 100% | 100% | +0.0% |
| Hunk overlap | 0.582 | 0.361 | -0.221 |
| Content similarity | 0.367 | 0.373 | +0.006 |
| Avg prompt tokens | 30,828 | 6,890 | -77.7% |
| Avg cost/problem | $0.488 | $0.136 | **-72.0%** |
**Key takeaways:**
- **File-level targeting is fully preserved** — the model edits the same files with or without compression.
- **Content similarity matches baseline** — the actual lines changed are comparable.
- **Hunk overlap drops modestly** (-0.221) — the model targets the right files but may edit slightly different line ranges with less surrounding context.
- **72% cost savings** with 78% token reduction.
### Metrics explained
| Metric | What it measures |
|---|---|
| **File overlap** | Fraction of gold-patch files present in the generated patch |
| **Exact file match** | Whether the generated patch touches exactly the same set of files |
| **Hunk overlap** | Fraction of gold hunk line ranges covered by generated hunks |
| **Content similarity** | Jaccard similarity of changed lines (added/removed) between gold and generated patches |
### Running the SWE-bench eval
```bash
# 5-problem quick check
python tests/eval_swe_bench.py --model claude-opus-4-20250514 --problems 5
# Custom trigger/target
python tests/eval_swe_bench.py --model gpt-4o --problems 20 \
--compression-trigger 15000 --compression-target 10000
# With embedding scoring
python tests/eval_swe_bench.py --model gpt-4o --problems 10 \
--embedding-model text-embedding-3-small
```
### Running the HumanEval-style eval
```bash
python scripts/eval_compression.py --model gpt-4o --problems 5
```

7
docs/my-website/package-lock.json generated
View File

@ -20403,6 +20403,13 @@
"url": "https://opencollective.com/webpack"
}
},
"node_modules/search-insights": {
"version": "2.17.3",
"resolved": "https://registry.npmjs.org/search-insights/-/search-insights-2.17.3.tgz",
"integrity": "sha512-RQPdCYTa8A68uM2jwxoY842xDhvx3E5LFL1LxvxCNMev4o5mLuokczhzjAgGwUZBAmOKZknArSxLKmXtIi2AxQ==",
"license": "MIT",
"peer": true
},
"node_modules/section-matter": {
"version": "1.0.0",
"resolved": "https://registry.npmjs.org/section-matter/-/section-matter-1.0.0.tgz",

6
docs/my-website/sidebars.js
View File

@ -254,6 +254,11 @@ const sidebars = {
id: "image_generation",
label: "image_generation()",
},
{
type: "doc",
id: "completion/prompt_compression",
label: "compress()",
},
{
type: "doc",
id: "audio_transcription",
@ -1280,6 +1285,7 @@ const learnSidebar = {
items: [
"completion/prefix",
"completion/predict_outputs",
"completion/prompt_compression",
"completion/message_trimming",
"completion/prompt_caching",
"completion/prompt_formatting",

1
litellm/__init__.py
View File

@ -1176,6 +1176,7 @@ from litellm.types.utils import LlmProviders
## Lazy loading this is not straightforward, will leave it here for now.
from .main import * # type: ignore
from .compression import compress # type: ignore[no-redef]
# Skills API
from .skills.main import (

3
litellm/compression/__init__.py Normal file
View File

@ -0,0 +1,3 @@
from litellm.compression.compress import compress
__all__ = ["compress"]

255
litellm/compression/compress.py Normal file
View File

@ -0,0 +1,255 @@
"""
Main compress() function — orchestrates BM25/embedding scoring, message stubbing,
and retrieval tool injection.
"""
from typing import Any, Dict, List, Optional, Set, Union, cast
from litellm.caching.dual_cache import DualCache
from litellm.compression.message_stubbing import (
extract_key,
stub_message,
truncate_message,
)
from litellm.compression.retrieval_tool import build_retrieval_tool
from litellm.compression.scoring.bm25 import bm25_score_messages
from litellm.litellm_core_utils.token_counter import token_counter
from litellm.types.compression import CompressedResult
from litellm.types.utils import AllMessageValues, Message
def _extract_last_user_message(messages: List[dict]) -> str:
"""Return the text content of the last user message."""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content", "")
if isinstance(content, str):
return content
if isinstance(content, list):
parts = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
parts.append(part.get("text", ""))
elif isinstance(part, str):
parts.append(part)
return " ".join(parts)
return ""
def _get_protected_indices(messages: List[dict]) -> List[int]:
"""
Return indices of messages that must never be compressed:
- All system messages
- The last user message
- The last assistant message
"""
protected: List[int] = []
last_user_idx = None
last_assistant_idx = None
for i, msg in enumerate(messages):
role = msg.get("role", "")
if role == "system":
protected.append(i)
elif role == "user":
last_user_idx = i
elif role == "assistant":
last_assistant_idx = i
if last_user_idx is not None:
protected.append(last_user_idx)
if last_assistant_idx is not None:
protected.append(last_assistant_idx)
return protected
def _combine_scores(
bm25_scores: List[float],
emb_scores: List[float],
bm25_weight: float = 0.4,
) -> List[float]:
"""Weighted average of BM25 and embedding scores, with min-max normalization."""
def _normalize(scores: List[float]) -> List[float]:
min_s = min(scores) if scores else 0.0
max_s = max(scores) if scores else 0.0
rng = max_s - min_s
if rng == 0:
return [0.0] * len(scores)
return [(s - min_s) / rng for s in scores]
norm_bm25 = _normalize(bm25_scores)
norm_emb = _normalize(emb_scores)
emb_weight = 1.0 - bm25_weight
return [bm25_weight * b + emb_weight * e for b, e in zip(norm_bm25, norm_emb)]
def compress(
messages: List[dict],
model: str,
compression_trigger: int = 200_000,
compression_target: Optional[int] = None,
embedding_model: Optional[str] = None,
embedding_model_params: Optional[Dict[str, Any]] = None,
compression_cache: Optional[DualCache] = None,
) -> CompressedResult:
"""
Compress a list of messages by replacing low-relevance content with stubs.
Messages below ``compression_trigger`` tokens pass through unchanged.
Messages above are scored with BM25 (and optionally embeddings), ranked,
and the lowest-relevance messages are replaced with stubs. Originals are
cached and a retrieval tool is injected so the model can recover dropped
content on demand.
Parameters:
messages: The conversation messages to (potentially) compress.
model: The LLM model name — used for token counting.
compression_trigger: Only compress if input exceeds this token count.
compression_target: Target token count after compression.
Defaults to ``compression_trigger // 2``.
embedding_model: If provided, use BM25 + embeddings for scoring.
If ``None``, BM25 only.
embedding_model_params: Optional kwargs forwarded to
``litellm.embedding()`` when ``embedding_model`` is set.
compression_cache: Passed through to ``litellm.embedding()`` for
cross-turn caching of embedding vectors.
Returns:
A ``CompressedResult`` dict containing compressed messages, token
counts, a cache of original content, and the retrieval tool definition.
"""
if compression_target is None:
compression_target = compression_trigger * 7 // 10
original_tokens = token_counter(
model=model, messages=cast(List[Union[AllMessageValues, Message]], messages)
)
# Pass through if below trigger
if original_tokens <= compression_trigger:
return CompressedResult(
messages=messages,
original_tokens=original_tokens,
compressed_tokens=original_tokens,
compression_ratio=0.0,
cache={},
tools=[],
)
# Extract query for relevance scoring
query = _extract_last_user_message(messages)
# Score each message
bm25_scores = bm25_score_messages(query, messages)
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,
embedding_model_params=embedding_model_params,
)
combined_scores = _combine_scores(bm25_scores, emb_scores, bm25_weight=0.4)
else:
combined_scores = bm25_scores
# Sort message indices by score descending
ranked_indices = sorted(
range(len(messages)),
key=lambda i: combined_scores[i],
reverse=True,
)
# Protected messages are never compressed
protected_indices = _get_protected_indices(messages)
kept_indices: Set[int] = set(protected_indices)
# 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 as much of the budget as possible.
# - Otherwise → stub it (pointer only, content goes to cache).
# Multiple messages may be truncated so we preserve partial content from
# several high-scoring messages rather than fully stubbing all but one.
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 remaining <= 0:
break # budget exhausted
if current_tokens + msg_tokens <= compression_target:
# Fits entirely
kept_indices.add(idx)
current_tokens += msg_tokens
elif remaining >= 100:
# Too large to fit whole, but we have budget — truncate it.
truncated = truncate_message(messages[idx], remaining)
truncated_tokens = token_counter(
model=model,
text=truncated.get("content", "") or "",
)
truncated_overrides[idx] = truncated
kept_indices.add(idx)
current_tokens += truncated_tokens
# Build compressed messages and cache
compressed_messages: List[dict] = []
cache: Dict[str, str] = {}
used_keys: Set[str] = set()
for i, msg in enumerate(messages):
if i in kept_indices:
# Use the truncated version if we made one, otherwise the original
compressed_messages.append(truncated_overrides.get(i, msg))
else:
key = extract_key(msg, fallback_index=i, used_keys=used_keys)
content = msg.get("content", "")
if isinstance(content, list):
content = " ".join(
p.get("text", "") if isinstance(p, dict) else str(p)
for p in content
)
cache[key] = content
compressed_messages.append(stub_message(msg, key))
# Build retrieval tool
tools = [build_retrieval_tool(list(cache.keys()))] if cache else []
compressed_tokens = token_counter(
model=model,
messages=cast(List[Union[AllMessageValues, Message]], compressed_messages),
)
return CompressedResult(
messages=compressed_messages,
original_tokens=original_tokens,
compressed_tokens=compressed_tokens,
compression_ratio=round(1 - (compressed_tokens / original_tokens), 4)
if original_tokens > 0
else 0.0,
cache=cache,
tools=tools,
)

45
litellm/compression/content_detection.py Normal file
View File

@ -0,0 +1,45 @@
"""
Auto-detect content type per message: code, JSON, or text.
"""
import json
import re
_CODE_KEYWORDS = re.compile(
r"\b(?:def |function |class |import |from |require\(|#include|fn |func |const |let |var |public |private |static )\b"
)
def detect_content_type(content: str) -> str:
"""
Detect whether content is code, JSON, or plain text.
Returns one of: "code", "json", "text"
"""
stripped = content.strip()
if not stripped:
return "text"
# Check JSON
if stripped[0] in ("{", "["):
try:
json.loads(stripped)
return "json"
except (json.JSONDecodeError, ValueError):
pass
# Check code indicators
# Sample first 5000 chars for performance
sample = stripped[:5000]
keyword_matches = len(_CODE_KEYWORDS.findall(sample))
lines = sample.split("\n")
indented_lines = sum(
1 for line in lines if line.startswith((" ", "\t")) and line.strip()
)
# If we see multiple code keywords or significant indentation, it's likely code
if keyword_matches >= 3 or (indented_lines > len(lines) * 0.3 and len(lines) > 5):
return "code"
return "text"

120
litellm/compression/message_stubbing.py Normal file
View File

@ -0,0 +1,120 @@
"""
Replace messages with compact stubs and extract human-readable keys.
"""
import re
from typing import Set
from litellm.compression.content_detection import detect_content_type
# Patterns for extracting file paths from content
_FILE_PATH_PATTERNS = [
re.compile(r"^#\s*(\S+\.\w+)", re.MULTILINE), # # filename.py
re.compile(r"^//\s*(\S+\.\w+)", re.MULTILINE), # // filename.js
re.compile(r"^File:\s*(\S+)", re.MULTILINE), # File: path/to/file
re.compile(r"^---\s*(\S+\.\w+)", re.MULTILINE), # --- filename.ext
re.compile(r"`(\S+\.\w{1,5})`"), # `filename.ext` in backticks
]
def extract_key(message: dict, fallback_index: int, used_keys: Set[str]) -> str:
"""
Extract a human-readable key for the message.
Looks for file path patterns in the content. Falls back to message_{index}.
Handles duplicates by appending _2, _3, etc.
"""
content = message.get("content", "")
if isinstance(content, list):
content = " ".join(
p.get("text", "") if isinstance(p, dict) else str(p) for p in content
)
key = None
for pattern in _FILE_PATH_PATTERNS:
match = pattern.search(content[:2000]) # Only search the beginning
if match:
# Use just the filename, not full path
path = match.group(1)
key = path.split("/")[-1]
break
if key is None:
key = f"message_{fallback_index}"
# Handle duplicates
base_key = key
counter = 2
while key in used_keys:
key = f"{base_key}_{counter}"
counter += 1
used_keys.add(key)
return key
def stub_message(message: dict, key: str) -> dict:
"""
Replace message content with a compact stub.
Returns a new message dict with the same role but content replaced
with a short description referencing the retrieval tool.
"""
content = message.get("content", "")
if isinstance(content, list):
content = " ".join(
p.get("text", "") if isinstance(p, dict) else str(p) for p in content
)
line_count = content.count("\n") + 1
content_type = detect_content_type(content)
stub_content = (
f"[Compressed: {key} — {line_count} lines, {content_type}. "
f"Use litellm_content_retrieve tool to get full content.]"
)
return {**message, "content": stub_content}
def truncate_message(message: dict, max_tokens: int) -> dict:
"""
Truncate a message's content to approximately max_tokens by keeping
the first 70% and last 30% of lines with a separator in between.
Uses line-based splitting to preserve code structure (function
boundaries, indentation) rather than word-based splitting which
mangles code.
Used when a message is too large to fit entirely in the budget but
too relevant to fully stub out.
"""
content = message.get("content", "")
if isinstance(content, list):
content = " ".join(
p.get("text", "") if isinstance(p, dict) else str(p) for p in content
)
# Rough conversion: 1 token ≈ 3 characters
target_chars = max(100, max_tokens * 3)
if len(content) <= target_chars:
return {**message, "content": content}
lines = content.split("\n")
# Estimate target line count from character budget
avg_line_len = max(1, len(content) // max(1, len(lines)))
target_lines = max(2, target_chars // avg_line_len)
if len(lines) <= target_lines:
return {**message, "content": content}
first_count = (target_lines * 7) // 10
last_count = target_lines - first_count
truncated = (
"\n".join(lines[:first_count])
+ "\n...[truncated for context window]...\n"
+ "\n".join(lines[-last_count:])
)
return {**message, "content": truncated}

35
litellm/compression/retrieval_tool.py Normal file
View File

@ -0,0 +1,35 @@
"""
Build the litellm_content_retrieve tool definition for the LLM.
"""
from typing import List
def build_retrieval_tool(available_keys: List[str]) -> dict:
"""
Return an OpenAI-format tool definition that lets the model
retrieve the full content of a compressed message.
"""
return {
"type": "function",
"function": {
"name": "litellm_content_retrieve",
"description": (
"Retrieve the full content of a file or message that was "
"compressed to save tokens. Use this when you need the complete "
"content to answer accurately. Available keys: "
+ ", ".join(available_keys)
),
"parameters": {
"type": "object",
"properties": {
"key": {
"type": "string",
"description": "The identifier of the content to retrieve",
"enum": available_keys,
}
},
"required": ["key"],
},
},
}

4
litellm/compression/scoring/__init__.py Normal file
View File

@ -0,0 +1,4 @@
from litellm.compression.scoring.bm25 import bm25_score_messages
from litellm.compression.scoring.embedding_scorer import embedding_score_messages
__all__ = ["bm25_score_messages", "embedding_score_messages"]

123
litellm/compression/scoring/bm25.py Normal file
View File

@ -0,0 +1,123 @@
"""
Pure Python BM25 (Okapi BM25) relevance scorer.
No external dependencies — uses only stdlib.
"""
import math
import re
from collections import Counter
from typing import Dict, List
def _tokenize(text: str) -> List[str]:
"""Split text into lowercase tokens on word boundaries."""
return re.findall(r"[a-z0-9_]+", text.lower())
def _extract_content(message: dict) -> str:
"""Extract text content from a message dict."""
content = message.get("content", "")
if isinstance(content, str):
return content
if isinstance(content, list):
parts = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
parts.append(part.get("text", ""))
elif isinstance(part, str):
parts.append(part)
return " ".join(parts)
return ""
def bm25_score_messages(
query: str,
messages: List[dict],
k1: float = 1.5,
b: float = 0.75,
) -> List[float]:
"""
Score each message's relevance to the query using BM25 (Okapi BM25).
Parameters:
query: The reference text to score against (typically the last user message).
messages: List of message dicts with "content" fields.
k1: Term frequency saturation parameter.
b: Length normalization parameter.
Returns:
List of float scores, one per message. Higher = more relevant.
"""
query_terms = _tokenize(query)
if not query_terms:
return [0.0] * len(messages)
# Tokenize all documents
doc_tokens: List[List[str]] = []
for msg in messages:
doc_tokens.append(_tokenize(_extract_content(msg)))
n = len(doc_tokens)
if n == 0:
return []
# Average document length
doc_lengths = [len(dt) for dt in doc_tokens]
avgdl = sum(doc_lengths) / n if n > 0 else 1.0
# Document frequency for each term
df: Dict[str, int] = {}
for dt in doc_tokens:
seen = set(dt)
for term in seen:
df[term] = df.get(term, 0) + 1
# IDF for query terms
idf: Dict[str, float] = {}
for term in set(query_terms):
term_df = df.get(term, 0)
# Standard BM25 IDF: log((N - df + 0.5) / (df + 0.5) + 1)
idf[term] = math.log((n - term_df + 0.5) / (term_df + 0.5) + 1.0)
# Build a prefix-expansion map per document: for each query term, find all
# document tokens that start with that term (min 4 chars match). This lets
# "cook" match "cooking" and "auth" match "authentication" without a full
# stemmer dependency.
def _expand_tf(query_term: str, tf_counts: Counter) -> int: # type: ignore[type-arg]
"""Sum TF across all doc tokens that are prefixed by query_term."""
exact = tf_counts.get(query_term, 0)
if exact:
return exact
if len(query_term) < 4:
return 0
return sum(
count
for token, count in tf_counts.items()
if token != query_term and token.startswith(query_term)
)
# Score each document
scores: List[float] = []
for i, dt in enumerate(doc_tokens):
if not dt:
scores.append(0.0)
continue
tf_counts = Counter(dt)
dl = doc_lengths[i]
score = 0.0
for term in query_terms:
if term not in idf:
continue
tf = _expand_tf(term, tf_counts)
if tf == 0:
continue
numerator = tf * (k1 + 1)
denominator = tf + k1 * (1 - b + b * dl / avgdl)
score += idf[term] * numerator / denominator
scores.append(score)
return scores

95
litellm/compression/scoring/embedding_scorer.py Normal file
View File

@ -0,0 +1,95 @@
"""
Semantic scoring via litellm.embedding().
Computes cosine similarity between the query embedding and each message embedding.
"""
import math
from typing import Any, Dict, List, Optional
from litellm.caching.dual_cache import DualCache
def _extract_content(message: dict) -> str:
"""Extract text content from a message dict."""
content = message.get("content", "")
if isinstance(content, str):
return content
if isinstance(content, list):
parts = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
parts.append(part.get("text", ""))
elif isinstance(part, str):
parts.append(part)
return " ".join(parts)
return ""
def _truncate_text(text: str, max_chars: int = 30000) -> str:
"""Truncate long text, keeping first and last portions."""
if len(text) <= max_chars:
return text
half = max_chars // 2
return text[:half] + "\n...\n" + text[-half:]
def _cosine_similarity(a: List[float], b: List[float]) -> float:
"""Compute cosine similarity between two vectors."""
dot = sum(x * y for x, y in zip(a, b))
norm_a = math.sqrt(sum(x * x for x in a))
norm_b = math.sqrt(sum(x * x for x in b))
if norm_a == 0 or norm_b == 0:
return 0.0
return dot / (norm_a * norm_b)
def embedding_score_messages(
query: str,
messages: List[dict],
model: str,
cache: Optional[DualCache] = None,
embedding_model_params: Optional[Dict[str, Any]] = None,
) -> List[float]:
"""
Score each message's semantic similarity to the query using embeddings.
Parameters:
query: The reference text to score against.
messages: List of message dicts with "content" fields.
model: The embedding model to use (e.g., "text-embedding-3-small").
cache: Optional DualCache for cross-turn embedding caching.
embedding_model_params: Optional additional kwargs forwarded to
``litellm.embedding()``.
Returns:
List of float scores (cosine similarity), one per message.
"""
import litellm
texts = [_truncate_text(query)]
for msg in messages:
texts.append(_truncate_text(_extract_content(msg)))
# Filter out empty texts — replace with a placeholder to maintain indexing
processed_texts = [t if t.strip() else "empty" for t in texts]
kwargs: Dict[str, Any] = {
"model": model,
"input": processed_texts,
"caching": cache is not None,
}
if embedding_model_params:
kwargs = {**kwargs, **embedding_model_params}
response = litellm.embedding(**kwargs)
# Extract embedding vectors
embeddings = [item["embedding"] for item in response.data]
query_embedding = embeddings[0]
scores: List[float] = []
for i in range(1, len(embeddings)):
scores.append(_cosine_similarity(query_embedding, embeddings[i]))
return scores

14
litellm/types/compression.py Normal file
View File

@ -0,0 +1,14 @@
"""
Type definitions for litellm.compress().
"""
from typing import Dict, List, TypedDict
class CompressedResult(TypedDict):
messages: List[dict] # compressed messages (stubs replace low-relevance messages)
original_tokens: int # token count before compression
compressed_tokens: int # token count after compression
compression_ratio: float # fraction reduced, e.g. 0.6 means 60% reduction
cache: Dict[str, str] # key -> original content (for retrieval tool responses)
tools: List[dict] # [litellm_content_retrieve tool definition]

1125
scripts/eval_compression.py Normal file
View File

File diff suppressed because it is too large Load Diff

751
tests/eval_swe_bench.py Normal file
View File

@ -0,0 +1,751 @@
"""
SWE-bench Compression Evaluation
==================================
Measures litellm.compress() impact on SWE-bench Lite problems.
Each instance includes ~27k tokens of BM25-retrieved repo context — large
enough to meaningfully stress compression without requiring Docker or GitHub
API calls.
Usage:
python tests/eval_swe_bench.py --model gpt-4o --problems 10
python tests/eval_swe_bench.py --model claude-sonnet-4-20250514 --problems 25
python tests/eval_swe_bench.py --model gpt-4o-mini --problems 50 --compression-trigger 8000
Requires:
pip install datasets
Proxy eval metrics (no Docker / test runner required):
- has_diff: model produced a valid unified diff
- file_overlap: fraction of gold-patch files present in generated patch
- exact_file_match: generated patch touches exactly the same files as gold patch
Full SWE-bench pass rate (FAIL_TO_PASS) requires the official evaluation
harness with Docker — not in scope here. The proxy metrics are a lightweight
signal for whether compression degrades patch quality.
"""
import argparse
import json
import os
import re
import statistics
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Optional
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import litellm # noqa: E402
from litellm.compression import compress as litellm_compress # noqa: E402
# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------
SYSTEM_MSG = (
"You are an expert software engineer resolving GitHub issues. "
"You will be given an issue description and relevant source files. "
"Produce a minimal unified diff patch that fixes the issue. "
"Your response must contain ONLY the patch in unified diff format. "
"Start with `diff --git a/path b/path`, then `---`, `+++`, and "
"`@@` hunks. Do NOT include any explanation, commentary, or markdown "
"fences — just the raw diff text."
)
# ---------------------------------------------------------------------------
# Dataset loading
# ---------------------------------------------------------------------------
def _load_via_datasets(n: int, split: str) -> list[dict]:
"""Load via the HuggingFace `datasets` library (preferred if available)."""
from datasets import load_dataset
ds = load_dataset("princeton-nlp/SWE-bench_Lite_bm25_27K", split=split)
problems = []
for i, item in enumerate(ds):
if n > 0 and i >= n:
break
problems.append(dict(item))
return problems
def _load_via_api(n: int, split: str) -> list[dict]:
"""Fallback: fetch rows directly from the HuggingFace dataset API (no deps).
The API returns at most 100 rows per request, so we paginate.
"""
import json
import urllib.request
# 0 means "all" — SWE-bench Lite has 300 test instances
target = n if n > 0 else 300
page_size = 100
all_rows: list[dict] = []
for offset in range(0, target, page_size):
length = min(page_size, target - offset)
url = (
"https://datasets-server.huggingface.co/rows"
"?dataset=princeton-nlp/SWE-bench_Lite_bm25_27K"
f"&config=default&split={split}&offset={offset}&length={length}"
)
req = urllib.request.Request(url, headers={"User-Agent": "litellm-eval"})
with urllib.request.urlopen(req, timeout=60) as resp:
data = json.loads(resp.read().decode())
rows = [row["row"] for row in data["rows"]]
all_rows.extend(rows)
if len(rows) < length:
break # no more data
return all_rows
def load_problems(n: int = 10, split: str = "test") -> list[dict]:
"""Load n problems from princeton-nlp/SWE-bench_Lite_bm25_27K."""
print("Loading SWE-bench_Lite_bm25_27K ...", flush=True)
# Try the HuggingFace API first — it's pure HTTP with no native deps,
# so it never triggers pyarrow/numpy binary incompatibilities that can
# poison the process. Fall back to the `datasets` library only if the
# API call fails.
try:
problems = _load_via_api(n, split)
except Exception:
try:
problems = _load_via_datasets(n, split)
except Exception as e:
print(f"ERROR: Could not load dataset ({type(e).__name__}: {e})")
sys.exit(1)
print(f"Loaded {len(problems)} problems.\n")
return problems
# ---------------------------------------------------------------------------
# Message construction
# ---------------------------------------------------------------------------
def build_messages(instance: dict) -> list[dict]:
"""
Build the message list for a SWE-bench instance.
Structure:
- system: instruction to produce a patch
- user: problem statement + hints (the issue)
- user: retrieved repo context (~27k tokens, the thing we compress)
- user: final instruction
"""
issue = instance["problem_statement"]
hints = instance.get("hints_text", "").strip()
context = instance["text"] # BM25-retrieved file contents
issue_content = f"## GitHub Issue\n\n{issue}"
if hints:
issue_content += f"\n\n## Hints\n\n{hints}"
return [
{"role": "system", "content": SYSTEM_MSG},
{"role": "user", "content": issue_content},
{
"role": "user",
"content": f"## Relevant source files\n\n{context}",
},
{
"role": "user",
"content": (
"Based on the issue and source files above, produce a minimal "
"unified diff patch. Output only the patch."
),
},
]
# ---------------------------------------------------------------------------
# Patch helpers
# ---------------------------------------------------------------------------
def parse_patch_files(patch: str) -> set[str]:
"""Extract modified file paths from a unified diff.
Tries `diff --git a/path b/path` first, then falls back to
`--- a/path` lines for diffs that omit the git header.
"""
files = set(re.findall(r"^diff --git a/(.*?) b/", patch, re.MULTILINE))
if not files:
# Fallback: extract from --- a/path lines
files = set(re.findall(r"^--- a/(.+)", patch, re.MULTILINE))
return files
def extract_patch(text: str) -> str:
"""Pull the diff out of an LLM response."""
# Prefer fenced code block
m = re.search(r"```(?:diff|patch)?\n(.*?)```", text, re.DOTALL)
if m:
return m.group(1).strip()
# Fall back to first `diff --git` line
idx = text.find("diff --git")
if idx != -1:
return text[idx:].strip()
return text.strip()
def is_valid_diff(patch: str) -> bool:
return bool(
re.search(r"^@@.*@@", patch, re.MULTILINE) and "---" in patch and "+++" in patch
)
# ---------------------------------------------------------------------------
# Proxy evaluation
# ---------------------------------------------------------------------------
def _parse_hunk_line_ranges(patch: str) -> dict[str, list[tuple[int, int]]]:
"""Parse a unified diff into {filepath: [(start, end), ...]} for modified line ranges."""
current_file = None
ranges: dict[str, list[tuple[int, int]]] = {}
for line in patch.split("\n"):
m = re.match(r"^diff --git a/(.*?) b/", line)
if m:
current_file = m.group(1)
if current_file not in ranges:
ranges[current_file] = []
continue
if not current_file:
m2 = re.match(r"^--- a/(.+)", line)
if m2:
current_file = m2.group(1)
if current_file not in ranges:
ranges[current_file] = []
continue
m3 = re.match(r"^@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@", line)
if m3 and current_file:
start = int(m3.group(1))
length = int(m3.group(2) or "1")
ranges[current_file].append((start, start + length))
return ranges
def _extract_changed_lines(patch: str) -> set[str]:
"""Extract the actual added/removed lines (stripped) from a diff."""
lines = set()
for line in patch.split("\n"):
if line.startswith(("+", "-")) and not line.startswith(("+++", "---")):
stripped = line[1:].strip()
if stripped:
lines.add(stripped)
return lines
def _line_range_overlap(
ranges_a: dict[str, list[tuple[int, int]]],
ranges_b: dict[str, list[tuple[int, int]]],
tolerance: int = 10,
) -> float:
"""Compute fraction of gold hunk line ranges that overlap with generated ranges.
Uses a tolerance window: a generated hunk counts as overlapping a gold hunk
if their line ranges are within ``tolerance`` lines of each other. This
accounts for LLM-generated patches having slightly different line numbers
than the gold patch (due to context window differences, reformatting, etc.)
while still targeting the same logical code region.
"""
shared_files = set(ranges_a.keys()) & set(ranges_b.keys())
if not shared_files:
return 0.0
total_gold_hunks = 0
overlapping_hunks = 0
for f in shared_files:
for g_start, g_end in ranges_a[f]:
total_gold_hunks += 1
for c_start, c_end in ranges_b[f]:
# Ranges overlap (with tolerance) if they're within tolerance
# lines of each other
if (c_start - tolerance) <= g_end and (c_end + tolerance) >= g_start:
overlapping_hunks += 1
break # count each gold hunk at most once
if total_gold_hunks == 0:
return 0.0
return min(overlapping_hunks / total_gold_hunks, 1.0)
def proxy_eval(generated_text: str, instance: dict) -> dict:
"""
Evaluate a generated patch without running the test suite.
Returns:
has_diff: bool — model produced a valid unified diff
file_overlap: float — fraction of gold files present in patch
exact_file_match: bool — generated patch touches exactly the right files
hunk_overlap: float — fraction of gold line ranges covered by generated hunks
content_similarity: float — Jaccard similarity of changed lines (added/removed)
"""
generated_patch = extract_patch(generated_text)
gold_patch = instance["patch"]
gold_files = parse_patch_files(gold_patch)
generated_files = parse_patch_files(generated_patch)
has_diff = is_valid_diff(generated_patch)
file_overlap = (
len(gold_files & generated_files) / len(gold_files) if gold_files else 0.0
)
exact_file_match = (gold_files == generated_files) and bool(gold_files)
# Hunk-level: do they modify the same line ranges?
gold_ranges = _parse_hunk_line_ranges(gold_patch)
gen_ranges = _parse_hunk_line_ranges(generated_patch)
hunk_overlap = _line_range_overlap(gold_ranges, gen_ranges)
# Content-level: Jaccard similarity of the actual changed lines
gold_lines = _extract_changed_lines(gold_patch)
gen_lines = _extract_changed_lines(generated_patch)
if gold_lines or gen_lines:
content_similarity = len(gold_lines & gen_lines) / len(gold_lines | gen_lines)
else:
content_similarity = 0.0
return {
"has_diff": has_diff,
"file_overlap": round(file_overlap, 3),
"exact_file_match": exact_file_match,
"hunk_overlap": round(hunk_overlap, 3),
"content_similarity": round(content_similarity, 3),
"gold_files": sorted(gold_files),
"generated_files": sorted(generated_files),
}
# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------
@dataclass
class SWERunResult:
instance_id: str
mode: str # "baseline" or "compressed"
has_diff: bool
file_overlap: float
exact_file_match: bool
hunk_overlap: float
content_similarity: float
prompt_tokens: int
completion_tokens: int
total_tokens: int
latency_ms: float
cost_usd: float = 0.0
compression_ratio: float = 0.0
error: str = ""
# ---------------------------------------------------------------------------
# Single instance evaluation
# ---------------------------------------------------------------------------
def _run_with_retrieval_loop(
model: str,
messages: list[dict],
tools: list[dict],
cache: dict[str, str],
max_retrievals: int = 5,
) -> tuple[str, object, float, float]:
"""
Call the model, and if it invokes litellm_content_retrieve, fulfill
the tool call from the cache and re-call until the model produces a
final text response (or we hit max_retrievals).
Returns (generated_text, final_usage, total_latency_ms, total_cost).
"""
total_latency = 0.0
total_cost = 0.0
total_usage = None
kwargs: dict = {
"model": model,
"messages": list(messages),
"temperature": 0.0,
"max_tokens": 4096,
}
if tools:
kwargs["tools"] = tools
for _ in range(max_retrievals + 1):
t0 = time.time()
resp = litellm.completion(**kwargs)
total_latency += (time.time() - t0) * 1000
total_cost += resp._hidden_params.get("response_cost", 0) or 0
total_usage = resp.usage
choice = resp.choices[0]
# If the model produced tool calls, fulfill them and loop
tool_calls = getattr(choice.message, "tool_calls", None)
if tool_calls:
# Append the assistant message with tool calls
kwargs["messages"].append(choice.message.model_dump())
for tc in tool_calls:
if tc.function.name == "litellm_content_retrieve":
import json as _json
args = _json.loads(tc.function.arguments)
key = args.get("key", "")
content = cache.get(key, f"[key {key!r} not found in cache]")
kwargs["messages"].append(
{
"role": "tool",
"tool_call_id": tc.id,
"content": content,
}
)
else:
kwargs["messages"].append(
{
"role": "tool",
"tool_call_id": tc.id,
"content": "[unknown tool]",
}
)
continue
# No tool calls — model produced a final text response
return choice.message.content or "", total_usage, total_latency, total_cost
# Exhausted retries — return whatever we have
return resp.choices[0].message.content or "", total_usage, total_latency, total_cost
def eval_instance(
instance: dict,
model: str,
use_compression: bool,
compression_trigger: int,
compression_target: Optional[int] = None,
embedding_model: Optional[str] = None,
) -> SWERunResult:
mode = "compressed" if use_compression else "baseline"
messages = build_messages(instance)
compression_ratio = 0.0
tools: list[dict] = []
cache: dict[str, str] = {}
if use_compression:
compress_kwargs: dict = {
"messages": messages,
"model": model,
"input_type": "openai_chat_completions",
"compression_trigger": compression_trigger,
"embedding_model": embedding_model,
}
if compression_target is not None:
compress_kwargs["compression_target"] = compression_target
result = litellm_compress(**compress_kwargs)
messages = result["messages"]
tools = result["tools"]
cache = result["cache"]
compression_ratio = result["compression_ratio"]
try:
generated_text, usage, latency_ms, cost = _run_with_retrieval_loop(
model=model,
messages=messages,
tools=tools,
cache=cache,
)
ev = proxy_eval(generated_text, instance)
return SWERunResult(
instance_id=instance["instance_id"],
mode=mode,
has_diff=ev["has_diff"],
file_overlap=ev["file_overlap"],
exact_file_match=ev["exact_file_match"],
hunk_overlap=ev["hunk_overlap"],
content_similarity=ev["content_similarity"],
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
latency_ms=latency_ms,
cost_usd=cost,
compression_ratio=compression_ratio,
)
except Exception as e:
return SWERunResult(
instance_id=instance["instance_id"],
mode=mode,
has_diff=False,
file_overlap=0.0,
exact_file_match=False,
hunk_overlap=0.0,
content_similarity=0.0,
prompt_tokens=0,
completion_tokens=0,
total_tokens=0,
latency_ms=0.0,
compression_ratio=0.0,
error=str(e)[:500],
)
# ---------------------------------------------------------------------------
# Aggregation
# ---------------------------------------------------------------------------
def aggregate(results: list[SWERunResult]) -> dict:
if not results:
return {}
valid = [r for r in results if not r.error]
errors = len(results) - len(valid)
return {
"total": len(results),
"errors": errors,
"has_diff_rate": round(
sum(r.has_diff for r in results) / len(results) * 100, 1
),
"avg_file_overlap": round(statistics.mean(r.file_overlap for r in results), 3),
"exact_file_match_rate": round(
sum(r.exact_file_match for r in results) / len(results) * 100, 1
),
"avg_hunk_overlap": round(statistics.mean(r.hunk_overlap for r in results), 3),
"avg_content_similarity": round(
statistics.mean(r.content_similarity for r in results), 3
),
"avg_prompt_tokens": round(statistics.mean(r.prompt_tokens for r in results)),
"avg_total_tokens": round(statistics.mean(r.total_tokens for r in results)),
"avg_latency_ms": round(statistics.mean(r.latency_ms for r in results), 1),
"avg_compression_ratio": round(
statistics.mean(r.compression_ratio for r in results), 4
),
"total_cost_usd": round(sum(r.cost_usd for r in results), 6),
"avg_cost_usd": round(statistics.mean(r.cost_usd for r in results), 6),
}
# ---------------------------------------------------------------------------
# Main benchmark
# ---------------------------------------------------------------------------
def run_benchmark(
model: str,
num_problems: int = 10,
compression_trigger: int = 10_000,
compression_target: Optional[int] = None,
embedding_model: Optional[str] = None,
) -> dict:
"""
Run baseline vs compressed evaluation on SWE-bench Lite problems.
Parameters:
model: LLM model name (litellm format).
num_problems: How many SWE-bench Lite problems to run.
compression_trigger: Token count above which compression activates.
The bm25_27K dataset has ~27k tokens of context
per problem, so a trigger of 10k–20k is sensible.
embedding_model: Optional embedding model for semantic scoring.
"""
problems = load_problems(n=num_problems)
print(f"{'=' * 60}")
print("SWE-bench Compression Eval")
print(f"{'=' * 60}")
print(f"Model: {model}")
print(f"Problems: {len(problems)}")
effective_target = (
compression_target
if compression_target is not None
else compression_trigger * 7 // 10
)
print(f"Compression trigger: {compression_trigger} tokens")
print(f"Compression target: {effective_target} tokens")
print(f"Embedding model: {embedding_model or 'None (BM25 only)'}")
print(f"{'=' * 60}\n")
baseline_results: list[SWERunResult] = []
compressed_results: list[SWERunResult] = []
for i, instance in enumerate(problems):
iid = instance["instance_id"]
print(f"[{i+1}/{len(problems)}] {iid}")
print(f" baseline ...", end=" ", flush=True)
r_base = eval_instance(
instance,
model,
use_compression=False,
compression_trigger=compression_trigger,
compression_target=compression_target,
)
baseline_results.append(r_base)
if r_base.error:
print(f"ERROR: {r_base.error[:80]}")
else:
print(
f"{'✓' if r_base.has_diff else '✗'} diff "
f"file_overlap={r_base.file_overlap:.2f} "
f"{r_base.prompt_tokens} tok "
f"${r_base.cost_usd:.4f}"
)
print(f" compressed ...", end=" ", flush=True)
r_comp = eval_instance(
instance,
model,
use_compression=True,
compression_trigger=compression_trigger,
compression_target=compression_target,
embedding_model=embedding_model,
)
compressed_results.append(r_comp)
if r_comp.error:
print(f"ERROR: {r_comp.error[:80]}")
else:
print(
f"{'✓' if r_comp.has_diff else '✗'} diff "
f"file_overlap={r_comp.file_overlap:.2f} "
f"{r_comp.prompt_tokens} tok "
f"${r_comp.cost_usd:.4f} "
f"(ratio: {r_comp.compression_ratio:.2%})"
)
base_agg = aggregate(baseline_results)
comp_agg = aggregate(compressed_results)
print(f"\n{'=' * 60}")
print("RESULTS")
print(f"{'=' * 60}")
print(f"\n Baseline:")
print(f" Has-diff rate: {base_agg['has_diff_rate']}%")
print(f" Avg file overlap: {base_agg['avg_file_overlap']:.3f}")
print(f" Exact file match: {base_agg['exact_file_match_rate']}%")
print(f" Avg hunk overlap: {base_agg['avg_hunk_overlap']:.3f}")
print(f" Avg content sim: {base_agg['avg_content_similarity']:.3f}")
print(f" Avg prompt tokens: {base_agg['avg_prompt_tokens']}")
print(f" Avg latency: {base_agg['avg_latency_ms']}ms")
print(f" Total cost: ${base_agg['total_cost_usd']:.4f}")
print(f" Avg cost/problem: ${base_agg['avg_cost_usd']:.6f}")
print(f"\n Compressed:")
print(f" Has-diff rate: {comp_agg['has_diff_rate']}%")
print(f" Avg file overlap: {comp_agg['avg_file_overlap']:.3f}")
print(f" Exact file match: {comp_agg['exact_file_match_rate']}%")
print(f" Avg hunk overlap: {comp_agg['avg_hunk_overlap']:.3f}")
print(f" Avg content sim: {comp_agg['avg_content_similarity']:.3f}")
print(f" Avg prompt tokens: {comp_agg['avg_prompt_tokens']}")
print(f" Avg latency: {comp_agg['avg_latency_ms']}ms")
print(f" Total cost: ${comp_agg['total_cost_usd']:.4f}")
print(f" Avg cost/problem: ${comp_agg['avg_cost_usd']:.6f}")
print(f" Avg compression: {comp_agg['avg_compression_ratio']:.2%}")
token_savings = base_agg["avg_prompt_tokens"] - comp_agg["avg_prompt_tokens"]
token_pct = (
round(token_savings / base_agg["avg_prompt_tokens"] * 100, 1)
if base_agg["avg_prompt_tokens"]
else 0
)
print(f"\n Delta (compressed vs baseline):")
print(f" Token savings: {token_savings} ({token_pct}%)")
print(
f" Latency delta: {base_agg['avg_latency_ms'] - comp_agg['avg_latency_ms']:+.1f}ms"
)
print(
f" Has-diff delta: {comp_agg['has_diff_rate'] - base_agg['has_diff_rate']:+.1f}%"
)
print(
f" File overlap delta: {comp_agg['avg_file_overlap'] - base_agg['avg_file_overlap']:+.3f}"
)
print(
f" Exact match delta: {comp_agg['exact_file_match_rate'] - base_agg['exact_file_match_rate']:+.1f}%"
)
print(
f" Hunk overlap delta: {comp_agg['avg_hunk_overlap'] - base_agg['avg_hunk_overlap']:+.3f}"
)
print(
f" Content sim delta: {comp_agg['avg_content_similarity'] - base_agg['avg_content_similarity']:+.3f}"
)
cost_savings = base_agg["total_cost_usd"] - comp_agg["total_cost_usd"]
cost_pct = (
round(cost_savings / base_agg["total_cost_usd"] * 100, 1)
if base_agg["total_cost_usd"]
else 0
)
print(f" Cost savings: ${cost_savings:.4f} ({cost_pct}%)")
ts = time.strftime("%Y-%m-%d_%H-%M-%S")
report_path = f"eval_swe_bench_report_{ts}.json"
report = {
"model": model,
"timestamp": ts,
"num_problems": len(problems),
"compression_trigger": compression_trigger,
"embedding_model": embedding_model,
"baseline": base_agg,
"compressed": comp_agg,
"baseline_results": [asdict(r) for r in baseline_results],
"compressed_results": [asdict(r) for r in compressed_results],
}
with open(report_path, "w") as f:
json.dump(report, f, indent=2)
print(f"\nFull report saved to: {report_path}")
return report
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="SWE-bench Compression Evaluation")
parser.add_argument(
"--model", default="gpt-4o-mini", help="Model name (litellm format)"
)
parser.add_argument(
"--problems",
type=int,
default=10,
help="Number of SWE-bench Lite problems to run (default: 10)",
)
parser.add_argument(
"--compression-trigger",
type=int,
default=10_000,
help="Token threshold to activate compression (default: 10000). "
"The bm25_27K dataset has ~27k tokens of context per problem.",
)
parser.add_argument(
"--compression-target",
type=int,
default=None,
help="Target token count after compression (default: 70%% of trigger). "
"Higher values preserve more context at the cost of less compression.",
)
parser.add_argument(
"--embedding-model",
type=str,
default=None,
help="Embedding model for semantic scoring (e.g. text-embedding-3-small)",
)
args = parser.parse_args()
run_benchmark(
model=args.model,
num_problems=args.problems,
compression_trigger=args.compression_trigger,
compression_target=args.compression_target,
embedding_model=args.embedding_model,
)

358
tests/test_litellm/test_compression.py Normal file
View File

@ -0,0 +1,358 @@
"""
Unit tests for litellm.compress().
"""
import os
import pytest
import litellm
from litellm.compression.scoring.bm25 import bm25_score_messages
from litellm.compression.scoring.embedding_scorer import embedding_score_messages
from litellm.compression.content_detection import detect_content_type
from litellm.compression.message_stubbing import extract_key, stub_message
from litellm.compression.retrieval_tool import build_retrieval_tool
# ---------------------------------------------------------------------------
# BM25 scorer
# ---------------------------------------------------------------------------
def test_bm25_relevance_ranking():
query = "Fix the authentication bug in the login handler"
messages = [
{
"role": "user",
"content": "def login_handler(): authentication check bug fix",
},
{"role": "user", "content": "def render_template(name): css styling layout"},
{"role": "user", "content": "def verify(): authentication token bug handler"},
]
scores = bm25_score_messages(query, messages)
# Messages sharing query terms should score higher than unrelated ones
assert scores[0] > scores[1]
assert scores[2] > scores[1]
def test_bm25_empty_query():
scores = bm25_score_messages("", [{"role": "user", "content": "hello"}])
assert scores == [0.0]
def test_bm25_empty_messages():
scores = bm25_score_messages("query", [])
assert scores == []
def test_bm25_empty_content():
scores = bm25_score_messages("query", [{"role": "user", "content": ""}])
assert scores == [0.0]
# ---------------------------------------------------------------------------
# Content detection
# ---------------------------------------------------------------------------
def test_detect_code():
code = """
import os
from pathlib import Path
def main():
class Foo:
pass
return Foo()
"""
assert detect_content_type(code) == "code"
def test_detect_json():
assert detect_content_type('{"key": "value", "num": 42}') == "json"
assert detect_content_type("[1, 2, 3]") == "json"
def test_detect_text():
assert detect_content_type("This is a plain text paragraph about dogs.") == "text"
def test_detect_empty():
assert detect_content_type("") == "text"
# ---------------------------------------------------------------------------
# Message stubbing
# ---------------------------------------------------------------------------
def test_extract_key_with_filename():
msg = {"role": "user", "content": "# auth.py\ndef authenticate():\n pass"}
used: set = set()
key = extract_key(msg, fallback_index=0, used_keys=used)
assert key == "auth.py"
def test_extract_key_fallback():
msg = {"role": "user", "content": "Some random content without a filename"}
used: set = set()
key = extract_key(msg, fallback_index=5, used_keys=used)
assert key == "message_5"
def test_extract_key_duplicates():
used: set = set()
msg = {"role": "user", "content": "# auth.py\ncode here"}
k1 = extract_key(msg, fallback_index=0, used_keys=used)
k2 = extract_key(msg, fallback_index=1, used_keys=used)
assert k1 == "auth.py"
assert k2 == "auth.py_2"
def test_stub_message():
msg = {"role": "user", "content": "line1\nline2\nline3"}
stubbed = stub_message(msg, "test_key")
assert stubbed["role"] == "user"
assert "test_key" in stubbed["content"]
assert "litellm_content_retrieve" in stubbed["content"]
assert "3 lines" in stubbed["content"]
# ---------------------------------------------------------------------------
# Retrieval tool
# ---------------------------------------------------------------------------
def test_retrieval_tool_schema():
tool = build_retrieval_tool(["auth.py", "utils.py"])
assert tool["type"] == "function"
assert tool["function"]["name"] == "litellm_content_retrieve"
assert "key" in tool["function"]["parameters"]["properties"]
assert tool["function"]["parameters"]["properties"]["key"]["enum"] == [
"auth.py",
"utils.py",
]
assert tool["function"]["parameters"]["required"] == ["key"]
def test_retrieval_tool_description_lists_keys():
tool = build_retrieval_tool(["foo.py", "bar.js"])
desc = tool["function"]["description"]
assert "foo.py" in desc
assert "bar.js" in desc
# ---------------------------------------------------------------------------
# compress() — end-to-end
# ---------------------------------------------------------------------------
def test_compress_below_trigger_passthrough():
messages = [{"role": "user", "content": "hello"}]
result = litellm.compress(messages, model="gpt-4o")
assert result["messages"] == messages
assert result["cache"] == {}
assert result["tools"] == []
assert result["compression_ratio"] == 0.0
assert result["original_tokens"] == result["compressed_tokens"]
def test_compress_above_trigger():
big_messages = [
{"role": "system", "content": "You are a coding assistant."},
{
"role": "user",
"content": "# auth.py\n" + "def authenticate():\n pass\n" * 2000,
},
{
"role": "user",
"content": "# utils.py\n" + "def helper():\n pass\n" * 2000,
},
{
"role": "user",
"content": "# readme.md\n" + "This is documentation. " * 2000,
},
{"role": "user", "content": "Fix the bug in auth.py"},
]
result = litellm.compress(
big_messages,
model="gpt-4o",
compression_trigger=1000,
compression_target=500,
)
assert result["compressed_tokens"] < result["original_tokens"]
assert result["compression_ratio"] > 0
assert len(result["cache"]) > 0
assert len(result["tools"]) == 1
assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve"
def test_compress_preserves_system_message():
messages = [
{"role": "system", "content": "System prompt. " * 500},
{"role": "user", "content": "Large file content. " * 5000},
{"role": "user", "content": "Fix the bug"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
assert result["messages"][0]["role"] == "system"
assert "System prompt" in result["messages"][0]["content"]
def test_compress_preserves_last_user_message():
messages = [
{"role": "user", "content": "Big context " * 5000},
{"role": "user", "content": "Fix the bug in auth.py"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
last_user = [m for m in result["messages"] if m["role"] == "user"][-1]
assert "Fix the bug in auth.py" in last_user["content"]
def test_compress_preserves_last_assistant_message():
messages = [
{"role": "user", "content": "Big context " * 5000},
{"role": "assistant", "content": "I'll help with that. " * 2000},
{"role": "user", "content": "Now fix the bug"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"]
assert len(assistant_msgs) >= 1
# The last assistant message should be preserved (not stubbed)
last_assistant = assistant_msgs[-1]
assert "I'll help with that" in last_assistant["content"]
def test_cache_keys_match_stubs():
messages = [
{"role": "user", "content": "# auth.py\n" + "code " * 5000},
{"role": "user", "content": "Fix it"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
if result["tools"]:
tool_desc = result["tools"][0]["function"]["description"]
for key in result["cache"]:
assert key in tool_desc
def test_compress_default_target():
"""compression_target defaults to compression_trigger // 2."""
messages = [
{"role": "user", "content": "content " * 5000},
{"role": "user", "content": "query"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=2000)
# Should have compressed — target = 1000
assert result["compressed_tokens"] <= result["original_tokens"]
def test_compress_forwards_embedding_model_params(monkeypatch):
captured = {}
def fake_embedding_score_messages(
query, messages, model, cache=None, embedding_model_params=None
):
captured["query"] = query
captured["model"] = model
captured["embedding_model_params"] = embedding_model_params
return [0.0] * len(messages)
monkeypatch.setattr(
"litellm.compression.scoring.embedding_scorer.embedding_score_messages",
fake_embedding_score_messages,
)
result = litellm.compress(
messages=[
{"role": "user", "content": "Authentication code " * 2000},
{"role": "user", "content": "Fix auth"},
],
model="gpt-4o",
compression_trigger=1000,
embedding_model="text-embedding-3-small",
embedding_model_params={"api_base": "https://example-embeddings.test"},
)
assert result["compressed_tokens"] <= result["original_tokens"]
assert captured["model"] == "text-embedding-3-small"
assert captured["embedding_model_params"] == {
"api_base": "https://example-embeddings.test"
}
def test_embedding_scorer_forwards_embedding_model_params(monkeypatch):
captured = {}
class _MockResponse:
data = [
{"embedding": [1.0, 0.0]},
{"embedding": [1.0, 0.0]},
{"embedding": [0.0, 1.0]},
]
def fake_embedding(**kwargs):
captured.update(kwargs)
return _MockResponse()
monkeypatch.setattr(litellm, "embedding", fake_embedding)
scores = embedding_score_messages(
query="auth",
messages=[
{"role": "user", "content": "auth code"},
{"role": "user", "content": "cooking recipe"},
],
model="text-embedding-3-small",
embedding_model_params={"api_base": "https://example-embeddings.test"},
)
assert len(scores) == 2
assert captured["model"] == "text-embedding-3-small"
assert captured["api_base"] == "https://example-embeddings.test"
# ---------------------------------------------------------------------------
# Embedding scorer — integration test (skipped without API key)
# ---------------------------------------------------------------------------
@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="Needs OPENAI_API_KEY")
def test_embedding_scorer():
result = litellm.compress(
messages=[
{"role": "user", "content": "Authentication code " * 2000},
{"role": "user", "content": "Unrelated cooking recipes " * 2000},
{"role": "user", "content": "Fix auth"},
],
model="gpt-4o",
compression_trigger=1000,
embedding_model="text-embedding-3-small",
)
assert result["compression_ratio"] > 0
assert len(result["cache"]) > 0
@pytest.mark.parametrize(
"final_user_message, expected_content",
[
("How to cook?", "Unrelated cooking recipes "),
("Fix auth", "Authentication code "),
],
)
def test_simple_compression(final_user_message, expected_content):
messages = [
{"role": "user", "content": "Authentication code " * 2000},
{"role": "user", "content": "Unrelated cooking recipes " * 2000},
{"role": "user", "content": final_user_message},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
print(result["messages"])
if expected_content == "Unrelated cooking recipes ":
assert "Unrelated cooking recipes " in result["messages"][1]["content"]
assert "Authentication code " not in result["messages"][0]["content"]
elif expected_content == "Authentication code ":
assert "Authentication code " in result["messages"][0]["content"]
assert "Unrelated cooking recipes " not in result["messages"][1]["content"]
else:
raise ValueError(f"Unexpected expected_content: {expected_content}")