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(