mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(router): validate encrypted classifiers after deployment selection
This commit is contained in:
moe-berri
parent
488bf6f596
commit
208d554c00
6 changed files with 168 additions and 35 deletions
|
|
@ -62,6 +62,9 @@ class BaseResponsesAPIConfig(ABC):
|
|||
"""
|
||||
return False
|
||||
|
||||
def supports_encrypted_agent_messages(self) -> bool:
|
||||
return False
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
|
|
|
|||
|
|
@ -110,6 +110,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
def supports_native_file_search(self) -> bool:
|
||||
return True
|
||||
|
||||
def supports_encrypted_agent_messages(self) -> bool:
|
||||
return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE)
|
||||
|
||||
@staticmethod
|
||||
def _is_gpt_5_model(model: str) -> bool:
|
||||
"""Return True only for actual OpenAI GPT-5 models.
|
||||
|
|
|
|||
|
|
@ -1078,6 +1078,7 @@ def responses(
|
|||
litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None)
|
||||
_is_async: Final = kwargs.pop("aresponses", False) is True
|
||||
skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False)
|
||||
require_encrypted_task_support: Final = kwargs.pop("_require_encrypted_task_support", False) is True
|
||||
use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs)
|
||||
|
||||
client_headers: Final = kwargs.get("headers")
|
||||
|
|
@ -1186,6 +1187,13 @@ def responses(
|
|||
model, custom_llm_provider, deployment_model_info
|
||||
)
|
||||
|
||||
if require_encrypted_task_support and (
|
||||
_bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api)
|
||||
or responses_api_provider_config is None
|
||||
or not responses_api_provider_config.supports_encrypted_agent_messages()
|
||||
):
|
||||
raise ValueError("Encrypted task classification requires a compatible native Responses deployment")
|
||||
|
||||
local_vars.update(kwargs)
|
||||
# Map reasoning_effort (from litellm_params/proxy config) to reasoning when not set
|
||||
if reasoning is None and "reasoning_effort" in local_vars:
|
||||
|
|
|
|||
|
|
@ -367,6 +367,9 @@ local scoring shortcut in `heuristic_first` and `hybrid` modes. The configured c
|
|||
a native OpenAI or Azure OpenAI Responses deployment with access to the encrypted content. The
|
||||
provider handles the encrypted task, and the classifier still chooses the tier dynamically
|
||||
|
||||
Compatibility is checked after normal deployment selection. A paused incompatible member of the
|
||||
classifier group does not prevent an eligible compatible deployment from classifying the task
|
||||
|
||||
Unsupported classifier deployments and provider decryption errors use the existing
|
||||
`classifier_fallback` policy. No fixed tier is introduced for encrypted tasks. Plaintext asks and
|
||||
requests carrying only historical encrypted reasoning retain the existing classifier path
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ from threading import Lock
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter, create_model
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError, create_model
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.constants import (
|
||||
|
|
@ -504,7 +504,10 @@ def _encrypted_classifier_task(
|
|||
raw_input: Final = (request_kwargs or EMPTY_MAPPING).get("input")
|
||||
if not isinstance(raw_input, list) or (request_kwargs or EMPTY_MAPPING).get("messages"):
|
||||
return None
|
||||
items: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(raw_input)
|
||||
try:
|
||||
items: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(raw_input)
|
||||
except ValidationError:
|
||||
return None
|
||||
current: Final = next(
|
||||
(
|
||||
item
|
||||
|
|
@ -516,7 +519,10 @@ def _encrypted_classifier_task(
|
|||
)
|
||||
if current is None or current.get("type") != "agent_message" or not isinstance(current.get("content"), list):
|
||||
return None
|
||||
parts: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(current["content"])
|
||||
try:
|
||||
parts: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(current["content"])
|
||||
except ValidationError:
|
||||
return None
|
||||
if not any(part.get("type") == "encrypted_content" and part.get("encrypted_content") for part in parts):
|
||||
return None
|
||||
return {
|
||||
|
|
@ -2058,7 +2064,7 @@ class ComplexityRouter(CustomLogger):
|
|||
classifier_call_params = MappingProxyType({"reasoning_effort": llm_config.reasoning_effort})
|
||||
|
||||
payload: Final = (
|
||||
self._native_classifier_payload(llm_config.model, messages_for_call, response_format, encrypted_task)
|
||||
self._native_classifier_payload(messages_for_call, response_format, encrypted_task)
|
||||
if encrypted_task is not None
|
||||
else {"messages": messages_for_call, "response_format": response_format, **classifier_call_params}
|
||||
)
|
||||
|
|
@ -2098,7 +2104,6 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
def _native_classifier_payload(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[AllMessageValues], # mutable-ok: existing transformation accepts the SDK message list
|
||||
response_format: Mapping[str, object],
|
||||
encrypted_task: Mapping[str, object],
|
||||
|
|
@ -2106,26 +2111,7 @@ class ComplexityRouter(CustomLogger):
|
|||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesTransformationHandler,
|
||||
)
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
|
||||
deployments: Final = self._group_deployments(model)
|
||||
if not deployments:
|
||||
raise ValueError("Encrypted task classification requires a native OpenAI Responses classifier deployment")
|
||||
for params in (LiteLLM_Params.model_validate(deployment.get("litellm_params")) for deployment in deployments):
|
||||
if declared_authenticating_provider(params.model, params.custom_llm_provider):
|
||||
raise ValueError(
|
||||
"Encrypted task classification requires a native OpenAI Responses classifier deployment"
|
||||
)
|
||||
_, provider, _, _ = get_llm_provider(model=params.model, litellm_params=params)
|
||||
if (
|
||||
provider not in ("openai", "azure")
|
||||
or params.use_chat_completions_api
|
||||
or params.model.startswith("openai/chat_completions/")
|
||||
):
|
||||
raise ValueError(
|
||||
"Encrypted task classification requires a native OpenAI Responses classifier deployment"
|
||||
)
|
||||
transformation: Final = LiteLLMResponsesTransformationHandler()
|
||||
input_items, instructions = transformation.convert_chat_completion_messages_to_responses_api(messages)
|
||||
llm_config: Final = self.config.classifier_llm_config
|
||||
|
|
@ -2139,6 +2125,7 @@ class ComplexityRouter(CustomLogger):
|
|||
"instructions": instructions,
|
||||
"text": transformation.transform_response_format_to_text_format(dict(response_format)),
|
||||
"store": False,
|
||||
"_require_encrypted_task_support": True,
|
||||
**reasoning,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -5,8 +5,10 @@ Tests the rule-based complexity scoring and tier assignment logic.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
import json
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
|
|
@ -14,6 +16,7 @@ from typing import Dict, Final, List
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import httpx
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
|
|
@ -69,6 +72,7 @@ from litellm.types.router import (
|
|||
TaggedPreRoutingStrategy,
|
||||
)
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
|
||||
requires_semantic_router = pytest.mark.skipif(
|
||||
|
|
@ -2520,11 +2524,21 @@ def _native_classifier_router(
|
|||
classifier_type: str = "llm",
|
||||
deployment_model: str = "openai/gpt-6-astra",
|
||||
failure: Exception | None = None,
|
||||
native_router: Router | None = None,
|
||||
http_handler: AsyncHTTPHandler | None = None,
|
||||
) -> tuple[ComplexityRouter, MagicMock]:
|
||||
dependency: Final = MagicMock(
|
||||
aresponses=AsyncMock(return_value=_native_classifier_response(output), side_effect=failure),
|
||||
aresponses=(
|
||||
native_router.factory_function(partial(litellm.aresponses, client=http_handler), call_type="aresponses")
|
||||
if native_router is not None
|
||||
else AsyncMock(return_value=_native_classifier_response(output), side_effect=failure)
|
||||
),
|
||||
acompletion=AsyncMock(return_value=_llm_response('{"tier":"SIMPLE"}')),
|
||||
get_model_list=MagicMock(return_value=[{"litellm_params": {"model": deployment_model}}]),
|
||||
get_model_list=(
|
||||
native_router.get_model_list
|
||||
if native_router is not None
|
||||
else MagicMock(return_value=[{"litellm_params": {"model": deployment_model}}])
|
||||
),
|
||||
)
|
||||
return (
|
||||
ComplexityRouter(
|
||||
|
|
@ -2533,7 +2547,11 @@ def _native_classifier_router(
|
|||
complexity_router_config={
|
||||
"tiers": {"SIMPLE": "cheap-model", "REASONING": "deep-model"},
|
||||
"classifier_type": classifier_type,
|
||||
"classifier_llm_config": {"model": "classifier", "timeout_ms": 100, "reasoning_effort": "low"},
|
||||
"classifier_llm_config": {
|
||||
"model": "classifier",
|
||||
"timeout_ms": 5000 if native_router is not None else 100,
|
||||
"reasoning_effort": "low",
|
||||
},
|
||||
"heuristic_first_max_tier": "SIMPLE" if classifier_type == "heuristic_first" else None,
|
||||
"hybrid_boundary_margin": 0.01 if classifier_type == "hybrid" else None,
|
||||
"classifier_fallback": "default_model",
|
||||
|
|
@ -2546,6 +2564,18 @@ def _native_classifier_router(
|
|||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def native_classifier_http() -> AsyncIterator[tuple[AsyncHTTPHandler, MagicMock]]:
|
||||
respond: Final = MagicMock(
|
||||
return_value=httpx.Response(200, json=_native_classifier_response('{"tier":"REASONING"}').model_dump())
|
||||
)
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
handler: Final = AsyncHTTPHandler()
|
||||
await handler.client.aclose()
|
||||
handler.client = client
|
||||
yield handler, respond
|
||||
|
||||
|
||||
class TestEncryptedTaskClassifier:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"])
|
||||
|
|
@ -2587,11 +2617,15 @@ class TestEncryptedTaskClassifier:
|
|||
assert "Caller constraints" not in call["instructions"]
|
||||
assert "SIMPLE" in call["instructions"] and "REASONING" in call["instructions"]
|
||||
assert call["text"]["format"]["schema"]["properties"]["tier"]["enum"] == [
|
||||
"SIMPLE", "MEDIUM", "COMPLEX", "REASONING"
|
||||
"SIMPLE",
|
||||
"MEDIUM",
|
||||
"COMPLEX",
|
||||
"REASONING",
|
||||
]
|
||||
assert call["text"]["format"]["strict"] is True
|
||||
assert call["reasoning"] == {"effort": "low"}
|
||||
assert call["store"] is False
|
||||
assert call["_require_encrypted_task_support"] is True
|
||||
assert call["stream"] is False
|
||||
assert "tools" not in call and "previous_response_id" not in call
|
||||
assert "messages" not in call and "response_format" not in call
|
||||
|
|
@ -2606,13 +2640,25 @@ class TestEncryptedTaskClassifier:
|
|||
@pytest.mark.parametrize(
|
||||
"items",
|
||||
[
|
||||
[{"type": "reasoning", "encrypted_content": "opaque-history", "summary": []}, {"role": "user", "content": "hi"}],
|
||||
[
|
||||
{"type": "reasoning", "encrypted_content": "opaque-history", "summary": []},
|
||||
{"role": "user", "content": "hi"},
|
||||
],
|
||||
[_encrypted_agent_task(), {"role": "user", "content": "hi"}],
|
||||
[{**_encrypted_agent_task(), "content": [{"type": "input_text", "text": "hi"}]}],
|
||||
[{"role": "user", "content": "gAAAA is plain text"}],
|
||||
[{"role": "user", "content": "hi"}, {"type": "function_call_output", "call_id": "call_1", "output": "opaque-provider-task"}],
|
||||
[
|
||||
{"role": "user", "content": "hi"},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "opaque-provider-task"},
|
||||
],
|
||||
],
|
||||
ids=[
|
||||
"historical-reasoning",
|
||||
"older-encrypted-task",
|
||||
"plaintext-agent",
|
||||
"ciphertext-looking-text",
|
||||
"tool-output",
|
||||
],
|
||||
ids=["historical-reasoning", "older-encrypted-task", "plaintext-agent", "ciphertext-looking-text", "tool-output"],
|
||||
)
|
||||
async def test_other_asks_keep_chat_classifier(self, items: list[dict[str, object]]):
|
||||
router, dependency = _native_classifier_router()
|
||||
|
|
@ -2639,9 +2685,28 @@ class TestEncryptedTaskClassifier:
|
|||
dependency.acompletion.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("deployment_model", ["anthropic/test-classifier", "openai/chat_completions/gpt-6-astra"])
|
||||
async def test_incompatible_classifier_does_not_flatten_encryption(self, deployment_model: str):
|
||||
router, dependency = _native_classifier_router(deployment_model=deployment_model)
|
||||
@pytest.mark.parametrize(
|
||||
"deployment_model",
|
||||
["anthropic/test-classifier", "openai/chat_completions/gpt-6-astra", "xai/test-classifier"],
|
||||
)
|
||||
async def test_incompatible_classifier_does_not_flatten_encryption(
|
||||
self, deployment_model: str, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock]
|
||||
):
|
||||
handler, respond = native_classifier_http
|
||||
native: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "classifier",
|
||||
"litellm_params": {
|
||||
"model": deployment_model,
|
||||
"api_key": "test-key",
|
||||
"api_base": "https://classifier.test/v1",
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
router, _ = _native_classifier_router(native_router=native, http_handler=handler)
|
||||
|
||||
result: Final = await router.async_pre_routing_hook(
|
||||
model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]}
|
||||
|
|
@ -2649,8 +2714,72 @@ class TestEncryptedTaskClassifier:
|
|||
|
||||
assert result.model == "deep-model"
|
||||
assert result.routing_decision["cause"] == "default_model_fallback"
|
||||
respond.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blocked", [True, False])
|
||||
async def test_native_classifier_validates_selected_deployment(
|
||||
self, blocked: bool, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock]
|
||||
):
|
||||
handler, respond = native_classifier_http
|
||||
native: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "classifier",
|
||||
"litellm_params": {"model": "anthropic/test-classifier", "api_key": "test-key", "order": 0},
|
||||
"model_info": {"id": "incompatible", "blocked": blocked},
|
||||
},
|
||||
{
|
||||
"model_name": "classifier",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-6-astra",
|
||||
"api_key": "test-key",
|
||||
"order": 1,
|
||||
"api_base": "https://classifier.test/v1",
|
||||
},
|
||||
"model_info": {"id": "compatible"},
|
||||
},
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
router, _ = _native_classifier_router(native_router=native, http_handler=handler)
|
||||
task: Final = _encrypted_agent_task()
|
||||
|
||||
result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs={"input": [task]})
|
||||
|
||||
assert result.model == "deep-model"
|
||||
assert result.routing_decision["cause"] == ("llm_classifier" if blocked else "default_model_fallback")
|
||||
if blocked:
|
||||
respond.assert_called_once()
|
||||
request: Final = respond.call_args.args[0]
|
||||
assert request.url.path == "/v1/responses"
|
||||
body: Final = json.loads(request.content)
|
||||
assert body["input"][-1] == task
|
||||
assert "_require_encrypted_task_support" not in body
|
||||
else:
|
||||
respond.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"])
|
||||
@pytest.mark.parametrize(
|
||||
"input_items",
|
||||
[
|
||||
["unsupported-input-item"],
|
||||
[{**_encrypted_agent_task(), "content": [{"type": "input_text", "text": "hi"}, None]}],
|
||||
],
|
||||
)
|
||||
async def test_encrypted_detection_does_not_reject_other_input_shapes(
|
||||
self, classifier_type: str, input_items: list[object]
|
||||
):
|
||||
router, dependency = _native_classifier_router(classifier_type=classifier_type)
|
||||
|
||||
result: Final = await router.aclassify("hi", request_kwargs={"input": input_items})
|
||||
|
||||
assert result.cause != "default_model_fallback"
|
||||
assert result.tier == ComplexityTier.SIMPLE
|
||||
dependency.aresponses.assert_not_called()
|
||||
dependency.acompletion.assert_not_called()
|
||||
if classifier_type == "llm":
|
||||
dependency.acompletion.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", [ValueError("invalid_encrypted_content"), TimeoutError("classifier timed out")])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue