Merge branch 'main' into litellm_jwt_auto_register_map_existing_key

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))

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 = "<system-reminder>"
 _REMINDER_CLOSE: Final = "</system-reminder>"
 _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
             )

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

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.

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(