From 0eade56aefd0efdfba0e9d0ed231599e747a3659 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 18 Apr 2026 15:23:14 -0700 Subject: [PATCH] refactor(compress): replace input_type with CallTypes call_type MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop the bespoke ``CompressionInputType`` literal and use the existing ``litellm.types.utils.CallTypes`` enum instead. ``litellm.compress()`` now takes ``call_type: Union[CallTypes, str]`` (default ``CallTypes.completion``) — no new concept to learn, and the enum is already the way the rest of the codebase talks about request shapes. Supported values: ``completion`` / ``acompletion`` (OpenAI chat-completions shape) and ``anthropic_messages`` (Anthropic structured content blocks). Updated: compress(), the compression_interception handler, tests, docs, and the two eval scripts. Co-Authored-By: Claude Opus 4.6 --- .../docs/completion/prompt_compression.md | 5 +- litellm/compression/compress.py | 82 ++++++++++++------- .../compression_interception/handler.py | 2 +- litellm/types/compression.py | 6 +- scripts/eval_compression.py | 3 +- tests/eval_swe_bench.py | 3 +- tests/test_litellm/test_compression.py | 45 +++++----- 7 files changed, 89 insertions(+), 57 deletions(-) diff --git a/docs/my-website/docs/completion/prompt_compression.md b/docs/my-website/docs/completion/prompt_compression.md index 88f8b52a79d..0d68ea2c101 100644 --- a/docs/my-website/docs/completion/prompt_compression.md +++ b/docs/my-website/docs/completion/prompt_compression.md @@ -8,6 +8,7 @@ The function keeps high-relevance and recent context, replaces low-relevance con ```python import litellm +from litellm.types.utils import CallTypes messages = [ {"role": "system", "content": "You are a coding assistant."}, @@ -19,7 +20,7 @@ messages = [ compressed = litellm.compress( messages=messages, model="gpt-4o", - input_type="openai_chat_completions", + call_type=CallTypes.completion, compression_trigger=1000, compression_target=500, ) @@ -46,7 +47,7 @@ response = litellm.completion( - `messages` (`List[dict]`, required): input conversation messages - `model` (`str`, required): model name used for token counting -- `input_type` (`Literal["anthropic_messages", "openai_chat_completions"]`, required): input message schema +- `call_type` (`CallTypes`, default `CallTypes.completion`): the LiteLLM call type whose message schema these messages follow. Supported values: `CallTypes.completion` / `CallTypes.acompletion` (OpenAI chat-completions shape) and `CallTypes.anthropic_messages` (Anthropic Messages shape) - `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 diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 4c4ce74fbbd..45795c9ca15 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -3,7 +3,7 @@ Main compress() function — normalizes input messages, orchestrates BM25/embedd scoring, message stubbing, and retrieval tool injection. """ -from typing import Any, Dict, List, Optional, Set, Tuple, cast +from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast from litellm.caching.dual_cache import DualCache from litellm.compression.message_stubbing import ( @@ -14,32 +14,55 @@ from litellm.compression.message_stubbing import ( 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, CompressionInputType +from litellm.types.compression import CompressedResult +from litellm.types.utils import CallTypes + +# CallTypes that produce Anthropic-shaped messages (structured content blocks). +# Everything else is treated as OpenAI chat-completions shape. +_ANTHROPIC_CALL_TYPES = frozenset({CallTypes.anthropic_messages.value}) +# CallTypes that are valid targets for compression. Compression operates on +# message-shaped inputs, so we only accept call types whose payload is a list +# of role/content messages. +_SUPPORTED_CALL_TYPES = frozenset( + { + CallTypes.completion.value, + CallTypes.acompletion.value, + CallTypes.anthropic_messages.value, + } +) -def _build_retrieval_tools(keys: List[str], input_type: CompressionInputType) -> List[dict]: +def _normalize_call_type(call_type: Union[CallTypes, str]) -> str: + """Return the string value for a ``CallTypes`` enum or a raw string.""" + if isinstance(call_type, CallTypes): + return call_type.value + return call_type + + +def _is_anthropic_call_type(call_type: str) -> bool: + return call_type in _ANTHROPIC_CALL_TYPES + + +def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]: """ Build retrieval tool definitions in the target request schema. - - OpenAI chat completions: keep OpenAI function-tool schema. - - Anthropic messages: remap OpenAI function-tool schema to Anthropic custom tool. + - Chat-completions call types: keep OpenAI function-tool schema. + - Anthropic messages call type: remap to Anthropic's custom tool schema. """ if not keys: return [] openai_tools = [build_retrieval_tool(keys)] - if input_type == "openai_chat_completions": + if not _is_anthropic_call_type(call_type): return openai_tools - if input_type == "anthropic_messages": - # Lazy import to avoid introducing provider transformation imports - # during module import for non-Anthropic call paths. - from litellm.llms.anthropic.chat.transformation import AnthropicConfig + # Lazy import to avoid introducing provider transformation imports during + # module import for non-Anthropic call paths. + from litellm.llms.anthropic.chat.transformation import AnthropicConfig - anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools) - return cast(List[dict], anthropic_tools) - - return openai_tools + anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools) + return cast(List[dict], anthropic_tools) def _content_to_text(content: Any) -> str: @@ -74,7 +97,7 @@ def _content_to_text(content: Any) -> str: def _normalize_messages_for_compression( messages: List[dict], - input_type: CompressionInputType, + call_type: str, ) -> Tuple[List[dict], List[dict]]: """ Normalize each original message to a text-surrogate content for scoring. @@ -82,10 +105,10 @@ def _normalize_messages_for_compression( Returns: (normalized_messages, original_messages_copy) """ - if input_type not in ("anthropic_messages", "openai_chat_completions"): + if call_type not in _SUPPORTED_CALL_TYPES: raise ValueError( - f"Unsupported input_type={input_type}. " - "Expected 'anthropic_messages' or 'openai_chat_completions'." + f"Unsupported call_type={call_type!r} for compression. " + f"Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}." ) original_messages: List[Dict[str, Any]] = [dict(m) for m in messages] @@ -324,7 +347,7 @@ def _get_dropped_tool_span_indices( def compress( messages: List[dict], model: str, - input_type: CompressionInputType = "openai_chat_completions", + call_type: Union[CallTypes, str] = CallTypes.completion, compression_trigger: int = 200_000, compression_target: Optional[int] = None, embedding_model: Optional[str] = None, @@ -343,10 +366,12 @@ def compress( Parameters: messages: The conversation messages to (potentially) compress. model: The LLM model name — used for token counting. - input_type: Message format of input messages. Must be either: - - "anthropic_messages" - - "openai_chat_completions" - Defaults to "openai_chat_completions" for backward compatibility. + call_type: The LiteLLM call type whose message schema these messages + follow. Supported values: + - ``CallTypes.completion`` / ``CallTypes.acompletion`` — OpenAI + chat-completions shape (default) + - ``CallTypes.anthropic_messages`` — Anthropic Messages shape + (structured content blocks + atomic tool exchanges) compression_trigger: Only compress if input exceeds this token count. compression_target: Target token count after compression. Defaults to ``compression_trigger // 2``. @@ -361,9 +386,10 @@ def compress( A ``CompressedResult`` dict containing compressed messages, token counts, a cache of original content, and the retrieval tool definition. """ + call_type_str = _normalize_call_type(call_type) normalized_messages, original_messages = _normalize_messages_for_compression( messages=messages, - input_type=input_type, + call_type=call_type_str, ) if compression_target is None: @@ -413,9 +439,9 @@ def compress( kept_indices: Set[int] = set(protected_indices) tool_exchange_spans: List[Set[int]] = [] - if input_type == "anthropic_messages": - tool_exchange_spans, tool_sequence_error = _extract_anthropic_tool_exchange_spans( - original_messages + if _is_anthropic_call_type(call_type_str): + tool_exchange_spans, tool_sequence_error = ( + _extract_anthropic_tool_exchange_spans(original_messages) ) if tool_sequence_error is not None: return CompressedResult( @@ -466,7 +492,7 @@ def compress( compressed_messages.append(stub_message(msg, key)) # Build retrieval tool in the target request schema - tools = _build_retrieval_tools(list(cache.keys()), input_type=input_type) + tools = _build_retrieval_tools(list(cache.keys()), call_type=call_type_str) compressed_tokens = token_counter( model=model, diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 925ae1e1042..c6ae7d9e82b 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -99,7 +99,7 @@ class CompressionInterceptionLogger(CustomLogger): compressed = compress( # type: ignore messages=messages, model=model, - input_type="anthropic_messages", + call_type=CallTypes.anthropic_messages, compression_trigger=self.compression_trigger, compression_target=self.compression_target, embedding_model=self.embedding_model, diff --git a/litellm/types/compression.py b/litellm/types/compression.py index 7b14db70942..5dae0c397f0 100644 --- a/litellm/types/compression.py +++ b/litellm/types/compression.py @@ -5,14 +5,12 @@ Type definitions for litellm.compress(). import sys if sys.version_info >= (3, 11): - from typing import Dict, List, Literal, NotRequired, TypedDict + from typing import Dict, List, NotRequired, TypedDict else: - from typing import Dict, List, Literal, TypedDict + from typing import Dict, List, TypedDict from typing_extensions import NotRequired -CompressionInputType = Literal["anthropic_messages", "openai_chat_completions"] - class CompressedResult(TypedDict): messages: List[dict] # compressed messages (stubs replace low-relevance messages) diff --git a/scripts/eval_compression.py b/scripts/eval_compression.py index c35742f6da6..a169cc02d74 100644 --- a/scripts/eval_compression.py +++ b/scripts/eval_compression.py @@ -33,6 +33,7 @@ from dataclasses import asdict, dataclass, field from typing import Optional import litellm +from litellm.types.utils import CallTypes # --------------------------------------------------------------------------- # Problem definitions (HumanEval-style) @@ -880,7 +881,7 @@ def eval_problem( result = litellm.compress( messages=messages, model=model, - input_type="openai_chat_completions", + call_type=CallTypes.completion, compression_trigger=compression_trigger, embedding_model=embedding_model, ) diff --git a/tests/eval_swe_bench.py b/tests/eval_swe_bench.py index 9c986283abd..6ae99f83ca1 100644 --- a/tests/eval_swe_bench.py +++ b/tests/eval_swe_bench.py @@ -40,6 +40,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) import litellm # noqa: E402 from litellm.compression import compress as litellm_compress # noqa: E402 +from litellm.types.utils import CallTypes # noqa: E402 # --------------------------------------------------------------------------- # Prompts @@ -445,7 +446,7 @@ def eval_instance( compress_kwargs: dict = { "messages": messages, "model": model, - "input_type": "openai_chat_completions", + "call_type": CallTypes.completion, "compression_trigger": compression_trigger, "embedding_model": embedding_model, } diff --git a/tests/test_litellm/test_compression.py b/tests/test_litellm/test_compression.py index 16fd9bb69ec..4fbcd4ed30d 100644 --- a/tests/test_litellm/test_compression.py +++ b/tests/test_litellm/test_compression.py @@ -13,9 +13,10 @@ from litellm.compression.scoring.embedding_scorer import embedding_score_message 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 +from litellm.types.utils import CallTypes -INPUT_TYPE = "openai_chat_completions" -ANTHROPIC_INPUT_TYPE = "anthropic_messages" +CALL_TYPE = CallTypes.completion +ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages # --------------------------------------------------------------------------- @@ -153,7 +154,7 @@ def test_retrieval_tool_description_lists_keys(): def test_compress_below_trigger_passthrough(): messages = [{"role": "user", "content": "hello"}] - result = litellm.compress(messages, model="gpt-4o", input_type=INPUT_TYPE) + result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE) assert result["messages"] == messages assert result["cache"] == {} assert result["tools"] == [] @@ -182,7 +183,7 @@ def test_compress_above_trigger(): result = litellm.compress( big_messages, model="gpt-4o", - input_type=INPUT_TYPE, + call_type=CALL_TYPE, compression_trigger=1000, compression_target=500, ) @@ -226,7 +227,7 @@ def test_compress_anthropic_list_content_is_boundary_stable(): result = litellm.compress( messages=messages, model="claude-sonnet-4-20250514", - input_type=ANTHROPIC_INPUT_TYPE, + call_type=ANTHROPIC_CALL_TYPE, compression_trigger=1000, compression_target=500, ) @@ -248,7 +249,7 @@ def test_compress_preserves_system_message(): {"role": "user", "content": "Fix the bug"}, ] result = litellm.compress( - messages, model="gpt-4o", input_type=INPUT_TYPE, compression_trigger=1000 + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 ) assert result["messages"][0]["role"] == "system" assert "System prompt" in result["messages"][0]["content"] @@ -260,7 +261,7 @@ def test_compress_preserves_last_user_message(): {"role": "user", "content": "Fix the bug in auth.py"}, ] result = litellm.compress( - messages, model="gpt-4o", input_type=INPUT_TYPE, compression_trigger=1000 + messages, model="gpt-4o", call_type=CALL_TYPE, 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"] @@ -273,7 +274,7 @@ def test_compress_preserves_last_assistant_message(): {"role": "user", "content": "Now fix the bug"}, ] result = litellm.compress( - messages, model="gpt-4o", input_type=INPUT_TYPE, compression_trigger=1000 + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 ) assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"] assert len(assistant_msgs) >= 1 @@ -288,7 +289,7 @@ def test_cache_keys_match_stubs(): {"role": "user", "content": "Fix it"}, ] result = litellm.compress( - messages, model="gpt-4o", input_type=INPUT_TYPE, compression_trigger=1000 + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 ) if result["tools"]: tool_desc = result["tools"][0]["function"]["description"] @@ -303,7 +304,7 @@ def test_compress_default_target(): {"role": "user", "content": "query"}, ] result = litellm.compress( - messages, model="gpt-4o", input_type=INPUT_TYPE, compression_trigger=2000 + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000 ) # Should have compressed — target = 1000 assert result["compressed_tokens"] <= result["original_tokens"] @@ -345,7 +346,7 @@ def test_compress_nested_tool_result_extracts_text_only(): result = litellm.compress( messages=messages, model="claude-sonnet-4-20250514", - input_type=ANTHROPIC_INPUT_TYPE, + call_type=ANTHROPIC_CALL_TYPE, compression_trigger=500, compression_target=100, ) @@ -356,7 +357,7 @@ def test_compress_nested_tool_result_extracts_text_only(): assert "https://example.com/top.png" not in cached_text -def test_compress_default_input_type_is_openai_chat_completions(): +def test_compress_default_call_type_is_completion(): result = litellm.compress( messages=[ {"role": "user", "content": "Large context " * 4000}, @@ -393,7 +394,7 @@ def test_compress_forwards_embedding_model_params(monkeypatch): {"role": "user", "content": "Fix auth"}, ], model="gpt-4o", - input_type=INPUT_TYPE, + call_type=CALL_TYPE, compression_trigger=1000, embedding_model="text-embedding-3-small", embedding_model_params={"api_base": "https://example-embeddings.test"}, @@ -451,7 +452,7 @@ def test_embedding_scorer(): {"role": "user", "content": "Fix auth"}, ], model="gpt-4o", - input_type=INPUT_TYPE, + call_type=CALL_TYPE, compression_trigger=1000, embedding_model="text-embedding-3-small", ) @@ -473,7 +474,7 @@ def test_simple_compression(final_user_message, expected_content): {"role": "user", "content": final_user_message}, ] result = litellm.compress( - messages, model="gpt-4o", input_type=INPUT_TYPE, compression_trigger=1000 + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 ) if expected_content == "Unrelated cooking recipes ": assert "Unrelated cooking recipes " in result["messages"][1]["content"] @@ -511,7 +512,9 @@ def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch): return 1 return 10 - monkeypatch.setattr(compress_module, "bm25_score_messages", fake_bm25_score_messages) + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) messages = [ @@ -544,7 +547,7 @@ def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch): result = litellm.compress( messages=messages, model="claude-sonnet-4-20250514", - input_type=ANTHROPIC_INPUT_TYPE, + call_type=ANTHROPIC_CALL_TYPE, compression_trigger=100, compression_target=280, ) @@ -584,7 +587,9 @@ def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch): return 1 return 10 - monkeypatch.setattr(compress_module, "bm25_score_messages", fake_bm25_score_messages) + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) messages = [ @@ -617,7 +622,7 @@ def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch): result = litellm.compress( messages=messages, model="claude-sonnet-4-20250514", - input_type=ANTHROPIC_INPUT_TYPE, + call_type=ANTHROPIC_CALL_TYPE, compression_trigger=100, compression_target=280, ) @@ -651,7 +656,7 @@ def test_compress_anthropic_malformed_tool_sequence_passes_through(): result = litellm.compress( messages=messages, model="claude-sonnet-4-20250514", - input_type=ANTHROPIC_INPUT_TYPE, + call_type=ANTHROPIC_CALL_TYPE, compression_trigger=100, compression_target=280, )