From 6d2b94261ade84225b8ea31e16acccb38723d9fe Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 14 Apr 2026 10:54:31 -0700 Subject: [PATCH] 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 8b0da380fd0..3b67d9e0021 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 718bc1c45c3..5baad460e14 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,