From 26c741233961eec03a49c1c68578286d58b188fb Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 13 Apr 2026 12:23:54 -0700 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20add=20litellm.compress()=20?= =?UTF-8?q?=E2=80=94=20BM25-based=20prompt=20compression=20with=20retrieva?= =?UTF-8?q?l=20tool=20(#25637)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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. Co-Authored-By: Claude Opus 4.6 * 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. Co-Authored-By: Claude Opus 4.6 * 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. Co-Authored-By: Claude Opus 4.6 * 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 Co-Authored-By: Claude Opus 4.6 * 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 Co-Authored-By: Claude Opus 4.6 * 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 Co-Authored-By: Claude Opus 4.6 * 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% Co-Authored-By: Claude Opus 4.6 * 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. Co-Authored-By: Claude Opus 4.6 * 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. Co-Authored-By: Claude Opus 4.6 * 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 Co-Authored-By: Claude Opus 4.6 * Revert "feat: add compression_interception callback for LiteLLM Proxy" This reverts commit 72bd5cb152ca1df07f14a14e14a2816e188874a8. --------- Co-authored-by: Claude Opus 4.6 --- .../docs/completion/prompt_compression.md | 123 ++ docs/my-website/package-lock.json | 7 + docs/my-website/sidebars.js | 6 + litellm/__init__.py | 1 + litellm/compression/__init__.py | 3 + litellm/compression/compress.py | 249 ++++ litellm/compression/content_detection.py | 45 + litellm/compression/message_stubbing.py | 120 ++ litellm/compression/retrieval_tool.py | 35 + litellm/compression/scoring/__init__.py | 4 + litellm/compression/scoring/bm25.py | 123 ++ .../compression/scoring/embedding_scorer.py | 95 ++ litellm/types/compression.py | 14 + scripts/eval_compression.py | 1125 +++++++++++++++++ tests/eval_swe_bench.py | 751 +++++++++++ tests/test_litellm/test_compression.py | 358 ++++++ 16 files changed, 3059 insertions(+) create mode 100644 docs/my-website/docs/completion/prompt_compression.md create mode 100644 litellm/compression/__init__.py create mode 100644 litellm/compression/compress.py create mode 100644 litellm/compression/content_detection.py create mode 100644 litellm/compression/message_stubbing.py create mode 100644 litellm/compression/retrieval_tool.py create mode 100644 litellm/compression/scoring/__init__.py create mode 100644 litellm/compression/scoring/bm25.py create mode 100644 litellm/compression/scoring/embedding_scorer.py create mode 100644 litellm/types/compression.py create mode 100644 scripts/eval_compression.py create mode 100644 tests/eval_swe_bench.py create mode 100644 tests/test_litellm/test_compression.py diff --git a/docs/my-website/docs/completion/prompt_compression.md b/docs/my-website/docs/completion/prompt_compression.md new file mode 100644 index 0000000000..2d999291af --- /dev/null +++ b/docs/my-website/docs/completion/prompt_compression.md @@ -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 +``` diff --git a/docs/my-website/package-lock.json b/docs/my-website/package-lock.json index 56684b737d..d14ca96cf5 100644 --- a/docs/my-website/package-lock.json +++ b/docs/my-website/package-lock.json @@ -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", diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index b2ac843391..46e392037a 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -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", diff --git a/litellm/__init__.py b/litellm/__init__.py index 8087e3f531..8b0da380fd 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 # Skills API from .skills.main import ( diff --git a/litellm/compression/__init__.py b/litellm/compression/__init__.py new file mode 100644 index 0000000000..11c5eaf84e --- /dev/null +++ b/litellm/compression/__init__.py @@ -0,0 +1,3 @@ +from litellm.compression.compress import compress + +__all__ = ["compress"] diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py new file mode 100644 index 0000000000..718bc1c45c --- /dev/null +++ b/litellm/compression/compress.py @@ -0,0 +1,249 @@ +""" +Main compress() function — orchestrates BM25/embedding scoring, message stubbing, +and retrieval tool injection. +""" + +from typing import Any, Dict, List, Optional, Set + +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 + + +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=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=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, + ) diff --git a/litellm/compression/content_detection.py b/litellm/compression/content_detection.py new file mode 100644 index 0000000000..0655a42daf --- /dev/null +++ b/litellm/compression/content_detection.py @@ -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" diff --git a/litellm/compression/message_stubbing.py b/litellm/compression/message_stubbing.py new file mode 100644 index 0000000000..2330f1bbc9 --- /dev/null +++ b/litellm/compression/message_stubbing.py @@ -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} diff --git a/litellm/compression/retrieval_tool.py b/litellm/compression/retrieval_tool.py new file mode 100644 index 0000000000..1ee24784a6 --- /dev/null +++ b/litellm/compression/retrieval_tool.py @@ -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"], + }, + }, + } diff --git a/litellm/compression/scoring/__init__.py b/litellm/compression/scoring/__init__.py new file mode 100644 index 0000000000..78bb434d17 --- /dev/null +++ b/litellm/compression/scoring/__init__.py @@ -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"] diff --git a/litellm/compression/scoring/bm25.py b/litellm/compression/scoring/bm25.py new file mode 100644 index 0000000000..e8e1bf631e --- /dev/null +++ b/litellm/compression/scoring/bm25.py @@ -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 diff --git a/litellm/compression/scoring/embedding_scorer.py b/litellm/compression/scoring/embedding_scorer.py new file mode 100644 index 0000000000..f3558ae8f5 --- /dev/null +++ b/litellm/compression/scoring/embedding_scorer.py @@ -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 diff --git a/litellm/types/compression.py b/litellm/types/compression.py new file mode 100644 index 0000000000..01d5a6dd4d --- /dev/null +++ b/litellm/types/compression.py @@ -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] diff --git a/scripts/eval_compression.py b/scripts/eval_compression.py new file mode 100644 index 0000000000..d7d90dacc2 --- /dev/null +++ b/scripts/eval_compression.py @@ -0,0 +1,1125 @@ +""" +Prompt Compression Evaluation Harness +====================================== +Compare model performance on coding tasks with and without prompt compression. + +Usage: + python scripts/eval_compression.py --model gpt-4o --problems 5 + python scripts/eval_compression.py --model claude-sonnet-4-20250514 --problems 12 --runs 3 + python scripts/eval_compression.py --model gpt-4o-mini --padding-factor 50 + +The harness runs each problem in two modes: + 1. **baseline** — raw prompt sent directly to the model. + 2. **compressed** — prompt is padded with distractor context, then + ``litellm.compress()`` removes the noise before sending. + +This measures whether compression preserves the signal the model needs +to solve the task while reducing token usage. + +Set --padding-factor to control how much distractor context is injected +(higher = more tokens to compress away). +""" + +import argparse +import json +import os +import statistics +import subprocess +import sys +import tempfile +import textwrap +import time +from dataclasses import asdict, dataclass, field +from typing import Optional + +import litellm + +# --------------------------------------------------------------------------- +# Problem definitions (HumanEval-style) +# --------------------------------------------------------------------------- + +PROBLEMS = [ + { + "id": "has_close_elements", + "prompt": textwrap.dedent( + """\ + from typing import List + + def has_close_elements(numbers: List[float], threshold: float) -> bool: + \"\"\"Check if in given list of numbers, are any two numbers closer to each other than + given threshold. + >>> has_close_elements([1.0, 2.0, 3.0], 0.5) + False + >>> has_close_elements([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3) + True + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert has_close_elements([1.0, 2.0, 3.9, 4.0, 5.0, 2.2], 0.3) == True + assert has_close_elements([1.0, 2.0, 3.9, 4.0, 5.0, 2.2], 0.05) == False + assert has_close_elements([1.0, 2.0, 5.9, 4.0, 5.0], 0.95) == True + assert has_close_elements([1.0, 2.0, 5.9, 4.0, 5.0], 0.8) == False + assert has_close_elements([1.0, 2.0, 3.0, 4.0, 5.0], 2.0) == True + assert has_close_elements([], 0.5) == False + print("PASSED") + """ + ), + }, + { + "id": "separate_paren_groups", + "prompt": textwrap.dedent( + """\ + from typing import List + + def separate_paren_groups(paren_string: str) -> List[str]: + \"\"\"Input to this function is a string containing multiple groups of nested parentheses. + Your goal is to separate those groups into separate strings and return the list of those. + Separate groups are balanced (each open brace is properly closed) and not nested within each other. + Ignore any spaces in the input string. + >>> separate_paren_groups('( ) (( )) (( )( ))') + ['()', '(())', '(()())'] + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert separate_paren_groups('(()()) ((())) () ((())()())') == ['(()())', '((()))', '()', '((())()())'] + assert separate_paren_groups('() (()) ((())) (((())))') == ['()', '(())', '((()))', '(((())))'] + assert separate_paren_groups('(()(()))') == ['(()(()))'] + assert separate_paren_groups('( ) (( )) (( )( ))') == ['()', '(())', '(()())'] + print("PASSED") + """ + ), + }, + { + "id": "truncate_number", + "prompt": textwrap.dedent( + """\ + def truncate_number(number: float) -> float: + \"\"\"Given a positive floating point number, it can be decomposed into + an integer part (largest integer smaller than given number) and decimals + (leftover part always smaller than 1). + Return the decimal part of the number. + >>> truncate_number(3.5) + 0.5 + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert truncate_number(3.5) == 0.5 + assert abs(truncate_number(1.33) - 0.33) < 1e-6 + assert abs(truncate_number(123.456) - 0.456) < 1e-6 + print("PASSED") + """ + ), + }, + { + "id": "below_zero", + "prompt": textwrap.dedent( + """\ + from typing import List + + def below_zero(operations: List[int]) -> bool: + \"\"\"You're given a list of deposit and withdrawal operations on a bank account that starts with + zero balance. Your task is to detect if at any point the balance of account falls below zero, and + at that point function should return True. Otherwise it should return False. + >>> below_zero([1, 2, 3]) + False + >>> below_zero([1, 2, -4, 5]) + True + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert below_zero([]) == False + assert below_zero([1, 2, -3, 1, 2, -3]) == False + assert below_zero([1, 2, -4, 5, 6]) == True + assert below_zero([1, -1, 2, -2, 5, -5, 4, -4]) == False + assert below_zero([1, -1, 2, -2, 5, -5, 4, -5]) == True + assert below_zero([1, -2]) == True + print("PASSED") + """ + ), + }, + { + "id": "mean_absolute_deviation", + "prompt": textwrap.dedent( + """\ + from typing import List + + def mean_absolute_deviation(numbers: List[float]) -> float: + \"\"\"For a given list of input numbers, calculate Mean Absolute Deviation + around the mean of this dataset. + Mean Absolute Deviation is the average absolute difference between each + element and a centerpoint (mean in this case): + MAD = average | x - x_mean | + >>> mean_absolute_deviation([1.0, 2.0, 3.0, 4.0]) + 1.0 + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert abs(mean_absolute_deviation([1.0, 2.0, 3.0, 4.0]) - 1.0) < 1e-6 + assert abs(mean_absolute_deviation([1.0, 2.0, 3.0, 4.0, 5.0]) - 1.2) < 1e-6 + assert abs(mean_absolute_deviation([1.0, 1.0, 1.0, 1.0]) - 0.0) < 1e-6 + print("PASSED") + """ + ), + }, + { + "id": "intersperse", + "prompt": textwrap.dedent( + """\ + from typing import List + + def intersperse(numbers: List[int], delimiter: int) -> List[int]: + \"\"\"Insert a number 'delimiter' between every two consecutive elements of input list `numbers`. + >>> intersperse([], 4) + [] + >>> intersperse([1, 2, 3], 4) + [1, 4, 2, 4, 3] + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert intersperse([], 7) == [] + assert intersperse([5, 6, 3, 2], 8) == [5, 8, 6, 8, 3, 8, 2] + assert intersperse([2, 2, 2], 2) == [2, 2, 2, 2, 2] + print("PASSED") + """ + ), + }, + { + "id": "parse_nested_parens", + "prompt": textwrap.dedent( + """\ + from typing import List + + def parse_nested_parens(paren_string: str) -> List[int]: + \"\"\"Input to this function is a string represented multiple groups of nested parentheses separated by spaces. + For each of the groups, output the deepest level of nesting of parentheses. + E.g. (()()) has maximum two levels of nesting while ((())) has three. + >>> parse_nested_parens('(()()) ((())) () ((())())') + [2, 3, 1, 3] + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert parse_nested_parens('(()()) ((())) () ((())())') == [2, 3, 1, 3] + assert parse_nested_parens('() (()) ((())) (((())))') == [1, 2, 3, 4] + assert parse_nested_parens('(()(())((())))') == [4] + print("PASSED") + """ + ), + }, + { + "id": "filter_by_substring", + "prompt": textwrap.dedent( + """\ + from typing import List + + def filter_by_substring(strings: List[str], substring: str) -> List[str]: + \"\"\"Filter an input list of strings only for ones that contain given substring. + >>> filter_by_substring([], 'a') + [] + >>> filter_by_substring(['abc', 'bacd', 'cde', 'array'], 'a') + ['abc', 'bacd', 'array'] + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert filter_by_substring([], 'john') == [] + assert filter_by_substring(['xxx', 'asd', 'xxy', 'john doe', 'xxxuj', 'xxx'], 'xxx') == ['xxx', 'xxxuj', 'xxx'] + assert filter_by_substring(['xxx', 'asd', 'aaber', 'john doe', 'xxxuj', 'xxx'], 'xx') == ['xxx', 'xxxuj', 'xxx'] + assert filter_by_substring(['grunt', 'hierarchial', 'abc', 'hierarchial'], 'hi') == ['hierarchial', 'hierarchial'] + print("PASSED") + """ + ), + }, + { + "id": "sum_product", + "prompt": textwrap.dedent( + """\ + from typing import List, Tuple + + def sum_product(numbers: List[int]) -> Tuple[int, int]: + \"\"\"For a given list of integers, return a tuple consisting of a sum and a product of all the integers in a list. + Empty sum should be equal to 0 and empty product should be equal to 1. + >>> sum_product([]) + (0, 1) + >>> sum_product([1, 2, 3, 4]) + (10, 24) + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert sum_product([]) == (0, 1) + assert sum_product([1, 1, 1]) == (3, 1) + assert sum_product([100, 0]) == (100, 0) + assert sum_product([3, 5, 7]) == (15, 105) + assert sum_product([10]) == (10, 10) + print("PASSED") + """ + ), + }, + { + "id": "max_element", + "prompt": textwrap.dedent( + """\ + from typing import List + + def max_element(l: List[int]) -> int: + \"\"\"Return maximum element in the list. + >>> max_element([1, 2, 3]) + 3 + >>> max_element([5, 3, -5, 2, -3, 3, 9, 0, 123, 1, -10]) + 123 + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert max_element([1, 2, 3]) == 3 + assert max_element([5, 3, -5, 2, -3, 3, 9, 0, 124, 1, -10]) == 124 + assert max_element([-1, -2, -3]) == -1 + print("PASSED") + """ + ), + }, + { + "id": "fizz_buzz", + "prompt": textwrap.dedent( + """\ + def fizz_buzz(n: int) -> int: + \"\"\"Return the number of times the digit 7 appears in integers less than n which are divisible by 11 or 13. + >>> fizz_buzz(50) + 0 + >>> fizz_buzz(78) + 2 + >>> fizz_buzz(79) + 3 + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert fizz_buzz(50) == 0 + assert fizz_buzz(78) == 2 + assert fizz_buzz(79) == 3 + assert fizz_buzz(100) == 3 + assert fizz_buzz(200) == 6 + assert fizz_buzz(4000) == 192 + print("PASSED") + """ + ), + }, + { + "id": "sort_by_binary_len", + "prompt": textwrap.dedent( + """\ + from typing import List + + def sort_array(arr: List[int]) -> List[int]: + \"\"\"Sort an array of non-negative integers according to number of ones in their binary + representation in ascending order. For equal number of ones, sort based on decimal value. + >>> sort_array([1, 5, 2, 3, 4]) + [1, 2, 4, 3, 5] + >>> sort_array([-2, -3, -4, -5, -6]) + [-6, -5, -4, -3, -2] + >>> sort_array([1, 0, 2, 3, 4]) + [0, 1, 2, 4, 3] + \"\"\" + """ + ), + "tests": textwrap.dedent( + """\ + assert sort_array([1, 5, 2, 3, 4]) == [1, 2, 4, 3, 5] + assert sort_array([-2, -3, -4, -5, -6]) == [-6, -5, -4, -3, -2] + assert sort_array([1, 0, 2, 3, 4]) == [0, 1, 2, 4, 3] + assert sort_array([]) == [] + assert sort_array([2, 5, 77, 4, 5, 3, 5, 7, 2, 3, 4]) == [2, 2, 4, 4, 3, 3, 5, 5, 5, 7, 77] + assert sort_array([3, 6, 44, 12, 32, 5]) == [32, 3, 5, 6, 12, 44] + print("PASSED") + """ + ), + }, +] + +# Distractor code snippets injected as prior conversation context. +# These are plausible but irrelevant to the actual task, forcing the +# compressor to identify and drop them. +DISTRACTOR_SNIPPETS = [ + # distractor 0 — database connection pool + textwrap.dedent( + """\ + # db_pool.py + import threading + from contextlib import contextmanager + + class ConnectionPool: + def __init__(self, dsn, min_size=2, max_size=10): + self._dsn = dsn + self._min_size = min_size + self._max_size = max_size + self._pool = [] + self._lock = threading.Lock() + self._initialize() + + def _initialize(self): + for _ in range(self._min_size): + self._pool.append(self._create_connection()) + + def _create_connection(self): + import psycopg2 + return psycopg2.connect(self._dsn) + + @contextmanager + def acquire(self): + conn = self._checkout() + try: + yield conn + finally: + self._checkin(conn) + + def _checkout(self): + with self._lock: + if self._pool: + return self._pool.pop() + if len(self._pool) < self._max_size: + return self._create_connection() + raise RuntimeError("Pool exhausted") + + def _checkin(self, conn): + with self._lock: + self._pool.append(conn) + + def close_all(self): + with self._lock: + for conn in self._pool: + conn.close() + self._pool.clear() + """ + ), + # distractor 1 — HTTP retry logic + textwrap.dedent( + """\ + # http_retry.py + import time + import random + import requests + from functools import wraps + + class RetryConfig: + def __init__(self, max_retries=3, base_delay=1.0, max_delay=60.0, backoff_factor=2.0): + self.max_retries = max_retries + self.base_delay = base_delay + self.max_delay = max_delay + self.backoff_factor = backoff_factor + + def retry_with_backoff(config=None): + if config is None: + config = RetryConfig() + + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + last_exception = None + for attempt in range(config.max_retries + 1): + try: + return func(*args, **kwargs) + except (requests.ConnectionError, requests.Timeout) as e: + last_exception = e + if attempt == config.max_retries: + break + delay = min( + config.base_delay * (config.backoff_factor ** attempt), + config.max_delay + ) + jitter = random.uniform(0, delay * 0.1) + time.sleep(delay + jitter) + raise last_exception + return wrapper + return decorator + + @retry_with_backoff(RetryConfig(max_retries=5)) + def fetch_data(url, params=None): + resp = requests.get(url, params=params, timeout=30) + resp.raise_for_status() + return resp.json() + """ + ), + # distractor 2 — LRU cache implementation + textwrap.dedent( + """\ + # lru_cache.py + from collections import OrderedDict + from threading import RLock + + class LRUCache: + def __init__(self, capacity=128): + self._capacity = capacity + self._cache = OrderedDict() + self._lock = RLock() + self._hits = 0 + self._misses = 0 + + def get(self, key, default=None): + with self._lock: + if key in self._cache: + self._cache.move_to_end(key) + self._hits += 1 + return self._cache[key] + self._misses += 1 + return default + + def put(self, key, value): + with self._lock: + if key in self._cache: + self._cache.move_to_end(key) + self._cache[key] = value + if len(self._cache) > self._capacity: + self._cache.popitem(last=False) + + def delete(self, key): + with self._lock: + self._cache.pop(key, None) + + def clear(self): + with self._lock: + self._cache.clear() + + @property + def stats(self): + total = self._hits + self._misses + hit_rate = self._hits / total if total else 0.0 + return {"hits": self._hits, "misses": self._misses, "hit_rate": hit_rate} + + def __len__(self): + return len(self._cache) + + def __contains__(self, key): + return key in self._cache + """ + ), + # distractor 3 — CSV report generator + textwrap.dedent( + """\ + # report_gen.py + import csv + import io + from datetime import datetime, timedelta + + class ReportGenerator: + def __init__(self, title, columns): + self.title = title + self.columns = columns + self.rows = [] + + def add_row(self, **kwargs): + row = {col: kwargs.get(col, "") for col in self.columns} + self.rows.append(row) + + def to_csv(self): + output = io.StringIO() + writer = csv.DictWriter(output, fieldnames=self.columns) + writer.writeheader() + writer.writerows(self.rows) + return output.getvalue() + + def summary(self): + numeric_cols = [] + for col in self.columns: + try: + vals = [float(r[col]) for r in self.rows if r[col] != ""] + if vals: + numeric_cols.append({ + "column": col, + "min": min(vals), + "max": max(vals), + "mean": sum(vals) / len(vals), + "count": len(vals), + }) + except (ValueError, TypeError): + continue + return numeric_cols + + def filter_rows(self, predicate): + gen = ReportGenerator(self.title, self.columns) + gen.rows = [r for r in self.rows if predicate(r)] + return gen + + def date_range_report(self, date_col, start, end): + def in_range(row): + try: + d = datetime.fromisoformat(row[date_col]) + return start <= d <= end + except (ValueError, KeyError): + return False + return self.filter_rows(in_range) + """ + ), + # distractor 4 — async task queue + textwrap.dedent( + """\ + # task_queue.py + import asyncio + import logging + from dataclasses import dataclass, field + from enum import Enum + from typing import Any, Callable, Coroutine + + logger = logging.getLogger(__name__) + + class TaskStatus(Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + + @dataclass + class Task: + id: str + func: Callable[..., Coroutine] + args: tuple = () + kwargs: dict = field(default_factory=dict) + status: TaskStatus = TaskStatus.PENDING + result: Any = None + error: str = "" + retries: int = 0 + max_retries: int = 3 + + class AsyncTaskQueue: + def __init__(self, concurrency=5): + self._queue = asyncio.Queue() + self._concurrency = concurrency + self._tasks = {} + self._workers = [] + + async def submit(self, task: Task): + self._tasks[task.id] = task + await self._queue.put(task) + + async def _worker(self): + while True: + task = await self._queue.get() + task.status = TaskStatus.RUNNING + try: + task.result = await task.func(*task.args, **task.kwargs) + task.status = TaskStatus.COMPLETED + except Exception as e: + task.retries += 1 + if task.retries <= task.max_retries: + task.status = TaskStatus.PENDING + await self._queue.put(task) + else: + task.status = TaskStatus.FAILED + task.error = str(e) + logger.error(f"Task {task.id} failed: {e}") + finally: + self._queue.task_done() + + async def start(self): + self._workers = [ + asyncio.create_task(self._worker()) + for _ in range(self._concurrency) + ] + + async def wait(self): + await self._queue.join() + + async def shutdown(self): + for w in self._workers: + w.cancel() + """ + ), + # distractor 5 — config parser with env var interpolation + textwrap.dedent( + """\ + # config_parser.py + import os + import re + import json + from pathlib import Path + + _ENV_PATTERN = re.compile(r'\\$\\{([A-Z_][A-Z0-9_]*)(?::-(.*?))?\\}') + + class ConfigError(Exception): + pass + + class Config: + def __init__(self, data=None): + self._data = data or {} + + @classmethod + def from_file(cls, path): + p = Path(path) + if not p.exists(): + raise ConfigError(f"Config file not found: {path}") + with open(p) as f: + raw = json.load(f) + return cls(cls._interpolate(raw)) + + @classmethod + def _interpolate(cls, obj): + if isinstance(obj, str): + return cls._interpolate_string(obj) + if isinstance(obj, dict): + return {k: cls._interpolate(v) for k, v in obj.items()} + if isinstance(obj, list): + return [cls._interpolate(item) for item in obj] + return obj + + @classmethod + def _interpolate_string(cls, s): + def replacer(match): + var_name = match.group(1) + default = match.group(2) + value = os.environ.get(var_name) + if value is None: + if default is not None: + return default + raise ConfigError(f"Required env var {var_name} is not set") + return value + return _ENV_PATTERN.sub(replacer, s) + + def get(self, key, default=None): + keys = key.split(".") + obj = self._data + for k in keys: + if isinstance(obj, dict) and k in obj: + obj = obj[k] + else: + return default + return obj + + def require(self, key): + val = self.get(key) + if val is None: + raise ConfigError(f"Required config key missing: {key}") + return val + """ + ), +] + + +# --------------------------------------------------------------------------- +# Data classes +# --------------------------------------------------------------------------- + + +@dataclass +class RunResult: + problem_id: str + mode: str # "baseline" or "compressed" + passed: bool + generated_code: str + prompt_tokens: int + completion_tokens: int + total_tokens: int + latency_ms: float + compression_ratio: float = 0.0 + error: str = "" + + +@dataclass +class BenchmarkReport: + model: str + timestamp: str + num_problems: int + num_runs: int + padding_factor: int + baseline: dict = field(default_factory=dict) + compressed: dict = field(default_factory=dict) + per_problem: list = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# LLM caller (uses litellm) +# --------------------------------------------------------------------------- + +SYSTEM_MSG = ( + "You are a Python coding assistant. Complete the function below. " + "Return ONLY the Python code (the complete function), no explanation, " + "no markdown fences." +) + + +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, + } + + +# --------------------------------------------------------------------------- +# Code extraction & execution +# --------------------------------------------------------------------------- + + +def extract_code(raw: str) -> str: + """Pull code out of the LLM response, stripping markdown fences if present.""" + text = raw.strip() + if text.startswith("```"): + lines = text.split("\n") + lines = [line for line in lines[1:] if not line.strip().startswith("```")] + text = "\n".join(lines) + return text.strip() + + +def run_tests(code: str, tests: str, timeout: int = 10) -> tuple[bool, str]: + """Execute generated code + tests in a subprocess. Returns (passed, error_msg).""" + full = code + "\n\n" + tests + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: + f.write(full) + f.flush() + try: + result = subprocess.run( + [sys.executable, f.name], + capture_output=True, + text=True, + timeout=timeout, + ) + if result.returncode == 0 and "PASSED" in result.stdout: + return True, "" + err = result.stderr.strip() or result.stdout.strip() + return False, err[:500] + except subprocess.TimeoutExpired: + return False, "TIMEOUT" + finally: + os.unlink(f.name) + + +# --------------------------------------------------------------------------- +# Context building — pad the prompt with distractors +# --------------------------------------------------------------------------- + + +def build_messages( + problem: dict, + padding_factor: int = 0, +) -> list[dict]: + """ + Build a message list for a problem. + + When ``padding_factor`` > 0, distractor code snippets are injected as + prior user messages (simulating a long coding session) so there is + enough context for compression to act on. + """ + messages: list[dict] = [{"role": "system", "content": SYSTEM_MSG}] + + if padding_factor > 0: + for i in range(padding_factor): + snippet = DISTRACTOR_SNIPPETS[i % len(DISTRACTOR_SNIPPETS)] + messages.append( + { + "role": "user", + "content": f"Here is some code from our codebase:\n\n{snippet}", + } + ) + messages.append( + { + "role": "assistant", + "content": "Got it, I've reviewed that code. What would you like me to help with?", + } + ) + + messages.append( + { + "role": "user", + "content": ( + "Complete the following Python function. Return ONLY the code.\n\n" + + problem["prompt"] + ), + } + ) + return messages + + +# --------------------------------------------------------------------------- +# Single problem evaluation +# --------------------------------------------------------------------------- + + +def eval_problem( + problem: dict, + model: str, + padding_factor: int, + use_compression: bool, + compression_trigger: int, + embedding_model: Optional[str], +) -> RunResult: + """Evaluate a single problem in either baseline or compressed mode.""" + mode = "compressed" if use_compression else "baseline" + messages = build_messages(problem, padding_factor=padding_factor) + + compression_ratio = 0.0 + + if use_compression: + result = litellm.compress( + messages=messages, + model=model, + compression_trigger=compression_trigger, + embedding_model=embedding_model, + ) + messages = result["messages"] + compression_ratio = result["compression_ratio"] + + try: + resp = call_llm(model, messages) + code = extract_code(resp["text"]) + passed, error = run_tests(code, problem["tests"]) + + return RunResult( + problem_id=problem["id"], + mode=mode, + passed=passed, + generated_code=code, + prompt_tokens=resp["prompt_tokens"], + completion_tokens=resp["completion_tokens"], + total_tokens=resp["total_tokens"], + latency_ms=resp["latency_ms"], + compression_ratio=compression_ratio, + error=error, + ) + except Exception as e: + return RunResult( + problem_id=problem["id"], + mode=mode, + passed=False, + generated_code="", + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + latency_ms=0, + compression_ratio=compression_ratio, + error=str(e)[:500], + ) + + +# --------------------------------------------------------------------------- +# Aggregation +# --------------------------------------------------------------------------- + + +def aggregate(results: list[RunResult]) -> dict: + """Compute aggregate stats from a list of RunResults.""" + if not results: + return {} + passed = sum(1 for r in results if r.passed) + total = len(results) + return { + "pass_rate": round(passed / total * 100, 1), + "passed": passed, + "total": total, + "avg_prompt_tokens": round(statistics.mean(r.prompt_tokens for r in results)), + "avg_completion_tokens": round( + statistics.mean(r.completion_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), + "median_latency_ms": round(statistics.median(r.latency_ms for r in results), 1), + "avg_compression_ratio": round( + statistics.mean(r.compression_ratio for r in results), 4 + ), + } + + +# --------------------------------------------------------------------------- +# Main harness +# --------------------------------------------------------------------------- + + +def run_benchmark( + model: str, + num_problems: int = 0, + num_runs: int = 1, + padding_factor: int = 20, + compression_trigger: int = 2000, + embedding_model: Optional[str] = None, +) -> dict: + """ + Run the full benchmark. + + Parameters: + model: LLM model name (litellm format). + num_problems: How many problems to run (0 = all). + num_runs: Number of runs per mode. + padding_factor: How many distractor snippets to inject. Each snippet + adds ~400-600 tokens. 20 snippets ≈ 10k tokens of noise. + compression_trigger: Token count above which compression activates. + embedding_model: Optional embedding model for semantic scoring. + """ + problems = PROBLEMS[:num_problems] if num_problems > 0 else PROBLEMS + + print(f"\n{'=' * 60}") + print("Prompt Compression Eval Harness") + print(f"{'=' * 60}") + print(f"Model: {model}") + print(f"Problems: {len(problems)}") + print(f"Runs per mode: {num_runs}") + print(f"Padding factor: {padding_factor}") + print(f"Compression trigger:{compression_trigger} tokens") + print(f"Embedding model: {embedding_model or 'None (BM25 only)'}") + print(f"{'=' * 60}\n") + + baseline_results: list[RunResult] = [] + compressed_results: list[RunResult] = [] + + for run_i in range(num_runs): + if num_runs > 1: + print(f"--- Run {run_i + 1}/{num_runs} ---") + + for p in problems: + # Baseline (with padding, but no compression) + print(f" [{p['id']}] baseline ... ", end="", flush=True) + r = eval_problem( + p, + model, + padding_factor=padding_factor, + use_compression=False, + compression_trigger=compression_trigger, + embedding_model=embedding_model, + ) + baseline_results.append(r) + print("PASS" if r.passed else f"FAIL ({r.error[:60]})") + + # Compressed + print(f" [{p['id']}] compressed ... ", end="", flush=True) + r = eval_problem( + p, + model, + padding_factor=padding_factor, + use_compression=True, + compression_trigger=compression_trigger, + embedding_model=embedding_model, + ) + compressed_results.append(r) + status = "PASS" if r.passed else f"FAIL ({r.error[:60]})" + print(f"{status} (ratio: {r.compression_ratio:.2%})") + + # Aggregate + base_agg = aggregate(baseline_results) + comp_agg = aggregate(compressed_results) + + print(f"\n{'=' * 60}") + print("RESULTS") + print(f"{'=' * 60}") + print(f"\n Baseline (with {padding_factor} distractor snippets, no compression):") + print( + f" Pass rate: {base_agg['pass_rate']}% ({base_agg['passed']}/{base_agg['total']})" + ) + print(f" Avg prompt tokens: {base_agg['avg_prompt_tokens']}") + print(f" Avg total tokens: {base_agg['avg_total_tokens']}") + print(f" Avg latency: {base_agg['avg_latency_ms']}ms") + + print(f"\n Compressed (litellm.compress → then call model):") + print( + f" Pass rate: {comp_agg['pass_rate']}% ({comp_agg['passed']}/{comp_agg['total']})" + ) + print(f" Avg prompt tokens: {comp_agg['avg_prompt_tokens']}") + print(f" Avg total tokens: {comp_agg['avg_total_tokens']}") + print(f" Avg latency: {comp_agg['avg_latency_ms']}ms") + 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 + ) + latency_diff = base_agg["avg_latency_ms"] - comp_agg["avg_latency_ms"] + pass_diff = comp_agg["pass_rate"] - base_agg["pass_rate"] + + print(f"\n Delta (compressed vs baseline):") + print(f" Token savings: {token_savings} tokens ({token_pct}%)") + print(f" Latency delta: {latency_diff:+.1f}ms") + print(f" Pass rate delta: {pass_diff:+.1f}%") + + # Save JSON report + ts = time.strftime("%Y-%m-%d_%H-%M-%S") + report_path = f"eval_report_{ts}.json" + report = { + "model": model, + "timestamp": ts, + "num_problems": len(problems), + "num_runs": num_runs, + "padding_factor": padding_factor, + "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="Prompt Compression Evaluation Harness" + ) + parser.add_argument( + "--model", default="gpt-4o-mini", help="Model name (litellm format)" + ) + parser.add_argument( + "--problems", type=int, default=0, help="Number of problems (0 = all)" + ) + parser.add_argument("--runs", type=int, default=1, help="Number of runs per mode") + parser.add_argument( + "--padding-factor", + type=int, + default=20, + help="Number of distractor snippets to inject (default: 20, ~10k tokens)", + ) + parser.add_argument( + "--compression-trigger", + type=int, + default=2000, + help="Token count threshold to trigger compression (default: 2000)", + ) + 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, + num_runs=args.runs, + padding_factor=args.padding_factor, + compression_trigger=args.compression_trigger, + embedding_model=args.embedding_model, + ) diff --git a/tests/eval_swe_bench.py b/tests/eval_swe_bench.py new file mode 100644 index 0000000000..9c986283ab --- /dev/null +++ b/tests/eval_swe_bench.py @@ -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, + ) diff --git a/tests/test_litellm/test_compression.py b/tests/test_litellm/test_compression.py new file mode 100644 index 0000000000..13dda0cbcb --- /dev/null +++ b/tests/test_litellm/test_compression.py @@ -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}") From 6d2b94261ade84225b8ea31e16acccb38723d9fe Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 14 Apr 2026 10:54:31 -0700 Subject: [PATCH 2/2] fix(mypy): resolve type errors in compression/compress.py and __init__.py Cast message lists to the expected `List[Union[AllMessageValues, Message]]` type at `token_counter` call sites, and suppress the `no-redef` warning for the `compress` import in `__init__.py` caused by the wildcard `main` import. Co-Authored-By: Claude Opus 4.6 --- litellm/__init__.py | 2 +- litellm/compression/compress.py | 12 +++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 8b0da380fd..3b67d9e002 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1176,7 +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 +from .compression import compress # type: ignore[no-redef] # Skills API from .skills.main import ( diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 718bc1c45c..5baad460e1 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -3,7 +3,7 @@ Main compress() function — orchestrates BM25/embedding scoring, message stubbi and retrieval tool injection. """ -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, Optional, Set, Union, cast from litellm.caching.dual_cache import DualCache from litellm.compression.message_stubbing import ( @@ -15,6 +15,7 @@ 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: @@ -124,7 +125,9 @@ def compress( if compression_target is None: compression_target = compression_trigger * 7 // 10 - original_tokens = token_counter(model=model, messages=messages) + original_tokens = token_counter( + model=model, messages=cast(List[Union[AllMessageValues, Message]], messages) + ) # Pass through if below trigger if original_tokens <= compression_trigger: @@ -235,7 +238,10 @@ def compress( # Build retrieval tool tools = [build_retrieval_tool(list(cache.keys()))] if cache else [] - compressed_tokens = token_counter(model=model, messages=compressed_messages) + compressed_tokens = token_counter( + model=model, + messages=cast(List[Union[AllMessageValues, Message]], compressed_messages), + ) return CompressedResult( messages=compressed_messages,