From d86c2e1f4238e0689e9732b1f531d29ea8eb434e Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
<158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Fri, 25 Sep 2026 16:06:44 -0700
Subject: [PATCH 1/6] fix(logging): redact raw_request when
turn_off_message_logging is set in the proxy config (#43219)
* fix(logging): redact raw_request when turn_off_message_logging is set in the proxy config
The raw request branch bound turn_off_message_logging by name at import, before the proxy config set it, so loggers kept receiving the prompt in metadata.raw_request and raw_request_typed_dict. It now runs the same per-request redaction check messages use, and json_logs is read at call time for the same reason
* fix(logging): keep raw_request_typed_dict for the explicit readers and tolerate missing headers in the json debug log
* refactor(logging): drop the stale comment above the raw request typed dict
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
litellm/litellm_core_utils/litellm_logging.py | 61 +++++++------------
.../test_litellm_logging.py | 61 ++++++++++++++++++-
tests/unit/test_main.py | 23 +++++++
3 files changed, 106 insertions(+), 39 deletions(-)
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 0ef5bbf807a..28d72702f3e 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -20,11 +20,7 @@ from httpx import Response
from pydantic import BaseModel, JsonValue
import litellm
-from litellm import (
- _custom_logger_compatible_callbacks_literal,
- json_logs,
- turn_off_message_logging,
-)
+from litellm import _custom_logger_compatible_callbacks_literal
from litellm._logging import (
_is_debugging_on,
_redact_string,
@@ -43,6 +39,7 @@ from litellm.constants import (
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
EMPTY_MAPPING,
PROVIDER_REQUEST_ID_HEADERS,
+ REDACTED_BY_LITELLM,
)
from litellm.cost_calculator import (
RealtimeAPITokenUsageProcessor,
@@ -1358,10 +1355,19 @@ class Logging(LiteLLMLoggingBaseClass):
_litellm_params: Final = self.model_call_details.get("litellm_params", {})
_metadata: Final = _litellm_params.get("metadata", {}) or {}
try:
- # [Non-blocking Extra Debug Information in metadata]
- if turn_off_message_logging is True:
- _metadata["raw_request"] = "redacted by litellm. \
- 'litellm.turn_off_message_logging=True'"
+ self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict(
+ raw_request_api_base=self._get_masked_api_base(str(additional_args.get("api_base") or "")),
+ raw_request_body=self._get_raw_request_body(additional_args.get("complete_input_dict", {})),
+ # NOTE: setting ignore_sensitive_headers to True will cause
+ # the Authorization header to be leaked when calls to the health
+ # endpoint are made and fail.
+ raw_request_headers=self._get_masked_headers(
+ additional_args.get("headers", {}) or {},
+ ),
+ error=None,
+ )
+ if should_redact_message_logging(self.model_call_details):
+ _metadata["raw_request"] = REDACTED_BY_LITELLM
else:
curl_command: Final = self._get_request_curl_command(
api_base=additional_args.get("api_base", ""),
@@ -1369,20 +1375,7 @@ class Logging(LiteLLMLoggingBaseClass):
additional_args=additional_args,
data=additional_args.get("complete_input_dict", {}),
)
-
_metadata["raw_request"] = _redact_string(str(curl_command))
- # split up, so it's easier to parse in the UI
- self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict(
- raw_request_api_base=self._get_masked_api_base(str(additional_args.get("api_base") or "")),
- raw_request_body=self._get_raw_request_body(additional_args.get("complete_input_dict", {})),
- # NOTE: setting ignore_sensitive_headers to True will cause
- # the Authorization header to be leaked when calls to the health
- # endpoint are made and fail.
- raw_request_headers=self._get_masked_headers(
- additional_args.get("headers", {}) or {},
- ),
- error=None,
- )
except Exception as e:
self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict(
error=str(e),
@@ -1476,7 +1469,7 @@ class Logging(LiteLLMLoggingBaseClass):
def _print_llm_call_debugging_log(
self,
api_base: str,
- headers: dict,
+ headers: dict | None,
additional_args: dict,
):
"""
@@ -1485,8 +1478,8 @@ class Logging(LiteLLMLoggingBaseClass):
Prints the RAW curl command sent from LiteLLM
"""
if _is_debugging_on() or self.litellm_request_debug:
- if json_logs:
- masked_headers: Final = self._get_masked_headers(headers)
+ if litellm.json_logs:
+ masked_headers: Final = self._get_masked_headers(headers or {})
masked_api_base: Final = self._get_masked_api_base(str(api_base or ""))
if self.litellm_request_debug:
verbose_logger.warning( # .warning ensures this shows up in all environments
@@ -1563,20 +1556,12 @@ class Logging(LiteLLMLoggingBaseClass):
else:
attr = "debug"
- if json_logs:
- callattr = verbose_logger.warning if attr == "warning" else verbose_logger.debug
- callattr(
- "RAW RESPONSE:\n{}\n\n".format(
- self.model_call_details.get("original_response", self.model_call_details)
- ),
- )
- else:
- callattr = verbose_logger.warning if attr == "warning" else verbose_logger.debug
- callattr(
- "RAW RESPONSE:\n{}\n\n".format(
- self.model_call_details.get("original_response", self.model_call_details)
- )
+ callattr: Final = verbose_logger.warning if attr == "warning" else verbose_logger.debug
+ callattr(
+ "RAW RESPONSE:\n{}\n\n".format(
+ self.model_call_details.get("original_response", self.model_call_details)
)
+ )
if getattr(self, "logger_fn", None) and callable(self.logger_fn):
try:
self.logger_fn(
diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
index 265fdb50836..d717718cba2 100644
--- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
@@ -20,7 +20,7 @@ from openai._legacy_response import HttpxBinaryResponseContent
import litellm
from litellm._logging import session_id_var, trace_id_var
-from litellm.constants import SENTRY_PII_DENYLIST
+from litellm.constants import REDACTED_BY_LITELLM, SENTRY_PII_DENYLIST
from litellm.cost_calculator import ocr_batch_cost
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
@@ -6615,6 +6615,65 @@ def test_pre_call_redacts_and_masks_raw_request(logging_obj):
assert "key=*****" in raw_api_base
+_PRIVATE_RAW_REQUEST_ARGS: Final = {
+ "api_base": "https://api.openai.com/v1/chat/completions",
+ "headers": {},
+ "complete_input_dict": {"messages": [{"role": "user", "content": "PRIVATE-PHRASE"}]},
+}
+
+
+def _pre_call_with_raw_request_logging(logging_obj) -> dict:
+ metadata: Final = {"user_api_key_alias": "qa-key"}
+ logging_obj.model_call_details["litellm_params"] = {"metadata": metadata}
+ logging_obj.log_raw_request_response = True
+ logging_obj.pre_call(input="hi", api_key="", additional_args=_PRIVATE_RAW_REQUEST_ARGS)
+ return metadata
+
+
+def _assert_raw_request_redacted_for_callbacks_only(logging_obj, metadata: dict) -> None:
+ assert metadata["raw_request"] == REDACTED_BY_LITELLM
+ typed_dict: Final = logging_obj.model_call_details["raw_request_typed_dict"]
+ assert typed_dict["raw_request_body"] == _PRIVATE_RAW_REQUEST_ARGS["complete_input_dict"]
+ assert typed_dict["error"] is None
+
+
+def test_pre_call_raw_request_honors_turn_off_message_logging_set_after_import(logging_obj, monkeypatch):
+ monkeypatch.setattr(litellm, "turn_off_message_logging", True)
+
+ metadata = _pre_call_with_raw_request_logging(logging_obj)
+
+ _assert_raw_request_redacted_for_callbacks_only(logging_obj, metadata)
+
+
+def test_pre_call_raw_request_honors_per_request_turn_off_message_logging(logging_obj, monkeypatch):
+ monkeypatch.setattr(litellm, "turn_off_message_logging", False)
+ logging_obj.model_call_details["standard_callback_dynamic_params"] = {"turn_off_message_logging": True}
+
+ metadata = _pre_call_with_raw_request_logging(logging_obj)
+
+ _assert_raw_request_redacted_for_callbacks_only(logging_obj, metadata)
+
+
+def test_debugging_log_honors_json_logs_set_after_import(logging_obj, monkeypatch):
+ monkeypatch.setattr(litellm, "json_logs", True)
+ logging_obj.litellm_request_debug = True
+
+ with patch("litellm.litellm_core_utils.litellm_logging.verbose_logger.warning") as warning:
+ logging_obj._print_llm_call_debugging_log(api_base="https://api.openai.com/v1", headers={}, additional_args={})
+
+ assert "https://api.openai.com/v1" in warning.call_args.kwargs["extra"]["api_base"]
+
+
+def test_debugging_log_with_json_logs_tolerates_missing_headers(logging_obj, monkeypatch):
+ monkeypatch.setattr(litellm, "json_logs", True)
+ logging_obj.litellm_request_debug = True
+
+ with patch("litellm.litellm_core_utils.litellm_logging.verbose_logger.warning") as warning:
+ logging_obj._print_llm_call_debugging_log(api_base="https://api.openai.com/v1", headers=None, additional_args={})
+
+ assert "https://api.openai.com/v1" in warning.call_args.kwargs["extra"]["api_base"]
+
+
def _streaming_logging_obj_with_callbacks(callbacks: list[CustomLogger]):
import datetime
diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py
index effc038f85b..c06216e4f4e 100644
--- a/tests/unit/test_main.py
+++ b/tests/unit/test_main.py
@@ -626,6 +626,29 @@ def test_return_raw_request_does_not_call_provider(respx_mock: respx.MockRouter)
]
+def test_return_raw_request_ignores_turn_off_message_logging(
+ respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+):
+ from litellm.types.utils import CallTypes
+ from litellm.utils import return_raw_request
+
+ model: Final = "gpt-4o"
+ messages: Final = [{"role": "user", "content": "PRIVATE-PHRASE"}]
+ route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
+ return_value=_mocked_openai_chat_response(model)
+ )
+ monkeypatch.setattr(litellm, "turn_off_message_logging", True)
+
+ request: Final = return_raw_request(
+ endpoint=CallTypes.completion,
+ kwargs={"model": model, "messages": messages},
+ )
+
+ assert route.call_count == 0
+ assert request.get("error") is None
+ assert request["raw_request_body"]["messages"] == messages
+
+
def test_completion_forwards_verbosity_in_raw_request(respx_mock: respx.MockRouter):
"""Regression test: completion() must forward the verbosity param to the provider request body."""
from litellm.types.utils import CallTypes
From cd1107aac481278561d6d0f9da054ba6bf4095a5 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
<158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Fri, 25 Sep 2026 16:18:51 -0700
Subject: [PATCH 2/6] fix(router): parse the classifier verdict out of
surrounding prose instead of falling to the default tier (#43215)
* fix(router): parse the classifier verdict out of prose and fences on every parse path
The complexity router's classifier parsers only tolerated a bare JSON object
(labeled tier) or a leading Markdown fence (capability, LLM V2), so a
json_object classifier that writes its verdict fenced and then explains it in
markdown, which Bedrock Haiku 4.5 does on nearly every Claude Code request,
failed validation and every request fell to the fallback tier.
All three parse paths now extract the first complete JSON object from the
reply with json.JSONDecoder.raw_decode, whatever prose or fence surrounds it,
and a reply that still fails validation is logged with the pydantic field
problems and the raw reply text, withheld when the request turns off message
logging. The capability failure reason names the exception type like the
labeled path does instead of interpolating str(e), which for a ValidationError
carried the whole reply as input_value and for TimeoutError was empty.
* fix(router): stop rejecting an LLM V2 forecast over a long explanation field
LLMV2Verdict capped crux and each forecast's likely_failure at 512 characters
through the ShortText alias, so a verdict whose explanation ran long failed
validation and the request fell to the capable tier, even though nothing
downstream reads either field. Five of nineteen real Claude Code replies from
Bedrock Haiku 4.5 tripped the cap. Both fields keep the strip and non-empty
constraints and lose the length cap; the operator-set calibration version keeps
ShortText.
* fix(complexity_router): withhold the rejected classifier reply under every message-logging opt-out and survive undecodable replies
* fix(complexity_router): withhold the rejected classifier reply when the redaction decision cannot be made
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
.../capability_classifier.py | 29 +++--
.../complexity_router/complexity_router.py | 57 ++++++++-
.../complexity_router/llm_v2.py | 5 +-
.../router_strategy/test_complexity_router.py | 81 ++++++++++++-
.../router_strategy/test_llm_v2.py | 113 +++++++++++++++++-
5 files changed, 261 insertions(+), 24 deletions(-)
diff --git a/litellm/router_strategy/complexity_router/capability_classifier.py b/litellm/router_strategy/complexity_router/capability_classifier.py
index 21046ff3421..93077af9e47 100644
--- a/litellm/router_strategy/complexity_router/capability_classifier.py
+++ b/litellm/router_strategy/complexity_router/capability_classifier.py
@@ -202,15 +202,26 @@ def capability_classifier_system_prompt(mode: Literal["json_schema", "json_objec
)
-def unwrap_classifier_json(content: str) -> str:
- """Remove the optional Markdown fence without repairing or weakening verdict JSON."""
- text: Final = content.strip()
- if not text.startswith("```"):
- return text
- unfenced: Final = text.removeprefix("```").removeprefix("json").lstrip("\n\r")
- return unfenced.removesuffix("```").strip()
+_JSON_DECODER: Final = json.JSONDecoder()
+
+
+def _complete_json_object_at(content: str, start: int) -> str | None:
+ try:
+ _, end = _JSON_DECODER.raw_decode(content, start)
+ except (ValueError, RecursionError):
+ return None
+ return content[start:end]
+
+
+def extract_classifier_json(content: str) -> str:
+ """Return the first complete JSON object in the reply, whatever prose or fence surrounds it.
+
+ A reply with no complete object comes back stripped so the caller's validation names the defect."""
+ object_starts: Final = (index for index, char in enumerate(content) if char == "{")
+ candidates: Final = (_complete_json_object_at(content, start) for start in object_starts)
+ return next((candidate for candidate in candidates if candidate is not None), content.strip())
def parse_capability_classifier_verdict(content: str) -> CapabilityClassifierVerdict:
- """Parse raw JSON or the fenced JSON shape tolerated by Switchyard."""
- return CapabilityClassifierVerdict.model_validate_json(unwrap_classifier_json(content))
+ """Parse the verdict object out of a bare, fenced, or prose-wrapped reply."""
+ return CapabilityClassifierVerdict.model_validate_json(extract_classifier_json(content))
diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py
index 0f252952a9d..9df6306436b 100644
--- a/litellm/router_strategy/complexity_router/complexity_router.py
+++ b/litellm/router_strategy/complexity_router/complexity_router.py
@@ -29,6 +29,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
from pydantic import BaseModel, TypeAdapter, ValidationError, create_model
+from pydantic_core import ErrorDetails
from litellm._logging import verbose_router_logger
from litellm.caching.affinity_cache import claim_affinity_pin
@@ -85,8 +86,8 @@ from .capability_classifier import (
CapabilityClassifierForecast,
capability_classifier_response_format,
capability_classifier_system_prompt,
+ extract_classifier_json,
parse_capability_classifier_verdict,
- unwrap_classifier_json,
)
from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section
from .config import (
@@ -427,6 +428,41 @@ def _effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | N
)
+def _classifier_reply_is_private(request_kwargs: Mapping[str, object] | None) -> bool:
+ from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
+ initialize_standard_callback_dynamic_params,
+ )
+ from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
+
+ kwargs: Final = dict(request_kwargs) if request_kwargs else {}
+ try:
+ return should_redact_message_logging(
+ {
+ "litellm_params": kwargs,
+ "standard_callback_dynamic_params": initialize_standard_callback_dynamic_params(kwargs),
+ }
+ )
+ except AttributeError:
+ return True
+
+
+def _validation_problem(detail: ErrorDetails) -> str:
+ location: Final = ".".join(str(part) for part in detail["loc"])
+ return f"{location}: {detail['msg']}" if location else detail["msg"]
+
+
+def _log_rejected_classifier_verdict(
+ error: ValidationError, content: str, request_kwargs: Mapping[str, object] | None
+) -> None:
+ problems: Final = "; ".join(_validation_problem(detail) for detail in error.errors())
+ reply: Final = (
+ "raw reply withheld (message logging is off)"
+ if _classifier_reply_is_private(request_kwargs)
+ else f"raw reply: {content!r}"
+ )
+ verbose_router_logger.warning("ComplexityRouter: classifier verdict rejected (%s); %s", problems, reply)
+
+
_REMINDER_OPEN: Final = ""
_REMINDER_CLOSE: Final = ""
_DEFAULT_REMINDER_MARKERS: Final = ((_REMINDER_OPEN, _REMINDER_CLOSE),)
@@ -2040,7 +2076,7 @@ class ComplexityRouter(CustomLogger):
except Exception as e: # noqa: BLE001 -- every unavailable or invalid judge verdict must fail closed
if breaker is not None and permit is not None:
breaker.record_failure(permit, is_timeout=_is_classifier_timeout(e))
- return self._capability_classifier_failure_outcome(f"capability classifier failed ({e})")
+ return self._capability_classifier_failure_outcome(f"capability classifier failed ({type(e).__name__})")
def _capability_classifier_failure_outcome(self, reason: str, signal: str | None = None) -> ClassificationOutcome:
"""Fail closed to the configured capable tier without consulting another taxonomy."""
@@ -2449,7 +2485,11 @@ class ComplexityRouter(CustomLogger):
content, classifier_cost = await self._call_classifier_model(
messages_for_call, request_kwargs, encrypted_task=encrypted_task
)
- raw_tier: Final = _LabeledTierClassification.model_validate_json(content).tier
+ try:
+ raw_tier: Final = _LabeledTierClassification.model_validate_json(extract_classifier_json(content)).tier
+ except ValidationError as error:
+ _log_rejected_classifier_verdict(error, content, request_kwargs)
+ raise
tier: Final = self.config.resolve_classified_tier(raw_tier)
if tier is None:
raise ValueError(f"LLM classifier returned an unrecognized tier: {raw_tier!r}")
@@ -2508,7 +2548,11 @@ class ComplexityRouter(CustomLogger):
max_output_tokens=capability.max_output_tokens,
encrypted_task=encrypted_task,
)
- verdict: Final = parse_capability_classifier_verdict(content)
+ try:
+ verdict: Final = parse_capability_classifier_verdict(content)
+ except ValidationError as error:
+ _log_rejected_classifier_verdict(error, content, request_kwargs)
+ raise
threshold: Final = verdict.routing_threshold(capability.base_threshold, capability.threshold_step)
calibration: Final = capability.calibration
forecast: Final = CapabilityClassifierForecast(
@@ -2563,8 +2607,9 @@ class ComplexityRouter(CustomLogger):
messages_for_call, request_kwargs, encrypted_task=encrypted, max_output_tokens=v2.max_output_tokens
)
try:
- verdict: Final = LLMV2Verdict.model_validate_json(unwrap_classifier_json(content))
- except ValidationError:
+ verdict: Final = LLMV2Verdict.model_validate_json(extract_classifier_json(content))
+ except ValidationError as error:
+ _log_rejected_classifier_verdict(error, content, request_kwargs)
return self._classifier_failure_outcome("Invalid LLM V2 forecast", prompt, system_prompt)._replace(
classifier_cost=classifier_cost
)
diff --git a/litellm/router_strategy/complexity_router/llm_v2.py b/litellm/router_strategy/complexity_router/llm_v2.py
index 18351237e65..8ef2f554ab2 100644
--- a/litellm/router_strategy/complexity_router/llm_v2.py
+++ b/litellm/router_strategy/complexity_router/llm_v2.py
@@ -16,6 +16,7 @@ from litellm.llms.base_llm.base_utils import (
from litellm.router_strategy.complexity_router.fuse_presets import ProfileText, resolve_fuse_profile
ShortText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=512)]
+VerdictText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
class _SolverProfile(TypedDict):
@@ -90,7 +91,7 @@ class LLMV2Demands(BaseModel):
class LLMV2SolverForecast(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
- likely_failure: ShortText
+ likely_failure: VerdictText
p_solve: StrictFloat = Field(ge=0.0, le=1.0)
@@ -104,7 +105,7 @@ class LLMV2SolverForecasts(BaseModel):
class LLMV2Verdict(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
- crux: ShortText
+ crux: VerdictText
demands: LLMV2Demands
verification: Literal["relevant", "partial", "unavailable", "unknown"]
forecasts: LLMV2SolverForecasts
diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py
index 4e146b59b61..3401a335b2f 100644
--- a/tests/test_litellm/router_strategy/test_complexity_router.py
+++ b/tests/test_litellm/router_strategy/test_complexity_router.py
@@ -2679,6 +2679,25 @@ def _llm_response(content: str, response_cost: float | None = None):
return response
+_REPLY_SHAPES: Final = ("fenced", "fenced-with-language", "prose-before", "prose-after", "fenced-then-prose")
+
+
+def _wrapped_reply(shape: str, verdict: str) -> str:
+ match shape:
+ case "fenced":
+ return f" ```\n{verdict}\n``` "
+ case "fenced-with-language":
+ return f"```json\n{verdict}\n```"
+ case "prose-before":
+ return f"Sure {{here}} is the verdict you asked for:\n\n{verdict}"
+ case "prose-after":
+ return f"{verdict}\n\nThe efficient solver should handle this {{well}}."
+ case "fenced-then-prose":
+ return f"```json\n{verdict}\n```\n\n## Reasoning\n\nThe task is coupled, so the forecasts differ."
+ case _:
+ raise AssertionError(shape)
+
+
@pytest.fixture
def llm_classifier_config() -> Dict:
"""Config with an LLM-based classifier wired to a 'haiku-classifier' model."""
@@ -3132,12 +3151,40 @@ class TestCapabilityClassifier:
assert outcome.capability_forecast.threshold == pytest.approx(expected_threshold)
@pytest.mark.asyncio
- async def test_fenced_json_verdict_is_accepted(self, mock_router_instance):
- reply = _capability_reply(p_solve=0.8)
- mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(f"```json\n{reply}\n```"))
+ @pytest.mark.parametrize("shape", _REPLY_SHAPES)
+ async def test_verdict_wrapped_in_fence_or_prose_is_accepted(self, mock_router_instance, shape: str):
+ reply = _wrapped_reply(shape, _capability_reply(p_solve=0.8))
+ mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply))
outcome = await self._router(mock_router_instance).aclassify("do the task")
assert outcome.tier == ComplexityTier.SIMPLE
assert outcome.cause == "capability_classifier"
+ assert outcome.capability_forecast is not None
+ assert outcome.capability_forecast.p_solve == 0.8
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("message_logging_off", (False, True))
+ async def test_unparseable_reply_is_logged_with_its_text_unless_message_logging_is_off(
+ self, mock_router_instance, caplog: pytest.LogCaptureFixture, message_logging_off: bool
+ ):
+ reply = "The task text is too {vague} for a forecast, sorry."
+ mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply))
+ outcome = await self._router(mock_router_instance).aclassify(
+ "do the task", request_kwargs={"turn_off_message_logging": message_logging_off}
+ )
+ assert outcome.cause == "capability_classifier_fallback"
+ assert "capability classifier failed (ValidationError)" in caplog.text
+ assert "classifier verdict rejected (" in caplog.text
+ assert ("raw reply withheld" in caplog.text) is message_logging_off
+ assert (reply in caplog.text) is not message_logging_off
+
+ @pytest.mark.asyncio
+ async def test_call_failure_reason_names_the_exception_type(
+ self, mock_router_instance, caplog: pytest.LogCaptureFixture
+ ):
+ mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError())
+ outcome = await self._router(mock_router_instance).aclassify("do the task")
+ assert outcome.cause == "capability_classifier_fallback"
+ assert "capability classifier failed (TimeoutError)" in caplog.text
@pytest.mark.asyncio
async def test_decimal_rounding_does_not_break_inclusive_threshold(self, mock_router_instance):
@@ -3954,6 +4001,34 @@ class TestLLMClassifier:
assert call_kwargs["model"] == "haiku-classifier"
assert call_kwargs["timeout"] == 0.4
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("shape", _REPLY_SHAPES)
+ async def test_aclassify_llm_verdict_wrapped_in_fence_or_prose_still_decides_the_tier(
+ self, llm_complexity_router, mock_router_instance, shape: str
+ ):
+ reply = _wrapped_reply(shape, '{"tier": "COMPLEX"}')
+ mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply))
+ outcome = await llm_complexity_router.aclassify("hi")
+ assert outcome.tier == ComplexityTier.COMPLEX
+ assert outcome.cause == "llm_classifier"
+ assert "llm-classifier:COMPLEX" in outcome.signals
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("message_logging_off", (False, True))
+ async def test_aclassify_llm_unparseable_reply_is_logged_with_its_text_unless_message_logging_is_off(
+ self, llm_complexity_router, mock_router_instance, caplog: pytest.LogCaptureFixture, message_logging_off: bool
+ ):
+ reply = "I would call this COMPLEX, the {tier} field is implied."
+ mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply))
+ outcome = await llm_complexity_router.aclassify(
+ "hi", request_kwargs={"turn_off_message_logging": message_logging_off}
+ )
+ assert outcome.cause != "llm_classifier"
+ assert "LLM classifier failed (ValidationError)" in caplog.text
+ assert "classifier verdict rejected (" in caplog.text
+ assert ("raw reply withheld" in caplog.text) is message_logging_off
+ assert (reply in caplog.text) is not message_logging_off
+
@pytest.mark.asyncio
async def test_aclassify_llm_success_captures_classifier_cost(self, llm_complexity_router, mock_router_instance):
"""The classifier call is billed, so its cost must ride the outcome.
diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/test_litellm/router_strategy/test_llm_v2.py
index 6fb6df3265d..3fd2e8808e7 100644
--- a/tests/test_litellm/router_strategy/test_llm_v2.py
+++ b/tests/test_litellm/router_strategy/test_llm_v2.py
@@ -66,6 +66,25 @@ def _response(content: str) -> ModelResponse:
return response
+_REPLY_SHAPES: Final = ("fenced", "fenced-with-language", "prose-before", "prose-after", "fenced-then-prose")
+
+
+def _wrapped_reply(shape: str, verdict: str) -> str:
+ match shape:
+ case "fenced":
+ return f" ```\n{verdict}\n``` "
+ case "fenced-with-language":
+ return f"```json\n{verdict}\n```"
+ case "prose-before":
+ return f"Sure {{here}} is the verdict you asked for:\n\n{verdict}"
+ case "prose-after":
+ return f"{verdict}\n\nThe efficient solver should handle this {{well}}."
+ case "fenced-then-prose":
+ return f"```json\n{verdict}\n```\n\n## Reasoning\n\nThe task is coupled, so the forecasts differ."
+ case _:
+ raise AssertionError(shape)
+
+
def _router(content: str, config: ComplexityRouterConfig | None = None) -> tuple[ComplexityRouter, MagicMock]:
client: Final = MagicMock(spec=Router)
client.acompletion = AsyncMock(return_value=_response(content))
@@ -334,13 +353,12 @@ async def test_json_object_mode_supplies_schema_in_prompt() -> None:
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ("json_schema", "json_object"))
-@pytest.mark.parametrize("fence", ("```json", "```"))
-async def test_fenced_forecast_routes_by_validated_probabilities(mode: str, fence: str) -> None:
+@pytest.mark.parametrize("shape", _REPLY_SHAPES)
+async def test_wrapped_forecast_routes_by_validated_probabilities(mode: str, shape: str) -> None:
base: Final = _config().llm_v2_config
assert base is not None
config: Final = _config(llm_v2_config={**base.model_dump(), "response_format": mode})
- content: Final = f" {fence}\n{_verdict().model_dump_json()}\n``` "
- router, client = _router(content, config)
+ router, client = _router(_wrapped_reply(shape, _verdict().model_dump_json()), config)
result: Final = await router.async_pre_routing_hook(
model="v2-router", messages=[{"role": "user", "content": "Fix nested behavior"}], request_kwargs={}
)
@@ -454,6 +472,93 @@ async def test_provider_failure_redacts_prompt_text_from_warning(caplog: pytest.
assert "private task text" not in caplog.text
+@pytest.mark.asyncio
+@pytest.mark.parametrize("field", ("crux", "likely_failure"))
+async def test_long_verdict_explanations_still_route_by_validated_probabilities(field: str) -> None:
+ explanation: Final = "The solver must keep the nested retry behavior intact while it edits. " * 12
+ assert len(explanation) > 512
+ verdict: Final = _verdict().model_dump()
+ if field == "crux":
+ content: Final = json.dumps({**verdict, "crux": explanation})
+ else:
+ forecasts: Final = {**verdict["forecasts"], "efficient": {**verdict["forecasts"]["efficient"], field: explanation}}
+ content = json.dumps({**verdict, "forecasts": forecasts})
+ router, _ = _router(content)
+ outcome: Final = await router.aclassify("Fix nested behavior")
+ assert outcome.cause == "llm_v2_classifier"
+ assert outcome.llm_v2_forecast is not None
+ assert outcome.llm_v2_forecast.use_efficient
+
+
+@pytest.mark.parametrize("field", ("crux", "likely_failure"))
+def test_blank_verdict_explanations_are_still_rejected(field: str) -> None:
+ verdict: Final = _verdict().model_dump()
+ blank: Final = (
+ {**verdict, "crux": " "}
+ if field == "crux"
+ else {**verdict, "forecasts": {**verdict["forecasts"], "capable": {"likely_failure": " ", "p_solve": 0.5}}}
+ )
+ with pytest.raises(ValidationError):
+ LLMV2Verdict.model_validate(blank)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("message_logging_off", (False, True))
+async def test_unparseable_reply_is_logged_with_its_text_unless_message_logging_is_off(
+ caplog: pytest.LogCaptureFixture, message_logging_off: bool
+) -> None:
+ reply: Final = "I cannot forecast this one, the task text is too {vague} to score."
+ router, _ = _router(reply)
+ outcome: Final = await router.aclassify("hi", request_kwargs={"turn_off_message_logging": message_logging_off})
+ assert outcome.cause == "llm_v2_fallback"
+ assert "classifier verdict rejected (" in caplog.text
+ assert "Invalid LLM V2 forecast" in caplog.text
+ assert ("raw reply withheld" in caplog.text) is message_logging_off
+ assert (reply in caplog.text) is not message_logging_off
+
+
+_MESSAGE_LOGGING_OPT_OUTS: Final = (
+ pytest.param({"turn_off_message_logging": "True"}, False, id="key-logging-settings-string"),
+ pytest.param({"metadata": {"headers": {"x-litellm-enable-message-redaction": "true"}}}, False, id="redaction-header"),
+ pytest.param({}, True, id="global-setting"),
+ pytest.param({"metadata": {"headers": None}}, False, id="undecidable-headers-fail-closed"),
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("request_kwargs", "global_off"), _MESSAGE_LOGGING_OPT_OUTS)
+async def test_unparseable_reply_text_is_withheld_under_every_message_logging_opt_out(
+ monkeypatch: pytest.MonkeyPatch,
+ caplog: pytest.LogCaptureFixture,
+ request_kwargs: dict[str, object],
+ global_off: bool,
+) -> None:
+ monkeypatch.setattr(litellm, "turn_off_message_logging", global_off)
+ reply: Final = "I cannot forecast this one, the task text is too {vague} to score."
+ router, _ = _router(reply)
+ outcome: Final = await router.aclassify("hi", request_kwargs=request_kwargs)
+ assert outcome.cause == "llm_v2_fallback"
+ assert "raw reply withheld" in caplog.text
+ assert reply not in caplog.text
+
+
+_REPLIES_THE_JSON_SCANNER_CANNOT_DECODE: Final = (
+ pytest.param('{"a":' * 3000, id="deeply-nested"),
+ pytest.param('{"capability_p": ' + "9" * 5000 + "}", id="integer-over-the-digit-limit"),
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("reply", _REPLIES_THE_JSON_SCANNER_CANNOT_DECODE)
+async def test_undecodable_reply_is_rejected_as_an_invalid_forecast(
+ caplog: pytest.LogCaptureFixture, reply: str
+) -> None:
+ router, _ = _router(reply)
+ outcome: Final = await router.aclassify("hi", request_kwargs={})
+ assert outcome.cause == "llm_v2_fallback"
+ assert "Invalid LLM V2 forecast" in caplog.text
+
+
def test_response_schema_requires_both_model_forecasts() -> None:
with pytest.raises(ValidationError):
LLMV2Verdict.model_validate(
From 474ab91c0972f5074eac297dc63606f607ce5dda Mon Sep 17 00:00:00 2001
From: tin-berri
Date: Fri, 25 Sep 2026 16:35:15 -0700
Subject: [PATCH 3/6] test(zerobus): move tests into active CI selection
(#43235)
* test(zerobus): move tests into the active CI selection
* test(zerobus): add package marker for unit test discovery
---
tests/unit/integrations/zerobus/__init__.py | 0
.../integrations/zerobus/test_zerobus_client.py | 0
.../integrations/zerobus/test_zerobus_logger.py | 0
.../integrations/zerobus/test_zerobus_row.py | 0
4 files changed, 0 insertions(+), 0 deletions(-)
create mode 100644 tests/unit/integrations/zerobus/__init__.py
rename tests/{test_litellm => unit}/integrations/zerobus/test_zerobus_client.py (100%)
rename tests/{test_litellm => unit}/integrations/zerobus/test_zerobus_logger.py (100%)
rename tests/{test_litellm => unit}/integrations/zerobus/test_zerobus_row.py (100%)
diff --git a/tests/unit/integrations/zerobus/__init__.py b/tests/unit/integrations/zerobus/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_client.py b/tests/unit/integrations/zerobus/test_zerobus_client.py
similarity index 100%
rename from tests/test_litellm/integrations/zerobus/test_zerobus_client.py
rename to tests/unit/integrations/zerobus/test_zerobus_client.py
diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py b/tests/unit/integrations/zerobus/test_zerobus_logger.py
similarity index 100%
rename from tests/test_litellm/integrations/zerobus/test_zerobus_logger.py
rename to tests/unit/integrations/zerobus/test_zerobus_logger.py
diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_row.py b/tests/unit/integrations/zerobus/test_zerobus_row.py
similarity index 100%
rename from tests/test_litellm/integrations/zerobus/test_zerobus_row.py
rename to tests/unit/integrations/zerobus/test_zerobus_row.py
From e0fb89bc82195a0d7ad4d7799286ffea61d2b377 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
<158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Fri, 25 Sep 2026 17:05:53 -0700
Subject: [PATCH 4/6] fix(proxy): keep the submitted body out of 422 validation
errors (#43231)
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
.../common_utils/validation_error_body.py | 13 ++++
litellm/proxy/list_api/common.py | 10 +--
litellm/proxy/proxy_server.py | 16 ++---
.../proxy_setting_endpoints.py | 3 +-
.../proxy_server/test_exception_handlers.py | 66 ++++++++++++++++++-
.../proxy_server/test_routes_onboarding.py | 19 ++++++
.../test_proxy_setting_endpoints.py | 16 +++++
.../test_validation_error_body.py | 46 +++++++++++++
8 files changed, 168 insertions(+), 21 deletions(-)
create mode 100644 litellm/proxy/common_utils/validation_error_body.py
create mode 100644 tests/unit/proxy/common_utils/test_validation_error_body.py
diff --git a/litellm/proxy/common_utils/validation_error_body.py b/litellm/proxy/common_utils/validation_error_body.py
new file mode 100644
index 00000000000..b21f33a2434
--- /dev/null
+++ b/litellm/proxy/common_utils/validation_error_body.py
@@ -0,0 +1,13 @@
+from collections.abc import Sequence
+
+from typing_extensions import ReadOnly, TypedDict
+
+
+class ValidationErrorDetail(TypedDict):
+ type: ReadOnly[str]
+ loc: ReadOnly[tuple[int | str, ...]]
+ msg: ReadOnly[str]
+
+
+def public_validation_errors(errors: Sequence[ValidationErrorDetail]) -> tuple[ValidationErrorDetail, ...]:
+ return tuple(ValidationErrorDetail(type=error["type"], loc=error["loc"], msg=error["msg"]) for error in errors)
diff --git a/litellm/proxy/list_api/common.py b/litellm/proxy/list_api/common.py
index daa6414fd94..efa8a271459 100644
--- a/litellm/proxy/list_api/common.py
+++ b/litellm/proxy/list_api/common.py
@@ -8,8 +8,8 @@ from fastapi import Request
from fastapi.dependencies.utils import get_flat_params
from fastapi.params import ParamTypes
from fastapi.responses import JSONResponse
-from typing_extensions import ReadOnly, TypedDict
+from litellm.proxy.common_utils.validation_error_body import ValidationErrorDetail
from litellm.types.proxy.management_endpoints.management_v1 import (
ListLinks,
PageLinks,
@@ -58,14 +58,6 @@ def escape_like(value: str) -> str:
return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
-class ValidationErrorDetail(TypedDict):
- """The keys of a pydantic/FastAPI validation error a problem document needs."""
-
- type: ReadOnly[str]
- loc: ReadOnly[tuple[int | str, ...]]
- msg: ReadOnly[str]
-
-
def _is_length_error_of_rejected_items(error: ValidationErrorDetail, errors: Sequence[ValidationErrorDetail]) -> bool:
"""pydantic counts only items that validated, so a bad item also trips the parent's min_length."""
return error["type"] == "too_short" and any(
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 61304f0d919..646eca071d1 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -476,6 +476,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
project_spend_counter_key,
tag_cache_key,
)
+from litellm.proxy.common_utils.validation_error_body import public_validation_errors
from litellm.proxy.config_resolvers import (
FieldSource,
SettingsStore,
@@ -551,7 +552,6 @@ from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger, run_sp
from litellm.proxy.image_endpoints.endpoints import router as image_router
from litellm.proxy.list_api.common import (
ManagementProblem,
- ValidationErrorDetail,
problem_response,
request_validation_problem,
)
@@ -1983,16 +1983,14 @@ class _ExceptionRow(TypedDict, total=False):
@app.exception_handler(RequestValidationError)
async def otel_request_validation_exception_handler(request: Request, exc: RequestValidationError):
+ public_errors: Final = public_validation_errors(exc.errors())
+ public_exc: Final = RequestValidationError(public_errors).with_traceback(exc.__traceback__)
if request.url.path.startswith(MANAGEMENT_V1_PREFIX):
- validation_errors: Final[Sequence[ValidationErrorDetail]] = exc.errors()
- problem: Final = request_validation_problem(validation_errors)
- _close_dangling_otel_server_span(request, problem.status, exc=exc)
+ problem: Final = request_validation_problem(public_errors)
+ _close_dangling_otel_server_span(request, problem.status, exc=public_exc)
return problem_response(problem)
- _close_dangling_otel_server_span(request, 422, exc=exc)
- return JSONResponse(
- status_code=422,
- content={"detail": jsonable_encoder(exc.errors())},
- )
+ _close_dangling_otel_server_span(request, 422, exc=public_exc)
+ return JSONResponse(status_code=422, content={"detail": public_errors})
@app.exception_handler(Exception)
diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
index c91b1afd64a..227f0e7f795 100644
--- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
+++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
@@ -25,6 +25,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys
from litellm.proxy._experimental.mcp_server.tool_search import MCP_TOOL_SEARCH_SETTINGS_KEY
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_utils.validation_error_body import public_validation_errors
from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for
from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError
from litellm.proxy.config_resolvers.sso import (
@@ -1875,7 +1876,7 @@ async def update_ui_settings(
try:
settings: Final = effective_cls.model_validate(settings_body)
except ValidationError as e:
- raise HTTPException(status_code=422, detail=e.errors())
+ raise HTTPException(status_code=422, detail=public_validation_errors(e.errors()))
unsupported_team_fields: Final = sorted(
frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_PERMISSIONS
diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py
index 089c2d57594..16cb1146ff5 100644
--- a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py
+++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py
@@ -265,7 +265,7 @@ def test_close_dangling_otel_server_span_logger_raises_state_cleared_error(monke
@pytest.mark.asyncio
async def test_otel_request_validation_exception_handler_returns_422_detail():
- errors = [{"loc": ["body", "model"], "msg": "field required", "type": "missing"}]
+ errors = [{"loc": ["body", "model"], "msg": "field required", "type": "missing", "input": {"messages": []}}]
exc = RequestValidationError(errors)
request = _make_request()
@@ -273,7 +273,69 @@ async def test_otel_request_validation_exception_handler_returns_422_detail():
body = json.loads(response.body)
assert response.status_code == 422
- assert normalize(body) == {"detail": exc.errors()}
+ assert body == {"detail": [{"type": "missing", "loc": ["body", "model"], "msg": "field required"}]}
+
+
+_SUBMITTED_PASSWORD: Final = "hunter2-Sup3rSecret!"
+_PASSWORD_LEAKING_ERRORS: Final = (
+ {
+ "type": "missing",
+ "loc": ["body", "new_password"],
+ "msg": "Field required",
+ "input": {"current_password": _SUBMITTED_PASSWORD},
+ },
+ {
+ "type": "value_error",
+ "loc": ["body", "password"],
+ "msg": "Value error, password cannot be set via /user/new",
+ "input": _SUBMITTED_PASSWORD,
+ "ctx": {"error": ValueError(_SUBMITTED_PASSWORD)},
+ },
+)
+_PUBLIC_ERRORS: Final = (
+ {"type": "missing", "loc": ["body", "new_password"], "msg": "Field required"},
+ {"type": "value_error", "loc": ["body", "password"], "msg": "Value error, password cannot be set via /user/new"},
+)
+
+
+@pytest.mark.asyncio
+async def test_otel_request_validation_exception_handler_never_echoes_the_submitted_body():
+ """A pydantic error carries the offending value as ``input`` (the whole body for a
+ ``missing`` error) and input-derived values in ``ctx``; a caller who mistyped a
+ request holding a password must not get that password back."""
+ exc = RequestValidationError(list(_PASSWORD_LEAKING_ERRORS))
+
+ response = await otel_request_validation_exception_handler(request=_make_request(), exc=exc)
+
+ assert response.status_code == 422
+ assert json.loads(response.body) == {"detail": list(_PUBLIC_ERRORS)}
+ assert _SUBMITTED_PASSWORD.encode() not in response.body
+
+
+@pytest.mark.asyncio
+async def test_otel_request_validation_exception_handler_hands_the_span_only_the_public_errors(monkeypatch):
+ """The OTEL SERVER span's error message is ``str(exc)``, which FastAPI builds from
+ every error dict ``input`` included, so the span gets the same public-only errors
+ the caller does, and keeps the traceback the original carried."""
+ import litellm.proxy.proxy_server as ps
+
+ fake_logger = MagicMock()
+ monkeypatch.setattr(ps, "open_telemetry_logger", fake_logger, raising=False)
+ exc = RequestValidationError(list(_PASSWORD_LEAKING_ERRORS))
+ try:
+ raise exc
+ except RequestValidationError as raised:
+ original_traceback = raised.__traceback__
+ request = _make_request(parent_otel_span=MagicMock())
+
+ await otel_request_validation_exception_handler(request=request, exc=exc)
+
+ (_span, span_exc, status_code) = fake_logger.record_error_attributes_on_span.call_args.args
+ assert status_code == 422
+ assert isinstance(span_exc, RequestValidationError)
+ assert list(span_exc.errors()) == list(_PUBLIC_ERRORS)
+ assert _SUBMITTED_PASSWORD not in str(span_exc)
+ assert span_exc.__traceback__ is original_traceback
@pytest.mark.asyncio
diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py
index 778acc1baab..6c1d869d113 100644
--- a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py
+++ b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py
@@ -317,6 +317,25 @@ def test_claim_onboarding_link_missing_field_422(client, monkeypatch, mock_prism
assert any("password" in str(item) for item in body["detail"])
+def test_claim_onboarding_link_422_never_echoes_the_submitted_password(client):
+ """A body that fails validation is answered with the field path and message only;
+ pydantic's ``input`` (the whole submitted body for a missing field, password
+ included) must never come back to the caller or land in whatever logs the response."""
+ password = "hunter2-Sup3rSecret!"
+
+ response = client.post(
+ "/onboarding/claim_token",
+ json={"invitation_link": "abc", "password": password},
+ )
+
+ assert response.status_code == 422
+ assert password.encode() not in response.content
+ detail = response.json()["detail"]
+ assert detail[0]["loc"] == ["body", "user_id"]
+ assert detail[0]["msg"]
+ assert set(detail[0]) == {"type", "loc", "msg"}
+
+
def test_claim_onboarding_link_bad_onboarding_jwt_401(
client, monkeypatch, mock_prisma
):
diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
index 08d542df16c..0d7a713a380 100644
--- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
+++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
@@ -3817,6 +3817,22 @@ class TestTeamAdminEditableTeamFieldsSetting:
assert response.status_code == 422
+ def test_patch_422_never_echoes_the_submitted_value(self, monkeypatch):
+ self._as_proxy_admin(monkeypatch)
+ submitted = "hunter2-Sup3rSecret!"
+
+ try:
+ response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": submitted})
+ finally:
+ app.dependency_overrides.clear()
+
+ assert response.status_code == 422
+ assert submitted.encode() not in response.content
+ detail = response.json()["detail"]
+ assert detail[0]["loc"] == ["team_admin_editable_team_fields"]
+ assert detail[0]["msg"]
+ assert set(detail[0]) == {"type", "loc", "msg"}
+
def test_patch_persists_and_syncs_the_list_to_general_settings(self, monkeypatch):
mock_prisma = self._as_proxy_admin(monkeypatch)
general_settings: dict = {"team_admin_editable_team_fields": []}
diff --git a/tests/unit/proxy/common_utils/test_validation_error_body.py b/tests/unit/proxy/common_utils/test_validation_error_body.py
new file mode 100644
index 00000000000..a86f17b7461
--- /dev/null
+++ b/tests/unit/proxy/common_utils/test_validation_error_body.py
@@ -0,0 +1,46 @@
+from typing import Final
+
+from litellm.proxy.common_utils.validation_error_body import public_validation_errors
+
+_PASSWORD: Final = "hunter2-Sup3rSecret!"
+
+
+def test_public_validation_errors_drops_input_ctx_and_url():
+ errors: Final = (
+ {
+ "type": "missing",
+ "loc": ("body", "user_id"),
+ "msg": "Field required",
+ "input": {"invitation_link": "abc", "password": _PASSWORD},
+ "url": "https://errors.pydantic.dev/2/v/missing",
+ },
+ {
+ "type": "value_error",
+ "loc": ("body", "password"),
+ "msg": "Value error, password cannot be set here",
+ "input": _PASSWORD,
+ "ctx": {"error": ValueError(_PASSWORD)},
+ },
+ )
+
+ public: Final = public_validation_errors(errors)
+
+ assert public == (
+ {"type": "missing", "loc": ("body", "user_id"), "msg": "Field required"},
+ {"type": "value_error", "loc": ("body", "password"), "msg": "Value error, password cannot be set here"},
+ )
+ assert _PASSWORD not in repr(public)
+
+
+def test_public_validation_errors_keeps_type_loc_and_msg_verbatim_in_order():
+ errors: Final = (
+ {"type": "int_parsing", "loc": ("body", "litellm_params", "rpm"), "msg": "Input should be a valid integer"},
+ {"type": "extra_forbidden", "loc": ("body", "users", 0, "user_emial"), "msg": "Extra inputs are not permitted"},
+ {"type": "too_short", "loc": ("body", "users"), "msg": "List should have at least 1 item"},
+ )
+
+ assert public_validation_errors(errors) == errors
+
+
+def test_public_validation_errors_empty_in_empty_out():
+ assert public_validation_errors(()) == ()
From a11a93f44a557dd91100c2e394006b5df0daed65 Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Fri, 25 Sep 2026 17:10:13 -0700
Subject: [PATCH 5/6] test: move tests/test_litellm core utils, routing,
responses, caching and rust_bridge into tests/unit (#43199)
* ci: run the unit_selection.sh shard files on every event instead of only fork pull requests
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* ci: rename fork-flag to unit-flag now that it applies on every event
* test: move tests/test_litellm root and small trees into tests/unit
Pure renames, no content changes. Follow-up commits in this PR fix
references, merge the three files that already existed in tests/unit,
keep live-provider tests in tests/test_litellm and wire CI.
* test: carry tests/test_litellm conftest isolation into tests/unit
Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS,
proxy-URL and keychain env, and session-end client cleanup now reset for
unit tests too. The environment isolation owns its MonkeyPatch so a test's
own monkeypatch is undone before the model-cost teardown runs.
* test: merge, split and prune the moved root and small-tree tests
Merge batches/test_batch_utils.py and the chat_completions and messages
dispatch tests into the files that already existed in tests/unit. Keep
the live Gemini interactions tests, the async image-fetch format test and
the OpenAI embedding scorer test in tests/test_litellm since they need
real network or keys. Put test_router.py under tests/unit/test_router so
the existing package no longer shadows it. Delete eight tests the audit
found superseded by stronger ones kept in this move.
* ci: run the moved root and small-tree tests under their legacy flags
Add the misc and responses-caching-types flags to unit_selection.sh and
CircleCI, extend enterprise-routing and mcp-integration, and point the
legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest
and change classifier at the new paths.
* test: make the new tests/unit directories packages
tests/unit/test_package_layout.py requires every directory to carry an
__init__.py, and without one the moved and retained
test_litellm_responses_bridge.py modules collide on import.
* test: scope the unit socket block to tests/unit in shared sessions
The GHA shards collect the legacy test-path and the unit selection in one
pytest session. The unit conftest's loopback-only block leaked into legacy
modules that reach the network at import. The legacy conftest now lifts the
restriction at collect and setup time, and the unit conftest re-applies it
when collecting its own modules.
* test: move tests/test_litellm/llms into tests/unit/llms
Rename-only. Moves the provider tests and the fine-tuning fixtures they
load, mirroring the old paths. Follow-up commits merge, split and wire them.
* test: merge, split and prune the moved llms tests
Merges the Databricks chat transformation tests into the existing unit
file, keeps the tests that need real keys or the network in
tests/test_litellm, deletes the audited tests a stronger unit test
already covers, and points imports at tests.unit.llms.
* ci: run the moved llms tests under their legacy flags
The Vertex AI and All Other Providers shards keep their legacy test-path
for the retained files and add the llm-vertex-ai and llm-other-providers
unit selections. CircleCI gets matching unit jobs.
* test: make the tests/unit/llms directories packages
Adds __init__.py to the moved dirs and drops the legacy ones whose
directories no longer hold tests.
* test: drop script runners and path hacks the llms split left dangling
The __main__ runners in the split openai_like files and the Databricks e2e
runner called tests that now live in the other half of the split or were
deleted. The retained legacy halves also no longer need sys.path edits.
* test: give the shard-script tests their own GITHUB_OUTPUT
They only passed where the runner set it. The CircleCI unit job's env
allowlist drops it, so the script's redirect failed there.
* test: point the router and module-deletion checks at tests/unit
router_code_coverage and code_qa_check_tests only searched tests/test_litellm,
so the moved router tests no longer counted. The two silent-experiment tests
the audit deleted were the only direct callers of those methods; they are
replaced with tests that assert the forwarded shadow request and the
recursion guard.
* test: move tests/test_litellm integrations and secret_managers into tests/unit
Rename-only. Mirrors the old paths, including the directory conftests
and the prompt and JSON fixtures. Follow-up commits prune and wire them.
* test: prune and repoint the moved integrations tests
Deletes the 7 audited tests a stronger test in the same tree already
covers, imports the TLS sink helpers from their new conftest path, and
restores os.environ after each integrations test. Some presets write
OTEL_EXPORTER_OTLP_HEADERS straight into os.environ, and without the
legacy tree's test ordering that header leaked into the AgentOps tests.
* ci: run the moved integrations tests under their legacy flag
The integrations GHA shard and a new CircleCI job run the integrations
unit selection. secret_managers joins the misc selection.
* docs: point integrations and secret_managers references at tests/unit
* test: make the moved integrations directories packages
* test: keep the Databricks manual e2e runner and fix the SageMaker Nova run path
The Databricks e2e file is a manual script whose main() calls the tests
that were pruned, so pruning them broke the documented run. It is back to
its main version. The SageMaker Nova docstring now points at the file's
real location in tests/local_testing.
* test: move tests/test_litellm core utils, routing, responses, caching and rust_bridge into tests/unit
Rename-only. Mirrors the old paths, including fixtures, the stubtest config
and the native-route wheel script. Two files that collide with existing unit
files are merged in a follow-up commit.
* test: merge, prune and repoint the moved core, routing, responses, caching and rust_bridge tests
Merges the two files that collided with existing unit files, folding the
legacy extra case into test_is_chat_completion_cached_dict, and deletes the
9 audited tests a stronger test in the same file already covers.
Keeps what needs the network in tests/test_litellm: test_tokenizers pulls a
tokenizer from the Hugging Face hub, and the gpt2 and r50k_base tokenizer
cases download their BPE files. The unit core_utils conftest points
TIKTOKEN_CACHE_DIR at litellm's bundled encodings so the rest never depend on
import order to stay offline, and FakeSecretVault moves to a shared module
so both trees can build it.
* ci: run the moved core, routing, responses, caching and rust_bridge tests under their flags
core_utils gets a core-utils flag and CircleCI job, and its GHA shard keeps
the legacy path for the retained network tests. router_utils and
router_strategy join enterprise-routing, responses joins
responses-caching-types (minus responses/mcp, which mcp-integration owns),
caching joins caching-local and rust_bridge joins misc. The redis-compat,
test-rust, stubtest and merge-smoke paths follow the move.
* docs: point the Rust crate references at tests/unit
* test: make the moved core, routing and rust_bridge directories packages
* test: keep the no-loop DualCache batch_get_cache regression test
It runs the sync path outside any event loop, which the inside-loop test
cannot, so a change that picks the Redis client by loop state would only
show up there.
* test: keep the job's UNIT_FLAG out of the shard-script tests
* fix(url_utils): block 192.0.0.0/24 on every Python patch release
* test: move the new budget limiter tests into tests/unit/router_strategy
* test: move the new sentry scrubbing tests into tests/unit/litellm_core_utils
* test: move the new zerobus tests into tests/unit/integrations
* test: make tests/unit/integrations/zerobus a package
* test: load litellm's own tiktoken cache setup once instead of resetting it per test
---------
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.circleci/scripts/unit_selection.sh | 9 +-
.circleci/tests.yml | 7 +
.github/merge-smoke-tests.json | 8 +-
.github/workflows/test-redis-compat.yml | 12 +-
.github/workflows/test-rust.yml | 6 +-
.github/workflows/test-unit.yml | 10 +-
Makefile | 6 +-
.../crates/callbacks-legacy-python/src/lib.rs | 2 +-
litellm-rust/crates/secrets/README.md | 2 +-
litellm/litellm_core_utils/url_utils.py | 10 +-
.../base_responses_api.py | 2 +-
tests/llm_translation/test_gemini.py | 2 +-
.../caching/test_caching_handler.py | 867 ---------
tests/test_litellm/conftest.py | 68 +-
.../litellm_core_utils/__init__.py | 1 -
.../litellm_core_utils/test_token_counter.py | 1572 +----------------
.../litellm_core_utils/test_tokenizer.py | 409 +----
tests/test_litellm/proxy/client/test_chat.py | 2 +-
.../proxy/hooks/test_tpm_concurrent.py | 2 +-
.../test_streaming_handler.py | 4 +-
tests/test_litellm/proxy/test_proxy_server.py | 4 +-
tests/test_litellm/proxy/test_proxy_utils.py | 2 +-
.../rust_bridge/messages/test_route_host.py | 124 --
.../rust_bridge/responses/__init__.py | 0
.../tokenizer/test_fast_count.py | 2 +-
.../test_a2a_streaming_iterator.py | 2 +-
tests/unit/a2a_protocol/test_main.py | 2 +-
.../caching/test_azure_blob_cache.py | 0
.../caching/test_caching.py | 0
tests/unit/caching/test_caching_handler.py | 804 +++++++++
...test_check_and_fix_namespace_none_guard.py | 0
.../caching/test_disk_cache.py | 0
.../caching/test_dual_cache.py | 0
.../caching/test_embedding_router.py | 0
.../caching/test_evicted_client_closer.py | 0
.../caching/test_gcs_cache.py | 0
.../caching/test_in_memory_cache.py | 0
.../caching/test_llm_caching_handler.py | 0
.../caching/test_llm_client_cache_e2e.py | 0
.../caching/test_qdrant_semantic_cache.py | 2 +-
.../caching/test_redis_cache.py | 0
.../caching/test_redis_cluster_cache.py | 0
.../test_redis_cluster_node_isolation.py | 0
.../caching/test_redis_connection_pool.py | 0
.../caching/test_redis_semantic_cache.py | 2 +-
.../caching/test_s3_cache.py | 0
.../caching/test_valkey_semantic_cache.py | 0
.../__init__.py | 0
.../azure_shell_tool.json | 0
.../context_management_and_shell.json | 0
.../test_compression_interception_handler.py | 2 +-
tests/unit/litellm_core_utils/conftest.py | 15 +
.../litellm_core_utils/event_loop_lag.py | 0
.../litellm_core_utils/fake_secret_vault.py | 67 +
.../llm_cost_calc}/__init__.py | 0
.../test_azure_assistant_cost_tracking.py | 0
.../llm_cost_calc/test_guardrail_cost.py | 0
.../llm_cost_calc/test_llm_cost_calc_utils.py | 0
.../test_openai_cache_write_cost.py | 0
.../test_responses_cache_cost_breakdown.py | 0
.../test_tool_call_cost_tracking.py | 0
...est_tool_call_cost_tracking_dict_safety.py | 0
.../test_usage_object_transformation.py | 0
.../test_zero_cost_diagnostic.py | 0
.../llm_response_utils/test_get_api_base.py | 0
.../messages_with_counts.py | 0
.../prompt_templates}/__init__.py | 0
...edrock_converse_strict_tools_opus_47_48.py | 0
...ore_utils_prompt_templates_common_utils.py | 0
...llm_core_utils_prompt_templates_factory.py | 0
...rompt_templates_mid_conversation_system.py | 0
.../specialty_caches}/__init__.py | 0
.../test_dynamic_logging_cache.py | 0
.../test_agentic_followup_kwargs.py | 0
.../test_anthropic_dedup_factory.py | 0
.../test_api_route_to_call_types.py | 0
.../litellm_core_utils/test_audio_utils.py | 0
.../litellm_core_utils/test_aws_partition.py | 0
.../test_bedrock_converse_dedup_factory.py | 0
.../litellm_core_utils/test_bug_report.py | 0
.../test_chat_completion_agentic_loop.py | 0
.../test_classifier_logging.py | 0
.../test_cli_token_utils.py | 0
.../test_cloud_storage_security.py | 0
.../test_codestral_provider_routing.py | 0
.../litellm_core_utils/test_core_helpers.py | 0
.../test_coroutine_checker.py | 0
.../litellm_core_utils/test_dd_tracing.py | 12 -
.../test_decode_special_tokens.py | 0
.../test_dot_notation_indexing.py | 0
.../test_duration_parser.py | 0
.../test_error_normalization.py | 0
.../test_exception_mapping_utils.py | 0
.../test_extract_base64_image.py | 0
.../test_fallback_generalizations.py | 0
.../litellm_core_utils/test_fallback_utils.py | 0
.../test_get_litellm_params.py | 0
.../test_get_llm_provider_endpoint_match.py | 0
.../test_get_llm_provider_logic.py | 0
.../test_get_model_cost_map.py | 0
.../test_get_supported_openai_params.py | 0
.../test_health_check_helpers.py | 0
.../litellm_core_utils/test_image_handling.py | 0
...test_initialize_dynamic_callback_params.py | 0
.../test_internal_call_metadata.py | 0
.../test_json_fragment_accumulator.py | 0
.../test_json_schema_validation.py | 0
.../test_litellm_logging.py | 0
.../litellm_core_utils/test_llm_judge.py | 0
.../test_llm_request_utils.py | 0
.../litellm_core_utils/test_logging_utils.py | 0
.../litellm_core_utils/test_logging_worker.py | 0
.../test_max_streaming_duration.py | 0
.../test_model_param_helper.py | 0
.../test_model_response_utils.py | 0
.../litellm_core_utils/test_private_json.py | 0
.../test_provider_affinity.py | 0
.../test_provider_specific_headers.py | 0
.../litellm_core_utils/test_ptu_pricing.py | 0
.../test_realtime_errors.py | 0
.../test_realtime_streaming.py | 0
.../test_redact_messages.py | 0
.../test_request_timeout_resolver.py | 0
.../test_retry_after_headers.py | 0
.../test_safe_divide_seconds.py | 0
.../test_safe_json_dumps.py | 0
.../test_sensitive_data_masker.py | 0
.../test_sentry_scrubbing.py | 0
.../test_served_output_texts.py | 0
.../test_streaming_chunk_builder_cursor.py | 0
...streaming_chunk_builder_server_tool_use.py | 0
.../test_streaming_chunk_builder_utils.py | 0
.../test_streaming_handler.py | 2 +-
.../test_streaming_overhead.py | 0
.../test_thread_pool_executor.py | 0
.../litellm_core_utils/test_token_counter.py | 1441 +++++++++++++++
.../test_token_counter_tool.py | 4 +-
.../test_token_counter_tool_data.py | 0
.../unit/litellm_core_utils/test_tokenizer.py | 411 +++++
.../test_tool_search_spend_logging.py | 0
.../litellm_core_utils/test_url_utils.py | 0
.../test_xai_oauth_routing.py | 0
.../context_management/test_compact.py | 2 +-
.../context_management/test_dispatcher.py | 2 +-
.../llms/test_polling_url_origin_match.py | 2 +-
.../__init__.py | 0
...test_function_call_output_normalization.py | 0
.../test_handler.py | 0
.../test_image_generation_output.py | 0
.../test_litellm_completion_responses.py | 0
.../test_reasoning_input_item_preservation.py | 0
.../test_session_handler.py | 0
.../test_session_handler_with_cold_storage.py | 0
.../test_streaming_iterator_transformation.py | 0
..._tool_output_order_preserved_for_gemini.py | 0
.../mcp/test_chat_completions_handler.py | 0
.../mcp/test_litellm_proxy_mcp_handler.py | 0
.../mcp/test_mcp_streaming_iterator.py | 0
.../responses/test_additional_tools.py | 0
.../responses/test_custom_tool_call.py | 0
.../responses/test_dispatch.py | 0
.../responses/test_metadata_codex_callback.py | 0
.../responses/test_no_duplicate_spend_logs.py | 29 -
.../responses/test_null_test_fix.py | 0
.../test_responses_api_bridge_flag.py | 0
.../test_responses_api_request_body.py | 2 +-
.../test_responses_prompt_management.py | 0
.../test_responses_router_cooldown.py | 0
.../test_responses_streaming_iterator.py | 0
...sponses_supported_endpoints_passthrough.py | 0
.../responses/test_responses_utils.py | 0
.../test_responses_websocket_all_providers.py | 91 -
.../responses/test_rust_bridge_websocket.py | 0
.../responses/test_sse_output_recovery.py | 0
.../responses/test_streaming_iterator.py | 0
.../test_streaming_iterator_error_events.py | 0
.../responses/test_text_format_conversion.py | 0
.../adaptive_router}/__init__.py | 0
.../adaptive_router/fixtures}/__init__.py | 0
.../fixtures/clean_no_signals.json | 0
.../fixtures/clean_satisfaction.json | 0
.../fixtures/disengagement_giveup.json | 0
.../fixtures/exhaustion_429.json | 0
.../fixtures/exhaustion_context_overflow.json | 0
.../fixtures/failure_tool_error.json | 0
.../fixtures/loop_same_tool.json | 0
.../fixtures/misalignment_rephrase.json | 0
.../mixed_failure_then_satisfaction.json | 0
.../fixtures/stagnation_repeat.json | 0
.../adaptive_router/test_adaptive_router.py | 0
.../adaptive_router/test_async_pre_routing.py | 0
.../adaptive_router/test_bandit.py | 0
.../adaptive_router/test_classifier.py | 0
.../adaptive_router/test_config.py | 0
.../test_e2e_adaptive_router.py | 0
.../adaptive_router/test_hooks.py | 0
.../adaptive_router/test_router_dispatch.py | 0
.../adaptive_router/test_signals.py | 0
.../adaptive_router/test_state_endpoint.py | 0
.../adaptive_router/test_update_queue.py | 0
.../test_context_compaction.py | 0
.../router_strategy/test_auto_router.py | 0
.../test_base_routing_strategy.py | 0
.../router_strategy/test_budget_limiter.py | 0
.../test_budget_limiter_hotpath.py | 0
.../router_strategy/test_complexity_router.py | 0
.../test_complexity_tier_predictor.py | 0
.../router_strategy/test_fuse_presets.py | 0
.../router_strategy/test_lar1_routing.py | 0
.../router_strategy/test_least_busy.py | 0
.../router_strategy/test_litellm_encoder.py | 0
.../router_strategy/test_llm_v2.py | 0
.../router_strategy/test_lowest_cost.py | 0
.../router_strategy/test_lowest_latency.py | 0
.../router_strategy/test_lowest_tpm_rpm.py | 0
.../router_strategy/test_quality_router.py | 0
.../test_router_routing_groups.py | 0
.../test_router_routing_plugins.py | 0
.../test_router_tag_regex_routing.py | 0
.../test_router_tag_routing.py | 0
.../router_strategy/test_savings_baseline.py | 0
.../router_strategy/test_simple_shuffle.py | 0
.../router_strategy/test_stall_detector.py | 0
.../test_prompt_caching_deployment_check.py | 4 +-
.../router_utils/test_access_windows.py | 0
.../test_add_retry_fallback_headers.py | 0
.../test_auto_router_model_naming.py | 0
.../test_auto_router_tuning_baseline.py | 0
.../test_client_initalization_utils.py | 0
.../router_utils/test_cooldown_cache.py | 0
.../router_utils/test_cooldown_handlers.py | 0
.../test_fallback_event_handlers.py | 0
.../test_get_retry_from_policy.py | 0
..._health_check_allowed_fails_integration.py | 0
.../router_utils/test_health_state_cache.py | 0
.../test_pattern_match_deployments.py | 0
.../test_reasoning_effort_capability.py | 0
.../test_router_health_check_routing.py | 0
.../test_router_interactions_endpoints.py | 0
.../test_router_utils_common_utils.py | 0
.../rust_bridge/AGENTS.md | 0
.../rust_bridge/messages/test_route_host.py | 122 ++
.../rust_bridge/messages/test_secrets.py | 0
.../rust_bridge/native_route_wheel_test.py | 0
.../rust_bridge/ocr/test_secrets.py | 0
.../rust_bridge/stubtest.ini | 0
.../rust_bridge/test_bindings.py | 0
.../test_callbacks_legacy_python.py | 0
.../rust_bridge/test_catalog.py | 0
.../rust_bridge/test_configuration.py | 0
.../rust_bridge/test_dispatch.py | 0
.../rust_bridge/test_failures.py | 0
.../rust_bridge/test_fork_guard.py | 0
.../rust_bridge/test_lifecycle.py | 0
.../rust_bridge/test_logger.py | 0
.../rust_bridge/test_runtime.py | 0
.../rust_bridge/test_secret_manager.py | 0
.../rust_bridge/test_settings.py | 0
.../rust_bridge/test_token_counter.py | 0
.../rust_bridge/test_tokenizer.py | 2 +-
.../test_verify_linux_native_wheel.py | 0
261 files changed, 2951 insertions(+), 3204 deletions(-)
delete mode 100644 tests/test_litellm/caching/test_caching_handler.py
delete mode 100644 tests/test_litellm/rust_bridge/messages/test_route_host.py
delete mode 100644 tests/test_litellm/rust_bridge/responses/__init__.py
rename tests/{test_litellm => unit}/caching/test_azure_blob_cache.py (100%)
rename tests/{test_litellm => unit}/caching/test_caching.py (100%)
rename tests/{test_litellm => unit}/caching/test_check_and_fix_namespace_none_guard.py (100%)
rename tests/{test_litellm => unit}/caching/test_disk_cache.py (100%)
rename tests/{test_litellm => unit}/caching/test_dual_cache.py (100%)
rename tests/{test_litellm => unit}/caching/test_embedding_router.py (100%)
rename tests/{test_litellm => unit}/caching/test_evicted_client_closer.py (100%)
rename tests/{test_litellm => unit}/caching/test_gcs_cache.py (100%)
rename tests/{test_litellm => unit}/caching/test_in_memory_cache.py (100%)
rename tests/{test_litellm => unit}/caching/test_llm_caching_handler.py (100%)
rename tests/{test_litellm => unit}/caching/test_llm_client_cache_e2e.py (100%)
rename tests/{test_litellm => unit}/caching/test_qdrant_semantic_cache.py (99%)
rename tests/{test_litellm => unit}/caching/test_redis_cache.py (100%)
rename tests/{test_litellm => unit}/caching/test_redis_cluster_cache.py (100%)
rename tests/{test_litellm => unit}/caching/test_redis_cluster_node_isolation.py (100%)
rename tests/{test_litellm => unit}/caching/test_redis_connection_pool.py (100%)
rename tests/{test_litellm => unit}/caching/test_redis_semantic_cache.py (99%)
rename tests/{test_litellm => unit}/caching/test_s3_cache.py (100%)
rename tests/{test_litellm => unit}/caching/test_valkey_semantic_cache.py (100%)
rename tests/{test_litellm/litellm_core_utils/audio_utils => unit/expected_responses_api_request}/__init__.py (100%)
rename tests/{test_litellm => unit}/expected_responses_api_request/azure_shell_tool.json (100%)
rename tests/{test_litellm => unit}/expected_responses_api_request/context_management_and_shell.json (100%)
create mode 100644 tests/unit/litellm_core_utils/conftest.py
rename tests/{test_litellm => unit}/litellm_core_utils/event_loop_lag.py (100%)
create mode 100644 tests/unit/litellm_core_utils/fake_secret_vault.py
rename tests/{test_litellm/litellm_core_utils/llm_response_utils => unit/litellm_core_utils/llm_cost_calc}/__init__.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/llm_response_utils/test_get_api_base.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/messages_with_counts.py (100%)
rename tests/{test_litellm/router_strategy/adaptive_router => unit/litellm_core_utils/prompt_templates}/__init__.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py (100%)
rename tests/{test_litellm/rust_bridge => unit/litellm_core_utils/specialty_caches}/__init__.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_agentic_followup_kwargs.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_anthropic_dedup_factory.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_api_route_to_call_types.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_audio_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_aws_partition.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_bedrock_converse_dedup_factory.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_bug_report.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_chat_completion_agentic_loop.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_classifier_logging.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_cli_token_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_cloud_storage_security.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_codestral_provider_routing.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_core_helpers.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_coroutine_checker.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_dd_tracing.py (85%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_decode_special_tokens.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_dot_notation_indexing.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_duration_parser.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_error_normalization.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_exception_mapping_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_extract_base64_image.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_fallback_generalizations.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_fallback_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_get_litellm_params.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_get_llm_provider_endpoint_match.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_get_llm_provider_logic.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_get_model_cost_map.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_get_supported_openai_params.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_health_check_helpers.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_image_handling.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_initialize_dynamic_callback_params.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_internal_call_metadata.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_json_fragment_accumulator.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_json_schema_validation.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_litellm_logging.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_llm_judge.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_llm_request_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_logging_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_logging_worker.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_max_streaming_duration.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_model_param_helper.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_model_response_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_private_json.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_provider_affinity.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_provider_specific_headers.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_ptu_pricing.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_realtime_errors.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_realtime_streaming.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_redact_messages.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_request_timeout_resolver.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_retry_after_headers.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_safe_divide_seconds.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_safe_json_dumps.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_sensitive_data_masker.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_sentry_scrubbing.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_served_output_texts.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_chunk_builder_cursor.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_chunk_builder_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_handler.py (99%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_streaming_overhead.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_thread_pool_executor.py (100%)
create mode 100644 tests/unit/litellm_core_utils/test_token_counter.py
rename tests/{test_litellm => unit}/litellm_core_utils/test_token_counter_tool.py (93%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_token_counter_tool_data.py (100%)
create mode 100644 tests/unit/litellm_core_utils/test_tokenizer.py
rename tests/{test_litellm => unit}/litellm_core_utils/test_tool_search_spend_logging.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_url_utils.py (100%)
rename tests/{test_litellm => unit}/litellm_core_utils/test_xai_oauth_routing.py (100%)
rename tests/{test_litellm/rust_bridge/chat_completions => unit/responses/litellm_completion_transformation}/__init__.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_function_call_output_normalization.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_handler.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_image_generation_output.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_litellm_completion_responses.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_session_handler.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py (100%)
rename tests/{test_litellm => unit}/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py (100%)
rename tests/{test_litellm => unit}/responses/mcp/test_chat_completions_handler.py (100%)
rename tests/{test_litellm => unit}/responses/mcp/test_litellm_proxy_mcp_handler.py (100%)
rename tests/{test_litellm => unit}/responses/mcp/test_mcp_streaming_iterator.py (100%)
rename tests/{test_litellm => unit}/responses/test_additional_tools.py (100%)
rename tests/{test_litellm => unit}/responses/test_custom_tool_call.py (100%)
rename tests/{test_litellm => unit}/responses/test_dispatch.py (100%)
rename tests/{test_litellm => unit}/responses/test_metadata_codex_callback.py (100%)
rename tests/{test_litellm => unit}/responses/test_no_duplicate_spend_logs.py (76%)
rename tests/{test_litellm => unit}/responses/test_null_test_fix.py (100%)
rename tests/{test_litellm => unit}/responses/test_responses_api_bridge_flag.py (100%)
rename tests/{test_litellm => unit}/responses/test_responses_api_request_body.py (99%)
rename tests/{test_litellm => unit}/responses/test_responses_prompt_management.py (100%)
rename tests/{test_litellm => unit}/responses/test_responses_router_cooldown.py (100%)
rename tests/{test_litellm => unit}/responses/test_responses_streaming_iterator.py (100%)
rename tests/{test_litellm => unit}/responses/test_responses_supported_endpoints_passthrough.py (100%)
rename tests/{test_litellm => unit}/responses/test_responses_utils.py (100%)
rename tests/{test_litellm => unit}/responses/test_responses_websocket_all_providers.py (97%)
rename tests/{test_litellm => unit}/responses/test_rust_bridge_websocket.py (100%)
rename tests/{test_litellm => unit}/responses/test_sse_output_recovery.py (100%)
rename tests/{test_litellm => unit}/responses/test_streaming_iterator.py (100%)
rename tests/{test_litellm => unit}/responses/test_streaming_iterator_error_events.py (100%)
rename tests/{test_litellm => unit}/responses/test_text_format_conversion.py (100%)
rename tests/{test_litellm/rust_bridge/messages => unit/router_strategy/adaptive_router}/__init__.py (100%)
rename tests/{test_litellm/rust_bridge/ocr => unit/router_strategy/adaptive_router/fixtures}/__init__.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/clean_no_signals.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/clean_satisfaction.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/disengagement_giveup.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/exhaustion_429.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/failure_tool_error.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/loop_same_tool.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/fixtures/stagnation_repeat.json (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_adaptive_router.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_async_pre_routing.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_bandit.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_classifier.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_config.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_e2e_adaptive_router.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_hooks.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_router_dispatch.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_signals.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_state_endpoint.py (100%)
rename tests/{test_litellm => unit}/router_strategy/adaptive_router/test_update_queue.py (100%)
rename tests/{test_litellm => unit}/router_strategy/complexity_router/test_context_compaction.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_auto_router.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_base_routing_strategy.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_budget_limiter.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_budget_limiter_hotpath.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_complexity_router.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_complexity_tier_predictor.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_fuse_presets.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_lar1_routing.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_least_busy.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_litellm_encoder.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_llm_v2.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_lowest_cost.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_lowest_latency.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_lowest_tpm_rpm.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_quality_router.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_router_routing_groups.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_router_routing_plugins.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_router_tag_regex_routing.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_router_tag_routing.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_savings_baseline.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_simple_shuffle.py (100%)
rename tests/{test_litellm => unit}/router_strategy/test_stall_detector.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_access_windows.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_add_retry_fallback_headers.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_auto_router_model_naming.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_auto_router_tuning_baseline.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_client_initalization_utils.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_cooldown_cache.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_cooldown_handlers.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_fallback_event_handlers.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_get_retry_from_policy.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_health_check_allowed_fails_integration.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_health_state_cache.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_pattern_match_deployments.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_reasoning_effort_capability.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_router_health_check_routing.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_router_interactions_endpoints.py (100%)
rename tests/{test_litellm => unit}/router_utils/test_router_utils_common_utils.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/AGENTS.md (100%)
rename tests/{test_litellm => unit}/rust_bridge/messages/test_secrets.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/native_route_wheel_test.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/ocr/test_secrets.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/stubtest.ini (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_bindings.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_callbacks_legacy_python.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_catalog.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_configuration.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_dispatch.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_failures.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_fork_guard.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_lifecycle.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_logger.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_runtime.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_secret_manager.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_settings.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_token_counter.py (100%)
rename tests/{test_litellm => unit}/rust_bridge/test_tokenizer.py (95%)
rename tests/{test_litellm => unit}/rust_bridge/test_verify_linux_native_wheel.py (100%)
diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh
index d56e29fb627..3f4f5620176 100755
--- a/.circleci/scripts/unit_selection.sh
+++ b/.circleci/scripts/unit_selection.sh
@@ -5,6 +5,7 @@ flag="${1:?usage: unit_selection.sh }"
legacy_flags=(
caching-local
+ core-utils
enterprise-package
enterprise-routing
integrations
@@ -32,6 +33,7 @@ legacy_flags=(
legacy_paths() {
case "$1" in
caching-local) echo tests/unit/caching ;;
+ core-utils) echo tests/unit/litellm_core_utils ;;
enterprise-package)
echo tests/unit/enterprise/integrations
echo tests/unit/enterprise/proxy/auth
@@ -42,6 +44,8 @@ legacy_paths() {
echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;;
enterprise-routing)
echo tests/unit/google_genai
+ echo tests/unit/router_strategy
+ echo tests/unit/router_utils
echo tests/unit/enterprise/enterprise_callbacks/send_emails
echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py
echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py
@@ -77,6 +81,7 @@ legacy_paths() {
echo tests/unit/messages
echo tests/unit/rag
echo tests/unit/rerank_api
+ echo tests/unit/rust_bridge
echo tests/unit/secret_managers
echo tests/unit/vector_stores
echo tests/unit/videos ;;
@@ -142,7 +147,9 @@ legacy_paths() {
proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;
proxy-extras) echo tests/unit/litellm_proxy_extras ;;
proxy-infra) echo tests/unit/gateway ;;
- responses-caching-types) echo tests/unit/types ;;
+ responses-caching-types)
+ find tests/unit/responses -name 'test_*.py' -not -path 'tests/unit/responses/mcp/*'
+ echo tests/unit/types ;;
*) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;;
esac
}
diff --git a/.circleci/tests.yml b/.circleci/tests.yml
index 41e9f11cefa..a9cd21bad5e 100644
--- a/.circleci/tests.yml
+++ b/.circleci/tests.yml
@@ -369,6 +369,13 @@ workflows:
reruns: 2
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
+ - unit:
+ name: unit-core-utils
+ flag: core-utils
+ shards: 2
+ reruns: 1
+ base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
+ pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
- unit:
name: unit-integrations
flag: integrations
diff --git a/.github/merge-smoke-tests.json b/.github/merge-smoke-tests.json
index a563424c230..727733fa954 100644
--- a/.github/merge-smoke-tests.json
+++ b/.github/merge-smoke-tests.json
@@ -7,9 +7,9 @@
"MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]",
"COST-EXPLICIT": "tests/unit/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones",
"COST-ZERO": "tests/unit/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero",
- "LOG-CONTENT-ON": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on",
- "LOG-CONTENT-OFF": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off",
- "CALLBACK-SUCCESS": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger",
- "CALLBACK-FAILURE": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_failure_handler_delivers_failure_payload_to_custom_logger"
+ "LOG-CONTENT-ON": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on",
+ "LOG-CONTENT-OFF": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off",
+ "CALLBACK-SUCCESS": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger",
+ "CALLBACK-FAILURE": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_async_failure_handler_delivers_failure_payload_to_custom_logger"
}
}
diff --git a/.github/workflows/test-redis-compat.yml b/.github/workflows/test-redis-compat.yml
index 2f5ce4d441a..0423b014ec5 100644
--- a/.github/workflows/test-redis-compat.yml
+++ b/.github/workflows/test-redis-compat.yml
@@ -12,9 +12,9 @@ on:
- "litellm/caching/evicted_client_closer.py"
- "tests/unit/test_redis.py"
- "tests/local_testing/test_caching.py"
- - "tests/test_litellm/caching/test_redis_connection_pool.py"
- - "tests/test_litellm/caching/test_redis_cluster_cache.py"
- - "tests/test_litellm/caching/test_evicted_client_closer.py"
+ - "tests/unit/caching/test_redis_connection_pool.py"
+ - "tests/unit/caching/test_redis_cluster_cache.py"
+ - "tests/unit/caching/test_evicted_client_closer.py"
- ".github/workflows/test-redis-compat.yml"
- "pyproject.toml"
- "uv.lock"
@@ -85,9 +85,9 @@ jobs:
redis-server --version
uv run --no-sync pytest \
tests/unit/test_redis.py \
- tests/test_litellm/caching/test_redis_connection_pool.py \
- tests/test_litellm/caching/test_redis_cluster_cache.py \
- tests/test_litellm/caching/test_evicted_client_closer.py \
+ tests/unit/caching/test_redis_connection_pool.py \
+ tests/unit/caching/test_redis_cluster_cache.py \
+ tests/unit/caching/test_evicted_client_closer.py \
tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_azure_credentials \
tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_gcp_credentials \
--tb=short -vv \
diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml
index 1f3b5c4d97c..808bb2afd08 100644
--- a/.github/workflows/test-rust.yml
+++ b/.github/workflows/test-rust.yml
@@ -24,7 +24,7 @@ on:
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- - "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
+ - "tests/unit/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
pull_request:
branches:
@@ -52,7 +52,7 @@ on:
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- - "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
+ - "tests/unit/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
permissions:
@@ -171,7 +171,7 @@ jobs:
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
- - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
+ - run: python tests/unit/rust_bridge/native_route_wheel_test.py dist/*.whl
- name: Run pytest tests/test_litellm_rust with the compiled extension
run: make test-rust-extension
diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index 4dca8075440..d75213d37ea 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -62,6 +62,7 @@ jobs:
- shard: core-utils
artifact-name: core-utils
test-path: "tests/test_litellm/litellm_core_utils"
+ unit-flag: core-utils
workers: 2
reruns: 1
timeout-minutes: 20
@@ -69,9 +70,7 @@ jobs:
- shard: enterprise-routing
artifact-name: enterprise-routing
- test-path: >-
- tests/test_litellm/router_utils
- tests/test_litellm/router_strategy
+ test-path: ""
unit-flag: enterprise-routing
workers: 2
reruns: 2
@@ -111,7 +110,6 @@ jobs:
tests/test_litellm/interactions
tests/test_litellm/ocr
tests/test_litellm/passthrough
- tests/test_litellm/rust_bridge
tests/test_litellm/test_*.py
unit-flag: misc
workers: 2
@@ -228,9 +226,7 @@ jobs:
- shard: responses-caching-types
artifact-name: responses-caching-types
- test-path: >-
- tests/test_litellm/responses
- tests/test_litellm/caching
+ test-path: ""
unit-flag: responses-caching-types
workers: 2
reruns: 2
diff --git a/Makefile b/Makefile
index f27525b58ff..311a7daef92 100644
--- a/Makefile
+++ b/Makefile
@@ -301,7 +301,7 @@ test-rust-extension:
UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
"$$temporary/venv/bin/python" -I -m mypy.stubtest \
- --mypy-config-file tests/test_litellm/rust_bridge/stubtest.ini \
+ --mypy-config-file tests/unit/rust_bridge/stubtest.ini \
litellm.rust_bridge._native && \
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
@@ -329,10 +329,10 @@ test-unit-integrations: install-test-deps
$(UV_RUN) pytest tests/unit/integrations --tb=short -vv -n 4 --durations=20
test-unit-core-utils: install-test-deps
- $(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20
+ $(UV_RUN) pytest tests/unit/litellm_core_utils --tb=short -vv -n 2 --durations=20
test-unit-other: install-test-deps
- $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20
+ $(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20
test-unit-root: install-test-deps
$(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20
diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs
index 030bf03d4ba..69f72fbc177 100644
--- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs
+++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs
@@ -33,7 +33,7 @@ mod test_support {
use crate::{LegacyLogging, LegacySurface, PublicCall};
/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
- /// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
+ /// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
/// signatures, and [`namespace`] binds every fake call against it.
pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json");
diff --git a/litellm-rust/crates/secrets/README.md b/litellm-rust/crates/secrets/README.md
index c8b01fe9b3a..10619613516 100644
--- a/litellm-rust/crates/secrets/README.md
+++ b/litellm-rust/crates/secrets/README.md
@@ -30,7 +30,7 @@ The HashiCorp Vault backend is enabled with the `hashicorp` feature and reads KV
Native backends consistently distinguish absence from failure instead of swallowing provider errors. Python-compatible resolution maps these results back to the Python handler contract before applying fallback
-`hosted_keys` excludes a name for every backend. Python's handler recognizes Azure `SecretClient` and Google `KeyManagementServiceClient` instances before the `local` branch, allowing excluded names to reach those providers. Rust treats that as a routing bug. `test_rust_hosted_keys_exclude_azure_sdk_clients_too` in `tests/test_litellm/rust_bridge/ocr/test_secrets.py` pins this behavior
+`hosted_keys` excludes a name for every backend. Python's handler recognizes Azure `SecretClient` and Google `KeyManagementServiceClient` instances before the `local` branch, allowing excluded names to reach those providers. Rust treats that as a routing bug. `test_rust_hosted_keys_exclude_azure_sdk_clients_too` in `tests/unit/rust_bridge/ocr/test_secrets.py` pins this behavior
Google rejects malformed base64 and mismatched CRC32C values instead of accepting corrupted payloads. Python currently ignores the checksum and uses permissive base64 decoding. Rust follows [RFC 4648](https://www.rfc-editor.org/rfc/rfc4648#section-3.3) and [Google's integrity guidance](https://docs.cloud.google.com/secret-manager/docs/data-integrity); `failed_or_missing_reads_are_not_cached` covers rejection and recovery
diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py
index 6c87ef4a3de..43e16599bf6 100644
--- a/litellm/litellm_core_utils/url_utils.py
+++ b/litellm/litellm_core_utils/url_utils.py
@@ -77,13 +77,13 @@ class _CallerHeadersView(TypedDict):
headers: ReadOnly[dict[str, str]]
-# Globally-routable IPs that are cloud-internal. Everything else
-# non-public is caught by ``not ip.is_global`` (RFC 6890, as implemented by
-# Python's ``ipaddress`` module). This list only holds IPs that are
-# publicly routable *and* point to cloud-fabric services reachable from
-# inside a VM via special in-fabric routing.
+# Cloud-internal IPs that ``ip.is_global`` can report as public. Everything
+# else non-public is caught by ``not ip.is_global`` (RFC 6890, as implemented
+# by Python's ``ipaddress`` module). Older Python patch releases (3.12.2, for
+# one) treat most of 192.0.0.0/24 as global, so it is listed to block it everywhere.
_CLOUD_METADATA_EXCEPTIONS: Final = [
ip_network("168.63.129.16/32"), # Azure Wire Server
+ ip_network("192.0.0.0/24"),
]
_ALLOWED_SCHEMES: Final = ("http", "https")
diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py
index 74c0478b08b..fbcf97839b9 100644
--- a/tests/llm_responses_api_testing/base_responses_api.py
+++ b/tests/llm_responses_api_testing/base_responses_api.py
@@ -742,7 +742,7 @@ class BaseResponsesAPITest(ABC):
Passes tools=[{"type": "shell", "environment": {"type": "container_auto"}}];
validates that the request is accepted and returns a valid response.
Only runs for OpenAI; offline coverage for the Azure route lives in
- tests/test_litellm/responses/test_responses_api_request_body.py.
+ tests/unit/responses/test_responses_api_request_body.py.
"""
base_completion_call_args = self.get_base_completion_call_args()
model = (
diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py
index 0c3eca52dde..1a34e404d7f 100644
--- a/tests/llm_translation/test_gemini.py
+++ b/tests/llm_translation/test_gemini.py
@@ -1800,7 +1800,7 @@ def test_gemini_image_size_limit_exceeded(monkeypatch):
that could cause memory issues and pod crashes.
The image fetch is mocked (mirroring the LargeImageClient pattern in
- tests/test_litellm/litellm_core_utils/test_image_handling.py) so the test
+ tests/unit/litellm_core_utils/test_image_handling.py) so the test
deterministically exercises the size-limit rejection path without any
external network dependency.
"""
diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py
deleted file mode 100644
index e5a7f1540ca..00000000000
--- a/tests/test_litellm/caching/test_caching_handler.py
+++ /dev/null
@@ -1,867 +0,0 @@
-import asyncio
-import json
-import time
-from unittest.mock import MagicMock, patch
-
-import httpx
-import pytest
-import respx
-from fastapi.testclient import TestClient
-
-from datetime import datetime
-from unittest.mock import AsyncMock
-
-from litellm.caching.caching_handler import _PENDING_CACHE_WRITES, LLMCachingHandler
-
-
-@pytest.mark.asyncio
-async def test_process_async_embedding_cached_response():
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- args = {
- "cached_result": [
- {
- "embedding": [-0.025122925639152527, -0.019487135112285614],
- "index": 0,
- "object": "embedding",
- }
- ]
- }
-
- mock_logging_obj = MagicMock()
- mock_logging_obj.async_success_handler = AsyncMock()
- response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
- final_embedding_cached_response=None,
- cached_result=args["cached_result"],
- kwargs={"model": "text-embedding-ada-002", "input": "test"},
- logging_obj=mock_logging_obj,
- start_time=datetime.now(),
- model="text-embedding-ada-002",
- )
-
- assert cache_hit
-
- print(f"response: {response}")
- assert len(response.data) == 1
-
-
-@pytest.mark.asyncio
-async def test_embedding_cache_preserves_prompt_tokens_details():
- """Test that prompt_tokens_details (including image_count) survives a full cache hit."""
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- cached_result = [
- {
- "embedding": [-0.025, -0.019],
- "index": 0,
- "object": "embedding",
- "model": "amazon.titan-embed-image-v1",
- "prompt_tokens_details": {"image_count": 1},
- }
- ]
-
- mock_logging_obj = MagicMock()
- mock_logging_obj.async_success_handler = AsyncMock()
- response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
- final_embedding_cached_response=None,
- cached_result=cached_result,
- kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"},
- logging_obj=mock_logging_obj,
- start_time=datetime.now(),
- model="amazon.titan-embed-image-v1",
- )
-
- assert cache_hit
- assert response.usage is not None
- assert response.usage.prompt_tokens_details is not None
- assert response.usage.prompt_tokens_details.image_count == 1
-
-
-@pytest.mark.asyncio
-async def test_embedding_cache_backward_compat_no_prompt_tokens_details():
- """Test that old cached items without prompt_tokens_details still work."""
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- # Old-format cached item — no prompt_tokens_details field
- cached_result = [
- {
- "embedding": [-0.025, -0.019],
- "index": 0,
- "object": "embedding",
- "model": "text-embedding-ada-002",
- }
- ]
-
- mock_logging_obj = MagicMock()
- mock_logging_obj.async_success_handler = AsyncMock()
- response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
- final_embedding_cached_response=None,
- cached_result=cached_result,
- kwargs={"model": "text-embedding-ada-002", "input": "test"},
- logging_obj=mock_logging_obj,
- start_time=datetime.now(),
- model="text-embedding-ada-002",
- )
-
- assert cache_hit
- assert response.usage is not None
- assert response.usage.prompt_tokens_details is None
-
-
-@pytest.mark.asyncio
-async def test_embedding_cache_aggregates_multiple_image_counts():
- """Test that image_count is summed correctly across multiple cached items."""
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- cached_result = [
- {
- "embedding": [-0.025, -0.019],
- "index": 0,
- "object": "embedding",
- "model": "amazon.titan-embed-image-v1",
- "prompt_tokens_details": {"image_count": 1},
- },
- {
- "embedding": [0.031, 0.042],
- "index": 1,
- "object": "embedding",
- "model": "amazon.titan-embed-image-v1",
- "prompt_tokens_details": {"image_count": 1},
- },
- ]
-
- mock_logging_obj = MagicMock()
- mock_logging_obj.async_success_handler = AsyncMock()
- response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
- final_embedding_cached_response=None,
- cached_result=cached_result,
- kwargs={
- "model": "amazon.titan-embed-image-v1",
- "input": ["img1", "img2"],
- },
- logging_obj=mock_logging_obj,
- start_time=datetime.now(),
- model="amazon.titan-embed-image-v1",
- )
-
- assert cache_hit
- assert response.usage.prompt_tokens_details is not None
- assert response.usage.prompt_tokens_details.image_count == 2
-
-
-def test_combine_usage_merges_prompt_tokens_details():
- """Test that combine_usage merges prompt_tokens_details from both Usage objects."""
- from litellm.types.utils import PromptTokensDetailsWrapper, Usage
-
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- usage1 = Usage(
- prompt_tokens=10,
- completion_tokens=0,
- total_tokens=10,
- prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
- )
- usage2 = Usage(
- prompt_tokens=20,
- completion_tokens=0,
- total_tokens=20,
- prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2),
- )
-
- combined = llm_caching_handler.combine_usage(usage1, usage2)
-
- assert combined.prompt_tokens == 30
- assert combined.total_tokens == 30
- assert combined.prompt_tokens_details is not None
- assert combined.prompt_tokens_details.image_count == 3
-
-
-def test_combine_usage_handles_none_details():
- """Test that combine_usage works when one or both sides have null prompt_tokens_details."""
- from litellm.types.utils import PromptTokensDetailsWrapper, Usage
-
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- # Both null
- usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10)
- usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20)
- combined = llm_caching_handler.combine_usage(usage_a, usage_b)
- assert combined.prompt_tokens_details is None
-
- # Only first has details
- usage_c = Usage(
- prompt_tokens=10,
- completion_tokens=0,
- total_tokens=10,
- prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
- )
- combined = llm_caching_handler.combine_usage(usage_c, usage_b)
- assert combined.prompt_tokens_details is not None
- assert combined.prompt_tokens_details.image_count == 1
-
- # Only second has details
- combined = llm_caching_handler.combine_usage(usage_a, usage_c)
- assert combined.prompt_tokens_details is not None
- assert combined.prompt_tokens_details.image_count == 1
-
-
-def test_is_chat_completion_cached_dict():
- from litellm.caching.caching_handler import _is_chat_completion_cached_dict
-
- assert _is_chat_completion_cached_dict(
- {"id": "chatcmpl-abc", "object": "chat.completion", "choices": []}
- )
- assert _is_chat_completion_cached_dict(
- {"id": "other", "object": "chat.completion.chunk", "choices": []}
- )
- assert _is_chat_completion_cached_dict(
- {"id": "no-object", "choices": [{"index": 0}]}
- )
- assert not _is_chat_completion_cached_dict(
- {"id": "resp_abc", "object": "response", "output": []}
- )
-
-
-def _build_logging_obj(call_type: str, stream: bool):
- import uuid as _uuid
-
- from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
-
- return LiteLLMLogging(
- litellm_call_id=str(datetime.now()),
- call_type=call_type,
- model="gpt-5.4",
- messages=[],
- function_id=str(_uuid.uuid4()),
- stream=stream,
- start_time=datetime.now(),
- )
-
-
-def test_convert_cached_aresponses_bridge_chat_completion_stream():
- """openai/responses chat-completions bridge: streaming cache hit replays as chat stream."""
- from litellm import aresponses
- from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
- from litellm.types.utils import CallTypes
-
- caching_handler = LLMCachingHandler(
- original_function=aresponses, request_kwargs={}, start_time=datetime.now()
- )
- cached_result = {
- "id": "chatcmpl-bridge-cache-test",
- "object": "chat.completion",
- "created": int(time.time()),
- "model": "gpt-5.4",
- "choices": [
- {
- "index": 0,
- "message": {"role": "assistant", "content": "Hi!"},
- "finish_reason": "stop",
- }
- ],
- "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18},
- }
-
- result = caching_handler._convert_cached_result_to_model_response(
- cached_result=cached_result,
- call_type=CallTypes.aresponses.value,
- kwargs={
- "model": "gpt-5.4",
- "stream": True,
- "messages": [{"role": "user", "content": "hi"}],
- },
- logging_obj=_build_logging_obj(CallTypes.aresponses.value, stream=True),
- model="gpt-5.4",
- args=(),
- )
-
- assert isinstance(result, CustomStreamWrapper)
-
-
-def test_convert_cached_responses_bridge_chat_completion_nonstream():
- """openai/responses chat-completions bridge: non-streaming cache hit replays as ModelResponse."""
- from litellm import responses
- from litellm.types.utils import CallTypes, ModelResponse
-
- caching_handler = LLMCachingHandler(
- original_function=responses, request_kwargs={}, start_time=datetime.now()
- )
- cached_result = {
- "id": "chatcmpl-bridge-nonstream",
- "object": "chat.completion",
- "created": int(time.time()),
- "model": "gpt-5.4",
- "choices": [
- {
- "index": 0,
- "message": {"role": "assistant", "content": "Hi!"},
- "finish_reason": "stop",
- }
- ],
- "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18},
- }
-
- result = caching_handler._convert_cached_result_to_model_response(
- cached_result=cached_result,
- call_type=CallTypes.responses.value,
- kwargs={
- "model": "gpt-5.4",
- "stream": False,
- "messages": [{"role": "user", "content": "hi"}],
- },
- logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False),
- model="gpt-5.4",
- args=(),
- )
-
- assert isinstance(result, ModelResponse)
- assert result.choices[0].message.content == "Hi!"
-
-
-def test_convert_cached_responses_legacy_nonstream_path():
- """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) falls through legacy path."""
- from litellm import responses
- from litellm.types.llms.openai import ResponsesAPIResponse
- from litellm.types.utils import CallTypes
-
- caching_handler = LLMCachingHandler(
- original_function=responses, request_kwargs={}, start_time=datetime.now()
- )
- cached_result = {
- "id": "resp_legacy_nonstream",
- "created_at": int(time.time()),
- "status": "completed",
- "model": "gpt-4o",
- "object": "response",
- "output": [
- {
- "type": "message",
- "id": "msg_legacy",
- "status": "completed",
- "role": "assistant",
- "content": [
- {
- "type": "output_text",
- "text": "legacy response",
- "annotations": [],
- }
- ],
- }
- ],
- }
-
- result = caching_handler._convert_cached_result_to_model_response(
- cached_result=cached_result,
- call_type=CallTypes.responses.value,
- kwargs={"model": "gpt-4o", "input": "hi", "stream": False},
- logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False),
- model="gpt-4o",
- args=(),
- )
-
- assert isinstance(result, ResponsesAPIResponse)
- assert result.id == "resp_legacy_nonstream"
-
-
-def test_convert_cached_responses_legacy_stream_path():
- """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) on stream falls through legacy path."""
- from litellm import responses
- from litellm.responses.streaming_iterator import (
- CachedResponsesAPIStreamingIterator,
- )
- from litellm.types.utils import CallTypes
-
- caching_handler = LLMCachingHandler(
- original_function=responses, request_kwargs={}, start_time=datetime.now()
- )
- cached_result = {
- "id": "resp_legacy_stream",
- "created_at": int(time.time()),
- "status": "completed",
- "model": "gpt-4o",
- "object": "response",
- "output": [
- {
- "type": "message",
- "id": "msg_legacy_stream",
- "status": "completed",
- "role": "assistant",
- "content": [
- {
- "type": "output_text",
- "text": "legacy stream",
- "annotations": [],
- }
- ],
- }
- ],
- }
-
- result = caching_handler._convert_cached_result_to_model_response(
- cached_result=cached_result,
- call_type=CallTypes.responses.value,
- kwargs={"model": "gpt-4o", "input": "hi", "stream": True},
- logging_obj=_build_logging_obj(CallTypes.responses.value, stream=True),
- model="gpt-4o",
- args=(),
- )
-
- assert isinstance(result, CachedResponsesAPIStreamingIterator)
-
-
-@pytest.mark.asyncio
-async def test_embedding_cache_restores_stored_prompt_tokens_for_image_input():
- """Image-embedding cache hit restores prompt_tokens=0 from the stored value
- instead of recomputing a bogus count by tokenizing the base64 input."""
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- # base64-like blob — token_counter over this would return a large nonzero count
- image_input = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" * 50
-
- cached_result = [
- {
- "embedding": [-0.025, -0.019],
- "index": 0,
- "object": "embedding",
- "model": "amazon.titan-embed-image-v1",
- "prompt_tokens": 0,
- "prompt_tokens_details": {"image_count": 1},
- }
- ]
-
- mock_logging_obj = MagicMock()
- mock_logging_obj.async_success_handler = AsyncMock()
- response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
- final_embedding_cached_response=None,
- cached_result=cached_result,
- kwargs={"model": "amazon.titan-embed-image-v1", "input": image_input},
- logging_obj=mock_logging_obj,
- start_time=datetime.now(),
- model="amazon.titan-embed-image-v1",
- )
-
- assert cache_hit
- assert response.usage is not None
- assert response.usage.prompt_tokens == 0
- assert response.usage.total_tokens == 0
- assert response.usage.prompt_tokens_details.image_count == 1
-
-
-@pytest.mark.asyncio
-async def test_embedding_cache_sums_stored_prompt_tokens_across_items():
- """A multi-item cache hit sums the stored per-item prompt_tokens back to the total."""
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- cached_result = [
- {
- "embedding": [-0.01],
- "index": 0,
- "object": "embedding",
- "model": "text-embedding-3-small",
- "prompt_tokens": 5,
- },
- {
- "embedding": [-0.02],
- "index": 1,
- "object": "embedding",
- "model": "text-embedding-3-small",
- "prompt_tokens": 4,
- },
- ]
-
- mock_logging_obj = MagicMock()
- mock_logging_obj.async_success_handler = AsyncMock()
- response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
- final_embedding_cached_response=None,
- cached_result=cached_result,
- kwargs={"model": "text-embedding-3-small", "input": ["hello world", "foo bar"]},
- logging_obj=mock_logging_obj,
- start_time=datetime.now(),
- model="text-embedding-3-small",
- )
-
- assert cache_hit
- assert response.usage.prompt_tokens == 9
- assert response.usage.total_tokens == 9
-
-
-@pytest.mark.asyncio
-async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries():
- """Legacy cache entries with no stored prompt_tokens still recompute via token_counter
- for str inputs (backward compatibility)."""
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- # No prompt_tokens key — pre-fix entry
- cached_result = [
- {
- "embedding": [-0.025, -0.019],
- "index": 0,
- "object": "embedding",
- "model": "text-embedding-ada-002",
- },
- ]
-
- mock_logging_obj = MagicMock()
- mock_logging_obj.async_success_handler = AsyncMock()
- response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
- final_embedding_cached_response=None,
- cached_result=cached_result,
- kwargs={"model": "text-embedding-ada-002", "input": "hello world"},
- logging_obj=mock_logging_obj,
- start_time=datetime.now(),
- model="text-embedding-ada-002",
- )
-
- assert cache_hit
- # token_counter over "hello world" yields a nonzero count — fallback path still runs
- assert response.usage.prompt_tokens > 0
-
-
-@pytest.mark.asyncio
-async def test_embedding_cache_hit_sets_custom_llm_provider_on_logging_obj():
- """A full embedding cache hit must stamp the resolved provider onto the logging
- obj so spend logs record the provider instead of None/unknown."""
- from litellm.types.utils import CallTypes
-
- llm_caching_handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs={},
- start_time=datetime.now(),
- )
-
- cached_result = [
- {
- "embedding": [-0.025, -0.019],
- "index": 0,
- "object": "embedding",
- "model": "text-embedding-3-small",
- "prompt_tokens": 5,
- }
- ]
-
- logging_obj = _build_logging_obj(CallTypes.aembedding.value, stream=False)
- logging_obj.async_success_handler = AsyncMock()
-
- response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
- final_embedding_cached_response=None,
- cached_result=cached_result,
- kwargs={"model": "text-embedding-3-small", "input": "hello world"},
- logging_obj=logging_obj,
- start_time=datetime.now(),
- model="text-embedding-3-small",
- )
-
- assert cache_hit
- assert logging_obj.model_call_details["custom_llm_provider"] == "openai"
-
-
-def test_sync_stream_responses_cache_hit_sets_custom_llm_provider_on_logging_obj(monkeypatch):
- import litellm
- from litellm.caching.caching import Cache
- from litellm.types.utils import CallTypes
-
- monkeypatch.setattr(litellm, "cache", Cache(type="local"))
- kwargs = {"model": "azure/gpt-5.4-mini", "input": "hello", "stream": True}
- cached_response = {
- "id": "resp_sync_stream",
- "created_at": int(time.time()),
- "status": "completed",
- "model": "gpt-5.4-mini",
- "object": "response",
- "output": [
- {
- "type": "message",
- "id": "msg_sync_stream",
- "status": "completed",
- "role": "assistant",
- "content": [{"type": "output_text", "text": "hi", "annotations": []}],
- }
- ],
- }
- litellm.cache.add_cache(json.dumps(cached_response), **kwargs)
- handler = LLMCachingHandler(original_function=litellm.responses, request_kwargs=kwargs, start_time=datetime.now())
- logging_obj = _build_logging_obj(CallTypes.responses.value, stream=True)
-
- hit = handler._sync_get_cache(
- model="azure/gpt-5.4-mini",
- original_function=litellm.responses,
- logging_obj=logging_obj,
- start_time=datetime.now(),
- call_type=CallTypes.responses.value,
- kwargs=kwargs,
- args=(),
- )
-
- assert hit.cached_result is not None
- assert logging_obj.model_call_details["custom_llm_provider"] == "azure"
- assert logging_obj.model_call_details["litellm_params"]["custom_llm_provider"] == "azure"
-
-
-def test_request_kwargs_does_not_retain_logging_obj():
- """
- The caching handler lives on logging_obj._llm_caching_handler, so keeping
- litellm_logging_obj inside request_kwargs closes a reference cycle
- (Logging -> LLMCachingHandler -> kwargs -> Logging). That cycle keeps the
- full request payload alive until a generational GC pass instead of being
- freed by refcount when the request finishes; under bursts of large-token
- requests this presents as stepwise RSS growth that never returns to
- baseline. Other kwargs (messages included) must be preserved.
- """
- logging_obj = MagicMock()
- kwargs = {
- "model": "gpt-4o",
- "messages": [{"role": "user", "content": "hello"}],
- "litellm_logging_obj": logging_obj,
- }
-
- handler = LLMCachingHandler(
- original_function=MagicMock(),
- request_kwargs=kwargs,
- start_time=datetime.now(),
- )
-
- assert "litellm_logging_obj" not in handler.request_kwargs
- assert handler.request_kwargs["messages"] == kwargs["messages"]
- assert handler.request_kwargs["model"] == "gpt-4o"
-
-
-def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch):
- """
- Regression test for the SDK losing async cache writes in short-lived scripts:
- async_set_cache dispatched the write as a bare fire-and-forget task, so
- asyncio.run cancelled it at loop close before the write landed (LIT-6184,
- deterministic with hiredis installed). The write must survive loop shutdown.
- """
- import litellm
-
- writes = []
-
- class _SlowWriteCache:
- supported_call_types = ["acompletion"]
- cache = None
-
- async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs):
- await asyncio.sleep(0.2)
- writes.append(result)
-
- async def acompletion(**kwargs):
- return None
-
- handler = LLMCachingHandler(
- original_function=acompletion,
- request_kwargs={},
- start_time=datetime.now(),
- )
- monkeypatch.setattr(litellm, "cache", _SlowWriteCache())
-
- async def _short_lived_script():
- await handler.async_set_cache(
- result=litellm.ModelResponse(),
- original_function=acompletion,
- kwargs={},
- )
-
- asyncio.run(_short_lived_script())
-
- assert len(writes) == 1
-
-
-@pytest.mark.asyncio
-async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monkeypatch):
- """The spend log for a cache hit must reuse the key the lookup already computed instead of hashing again."""
- import litellm
- from litellm.caching.caching import Cache
- from litellm.types.utils import CallTypes
-
- async def acompletion(**kwargs):
- return None
-
- monkeypatch.setattr(litellm, "cache", Cache(type="local"))
- kwargs = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "caching": True}
- await litellm.cache.async_add_cache(
- litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}]), **kwargs
- )
- handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now())
- logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False)
- logging_obj.async_success_handler = AsyncMock()
-
- hit = await handler._async_get_cache(
- model="gpt-5.4",
- original_function=acompletion,
- logging_obj=logging_obj,
- start_time=datetime.now(),
- call_type=CallTypes.acompletion.value,
- kwargs=kwargs,
- args=(),
- )
-
- assert hit is not None and hit.cached_result is not None
- assert handler.preset_cache_key is not None
- assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key
- assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key
-
-
-@pytest.mark.asyncio
-async def test_converted_stream_cache_hit_replayed_as_plain_object_logs_at_hit_time(monkeypatch):
- import litellm
- from litellm.caching.caching import Cache
- from litellm.types.utils import CallTypes
-
- async def aanthropic_messages(**kwargs):
- return None
-
- monkeypatch.setattr(litellm, "cache", Cache(type="local"))
- kwargs = {
- "model": "claude-sonnet-5",
- "messages": [{"role": "user", "content": "hello"}],
- "max_tokens": 16,
- "caching": True,
- "stream": False,
- "_websearch_interception_converted_stream": True,
- }
- cached_message = {
- "id": "msg_1",
- "type": "message",
- "role": "assistant",
- "content": [{"type": "text", "text": "hi"}],
- }
- await litellm.cache.async_add_cache(cached_message, **kwargs)
- handler = LLMCachingHandler(original_function=aanthropic_messages, request_kwargs=kwargs, start_time=datetime.now())
- logging_obj = _build_logging_obj(CallTypes.aanthropic_messages.value, stream=False)
- logging_obj.async_success_handler = AsyncMock()
- logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
-
- hit = await handler._async_get_cache(
- model="claude-sonnet-5",
- original_function=aanthropic_messages,
- logging_obj=logging_obj,
- start_time=datetime.now(),
- call_type=CallTypes.aanthropic_messages.value,
- kwargs=kwargs,
- args=(),
- )
-
- assert hit is not None and hit.cached_result == cached_message
- logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
- assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
-
-
-@pytest.mark.asyncio
-async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_replays_as_plain_object(monkeypatch):
- import litellm
- from litellm.caching.caching import Cache
- from litellm.types.utils import CallTypes
-
- async def acompletion(**kwargs):
- return None
-
- monkeypatch.setattr(litellm, "cache", Cache(type="local"))
- kwargs = {
- "model": "gpt-5.6",
- "messages": [{"role": "user", "content": "run the code"}],
- "caching": True,
- "stream": False,
- "_code_interpreter_interception_converted_stream": True,
- "_agentic_loop_depth": 1,
- }
- await litellm.cache.async_add_cache(
- litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]), **kwargs
- )
- handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now())
- logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False)
- logging_obj.async_success_handler = AsyncMock()
- logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
-
- hit = await handler._async_get_cache(
- model="gpt-5.6",
- original_function=acompletion,
- logging_obj=logging_obj,
- start_time=datetime.now(),
- call_type=CallTypes.acompletion.value,
- kwargs=kwargs,
- args=(),
- )
-
- assert hit is not None and isinstance(hit.cached_result, litellm.ModelResponse)
- assert hit.cached_result.choices[0].message.content == "done"
- logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
- assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
-
-
-@pytest.mark.asyncio
-async def test_partial_embedding_cache_hit_sends_only_misses_and_keeps_input_order(monkeypatch):
- import litellm
- from litellm import CustomLLM
- from litellm.caching.caching import Cache
- from litellm.types.utils import Embedding, EmbeddingResponse
-
- class RecordingEmbedder(CustomLLM):
- provider_inputs: tuple[tuple[str, ...], ...] = ()
-
- async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse:
- self.provider_inputs = (*self.provider_inputs, tuple(input))
- return EmbeddingResponse(
- model=model,
- data=[
- Embedding(embedding=[float(len(text))], index=idx, object="embedding")
- for idx, text in enumerate(input)
- ],
- )
-
- embedder = RecordingEmbedder()
- monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "recording-embedder", "custom_handler": embedder}])
- monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "recording-embedder"])
- monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "recording-embedder"])
- monkeypatch.setattr(litellm, "cache", Cache(type="local"))
-
- await litellm.aembedding(model="recording-embedder/m", input=["aa", "bbbb"])
- await asyncio.gather(*_PENDING_CACHE_WRITES)
- mixed_input = ["c", "aa", "ddd", "bbbb", "eeeee"]
- response = await litellm.aembedding(model="recording-embedder/m", input=mixed_input)
- await asyncio.gather(*_PENDING_CACHE_WRITES)
-
- assert embedder.provider_inputs == (("aa", "bbbb"), ("c", "ddd", "eeeee")), embedder.provider_inputs
- assert [item["index"] for item in response.data] == [0, 1, 2, 3, 4]
- assert [item["embedding"] for item in response.data] == [[float(len(text))] for text in mixed_input]
- assert response._hidden_params["cache_hit"] is True, "a partial hit must still be reported as a cache hit"
-
- repeat = await litellm.aembedding(model="recording-embedder/m", input=mixed_input)
-
- assert len(embedder.provider_inputs) == 2, embedder.provider_inputs
- assert [item["embedding"] for item in repeat.data] == [[float(len(text))] for text in mixed_input]
diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py
index f8c7d5273d1..f83c1e76b3a 100644
--- a/tests/test_litellm/conftest.py
+++ b/tests/test_litellm/conftest.py
@@ -22,19 +22,6 @@ import litellm
from litellm import router as litellm_router_module
from litellm import utils as litellm_utils_module
from litellm._logging import ALL_LOGGERS
-from litellm.litellm_core_utils.cli_keyring import (
- KeyringDiscardsWrites,
- KeyringUnreachable,
- KeyringUnusable,
- SecretErase,
- SecretErased,
- SecretFound,
- SecretMissing,
- SecretRead,
- SecretStored,
- SecretStranded,
- SecretWrite,
-)
from litellm.litellm_core_utils.prompt_templates import (
image_handling as image_handling_module,
)
@@ -42,6 +29,7 @@ from litellm.llms.custom_httpx.async_client_cleanup import (
close_litellm_async_clients,
)
from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module
+from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault
def _reset_module_level_aws_auth_caches():
@@ -128,60 +116,6 @@ def isolate_host_os_keychain(monkeypatch):
monkeypatch.setenv("LITELLM_CLI_DISABLE_KEYRING", "1")
-class FakeSecretVault:
- """In-memory stand-in for the OS keychain, injected wherever CLI credential storage is exercised.
-
- `available=False` models a keychain that is locked or has no backend, `writable=False` one that
- refuses to store, `erasable=False` one that will not release what it already holds, and `failure`
- picks which unusable state those report. `discards=True` is keyring's null backend, which answers
- reads and erases like any other yet keeps nothing it is given, so only writes report it.
- """
-
- def __init__(
- self,
- blob: str | None = None,
- *,
- available: bool = True,
- writable: bool = True,
- erasable: bool = True,
- discards: bool = False,
- failure: KeyringUnusable = KeyringUnreachable(),
- ) -> None:
- self.blob: str | None = blob
- self.available: bool = available
- self.writable: bool = writable
- self.erasable: bool = erasable
- self.discards: bool = discards
- self.failure: KeyringUnusable = failure
- self.reads: int = 0
- self.writes: list[str] = []
- self.erases: int = 0
-
- def read(self) -> SecretRead:
- self.reads += 1
- if not self.available:
- return self.failure
- return SecretMissing() if self.blob is None else SecretFound(self.blob)
-
- def write(self, blob: str) -> SecretWrite:
- self.writes.append(blob)
- if not (self.available and self.writable):
- return self.failure
- if self.discards:
- return KeyringDiscardsWrites()
- self.blob = blob
- return SecretStored()
-
- def erase(self) -> SecretErase:
- self.erases += 1
- if not self.available:
- return self.failure
- if not self.erasable:
- return SecretStranded() if self.blob is not None else SecretErased()
- self.blob = None
- return SecretErased()
-
-
@pytest.fixture
def secret_vault_factory():
"""Build FakeSecretVault instances; see its docstring for the failure modes it can model."""
diff --git a/tests/test_litellm/litellm_core_utils/__init__.py b/tests/test_litellm/litellm_core_utils/__init__.py
index 8c64613a5da..e69de29bb2d 100644
--- a/tests/test_litellm/litellm_core_utils/__init__.py
+++ b/tests/test_litellm/litellm_core_utils/__init__.py
@@ -1 +0,0 @@
-# This file makes the tests/litellm/litellm_core_utils directory a Python package
diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py
index eccf44a1bda..1e10b7e82b1 100644
--- a/tests/test_litellm/litellm_core_utils/test_token_counter.py
+++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py
@@ -1,442 +1,6 @@
-#### What this tests ####
-# This tests litellm.token_counter.token_counter() function
-import asyncio
-import base64
-import importlib
-import threading
-import time
-import traceback
-from concurrent.futures import Future, wait
-from typing import Final
-from unittest.mock import MagicMock
-
-import anyio.to_thread
import pytest
-import tiktoken
-
-from unittest.mock import AsyncMock, patch
-
-import litellm
-from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens
-from litellm import token_counter as token_counter_old
-import litellm.constants
-from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
-from litellm.litellm_core_utils.asyncify import asyncify
-from litellm.litellm_core_utils.token_counter import (
- _get_exact_count_function,
- _get_extrapolating_count_function,
- _get_tiktoken_count_function,
- calculate_img_tokens,
- high_detail_image_token_upper_bound,
- offload_token_count,
-)
-from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
-from tests.large_text import text
-from tests.test_litellm.litellm_core_utils.event_loop_lag import (
- assert_loop_stayed_free,
- timed_with_loop_lags,
- warm_tokenizer,
-)
-from tests.test_litellm.litellm_core_utils.messages_with_counts import (
- MESSAGES_TEXT,
- MESSAGES_WITH_IMAGES,
- MESSAGES_WITH_TOOLS,
-)
-
-
-def token_counter_both_assert_same(**args):
- new = token_counter_new(**args)
- old = token_counter_old(**args)
- assert new == old, f"New token counter {new} does not match old token counter {old}"
- return new
-
-
-## Choose which token_counter the test will use.
-
-# token_counter = token_counter_new
-# token_counter = token_counter_old
-token_counter = token_counter_both_assert_same
-
-
-def test_token_counter_basic():
- assert (
- token_counter(
- model="claude-2",
- messages=[
- {
- "role": "user",
- "content": "This is a long message that definitely exceeds the token limit.",
- }
- ],
- )
- == 19
- )
-
-
-def test_token_counter_large_repeated_text_is_fast():
- messages = [{"role": "user", "content": [{"type": "text", "text": "A" * 1024 * 1024}]}]
-
- start_time = time.perf_counter()
- tokens = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages)
- elapsed = time.perf_counter() - start_time
-
- assert elapsed < 2, f"Token counting took too long: {elapsed:.2f}s"
- assert tokens > 0
-
-
-@pytest.mark.parametrize(
- "text",
- [
- "Short text",
- "This is a normal message with punctuation, numbers, and a few words.",
- ],
-)
-def test_token_counter_short_text_matches_tiktoken(text):
- encoding = tiktoken.get_encoding("cl100k_base")
- expected = len(encoding.encode(text, disallowed_special=()))
-
- assert token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text) == expected
-
-
-def test_token_counter_default_encoding_matches_cl100k():
- encoding: Final = tiktoken.get_encoding("cl100k_base")
- expected: Final = len(encoding.encode("hello world", disallowed_special=()))
-
- assert token_counter_new(model=None, text="hello world") == expected
-
-
-def test_token_counter_text_over_chunk_boundary_stays_close_to_tiktoken():
- text = ("The quick brown fox jumps over the lazy dog. " * 30)[:1025]
- encoding = tiktoken.get_encoding("cl100k_base")
- expected = len(encoding.encode(text, disallowed_special=()))
-
- actual = token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text)
-
- assert abs(actual - expected) <= 4
-
-
-@pytest.mark.parametrize(
- "configured",
- ["0", "-1", "-1024", "not-an-int", "", " ", "999999999", "inf", "1e9"],
-)
-def test_invalid_chunk_size_config_stays_usable(monkeypatch, configured):
- """A misconfigured chunk size must not raise, count zero, or restore the quadratic encode cost."""
- monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", configured)
- try:
- reloaded = importlib.reload(litellm.constants)
- chunk_size = reloaded.TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
- assert 1 <= chunk_size <= reloaded.TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS
-
- encoding = tiktoken.get_encoding("cl100k_base")
- count_tokens = _get_tiktoken_count_function(
- lambda text: len(encoding.encode(text, disallowed_special=())),
- chunk_size=chunk_size,
- )
- assert count_tokens("The quick brown fox jumps over the lazy dog. " * 40) > 0
- finally:
- monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS")
- importlib.reload(litellm.constants)
-
-
-def test_valid_chunk_size_config_is_honoured(monkeypatch):
- monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", "2048")
- try:
- assert importlib.reload(litellm.constants).TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS == 2048
- finally:
- monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS")
- importlib.reload(litellm.constants)
-
-
-async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free():
- warm_tokenizer("claude-fable-5")
-
- tokens, took, lags = await timed_with_loop_lags(
- lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100)
- )
-
- assert tokens > 0
- assert_loop_stayed_free(took, lags)
-
-
-@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500])
-def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int):
- count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk))
- front_heavy: Final = "a" * 1_000 + "b" * 4_000
- exact: Final = 1_000 + len(front_heavy)
-
- estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy)
-
- assert abs(estimate - exact) <= exact // 100
- assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars
-
-
-def test_count_at_or_below_the_cap_is_exact():
- count_exactly: Final = MagicMock(side_effect=len)
-
- assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000
- assert count_exactly.call_args_list == [(("a" * 5_000,),)]
-
-
-class _SlowEncoder:
- def __init__(self) -> None:
- self._lock: Final = threading.Lock()
- self.in_flight = 0
- self.peak_in_flight = 0
-
- def encode_batch_fast(self, texts: list[str]) -> list[list[int]]:
- with self._lock:
- self.in_flight += 1
- self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
- time.sleep(0.1)
- with self._lock:
- self.in_flight -= 1
- return [[0] * len(text) for text in texts]
-
-
-@pytest.mark.asyncio
-async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool():
- encoder: Final = _SlowEncoder()
- count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder})
- shared_pool: Final = anyio.to_thread.current_default_thread_limiter()
- burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
-
- async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]:
- if counting.done():
- return ()
- await asyncio.sleep(0.01)
- return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting))
-
- counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst))))
- borrowed: Final = await shared_pool_borrowed_until_done(counting)
-
- assert await counting == [3] * burst
- assert len(borrowed) > 1 and max(borrowed) == 0
- assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
-
-
-def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None:
- def slow_count(counted: str) -> int:
- time.sleep(0.1)
- return len(counted)
-
- result.set_result(asyncio.run(offload_token_count(slow_count)(text)))
-
-
-def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process():
- loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
- results: Final = tuple(Future[int]() for _ in range(loops))
- threads: Final = tuple(
- threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True)
- for size, result in enumerate(results, start=1)
- )
- for thread in threads:
- thread.start()
-
- _, pending = wait(results, timeout=5)
-
- assert not pending
- assert tuple(result.result() for result in results) == tuple(range(1, loops + 1))
-
-
-@pytest.mark.parametrize(
- ("configured", "expected"),
- [("8", 8), ("0", 4), ("not-an-int", 4)],
-)
-def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int):
- monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured)
- try:
- assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected
- finally:
- monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS")
- importlib.reload(litellm.constants)
-
-
-def test_token_counter_applies_the_default_cap():
- max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS
- prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars]
- over_the_cap: Final = prose + "a" * 200_000
- exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap)
-
- estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap)
-
- assert estimate != exact
- assert abs(estimate - exact) <= exact // 100
-
-
-@pytest.mark.parametrize(
- ("configured", "expected"),
- [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)],
-)
-def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int):
- monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured)
- try:
- assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected
- finally:
- monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS")
- importlib.reload(litellm.constants)
-
-
-def test_token_counter_with_prefix():
- messages = [
- {"role": "user", "content": "Who won the world cup in 2022?"},
- {"role": "assistant", "content": "Argentina", "prefix": True},
- ]
- tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
- assert tokens == 22, f"Expected 22 tokens, got {tokens}"
-
-
-def test_token_counter_normal_plus_function_calling():
- messages = [
- {"role": "system", "content": "System prompt"},
- {"role": "user", "content": "content1"},
- {"role": "assistant", "content": "content2"},
- {"role": "user", "content": "conten3"},
- {
- "role": "assistant",
- "content": None,
- "tool_calls": [
- {
- "id": "call_E0lOb1h6qtmflUyok4L06TgY",
- "function": {
- "arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}',
- "name": "SearchInternet",
- },
- "type": "function",
- }
- ],
- },
- {
- "tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY",
- "role": "tool",
- "name": "SearchInternet",
- "content": "tool content",
- },
- ]
- tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
- assert tokens == 80
-
-
-# test_token_counter_normal_plus_function_calling()
-
-
-def test_token_counter_legacy_function_call_counts_arguments():
- """
- Regression for VERIA-492 (Token-counter function_call bypass).
-
- The legacy OpenAI assistant `function_call` field carries arbitrary text in
- `arguments`. Before the fix, `_count_messages` had no branch for
- `function_call` and fell through to the unsupported-key `continue`, so an
- assistant turn could smuggle unlimited text past `token_counter` and the
- proxy `/utils/token_counter` endpoint (and downstream pre-call budget /
- `get_modified_max_tokens` math). After the fix it must be counted the
- same as the equivalent `tool_calls` payload.
- """
- long_arg = "A" * 4000
- fc_messages = [
- {"role": "user", "content": "hi"},
- {
- "role": "assistant",
- "content": None,
- "function_call": {"name": "search", "arguments": long_arg},
- },
- ]
- tc_messages = [
- {"role": "user", "content": "hi"},
- {
- "role": "assistant",
- "content": None,
- "tool_calls": [
- {
- "id": "call_1",
- "type": "function",
- "function": {"name": "search", "arguments": long_arg},
- }
- ],
- },
- ]
- fc_tokens = token_counter(model="gpt-3.5-turbo", messages=fc_messages)
- tc_tokens = token_counter(model="gpt-3.5-turbo", messages=tc_messages)
- assert fc_tokens == tc_tokens, (
- f"function_call arguments must count like tool_calls arguments; "
- f"got function_call={fc_tokens}, tool_calls={tc_tokens}"
- )
- assert fc_tokens > 500, f"4000-char arguments payload must contribute real tokens, got {fc_tokens}"
-
-
-@pytest.mark.parametrize(
- "message_count_pair",
- MESSAGES_TEXT,
-)
-def test_token_counter_textonly(message_count_pair):
- counted_tokens = token_counter(
- model="gpt-35-turbo", messages=[message_count_pair["message"]]
- )
- assert counted_tokens == message_count_pair["count"]
-
-
-@pytest.mark.parametrize(
- "message_count_pair",
- MESSAGES_TEXT,
-)
-def test_token_counter_count_response_tokens(message_count_pair):
- counted_tokens = token_counter(
- model="gpt-35-turbo",
- messages=[message_count_pair["message"]],
- count_response_tokens=True,
- )
- # 3 tokens are not added because of count_response_tokens=True
- expected = message_count_pair["count"] - 3
- assert counted_tokens == expected
-
-
-@pytest.mark.parametrize(
- "message_count_pair",
- MESSAGES_WITH_IMAGES,
-)
-def test_token_counter_with_images(message_count_pair):
- counted_tokens = token_counter(
- model="gpt-4o", messages=[message_count_pair["message"]]
- )
- assert counted_tokens == message_count_pair["count"]
-
-
-@pytest.mark.parametrize(
- "message_count_pair",
- MESSAGES_WITH_TOOLS,
-)
-def test_token_counter_with_tools(message_count_pair):
- counted_tokens = token_counter(
- model="gpt-35-turbo",
- messages=[message_count_pair["system_message"]],
- tools=message_count_pair["tools"],
- tool_choice=message_count_pair["tool_choice"],
- )
- expected_tokens = message_count_pair["count"]
- actual_diff = counted_tokens - expected_tokens
-
- if "count-tolerate" in message_count_pair:
- if message_count_pair["count-tolerate"] == counted_tokens:
- pass # expected
- else:
- tolerated_diff = message_count_pair["count-tolerate"] - expected_tokens
- assert (
- actual_diff <= tolerated_diff
- ), f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by {tolerated_diff} in the over-counting direction."
- if actual_diff != tolerated_diff:
- raise NeedsToleranceUpdateError(
- f"SOMETHING BROKEN GOT FIXED! THIS is good! Adjust 'count-tolerate' from {message_count_pair['count-tolerate']} to {counted_tokens}"
- )
-
- else:
- assert (
- expected_tokens == counted_tokens
- ), f"Expected {expected_tokens} tokens, got {counted_tokens}."
-
-
-class NeedsToleranceUpdateError(Exception):
- """Custom exception to mark tests that have improved"""
-
- pass
+from litellm import create_pretrained_tokenizer
+from tests.unit.litellm_core_utils.test_token_counter import token_counter
def test_tokenizers():
@@ -449,32 +13,22 @@ def test_tokenizers():
openai_tokens = token_counter(model="gpt-3.5-turbo", text=sample_text)
# claude tokenizer
- claude_tokens = token_counter(
- model="claude-3-5-haiku-20241022", text=sample_text
- )
+ claude_tokens = token_counter(model="claude-3-5-haiku-20241022", text=sample_text)
# cohere tokenizer
cohere_tokens = token_counter(model="command-nightly", text=sample_text)
# llama2 tokenizer
- llama2_tokens = token_counter(
- model="meta-llama/Llama-2-7b-chat", text=sample_text
- )
+ llama2_tokens = token_counter(model="meta-llama/Llama-2-7b-chat", text=sample_text)
# llama3 tokenizer (also testing custom tokenizer)
- llama3_tokens_1 = token_counter(
- model="meta-llama/llama-3-70b-instruct", text=sample_text
- )
+ llama3_tokens_1 = token_counter(model="meta-llama/llama-3-70b-instruct", text=sample_text)
try:
llama3_tokenizer = create_pretrained_tokenizer("Xenova/llama-3-tokenizer")
except Exception as e:
- pytest.skip(
- f"custom tokenizer download failed (HF hub unreachable): {e}"
- )
- llama3_tokens_2 = token_counter(
- custom_tokenizer=llama3_tokenizer, text=sample_text
- )
+ pytest.skip(f"custom tokenizer download failed (HF hub unreachable): {e}")
+ llama3_tokens_2 = token_counter(custom_tokenizer=llama3_tokenizer, text=sample_text)
print(
f"openai tokens: {openai_tokens}; claude tokens: {claude_tokens}; cohere tokens: {cohere_tokens}; llama2 tokens: {llama2_tokens}; llama3 tokens: {llama3_tokens_1}"
@@ -485,1117 +39,13 @@ def test_tokenizers():
# model hub is unreachable (e.g. in CI). In that case the count will
# equal the openai count and the differentiation assertion is skipped.
if openai_tokens == llama2_tokens:
- pytest.skip(
- "llama2 fell back to tiktoken (HF hub unreachable); skipping differentiation assertion"
- )
+ pytest.skip("llama2 fell back to tiktoken (HF hub unreachable); skipping differentiation assertion")
assert llama2_tokens != llama3_tokens_1, "Token values are not different."
- assert (
- llama3_tokens_1 == llama3_tokens_2
- ), "Custom tokenizer is not being used! It has been configured to use the same tokenizer as the built in llama3 tokenizer and the results should be the same."
+ assert llama3_tokens_1 == llama3_tokens_2, (
+ "Custom tokenizer is not being used! It has been configured to use the same tokenizer as the built in llama3 tokenizer and the results should be the same."
+ )
print("test tokenizer: It worked!")
except Exception as e:
pytest.fail(f"An exception occured: {e}")
-
-
-# test_tokenizers()
-
-
-def test_encoding_and_decoding():
- try:
- sample_text = "Hellö World, this is my input string!"
- # openai encoding + decoding
- openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text)
- openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens)
-
- assert openai_text == sample_text
-
- # claude encoding + decoding
- claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text)
-
- claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens)
-
- assert claude_text == sample_text
-
- # cohere encoding + decoding
- cohere_tokens = encode(model="command-nightly", text=sample_text)
- cohere_text = decode(model="command-nightly", tokens=cohere_tokens)
-
- assert cohere_text == sample_text
-
- # llama2 encoding + decoding
- llama2_tokens = encode(model="meta-llama/Llama-2-7b-chat", text=sample_text)
- llama2_text = decode(model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens)
-
- assert llama2_text == sample_text
- except Exception as e:
- pytest.fail(f"An exception occured: {e}\n{traceback.format_exc()}")
-
-
-# test_encoding_and_decoding()
-
-
-def test_gpt_vision_token_counting():
- messages = [
- {
- "role": "user",
- "content": [
- {"type": "text", "text": "What’s in this image?"},
- {
- "type": "image_url",
- "image_url": "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png",
- },
- ],
- }
- ]
- tokens = token_counter(model="gpt-4-vision-preview", messages=messages)
- print(f"tokens: {tokens}")
-
-
-# test_gpt_vision_token_counting()
-
-
-@pytest.mark.parametrize(
- "model",
- [
- "gpt-4-vision-preview",
- "gpt-4o",
- "claude-3-opus-20240229",
- "command-nightly",
- "mistral/mistral-tiny",
- ],
-)
-def test_load_test_token_counter(model):
- """
- Token count large prompt 100 times.
-
- Assert time taken is < 1.5s.
- """
- import tiktoken
-
- messages = [{"role": "user", "content": text}] * 10
-
- start_time = time.time()
- for _ in range(10):
- _ = token_counter(model=model, messages=messages)
- # enc.encode("".join(m["content"] for m in messages))
-
- end_time = time.time()
-
- total_time = end_time - start_time
- print("model={}, total test time={}".format(model, total_time))
- assert total_time < 10, f"Total encoding time > 10s, {total_time}"
-
-
-def test_openai_token_with_image_and_text():
- model = "gpt-4o"
- full_request = {
- "model": "gpt-4o",
- "tools": [
- {
- "type": "function",
- "function": {
- "name": "json",
- "parameters": {
- "type": "object",
- "required": ["clause"],
- "properties": {"clause": {"type": "string"}},
- },
- "description": "Respond with a JSON object.",
- },
- }
- ],
- "logprobs": False,
- "messages": [
- {
- "role": "user",
- "content": [
- {
- "text": "\n Just some long text, long long text, and you know it will be longer than 7 tokens definetly.",
- "type": "text",
- }
- ],
- }
- ],
- "tool_choice": {"type": "function", "function": {"name": "json"}},
- "exclude_models": [],
- "disable_fallback": False,
- "exclude_providers": [],
- }
- messages = full_request.get("messages", [])
-
- token_count = token_counter(model=model, messages=messages)
- print(token_count)
-
-
-@pytest.mark.parametrize(
- "model, base_model, input_tokens, user_max_tokens, expected_value",
- [
- ("random-model", "random-model", 1024, 1024, 1024),
- ("gpt-3.5-turbo", "gpt-3.5-turbo", 4000, 5000, 4096), # model max output = 4096
- ],
-)
-def test_get_modified_max_tokens(
- model, base_model, input_tokens, user_max_tokens, expected_value
-):
- """
- - Test when max_output is not known => expect user_max_tokens
- - Test when max_output == max_input,
- - input > max_output, no max_tokens => expect None
- - input + max_tokens > max_output => expect remainder
- - input + max_tokens < max_output => expect max_tokens
- - Test when max_tokens > max_output => expect max_output
- """
- args = locals()
- import litellm
-
- litellm.token_counter = MagicMock()
-
- def _mock_token_counter(*args, **kwargs):
- return input_tokens
-
- litellm.token_counter.side_effect = _mock_token_counter
- print(f"_mock_token_counter: {_mock_token_counter()}")
- messages = [{"role": "user", "content": "Hello world!"}]
-
- calculated_value = get_modified_max_tokens(
- model=model,
- base_model=base_model,
- messages=messages,
- user_max_tokens=user_max_tokens,
- buffer_perc=0,
- buffer_num=0,
- )
-
- if expected_value is None:
- assert calculated_value is None
- else:
- assert (
- calculated_value == expected_value
- ), "Got={}, Expected={}, Params={}".format(
- calculated_value, expected_value, args
- )
-
-
-def test_empty_tools():
- messages = [{"role": "user", "content": "hey, how's it going?", "tool_calls": None}]
-
- result = token_counter(
- messages=messages,
- )
-
- print(result)
-
-
-@pytest.mark.skip(
- reason="Skipping this test temporarily because it relies on a function being called that I am removing."
-)
-def test_gpt_4o_token_counter():
- with patch.object(
- litellm.utils, "openai_token_counter", new=MagicMock()
- ) as mock_client:
- token_counter(
- model="gpt-4o-2024-05-13", messages=[{"role": "user", "content": "Hey!"}]
- )
-
- mock_client.assert_called()
-
-
-@pytest.mark.parametrize(
- "img_url",
- [
- "https://example.com/test-image.png",
- "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAL0AAAC9CAMAAADRCYwCAAAAh1BMVEX///8AAAD8/Pz5+fkEBAT39/cJCQn09PRNTU3y8vIMDAwzMzPe3t7v7+8QEBCOjo7FxcXR0dHn5+elpaWGhoYYGBivr686OjocHBy0tLQtLS1TU1PY2Ni6urpaWlpERER3d3ecnJxoaGiUlJRiYmIlJSU4ODhBQUFycnKAgIDBwcFnZ2chISE7EjuwAAAI/UlEQVR4nO1caXfiOgz1bhJIyAJhX1JoSzv8/9/3LNlpYd4rhX6o4/N8Z2lKM2cURZau5JsQEhERERERERERERERERERERHx/wBjhDPC3OGN8+Cc5JeMuheaETSdO8vZFyCScHtmz2CsktoeMn7rLM1u3h0PMAEhyYX7v/Q9wQvoGdB0hlbzm45lEq/wd6y6G9aezvBk9AXwp1r3LHJIRsh6s2maxaJpmvqgvkC7WFS3loUnaFJtKRVUCEoV/RpCnHRvAsesVQ1hw+vd7Mpo+424tLs72NplkvQgcdrsvXkW/zJWqH/fA0FT84M/xnQJt4to3+ZLuanbM6X5lfXKHosO9COgREqpCR5i86pf2zPS7j9tTj+9nO7bQz3+xGEyGW9zqgQ1tyQ/VsxEDvce/4dcUPNb5OD9yXvR4Z2QisuP0xiGWPnemgugU5q/troHhGEjIF5sTOyW648aC0TssuaaCEsYEIkGzjWXOp3A0vVsf6kgRyqaDk+T7DIVWrb58b2tT5xpUucKwodOD/5LbrZC1ws6YSaBZJ/8xlh+XZSYXaMJ2ezNqjB3IPXuehPcx2U6b4t1dS/xNdFzguUt8ie7arnPeyCZroxLHzGgGdqVcspwafizPWEXBee+9G1OaufGdvNng/9C+gwgZ3PH3r87G6zXTZ5D5De2G2DeFoANXfbACkT+fxBQ22YFsTTJF9hjFVO6VbqxZXko4WJ8s52P4PnuxO5KRzu0/hlix1ySt8iXjgaQ+4IHPA9nVzNkdduM9LFT/Aacj4FtKrHA7iAw602Vnht6R8Vq1IOS+wNMKLYqayAYfRuufQPGeGb7sZogQQoLZrGPgZ6KoYn70Iw30O92BNEDpvwouCFn6wH2uS+EhRb3WF/HObZk3HuxfRQM3Y/Of/VH0n4MKNHZDiZvO9+m/ABALfkOcuar/7nOo7B95ACGVAFaz4jMiJwJhdaHBkySmzlGTu82gr6FSTik2kJvLnY9nOd/D90qcH268m3I/cgI1xg1maE5CuZYaWLH+UHANCIck0yt7Mx5zBm5vVHXHwChsZ35kKqUpmo5Svq5/fzfAI5g2vDtFPYo1HiEA85QrDeGm9g//LG7K0scO3sdpj2CBDgCa+0OFs0bkvVgnnM/QBDwllOMm+cN7vMSHlB7Uu4haHKaTwgGkv8tlK+hP8fzmFuK/RQTpaLPWvbd58yWIo66HHM0OsPoPhVqmtaEVL7N+wYcTLTbb0DLdgp23Eyy2VYJ2N7bkLFAAibtoLPe5sLt6Oa2bvU+zyeMa8wrixO0gRTn9tO9NCSThTLGqcqtsDvphlfmx/cPBZVvw24jg1LE2lPuEo35Mhi58U0I/Ga8n5w+NS8i34MAQLos5B1u0xL1ZvCVYVRw/Fs2q53KLaXJMWwOZZ/4MPYV19bAHmgGDKB6f01xoeJKFbl63q9J34KdaVNPJWztQyRkzA3KNs1AdAEDowMxh10emXTCx75CkurtbY/ZpdNDGdsn2UcHKHsQ8Ai3WZi48IfkvtjOhsLpuIRSKZTX9FA4o+0d6o/zOWqQzVJMynL9NsxhSJOaourq6nBVQBueMSyubsX2xHrmuABZN2Ns9jr5nwLFlLF/2R6atjW/67Yd11YQ1Z+kA9Zk9dPTM/o6dVo6HHVgC0JR8oUfmI93T9u3gvTG94bAH02Y5xeqRcjuwnKCK6Q2+ajl8KXJ3GSh22P3Zfx6S+n008ROhJn+JRIUVu6o7OXl8w1SeyhuqNDwNI7SjbK08QrqPxS95jy4G7nCXVq6G3HNu0LtK5J0e226CfC005WKK9sVvfxI0eUbcnzutfhWe3rpZHM0nZ/ny/N8tanKYlQ6VEW5Xuym8yV1zZX58vwGhZp/5tFfhybZabdbrQYOs8F+xEhmPsb0/nki6kIyVvzZzUASiOrTfF+Sj9bXC7DoJxeiV8tjQL6loSd0yCx7YyB6rPdLx31U2qCG3F/oXIuDuqd6LFO+4DNIJuxFZqSsU0ea88avovFnWKRYFYRQDfCfcGaBCLn4M4A1ntJ5E57vicwqq2enaZEF5nokCYu9TbKqCC5yCDfL+GhLxT4w4xEJs+anqgou8DOY2q8FMryjb2MehC1dRJ9s4g9NXeTwPkWON4RH+FhIe0AWR/S9ekvQ+t70XHeimGF78LzuU7d7PwrswdIG2VpgF8C53qVQsTDtBJc4CdnkQPbnZY9mbPdDFra3PCXBBQ5QBn2aQqtyhvlyYM4Hb2/mdhsxCUen04GZVvIJZw5PAamMOmjzq8Q+dzAKLXDQ3RUZItWsg4t7W2DP+JDrJDymoMH7E5zQtuEpG03GTIjGCW3LQqOYEsXgFc78x76NeRwY6SNM+IfQoh6myJKRBIcLYxZcwscJ/gI2isTBty2Po9IkYzP0/SS4hGlxRjFAG5z1Jt1LckiB57yWvo35EaolbvA+6fBa24xodL2YjsPpTnj3JgJOqhcgOeLVsYYwoK0wjY+m1D3rGc40CukkaHnkEjarlXrF1B9M6ECQ6Ow0V7R7N4G3LfOHAXtymoyXOb4QhaYHJ/gNBJUkxclpSs7DNcgWWDDmM7Ke5MJpGuioe7w5EOvfTunUKRzOh7G2ylL+6ynHrD54oQO3//cN3yVO+5qMVsPZq0CZIOx4TlcJ8+Vz7V5waL+7WekzUpRFMTnnTlSCq3X5usi8qmIleW/rit1+oQZn1WGSU/sKBYEqMNh1mBOc6PhK8yCfKHdUNQk8o/G19ZPTs5MYfai+DLs5vmee37zEyyH48WW3XA6Xw6+Az8lMhci7N/KleToo7PtTKm+RA887Kqc6E9dyqL/QPTugzMHLbLZtJKqKLFfzVWRNJ63c+95uWT/F7R0U5dDVvuS409AJXhJvD0EwWaWdW8UN11u/7+umaYjT8mJtzZwP/MD4r57fihiHlC5fylHfaqnJdro+Dr7DajvO+vi2EwyD70s8nCH71nzIO1l5Zl+v1DMCb5ebvCMkGHvobXy/hPumGLyX0218/3RyD1GRLOuf9u/OGQyDmto32yMiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIv7GP8YjWPR/czH2AAAAAElFTkSuQmCC",
- ],
-)
-def test_img_url_token_counter(img_url, monkeypatch):
- """
- Verify get_image_dimensions returns valid (width, height) for both an
- HTTPS URL and a base64 data URI. The HTTPS branch is exercised with a
- mocked HTTP fetch so the test is hermetic - it can't break when a
- third-party image URL goes away.
- """
- import base64
- from litellm.litellm_core_utils.token_counter import get_image_dimensions
-
- # Minimal valid 1x1 PNG, served by the mocked safe_get for the URL case.
- _tiny_png = base64.b64decode(
- "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
- )
-
- if img_url.startswith(("http://", "https://")):
-
- class _FakeResponse:
- headers = {"Content-Length": str(len(_tiny_png))}
-
- def read(self):
- return _tiny_png
-
- monkeypatch.setattr(
- "litellm.litellm_core_utils.token_counter.safe_get",
- lambda client, url, **kw: _FakeResponse(),
- )
-
- width, height = get_image_dimensions(data=img_url)
-
- print(width, height)
-
- assert width is not None
- assert height is not None
-
-
-def test_token_encode_disallowed_special():
- encode(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>")
- token_counter(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>")
-
-
-def test_token_counter():
- try:
- messages = [{"role": "user", "content": "hi how are you what time is it"}]
- tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
- print("gpt-35-turbo")
- print(tokens)
- assert tokens > 0
-
- tokens = token_counter(model="claude-2", messages=messages)
- print("claude-2")
- print(tokens)
- assert tokens > 0
-
- tokens = token_counter(model="gemini/chat-bison", messages=messages)
- print("gemini/chat-bison")
- print(tokens)
- assert tokens > 0
-
- tokens = token_counter(model="ollama/llama2", messages=messages)
- print("ollama/llama2")
- print(tokens)
- assert tokens > 0
-
- tokens = token_counter(model="anthropic.claude-instant-v1", messages=messages)
- print("anthropic.claude-instant-v1")
- print(tokens)
- assert tokens > 0
- except Exception as e:
- pytest.fail(f"Error occurred: {e}")
-
-
-import unittest
-
-from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding
-
-# Clear the cache at module load to ensure clean state
-_load_huggingface_tokenizer.cache_clear()
-
-
-class TestTokenizerSelection(unittest.TestCase):
- def setUp(self):
- """Clear the LRU cache before each test method.
-
- The HuggingFace tokenizers behind _select_tokenizer_helper are cached with
- @lru_cache, which can cause cache hits from previous tests when running with
- --dist=loadscope (tests from same file run on same worker).
- """
- _load_huggingface_tokenizer.cache_clear()
-
- @patch("litellm.utils.tokenizer_dispatch.from_pretrained")
- def test_llama3_tokenizer_api_failure(self, mock_from_pretrained):
- # Setup mock to raise an error
- mock_from_pretrained.side_effect = Exception("Failed to load tokenizer")
-
- # Test with llama-3 model
- result = _select_tokenizer_helper("llama-3-7b")
-
- # Verify the attempt to load Llama-3 tokenizer
- mock_from_pretrained.assert_called_once_with("Xenova/llama-3-tokenizer")
-
- # Verify fallback to OpenAI tokenizer
- self.assertEqual(result["type"], "openai_tokenizer")
- self.assertEqual(result["tokenizer"], encoding)
-
- @patch("litellm.utils.tokenizer_dispatch.from_pretrained")
- def test_cohere_tokenizer_api_failure(self, mock_from_pretrained):
- # Setup mock to raise an error
- mock_from_pretrained.side_effect = Exception("Failed to load tokenizer")
-
- # Add Cohere model to the list for testing
- litellm.cohere_models = ["command-r-v1"]
-
- # Test with Cohere model
- result = _select_tokenizer_helper("command-r-v1")
-
- # Verify the attempt to load Cohere tokenizer
- mock_from_pretrained.assert_called_once_with(
- "Xenova/c4ai-command-r-v01-tokenizer"
- )
-
- # Verify fallback to OpenAI tokenizer
- self.assertEqual(result["type"], "openai_tokenizer")
- self.assertEqual(result["tokenizer"], encoding)
-
- @patch("litellm.utils.tokenizer_dispatch.anthropic")
- def test_claude_tokenizer_api_failure(self, mock_anthropic):
- # Setup mock to raise an error
- mock_anthropic.side_effect = Exception("Failed to load tokenizer")
-
- # Add Claude model to the list for testing
- litellm.anthropic_models = ["claude-2"]
-
- # Test with Claude model
- result = _select_tokenizer_helper("claude-2")
-
- # Verify the attempt to load Claude tokenizer
- mock_anthropic.assert_called_once_with()
-
- # Verify fallback to OpenAI tokenizer
- self.assertEqual(result["type"], "openai_tokenizer")
- self.assertEqual(result["tokenizer"], encoding)
-
- @patch("litellm.utils.tokenizer_dispatch.from_pretrained")
- def test_llama2_tokenizer_api_failure(self, mock_from_pretrained):
- # Setup mock to raise an error
- mock_from_pretrained.side_effect = Exception("Failed to load tokenizer")
-
- # Test with Llama-2 model
- result = _select_tokenizer_helper("llama-2-7b")
-
- # Verify the attempt to load Llama-2 tokenizer
- mock_from_pretrained.assert_called_once_with(
- "hf-internal-testing/llama-tokenizer"
- )
-
- # Verify fallback to OpenAI tokenizer
- self.assertEqual(result["type"], "openai_tokenizer")
- self.assertEqual(result["tokenizer"], encoding)
-
- @patch("litellm.utils._return_huggingface_tokenizer")
- def test_disable_hf_tokenizer_download(self, mock_return_huggingface_tokenizer):
- monkeypatch = pytest.MonkeyPatch()
- monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
- try:
- result = _select_tokenizer_helper("grok-32r22r")
- mock_return_huggingface_tokenizer.assert_not_called()
- assert result["type"] == "openai_tokenizer"
- assert result["tokenizer"] == encoding
- finally:
- monkeypatch.undo()
-
-
-@pytest.mark.parametrize(
- "model",
- [
- "gpt-4o",
- "claude-3-opus-20240229",
- ],
-)
-@pytest.mark.parametrize(
- "messages",
- [
- [
- {
- "role": "user",
- "content": [
- {
- "type": "text",
- "text": "These are some sample images from a movie. Based on these images, what do you think the tone of the movie is?",
- },
- {
- "type": "text",
- "image_url": {
- "url": "https://gratisography.com/wp-content/uploads/2024/11/gratisography-augmented-reality-800x525.jpg",
- "detail": "high",
- },
- },
- ],
- }
- ],
- ],
-)
-def test_bad_input_token_counter(model, messages):
- """
- Safely handle bad input for token counter.
- """
- token_counter(
- model=model,
- messages=messages,
- default_token_count=1000,
- )
-
-
-def test_token_counter_with_anthropic_tool_use():
- """
- Test that _count_anthropic_content() correctly handles tool_use blocks.
-
- Validates that:
- - 'name' field is counted (string)
- - 'input' field is counted (dict serialized to string)
- - Metadata fields ('type', 'id') are skipped
- """
- messages = [
- {"role": "user", "content": "What's the weather in San Francisco?"},
- {
- "role": "assistant",
- "content": [
- {"type": "text", "text": "I'll check the weather for you."},
- {
- "type": "tool_use",
- "id": "toolu_01234567890", # Should be skipped
- "name": "get_weather", # Should be counted
- "input": { # Should be counted (serialized)
- "location": "San Francisco, CA",
- "unit": "fahrenheit",
- },
- },
- ],
- },
- ]
-
- tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
- assert tokens > 0, f"Expected positive token count, got {tokens}"
- # Should count: user message + "I'll check" text + "get_weather" name + input dict
- assert (
- tokens > 15
- ), f"Expected reasonable token count for message with tool_use, got {tokens}"
-
-
-def test_token_counter_with_anthropic_tool_result():
- """
- Test that _count_anthropic_content() correctly handles tool_result blocks.
-
- Validates that:
- - 'content' field (when string) is counted
- - Metadata fields ('type', 'tool_use_id') are skipped
- - Full conversation with tool_use → tool_result flow works
- """
- messages = [
- {"role": "user", "content": "What's the weather in San Francisco?"},
- {
- "role": "assistant",
- "content": [
- {
- "type": "tool_use",
- "id": "toolu_01234567890",
- "name": "get_weather",
- "input": {"location": "San Francisco, CA"},
- }
- ],
- },
- {
- "role": "user",
- "content": [
- {
- "type": "tool_result",
- "tool_use_id": "toolu_01234567890", # Should be skipped
- "content": "The weather in San Francisco is 65°F and sunny.", # Should be counted
- }
- ],
- },
- ]
-
- tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
- assert tokens > 0, f"Expected positive token count, got {tokens}"
- assert (
- tokens > 25
- ), f"Expected reasonable token count for conversation with tool_result, got {tokens}"
-
-
-def test_token_counter_with_nested_tool_result():
- """
- Test that _count_anthropic_content() recursively handles nested content lists.
-
- Validates that:
- - tool_result with 'content' as a list (not string) is handled
- - Nested content blocks are recursively counted via _count_content_list()
- - TypedDict inference correctly identifies list fields
- """
- messages = [
- {
- "role": "user",
- "content": [
- {
- "type": "tool_result",
- "tool_use_id": "toolu_01234567890",
- "content": [ # Nested list - should recursively count
- {
- "type": "text",
- "text": "The weather in San Francisco is 65°F and sunny.",
- },
- {"type": "text", "text": "UV index is moderate."},
- ],
- }
- ],
- }
- ]
-
- tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
- assert tokens > 0, f"Expected positive token count, got {tokens}"
- # Should count both nested text blocks
- assert (
- tokens > 15
- ), f"Expected reasonable token count for nested tool_result, got {tokens}"
-
-
-def test_token_counter_tool_use_and_result_combined():
- """
- Test dynamic field inference with multiple tool_use and tool_result blocks.
-
- Validates that:
- - Multiple tool_use blocks in same message are handled
- - Multiple tool_result blocks in same message are handled
- - skip_fields correctly filters metadata across all blocks
- - Full realistic conversation flow works end-to-end
- """
- messages = [
- {
- "role": "user",
- "content": "What's the weather in San Francisco and New York?",
- },
- {
- "role": "assistant",
- "content": [
- {
- "type": "text",
- "text": "I'll check the weather in both cities for you.",
- },
- {
- "type": "tool_use",
- "id": "toolu_01A",
- "name": "get_weather",
- "input": {"location": "San Francisco, CA"},
- },
- {
- "type": "tool_use",
- "id": "toolu_01B",
- "name": "get_weather",
- "input": {"location": "New York, NY"},
- },
- ],
- },
- {
- "role": "user",
- "content": [
- {
- "type": "tool_result",
- "tool_use_id": "toolu_01A",
- "content": "San Francisco: 65°F, sunny",
- },
- {
- "type": "tool_result",
- "tool_use_id": "toolu_01B",
- "content": "New York: 45°F, cloudy",
- },
- ],
- },
- {
- "role": "assistant",
- "content": "The weather in San Francisco is 65°F and sunny, while New York is cooler at 45°F and cloudy.",
- },
- ]
-
- tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
- assert tokens > 0, f"Expected positive token count, got {tokens}"
- # Should count all text, tool names, inputs, and results
- assert (
- tokens > 60
- ), f"Expected substantial token count for full tool conversation, got {tokens}"
-
-
-def test_token_counter_with_image_url():
- """
- Test that _count_image_tokens() correctly handles image_url content blocks.
-
- Validates that:
- - image_url as dict with 'url' and 'detail' is handled
- - image_url as string is handled
- - 'detail' field validation works ('low', 'high', 'auto')
- - calculate_img_tokens is called with correct parameters
- """
- # Test with dict format (detail: low)
- messages_dict = [
- {
- "role": "user",
- "content": [
- {"type": "text", "text": "What's in this image?"},
- {
- "type": "image_url",
- "image_url": {
- "url": "https://example.com/image.jpg",
- "detail": "low", # Should use low token count (85 base tokens)
- },
- },
- ],
- }
- ]
-
- tokens_dict = token_counter(
- model="gpt-3.5-turbo",
- messages=messages_dict,
- use_default_image_token_count=True, # Avoid actual HTTP request
- )
- assert tokens_dict > 0, f"Expected positive token count, got {tokens_dict}"
- assert tokens_dict > 85, f"Expected at least base image tokens, got {tokens_dict}"
-
- # Test with string format (defaults to auto/low)
- messages_str = [
- {
- "role": "user",
- "content": [
- {
- "type": "image_url",
- "image_url": "https://example.com/image.jpg", # String format
- }
- ],
- }
- ]
-
- tokens_str = token_counter(
- model="gpt-3.5-turbo", messages=messages_str, use_default_image_token_count=True
- )
- assert (
- tokens_str > 0
- ), f"Expected positive token count for string image_url, got {tokens_str}"
-
- # Test invalid detail value raises error
- messages_invalid = [
- {
- "role": "user",
- "content": [
- {
- "type": "image_url",
- "image_url": {
- "url": "https://example.com/image.jpg",
- "detail": "invalid", # Should raise ValueError
- },
- }
- ],
- }
- ]
-
- with pytest.raises(ValueError, match="Invalid detail value") as exc_info:
- token_counter(model="gpt-3.5-turbo", messages=messages_invalid)
- e = exc_info.value
- assert "Invalid detail value" in str(
- e
- ), f"Expected detail validation error, got: {e}"
-
-
-def test_token_counter_with_thinking_content():
- """
- Test that _count_content_list() correctly handles Claude's extended thinking content blocks.
-
- Validates that:
- - 'thinking' content type is recognized and counted
- - 'thinking' text field is counted
- - 'signature' field is skipped (opaque signature blob)
- - Full conversation with thinking blocks works
- """
- messages = [
- {
- "role": "user",
- "content": [
- {
- "type": "text",
- "text": "Analyze this complex problem: who came first, chicken or egg",
- }
- ],
- },
- {
- "role": "assistant",
- "content": [
- {
- "type": "thinking",
- "thinking": "This is actually a fascinating question that touches on philosophy, biology, and semantics. Let me break this down: The egg came first from an evolutionary biology perspective.",
- "signature": "EqcLCkYICxgCKkCrqu6lP...", # Should be skipped
- },
- {
- "type": "text",
- "text": "# The Chicken-or-Egg Question: A Multi-Layered Answer\n\n## **The Short Answer: The Egg Came First**",
- },
- ],
- },
- {"role": "user", "content": [{"type": "text", "text": "Thanks"}]},
- ]
-
- tokens = token_counter(
- model="anthropic/claude-sonnet-4-5-20250929", messages=messages
- )
- assert tokens > 0, f"Expected positive token count, got {tokens}"
- # Should count: user message + thinking text + response text + "Thanks"
- # The thinking text alone is ~30 tokens, plus other content should be > 50 total
- assert (
- tokens > 50
- ), f"Expected substantial token count for message with thinking, got {tokens}"
-
- # Test that thinking block without 'thinking' field doesn't crash (edge case)
- messages_no_thinking = [
- {
- "role": "assistant",
- "content": [
- {
- "type": "thinking",
- # No 'thinking' field - should count as 0 tokens
- "signature": "EqcLCkYICxgCKkCrqu6lP...",
- },
- {"type": "text", "text": "Response"},
- ],
- }
- ]
-
- tokens_no_thinking = token_counter(
- model="anthropic/claude-sonnet-4-5-20250929", messages=messages_no_thinking
- )
- assert (
- tokens_no_thinking > 0
- ), f"Expected positive token count even with empty thinking, got {tokens_no_thinking}"
- # Should only count "Response" and message overhead
- assert (
- tokens_no_thinking < 15
- ), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}"
-
-
-
-def test_token_counter_with_redacted_thinking_content():
- """
- A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in
- for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking
- block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the
- prompt_caching pre-call check stop pinning the deployment that held the cached prefix.
- """
- model = "anthropic/claude-sonnet-4-5-20250929"
- reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."}
- redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30}
- user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]}
- follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]}
-
- without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up]
- with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up]
-
- assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block)
-
-def test_token_counter_with_tool_reference_block():
- """
- Regression test: a message containing an Anthropic tool-search
- `tool_reference` content block must NOT raise.
-
- Before the fix, token_counter raised
- `Invalid content item type: tool_reference`. On the streaming
- anthropic_messages proxy path this nulled response_cost and caused the
- SpendLogs row to be dropped, silently undercounting cost. token_counter
- must instead count the referenced tool name and return a positive count.
- """
- messages = [
- {
- "role": "assistant",
- "content": [
- {"type": "text", "text": "Let me look up the right tool."},
- {"type": "tool_reference", "tool_name": "search_knowledge_base"},
- ],
- }
- ]
-
- # Must not raise, and must produce a positive token count.
- tokens = token_counter_new(
- model="anthropic/claude-sonnet-4-5-20250929", messages=messages
- )
- assert tokens > 0, f"Expected positive token count, got {tokens}"
-
- # A tool_reference with no/empty tool_name must also be handled gracefully.
- messages_empty = [
- {
- "role": "assistant",
- "content": [{"type": "tool_reference", "tool_name": ""}],
- }
- ]
- tokens_empty = token_counter_new(
- model="anthropic/claude-sonnet-4-5-20250929", messages=messages_empty
- )
- assert tokens_empty >= 0
-
-
-def test_count_content_list_rejects_unknown_type():
- """
- An unrecognized content block type must raise, and the error message must
- enumerate the supported types (including `tool_reference`). This pins the
- catch-all contract so a future block type isn't silently dropped.
- """
- from litellm.litellm_core_utils.token_counter import _count_content_list
-
- with pytest.raises(ValueError, match='Error getting number of tokens from content list: Invalid') as exc_info:
- _count_content_list(
- count_function=len,
- content_list=[{"type": "totally_unknown_block"}],
- use_default_image_token_count=False,
- default_token_count=None,
- )
-
- message = str(exc_info.value)
- assert "Invalid content item type: totally_unknown_block" in message
- assert "tool_reference" in message
-
-
-@pytest.mark.parametrize(
- "source",
- [
- {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="},
- {"type": "url", "url": "https://example.com/image.png"},
- {"type": "file", "file_id": "file-abc123"},
- ],
- ids=["base64", "url", "file"],
-)
-def test_token_counter_with_anthropic_image_block(source: dict[str, str]):
- """Anthropic `image` blocks must count for every source variant, not raise `Invalid content item type` (which the router's context-window pre-call check swallows into an unfiltered dispatch)."""
- from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT
-
- messages = [
- {
- "role": "user",
- "content": [
- {"type": "text", "text": "What is in this image?"},
- {"type": "image", "source": source},
- ],
- }
- ]
-
- tokens = token_counter(
- model="anthropic/claude-sonnet-4-5-20250929",
- messages=messages,
- use_default_image_token_count=True,
- )
- assert tokens > DEFAULT_IMAGE_TOKEN_COUNT, (
- f"Expected the image block to contribute tokens, got {tokens}"
- )
-
-
-def test_anthropic_image_block_matches_equivalent_image_url():
- """An Anthropic `image` block prices identically to the OpenAI `image_url` carrying the same bytes."""
- anthropic_messages = [
- {
- "role": "user",
- "content": [
- {
- "type": "image",
- "source": {
- "type": "base64",
- "media_type": "image/png",
- "data": "iVBORw0KGgo=",
- },
- }
- ],
- }
- ]
- openai_messages = [
- {
- "role": "user",
- "content": [
- {
- "type": "image_url",
- "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
- }
- ],
- }
- ]
-
- anthropic_tokens = token_counter(
- model="anthropic/claude-sonnet-4-5-20250929", messages=anthropic_messages
- )
- openai_tokens = token_counter(
- model="anthropic/claude-sonnet-4-5-20250929", messages=openai_messages
- )
- assert anthropic_tokens == openai_tokens
-
-
-def test_anthropic_image_block_nested_in_tool_result():
- """An `image` block nested in a `tool_result.content` list is counted through the same recursion."""
- messages = [
- {
- "role": "user",
- "content": [
- {
- "type": "tool_result",
- "tool_use_id": "toolu_01",
- "content": [
- {
- "type": "image",
- "source": {
- "type": "base64",
- "media_type": "image/png",
- "data": "iVBORw0KGgo=",
- },
- }
- ],
- }
- ],
- }
- ]
-
- tokens = token_counter(
- model="anthropic/claude-sonnet-4-5-20250929",
- messages=messages,
- use_default_image_token_count=True,
- )
- assert tokens > 0
-
-
-@pytest.mark.parametrize(
- ("source", "expected"),
- [
- ({"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQ"}, "data:image/jpeg;base64,/9j/4AAQ"),
- ({"type": "url", "url": "https://example.com/image.png"}, "https://example.com/image.png"),
- ({"type": "file", "file_id": "file-abc123"}, ""),
- ],
- ids=["base64", "url", "file"],
-)
-def test_anthropic_image_source_resolves_to_what_the_image_pricer_reads(source: dict[str, str], expected: str):
- """base64 sources become a data URI, url sources pass through, file sources resolve to an empty string."""
- from litellm.litellm_core_utils.token_counter import _anthropic_image_source_data
-
- assert _anthropic_image_source_data(source) == expected
-
-
-def test_anthropic_image_block_with_empty_base64_data():
- """A base64 source with empty `data` prices as an image rather than raising."""
- from litellm.litellm_core_utils.token_counter import _count_content_list
-
- tokens = _count_content_list(
- count_function=len,
- content_list=[
- {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}}
- ],
- use_default_image_token_count=False,
- default_token_count=None,
- )
- assert tokens > 0
-
-
-def test_anthropic_image_block_without_source_raises():
- """An `image` block with no `source` raises, matching the OpenAI `image_url`-without-`url` behavior."""
- from litellm.litellm_core_utils.token_counter import _count_content_list
-
- with pytest.raises(ValueError, match="Error getting number of tokens from content list"):
- _count_content_list(
- count_function=len,
- content_list=[{"type": "image"}],
- use_default_image_token_count=False,
- default_token_count=None,
- )
-
- # ... and `default_token_count`, the caller's opt-out from raising, still wins.
- assert (
- _count_content_list(
- count_function=len,
- content_list=[{"type": "image"}],
- use_default_image_token_count=False,
- default_token_count=7,
- )
- == 7
- )
-
-
-def _count_user_content(content: list[dict]) -> int:
- from litellm.litellm_core_utils.token_counter import token_counter
-
- return token_counter(
- model="anthropic/claude-fable-5",
- messages=[{"role": "user", "content": content}],
- use_default_image_token_count=True,
- )
-
-
-@pytest.mark.parametrize(
- "source",
- [
- {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
- {"type": "url", "url": "https://example.com/report.pdf"},
- {"type": "file", "file_id": "file-abc123"},
- ],
- ids=["base64", "url", "file"],
-)
-def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]):
- """A `document` whose bytes can't be tokenized locally is priced like an `image`, not raised on."""
- prompt = {"type": "text", "text": "Summarize this file."}
-
- assert _count_user_content([prompt, {"type": "document", "source": source}]) == _count_user_content(
- [prompt, {"type": "image", "source": source}]
- )
-
-
-def test_anthropic_document_block_text_sources_count_their_text():
- """`text` and `content` document sources count the text they carry, as inline text blocks would."""
- prompt = {"type": "text", "text": "Summarize this file."}
- body = {"type": "text", "text": "Revenue grew eleven percent while churn fell to two percent."}
- picture = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}}
-
- text_source = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": body["text"]}}
- assert _count_user_content([prompt, text_source]) == _count_user_content([prompt, body])
-
- string_content = {"type": "document", "source": {"type": "content", "content": body["text"]}}
- assert _count_user_content([prompt, string_content]) == _count_user_content([prompt, body])
-
- block_content = {"type": "document", "source": {"type": "content", "content": [body, picture]}}
- assert _count_user_content([prompt, block_content]) == _count_user_content([prompt, body, picture])
-
-
-def test_anthropic_document_title_and_context_add_their_tokens():
- prompt = {"type": "text", "text": "Summarize this file."}
- source = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}
- described = {"type": "document", "source": source, "title": "Q3 board packet", "context": "Shared by finance"}
-
- assert _count_user_content([prompt, described]) == _count_user_content(
- [
- prompt,
- {"type": "text", "text": "Q3 board packet"},
- {"type": "text", "text": "Shared by finance"},
- {"type": "document", "source": source},
- ]
- )
-
-
-def test_openai_file_block_prices_like_the_equivalent_anthropic_document():
- """An inline `file` is a `document` in the chat-completions dialect, so it must price identically, not raise.
-
- Before the fix `file` was missing from the content-block match even though `ChatCompletionFileObject`
- is in the union this counter accepts, so every local count of a Responses `input_file` raised
- `Invalid content item type: file` and surfaced as a 500 on /v1/responses/input_tokens.
- """
- prompt = {"type": "text", "text": "Summarize this file."}
- inline_file = {
- "type": "file",
- "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"},
- }
- document = {
- "type": "document",
- "title": "report.pdf",
- "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
- }
-
- assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document])
- assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt])
-
-
-def test_openai_file_block_without_inline_bytes_counts_what_it_carries():
- """A `file` block naming an uploaded file has no bytes to price, so it adds only the filename's tokens."""
- prompt = {"type": "text", "text": "Summarize this file."}
-
- by_id = {"type": "file", "file": {"file_id": "file-abc123"}}
- assert _count_user_content([prompt, by_id]) == _count_user_content([prompt])
-
- named = {"type": "file", "file": {"file_id": "file-abc123", "filename": "report.pdf"}}
- assert _count_user_content([prompt, named]) == _count_user_content(
- [prompt, {"type": "text", "text": "report.pdf"}]
- )
-
-
-def _png_data_url(width: int, height: int) -> str:
- ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big")
- return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode()
-
-
-@pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)])
-def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None:
- assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound()
-
-
-def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None:
- assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound()
- assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound()
diff --git a/tests/test_litellm/litellm_core_utils/test_tokenizer.py b/tests/test_litellm/litellm_core_utils/test_tokenizer.py
index aa4a0fc6a1c..2171044970c 100644
--- a/tests/test_litellm/litellm_core_utils/test_tokenizer.py
+++ b/tests/test_litellm/litellm_core_utils/test_tokenizer.py
@@ -1,403 +1,20 @@
-import copy
-import os
-import pickle
-import subprocess
-import sys
-from pathlib import Path
-from typing import Final, Literal
-
import pytest
-import tiktoken
-from tokenizers import Tokenizer as ReferenceTokenizer
-import litellm
-from litellm.caching._embedding_router import truncate_embedding_input
-from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding
-from litellm.utils import claude_json_str
-from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
-
-
-@pytest.mark.parametrize(
- "name", ("cl100k_base", "o200k_base", "p50k_base", "p50k_edit", "r50k_base", "gpt2", "o200k_harmony")
-)
-@pytest.mark.parametrize(
- "text", ("hello world", "café æ¼¢å— ðŸ™‚", "", "a\ud800b", "\ud83d\ude42", "🙂\ud83d\ude42\udfff", " " * 64)
+from tests.unit.litellm_core_utils.test_tokenizer import (
+ UNICODE_TEXTS,
+ assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface,
+ assert_openai_encoding_matches_python,
)
+
+NETWORK_ENCODINGS = ("r50k_base", "gpt2")
+
+
+@pytest.mark.parametrize("name", NETWORK_ENCODINGS)
+@pytest.mark.parametrize("text", UNICODE_TEXTS)
def test_openai_encoding_matches_python_unicode_and_batches(name: str, text: str) -> None:
- reference: Final = tiktoken.get_encoding(name)
- encoding: Final = OpenAIEncoding.from_tiktoken(name)
- expected: Final = reference.encode(text)
-
- assert encoding.encode(text) == expected
- assert encoding.count(text) == len(expected)
- assert encoding.encode_batch([text], num_threads=2) == reference.encode_batch([text], num_threads=2)
- assert encoding.encode_ordinary_batch([text]) == reference.encode_ordinary_batch([text])
- assert encoding.decode_batch([expected]) == reference.decode_batch([expected])
- assert encoding.decode_bytes_batch([expected]) == reference.decode_bytes_batch([expected])
+ assert_openai_encoding_matches_python(name, text)
-@pytest.mark.parametrize("allowed", (frozenset(), frozenset({"<|endoftext|>"}), "all"))
-@pytest.mark.parametrize("disallowed", (frozenset(), frozenset({"<|fim_prefix|>"}), "all"))
-def test_openai_special_token_options_match_python(
- allowed: frozenset[str] | Literal["all"], disallowed: frozenset[str] | Literal["all"]
-) -> None:
- reference: Final = tiktoken.get_encoding("cl100k_base")
- encoding: Final = OpenAIEncoding.from_tiktoken(reference.name)
- text: Final = "hello<|endoftext|><|fim_prefix|>world"
- allowed_set: Final = reference.special_tokens_set if allowed == "all" else allowed
- disallowed_set: Final = reference.special_tokens_set - allowed_set if disallowed == "all" else disallowed
- if any(token in text for token in disallowed_set):
- with pytest.raises(ValueError, match="disallowed special token"):
- encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed)
- return
- assert encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed) == reference.encode(
- text, allowed_special=allowed, disallowed_special=disallowed
- )
- assert encoding.special_tokens_set == reference.special_tokens_set
- assert encoding.eot_token == reference.eot_token
-
-
-@pytest.mark.parametrize("errors", ("replace", "ignore", "backslashreplace", "strict"))
-def test_openai_partial_token_decoding_preserves_error_policy(errors: str) -> None:
- reference: Final = tiktoken.get_encoding("cl100k_base")
- encoding: Final = OpenAIEncoding.from_tiktoken(reference.name)
- tokens: Final = reference.encode("🙂")[:1]
- assert encoding.decode_bytes(tokens) == reference.decode_bytes(tokens)
- if errors == "strict":
- with pytest.raises(UnicodeDecodeError):
- encoding.decode(tokens, errors=errors)
- return
- assert encoding.decode(tokens, errors=errors) == reference.decode(tokens, errors=errors)
- assert encoding.decode_tokens_bytes(tokens) == reference.decode_tokens_bytes(tokens)
-
-
-def test_public_encoding_and_semantic_cache_preserve_truncated_unicode() -> None:
- reference: Final = tiktoken.get_encoding(litellm.encoding.name)
- text: Final = "🙂"
- tokens: Final = reference.encode(text)
-
- assert litellm.encoding.encode(text, disallowed_special=()) == tokens
- assert litellm.encoding.encode_batch([text]) == [tokens]
- assert litellm.decode(tokens=tokens[:1]) == reference.decode(tokens[:1])
- assert truncate_embedding_input(text, "", 1) == reference.decode(tokens[:1])
-
-
-@pytest.mark.parametrize("add_special_tokens", (True, False))
-def test_huggingface_encoding_preserves_result_fields_and_serialization(add_special_tokens: bool) -> None:
- reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
- tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON)
- expected: Final = reference.encode("Hello World", add_special_tokens=add_special_tokens)
- actual: Final = tokenizer.encode("Hello World", add_special_tokens=add_special_tokens)
-
- assert (actual.ids, actual.tokens, actual.type_ids, actual.offsets, actual.word_ids, actual.sequence_ids) == (
- expected.ids,
- expected.tokens,
- expected.type_ids,
- expected.offsets,
- expected.word_ids,
- expected.sequence_ids,
- )
- assert (actual.attention_mask, actual.special_tokens_mask, actual.n_sequences, len(actual)) == (
- expected.attention_mask,
- expected.special_tokens_mask,
- expected.n_sequences,
- len(expected),
- )
- assert copy.deepcopy(actual).ids == expected.ids
- assert pickle.loads(pickle.dumps(actual)).offsets == expected.offsets
- assert tokenizer.decode(actual.ids, skip_special_tokens=False) == reference.decode(
- expected.ids, skip_special_tokens=False
- )
-
-
-def test_huggingface_character_offsets_and_pretokenized_pairs_match_python() -> None:
- reference: Final = ReferenceTokenizer.from_str(claude_json_str)
- tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str)
- text: Final = "café æ¼¢å— ðŸ™‚"
- actual: Final = tokenizer.encode(text)
- expected: Final = reference.encode(text)
-
- assert actual.offsets == expected.offsets
- assert actual.ids == expected.ids
- assert (
- tokenizer.encode(["hello", "world"], ["again"], is_pretokenized=True).ids
- == reference.encode(["hello", "world"], ["again"], is_pretokenized=True).ids
- )
-
-
-def test_huggingface_batches_apply_padding_across_inputs() -> None:
- reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
- reference.enable_padding(pad_id=0, pad_token="[UNK]")
- tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str())
- inputs: Final = ["Hello", ("Hello World", "World")]
- expected: Final = reference.encode_batch(inputs)
- actual: Final = tokenizer.encode_batch(inputs)
- fast: Final = tokenizer.encode_batch_fast(inputs)
-
- assert [(item.ids, item.attention_mask, item.offsets) for item in actual] == [
- (item.ids, item.attention_mask, item.offsets) for item in expected
- ]
- assert [item.ids for item in fast] == [item.ids for item in expected]
- assert tokenizer.decode_batch([item.ids for item in actual]) == reference.decode_batch(
- [item.ids for item in expected]
- )
-
-
-def test_caller_supplied_huggingface_tokenizer_preserves_public_encode_and_count() -> None:
- tokenizer: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
- custom: Final = {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
- expected: Final = tokenizer.encode("Hello World").ids
-
- assert litellm.encode(text="Hello World", custom_tokenizer=custom) == expected
- assert litellm.token_counter(text="Hello World", custom_tokenizer=custom) == len(expected)
- assert litellm.decode(tokens=expected, custom_tokenizer=custom) == "Hello World"
-
-
-def test_caller_supplied_tiktoken_treats_special_spellings_as_text() -> None:
- tokenizer: Final = tiktoken.get_encoding("cl100k_base")
- custom: Final = {"type": "openai_tokenizer", "tokenizer": tokenizer}
- text: Final = "<|endoftext|>"
-
- assert litellm.encode(text=text, custom_tokenizer=custom) == tokenizer.encode(text, disallowed_special=())
-
-
-def test_public_tokenizer_objects_survive_pickle_and_deepcopy(tmp_path: Path) -> None:
- custom: Final = litellm.create_tokenizer(TOKENIZER_JSON)
- tokenizer: Final = custom["tokenizer"]
- path: Final = tmp_path / "tokenizer.json"
- tokenizer.save(str(path))
-
- assert copy.deepcopy(custom)["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids
- assert (
- pickle.loads(pickle.dumps(custom))["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids
- )
- assert HuggingFaceTokenizer.from_file(str(path)).encode("Hello World").ids == tokenizer.encode("Hello World").ids
- assert copy.deepcopy(litellm.encoding).encode("hello") == litellm.encoding.encode("hello")
- assert pickle.loads(pickle.dumps(litellm.encoding)).encode("hello") == litellm.encoding.encode("hello")
-
-
-@pytest.mark.parametrize("offline", ("0", "1"))
-def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None:
- script: Final = """
-import json
-import sys
-from pathlib import Path
-sys.path.insert(0, sys.argv[1])
-import httpx
-import huggingface_hub
-from huggingface_hub.errors import LocalEntryNotFoundError
-import litellm
-payload = sys.argv[2].encode()
-offline = sys.argv[3] == "1"
-observed = []
-def handle(request):
- assert not offline, "offline loading issued a request"
- if request.url.path.endswith("/tokenizer.json"):
- observed.append(request.headers.get("authorization"))
- if request.headers.get("authorization") != "Bearer audit-fixture-token":
- return httpx.Response(401)
- return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"")
-if not offline:
- huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle)))
-try:
- tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"]
-except LocalEntryNotFoundError:
- assert offline
- assert observed == []
-else:
- assert not offline
- assert "Bearer audit-fixture-token" in observed
- assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World"
- assert tuple(Path(sys.argv[4]).rglob("tokenizer.json"))
-print("compatible")
-"""
- result: Final = subprocess.run(
- [
- sys.executable,
- "-I",
- "-c",
- script,
- str(Path(litellm.__file__).parent.parent),
- TOKENIZER_JSON,
- offline,
- str(tmp_path / "cache"),
- ],
- capture_output=True,
- text=True,
- timeout=30,
- env={
- **os.environ,
- "HF_HOME": str(tmp_path / "home"),
- "HF_HUB_CACHE": str(tmp_path / "cache"),
- "HF_ENDPOINT": "http://127.0.0.1:9",
- "HF_TOKEN": "audit-fixture-token",
- "HF_HUB_OFFLINE": offline,
- "HF_HUB_DISABLE_IMPLICIT_TOKEN": "0",
- "LITELLM_LOCAL_MODEL_COST_MAP": "True",
- },
- )
- assert result.returncode == 0, result.stdout + result.stderr
- assert result.stdout.strip() == "compatible"
-
-
-@pytest.mark.parametrize("rust", (None, "0", "1"))
-def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None:
- script: Final = """
-import importlib.abc
-import sys
-sys.path.insert(0, sys.argv[1])
-def reject_network(event, args):
- if event == "socket.connect":
- raise AssertionError("tokenizer attempted a network connection")
-sys.addaudithook(reject_network)
-class Block(importlib.abc.MetaPathFinder):
- def find_spec(self, fullname, path=None, target=None):
- if fullname == "litellm.rust_bridge._native":
- raise ImportError("native extension is unavailable")
-sys.meta_path.insert(0, Block())
-import litellm
-from litellm.rust_bridge.tokenizer import get_encoding
-import tiktoken
-from tokenizers import Tokenizer
-assert isinstance(litellm.encoding, tiktoken.Encoding)
-for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"):
- encoding = get_encoding(name)
- text = "offline café æ¼¢å— ðŸ™‚" + " " * 64
- assert encoding.decode(encoding.encode(text)) == text
-ids = litellm.encode(text="hello world")
-assert litellm.decode(tokens=ids) == "hello world"
-assert litellm.token_counter(model=None, text="hello world") == len(ids)
-custom = litellm.create_tokenizer(sys.argv[2])
-assert isinstance(custom["tokenizer"], Tokenizer)
-custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]")
-assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World"
-print("compatible")
-"""
- result: Final = subprocess.run(
- [sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON],
- capture_output=True,
- text=True,
- timeout=30,
- cwd=tmp_path,
- env={
- **{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"},
- **({"LITELLM_RUST": rust} if rust is not None else {}),
- "LITELLM_LOCAL_MODEL_COST_MAP": "True",
- "TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"),
- },
- )
- assert result.returncode == 0, result.stdout + result.stderr
- assert result.stdout.strip() == "compatible"
- assert not (tmp_path / "unused-tokenizer-cache").exists()
-
-
-@pytest.mark.parametrize("is_pretokenized", (False, True))
-def test_huggingface_batch_sequence_containers_match_python(is_pretokenized: bool) -> None:
- reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
- tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON)
- inputs: Final = [["Hello", "World"], ("Hello", "World")]
- actual: Final = tokenizer.encode_batch(inputs, is_pretokenized=is_pretokenized)
- expected: Final = reference.encode_batch(inputs, is_pretokenized=is_pretokenized)
- assert [(item.ids, item.type_ids, item.sequence_ids) for item in actual] == [
- (item.ids, item.type_ids, item.sequence_ids) for item in expected
- ]
-
-
-@pytest.mark.parametrize("name", ("cl100k_base", "o200k_base", "p50k_edit", "gpt2"))
+@pytest.mark.parametrize("name", ("gpt2",))
def test_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None:
- reference: Final = tiktoken.get_encoding(name)
- encoding: Final = OpenAIEncoding.from_tiktoken(name)
- text: Final = "hello fanta"
-
- assert repr(encoding) == repr(reference) == f""
- assert (encoding.name, encoding.n_vocab, encoding.max_token_value) == (
- reference.name,
- reference.n_vocab,
- reference.max_token_value,
- )
- assert encoding.token_byte_values() == reference.token_byte_values()
- assert encoding.encode_single_token("hello") == reference.encode_single_token("hello")
- assert encoding.encode_single_token(b"<|endoftext|>") == reference.eot_token
- assert [encoding.is_special_token(token) for token in (0, reference.eot_token)] == [False, True]
- assert encoding.decode_with_offsets(reference.encode(text)) == reference.decode_with_offsets(reference.encode(text))
- assert encoding.encode_to_numpy(text).tolist() == reference.encode_to_numpy(text).tolist()
- stable, completions = encoding.encode_with_unstable(text)
- expected_stable, expected_completions = reference.encode_with_unstable(text)
- assert (stable, sorted(completions)) == (expected_stable, sorted(expected_completions))
- with pytest.raises(KeyError):
- encoding.encode_single_token("<|not-a-token|>")
-
-
-def test_huggingface_tokenizer_exposes_the_tokenizers_vocabulary_surface() -> None:
- reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
- reference.enable_padding(pad_id=0, pad_token="[UNK]", length=4)
- reference.enable_truncation(max_length=3, stride=1, strategy="only_first", direction="left")
- tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str())
-
- assert tokenizer.token_to_id("Hello") == reference.token_to_id("Hello") == 1
- assert tokenizer.id_to_token(3) == reference.id_to_token(3) == "[BOS]"
- assert tokenizer.id_to_token(99) is None
- assert tokenizer.get_vocab() == reference.get_vocab()
- assert tokenizer.get_vocab(with_added_tokens=False) == reference.get_vocab(with_added_tokens=False)
- assert tokenizer.get_vocab_size() == reference.get_vocab_size() == 4
- assert tokenizer.get_vocab_size(with_added_tokens=False) == reference.get_vocab_size(with_added_tokens=False)
- added: Final = tokenizer.get_added_tokens_decoder()
- expected_added: Final = reference.get_added_tokens_decoder()
- assert {token_id: str(token) for token_id, token in added.items()} == {
- token_id: str(token) for token_id, token in expected_added.items()
- }
- assert added[3].special == expected_added[3].special
- assert tokenizer.num_special_tokens_to_add(False) == reference.num_special_tokens_to_add(False) == 1
- assert tokenizer.num_special_tokens_to_add(True) == reference.num_special_tokens_to_add(True) == 0
- assert tokenizer.padding == reference.padding
- assert tokenizer.truncation == reference.truncation
- assert tokenizer.encode_special_tokens == reference.encode_special_tokens is False
- assert HuggingFaceTokenizer.from_buffer(TOKENIZER_JSON.encode()).encode("Hello").ids == [3, 1]
- assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).padding is None
- assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).truncation is None
-
-
-def test_huggingface_encoding_exposes_the_tokenizers_lookup_and_mutation_surface() -> None:
- reference: Final = ReferenceTokenizer.from_str(claude_json_str)
- tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str)
- text: Final = "hello wide world"
- actual: Final = tokenizer.encode(text, "again")
- expected: Final = reference.encode(text, "again")
-
- lookups: Final = (
- lambda encoding: [encoding.token_to_chars(index) for index in range(len(encoding))],
- lambda encoding: [encoding.token_to_word(index) for index in range(len(encoding))],
- lambda encoding: [encoding.token_to_sequence(index) for index in range(len(encoding))],
- lambda encoding: [encoding.char_to_token(position) for position in range(len(text))],
- lambda encoding: [encoding.char_to_word(position) for position in range(len(text))],
- lambda encoding: [encoding.char_to_token(position, 1) for position in range(5)],
- lambda encoding: [encoding.word_to_tokens(word) for word in range(3)],
- lambda encoding: [encoding.word_to_chars(word) for word in range(3)],
- lambda encoding: [encoding.word_to_tokens(0, 1), encoding.word_to_chars(0, 1)],
- )
- for lookup in lookups:
- assert lookup(actual) == lookup(expected)
- assert repr(actual) == repr(expected)
-
- actual.truncate(4, stride=1, direction="left")
- expected.truncate(4, stride=1, direction="left")
- assert (actual.ids, [item.ids for item in actual.overflowing]) == (
- expected.ids,
- [item.ids for item in expected.overflowing],
- )
- actual.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="")
- expected.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="")
- assert (actual.ids, actual.attention_mask, actual.type_ids, actual.tokens) == (
- expected.ids,
- expected.attention_mask,
- expected.type_ids,
- expected.tokens,
- )
- actual.set_sequence_id(3)
- expected.set_sequence_id(3)
- assert actual.sequence_ids == expected.sequence_ids
- merged: Final = type(actual).merge([actual, tokenizer.encode("more")])
- assert merged.ids == type(expected).merge([expected, reference.encode("more")]).ids
- assert merged.offsets == type(expected).merge([expected, reference.encode("more")]).offsets
- with pytest.raises(ValueError, match="direction"):
- actual.pad(8, direction="sideways")
+ assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name)
diff --git a/tests/test_litellm/proxy/client/test_chat.py b/tests/test_litellm/proxy/client/test_chat.py
index 67b6ee833f2..8fe1bfcbb2f 100644
--- a/tests/test_litellm/proxy/client/test_chat.py
+++ b/tests/test_litellm/proxy/client/test_chat.py
@@ -13,7 +13,7 @@ from litellm.proxy.client.exceptions import UnauthorizedError
def _load_http_mocking_responses():
"""Load the third-party `responses` package even if test collection creates
- a top-level `responses` namespace package from `tests/test_litellm/responses`.
+ a top-level `responses` namespace package from `tests/unit/responses`.
"""
module = importlib.import_module("responses")
if hasattr(module, "activate"):
diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py
index e6795bb22f3..42c1f489bdd 100644
--- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py
+++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py
@@ -3674,7 +3674,7 @@ async def test_post_call_success_hook_contains_header_merge_failures(
@pytest.mark.asyncio
async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loop(rate_limiter):
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py
index e91b7ef970c..9d4532df49a 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py
@@ -162,7 +162,7 @@ async def test_interrupted_anthropic_stream_recovers_output_tokens_off_the_event
from unittest.mock import AsyncMock
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
@@ -201,7 +201,7 @@ async def test_failed_anthropic_stream_records_partial_usage_off_the_event_loop(
from unittest.mock import AsyncMock
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 884a9c81500..89156cd19a0 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -14974,7 +14974,7 @@ def test_settings_store_exposes_dashboard_saved_mcp_client_allowlist_to_the_mcp_
async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch):
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
@@ -14995,7 +14995,7 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
from litellm.rust_bridge._native import Tokenizer
from litellm import Router
- from tests.test_litellm.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags
+ from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags
claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"]
diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py
index 0fc7295a717..ea1870d3b73 100644
--- a/tests/test_litellm/proxy/test_proxy_utils.py
+++ b/tests/test_litellm/proxy/test_proxy_utils.py
@@ -2021,7 +2021,7 @@ async def test_a_dispatched_failure_is_counted_off_the_event_loop():
from unittest.mock import AsyncMock, patch
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/test_litellm/rust_bridge/messages/test_route_host.py b/tests/test_litellm/rust_bridge/messages/test_route_host.py
deleted file mode 100644
index c5a442e0709..00000000000
--- a/tests/test_litellm/rust_bridge/messages/test_route_host.py
+++ /dev/null
@@ -1,124 +0,0 @@
-from dataclasses import astuple
-from typing import Final
-
-import pytest
-
-import litellm
-from litellm.rust_bridge.messages import route_host
-
-pytestmark = pytest.mark.usefixtures("local_model_cost_map")
-
-
-def _flag_model(monkeypatch: pytest.MonkeyPatch, name: str, **flags: bool) -> None:
- monkeypatch.setitem(
- litellm.model_cost,
- name,
- {
- "litellm_provider": "anthropic",
- "mode": "chat",
- "input_cost_per_token": 0,
- "output_cost_per_token": 0,
- **flags,
- },
- )
-
-
-def test_capabilities_come_from_the_model_map_under_the_callers_provider(monkeypatch: pytest.MonkeyPatch) -> None:
- _flag_model(
- monkeypatch,
- "claude-test-adaptive",
- supports_reasoning=True,
- supports_adaptive_thinking=True,
- supports_output_config=True,
- supports_xhigh_reasoning_effort=True,
- supports_sampling_params=False,
- )
-
- capabilities: Final = route_host.model_capabilities("anthropic/claude-test-adaptive", None)
-
- assert capabilities.supports_adaptive_thinking
- assert capabilities.supports_output_config
- assert not capabilities.supports_legacy_thinking
- assert not capabilities.supports_sampling_params
- assert capabilities.effort_tiers.xhigh
- assert not capabilities.effort_tiers.max
-
-
-def test_unmapped_model_keeps_sampling_params_and_no_reasoning_features() -> None:
- capabilities: Final = route_host.model_capabilities("anthropic/not-a-real-model", None)
-
- assert capabilities.supports_sampling_params
- assert not capabilities.supports_reasoning
- assert not capabilities.supports_adaptive_thinking
- assert not any(astuple(capabilities.effort_tiers))
-
-
-@pytest.mark.parametrize(
- ("global_flag", "kwargs", "expected"),
- [
- (False, {}, False),
- (True, {}, True),
- (False, {"drop_params": "true"}, True),
- (False, {"drop_params": "nonsense"}, False),
- (False, {"drop_params": False}, False),
- ],
-)
-def test_drop_params_merges_the_global_flag_with_the_request(
- monkeypatch: pytest.MonkeyPatch, global_flag: bool, kwargs: dict[str, object], expected: bool
-) -> None:
- monkeypatch.setattr(litellm, "drop_params", global_flag)
-
- assert route_host.shaping("anthropic/not-a-real-model", None, kwargs)["drop_params"] is expected
-
-
-@pytest.mark.parametrize(
- ("configured", "expected"),
- [
- (["tools[*].input_examples", 3, "metadata.user_id"], ("tools[*].input_examples", "metadata.user_id")),
- ("tools", ()),
- (None, ()),
- ],
-)
-def test_additional_drop_params_keep_only_string_paths(configured: object, expected: tuple[str, ...]) -> None:
- shaping: Final = route_host.shaping("anthropic/not-a-real-model", None, {"additional_drop_params": configured})
-
- assert shaping["additional_drop_params"] == expected
-
-
-def test_native_request_rejections_map_to_the_public_400() -> None:
- from types import MappingProxyType
-
- from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
-
- request: Final = LiteLLMMessagesRequest(
- model="anthropic/claude-sonnet-5",
- messages=(),
- max_tokens=8,
- stream=None,
- api_key=None,
- api_base=None,
- custom_llm_provider=None,
- kwargs=MappingProxyType({}),
- )
- rejected: Final = ValueError("claude-sonnet-5 does not support top_k=5")
- rejected.messages_request_error = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets
-
- mapped: Final = route_host.map_failure(rejected, request, "anthropic")
-
- assert isinstance(mapped, litellm.BadRequestError)
- assert mapped.status_code == 400
- assert "does not support top_k=5" in mapped.message
- assert mapped.model == "claude-sonnet-5"
- assert not isinstance(route_host.map_failure(ValueError("plain"), request, "anthropic"), litellm.BadRequestError)
-
-
-def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handler_does() -> None:
- hidden: Final = route_host.stream_hidden_params(
- (("request-id", "req_upstream_123"), ("x-ratelimit-remaining-requests", "41"))
- )
-
- additional: Final = hidden["additional_headers"]
- assert isinstance(additional, dict)
- assert additional["llm_provider-request-id"] == "req_upstream_123"
- assert additional["x-ratelimit-remaining-requests"] == "41"
- assert "request-id" not in additional
diff --git a/tests/test_litellm/rust_bridge/responses/__init__.py b/tests/test_litellm/rust_bridge/responses/__init__.py
deleted file mode 100644
index e69de29bb2d..00000000000
diff --git a/tests/test_litellm_rust/tokenizer/test_fast_count.py b/tests/test_litellm_rust/tokenizer/test_fast_count.py
index 2902b79dca8..f91f47e4b86 100644
--- a/tests/test_litellm_rust/tokenizer/test_fast_count.py
+++ b/tests/test_litellm_rust/tokenizer/test_fast_count.py
@@ -7,7 +7,7 @@ from tokenizers import Tokenizer as ReferenceTokenizer
from litellm.rust_bridge import _native
from litellm.utils import claude_json_str
-from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
+from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
pytestmark = pytest.mark.requires_rust_extension
diff --git a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py
index abf6a6dda31..2e883e91fda 100644
--- a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py
+++ b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py
@@ -94,7 +94,7 @@ class _AgentChunk:
@pytest.mark.asyncio
async def test_stream_completion_counts_tokens_off_the_event_loop(monkeypatch):
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/unit/a2a_protocol/test_main.py b/tests/unit/a2a_protocol/test_main.py
index c65d171246d..4ba0ef8fa04 100644
--- a/tests/unit/a2a_protocol/test_main.py
+++ b/tests/unit/a2a_protocol/test_main.py
@@ -469,7 +469,7 @@ class _UsageRecorder(CustomLogger):
@pytest.mark.asyncio
async def test_asend_message_counts_usage_off_the_event_loop(monkeypatch):
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/test_litellm/caching/test_azure_blob_cache.py b/tests/unit/caching/test_azure_blob_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_azure_blob_cache.py
rename to tests/unit/caching/test_azure_blob_cache.py
diff --git a/tests/test_litellm/caching/test_caching.py b/tests/unit/caching/test_caching.py
similarity index 100%
rename from tests/test_litellm/caching/test_caching.py
rename to tests/unit/caching/test_caching.py
diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py
index a181ef89fe0..425d657312a 100644
--- a/tests/unit/caching/test_caching_handler.py
+++ b/tests/unit/caching/test_caching_handler.py
@@ -39,6 +39,11 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm._logging import verbose_logger
import logging
+import json
+import httpx
+import respx
+from fastapi.testclient import TestClient
+from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
def setup_cache():
@@ -1062,6 +1067,9 @@ def test_is_chat_completion_cached_dict():
assert _is_chat_completion_cached_dict(
{"id": "other", "object": "chat.completion.chunk", "choices": []}
)
+ assert _is_chat_completion_cached_dict(
+ {"id": "no-object", "choices": [{"index": 0}]}
+ )
assert not _is_chat_completion_cached_dict(
{"id": "resp_abc", "object": "response", "output": []}
)
@@ -1432,3 +1440,799 @@ def test_convert_cached_responses_result_parameterized(
assert result is not None
assert result.id == cached_result["id"]
assert result.status == cached_result["status"]
+
+
+@pytest.mark.asyncio
+async def test_process_async_embedding_cached_response():
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ args = {
+ "cached_result": [
+ {
+ "embedding": [-0.025122925639152527, -0.019487135112285614],
+ "index": 0,
+ "object": "embedding",
+ }
+ ]
+ }
+
+ mock_logging_obj = MagicMock()
+ mock_logging_obj.async_success_handler = AsyncMock()
+ response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
+ final_embedding_cached_response=None,
+ cached_result=args["cached_result"],
+ kwargs={"model": "text-embedding-ada-002", "input": "test"},
+ logging_obj=mock_logging_obj,
+ start_time=datetime.now(),
+ model="text-embedding-ada-002",
+ )
+
+ assert cache_hit
+
+ print(f"response: {response}")
+ assert len(response.data) == 1
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_preserves_prompt_tokens_details():
+ """Test that prompt_tokens_details (including image_count) survives a full cache hit."""
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ cached_result = [
+ {
+ "embedding": [-0.025, -0.019],
+ "index": 0,
+ "object": "embedding",
+ "model": "amazon.titan-embed-image-v1",
+ "prompt_tokens_details": {"image_count": 1},
+ }
+ ]
+
+ mock_logging_obj = MagicMock()
+ mock_logging_obj.async_success_handler = AsyncMock()
+ response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
+ final_embedding_cached_response=None,
+ cached_result=cached_result,
+ kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"},
+ logging_obj=mock_logging_obj,
+ start_time=datetime.now(),
+ model="amazon.titan-embed-image-v1",
+ )
+
+ assert cache_hit
+ assert response.usage is not None
+ assert response.usage.prompt_tokens_details is not None
+ assert response.usage.prompt_tokens_details.image_count == 1
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_backward_compat_no_prompt_tokens_details():
+ """Test that old cached items without prompt_tokens_details still work."""
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ # Old-format cached item — no prompt_tokens_details field
+ cached_result = [
+ {
+ "embedding": [-0.025, -0.019],
+ "index": 0,
+ "object": "embedding",
+ "model": "text-embedding-ada-002",
+ }
+ ]
+
+ mock_logging_obj = MagicMock()
+ mock_logging_obj.async_success_handler = AsyncMock()
+ response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
+ final_embedding_cached_response=None,
+ cached_result=cached_result,
+ kwargs={"model": "text-embedding-ada-002", "input": "test"},
+ logging_obj=mock_logging_obj,
+ start_time=datetime.now(),
+ model="text-embedding-ada-002",
+ )
+
+ assert cache_hit
+ assert response.usage is not None
+ assert response.usage.prompt_tokens_details is None
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_aggregates_multiple_image_counts():
+ """Test that image_count is summed correctly across multiple cached items."""
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ cached_result = [
+ {
+ "embedding": [-0.025, -0.019],
+ "index": 0,
+ "object": "embedding",
+ "model": "amazon.titan-embed-image-v1",
+ "prompt_tokens_details": {"image_count": 1},
+ },
+ {
+ "embedding": [0.031, 0.042],
+ "index": 1,
+ "object": "embedding",
+ "model": "amazon.titan-embed-image-v1",
+ "prompt_tokens_details": {"image_count": 1},
+ },
+ ]
+
+ mock_logging_obj = MagicMock()
+ mock_logging_obj.async_success_handler = AsyncMock()
+ response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
+ final_embedding_cached_response=None,
+ cached_result=cached_result,
+ kwargs={
+ "model": "amazon.titan-embed-image-v1",
+ "input": ["img1", "img2"],
+ },
+ logging_obj=mock_logging_obj,
+ start_time=datetime.now(),
+ model="amazon.titan-embed-image-v1",
+ )
+
+ assert cache_hit
+ assert response.usage.prompt_tokens_details is not None
+ assert response.usage.prompt_tokens_details.image_count == 2
+
+
+def test_combine_usage_merges_prompt_tokens_details():
+ """Test that combine_usage merges prompt_tokens_details from both Usage objects."""
+ from litellm.types.utils import PromptTokensDetailsWrapper, Usage
+
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ usage1 = Usage(
+ prompt_tokens=10,
+ completion_tokens=0,
+ total_tokens=10,
+ prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
+ )
+ usage2 = Usage(
+ prompt_tokens=20,
+ completion_tokens=0,
+ total_tokens=20,
+ prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2),
+ )
+
+ combined = llm_caching_handler.combine_usage(usage1, usage2)
+
+ assert combined.prompt_tokens == 30
+ assert combined.total_tokens == 30
+ assert combined.prompt_tokens_details is not None
+ assert combined.prompt_tokens_details.image_count == 3
+
+
+def test_combine_usage_handles_none_details():
+ """Test that combine_usage works when one or both sides have null prompt_tokens_details."""
+ from litellm.types.utils import PromptTokensDetailsWrapper, Usage
+
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ # Both null
+ usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10)
+ usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20)
+ combined = llm_caching_handler.combine_usage(usage_a, usage_b)
+ assert combined.prompt_tokens_details is None
+
+ # Only first has details
+ usage_c = Usage(
+ prompt_tokens=10,
+ completion_tokens=0,
+ total_tokens=10,
+ prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
+ )
+ combined = llm_caching_handler.combine_usage(usage_c, usage_b)
+ assert combined.prompt_tokens_details is not None
+ assert combined.prompt_tokens_details.image_count == 1
+
+ # Only second has details
+ combined = llm_caching_handler.combine_usage(usage_a, usage_c)
+ assert combined.prompt_tokens_details is not None
+ assert combined.prompt_tokens_details.image_count == 1
+
+
+def _build_logging_obj(call_type: str, stream: bool):
+ import uuid as _uuid
+
+ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
+
+ return LiteLLMLogging(
+ litellm_call_id=str(datetime.now()),
+ call_type=call_type,
+ model="gpt-5.4",
+ messages=[],
+ function_id=str(_uuid.uuid4()),
+ stream=stream,
+ start_time=datetime.now(),
+ )
+
+
+def test_convert_cached_responses_bridge_chat_completion_nonstream():
+ """openai/responses chat-completions bridge: non-streaming cache hit replays as ModelResponse."""
+ from litellm import responses
+ from litellm.types.utils import CallTypes, ModelResponse
+
+ caching_handler = LLMCachingHandler(
+ original_function=responses, request_kwargs={}, start_time=datetime.now()
+ )
+ cached_result = {
+ "id": "chatcmpl-bridge-nonstream",
+ "object": "chat.completion",
+ "created": int(time.time()),
+ "model": "gpt-5.4",
+ "choices": [
+ {
+ "index": 0,
+ "message": {"role": "assistant", "content": "Hi!"},
+ "finish_reason": "stop",
+ }
+ ],
+ "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18},
+ }
+
+ result = caching_handler._convert_cached_result_to_model_response(
+ cached_result=cached_result,
+ call_type=CallTypes.responses.value,
+ kwargs={
+ "model": "gpt-5.4",
+ "stream": False,
+ "messages": [{"role": "user", "content": "hi"}],
+ },
+ logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False),
+ model="gpt-5.4",
+ args=(),
+ )
+
+ assert isinstance(result, ModelResponse)
+ assert result.choices[0].message.content == "Hi!"
+
+
+def test_convert_cached_responses_legacy_nonstream_path():
+ """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) falls through legacy path."""
+ from litellm import responses
+ from litellm.types.llms.openai import ResponsesAPIResponse
+ from litellm.types.utils import CallTypes
+
+ caching_handler = LLMCachingHandler(
+ original_function=responses, request_kwargs={}, start_time=datetime.now()
+ )
+ cached_result = {
+ "id": "resp_legacy_nonstream",
+ "created_at": int(time.time()),
+ "status": "completed",
+ "model": "gpt-4o",
+ "object": "response",
+ "output": [
+ {
+ "type": "message",
+ "id": "msg_legacy",
+ "status": "completed",
+ "role": "assistant",
+ "content": [
+ {
+ "type": "output_text",
+ "text": "legacy response",
+ "annotations": [],
+ }
+ ],
+ }
+ ],
+ }
+
+ result = caching_handler._convert_cached_result_to_model_response(
+ cached_result=cached_result,
+ call_type=CallTypes.responses.value,
+ kwargs={"model": "gpt-4o", "input": "hi", "stream": False},
+ logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False),
+ model="gpt-4o",
+ args=(),
+ )
+
+ assert isinstance(result, ResponsesAPIResponse)
+ assert result.id == "resp_legacy_nonstream"
+
+
+def test_convert_cached_responses_legacy_stream_path():
+ """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) on stream falls through legacy path."""
+ from litellm import responses
+ from litellm.responses.streaming_iterator import (
+ CachedResponsesAPIStreamingIterator,
+ )
+ from litellm.types.utils import CallTypes
+
+ caching_handler = LLMCachingHandler(
+ original_function=responses, request_kwargs={}, start_time=datetime.now()
+ )
+ cached_result = {
+ "id": "resp_legacy_stream",
+ "created_at": int(time.time()),
+ "status": "completed",
+ "model": "gpt-4o",
+ "object": "response",
+ "output": [
+ {
+ "type": "message",
+ "id": "msg_legacy_stream",
+ "status": "completed",
+ "role": "assistant",
+ "content": [
+ {
+ "type": "output_text",
+ "text": "legacy stream",
+ "annotations": [],
+ }
+ ],
+ }
+ ],
+ }
+
+ result = caching_handler._convert_cached_result_to_model_response(
+ cached_result=cached_result,
+ call_type=CallTypes.responses.value,
+ kwargs={"model": "gpt-4o", "input": "hi", "stream": True},
+ logging_obj=_build_logging_obj(CallTypes.responses.value, stream=True),
+ model="gpt-4o",
+ args=(),
+ )
+
+ assert isinstance(result, CachedResponsesAPIStreamingIterator)
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_restores_stored_prompt_tokens_for_image_input():
+ """Image-embedding cache hit restores prompt_tokens=0 from the stored value
+ instead of recomputing a bogus count by tokenizing the base64 input."""
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ # base64-like blob — token_counter over this would return a large nonzero count
+ image_input = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" * 50
+
+ cached_result = [
+ {
+ "embedding": [-0.025, -0.019],
+ "index": 0,
+ "object": "embedding",
+ "model": "amazon.titan-embed-image-v1",
+ "prompt_tokens": 0,
+ "prompt_tokens_details": {"image_count": 1},
+ }
+ ]
+
+ mock_logging_obj = MagicMock()
+ mock_logging_obj.async_success_handler = AsyncMock()
+ response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
+ final_embedding_cached_response=None,
+ cached_result=cached_result,
+ kwargs={"model": "amazon.titan-embed-image-v1", "input": image_input},
+ logging_obj=mock_logging_obj,
+ start_time=datetime.now(),
+ model="amazon.titan-embed-image-v1",
+ )
+
+ assert cache_hit
+ assert response.usage is not None
+ assert response.usage.prompt_tokens == 0
+ assert response.usage.total_tokens == 0
+ assert response.usage.prompt_tokens_details.image_count == 1
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_sums_stored_prompt_tokens_across_items():
+ """A multi-item cache hit sums the stored per-item prompt_tokens back to the total."""
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ cached_result = [
+ {
+ "embedding": [-0.01],
+ "index": 0,
+ "object": "embedding",
+ "model": "text-embedding-3-small",
+ "prompt_tokens": 5,
+ },
+ {
+ "embedding": [-0.02],
+ "index": 1,
+ "object": "embedding",
+ "model": "text-embedding-3-small",
+ "prompt_tokens": 4,
+ },
+ ]
+
+ mock_logging_obj = MagicMock()
+ mock_logging_obj.async_success_handler = AsyncMock()
+ response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
+ final_embedding_cached_response=None,
+ cached_result=cached_result,
+ kwargs={"model": "text-embedding-3-small", "input": ["hello world", "foo bar"]},
+ logging_obj=mock_logging_obj,
+ start_time=datetime.now(),
+ model="text-embedding-3-small",
+ )
+
+ assert cache_hit
+ assert response.usage.prompt_tokens == 9
+ assert response.usage.total_tokens == 9
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries():
+ """Legacy cache entries with no stored prompt_tokens still recompute via token_counter
+ for str inputs (backward compatibility)."""
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ # No prompt_tokens key — pre-fix entry
+ cached_result = [
+ {
+ "embedding": [-0.025, -0.019],
+ "index": 0,
+ "object": "embedding",
+ "model": "text-embedding-ada-002",
+ },
+ ]
+
+ mock_logging_obj = MagicMock()
+ mock_logging_obj.async_success_handler = AsyncMock()
+ response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
+ final_embedding_cached_response=None,
+ cached_result=cached_result,
+ kwargs={"model": "text-embedding-ada-002", "input": "hello world"},
+ logging_obj=mock_logging_obj,
+ start_time=datetime.now(),
+ model="text-embedding-ada-002",
+ )
+
+ assert cache_hit
+ # token_counter over "hello world" yields a nonzero count — fallback path still runs
+ assert response.usage.prompt_tokens > 0
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_hit_sets_custom_llm_provider_on_logging_obj():
+ """A full embedding cache hit must stamp the resolved provider onto the logging
+ obj so spend logs record the provider instead of None/unknown."""
+ from litellm.types.utils import CallTypes
+
+ llm_caching_handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+
+ cached_result = [
+ {
+ "embedding": [-0.025, -0.019],
+ "index": 0,
+ "object": "embedding",
+ "model": "text-embedding-3-small",
+ "prompt_tokens": 5,
+ }
+ ]
+
+ logging_obj = _build_logging_obj(CallTypes.aembedding.value, stream=False)
+ logging_obj.async_success_handler = AsyncMock()
+
+ response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
+ final_embedding_cached_response=None,
+ cached_result=cached_result,
+ kwargs={"model": "text-embedding-3-small", "input": "hello world"},
+ logging_obj=logging_obj,
+ start_time=datetime.now(),
+ model="text-embedding-3-small",
+ )
+
+ assert cache_hit
+ assert logging_obj.model_call_details["custom_llm_provider"] == "openai"
+
+
+def test_sync_stream_responses_cache_hit_sets_custom_llm_provider_on_logging_obj(monkeypatch):
+ import litellm
+ from litellm.caching.caching import Cache
+ from litellm.types.utils import CallTypes
+
+ monkeypatch.setattr(litellm, "cache", Cache(type="local"))
+ kwargs = {"model": "azure/gpt-5.4-mini", "input": "hello", "stream": True}
+ cached_response = {
+ "id": "resp_sync_stream",
+ "created_at": int(time.time()),
+ "status": "completed",
+ "model": "gpt-5.4-mini",
+ "object": "response",
+ "output": [
+ {
+ "type": "message",
+ "id": "msg_sync_stream",
+ "status": "completed",
+ "role": "assistant",
+ "content": [{"type": "output_text", "text": "hi", "annotations": []}],
+ }
+ ],
+ }
+ litellm.cache.add_cache(json.dumps(cached_response), **kwargs)
+ handler = LLMCachingHandler(original_function=litellm.responses, request_kwargs=kwargs, start_time=datetime.now())
+ logging_obj = _build_logging_obj(CallTypes.responses.value, stream=True)
+
+ hit = handler._sync_get_cache(
+ model="azure/gpt-5.4-mini",
+ original_function=litellm.responses,
+ logging_obj=logging_obj,
+ start_time=datetime.now(),
+ call_type=CallTypes.responses.value,
+ kwargs=kwargs,
+ args=(),
+ )
+
+ assert hit.cached_result is not None
+ assert logging_obj.model_call_details["custom_llm_provider"] == "azure"
+ assert logging_obj.model_call_details["litellm_params"]["custom_llm_provider"] == "azure"
+
+
+def test_request_kwargs_does_not_retain_logging_obj():
+ """
+ The caching handler lives on logging_obj._llm_caching_handler, so keeping
+ litellm_logging_obj inside request_kwargs closes a reference cycle
+ (Logging -> LLMCachingHandler -> kwargs -> Logging). That cycle keeps the
+ full request payload alive until a generational GC pass instead of being
+ freed by refcount when the request finishes; under bursts of large-token
+ requests this presents as stepwise RSS growth that never returns to
+ baseline. Other kwargs (messages included) must be preserved.
+ """
+ logging_obj = MagicMock()
+ kwargs = {
+ "model": "gpt-4o",
+ "messages": [{"role": "user", "content": "hello"}],
+ "litellm_logging_obj": logging_obj,
+ }
+
+ handler = LLMCachingHandler(
+ original_function=MagicMock(),
+ request_kwargs=kwargs,
+ start_time=datetime.now(),
+ )
+
+ assert "litellm_logging_obj" not in handler.request_kwargs
+ assert handler.request_kwargs["messages"] == kwargs["messages"]
+ assert handler.request_kwargs["model"] == "gpt-4o"
+
+
+def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch):
+ """
+ Regression test for the SDK losing async cache writes in short-lived scripts:
+ async_set_cache dispatched the write as a bare fire-and-forget task, so
+ asyncio.run cancelled it at loop close before the write landed (LIT-6184,
+ deterministic with hiredis installed). The write must survive loop shutdown.
+ """
+ import litellm
+
+ writes = []
+
+ class _SlowWriteCache:
+ supported_call_types = ["acompletion"]
+ cache = None
+
+ async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs):
+ await asyncio.sleep(0.2)
+ writes.append(result)
+
+ async def acompletion(**kwargs):
+ return None
+
+ handler = LLMCachingHandler(
+ original_function=acompletion,
+ request_kwargs={},
+ start_time=datetime.now(),
+ )
+ monkeypatch.setattr(litellm, "cache", _SlowWriteCache())
+
+ async def _short_lived_script():
+ await handler.async_set_cache(
+ result=litellm.ModelResponse(),
+ original_function=acompletion,
+ kwargs={},
+ )
+
+ asyncio.run(_short_lived_script())
+
+ assert len(writes) == 1
+
+
+@pytest.mark.asyncio
+async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monkeypatch):
+ """The spend log for a cache hit must reuse the key the lookup already computed instead of hashing again."""
+ import litellm
+ from litellm.caching.caching import Cache
+ from litellm.types.utils import CallTypes
+
+ async def acompletion(**kwargs):
+ return None
+
+ monkeypatch.setattr(litellm, "cache", Cache(type="local"))
+ kwargs = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "caching": True}
+ await litellm.cache.async_add_cache(
+ litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}]), **kwargs
+ )
+ handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now())
+ logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False)
+ logging_obj.async_success_handler = AsyncMock()
+
+ hit = await handler._async_get_cache(
+ model="gpt-5.4",
+ original_function=acompletion,
+ logging_obj=logging_obj,
+ start_time=datetime.now(),
+ call_type=CallTypes.acompletion.value,
+ kwargs=kwargs,
+ args=(),
+ )
+
+ assert hit is not None and hit.cached_result is not None
+ assert handler.preset_cache_key is not None
+ assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key
+ assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key
+
+
+@pytest.mark.asyncio
+async def test_converted_stream_cache_hit_replayed_as_plain_object_logs_at_hit_time(monkeypatch):
+ import litellm
+ from litellm.caching.caching import Cache
+ from litellm.types.utils import CallTypes
+
+ async def aanthropic_messages(**kwargs):
+ return None
+
+ monkeypatch.setattr(litellm, "cache", Cache(type="local"))
+ kwargs = {
+ "model": "claude-sonnet-5",
+ "messages": [{"role": "user", "content": "hello"}],
+ "max_tokens": 16,
+ "caching": True,
+ "stream": False,
+ "_websearch_interception_converted_stream": True,
+ }
+ cached_message = {
+ "id": "msg_1",
+ "type": "message",
+ "role": "assistant",
+ "content": [{"type": "text", "text": "hi"}],
+ }
+ await litellm.cache.async_add_cache(cached_message, **kwargs)
+ handler = LLMCachingHandler(original_function=aanthropic_messages, request_kwargs=kwargs, start_time=datetime.now())
+ logging_obj = _build_logging_obj(CallTypes.aanthropic_messages.value, stream=False)
+ logging_obj.async_success_handler = AsyncMock()
+ logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
+
+ hit = await handler._async_get_cache(
+ model="claude-sonnet-5",
+ original_function=aanthropic_messages,
+ logging_obj=logging_obj,
+ start_time=datetime.now(),
+ call_type=CallTypes.aanthropic_messages.value,
+ kwargs=kwargs,
+ args=(),
+ )
+
+ assert hit is not None and hit.cached_result == cached_message
+ logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
+ assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
+
+
+@pytest.mark.asyncio
+async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_replays_as_plain_object(monkeypatch):
+ import litellm
+ from litellm.caching.caching import Cache
+ from litellm.types.utils import CallTypes
+
+ async def acompletion(**kwargs):
+ return None
+
+ monkeypatch.setattr(litellm, "cache", Cache(type="local"))
+ kwargs = {
+ "model": "gpt-5.6",
+ "messages": [{"role": "user", "content": "run the code"}],
+ "caching": True,
+ "stream": False,
+ "_code_interpreter_interception_converted_stream": True,
+ "_agentic_loop_depth": 1,
+ }
+ await litellm.cache.async_add_cache(
+ litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]), **kwargs
+ )
+ handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now())
+ logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False)
+ logging_obj.async_success_handler = AsyncMock()
+ logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
+
+ hit = await handler._async_get_cache(
+ model="gpt-5.6",
+ original_function=acompletion,
+ logging_obj=logging_obj,
+ start_time=datetime.now(),
+ call_type=CallTypes.acompletion.value,
+ kwargs=kwargs,
+ args=(),
+ )
+
+ assert hit is not None and isinstance(hit.cached_result, litellm.ModelResponse)
+ assert hit.cached_result.choices[0].message.content == "done"
+ logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
+ assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
+
+
+@pytest.mark.asyncio
+async def test_partial_embedding_cache_hit_sends_only_misses_and_keeps_input_order(monkeypatch):
+ import litellm
+ from litellm import CustomLLM
+ from litellm.caching.caching import Cache
+ from litellm.types.utils import Embedding, EmbeddingResponse
+
+ class RecordingEmbedder(CustomLLM):
+ provider_inputs: tuple[tuple[str, ...], ...] = ()
+
+ async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse:
+ self.provider_inputs = (*self.provider_inputs, tuple(input))
+ return EmbeddingResponse(
+ model=model,
+ data=[
+ Embedding(embedding=[float(len(text))], index=idx, object="embedding")
+ for idx, text in enumerate(input)
+ ],
+ )
+
+ embedder = RecordingEmbedder()
+ monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "recording-embedder", "custom_handler": embedder}])
+ monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "recording-embedder"])
+ monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "recording-embedder"])
+ monkeypatch.setattr(litellm, "cache", Cache(type="local"))
+
+ await litellm.aembedding(model="recording-embedder/m", input=["aa", "bbbb"])
+ await asyncio.gather(*_PENDING_CACHE_WRITES)
+ mixed_input = ["c", "aa", "ddd", "bbbb", "eeeee"]
+ response = await litellm.aembedding(model="recording-embedder/m", input=mixed_input)
+ await asyncio.gather(*_PENDING_CACHE_WRITES)
+
+ assert embedder.provider_inputs == (("aa", "bbbb"), ("c", "ddd", "eeeee")), embedder.provider_inputs
+ assert [item["index"] for item in response.data] == [0, 1, 2, 3, 4]
+ assert [item["embedding"] for item in response.data] == [[float(len(text))] for text in mixed_input]
+ assert response._hidden_params["cache_hit"] is True, "a partial hit must still be reported as a cache hit"
+
+ repeat = await litellm.aembedding(model="recording-embedder/m", input=mixed_input)
+
+ assert len(embedder.provider_inputs) == 2, embedder.provider_inputs
+ assert [item["embedding"] for item in repeat.data] == [[float(len(text))] for text in mixed_input]
diff --git a/tests/test_litellm/caching/test_check_and_fix_namespace_none_guard.py b/tests/unit/caching/test_check_and_fix_namespace_none_guard.py
similarity index 100%
rename from tests/test_litellm/caching/test_check_and_fix_namespace_none_guard.py
rename to tests/unit/caching/test_check_and_fix_namespace_none_guard.py
diff --git a/tests/test_litellm/caching/test_disk_cache.py b/tests/unit/caching/test_disk_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_disk_cache.py
rename to tests/unit/caching/test_disk_cache.py
diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_dual_cache.py
rename to tests/unit/caching/test_dual_cache.py
diff --git a/tests/test_litellm/caching/test_embedding_router.py b/tests/unit/caching/test_embedding_router.py
similarity index 100%
rename from tests/test_litellm/caching/test_embedding_router.py
rename to tests/unit/caching/test_embedding_router.py
diff --git a/tests/test_litellm/caching/test_evicted_client_closer.py b/tests/unit/caching/test_evicted_client_closer.py
similarity index 100%
rename from tests/test_litellm/caching/test_evicted_client_closer.py
rename to tests/unit/caching/test_evicted_client_closer.py
diff --git a/tests/test_litellm/caching/test_gcs_cache.py b/tests/unit/caching/test_gcs_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_gcs_cache.py
rename to tests/unit/caching/test_gcs_cache.py
diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/unit/caching/test_in_memory_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_in_memory_cache.py
rename to tests/unit/caching/test_in_memory_cache.py
diff --git a/tests/test_litellm/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py
similarity index 100%
rename from tests/test_litellm/caching/test_llm_caching_handler.py
rename to tests/unit/caching/test_llm_caching_handler.py
diff --git a/tests/test_litellm/caching/test_llm_client_cache_e2e.py b/tests/unit/caching/test_llm_client_cache_e2e.py
similarity index 100%
rename from tests/test_litellm/caching/test_llm_client_cache_e2e.py
rename to tests/unit/caching/test_llm_client_cache_e2e.py
diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/unit/caching/test_qdrant_semantic_cache.py
similarity index 99%
rename from tests/test_litellm/caching/test_qdrant_semantic_cache.py
rename to tests/unit/caching/test_qdrant_semantic_cache.py
index ca7303e4c6d..4f18fb1bca6 100644
--- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py
+++ b/tests/unit/caching/test_qdrant_semantic_cache.py
@@ -1033,7 +1033,7 @@ def test_qdrant_semantic_cache_defaults_embedding_timeout():
@pytest.mark.asyncio
async def test_qdrant_async_embedding_truncates_off_the_event_loop(monkeypatch):
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_redis_cache.py
rename to tests/unit/caching/test_redis_cache.py
diff --git a/tests/test_litellm/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_redis_cluster_cache.py
rename to tests/unit/caching/test_redis_cluster_cache.py
diff --git a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py b/tests/unit/caching/test_redis_cluster_node_isolation.py
similarity index 100%
rename from tests/test_litellm/caching/test_redis_cluster_node_isolation.py
rename to tests/unit/caching/test_redis_cluster_node_isolation.py
diff --git a/tests/test_litellm/caching/test_redis_connection_pool.py b/tests/unit/caching/test_redis_connection_pool.py
similarity index 100%
rename from tests/test_litellm/caching/test_redis_connection_pool.py
rename to tests/unit/caching/test_redis_connection_pool.py
diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/unit/caching/test_redis_semantic_cache.py
similarity index 99%
rename from tests/test_litellm/caching/test_redis_semantic_cache.py
rename to tests/unit/caching/test_redis_semantic_cache.py
index de253b4f10b..461689165bb 100644
--- a/tests/test_litellm/caching/test_redis_semantic_cache.py
+++ b/tests/unit/caching/test_redis_semantic_cache.py
@@ -1392,7 +1392,7 @@ def test_redis_semantic_cache_defaults_embedding_timeout():
@pytest.mark.asyncio
async def test_redis_async_embedding_truncates_off_the_event_loop(monkeypatch):
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/test_litellm/caching/test_s3_cache.py b/tests/unit/caching/test_s3_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_s3_cache.py
rename to tests/unit/caching/test_s3_cache.py
diff --git a/tests/test_litellm/caching/test_valkey_semantic_cache.py b/tests/unit/caching/test_valkey_semantic_cache.py
similarity index 100%
rename from tests/test_litellm/caching/test_valkey_semantic_cache.py
rename to tests/unit/caching/test_valkey_semantic_cache.py
diff --git a/tests/test_litellm/litellm_core_utils/audio_utils/__init__.py b/tests/unit/expected_responses_api_request/__init__.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/audio_utils/__init__.py
rename to tests/unit/expected_responses_api_request/__init__.py
diff --git a/tests/test_litellm/expected_responses_api_request/azure_shell_tool.json b/tests/unit/expected_responses_api_request/azure_shell_tool.json
similarity index 100%
rename from tests/test_litellm/expected_responses_api_request/azure_shell_tool.json
rename to tests/unit/expected_responses_api_request/azure_shell_tool.json
diff --git a/tests/test_litellm/expected_responses_api_request/context_management_and_shell.json b/tests/unit/expected_responses_api_request/context_management_and_shell.json
similarity index 100%
rename from tests/test_litellm/expected_responses_api_request/context_management_and_shell.json
rename to tests/unit/expected_responses_api_request/context_management_and_shell.json
diff --git a/tests/unit/integrations/compression_interception/test_compression_interception_handler.py b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py
index e66cd654f93..d7a1d6f14e1 100644
--- a/tests/unit/integrations/compression_interception/test_compression_interception_handler.py
+++ b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py
@@ -528,7 +528,7 @@ async def test_pre_call_hook_no_compression_records_no_savings(monkeypatch):
@pytest.mark.asyncio
async def test_pre_call_hook_counts_tokens_off_the_event_loop():
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/unit/litellm_core_utils/conftest.py b/tests/unit/litellm_core_utils/conftest.py
new file mode 100644
index 00000000000..2a1e1f6382c
--- /dev/null
+++ b/tests/unit/litellm_core_utils/conftest.py
@@ -0,0 +1,15 @@
+import importlib
+
+import pytest
+
+from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault
+
+
+@pytest.fixture(autouse=True, scope="session")
+def bundled_tiktoken_cache() -> None:
+ importlib.import_module("litellm.litellm_core_utils.default_encoding")
+
+
+@pytest.fixture
+def secret_vault_factory() -> type[FakeSecretVault]:
+ return FakeSecretVault
diff --git a/tests/test_litellm/litellm_core_utils/event_loop_lag.py b/tests/unit/litellm_core_utils/event_loop_lag.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/event_loop_lag.py
rename to tests/unit/litellm_core_utils/event_loop_lag.py
diff --git a/tests/unit/litellm_core_utils/fake_secret_vault.py b/tests/unit/litellm_core_utils/fake_secret_vault.py
new file mode 100644
index 00000000000..75e9d16e9ed
--- /dev/null
+++ b/tests/unit/litellm_core_utils/fake_secret_vault.py
@@ -0,0 +1,67 @@
+from litellm.litellm_core_utils.cli_keyring import (
+ KeyringDiscardsWrites,
+ KeyringUnreachable,
+ KeyringUnusable,
+ SecretErase,
+ SecretErased,
+ SecretFound,
+ SecretMissing,
+ SecretRead,
+ SecretStored,
+ SecretStranded,
+ SecretWrite,
+)
+
+
+class FakeSecretVault:
+ """In-memory stand-in for the OS keychain, injected wherever CLI credential storage is exercised.
+
+ `available=False` models a keychain that is locked or has no backend, `writable=False` one that
+ refuses to store, `erasable=False` one that will not release what it already holds, and `failure`
+ picks which unusable state those report. `discards=True` is keyring's null backend, which answers
+ reads and erases like any other yet keeps nothing it is given, so only writes report it.
+ """
+
+ def __init__(
+ self,
+ blob: str | None = None,
+ *,
+ available: bool = True,
+ writable: bool = True,
+ erasable: bool = True,
+ discards: bool = False,
+ failure: KeyringUnusable = KeyringUnreachable(),
+ ) -> None:
+ self.blob: str | None = blob
+ self.available: bool = available
+ self.writable: bool = writable
+ self.erasable: bool = erasable
+ self.discards: bool = discards
+ self.failure: KeyringUnusable = failure
+ self.reads: int = 0
+ self.writes: list[str] = []
+ self.erases: int = 0
+
+ def read(self) -> SecretRead:
+ self.reads += 1
+ if not self.available:
+ return self.failure
+ return SecretMissing() if self.blob is None else SecretFound(self.blob)
+
+ def write(self, blob: str) -> SecretWrite:
+ self.writes.append(blob)
+ if not (self.available and self.writable):
+ return self.failure
+ if self.discards:
+ return KeyringDiscardsWrites()
+ self.blob = blob
+ return SecretStored()
+
+ def erase(self) -> SecretErase:
+ self.erases += 1
+ if not self.available:
+ return self.failure
+ if not self.erasable:
+ return SecretStranded() if self.blob is not None else SecretErased()
+ self.blob = None
+ return SecretErased()
diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/__init__.py b/tests/unit/litellm_core_utils/llm_cost_calc/__init__.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_response_utils/__init__.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/__init__.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_responses_cache_cost_breakdown.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py
rename to tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py
diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py
rename to tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py
diff --git a/tests/test_litellm/litellm_core_utils/messages_with_counts.py b/tests/unit/litellm_core_utils/messages_with_counts.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/messages_with_counts.py
rename to tests/unit/litellm_core_utils/messages_with_counts.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/__init__.py b/tests/unit/litellm_core_utils/prompt_templates/__init__.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/__init__.py
rename to tests/unit/litellm_core_utils/prompt_templates/__init__.py
diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py b/tests/unit/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py
rename to tests/unit/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py
diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py
rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
diff --git a/tests/test_litellm/rust_bridge/__init__.py b/tests/unit/litellm_core_utils/specialty_caches/__init__.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/__init__.py
rename to tests/unit/litellm_core_utils/specialty_caches/__init__.py
diff --git a/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py b/tests/unit/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py
rename to tests/unit/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py
diff --git a/tests/test_litellm/litellm_core_utils/test_agentic_followup_kwargs.py b/tests/unit/litellm_core_utils/test_agentic_followup_kwargs.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_agentic_followup_kwargs.py
rename to tests/unit/litellm_core_utils/test_agentic_followup_kwargs.py
diff --git a/tests/test_litellm/litellm_core_utils/test_anthropic_dedup_factory.py b/tests/unit/litellm_core_utils/test_anthropic_dedup_factory.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_anthropic_dedup_factory.py
rename to tests/unit/litellm_core_utils/test_anthropic_dedup_factory.py
diff --git a/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py b/tests/unit/litellm_core_utils/test_api_route_to_call_types.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py
rename to tests/unit/litellm_core_utils/test_api_route_to_call_types.py
diff --git a/tests/test_litellm/litellm_core_utils/test_audio_utils.py b/tests/unit/litellm_core_utils/test_audio_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_audio_utils.py
rename to tests/unit/litellm_core_utils/test_audio_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_aws_partition.py b/tests/unit/litellm_core_utils/test_aws_partition.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_aws_partition.py
rename to tests/unit/litellm_core_utils/test_aws_partition.py
diff --git a/tests/test_litellm/litellm_core_utils/test_bedrock_converse_dedup_factory.py b/tests/unit/litellm_core_utils/test_bedrock_converse_dedup_factory.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_bedrock_converse_dedup_factory.py
rename to tests/unit/litellm_core_utils/test_bedrock_converse_dedup_factory.py
diff --git a/tests/test_litellm/litellm_core_utils/test_bug_report.py b/tests/unit/litellm_core_utils/test_bug_report.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_bug_report.py
rename to tests/unit/litellm_core_utils/test_bug_report.py
diff --git a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py
rename to tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py
diff --git a/tests/test_litellm/litellm_core_utils/test_classifier_logging.py b/tests/unit/litellm_core_utils/test_classifier_logging.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_classifier_logging.py
rename to tests/unit/litellm_core_utils/test_classifier_logging.py
diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/unit/litellm_core_utils/test_cli_token_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_cli_token_utils.py
rename to tests/unit/litellm_core_utils/test_cli_token_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py b/tests/unit/litellm_core_utils/test_cloud_storage_security.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py
rename to tests/unit/litellm_core_utils/test_cloud_storage_security.py
diff --git a/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py b/tests/unit/litellm_core_utils/test_codestral_provider_routing.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py
rename to tests/unit/litellm_core_utils/test_codestral_provider_routing.py
diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_core_helpers.py
rename to tests/unit/litellm_core_utils/test_core_helpers.py
diff --git a/tests/test_litellm/litellm_core_utils/test_coroutine_checker.py b/tests/unit/litellm_core_utils/test_coroutine_checker.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_coroutine_checker.py
rename to tests/unit/litellm_core_utils/test_coroutine_checker.py
diff --git a/tests/test_litellm/litellm_core_utils/test_dd_tracing.py b/tests/unit/litellm_core_utils/test_dd_tracing.py
similarity index 85%
rename from tests/test_litellm/litellm_core_utils/test_dd_tracing.py
rename to tests/unit/litellm_core_utils/test_dd_tracing.py
index b55ade5225d..30cae45e250 100644
--- a/tests/test_litellm/litellm_core_utils/test_dd_tracing.py
+++ b/tests/unit/litellm_core_utils/test_dd_tracing.py
@@ -55,18 +55,6 @@ def test_dd_tracer_when_package_not_exists():
assert result == "test"
-def test_null_tracer_context_manager():
- """
- Test that the context manager works without raising exceptions when should_use_dd_tracer is False
- """
- with patch("litellm.litellm_core_utils.dd_tracing.should_use_dd_tracer", False):
- # Test that the context manager works without raising exceptions
- with dd_tracer.trace("test_operation") as span:
- # Test that we can call methods on the null span
- span.finish()
- assert True # If we get here without exceptions, the test passes
-
-
def test_should_use_dd_tracer():
"""
Test that the should_use_dd_tracer function works as expected
diff --git a/tests/test_litellm/litellm_core_utils/test_decode_special_tokens.py b/tests/unit/litellm_core_utils/test_decode_special_tokens.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_decode_special_tokens.py
rename to tests/unit/litellm_core_utils/test_decode_special_tokens.py
diff --git a/tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py b/tests/unit/litellm_core_utils/test_dot_notation_indexing.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py
rename to tests/unit/litellm_core_utils/test_dot_notation_indexing.py
diff --git a/tests/test_litellm/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_duration_parser.py
rename to tests/unit/litellm_core_utils/test_duration_parser.py
diff --git a/tests/test_litellm/litellm_core_utils/test_error_normalization.py b/tests/unit/litellm_core_utils/test_error_normalization.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_error_normalization.py
rename to tests/unit/litellm_core_utils/test_error_normalization.py
diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py
rename to tests/unit/litellm_core_utils/test_exception_mapping_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py b/tests/unit/litellm_core_utils/test_extract_base64_image.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_extract_base64_image.py
rename to tests/unit/litellm_core_utils/test_extract_base64_image.py
diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py
rename to tests/unit/litellm_core_utils/test_fallback_generalizations.py
diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_utils.py b/tests/unit/litellm_core_utils/test_fallback_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_fallback_utils.py
rename to tests/unit/litellm_core_utils/test_fallback_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/unit/litellm_core_utils/test_get_litellm_params.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_get_litellm_params.py
rename to tests/unit/litellm_core_utils/test_get_litellm_params.py
diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/unit/litellm_core_utils/test_get_llm_provider_endpoint_match.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
rename to tests/unit/litellm_core_utils/test_get_llm_provider_endpoint_match.py
diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_get_llm_provider_logic.py
rename to tests/unit/litellm_core_utils/test_get_llm_provider_logic.py
diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py
rename to tests/unit/litellm_core_utils/test_get_model_cost_map.py
diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/unit/litellm_core_utils/test_get_supported_openai_params.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py
rename to tests/unit/litellm_core_utils/test_get_supported_openai_params.py
diff --git a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_health_check_helpers.py
rename to tests/unit/litellm_core_utils/test_health_check_helpers.py
diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/unit/litellm_core_utils/test_image_handling.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_image_handling.py
rename to tests/unit/litellm_core_utils/test_image_handling.py
diff --git a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/unit/litellm_core_utils/test_initialize_dynamic_callback_params.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py
rename to tests/unit/litellm_core_utils/test_initialize_dynamic_callback_params.py
diff --git a/tests/test_litellm/litellm_core_utils/test_internal_call_metadata.py b/tests/unit/litellm_core_utils/test_internal_call_metadata.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_internal_call_metadata.py
rename to tests/unit/litellm_core_utils/test_internal_call_metadata.py
diff --git a/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py b/tests/unit/litellm_core_utils/test_json_fragment_accumulator.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py
rename to tests/unit/litellm_core_utils/test_json_fragment_accumulator.py
diff --git a/tests/test_litellm/litellm_core_utils/test_json_schema_validation.py b/tests/unit/litellm_core_utils/test_json_schema_validation.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_json_schema_validation.py
rename to tests/unit/litellm_core_utils/test_json_schema_validation.py
diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_litellm_logging.py
rename to tests/unit/litellm_core_utils/test_litellm_logging.py
diff --git a/tests/test_litellm/litellm_core_utils/test_llm_judge.py b/tests/unit/litellm_core_utils/test_llm_judge.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_llm_judge.py
rename to tests/unit/litellm_core_utils/test_llm_judge.py
diff --git a/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py b/tests/unit/litellm_core_utils/test_llm_request_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_llm_request_utils.py
rename to tests/unit/litellm_core_utils/test_llm_request_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_logging_utils.py b/tests/unit/litellm_core_utils/test_logging_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_logging_utils.py
rename to tests/unit/litellm_core_utils/test_logging_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_logging_worker.py
rename to tests/unit/litellm_core_utils/test_logging_worker.py
diff --git a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py b/tests/unit/litellm_core_utils/test_max_streaming_duration.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py
rename to tests/unit/litellm_core_utils/test_max_streaming_duration.py
diff --git a/tests/test_litellm/litellm_core_utils/test_model_param_helper.py b/tests/unit/litellm_core_utils/test_model_param_helper.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_model_param_helper.py
rename to tests/unit/litellm_core_utils/test_model_param_helper.py
diff --git a/tests/test_litellm/litellm_core_utils/test_model_response_utils.py b/tests/unit/litellm_core_utils/test_model_response_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_model_response_utils.py
rename to tests/unit/litellm_core_utils/test_model_response_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_private_json.py b/tests/unit/litellm_core_utils/test_private_json.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_private_json.py
rename to tests/unit/litellm_core_utils/test_private_json.py
diff --git a/tests/test_litellm/litellm_core_utils/test_provider_affinity.py b/tests/unit/litellm_core_utils/test_provider_affinity.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_provider_affinity.py
rename to tests/unit/litellm_core_utils/test_provider_affinity.py
diff --git a/tests/test_litellm/litellm_core_utils/test_provider_specific_headers.py b/tests/unit/litellm_core_utils/test_provider_specific_headers.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_provider_specific_headers.py
rename to tests/unit/litellm_core_utils/test_provider_specific_headers.py
diff --git a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py b/tests/unit/litellm_core_utils/test_ptu_pricing.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_ptu_pricing.py
rename to tests/unit/litellm_core_utils/test_ptu_pricing.py
diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py b/tests/unit/litellm_core_utils/test_realtime_errors.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_realtime_errors.py
rename to tests/unit/litellm_core_utils/test_realtime_errors.py
diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_realtime_streaming.py
rename to tests/unit/litellm_core_utils/test_realtime_streaming.py
diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_redact_messages.py
rename to tests/unit/litellm_core_utils/test_redact_messages.py
diff --git a/tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py b/tests/unit/litellm_core_utils/test_request_timeout_resolver.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py
rename to tests/unit/litellm_core_utils/test_request_timeout_resolver.py
diff --git a/tests/test_litellm/litellm_core_utils/test_retry_after_headers.py b/tests/unit/litellm_core_utils/test_retry_after_headers.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_retry_after_headers.py
rename to tests/unit/litellm_core_utils/test_retry_after_headers.py
diff --git a/tests/test_litellm/litellm_core_utils/test_safe_divide_seconds.py b/tests/unit/litellm_core_utils/test_safe_divide_seconds.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_safe_divide_seconds.py
rename to tests/unit/litellm_core_utils/test_safe_divide_seconds.py
diff --git a/tests/test_litellm/litellm_core_utils/test_safe_json_dumps.py b/tests/unit/litellm_core_utils/test_safe_json_dumps.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_safe_json_dumps.py
rename to tests/unit/litellm_core_utils/test_safe_json_dumps.py
diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/unit/litellm_core_utils/test_sensitive_data_masker.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py
rename to tests/unit/litellm_core_utils/test_sensitive_data_masker.py
diff --git a/tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py b/tests/unit/litellm_core_utils/test_sentry_scrubbing.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py
rename to tests/unit/litellm_core_utils/test_sentry_scrubbing.py
diff --git a/tests/test_litellm/litellm_core_utils/test_served_output_texts.py b/tests/unit/litellm_core_utils/test_served_output_texts.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_served_output_texts.py
rename to tests/unit/litellm_core_utils/test_served_output_texts.py
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_cursor.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py
rename to tests/unit/litellm_core_utils/test_streaming_chunk_builder_cursor.py
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py
rename to tests/unit/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
rename to tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py
similarity index 99%
rename from tests/test_litellm/litellm_core_utils/test_streaming_handler.py
rename to tests/unit/litellm_core_utils/test_streaming_handler.py
index 3af79c709cc..6557811b530 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/unit/litellm_core_utils/test_streaming_handler.py
@@ -4900,7 +4900,7 @@ class TestStableStreamingResponseId:
@pytest.mark.asyncio
async def test_async_stream_without_usage_counts_tokens_off_the_event_loop():
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_overhead.py b/tests/unit/litellm_core_utils/test_streaming_overhead.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_streaming_overhead.py
rename to tests/unit/litellm_core_utils/test_streaming_overhead.py
diff --git a/tests/test_litellm/litellm_core_utils/test_thread_pool_executor.py b/tests/unit/litellm_core_utils/test_thread_pool_executor.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_thread_pool_executor.py
rename to tests/unit/litellm_core_utils/test_thread_pool_executor.py
diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py
new file mode 100644
index 00000000000..b1a14e61b96
--- /dev/null
+++ b/tests/unit/litellm_core_utils/test_token_counter.py
@@ -0,0 +1,1441 @@
+#### What this tests ####
+# This tests litellm.token_counter.token_counter() function
+import asyncio
+import base64
+import importlib
+import threading
+import time
+import traceback
+from concurrent.futures import Future, wait
+from typing import Final
+from unittest.mock import MagicMock
+
+import anyio.to_thread
+import pytest
+import tiktoken
+
+from unittest.mock import AsyncMock, patch
+
+import litellm
+from litellm import decode, encode, get_modified_max_tokens
+from litellm import token_counter as token_counter_old
+import litellm.constants
+from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+from litellm.litellm_core_utils.asyncify import asyncify
+from litellm.litellm_core_utils.token_counter import (
+ _get_exact_count_function,
+ _get_extrapolating_count_function,
+ _get_tiktoken_count_function,
+ calculate_img_tokens,
+ high_detail_image_token_upper_bound,
+ offload_token_count,
+)
+from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
+from tests.large_text import text
+from tests.unit.litellm_core_utils.event_loop_lag import (
+ assert_loop_stayed_free,
+ timed_with_loop_lags,
+ warm_tokenizer,
+)
+from tests.unit.litellm_core_utils.messages_with_counts import (
+ MESSAGES_TEXT,
+ MESSAGES_WITH_IMAGES,
+ MESSAGES_WITH_TOOLS,
+)
+
+
+def token_counter_both_assert_same(**args):
+ new = token_counter_new(**args)
+ old = token_counter_old(**args)
+ assert new == old, f"New token counter {new} does not match old token counter {old}"
+ return new
+
+
+## Choose which token_counter the test will use.
+
+# token_counter = token_counter_new
+# token_counter = token_counter_old
+token_counter = token_counter_both_assert_same
+
+
+def test_token_counter_basic():
+ assert (
+ token_counter(
+ model="claude-2",
+ messages=[
+ {
+ "role": "user",
+ "content": "This is a long message that definitely exceeds the token limit.",
+ }
+ ],
+ )
+ == 19
+ )
+
+
+def test_token_counter_large_repeated_text_is_fast():
+ messages = [{"role": "user", "content": [{"type": "text", "text": "A" * 1024 * 1024}]}]
+
+ start_time = time.perf_counter()
+ tokens = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages)
+ elapsed = time.perf_counter() - start_time
+
+ assert elapsed < 2, f"Token counting took too long: {elapsed:.2f}s"
+ assert tokens > 0
+
+
+@pytest.mark.parametrize(
+ "text",
+ [
+ "Short text",
+ "This is a normal message with punctuation, numbers, and a few words.",
+ ],
+)
+def test_token_counter_short_text_matches_tiktoken(text):
+ encoding = tiktoken.get_encoding("cl100k_base")
+ expected = len(encoding.encode(text, disallowed_special=()))
+
+ assert token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text) == expected
+
+
+def test_token_counter_default_encoding_matches_cl100k():
+ encoding: Final = tiktoken.get_encoding("cl100k_base")
+ expected: Final = len(encoding.encode("hello world", disallowed_special=()))
+
+ assert token_counter_new(model=None, text="hello world") == expected
+
+
+def test_token_counter_text_over_chunk_boundary_stays_close_to_tiktoken():
+ text = ("The quick brown fox jumps over the lazy dog. " * 30)[:1025]
+ encoding = tiktoken.get_encoding("cl100k_base")
+ expected = len(encoding.encode(text, disallowed_special=()))
+
+ actual = token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text)
+
+ assert abs(actual - expected) <= 4
+
+
+@pytest.mark.parametrize(
+ "configured",
+ ["0", "-1", "-1024", "not-an-int", "", " ", "999999999", "inf", "1e9"],
+)
+def test_invalid_chunk_size_config_stays_usable(monkeypatch, configured):
+ """A misconfigured chunk size must not raise, count zero, or restore the quadratic encode cost."""
+ monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", configured)
+ try:
+ reloaded = importlib.reload(litellm.constants)
+ chunk_size = reloaded.TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
+ assert 1 <= chunk_size <= reloaded.TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS
+
+ encoding = tiktoken.get_encoding("cl100k_base")
+ count_tokens = _get_tiktoken_count_function(
+ lambda text: len(encoding.encode(text, disallowed_special=())),
+ chunk_size=chunk_size,
+ )
+ assert count_tokens("The quick brown fox jumps over the lazy dog. " * 40) > 0
+ finally:
+ monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS")
+ importlib.reload(litellm.constants)
+
+
+def test_valid_chunk_size_config_is_honoured(monkeypatch):
+ monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", "2048")
+ try:
+ assert importlib.reload(litellm.constants).TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS == 2048
+ finally:
+ monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS")
+ importlib.reload(litellm.constants)
+
+
+async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free():
+ warm_tokenizer("claude-fable-5")
+
+ tokens, took, lags = await timed_with_loop_lags(
+ lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100)
+ )
+
+ assert tokens > 0
+ assert_loop_stayed_free(took, lags)
+
+
+@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500])
+def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int):
+ count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk))
+ front_heavy: Final = "a" * 1_000 + "b" * 4_000
+ exact: Final = 1_000 + len(front_heavy)
+
+ estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy)
+
+ assert abs(estimate - exact) <= exact // 100
+ assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars
+
+
+def test_count_at_or_below_the_cap_is_exact():
+ count_exactly: Final = MagicMock(side_effect=len)
+
+ assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000
+ assert count_exactly.call_args_list == [(("a" * 5_000,),)]
+
+
+class _SlowEncoder:
+ def __init__(self) -> None:
+ self._lock: Final = threading.Lock()
+ self.in_flight = 0
+ self.peak_in_flight = 0
+
+ def encode_batch_fast(self, texts: list[str]) -> list[list[int]]:
+ with self._lock:
+ self.in_flight += 1
+ self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
+ time.sleep(0.1)
+ with self._lock:
+ self.in_flight -= 1
+ return [[0] * len(text) for text in texts]
+
+
+@pytest.mark.asyncio
+async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool():
+ encoder: Final = _SlowEncoder()
+ count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder})
+ shared_pool: Final = anyio.to_thread.current_default_thread_limiter()
+ burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+
+ async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]:
+ if counting.done():
+ return ()
+ await asyncio.sleep(0.01)
+ return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting))
+
+ counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst))))
+ borrowed: Final = await shared_pool_borrowed_until_done(counting)
+
+ assert await counting == [3] * burst
+ assert len(borrowed) > 1 and max(borrowed) == 0
+ assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+
+
+def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None:
+ def slow_count(counted: str) -> int:
+ time.sleep(0.1)
+ return len(counted)
+
+ result.set_result(asyncio.run(offload_token_count(slow_count)(text)))
+
+
+def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process():
+ loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+ results: Final = tuple(Future[int]() for _ in range(loops))
+ threads: Final = tuple(
+ threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True)
+ for size, result in enumerate(results, start=1)
+ )
+ for thread in threads:
+ thread.start()
+
+ _, pending = wait(results, timeout=5)
+
+ assert not pending
+ assert tuple(result.result() for result in results) == tuple(range(1, loops + 1))
+
+
+@pytest.mark.parametrize(
+ ("configured", "expected"),
+ [("8", 8), ("0", 4), ("not-an-int", 4)],
+)
+def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int):
+ monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured)
+ try:
+ assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected
+ finally:
+ monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS")
+ importlib.reload(litellm.constants)
+
+
+def test_token_counter_applies_the_default_cap():
+ max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS
+ prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars]
+ over_the_cap: Final = prose + "a" * 200_000
+ exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap)
+
+ estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap)
+
+ assert estimate != exact
+ assert abs(estimate - exact) <= exact // 100
+
+
+@pytest.mark.parametrize(
+ ("configured", "expected"),
+ [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)],
+)
+def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int):
+ monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured)
+ try:
+ assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected
+ finally:
+ monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS")
+ importlib.reload(litellm.constants)
+
+
+def test_token_counter_with_prefix():
+ messages = [
+ {"role": "user", "content": "Who won the world cup in 2022?"},
+ {"role": "assistant", "content": "Argentina", "prefix": True},
+ ]
+ tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
+ assert tokens == 22, f"Expected 22 tokens, got {tokens}"
+
+
+def test_token_counter_normal_plus_function_calling():
+ messages = [
+ {"role": "system", "content": "System prompt"},
+ {"role": "user", "content": "content1"},
+ {"role": "assistant", "content": "content2"},
+ {"role": "user", "content": "conten3"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {
+ "id": "call_E0lOb1h6qtmflUyok4L06TgY",
+ "function": {
+ "arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}',
+ "name": "SearchInternet",
+ },
+ "type": "function",
+ }
+ ],
+ },
+ {
+ "tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY",
+ "role": "tool",
+ "name": "SearchInternet",
+ "content": "tool content",
+ },
+ ]
+ tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
+ assert tokens == 80
+
+
+# test_token_counter_normal_plus_function_calling()
+
+
+def test_token_counter_legacy_function_call_counts_arguments():
+ """
+ Regression for VERIA-492 (Token-counter function_call bypass).
+
+ The legacy OpenAI assistant `function_call` field carries arbitrary text in
+ `arguments`. Before the fix, `_count_messages` had no branch for
+ `function_call` and fell through to the unsupported-key `continue`, so an
+ assistant turn could smuggle unlimited text past `token_counter` and the
+ proxy `/utils/token_counter` endpoint (and downstream pre-call budget /
+ `get_modified_max_tokens` math). After the fix it must be counted the
+ same as the equivalent `tool_calls` payload.
+ """
+ long_arg = "A" * 4000
+ fc_messages = [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": None,
+ "function_call": {"name": "search", "arguments": long_arg},
+ },
+ ]
+ tc_messages = [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {
+ "id": "call_1",
+ "type": "function",
+ "function": {"name": "search", "arguments": long_arg},
+ }
+ ],
+ },
+ ]
+ fc_tokens = token_counter(model="gpt-3.5-turbo", messages=fc_messages)
+ tc_tokens = token_counter(model="gpt-3.5-turbo", messages=tc_messages)
+ assert fc_tokens == tc_tokens, (
+ f"function_call arguments must count like tool_calls arguments; "
+ f"got function_call={fc_tokens}, tool_calls={tc_tokens}"
+ )
+ assert fc_tokens > 500, f"4000-char arguments payload must contribute real tokens, got {fc_tokens}"
+
+
+@pytest.mark.parametrize(
+ "message_count_pair",
+ MESSAGES_TEXT,
+)
+def test_token_counter_textonly(message_count_pair):
+ counted_tokens = token_counter(
+ model="gpt-35-turbo", messages=[message_count_pair["message"]]
+ )
+ assert counted_tokens == message_count_pair["count"]
+
+
+@pytest.mark.parametrize(
+ "message_count_pair",
+ MESSAGES_TEXT,
+)
+def test_token_counter_count_response_tokens(message_count_pair):
+ counted_tokens = token_counter(
+ model="gpt-35-turbo",
+ messages=[message_count_pair["message"]],
+ count_response_tokens=True,
+ )
+ # 3 tokens are not added because of count_response_tokens=True
+ expected = message_count_pair["count"] - 3
+ assert counted_tokens == expected
+
+
+@pytest.mark.parametrize(
+ "message_count_pair",
+ MESSAGES_WITH_IMAGES,
+)
+def test_token_counter_with_images(message_count_pair):
+ counted_tokens = token_counter(
+ model="gpt-4o", messages=[message_count_pair["message"]]
+ )
+ assert counted_tokens == message_count_pair["count"]
+
+
+@pytest.mark.parametrize(
+ "message_count_pair",
+ MESSAGES_WITH_TOOLS,
+)
+def test_token_counter_with_tools(message_count_pair):
+ counted_tokens = token_counter(
+ model="gpt-35-turbo",
+ messages=[message_count_pair["system_message"]],
+ tools=message_count_pair["tools"],
+ tool_choice=message_count_pair["tool_choice"],
+ )
+ expected_tokens = message_count_pair["count"]
+ actual_diff = counted_tokens - expected_tokens
+
+ if "count-tolerate" in message_count_pair:
+ if message_count_pair["count-tolerate"] == counted_tokens:
+ pass # expected
+ else:
+ tolerated_diff = message_count_pair["count-tolerate"] - expected_tokens
+ assert (
+ actual_diff <= tolerated_diff
+ ), f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by {tolerated_diff} in the over-counting direction."
+ if actual_diff != tolerated_diff:
+ raise NeedsToleranceUpdateError(
+ f"SOMETHING BROKEN GOT FIXED! THIS is good! Adjust 'count-tolerate' from {message_count_pair['count-tolerate']} to {counted_tokens}"
+ )
+
+ else:
+ assert (
+ expected_tokens == counted_tokens
+ ), f"Expected {expected_tokens} tokens, got {counted_tokens}."
+
+
+class NeedsToleranceUpdateError(Exception):
+ """Custom exception to mark tests that have improved"""
+
+ pass
+
+
+# test_tokenizers()
+
+
+def test_encoding_and_decoding():
+ try:
+ sample_text = "Hellö World, this is my input string!"
+ # openai encoding + decoding
+ openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text)
+ openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens)
+
+ assert openai_text == sample_text
+
+ # claude encoding + decoding
+ claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text)
+
+ claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens)
+
+ assert claude_text == sample_text
+
+ # cohere encoding + decoding
+ cohere_tokens = encode(model="command-nightly", text=sample_text)
+ cohere_text = decode(model="command-nightly", tokens=cohere_tokens)
+
+ assert cohere_text == sample_text
+
+ # llama2 encoding + decoding
+ llama2_tokens = encode(model="meta-llama/Llama-2-7b-chat", text=sample_text)
+ llama2_text = decode(model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens)
+
+ assert llama2_text == sample_text
+ except Exception as e:
+ pytest.fail(f"An exception occured: {e}\n{traceback.format_exc()}")
+
+
+# test_encoding_and_decoding()
+
+
+# test_gpt_vision_token_counting()
+
+
+@pytest.mark.parametrize(
+ "model",
+ [
+ "gpt-4-vision-preview",
+ "gpt-4o",
+ "claude-3-opus-20240229",
+ "command-nightly",
+ "mistral/mistral-tiny",
+ ],
+)
+def test_load_test_token_counter(model):
+ """
+ Token count large prompt 100 times.
+
+ Assert time taken is < 1.5s.
+ """
+ import tiktoken
+
+ messages = [{"role": "user", "content": text}] * 10
+
+ start_time = time.time()
+ for _ in range(10):
+ _ = token_counter(model=model, messages=messages)
+ # enc.encode("".join(m["content"] for m in messages))
+
+ end_time = time.time()
+
+ total_time = end_time - start_time
+ print("model={}, total test time={}".format(model, total_time))
+ assert total_time < 10, f"Total encoding time > 10s, {total_time}"
+
+
+@pytest.mark.parametrize(
+ "model, base_model, input_tokens, user_max_tokens, expected_value",
+ [
+ ("random-model", "random-model", 1024, 1024, 1024),
+ ("gpt-3.5-turbo", "gpt-3.5-turbo", 4000, 5000, 4096), # model max output = 4096
+ ],
+)
+def test_get_modified_max_tokens(
+ model, base_model, input_tokens, user_max_tokens, expected_value
+):
+ """
+ - Test when max_output is not known => expect user_max_tokens
+ - Test when max_output == max_input,
+ - input > max_output, no max_tokens => expect None
+ - input + max_tokens > max_output => expect remainder
+ - input + max_tokens < max_output => expect max_tokens
+ - Test when max_tokens > max_output => expect max_output
+ """
+ args = locals()
+ import litellm
+
+ litellm.token_counter = MagicMock()
+
+ def _mock_token_counter(*args, **kwargs):
+ return input_tokens
+
+ litellm.token_counter.side_effect = _mock_token_counter
+ print(f"_mock_token_counter: {_mock_token_counter()}")
+ messages = [{"role": "user", "content": "Hello world!"}]
+
+ calculated_value = get_modified_max_tokens(
+ model=model,
+ base_model=base_model,
+ messages=messages,
+ user_max_tokens=user_max_tokens,
+ buffer_perc=0,
+ buffer_num=0,
+ )
+
+ if expected_value is None:
+ assert calculated_value is None
+ else:
+ assert (
+ calculated_value == expected_value
+ ), "Got={}, Expected={}, Params={}".format(
+ calculated_value, expected_value, args
+ )
+
+
+def test_empty_tools():
+ messages = [{"role": "user", "content": "hey, how's it going?", "tool_calls": None}]
+
+ result = token_counter(
+ messages=messages,
+ )
+
+ print(result)
+
+
+@pytest.mark.skip(
+ reason="Skipping this test temporarily because it relies on a function being called that I am removing."
+)
+def test_gpt_4o_token_counter():
+ with patch.object(
+ litellm.utils, "openai_token_counter", new=MagicMock()
+ ) as mock_client:
+ token_counter(
+ model="gpt-4o-2024-05-13", messages=[{"role": "user", "content": "Hey!"}]
+ )
+
+ mock_client.assert_called()
+
+
+@pytest.mark.parametrize(
+ "img_url",
+ [
+ "https://example.com/test-image.png",
+ "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAL0AAAC9CAMAAADRCYwCAAAAh1BMVEX///8AAAD8/Pz5+fkEBAT39/cJCQn09PRNTU3y8vIMDAwzMzPe3t7v7+8QEBCOjo7FxcXR0dHn5+elpaWGhoYYGBivr686OjocHBy0tLQtLS1TU1PY2Ni6urpaWlpERER3d3ecnJxoaGiUlJRiYmIlJSU4ODhBQUFycnKAgIDBwcFnZ2chISE7EjuwAAAI/UlEQVR4nO1caXfiOgz1bhJIyAJhX1JoSzv8/9/3LNlpYd4rhX6o4/N8Z2lKM2cURZau5JsQEhERERERERERERERERERERHx/wBjhDPC3OGN8+Cc5JeMuheaETSdO8vZFyCScHtmz2CsktoeMn7rLM1u3h0PMAEhyYX7v/Q9wQvoGdB0hlbzm45lEq/wd6y6G9aezvBk9AXwp1r3LHJIRsh6s2maxaJpmvqgvkC7WFS3loUnaFJtKRVUCEoV/RpCnHRvAsesVQ1hw+vd7Mpo+424tLs72NplkvQgcdrsvXkW/zJWqH/fA0FT84M/xnQJt4to3+ZLuanbM6X5lfXKHosO9COgREqpCR5i86pf2zPS7j9tTj+9nO7bQz3+xGEyGW9zqgQ1tyQ/VsxEDvce/4dcUPNb5OD9yXvR4Z2QisuP0xiGWPnemgugU5q/troHhGEjIF5sTOyW648aC0TssuaaCEsYEIkGzjWXOp3A0vVsf6kgRyqaDk+T7DIVWrb58b2tT5xpUucKwodOD/5LbrZC1ws6YSaBZJ/8xlh+XZSYXaMJ2ezNqjB3IPXuehPcx2U6b4t1dS/xNdFzguUt8ie7arnPeyCZroxLHzGgGdqVcspwafizPWEXBee+9G1OaufGdvNng/9C+gwgZ3PH3r87G6zXTZ5D5De2G2DeFoANXfbACkT+fxBQ22YFsTTJF9hjFVO6VbqxZXko4WJ8s52P4PnuxO5KRzu0/hlix1ySt8iXjgaQ+4IHPA9nVzNkdduM9LFT/Aacj4FtKrHA7iAw602Vnht6R8Vq1IOS+wNMKLYqayAYfRuufQPGeGb7sZogQQoLZrGPgZ6KoYn70Iw30O92BNEDpvwouCFn6wH2uS+EhRb3WF/HObZk3HuxfRQM3Y/Of/VH0n4MKNHZDiZvO9+m/ABALfkOcuar/7nOo7B95ACGVAFaz4jMiJwJhdaHBkySmzlGTu82gr6FSTik2kJvLnY9nOd/D90qcH268m3I/cgI1xg1maE5CuZYaWLH+UHANCIck0yt7Mx5zBm5vVHXHwChsZ35kKqUpmo5Svq5/fzfAI5g2vDtFPYo1HiEA85QrDeGm9g//LG7K0scO3sdpj2CBDgCa+0OFs0bkvVgnnM/QBDwllOMm+cN7vMSHlB7Uu4haHKaTwgGkv8tlK+hP8fzmFuK/RQTpaLPWvbd58yWIo66HHM0OsPoPhVqmtaEVL7N+wYcTLTbb0DLdgp23Eyy2VYJ2N7bkLFAAibtoLPe5sLt6Oa2bvU+zyeMa8wrixO0gRTn9tO9NCSThTLGqcqtsDvphlfmx/cPBZVvw24jg1LE2lPuEo35Mhi58U0I/Ga8n5w+NS8i34MAQLos5B1u0xL1ZvCVYVRw/Fs2q53KLaXJMWwOZZ/4MPYV19bAHmgGDKB6f01xoeJKFbl63q9J34KdaVNPJWztQyRkzA3KNs1AdAEDowMxh10emXTCx75CkurtbY/ZpdNDGdsn2UcHKHsQ8Ai3WZi48IfkvtjOhsLpuIRSKZTX9FA4o+0d6o/zOWqQzVJMynL9NsxhSJOaourq6nBVQBueMSyubsX2xHrmuABZN2Ns9jr5nwLFlLF/2R6atjW/67Yd11YQ1Z+kA9Zk9dPTM/o6dVo6HHVgC0JR8oUfmI93T9u3gvTG94bAH02Y5xeqRcjuwnKCK6Q2+ajl8KXJ3GSh22P3Zfx6S+n008ROhJn+JRIUVu6o7OXl8w1SeyhuqNDwNI7SjbK08QrqPxS95jy4G7nCXVq6G3HNu0LtK5J0e226CfC005WKK9sVvfxI0eUbcnzutfhWe3rpZHM0nZ/ny/N8tanKYlQ6VEW5Xuym8yV1zZX58vwGhZp/5tFfhybZabdbrQYOs8F+xEhmPsb0/nki6kIyVvzZzUASiOrTfF+Sj9bXC7DoJxeiV8tjQL6loSd0yCx7YyB6rPdLx31U2qCG3F/oXIuDuqd6LFO+4DNIJuxFZqSsU0ea88avovFnWKRYFYRQDfCfcGaBCLn4M4A1ntJ5E57vicwqq2enaZEF5nokCYu9TbKqCC5yCDfL+GhLxT4w4xEJs+anqgou8DOY2q8FMryjb2MehC1dRJ9s4g9NXeTwPkWON4RH+FhIe0AWR/S9ekvQ+t70XHeimGF78LzuU7d7PwrswdIG2VpgF8C53qVQsTDtBJc4CdnkQPbnZY9mbPdDFra3PCXBBQ5QBn2aQqtyhvlyYM4Hb2/mdhsxCUen04GZVvIJZw5PAamMOmjzq8Q+dzAKLXDQ3RUZItWsg4t7W2DP+JDrJDymoMH7E5zQtuEpG03GTIjGCW3LQqOYEsXgFc78x76NeRwY6SNM+IfQoh6myJKRBIcLYxZcwscJ/gI2isTBty2Po9IkYzP0/SS4hGlxRjFAG5z1Jt1LckiB57yWvo35EaolbvA+6fBa24xodL2YjsPpTnj3JgJOqhcgOeLVsYYwoK0wjY+m1D3rGc40CukkaHnkEjarlXrF1B9M6ECQ6Ow0V7R7N4G3LfOHAXtymoyXOb4QhaYHJ/gNBJUkxclpSs7DNcgWWDDmM7Ke5MJpGuioe7w5EOvfTunUKRzOh7G2ylL+6ynHrD54oQO3//cN3yVO+5qMVsPZq0CZIOx4TlcJ8+Vz7V5waL+7WekzUpRFMTnnTlSCq3X5usi8qmIleW/rit1+oQZn1WGSU/sKBYEqMNh1mBOc6PhK8yCfKHdUNQk8o/G19ZPTs5MYfai+DLs5vmee37zEyyH48WW3XA6Xw6+Az8lMhci7N/KleToo7PtTKm+RA887Kqc6E9dyqL/QPTugzMHLbLZtJKqKLFfzVWRNJ63c+95uWT/F7R0U5dDVvuS409AJXhJvD0EwWaWdW8UN11u/7+umaYjT8mJtzZwP/MD4r57fihiHlC5fylHfaqnJdro+Dr7DajvO+vi2EwyD70s8nCH71nzIO1l5Zl+v1DMCb5ebvCMkGHvobXy/hPumGLyX0218/3RyD1GRLOuf9u/OGQyDmto32yMiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIiIv7GP8YjWPR/czH2AAAAAElFTkSuQmCC",
+ ],
+)
+def test_img_url_token_counter(img_url, monkeypatch):
+ """
+ Verify get_image_dimensions returns valid (width, height) for both an
+ HTTPS URL and a base64 data URI. The HTTPS branch is exercised with a
+ mocked HTTP fetch so the test is hermetic - it can't break when a
+ third-party image URL goes away.
+ """
+ import base64
+ from litellm.litellm_core_utils.token_counter import get_image_dimensions
+
+ # Minimal valid 1x1 PNG, served by the mocked safe_get for the URL case.
+ _tiny_png = base64.b64decode(
+ "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
+ )
+
+ if img_url.startswith(("http://", "https://")):
+
+ class _FakeResponse:
+ headers = {"Content-Length": str(len(_tiny_png))}
+
+ def read(self):
+ return _tiny_png
+
+ monkeypatch.setattr(
+ "litellm.litellm_core_utils.token_counter.safe_get",
+ lambda client, url, **kw: _FakeResponse(),
+ )
+
+ width, height = get_image_dimensions(data=img_url)
+
+ print(width, height)
+
+ assert width is not None
+ assert height is not None
+
+
+def test_token_encode_disallowed_special():
+ encode(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>")
+ token_counter(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>")
+
+
+def test_token_counter():
+ try:
+ messages = [{"role": "user", "content": "hi how are you what time is it"}]
+ tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
+ print("gpt-35-turbo")
+ print(tokens)
+ assert tokens > 0
+
+ tokens = token_counter(model="claude-2", messages=messages)
+ print("claude-2")
+ print(tokens)
+ assert tokens > 0
+
+ tokens = token_counter(model="gemini/chat-bison", messages=messages)
+ print("gemini/chat-bison")
+ print(tokens)
+ assert tokens > 0
+
+ tokens = token_counter(model="ollama/llama2", messages=messages)
+ print("ollama/llama2")
+ print(tokens)
+ assert tokens > 0
+
+ tokens = token_counter(model="anthropic.claude-instant-v1", messages=messages)
+ print("anthropic.claude-instant-v1")
+ print(tokens)
+ assert tokens > 0
+ except Exception as e:
+ pytest.fail(f"Error occurred: {e}")
+
+
+import unittest
+
+from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding
+
+# Clear the cache at module load to ensure clean state
+_load_huggingface_tokenizer.cache_clear()
+
+
+class TestTokenizerSelection(unittest.TestCase):
+ def setUp(self):
+ """Clear the LRU cache before each test method.
+
+ The HuggingFace tokenizers behind _select_tokenizer_helper are cached with
+ @lru_cache, which can cause cache hits from previous tests when running with
+ --dist=loadscope (tests from same file run on same worker).
+ """
+ _load_huggingface_tokenizer.cache_clear()
+
+ @patch("litellm.utils.tokenizer_dispatch.from_pretrained")
+ def test_llama3_tokenizer_api_failure(self, mock_from_pretrained):
+ # Setup mock to raise an error
+ mock_from_pretrained.side_effect = Exception("Failed to load tokenizer")
+
+ # Test with llama-3 model
+ result = _select_tokenizer_helper("llama-3-7b")
+
+ # Verify the attempt to load Llama-3 tokenizer
+ mock_from_pretrained.assert_called_once_with("Xenova/llama-3-tokenizer")
+
+ # Verify fallback to OpenAI tokenizer
+ self.assertEqual(result["type"], "openai_tokenizer")
+ self.assertEqual(result["tokenizer"], encoding)
+
+ @patch("litellm.utils.tokenizer_dispatch.from_pretrained")
+ def test_cohere_tokenizer_api_failure(self, mock_from_pretrained):
+ # Setup mock to raise an error
+ mock_from_pretrained.side_effect = Exception("Failed to load tokenizer")
+
+ # Add Cohere model to the list for testing
+ litellm.cohere_models = ["command-r-v1"]
+
+ # Test with Cohere model
+ result = _select_tokenizer_helper("command-r-v1")
+
+ # Verify the attempt to load Cohere tokenizer
+ mock_from_pretrained.assert_called_once_with(
+ "Xenova/c4ai-command-r-v01-tokenizer"
+ )
+
+ # Verify fallback to OpenAI tokenizer
+ self.assertEqual(result["type"], "openai_tokenizer")
+ self.assertEqual(result["tokenizer"], encoding)
+
+ @patch("litellm.utils.tokenizer_dispatch.anthropic")
+ def test_claude_tokenizer_api_failure(self, mock_anthropic):
+ # Setup mock to raise an error
+ mock_anthropic.side_effect = Exception("Failed to load tokenizer")
+
+ # Add Claude model to the list for testing
+ litellm.anthropic_models = ["claude-2"]
+
+ # Test with Claude model
+ result = _select_tokenizer_helper("claude-2")
+
+ # Verify the attempt to load Claude tokenizer
+ mock_anthropic.assert_called_once_with()
+
+ # Verify fallback to OpenAI tokenizer
+ self.assertEqual(result["type"], "openai_tokenizer")
+ self.assertEqual(result["tokenizer"], encoding)
+
+ @patch("litellm.utils.tokenizer_dispatch.from_pretrained")
+ def test_llama2_tokenizer_api_failure(self, mock_from_pretrained):
+ # Setup mock to raise an error
+ mock_from_pretrained.side_effect = Exception("Failed to load tokenizer")
+
+ # Test with Llama-2 model
+ result = _select_tokenizer_helper("llama-2-7b")
+
+ # Verify the attempt to load Llama-2 tokenizer
+ mock_from_pretrained.assert_called_once_with(
+ "hf-internal-testing/llama-tokenizer"
+ )
+
+ # Verify fallback to OpenAI tokenizer
+ self.assertEqual(result["type"], "openai_tokenizer")
+ self.assertEqual(result["tokenizer"], encoding)
+
+ @patch("litellm.utils._return_huggingface_tokenizer")
+ def test_disable_hf_tokenizer_download(self, mock_return_huggingface_tokenizer):
+ monkeypatch = pytest.MonkeyPatch()
+ monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
+ try:
+ result = _select_tokenizer_helper("grok-32r22r")
+ mock_return_huggingface_tokenizer.assert_not_called()
+ assert result["type"] == "openai_tokenizer"
+ assert result["tokenizer"] == encoding
+ finally:
+ monkeypatch.undo()
+
+
+def test_token_counter_with_anthropic_tool_use():
+ """
+ Test that _count_anthropic_content() correctly handles tool_use blocks.
+
+ Validates that:
+ - 'name' field is counted (string)
+ - 'input' field is counted (dict serialized to string)
+ - Metadata fields ('type', 'id') are skipped
+ """
+ messages = [
+ {"role": "user", "content": "What's the weather in San Francisco?"},
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "text", "text": "I'll check the weather for you."},
+ {
+ "type": "tool_use",
+ "id": "toolu_01234567890", # Should be skipped
+ "name": "get_weather", # Should be counted
+ "input": { # Should be counted (serialized)
+ "location": "San Francisco, CA",
+ "unit": "fahrenheit",
+ },
+ },
+ ],
+ },
+ ]
+
+ tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
+ assert tokens > 0, f"Expected positive token count, got {tokens}"
+ # Should count: user message + "I'll check" text + "get_weather" name + input dict
+ assert (
+ tokens > 15
+ ), f"Expected reasonable token count for message with tool_use, got {tokens}"
+
+
+def test_token_counter_with_anthropic_tool_result():
+ """
+ Test that _count_anthropic_content() correctly handles tool_result blocks.
+
+ Validates that:
+ - 'content' field (when string) is counted
+ - Metadata fields ('type', 'tool_use_id') are skipped
+ - Full conversation with tool_use → tool_result flow works
+ """
+ messages = [
+ {"role": "user", "content": "What's the weather in San Francisco?"},
+ {
+ "role": "assistant",
+ "content": [
+ {
+ "type": "tool_use",
+ "id": "toolu_01234567890",
+ "name": "get_weather",
+ "input": {"location": "San Francisco, CA"},
+ }
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_01234567890", # Should be skipped
+ "content": "The weather in San Francisco is 65°F and sunny.", # Should be counted
+ }
+ ],
+ },
+ ]
+
+ tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
+ assert tokens > 0, f"Expected positive token count, got {tokens}"
+ assert (
+ tokens > 25
+ ), f"Expected reasonable token count for conversation with tool_result, got {tokens}"
+
+
+def test_token_counter_with_nested_tool_result():
+ """
+ Test that _count_anthropic_content() recursively handles nested content lists.
+
+ Validates that:
+ - tool_result with 'content' as a list (not string) is handled
+ - Nested content blocks are recursively counted via _count_content_list()
+ - TypedDict inference correctly identifies list fields
+ """
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_01234567890",
+ "content": [ # Nested list - should recursively count
+ {
+ "type": "text",
+ "text": "The weather in San Francisco is 65°F and sunny.",
+ },
+ {"type": "text", "text": "UV index is moderate."},
+ ],
+ }
+ ],
+ }
+ ]
+
+ tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
+ assert tokens > 0, f"Expected positive token count, got {tokens}"
+ # Should count both nested text blocks
+ assert (
+ tokens > 15
+ ), f"Expected reasonable token count for nested tool_result, got {tokens}"
+
+
+def test_token_counter_tool_use_and_result_combined():
+ """
+ Test dynamic field inference with multiple tool_use and tool_result blocks.
+
+ Validates that:
+ - Multiple tool_use blocks in same message are handled
+ - Multiple tool_result blocks in same message are handled
+ - skip_fields correctly filters metadata across all blocks
+ - Full realistic conversation flow works end-to-end
+ """
+ messages = [
+ {
+ "role": "user",
+ "content": "What's the weather in San Francisco and New York?",
+ },
+ {
+ "role": "assistant",
+ "content": [
+ {
+ "type": "text",
+ "text": "I'll check the weather in both cities for you.",
+ },
+ {
+ "type": "tool_use",
+ "id": "toolu_01A",
+ "name": "get_weather",
+ "input": {"location": "San Francisco, CA"},
+ },
+ {
+ "type": "tool_use",
+ "id": "toolu_01B",
+ "name": "get_weather",
+ "input": {"location": "New York, NY"},
+ },
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_01A",
+ "content": "San Francisco: 65°F, sunny",
+ },
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_01B",
+ "content": "New York: 45°F, cloudy",
+ },
+ ],
+ },
+ {
+ "role": "assistant",
+ "content": "The weather in San Francisco is 65°F and sunny, while New York is cooler at 45°F and cloudy.",
+ },
+ ]
+
+ tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
+ assert tokens > 0, f"Expected positive token count, got {tokens}"
+ # Should count all text, tool names, inputs, and results
+ assert (
+ tokens > 60
+ ), f"Expected substantial token count for full tool conversation, got {tokens}"
+
+
+def test_token_counter_with_image_url():
+ """
+ Test that _count_image_tokens() correctly handles image_url content blocks.
+
+ Validates that:
+ - image_url as dict with 'url' and 'detail' is handled
+ - image_url as string is handled
+ - 'detail' field validation works ('low', 'high', 'auto')
+ - calculate_img_tokens is called with correct parameters
+ """
+ # Test with dict format (detail: low)
+ messages_dict = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "What's in this image?"},
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "https://example.com/image.jpg",
+ "detail": "low", # Should use low token count (85 base tokens)
+ },
+ },
+ ],
+ }
+ ]
+
+ tokens_dict = token_counter(
+ model="gpt-3.5-turbo",
+ messages=messages_dict,
+ use_default_image_token_count=True, # Avoid actual HTTP request
+ )
+ assert tokens_dict > 0, f"Expected positive token count, got {tokens_dict}"
+ assert tokens_dict > 85, f"Expected at least base image tokens, got {tokens_dict}"
+
+ # Test with string format (defaults to auto/low)
+ messages_str = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "image_url",
+ "image_url": "https://example.com/image.jpg", # String format
+ }
+ ],
+ }
+ ]
+
+ tokens_str = token_counter(
+ model="gpt-3.5-turbo", messages=messages_str, use_default_image_token_count=True
+ )
+ assert (
+ tokens_str > 0
+ ), f"Expected positive token count for string image_url, got {tokens_str}"
+
+ # Test invalid detail value raises error
+ messages_invalid = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "https://example.com/image.jpg",
+ "detail": "invalid", # Should raise ValueError
+ },
+ }
+ ],
+ }
+ ]
+
+ with pytest.raises(ValueError, match="Invalid detail value") as exc_info:
+ token_counter(model="gpt-3.5-turbo", messages=messages_invalid)
+ e = exc_info.value
+ assert "Invalid detail value" in str(
+ e
+ ), f"Expected detail validation error, got: {e}"
+
+
+def test_token_counter_with_thinking_content():
+ """
+ Test that _count_content_list() correctly handles Claude's extended thinking content blocks.
+
+ Validates that:
+ - 'thinking' content type is recognized and counted
+ - 'thinking' text field is counted
+ - 'signature' field is skipped (opaque signature blob)
+ - Full conversation with thinking blocks works
+ """
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "text",
+ "text": "Analyze this complex problem: who came first, chicken or egg",
+ }
+ ],
+ },
+ {
+ "role": "assistant",
+ "content": [
+ {
+ "type": "thinking",
+ "thinking": "This is actually a fascinating question that touches on philosophy, biology, and semantics. Let me break this down: The egg came first from an evolutionary biology perspective.",
+ "signature": "EqcLCkYICxgCKkCrqu6lP...", # Should be skipped
+ },
+ {
+ "type": "text",
+ "text": "# The Chicken-or-Egg Question: A Multi-Layered Answer\n\n## **The Short Answer: The Egg Came First**",
+ },
+ ],
+ },
+ {"role": "user", "content": [{"type": "text", "text": "Thanks"}]},
+ ]
+
+ tokens = token_counter(
+ model="anthropic/claude-sonnet-4-5-20250929", messages=messages
+ )
+ assert tokens > 0, f"Expected positive token count, got {tokens}"
+ # Should count: user message + thinking text + response text + "Thanks"
+ # The thinking text alone is ~30 tokens, plus other content should be > 50 total
+ assert (
+ tokens > 50
+ ), f"Expected substantial token count for message with thinking, got {tokens}"
+
+ # Test that thinking block without 'thinking' field doesn't crash (edge case)
+ messages_no_thinking = [
+ {
+ "role": "assistant",
+ "content": [
+ {
+ "type": "thinking",
+ # No 'thinking' field - should count as 0 tokens
+ "signature": "EqcLCkYICxgCKkCrqu6lP...",
+ },
+ {"type": "text", "text": "Response"},
+ ],
+ }
+ ]
+
+ tokens_no_thinking = token_counter(
+ model="anthropic/claude-sonnet-4-5-20250929", messages=messages_no_thinking
+ )
+ assert (
+ tokens_no_thinking > 0
+ ), f"Expected positive token count even with empty thinking, got {tokens_no_thinking}"
+ # Should only count "Response" and message overhead
+ assert (
+ tokens_no_thinking < 15
+ ), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}"
+
+
+def test_token_counter_with_redacted_thinking_content():
+ """
+ A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in
+ for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking
+ block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the
+ prompt_caching pre-call check stop pinning the deployment that held the cached prefix.
+ """
+ model = "anthropic/claude-sonnet-4-5-20250929"
+ reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."}
+ redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30}
+ user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]}
+ follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]}
+
+ without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up]
+ with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up]
+
+ assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block)
+
+def test_token_counter_with_tool_reference_block():
+ """
+ Regression test: a message containing an Anthropic tool-search
+ `tool_reference` content block must NOT raise.
+
+ Before the fix, token_counter raised
+ `Invalid content item type: tool_reference`. On the streaming
+ anthropic_messages proxy path this nulled response_cost and caused the
+ SpendLogs row to be dropped, silently undercounting cost. token_counter
+ must instead count the referenced tool name and return a positive count.
+ """
+ messages = [
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "text", "text": "Let me look up the right tool."},
+ {"type": "tool_reference", "tool_name": "search_knowledge_base"},
+ ],
+ }
+ ]
+
+ # Must not raise, and must produce a positive token count.
+ tokens = token_counter_new(
+ model="anthropic/claude-sonnet-4-5-20250929", messages=messages
+ )
+ assert tokens > 0, f"Expected positive token count, got {tokens}"
+
+ # A tool_reference with no/empty tool_name must also be handled gracefully.
+ messages_empty = [
+ {
+ "role": "assistant",
+ "content": [{"type": "tool_reference", "tool_name": ""}],
+ }
+ ]
+ tokens_empty = token_counter_new(
+ model="anthropic/claude-sonnet-4-5-20250929", messages=messages_empty
+ )
+ assert tokens_empty >= 0
+
+
+def test_count_content_list_rejects_unknown_type():
+ """
+ An unrecognized content block type must raise, and the error message must
+ enumerate the supported types (including `tool_reference`). This pins the
+ catch-all contract so a future block type isn't silently dropped.
+ """
+ from litellm.litellm_core_utils.token_counter import _count_content_list
+
+ with pytest.raises(ValueError, match='Error getting number of tokens from content list: Invalid') as exc_info:
+ _count_content_list(
+ count_function=len,
+ content_list=[{"type": "totally_unknown_block"}],
+ use_default_image_token_count=False,
+ default_token_count=None,
+ )
+
+ message = str(exc_info.value)
+ assert "Invalid content item type: totally_unknown_block" in message
+ assert "tool_reference" in message
+
+
+@pytest.mark.parametrize(
+ "source",
+ [
+ {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="},
+ {"type": "url", "url": "https://example.com/image.png"},
+ {"type": "file", "file_id": "file-abc123"},
+ ],
+ ids=["base64", "url", "file"],
+)
+def test_token_counter_with_anthropic_image_block(source: dict[str, str]):
+ """Anthropic `image` blocks must count for every source variant, not raise `Invalid content item type` (which the router's context-window pre-call check swallows into an unfiltered dispatch)."""
+ from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT
+
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "What is in this image?"},
+ {"type": "image", "source": source},
+ ],
+ }
+ ]
+
+ tokens = token_counter(
+ model="anthropic/claude-sonnet-4-5-20250929",
+ messages=messages,
+ use_default_image_token_count=True,
+ )
+ assert tokens > DEFAULT_IMAGE_TOKEN_COUNT, (
+ f"Expected the image block to contribute tokens, got {tokens}"
+ )
+
+
+def test_anthropic_image_block_matches_equivalent_image_url():
+ """An Anthropic `image` block prices identically to the OpenAI `image_url` carrying the same bytes."""
+ anthropic_messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/png",
+ "data": "iVBORw0KGgo=",
+ },
+ }
+ ],
+ }
+ ]
+ openai_messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "image_url",
+ "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
+ }
+ ],
+ }
+ ]
+
+ anthropic_tokens = token_counter(
+ model="anthropic/claude-sonnet-4-5-20250929", messages=anthropic_messages
+ )
+ openai_tokens = token_counter(
+ model="anthropic/claude-sonnet-4-5-20250929", messages=openai_messages
+ )
+ assert anthropic_tokens == openai_tokens
+
+
+def test_anthropic_image_block_nested_in_tool_result():
+ """An `image` block nested in a `tool_result.content` list is counted through the same recursion."""
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_01",
+ "content": [
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/png",
+ "data": "iVBORw0KGgo=",
+ },
+ }
+ ],
+ }
+ ],
+ }
+ ]
+
+ tokens = token_counter(
+ model="anthropic/claude-sonnet-4-5-20250929",
+ messages=messages,
+ use_default_image_token_count=True,
+ )
+ assert tokens > 0
+
+
+@pytest.mark.parametrize(
+ ("source", "expected"),
+ [
+ ({"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQ"}, "data:image/jpeg;base64,/9j/4AAQ"),
+ ({"type": "url", "url": "https://example.com/image.png"}, "https://example.com/image.png"),
+ ({"type": "file", "file_id": "file-abc123"}, ""),
+ ],
+ ids=["base64", "url", "file"],
+)
+def test_anthropic_image_source_resolves_to_what_the_image_pricer_reads(source: dict[str, str], expected: str):
+ """base64 sources become a data URI, url sources pass through, file sources resolve to an empty string."""
+ from litellm.litellm_core_utils.token_counter import _anthropic_image_source_data
+
+ assert _anthropic_image_source_data(source) == expected
+
+
+def test_anthropic_image_block_with_empty_base64_data():
+ """A base64 source with empty `data` prices as an image rather than raising."""
+ from litellm.litellm_core_utils.token_counter import _count_content_list
+
+ tokens = _count_content_list(
+ count_function=len,
+ content_list=[
+ {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}}
+ ],
+ use_default_image_token_count=False,
+ default_token_count=None,
+ )
+ assert tokens > 0
+
+
+def test_anthropic_image_block_without_source_raises():
+ """An `image` block with no `source` raises, matching the OpenAI `image_url`-without-`url` behavior."""
+ from litellm.litellm_core_utils.token_counter import _count_content_list
+
+ with pytest.raises(ValueError, match="Error getting number of tokens from content list"):
+ _count_content_list(
+ count_function=len,
+ content_list=[{"type": "image"}],
+ use_default_image_token_count=False,
+ default_token_count=None,
+ )
+
+ # ... and `default_token_count`, the caller's opt-out from raising, still wins.
+ assert (
+ _count_content_list(
+ count_function=len,
+ content_list=[{"type": "image"}],
+ use_default_image_token_count=False,
+ default_token_count=7,
+ )
+ == 7
+ )
+
+
+def _count_user_content(content: list[dict]) -> int:
+ from litellm.litellm_core_utils.token_counter import token_counter
+
+ return token_counter(
+ model="anthropic/claude-fable-5",
+ messages=[{"role": "user", "content": content}],
+ use_default_image_token_count=True,
+ )
+
+
+@pytest.mark.parametrize(
+ "source",
+ [
+ {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
+ {"type": "url", "url": "https://example.com/report.pdf"},
+ {"type": "file", "file_id": "file-abc123"},
+ ],
+ ids=["base64", "url", "file"],
+)
+def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]):
+ """A `document` whose bytes can't be tokenized locally is priced like an `image`, not raised on."""
+ prompt = {"type": "text", "text": "Summarize this file."}
+
+ assert _count_user_content([prompt, {"type": "document", "source": source}]) == _count_user_content(
+ [prompt, {"type": "image", "source": source}]
+ )
+
+
+def test_anthropic_document_block_text_sources_count_their_text():
+ """`text` and `content` document sources count the text they carry, as inline text blocks would."""
+ prompt = {"type": "text", "text": "Summarize this file."}
+ body = {"type": "text", "text": "Revenue grew eleven percent while churn fell to two percent."}
+ picture = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}}
+
+ text_source = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": body["text"]}}
+ assert _count_user_content([prompt, text_source]) == _count_user_content([prompt, body])
+
+ string_content = {"type": "document", "source": {"type": "content", "content": body["text"]}}
+ assert _count_user_content([prompt, string_content]) == _count_user_content([prompt, body])
+
+ block_content = {"type": "document", "source": {"type": "content", "content": [body, picture]}}
+ assert _count_user_content([prompt, block_content]) == _count_user_content([prompt, body, picture])
+
+
+def test_anthropic_document_title_and_context_add_their_tokens():
+ prompt = {"type": "text", "text": "Summarize this file."}
+ source = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}
+ described = {"type": "document", "source": source, "title": "Q3 board packet", "context": "Shared by finance"}
+
+ assert _count_user_content([prompt, described]) == _count_user_content(
+ [
+ prompt,
+ {"type": "text", "text": "Q3 board packet"},
+ {"type": "text", "text": "Shared by finance"},
+ {"type": "document", "source": source},
+ ]
+ )
+
+
+def test_openai_file_block_prices_like_the_equivalent_anthropic_document():
+ """An inline `file` is a `document` in the chat-completions dialect, so it must price identically, not raise.
+
+ Before the fix `file` was missing from the content-block match even though `ChatCompletionFileObject`
+ is in the union this counter accepts, so every local count of a Responses `input_file` raised
+ `Invalid content item type: file` and surfaced as a 500 on /v1/responses/input_tokens.
+ """
+ prompt = {"type": "text", "text": "Summarize this file."}
+ inline_file = {
+ "type": "file",
+ "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"},
+ }
+ document = {
+ "type": "document",
+ "title": "report.pdf",
+ "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
+ }
+
+ assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document])
+ assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt])
+
+
+def test_openai_file_block_without_inline_bytes_counts_what_it_carries():
+ """A `file` block naming an uploaded file has no bytes to price, so it adds only the filename's tokens."""
+ prompt = {"type": "text", "text": "Summarize this file."}
+
+ by_id = {"type": "file", "file": {"file_id": "file-abc123"}}
+ assert _count_user_content([prompt, by_id]) == _count_user_content([prompt])
+
+ named = {"type": "file", "file": {"file_id": "file-abc123", "filename": "report.pdf"}}
+ assert _count_user_content([prompt, named]) == _count_user_content(
+ [prompt, {"type": "text", "text": "report.pdf"}]
+ )
+
+
+def _png_data_url(width: int, height: int) -> str:
+ ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big")
+ return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode()
+
+
+@pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)])
+def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None:
+ assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound()
+
+
+def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None:
+ assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound()
+ assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound()
diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter_tool.py b/tests/unit/litellm_core_utils/test_token_counter_tool.py
similarity index 93%
rename from tests/test_litellm/litellm_core_utils/test_token_counter_tool.py
rename to tests/unit/litellm_core_utils/test_token_counter_tool.py
index 9f8c1070a47..f61b7d335c1 100644
--- a/tests/test_litellm/litellm_core_utils/test_token_counter_tool.py
+++ b/tests/unit/litellm_core_utils/test_token_counter_tool.py
@@ -5,8 +5,8 @@ import pytest
# Use the same token_counter as the main test.
-from tests.test_litellm.litellm_core_utils.test_token_counter import token_counter
-from tests.test_litellm.litellm_core_utils.test_token_counter_tool_data import *
+from tests.unit.litellm_core_utils.test_token_counter import token_counter
+from tests.unit.litellm_core_utils.test_token_counter_tool_data import *
@pytest.mark.parametrize(
diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter_tool_data.py b/tests/unit/litellm_core_utils/test_token_counter_tool_data.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_token_counter_tool_data.py
rename to tests/unit/litellm_core_utils/test_token_counter_tool_data.py
diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py
new file mode 100644
index 00000000000..a9005ff6a86
--- /dev/null
+++ b/tests/unit/litellm_core_utils/test_tokenizer.py
@@ -0,0 +1,411 @@
+import copy
+import os
+import pickle
+import subprocess
+import sys
+from pathlib import Path
+from typing import Final, Literal
+
+import pytest
+import tiktoken
+from tokenizers import Tokenizer as ReferenceTokenizer
+
+import litellm
+from litellm.caching._embedding_router import truncate_embedding_input
+from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding
+from litellm.utils import claude_json_str
+from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
+
+
+OFFLINE_ENCODINGS: Final = ("cl100k_base", "o200k_base", "p50k_base", "p50k_edit", "o200k_harmony")
+UNICODE_TEXTS: Final = ("hello world", "café æ¼¢å— ðŸ™‚", "", "a\ud800b", "\ud83d\ude42", "🙂\ud83d\ude42\udfff", " " * 64)
+
+
+@pytest.mark.parametrize("name", OFFLINE_ENCODINGS)
+@pytest.mark.parametrize("text", UNICODE_TEXTS)
+def test_openai_encoding_matches_python_unicode_and_batches(name: str, text: str) -> None:
+ assert_openai_encoding_matches_python(name, text)
+
+
+def assert_openai_encoding_matches_python(name: str, text: str) -> None:
+ reference: Final = tiktoken.get_encoding(name)
+ encoding: Final = OpenAIEncoding.from_tiktoken(name)
+ expected: Final = reference.encode(text)
+
+ assert encoding.encode(text) == expected
+ assert encoding.count(text) == len(expected)
+ assert encoding.encode_batch([text], num_threads=2) == reference.encode_batch([text], num_threads=2)
+ assert encoding.encode_ordinary_batch([text]) == reference.encode_ordinary_batch([text])
+ assert encoding.decode_batch([expected]) == reference.decode_batch([expected])
+ assert encoding.decode_bytes_batch([expected]) == reference.decode_bytes_batch([expected])
+
+
+@pytest.mark.parametrize("allowed", (frozenset(), frozenset({"<|endoftext|>"}), "all"))
+@pytest.mark.parametrize("disallowed", (frozenset(), frozenset({"<|fim_prefix|>"}), "all"))
+def test_openai_special_token_options_match_python(
+ allowed: frozenset[str] | Literal["all"], disallowed: frozenset[str] | Literal["all"]
+) -> None:
+ reference: Final = tiktoken.get_encoding("cl100k_base")
+ encoding: Final = OpenAIEncoding.from_tiktoken(reference.name)
+ text: Final = "hello<|endoftext|><|fim_prefix|>world"
+ allowed_set: Final = reference.special_tokens_set if allowed == "all" else allowed
+ disallowed_set: Final = reference.special_tokens_set - allowed_set if disallowed == "all" else disallowed
+ if any(token in text for token in disallowed_set):
+ with pytest.raises(ValueError, match="disallowed special token"):
+ encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed)
+ return
+ assert encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed) == reference.encode(
+ text, allowed_special=allowed, disallowed_special=disallowed
+ )
+ assert encoding.special_tokens_set == reference.special_tokens_set
+ assert encoding.eot_token == reference.eot_token
+
+
+@pytest.mark.parametrize("errors", ("replace", "ignore", "backslashreplace", "strict"))
+def test_openai_partial_token_decoding_preserves_error_policy(errors: str) -> None:
+ reference: Final = tiktoken.get_encoding("cl100k_base")
+ encoding: Final = OpenAIEncoding.from_tiktoken(reference.name)
+ tokens: Final = reference.encode("🙂")[:1]
+ assert encoding.decode_bytes(tokens) == reference.decode_bytes(tokens)
+ if errors == "strict":
+ with pytest.raises(UnicodeDecodeError):
+ encoding.decode(tokens, errors=errors)
+ return
+ assert encoding.decode(tokens, errors=errors) == reference.decode(tokens, errors=errors)
+ assert encoding.decode_tokens_bytes(tokens) == reference.decode_tokens_bytes(tokens)
+
+
+def test_public_encoding_and_semantic_cache_preserve_truncated_unicode() -> None:
+ reference: Final = tiktoken.get_encoding(litellm.encoding.name)
+ text: Final = "🙂"
+ tokens: Final = reference.encode(text)
+
+ assert litellm.encoding.encode(text, disallowed_special=()) == tokens
+ assert litellm.encoding.encode_batch([text]) == [tokens]
+ assert litellm.decode(tokens=tokens[:1]) == reference.decode(tokens[:1])
+ assert truncate_embedding_input(text, "", 1) == reference.decode(tokens[:1])
+
+
+@pytest.mark.parametrize("add_special_tokens", (True, False))
+def test_huggingface_encoding_preserves_result_fields_and_serialization(add_special_tokens: bool) -> None:
+ reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
+ tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON)
+ expected: Final = reference.encode("Hello World", add_special_tokens=add_special_tokens)
+ actual: Final = tokenizer.encode("Hello World", add_special_tokens=add_special_tokens)
+
+ assert (actual.ids, actual.tokens, actual.type_ids, actual.offsets, actual.word_ids, actual.sequence_ids) == (
+ expected.ids,
+ expected.tokens,
+ expected.type_ids,
+ expected.offsets,
+ expected.word_ids,
+ expected.sequence_ids,
+ )
+ assert (actual.attention_mask, actual.special_tokens_mask, actual.n_sequences, len(actual)) == (
+ expected.attention_mask,
+ expected.special_tokens_mask,
+ expected.n_sequences,
+ len(expected),
+ )
+ assert copy.deepcopy(actual).ids == expected.ids
+ assert pickle.loads(pickle.dumps(actual)).offsets == expected.offsets
+ assert tokenizer.decode(actual.ids, skip_special_tokens=False) == reference.decode(
+ expected.ids, skip_special_tokens=False
+ )
+
+
+def test_huggingface_character_offsets_and_pretokenized_pairs_match_python() -> None:
+ reference: Final = ReferenceTokenizer.from_str(claude_json_str)
+ tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str)
+ text: Final = "café æ¼¢å— ðŸ™‚"
+ actual: Final = tokenizer.encode(text)
+ expected: Final = reference.encode(text)
+
+ assert actual.offsets == expected.offsets
+ assert actual.ids == expected.ids
+ assert (
+ tokenizer.encode(["hello", "world"], ["again"], is_pretokenized=True).ids
+ == reference.encode(["hello", "world"], ["again"], is_pretokenized=True).ids
+ )
+
+
+def test_huggingface_batches_apply_padding_across_inputs() -> None:
+ reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
+ reference.enable_padding(pad_id=0, pad_token="[UNK]")
+ tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str())
+ inputs: Final = ["Hello", ("Hello World", "World")]
+ expected: Final = reference.encode_batch(inputs)
+ actual: Final = tokenizer.encode_batch(inputs)
+ fast: Final = tokenizer.encode_batch_fast(inputs)
+
+ assert [(item.ids, item.attention_mask, item.offsets) for item in actual] == [
+ (item.ids, item.attention_mask, item.offsets) for item in expected
+ ]
+ assert [item.ids for item in fast] == [item.ids for item in expected]
+ assert tokenizer.decode_batch([item.ids for item in actual]) == reference.decode_batch(
+ [item.ids for item in expected]
+ )
+
+
+def test_caller_supplied_huggingface_tokenizer_preserves_public_encode_and_count() -> None:
+ tokenizer: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
+ custom: Final = {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
+ expected: Final = tokenizer.encode("Hello World").ids
+
+ assert litellm.encode(text="Hello World", custom_tokenizer=custom) == expected
+ assert litellm.token_counter(text="Hello World", custom_tokenizer=custom) == len(expected)
+ assert litellm.decode(tokens=expected, custom_tokenizer=custom) == "Hello World"
+
+
+def test_caller_supplied_tiktoken_treats_special_spellings_as_text() -> None:
+ tokenizer: Final = tiktoken.get_encoding("cl100k_base")
+ custom: Final = {"type": "openai_tokenizer", "tokenizer": tokenizer}
+ text: Final = "<|endoftext|>"
+
+ assert litellm.encode(text=text, custom_tokenizer=custom) == tokenizer.encode(text, disallowed_special=())
+
+
+def test_public_tokenizer_objects_survive_pickle_and_deepcopy(tmp_path: Path) -> None:
+ custom: Final = litellm.create_tokenizer(TOKENIZER_JSON)
+ tokenizer: Final = custom["tokenizer"]
+ path: Final = tmp_path / "tokenizer.json"
+ tokenizer.save(str(path))
+
+ assert copy.deepcopy(custom)["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids
+ assert (
+ pickle.loads(pickle.dumps(custom))["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids
+ )
+ assert HuggingFaceTokenizer.from_file(str(path)).encode("Hello World").ids == tokenizer.encode("Hello World").ids
+ assert copy.deepcopy(litellm.encoding).encode("hello") == litellm.encoding.encode("hello")
+ assert pickle.loads(pickle.dumps(litellm.encoding)).encode("hello") == litellm.encoding.encode("hello")
+
+
+@pytest.mark.parametrize("offline", ("0", "1"))
+def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None:
+ script: Final = """
+import json
+import sys
+from pathlib import Path
+sys.path.insert(0, sys.argv[1])
+import httpx
+import huggingface_hub
+from huggingface_hub.errors import LocalEntryNotFoundError
+import litellm
+payload = sys.argv[2].encode()
+offline = sys.argv[3] == "1"
+observed = []
+def handle(request):
+ assert not offline, "offline loading issued a request"
+ if request.url.path.endswith("/tokenizer.json"):
+ observed.append(request.headers.get("authorization"))
+ if request.headers.get("authorization") != "Bearer audit-fixture-token":
+ return httpx.Response(401)
+ return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"")
+if not offline:
+ huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle)))
+try:
+ tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"]
+except LocalEntryNotFoundError:
+ assert offline
+ assert observed == []
+else:
+ assert not offline
+ assert "Bearer audit-fixture-token" in observed
+ assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World"
+ assert tuple(Path(sys.argv[4]).rglob("tokenizer.json"))
+print("compatible")
+"""
+ result: Final = subprocess.run(
+ [
+ sys.executable,
+ "-I",
+ "-c",
+ script,
+ str(Path(litellm.__file__).parent.parent),
+ TOKENIZER_JSON,
+ offline,
+ str(tmp_path / "cache"),
+ ],
+ capture_output=True,
+ text=True,
+ timeout=30,
+ env={
+ **os.environ,
+ "HF_HOME": str(tmp_path / "home"),
+ "HF_HUB_CACHE": str(tmp_path / "cache"),
+ "HF_ENDPOINT": "http://127.0.0.1:9",
+ "HF_TOKEN": "audit-fixture-token",
+ "HF_HUB_OFFLINE": offline,
+ "HF_HUB_DISABLE_IMPLICIT_TOKEN": "0",
+ "LITELLM_LOCAL_MODEL_COST_MAP": "True",
+ },
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+ assert result.stdout.strip() == "compatible"
+
+
+@pytest.mark.parametrize("rust", (None, "0", "1"))
+def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None:
+ script: Final = """
+import importlib.abc
+import sys
+sys.path.insert(0, sys.argv[1])
+def reject_network(event, args):
+ if event == "socket.connect":
+ raise AssertionError("tokenizer attempted a network connection")
+sys.addaudithook(reject_network)
+class Block(importlib.abc.MetaPathFinder):
+ def find_spec(self, fullname, path=None, target=None):
+ if fullname == "litellm.rust_bridge._native":
+ raise ImportError("native extension is unavailable")
+sys.meta_path.insert(0, Block())
+import litellm
+from litellm.rust_bridge.tokenizer import get_encoding
+import tiktoken
+from tokenizers import Tokenizer
+assert isinstance(litellm.encoding, tiktoken.Encoding)
+for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"):
+ encoding = get_encoding(name)
+ text = "offline café æ¼¢å— ðŸ™‚" + " " * 64
+ assert encoding.decode(encoding.encode(text)) == text
+ids = litellm.encode(text="hello world")
+assert litellm.decode(tokens=ids) == "hello world"
+assert litellm.token_counter(model=None, text="hello world") == len(ids)
+custom = litellm.create_tokenizer(sys.argv[2])
+assert isinstance(custom["tokenizer"], Tokenizer)
+custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]")
+assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World"
+print("compatible")
+"""
+ result: Final = subprocess.run(
+ [sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON],
+ capture_output=True,
+ text=True,
+ timeout=30,
+ cwd=tmp_path,
+ env={
+ **{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"},
+ **({"LITELLM_RUST": rust} if rust is not None else {}),
+ "LITELLM_LOCAL_MODEL_COST_MAP": "True",
+ "TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"),
+ },
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+ assert result.stdout.strip() == "compatible"
+ assert not (tmp_path / "unused-tokenizer-cache").exists()
+
+
+@pytest.mark.parametrize("is_pretokenized", (False, True))
+def test_huggingface_batch_sequence_containers_match_python(is_pretokenized: bool) -> None:
+ reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
+ tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON)
+ inputs: Final = [["Hello", "World"], ("Hello", "World")]
+ actual: Final = tokenizer.encode_batch(inputs, is_pretokenized=is_pretokenized)
+ expected: Final = reference.encode_batch(inputs, is_pretokenized=is_pretokenized)
+ assert [(item.ids, item.type_ids, item.sequence_ids) for item in actual] == [
+ (item.ids, item.type_ids, item.sequence_ids) for item in expected
+ ]
+
+
+@pytest.mark.parametrize("name", ("cl100k_base", "o200k_base", "p50k_edit"))
+def test_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None:
+ assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name)
+
+
+def assert_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None:
+ reference: Final = tiktoken.get_encoding(name)
+ encoding: Final = OpenAIEncoding.from_tiktoken(name)
+ text: Final = "hello fanta"
+
+ assert repr(encoding) == repr(reference) == f""
+ assert (encoding.name, encoding.n_vocab, encoding.max_token_value) == (
+ reference.name,
+ reference.n_vocab,
+ reference.max_token_value,
+ )
+ assert encoding.token_byte_values() == reference.token_byte_values()
+ assert encoding.encode_single_token("hello") == reference.encode_single_token("hello")
+ assert encoding.encode_single_token(b"<|endoftext|>") == reference.eot_token
+ assert [encoding.is_special_token(token) for token in (0, reference.eot_token)] == [False, True]
+ assert encoding.decode_with_offsets(reference.encode(text)) == reference.decode_with_offsets(reference.encode(text))
+ assert encoding.encode_to_numpy(text).tolist() == reference.encode_to_numpy(text).tolist()
+ stable, completions = encoding.encode_with_unstable(text)
+ expected_stable, expected_completions = reference.encode_with_unstable(text)
+ assert (stable, sorted(completions)) == (expected_stable, sorted(expected_completions))
+ with pytest.raises(KeyError):
+ encoding.encode_single_token("<|not-a-token|>")
+
+
+def test_huggingface_tokenizer_exposes_the_tokenizers_vocabulary_surface() -> None:
+ reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
+ reference.enable_padding(pad_id=0, pad_token="[UNK]", length=4)
+ reference.enable_truncation(max_length=3, stride=1, strategy="only_first", direction="left")
+ tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str())
+
+ assert tokenizer.token_to_id("Hello") == reference.token_to_id("Hello") == 1
+ assert tokenizer.id_to_token(3) == reference.id_to_token(3) == "[BOS]"
+ assert tokenizer.id_to_token(99) is None
+ assert tokenizer.get_vocab() == reference.get_vocab()
+ assert tokenizer.get_vocab(with_added_tokens=False) == reference.get_vocab(with_added_tokens=False)
+ assert tokenizer.get_vocab_size() == reference.get_vocab_size() == 4
+ assert tokenizer.get_vocab_size(with_added_tokens=False) == reference.get_vocab_size(with_added_tokens=False)
+ added: Final = tokenizer.get_added_tokens_decoder()
+ expected_added: Final = reference.get_added_tokens_decoder()
+ assert {token_id: str(token) for token_id, token in added.items()} == {
+ token_id: str(token) for token_id, token in expected_added.items()
+ }
+ assert added[3].special == expected_added[3].special
+ assert tokenizer.num_special_tokens_to_add(False) == reference.num_special_tokens_to_add(False) == 1
+ assert tokenizer.num_special_tokens_to_add(True) == reference.num_special_tokens_to_add(True) == 0
+ assert tokenizer.padding == reference.padding
+ assert tokenizer.truncation == reference.truncation
+ assert tokenizer.encode_special_tokens == reference.encode_special_tokens is False
+ assert HuggingFaceTokenizer.from_buffer(TOKENIZER_JSON.encode()).encode("Hello").ids == [3, 1]
+ assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).padding is None
+ assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).truncation is None
+
+
+def test_huggingface_encoding_exposes_the_tokenizers_lookup_and_mutation_surface() -> None:
+ reference: Final = ReferenceTokenizer.from_str(claude_json_str)
+ tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str)
+ text: Final = "hello wide world"
+ actual: Final = tokenizer.encode(text, "again")
+ expected: Final = reference.encode(text, "again")
+
+ lookups: Final = (
+ lambda encoding: [encoding.token_to_chars(index) for index in range(len(encoding))],
+ lambda encoding: [encoding.token_to_word(index) for index in range(len(encoding))],
+ lambda encoding: [encoding.token_to_sequence(index) for index in range(len(encoding))],
+ lambda encoding: [encoding.char_to_token(position) for position in range(len(text))],
+ lambda encoding: [encoding.char_to_word(position) for position in range(len(text))],
+ lambda encoding: [encoding.char_to_token(position, 1) for position in range(5)],
+ lambda encoding: [encoding.word_to_tokens(word) for word in range(3)],
+ lambda encoding: [encoding.word_to_chars(word) for word in range(3)],
+ lambda encoding: [encoding.word_to_tokens(0, 1), encoding.word_to_chars(0, 1)],
+ )
+ for lookup in lookups:
+ assert lookup(actual) == lookup(expected)
+ assert repr(actual) == repr(expected)
+
+ actual.truncate(4, stride=1, direction="left")
+ expected.truncate(4, stride=1, direction="left")
+ assert (actual.ids, [item.ids for item in actual.overflowing]) == (
+ expected.ids,
+ [item.ids for item in expected.overflowing],
+ )
+ actual.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="")
+ expected.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="")
+ assert (actual.ids, actual.attention_mask, actual.type_ids, actual.tokens) == (
+ expected.ids,
+ expected.attention_mask,
+ expected.type_ids,
+ expected.tokens,
+ )
+ actual.set_sequence_id(3)
+ expected.set_sequence_id(3)
+ assert actual.sequence_ids == expected.sequence_ids
+ merged: Final = type(actual).merge([actual, tokenizer.encode("more")])
+ assert merged.ids == type(expected).merge([expected, reference.encode("more")]).ids
+ assert merged.offsets == type(expected).merge([expected, reference.encode("more")]).offsets
+ with pytest.raises(ValueError, match="direction"):
+ actual.pad(8, direction="sideways")
diff --git a/tests/test_litellm/litellm_core_utils/test_tool_search_spend_logging.py b/tests/unit/litellm_core_utils/test_tool_search_spend_logging.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_tool_search_spend_logging.py
rename to tests/unit/litellm_core_utils/test_tool_search_spend_logging.py
diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/unit/litellm_core_utils/test_url_utils.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_url_utils.py
rename to tests/unit/litellm_core_utils/test_url_utils.py
diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/unit/litellm_core_utils/test_xai_oauth_routing.py
similarity index 100%
rename from tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py
rename to tests/unit/litellm_core_utils/test_xai_oauth_routing.py
diff --git a/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py
index d835db63d83..5e2956b532a 100644
--- a/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py
+++ b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py
@@ -2733,7 +2733,7 @@ def test_build_summary_messages_keeps_midturn_system_correction_in_place():
async def test_threshold_check_counts_tokens_off_the_event_loop(monkeypatch):
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py
index a21c22cf5fa..9fad6ca5e66 100644
--- a/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py
+++ b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py
@@ -133,7 +133,7 @@ async def test_malformed_edit_entries_are_skipped():
async def test_sync_editor_counts_tokens_off_the_event_loop():
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/unit/llms/test_polling_url_origin_match.py b/tests/unit/llms/test_polling_url_origin_match.py
index ab5f41c757f..2df35131e3d 100644
--- a/tests/unit/llms/test_polling_url_origin_match.py
+++ b/tests/unit/llms/test_polling_url_origin_match.py
@@ -18,7 +18,7 @@ import pytest
# Azure DALL-E sync + async paths route through ``assert_same_origin``
# the same way as the case below. The helper itself is unit-tested in
-# ``tests/test_litellm/litellm_core_utils/test_url_utils.py``.
+# ``tests/unit/litellm_core_utils/test_url_utils.py``.
# ── Black Forest Labs polling ─────────────────────────────────────────────────
diff --git a/tests/test_litellm/rust_bridge/chat_completions/__init__.py b/tests/unit/responses/litellm_completion_transformation/__init__.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/chat_completions/__init__.py
rename to tests/unit/responses/litellm_completion_transformation/__init__.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_function_call_output_normalization.py b/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_function_call_output_normalization.py
rename to tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_handler.py
rename to tests/unit/responses/litellm_completion_transformation/test_handler.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_image_generation_output.py b/tests/unit/responses/litellm_completion_transformation/test_image_generation_output.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_image_generation_output.py
rename to tests/unit/responses/litellm_completion_transformation/test_image_generation_output.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py
rename to tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py
rename to tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py
rename to tests/unit/responses/litellm_completion_transformation/test_session_handler.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py
rename to tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
rename to tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py b/tests/unit/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py
similarity index 100%
rename from tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py
rename to tests/unit/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py
diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py
similarity index 100%
rename from tests/test_litellm/responses/mcp/test_chat_completions_handler.py
rename to tests/unit/responses/mcp/test_chat_completions_handler.py
diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py
similarity index 100%
rename from tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
rename to tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py
diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py
similarity index 100%
rename from tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py
rename to tests/unit/responses/mcp/test_mcp_streaming_iterator.py
diff --git a/tests/test_litellm/responses/test_additional_tools.py b/tests/unit/responses/test_additional_tools.py
similarity index 100%
rename from tests/test_litellm/responses/test_additional_tools.py
rename to tests/unit/responses/test_additional_tools.py
diff --git a/tests/test_litellm/responses/test_custom_tool_call.py b/tests/unit/responses/test_custom_tool_call.py
similarity index 100%
rename from tests/test_litellm/responses/test_custom_tool_call.py
rename to tests/unit/responses/test_custom_tool_call.py
diff --git a/tests/test_litellm/responses/test_dispatch.py b/tests/unit/responses/test_dispatch.py
similarity index 100%
rename from tests/test_litellm/responses/test_dispatch.py
rename to tests/unit/responses/test_dispatch.py
diff --git a/tests/test_litellm/responses/test_metadata_codex_callback.py b/tests/unit/responses/test_metadata_codex_callback.py
similarity index 100%
rename from tests/test_litellm/responses/test_metadata_codex_callback.py
rename to tests/unit/responses/test_metadata_codex_callback.py
diff --git a/tests/test_litellm/responses/test_no_duplicate_spend_logs.py b/tests/unit/responses/test_no_duplicate_spend_logs.py
similarity index 76%
rename from tests/test_litellm/responses/test_no_duplicate_spend_logs.py
rename to tests/unit/responses/test_no_duplicate_spend_logs.py
index c98b519ae67..7e4bef5812c 100644
--- a/tests/test_litellm/responses/test_no_duplicate_spend_logs.py
+++ b/tests/unit/responses/test_no_duplicate_spend_logs.py
@@ -15,35 +15,6 @@ import litellm
from litellm.integrations.custom_logger import CustomLogger
-def test_logging_object_not_popped():
- """
- Test that litellm_logging_obj is not popped from kwargs.
-
- This is a regression test for issue #15740. The bug was using
- kwargs.pop() which removed the logging object, causing duplicate
- spend logs for non-OpenAI providers.
- """
- import inspect
-
- from litellm.responses import main as responses_module
-
- # Get the source code of the responses function
- source = inspect.getsource(responses_module.responses)
-
- # Check that .pop("litellm_logging_obj") is NOT used
- # The bug was using kwargs.pop("litellm_logging_obj") which removes it
- assert 'kwargs.pop("litellm_logging_obj")' not in source, (
- "FAIL: Found kwargs.pop('litellm_logging_obj') in responses() function. "
- "This causes duplicate spend logs. Use kwargs.get('litellm_logging_obj') instead."
- )
-
- # Check that .get("litellm_logging_obj") IS used
- assert 'kwargs.get("litellm_logging_obj")' in source, (
- "FAIL: Expected kwargs.get('litellm_logging_obj') but not found. "
- "The logging object must be accessed with .get() not .pop() to prevent duplication."
- )
-
-
@pytest.mark.asyncio
async def test_async_no_duplicate_spend_logs():
"""
diff --git a/tests/test_litellm/responses/test_null_test_fix.py b/tests/unit/responses/test_null_test_fix.py
similarity index 100%
rename from tests/test_litellm/responses/test_null_test_fix.py
rename to tests/unit/responses/test_null_test_fix.py
diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py
similarity index 100%
rename from tests/test_litellm/responses/test_responses_api_bridge_flag.py
rename to tests/unit/responses/test_responses_api_bridge_flag.py
diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py
similarity index 99%
rename from tests/test_litellm/responses/test_responses_api_request_body.py
rename to tests/unit/responses/test_responses_api_request_body.py
index 98e74955c6f..b27401d693a 100644
--- a/tests/test_litellm/responses/test_responses_api_request_body.py
+++ b/tests/unit/responses/test_responses_api_request_body.py
@@ -20,7 +20,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
def _expected_dir() -> Path:
- """Path to expected_responses_api_request folder (sibling of test_litellm/responses)."""
+ """Path to expected_responses_api_request folder (sibling of tests/unit/responses)."""
return Path(__file__).resolve().parent.parent / "expected_responses_api_request"
diff --git a/tests/test_litellm/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py
similarity index 100%
rename from tests/test_litellm/responses/test_responses_prompt_management.py
rename to tests/unit/responses/test_responses_prompt_management.py
diff --git a/tests/test_litellm/responses/test_responses_router_cooldown.py b/tests/unit/responses/test_responses_router_cooldown.py
similarity index 100%
rename from tests/test_litellm/responses/test_responses_router_cooldown.py
rename to tests/unit/responses/test_responses_router_cooldown.py
diff --git a/tests/test_litellm/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py
similarity index 100%
rename from tests/test_litellm/responses/test_responses_streaming_iterator.py
rename to tests/unit/responses/test_responses_streaming_iterator.py
diff --git a/tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py b/tests/unit/responses/test_responses_supported_endpoints_passthrough.py
similarity index 100%
rename from tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py
rename to tests/unit/responses/test_responses_supported_endpoints_passthrough.py
diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py
similarity index 100%
rename from tests/test_litellm/responses/test_responses_utils.py
rename to tests/unit/responses/test_responses_utils.py
diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py
similarity index 97%
rename from tests/test_litellm/responses/test_responses_websocket_all_providers.py
rename to tests/unit/responses/test_responses_websocket_all_providers.py
index 3888a84fb5d..6f346a25d9c 100644
--- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py
+++ b/tests/unit/responses/test_responses_websocket_all_providers.py
@@ -2718,97 +2718,6 @@ class TestWebSocketChunkTypes:
assert "response.reasoning_content.done" in serialized
assert "Complete reasoning" in serialized
- def test_extract_output_messages_preserves_multiple_messages(self):
- """Test that multiple output messages are all preserved"""
- from litellm.responses.streaming_iterator import (
- ManagedResponsesWebSocketHandler,
- )
-
- completed_event = {
- "type": "response.completed",
- "response": {
- "id": "resp_123",
- "output": [
- {
- "type": "message",
- "role": "assistant",
- "content": [{"type": "output_text", "text": "First message"}],
- },
- {
- "type": "function_call",
- "id": "call_123",
- "name": "get_weather",
- "arguments": "{}",
- },
- {
- "type": "message",
- "role": "assistant",
- "content": [{"type": "output_text", "text": "Second message"}],
- },
- ],
- },
- }
-
- messages = ManagedResponsesWebSocketHandler._extract_output_messages(
- completed_event
- )
- assert len(messages) == 3
- assert messages[0]["content"][0]["text"] == "First message"
- assert messages[1]["type"] == "function_call"
- assert messages[2]["content"][0]["text"] == "Second message"
-
- def test_input_to_messages_with_mixed_content_types(self):
- """Test input conversion with mixed content types"""
- from litellm.responses.streaming_iterator import (
- ManagedResponsesWebSocketHandler,
- )
-
- input_list = [
- {
- "type": "message",
- "role": "user",
- "content": [
- {"type": "input_text", "text": "Question"},
- {"type": "input_image", "image_url": "https://example.com/img.png"},
- ],
- }
- ]
-
- messages = ManagedResponsesWebSocketHandler._input_to_messages(input_list)
- assert len(messages) == 1
- assert len(messages[0]["content"]) == 2
- assert messages[0]["content"][0]["type"] == "input_text"
- assert messages[0]["content"][1]["type"] == "input_image"
-
- def test_extract_output_messages_with_mixed_text_types(self):
- """Test that both 'output_text' and 'text' types are extracted"""
- from litellm.responses.streaming_iterator import (
- ManagedResponsesWebSocketHandler,
- )
-
- completed_event = {
- "type": "response.completed",
- "response": {
- "id": "resp_123",
- "output": [
- {
- "type": "message",
- "role": "assistant",
- "content": [
- {"type": "output_text", "text": "Part 1"},
- {"type": "text", "text": "Part 2"},
- ],
- }
- ],
- },
- }
-
- messages = ManagedResponsesWebSocketHandler._extract_output_messages(
- completed_event
- )
- assert len(messages) == 1
- assert messages[0]["content"][0]["text"] == "Part 1Part 2"
-
class TestNativeWebSocketUrlConstruction:
"""Test that native WebSocket URLs include the model query parameter.
diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/unit/responses/test_rust_bridge_websocket.py
similarity index 100%
rename from tests/test_litellm/responses/test_rust_bridge_websocket.py
rename to tests/unit/responses/test_rust_bridge_websocket.py
diff --git a/tests/test_litellm/responses/test_sse_output_recovery.py b/tests/unit/responses/test_sse_output_recovery.py
similarity index 100%
rename from tests/test_litellm/responses/test_sse_output_recovery.py
rename to tests/unit/responses/test_sse_output_recovery.py
diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py
similarity index 100%
rename from tests/test_litellm/responses/test_streaming_iterator.py
rename to tests/unit/responses/test_streaming_iterator.py
diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py
similarity index 100%
rename from tests/test_litellm/responses/test_streaming_iterator_error_events.py
rename to tests/unit/responses/test_streaming_iterator_error_events.py
diff --git a/tests/test_litellm/responses/test_text_format_conversion.py b/tests/unit/responses/test_text_format_conversion.py
similarity index 100%
rename from tests/test_litellm/responses/test_text_format_conversion.py
rename to tests/unit/responses/test_text_format_conversion.py
diff --git a/tests/test_litellm/rust_bridge/messages/__init__.py b/tests/unit/router_strategy/adaptive_router/__init__.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/messages/__init__.py
rename to tests/unit/router_strategy/adaptive_router/__init__.py
diff --git a/tests/test_litellm/rust_bridge/ocr/__init__.py b/tests/unit/router_strategy/adaptive_router/fixtures/__init__.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/ocr/__init__.py
rename to tests/unit/router_strategy/adaptive_router/fixtures/__init__.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_no_signals.json b/tests/unit/router_strategy/adaptive_router/fixtures/clean_no_signals.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_no_signals.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/clean_no_signals.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_satisfaction.json b/tests/unit/router_strategy/adaptive_router/fixtures/clean_satisfaction.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_satisfaction.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/clean_satisfaction.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/disengagement_giveup.json b/tests/unit/router_strategy/adaptive_router/fixtures/disengagement_giveup.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/disengagement_giveup.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/disengagement_giveup.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_429.json b/tests/unit/router_strategy/adaptive_router/fixtures/exhaustion_429.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_429.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/exhaustion_429.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json b/tests/unit/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/failure_tool_error.json b/tests/unit/router_strategy/adaptive_router/fixtures/failure_tool_error.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/failure_tool_error.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/failure_tool_error.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/loop_same_tool.json b/tests/unit/router_strategy/adaptive_router/fixtures/loop_same_tool.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/loop_same_tool.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/loop_same_tool.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json b/tests/unit/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json b/tests/unit/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/stagnation_repeat.json b/tests/unit/router_strategy/adaptive_router/fixtures/stagnation_repeat.json
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/fixtures/stagnation_repeat.json
rename to tests/unit/router_strategy/adaptive_router/fixtures/stagnation_repeat.json
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/unit/router_strategy/adaptive_router/test_adaptive_router.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py
rename to tests/unit/router_strategy/adaptive_router/test_adaptive_router.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_async_pre_routing.py b/tests/unit/router_strategy/adaptive_router/test_async_pre_routing.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_async_pre_routing.py
rename to tests/unit/router_strategy/adaptive_router/test_async_pre_routing.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py b/tests/unit/router_strategy/adaptive_router/test_bandit.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_bandit.py
rename to tests/unit/router_strategy/adaptive_router/test_bandit.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_classifier.py b/tests/unit/router_strategy/adaptive_router/test_classifier.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_classifier.py
rename to tests/unit/router_strategy/adaptive_router/test_classifier.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_config.py b/tests/unit/router_strategy/adaptive_router/test_config.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_config.py
rename to tests/unit/router_strategy/adaptive_router/test_config.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py b/tests/unit/router_strategy/adaptive_router/test_e2e_adaptive_router.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py
rename to tests/unit/router_strategy/adaptive_router/test_e2e_adaptive_router.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_hooks.py
rename to tests/unit/router_strategy/adaptive_router/test_hooks.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py b/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py
rename to tests/unit/router_strategy/adaptive_router/test_router_dispatch.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_signals.py b/tests/unit/router_strategy/adaptive_router/test_signals.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_signals.py
rename to tests/unit/router_strategy/adaptive_router/test_signals.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py b/tests/unit/router_strategy/adaptive_router/test_state_endpoint.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py
rename to tests/unit/router_strategy/adaptive_router/test_state_endpoint.py
diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py b/tests/unit/router_strategy/adaptive_router/test_update_queue.py
similarity index 100%
rename from tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py
rename to tests/unit/router_strategy/adaptive_router/test_update_queue.py
diff --git a/tests/test_litellm/router_strategy/complexity_router/test_context_compaction.py b/tests/unit/router_strategy/complexity_router/test_context_compaction.py
similarity index 100%
rename from tests/test_litellm/router_strategy/complexity_router/test_context_compaction.py
rename to tests/unit/router_strategy/complexity_router/test_context_compaction.py
diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/unit/router_strategy/test_auto_router.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_auto_router.py
rename to tests/unit/router_strategy/test_auto_router.py
diff --git a/tests/test_litellm/router_strategy/test_base_routing_strategy.py b/tests/unit/router_strategy/test_base_routing_strategy.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_base_routing_strategy.py
rename to tests/unit/router_strategy/test_base_routing_strategy.py
diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_budget_limiter.py
rename to tests/unit/router_strategy/test_budget_limiter.py
diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py
rename to tests/unit/router_strategy/test_budget_limiter_hotpath.py
diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_complexity_router.py
rename to tests/unit/router_strategy/test_complexity_router.py
diff --git a/tests/test_litellm/router_strategy/test_complexity_tier_predictor.py b/tests/unit/router_strategy/test_complexity_tier_predictor.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_complexity_tier_predictor.py
rename to tests/unit/router_strategy/test_complexity_tier_predictor.py
diff --git a/tests/test_litellm/router_strategy/test_fuse_presets.py b/tests/unit/router_strategy/test_fuse_presets.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_fuse_presets.py
rename to tests/unit/router_strategy/test_fuse_presets.py
diff --git a/tests/test_litellm/router_strategy/test_lar1_routing.py b/tests/unit/router_strategy/test_lar1_routing.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_lar1_routing.py
rename to tests/unit/router_strategy/test_lar1_routing.py
diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_least_busy.py
rename to tests/unit/router_strategy/test_least_busy.py
diff --git a/tests/test_litellm/router_strategy/test_litellm_encoder.py b/tests/unit/router_strategy/test_litellm_encoder.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_litellm_encoder.py
rename to tests/unit/router_strategy/test_litellm_encoder.py
diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_llm_v2.py
rename to tests/unit/router_strategy/test_llm_v2.py
diff --git a/tests/test_litellm/router_strategy/test_lowest_cost.py b/tests/unit/router_strategy/test_lowest_cost.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_lowest_cost.py
rename to tests/unit/router_strategy/test_lowest_cost.py
diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/unit/router_strategy/test_lowest_latency.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_lowest_latency.py
rename to tests/unit/router_strategy/test_lowest_latency.py
diff --git a/tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py
rename to tests/unit/router_strategy/test_lowest_tpm_rpm.py
diff --git a/tests/test_litellm/router_strategy/test_quality_router.py b/tests/unit/router_strategy/test_quality_router.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_quality_router.py
rename to tests/unit/router_strategy/test_quality_router.py
diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_router_routing_groups.py
rename to tests/unit/router_strategy/test_router_routing_groups.py
diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/unit/router_strategy/test_router_routing_plugins.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_router_routing_plugins.py
rename to tests/unit/router_strategy/test_router_routing_plugins.py
diff --git a/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py b/tests/unit/router_strategy/test_router_tag_regex_routing.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_router_tag_regex_routing.py
rename to tests/unit/router_strategy/test_router_tag_regex_routing.py
diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_router_tag_routing.py
rename to tests/unit/router_strategy/test_router_tag_routing.py
diff --git a/tests/test_litellm/router_strategy/test_savings_baseline.py b/tests/unit/router_strategy/test_savings_baseline.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_savings_baseline.py
rename to tests/unit/router_strategy/test_savings_baseline.py
diff --git a/tests/test_litellm/router_strategy/test_simple_shuffle.py b/tests/unit/router_strategy/test_simple_shuffle.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_simple_shuffle.py
rename to tests/unit/router_strategy/test_simple_shuffle.py
diff --git a/tests/test_litellm/router_strategy/test_stall_detector.py b/tests/unit/router_strategy/test_stall_detector.py
similarity index 100%
rename from tests/test_litellm/router_strategy/test_stall_detector.py
rename to tests/unit/router_strategy/test_stall_detector.py
diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
index a7006c62438..00462b65bc2 100644
--- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
+++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
@@ -565,7 +565,7 @@ async def test_wildcard_route_resolves_underlying_model_minimum(local_model_cost
@pytest.mark.asyncio
async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop():
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
@@ -589,7 +589,7 @@ async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop():
@pytest.mark.asyncio
async def test_async_log_success_event_counts_the_prompt_off_the_event_loop():
from tests.large_text import text
- from tests.test_litellm.litellm_core_utils.event_loop_lag import (
+ from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
diff --git a/tests/test_litellm/router_utils/test_access_windows.py b/tests/unit/router_utils/test_access_windows.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_access_windows.py
rename to tests/unit/router_utils/test_access_windows.py
diff --git a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py b/tests/unit/router_utils/test_add_retry_fallback_headers.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_add_retry_fallback_headers.py
rename to tests/unit/router_utils/test_add_retry_fallback_headers.py
diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/unit/router_utils/test_auto_router_model_naming.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_auto_router_model_naming.py
rename to tests/unit/router_utils/test_auto_router_model_naming.py
diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/unit/router_utils/test_auto_router_tuning_baseline.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py
rename to tests/unit/router_utils/test_auto_router_tuning_baseline.py
diff --git a/tests/test_litellm/router_utils/test_client_initalization_utils.py b/tests/unit/router_utils/test_client_initalization_utils.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_client_initalization_utils.py
rename to tests/unit/router_utils/test_client_initalization_utils.py
diff --git a/tests/test_litellm/router_utils/test_cooldown_cache.py b/tests/unit/router_utils/test_cooldown_cache.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_cooldown_cache.py
rename to tests/unit/router_utils/test_cooldown_cache.py
diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_cooldown_handlers.py
rename to tests/unit/router_utils/test_cooldown_handlers.py
diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_fallback_event_handlers.py
rename to tests/unit/router_utils/test_fallback_event_handlers.py
diff --git a/tests/test_litellm/router_utils/test_get_retry_from_policy.py b/tests/unit/router_utils/test_get_retry_from_policy.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_get_retry_from_policy.py
rename to tests/unit/router_utils/test_get_retry_from_policy.py
diff --git a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py
rename to tests/unit/router_utils/test_health_check_allowed_fails_integration.py
diff --git a/tests/test_litellm/router_utils/test_health_state_cache.py b/tests/unit/router_utils/test_health_state_cache.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_health_state_cache.py
rename to tests/unit/router_utils/test_health_state_cache.py
diff --git a/tests/test_litellm/router_utils/test_pattern_match_deployments.py b/tests/unit/router_utils/test_pattern_match_deployments.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_pattern_match_deployments.py
rename to tests/unit/router_utils/test_pattern_match_deployments.py
diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/unit/router_utils/test_reasoning_effort_capability.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_reasoning_effort_capability.py
rename to tests/unit/router_utils/test_reasoning_effort_capability.py
diff --git a/tests/test_litellm/router_utils/test_router_health_check_routing.py b/tests/unit/router_utils/test_router_health_check_routing.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_router_health_check_routing.py
rename to tests/unit/router_utils/test_router_health_check_routing.py
diff --git a/tests/test_litellm/router_utils/test_router_interactions_endpoints.py b/tests/unit/router_utils/test_router_interactions_endpoints.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_router_interactions_endpoints.py
rename to tests/unit/router_utils/test_router_interactions_endpoints.py
diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py
similarity index 100%
rename from tests/test_litellm/router_utils/test_router_utils_common_utils.py
rename to tests/unit/router_utils/test_router_utils_common_utils.py
diff --git a/tests/test_litellm/rust_bridge/AGENTS.md b/tests/unit/rust_bridge/AGENTS.md
similarity index 100%
rename from tests/test_litellm/rust_bridge/AGENTS.md
rename to tests/unit/rust_bridge/AGENTS.md
diff --git a/tests/unit/rust_bridge/messages/test_route_host.py b/tests/unit/rust_bridge/messages/test_route_host.py
index a880cfe3588..1be42e2249d 100644
--- a/tests/unit/rust_bridge/messages/test_route_host.py
+++ b/tests/unit/rust_bridge/messages/test_route_host.py
@@ -3,6 +3,10 @@ from typing import Final
from litellm.rust_bridge.messages.route_host import arguments, response
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
+from dataclasses import astuple
+import pytest
+import litellm
+from litellm.rust_bridge.messages import route_host
def test_response_is_a_detached_public_messages_dict() -> None:
@@ -40,3 +44,121 @@ def test_arguments_are_the_public_kwargs_view() -> None:
)
assert arguments(request) is kwargs
+
+
+pytestmark = pytest.mark.usefixtures("local_model_cost_map")
+
+
+def _flag_model(monkeypatch: pytest.MonkeyPatch, name: str, **flags: bool) -> None:
+ monkeypatch.setitem(
+ litellm.model_cost,
+ name,
+ {
+ "litellm_provider": "anthropic",
+ "mode": "chat",
+ "input_cost_per_token": 0,
+ "output_cost_per_token": 0,
+ **flags,
+ },
+ )
+
+
+def test_capabilities_come_from_the_model_map_under_the_callers_provider(monkeypatch: pytest.MonkeyPatch) -> None:
+ _flag_model(
+ monkeypatch,
+ "claude-test-adaptive",
+ supports_reasoning=True,
+ supports_adaptive_thinking=True,
+ supports_output_config=True,
+ supports_xhigh_reasoning_effort=True,
+ supports_sampling_params=False,
+ )
+
+ capabilities: Final = route_host.model_capabilities("anthropic/claude-test-adaptive", None)
+
+ assert capabilities.supports_adaptive_thinking
+ assert capabilities.supports_output_config
+ assert not capabilities.supports_legacy_thinking
+ assert not capabilities.supports_sampling_params
+ assert capabilities.effort_tiers.xhigh
+ assert not capabilities.effort_tiers.max
+
+
+def test_unmapped_model_keeps_sampling_params_and_no_reasoning_features() -> None:
+ capabilities: Final = route_host.model_capabilities("anthropic/not-a-real-model", None)
+
+ assert capabilities.supports_sampling_params
+ assert not capabilities.supports_reasoning
+ assert not capabilities.supports_adaptive_thinking
+ assert not any(astuple(capabilities.effort_tiers))
+
+
+@pytest.mark.parametrize(
+ ("global_flag", "kwargs", "expected"),
+ [
+ (False, {}, False),
+ (True, {}, True),
+ (False, {"drop_params": "true"}, True),
+ (False, {"drop_params": "nonsense"}, False),
+ (False, {"drop_params": False}, False),
+ ],
+)
+def test_drop_params_merges_the_global_flag_with_the_request(
+ monkeypatch: pytest.MonkeyPatch, global_flag: bool, kwargs: dict[str, object], expected: bool
+) -> None:
+ monkeypatch.setattr(litellm, "drop_params", global_flag)
+
+ assert route_host.shaping("anthropic/not-a-real-model", None, kwargs)["drop_params"] is expected
+
+
+@pytest.mark.parametrize(
+ ("configured", "expected"),
+ [
+ (["tools[*].input_examples", 3, "metadata.user_id"], ("tools[*].input_examples", "metadata.user_id")),
+ ("tools", ()),
+ (None, ()),
+ ],
+)
+def test_additional_drop_params_keep_only_string_paths(configured: object, expected: tuple[str, ...]) -> None:
+ shaping: Final = route_host.shaping("anthropic/not-a-real-model", None, {"additional_drop_params": configured})
+
+ assert shaping["additional_drop_params"] == expected
+
+
+def test_native_request_rejections_map_to_the_public_400() -> None:
+ from types import MappingProxyType
+
+ from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
+
+ request: Final = LiteLLMMessagesRequest(
+ model="anthropic/claude-sonnet-5",
+ messages=(),
+ max_tokens=8,
+ stream=None,
+ api_key=None,
+ api_base=None,
+ custom_llm_provider=None,
+ kwargs=MappingProxyType({}),
+ )
+ rejected: Final = ValueError("claude-sonnet-5 does not support top_k=5")
+ rejected.messages_request_error = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets
+
+ mapped: Final = route_host.map_failure(rejected, request, "anthropic")
+
+ assert isinstance(mapped, litellm.BadRequestError)
+ assert mapped.status_code == 400
+ assert "does not support top_k=5" in mapped.message
+ assert mapped.model == "claude-sonnet-5"
+ assert not isinstance(route_host.map_failure(ValueError("plain"), request, "anthropic"), litellm.BadRequestError)
+
+
+def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handler_does() -> None:
+ hidden: Final = route_host.stream_hidden_params(
+ (("request-id", "req_upstream_123"), ("x-ratelimit-remaining-requests", "41"))
+ )
+
+ additional: Final = hidden["additional_headers"]
+ assert isinstance(additional, dict)
+ assert additional["llm_provider-request-id"] == "req_upstream_123"
+ assert additional["x-ratelimit-remaining-requests"] == "41"
+ assert "request-id" not in additional
diff --git a/tests/test_litellm/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/messages/test_secrets.py
rename to tests/unit/rust_bridge/messages/test_secrets.py
diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/unit/rust_bridge/native_route_wheel_test.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/native_route_wheel_test.py
rename to tests/unit/rust_bridge/native_route_wheel_test.py
diff --git a/tests/test_litellm/rust_bridge/ocr/test_secrets.py b/tests/unit/rust_bridge/ocr/test_secrets.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/ocr/test_secrets.py
rename to tests/unit/rust_bridge/ocr/test_secrets.py
diff --git a/tests/test_litellm/rust_bridge/stubtest.ini b/tests/unit/rust_bridge/stubtest.ini
similarity index 100%
rename from tests/test_litellm/rust_bridge/stubtest.ini
rename to tests/unit/rust_bridge/stubtest.ini
diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/unit/rust_bridge/test_bindings.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_bindings.py
rename to tests/unit/rust_bridge/test_bindings.py
diff --git a/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/test_callbacks_legacy_python.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py
rename to tests/unit/rust_bridge/test_callbacks_legacy_python.py
diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_catalog.py
rename to tests/unit/rust_bridge/test_catalog.py
diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/unit/rust_bridge/test_configuration.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_configuration.py
rename to tests/unit/rust_bridge/test_configuration.py
diff --git a/tests/test_litellm/rust_bridge/test_dispatch.py b/tests/unit/rust_bridge/test_dispatch.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_dispatch.py
rename to tests/unit/rust_bridge/test_dispatch.py
diff --git a/tests/test_litellm/rust_bridge/test_failures.py b/tests/unit/rust_bridge/test_failures.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_failures.py
rename to tests/unit/rust_bridge/test_failures.py
diff --git a/tests/test_litellm/rust_bridge/test_fork_guard.py b/tests/unit/rust_bridge/test_fork_guard.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_fork_guard.py
rename to tests/unit/rust_bridge/test_fork_guard.py
diff --git a/tests/test_litellm/rust_bridge/test_lifecycle.py b/tests/unit/rust_bridge/test_lifecycle.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_lifecycle.py
rename to tests/unit/rust_bridge/test_lifecycle.py
diff --git a/tests/test_litellm/rust_bridge/test_logger.py b/tests/unit/rust_bridge/test_logger.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_logger.py
rename to tests/unit/rust_bridge/test_logger.py
diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/unit/rust_bridge/test_runtime.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_runtime.py
rename to tests/unit/rust_bridge/test_runtime.py
diff --git a/tests/test_litellm/rust_bridge/test_secret_manager.py b/tests/unit/rust_bridge/test_secret_manager.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_secret_manager.py
rename to tests/unit/rust_bridge/test_secret_manager.py
diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/unit/rust_bridge/test_settings.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_settings.py
rename to tests/unit/rust_bridge/test_settings.py
diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/unit/rust_bridge/test_token_counter.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_token_counter.py
rename to tests/unit/rust_bridge/test_token_counter.py
diff --git a/tests/test_litellm/rust_bridge/test_tokenizer.py b/tests/unit/rust_bridge/test_tokenizer.py
similarity index 95%
rename from tests/test_litellm/rust_bridge/test_tokenizer.py
rename to tests/unit/rust_bridge/test_tokenizer.py
index 188aa81093f..c5093cdb0ce 100644
--- a/tests/test_litellm/rust_bridge/test_tokenizer.py
+++ b/tests/unit/rust_bridge/test_tokenizer.py
@@ -7,7 +7,7 @@ from tokenizers import Tokenizer
from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding
from litellm.rust_bridge import tokenizer
from litellm.utils import claude_json_str
-from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
+from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
TEXTS: Final = ("hello <|endoftext|> world", "café æ¼¢å— ðŸ™‚", " def f():\n return 1\n", "hello again")
diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/unit/rust_bridge/test_verify_linux_native_wheel.py
similarity index 100%
rename from tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py
rename to tests/unit/rust_bridge/test_verify_linux_native_wheel.py
From a8fe84bb3e298549edaf6a7abd845d54356b703f Mon Sep 17 00:00:00 2001
From: "berriai-litellm-provider-info-sync[bot]"
<328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com>
Date: Fri, 25 Sep 2026 17:20:52 -0700
Subject: [PATCH 6/6] chore(cost-map): add fireworks us-only deepseek v4.1
flash priority prices (#43247)
Price-Sync: litellm-providers
Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com>
---
litellm/model_prices_and_context_window_backup.json | 6 ++++++
model_prices_and_context_window.json | 6 ++++++
2 files changed, 12 insertions(+)
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index a670200c132..792e7316c67 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -60264,13 +60264,16 @@
},
"fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": {
"cache_read_input_token_cost": 9e-09,
+ "cache_read_input_token_cost_priority": 1.125e-08,
"input_cost_per_token": 4.5e-07,
+ "input_cost_per_token_priority": 5.625e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"mode": "chat",
"output_cost_per_token": 1.8e-06,
+ "output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
@@ -60360,13 +60363,16 @@
},
"fireworks_ai/deepseek-v4p1-flash-us": {
"cache_read_input_token_cost": 9e-09,
+ "cache_read_input_token_cost_priority": 1.125e-08,
"input_cost_per_token": 4.5e-07,
+ "input_cost_per_token_priority": 5.625e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"mode": "chat",
"output_cost_per_token": 1.8e-06,
+ "output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index a670200c132..792e7316c67 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -60264,13 +60264,16 @@
},
"fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": {
"cache_read_input_token_cost": 9e-09,
+ "cache_read_input_token_cost_priority": 1.125e-08,
"input_cost_per_token": 4.5e-07,
+ "input_cost_per_token_priority": 5.625e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"mode": "chat",
"output_cost_per_token": 1.8e-06,
+ "output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
@@ -60360,13 +60363,16 @@
},
"fireworks_ai/deepseek-v4p1-flash-us": {
"cache_read_input_token_cost": 9e-09,
+ "cache_read_input_token_cost_priority": 1.125e-08,
"input_cost_per_token": 4.5e-07,
+ "input_cost_per_token_priority": 5.625e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"mode": "chat",
"output_cost_per_token": 1.8e-06,
+ "output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,