scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(mypy): resolve type errors in compression/compress.py and __init__.py

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia 2026-04-14 10:54:31 -07:00
parent 26c7412339
commit 6d2b94261a
2 changed files with 10 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/__init__.py
View file

@ -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 (

12
litellm/compression/compress.py
View file

@ -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,