From 6c23c27bca8a6b217b0aca79d7e3ce4b9584f21e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:29:01 -0700 Subject: [PATCH 01/18] feat(decisions): add hosted_vllm provider (#45501) Backport of #45501 to rc/1.106.0. Cherry-picked from 61e5f2dd3c1e0671c07f25fa44f10097d161740e (main). --- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 5 + litellm/constants.py | 1 + .../health_check_helpers.py | 21 ++- .../llms/base_llm/decisions/transformation.py | 3 + .../hosted_vllm/decisions/transformation.py | 26 ++++ litellm/utils.py | 2 + provider_endpoints_support.json | 3 +- .../management/test_model_health_check.py | 39 +++++ .../providers/test_decisions_wire.py | 3 + tests/integration/proxy_config.yaml | 4 + .../decisions/bases/hosted_vllm.py | 67 +++++++++ .../basic/test_decisions_basic_hosted_vllm.py | 11 ++ tests/unit/decisions/test_main.py | 133 ++++++++++++++++++ .../test_health_check_helpers.py | 75 ++++++++++ 15 files changed, 393 insertions(+), 3 deletions(-) create mode 100644 litellm/llms/hosted_vllm/decisions/transformation.py create mode 100644 tests/integration/translation/decisions/bases/hosted_vllm.py create mode 100644 tests/integration/translation/decisions/basic/test_decisions_basic_hosted_vllm.py diff --git a/litellm/__init__.py b/litellm/__init__.py index d6133942542f..ebaef292e2ba 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1681,6 +1681,9 @@ def set_global_gitlab_config(config: Mapping[str, object]) -> None: from .llms.strands_decider.decisions.transformation import ( StrandsDeciderDecisionsConfig as StrandsDeciderDecisionsConfig, ) + from .llms.hosted_vllm.decisions.transformation import ( + HostedVLLMDecisionsConfig as HostedVLLMDecisionsConfig, + ) from .llms.openai.decisions.transformation import ( OpenAIDecisionsConfig as OpenAIDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index ca8bab010f3d..9545ed16851c 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -162,6 +162,7 @@ "OpenRouterDecisionsConfig", "CloudflareDecisionsConfig", "StrandsDeciderDecisionsConfig", + "HostedVLLMDecisionsConfig", "OpenAIDecisionsConfig", "AzureAIDecisionsConfig", "NvidiaNimRerankConfig", @@ -724,6 +725,10 @@ ".llms.strands_decider.decisions.transformation", "StrandsDeciderDecisionsConfig", ), + "HostedVLLMDecisionsConfig": ( + ".llms.hosted_vllm.decisions.transformation", + "HostedVLLMDecisionsConfig", + ), "OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"), "AzureAIDecisionsConfig": (".llms.azure_ai.decisions.transformation", "AzureAIDecisionsConfig"), "NvidiaNimRerankConfig": ( diff --git a/litellm/constants.py b/litellm/constants.py index 699872f32cc1..290749e562fc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -285,6 +285,7 @@ class ServerStreamingClassification(str, Enum): # Provider-specific API base URLs XAI_API_BASE: Final = "https://api.x.ai/v1" +PLACEHOLDER_API_KEY: Final = "fake-api-key" OPEN_SANDBOX_API_BASE_ENV_VAR: Final = "OPEN_SANDBOX_API_BASE" OPEN_SANDBOX_API_KEY_ENV_VAR: Final = "OPEN_SANDBOX_API_KEY" OPEN_SANDBOX_DEFAULT_TEMPLATE: Final = "opensandbox/code-interpreter:v1.1.0" diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 60e4b7495ede..5d14a9761d50 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -3,7 +3,7 @@ """ import base64 -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from typing import TYPE_CHECKING, Final, Literal from litellm.llms.base_llm.ocr.transformation import DocumentType @@ -56,6 +56,21 @@ def get_image_file_for_health_check() -> bytes: return base64.b64decode(TEST_IMAGE_BASE64) +def _decisions_health_check_questions(model: str, custom_llm_provider: str) -> Mapping[str, Mapping[str, object]]: + import litellm + from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + from litellm.utils import ProviderConfigManager + + provider: Final = next((member for member in litellm.LlmProviders if member.value == custom_llm_provider), None) + config: Final = ( + None + if provider is None + else ProviderConfigManager.get_provider_decisions_config(model=model, provider=provider) + ) + questions: Final = BaseDecisionsConfig.health_check_questions if config is None else config.health_check_questions + return {name: dict(question) for name, question in questions.items()} + + def _ocr_health_check_document(model: str, custom_llm_provider: str) -> DocumentType: native: Final = NATIVE_OCR_HEALTH_CHECK_DOCUMENT.load() if native is None: @@ -300,7 +315,9 @@ def get_mode_handlers( **DECISIONS_CALL_PARAMS.validate_python( { "state": prompt or "health check", - "questions": {"reachable": {"type": "noul", "instructions": "Is the service reachable?"}}, + "questions": _decisions_health_check_questions( + model=model, custom_llm_provider=custom_llm_provider + ), **_filter_model_params(model_params=model_params), } ) diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index a8d14ca59a3d..9d1ed188d7b1 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -301,6 +301,9 @@ class BaseDecisionsConfig(ABC): api_key_env: tuple[str, ...] = () api_base_env: tuple[str, ...] = () api_key_required: bool = True + health_check_questions: Mapping[str, Mapping[str, object]] = MappingProxyType( + {"reachable": MappingProxyType({"type": "noul", "instructions": "Is the service reachable?"})} + ) def get_default_api_base(self) -> str | None: return None diff --git a/litellm/llms/hosted_vllm/decisions/transformation.py b/litellm/llms/hosted_vllm/decisions/transformation.py new file mode 100644 index 000000000000..ef99f2de1737 --- /dev/null +++ b/litellm/llms/hosted_vllm/decisions/transformation.py @@ -0,0 +1,26 @@ +from types import MappingProxyType +from typing import Final + +from litellm.constants import PLACEHOLDER_API_KEY +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + + +class HostedVLLMDecisionsConfig(BaseDecisionsConfig): + api_key_env = ("HOSTED_VLLM_API_KEY",) + api_base_env = ("HOSTED_VLLM_API_BASE",) + api_key_required = False + health_check_questions = MappingProxyType( + { + "reachable": MappingProxyType( + { + "type": "choice", + "instructions": "Is the service reachable?", + "criteria": {"yes": None, "no": None}, + } + ) + } + ) + + def resolve_api_key(self, api_key: str | None) -> str | None: + resolved: Final = super().resolve_api_key(api_key) + return None if resolved == PLACEHOLDER_API_KEY else resolved diff --git a/litellm/utils.py b/litellm/utils.py index 5215e079b990..35bb8df37d5f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8989,6 +8989,8 @@ def get_provider_decisions_config(model: str, provider: LlmProviders) -> BaseDec return litellm.CloudflareDecisionsConfig() if provider == LlmProviders.STRANDS_DECIDER: return litellm.StrandsDeciderDecisionsConfig() + if provider == LlmProviders.HOSTED_VLLM: + return litellm.HostedVLLMDecisionsConfig() if provider == LlmProviders.OPENAI: return litellm.OpenAIDecisionsConfig() if provider == LlmProviders.AZURE_AI: diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 62e233885817..ef1fa270b065 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1370,7 +1370,8 @@ "rerank": true, "a2a": true, "interactions": true, - "video_generations": true + "video_generations": true, + "systemone": true } }, "huggingface": { diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index 17bd1c2ebfe3..11b902b540a8 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -24,6 +24,16 @@ "usage": {"input_tokens": 10, "output_tokens": 1}, }, ) +_VLLM_PROBE_REPLY: Final = JsonResponse( + content_type="application/json", + body={ + "model": "Qwen/Qwen3-0.6B", + "answers": { + "reachable": {"type": "choice", "choice": "yes", "confidence": 1.0, "probabilities": {"yes": 1.0, "no": 0.0}} + }, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, +) _STRANDS_PROBE_REPLY: Final = JsonResponse( content_type="application/json", body={ @@ -147,3 +157,32 @@ def test_evaluation_mode_health_check_of_the_self_hosted_strands_model_resolves_ }, ) ] + + +def test_evaluation_mode_health_check_of_hosted_vllm_sends_a_choice_probe(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = register_scenario(f"health-decisions-{uuid.uuid4().hex[:12]}", _VLLM_PROBE_REPLY) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + model="hosted_vllm/Qwen/Qwen3-0.6B", + api_base=handle.api_base(), + model_info={"mode": "evaluation"}, + ) + report: Final = _health_report(gateway, model) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert _probes_sent_to(gateway, handle) == [ + ( + f"/{handle.scenario_id}/v1/systemone", + { + "model": "Qwen/Qwen3-0.6B", + "state": os.environ.get("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"), + "questions": { + "reachable": { + "type": "choice", + "instructions": "Is the service reachable?", + "criteria": {"yes": None, "no": None}, + } + }, + }, + ) + ] diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 2a2babc030dd..e61739262697 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -94,6 +94,9 @@ class _Provider: _Provider( "strands_decider", "strands_decider/systemone-decider", "/v1/systemone", "systemone-decider", None, False, None ), + _Provider( + "hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None + ), _Provider( "cloudflare", "cloudflare/clef", diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index 4c03d82ac99c..f958984bae56 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -346,6 +346,10 @@ model_list: litellm_params: model: strands_decider/strands-decider-2B-hobson-v19 api_base: http://127.0.0.1:8191 + - model_name: hosted_vllm/Qwen/Qwen3-0.6B + litellm_params: + model: hosted_vllm/Qwen/Qwen3-0.6B + api_base: http://127.0.0.1:8191 - model_name: azure_ai/decision-1 litellm_params: model: azure_ai/decision-1 diff --git a/tests/integration/translation/decisions/bases/hosted_vllm.py b/tests/integration/translation/decisions/bases/hosted_vllm.py new file mode 100644 index 000000000000..2da37f83a861 --- /dev/null +++ b/tests/integration/translation/decisions/bases/hosted_vllm.py @@ -0,0 +1,67 @@ +from typing import Final + +from integration.translation.case import TranslationTestCase + +"""Provider request and reply shape from a real vLLM serve (POST /v1/systemone, no auth). Mock reply captured live on 2026-10-09 from a vLLM CPU build of vllm main 702ce313f serving Qwen/Qwen3-0.6B. +""" +QWEN3_0_6B_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/systemone", + litellm_request={ + "model": "hosted_vllm/Qwen/Qwen3-0.6B", + "state": "Help! My payouts have been failing for 3 days!", + "questions": { + "is_urgent": { + "type": "choice", + "instructions": "Does this convey urgency?", + "criteria": {"yes": "the message conveys urgency", "no": "the message does not convey urgency"}, + } + }, + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/systemone", + expected_provider_headers={"content-type": "application/json"}, + expected_provider_request={ + "model": "Qwen/Qwen3-0.6B", + "state": "Help! My payouts have been failing for 3 days!", + "questions": { + "is_urgent": { + "type": "choice", + "instructions": "Does this convey urgency?", + "criteria": {"yes": "the message conveys urgency", "no": "the message does not convey urgency"}, + } + }, + }, + mock_provider_response={ + "id": "decision-9c26c028d7f4457e", + "object": "structured_decision", + "created": 1791517765, + "model": "Qwen/Qwen3-0.6B", + "answers": { + "is_urgent": { + "type": "choice", + "choice": "yes", + "probabilities": {"yes": 0.9669140238756562, "no": 0.03308597612434374}, + "confidence": 0.9649145337527952, + } + }, + "usage": {"input_tokens": 61, "output_tokens": 1}, + "diagnostics": {"is_urgent": {"label_mass": 0.9979320910923947, "argmax_is_label": True}}, + }, + expected_litellm_response={ + "id": "decision-9c26c028d7f4457e", + "object": "structured_decision", + "created": 1791517765, + "model": "Qwen/Qwen3-0.6B", + "answers": { + "is_urgent": { + "type": "choice", + "choice": "yes", + "probabilities": {"yes": 0.9669140238756562, "no": 0.03308597612434374}, + "confidence": 0.9649145337527952, + } + }, + "usage": {"input_tokens": 61, "output_tokens": 1}, + "diagnostics": {"is_urgent": {"label_mass": 0.9979320910923947, "argmax_is_label": True}}, + }, +) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_hosted_vllm.py b/tests/integration/translation/decisions/basic/test_decisions_basic_hosted_vllm.py new file mode 100644 index 000000000000..f52d9147273d --- /dev/null +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_hosted_vllm.py @@ -0,0 +1,11 @@ +import pytest +from integration._support.client import Gateway +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.decisions.bases.hosted_vllm import QWEN3_0_6B_TEST_CASE +from integration.translation.runner import assert_translation + + +@pytest.mark.parametrize("case", [QWEN3_0_6B_TEST_CASE], ids=lambda case: case.id) +def test_decisions_basic_hosted_vllm(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 41c01ccf15e2..bf90364fbbb8 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -733,6 +733,139 @@ async def test_strands_decider_provider_resolution_and_router_dispatch( assert response.model == _STRANDS_RESPONSE["model"] +_VLLM_RESPONSE: Final[Mapping[str, object]] = { + "id": "systemone-4af3d2c1", + "object": "structured_decision", + "created": 1760012345, + "model": "Qwen/Qwen3-0.6B", + "answers": { + "is_defect": { + "type": "choice", + "choice": "yes", + "confidence": 0.91, + "probabilities": {"yes": 0.91, "no": 0.09}, + } + }, + "usage": {"input_tokens": 154, "output_tokens": 2}, + "diagnostics": {"engine": "vllm", "logprobs_mode": "raw_logprobs"}, +} + + +@pytest.mark.asyncio +async def test_hosted_vllm_requires_api_base_before_http( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + + with pytest.raises(litellm.BadRequestError, match="api_base is required"): + await litellm.adecisions( + model="hosted_vllm/Qwen/Qwen3-0.6B", + state="review", + questions={"is_defect": {"type": "choice", "criteria": {"yes": None, "no": None}}}, + ) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +async def test_hosted_vllm_without_key_sends_no_authorization_and_preserves_response_extras( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond(json=_VLLM_RESPONSE) + + response: Final = await litellm.adecisions( + model="hosted_vllm/Qwen/Qwen3-0.6B", + state="The package arrived broken.", + questions={"is_defect": {"type": "choice", "criteria": {"yes": None, "no": None}}}, + api_base="http://vllm.local:8000", + ) + + assert route.called + assert "authorization" not in respx_mock.calls[0].request.headers + assert response.model_extra["id"] == _VLLM_RESPONSE["id"] + assert response.model_extra["object"] == _VLLM_RESPONSE["object"] + assert response.model_extra["diagnostics"] == _VLLM_RESPONSE["diagnostics"] + + +@pytest.mark.asyncio +async def test_hosted_vllm_uses_key_and_base_from_environment( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://vllm.local:8000/v1") + monkeypatch.setenv("HOSTED_VLLM_API_KEY", "vllm-key") + route: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond(json=_VLLM_RESPONSE) + + await litellm.adecisions( + model="hosted_vllm/Qwen/Qwen3-0.6B", + state="The package arrived broken.", + questions={"is_defect": {"type": "choice", "criteria": {"yes": None, "no": None}}}, + ) + + assert route.called + assert respx_mock.calls[0].request.headers["authorization"] == "Bearer vllm-key" + + +@pytest.mark.asyncio +async def test_hosted_vllm_api_base_with_v1_suffix_posts_to_v1_systemone( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond(json=_VLLM_RESPONSE) + + await litellm.adecisions( + model="hosted_vllm/Qwen/Qwen3-0.6B", + state="The package arrived broken.", + questions={"is_defect": {"type": "choice", "criteria": {"yes": None, "no": None}}}, + api_base="http://vllm.local:8000/v1", + ) + + assert route.called + + +@pytest.mark.asyncio +async def test_hosted_vllm_choice_question_goes_out_as_a_jev_choice_with_the_served_model_name( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond(json=_VLLM_RESPONSE) + + await litellm.adecisions( + model="hosted_vllm/Qwen/Qwen3-0.6B", + state="The package arrived broken.", + questions={ + "is_defect": { + "type": "choice", + "instructions": "Is this a defect?", + "criteria": {"yes": "defective", "no": "intact"}, + } + }, + api_base="http://vllm.local:8000", + ) + + assert route.called + assert json.loads(respx_mock.calls[0].request.content) == { + "model": "Qwen/Qwen3-0.6B", + "state": "The package arrived broken.", + "questions": { + "is_defect": { + "type": "choice", + "instructions": "Is this a defect?", + "criteria": {"yes": "defective", "no": "intact"}, + } + }, + } + + @pytest.mark.asyncio @pytest.mark.parametrize( ("api_base", "url"), diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index c0298bc1c809..9903ab71876f 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -679,6 +679,81 @@ async def test_ahealth_check_probes_strands_through_decisions_without_mode( assert "authorization" not in upstream.calls[0].request.headers +@pytest.mark.asyncio +async def test_ahealth_check_probes_hosted_vllm_with_a_choice_question( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond( + json={ + "model": "Qwen/Qwen3-0.6B", + "answers": { + "reachable": { + "type": "choice", + "choice": "yes", + "confidence": 1.0, + "probabilities": {"yes": 1.0, "no": 0.0}, + } + }, + "usage": {"input_tokens": 12, "output_tokens": 1}, + } + ) + + result: Final = await ahealth_check( + { + "model": "hosted_vllm/Qwen/Qwen3-0.6B", + "api_base": "http://vllm.local:8000", + }, + mode="evaluation", + ) + + assert "error" not in result, result + assert upstream.called + sent: Final = json.loads(upstream.calls[0].request.content) + assert sent["questions"] == { + "reachable": { + "type": "choice", + "instructions": "Is the service reachable?", + "criteria": {"yes": None, "no": None}, + } + } + + +@pytest.mark.asyncio +async def test_ahealth_check_keeps_the_noul_probe_for_other_decisions_providers( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.typesafe.ai/v1/systemone").respond( + json={ + "model": "jev-1.13", + "answers": {"reachable": {"type": "noul", "noul": 1.0}}, + "usage": {"input_tokens": 12, "output_tokens": 1}, + } + ) + + result: Final = await ahealth_check( + { + "model": "typesafe/jev-1.13", + "api_key": "sk-test", + }, + mode="evaluation", + ) + + assert "error" not in result, result + assert upstream.called + sent: Final = json.loads(upstream.calls[0].request.content) + assert sent["questions"] == { + "reachable": {"type": "noul", "instructions": "Is the service reachable?"} + } + + @pytest.mark.parametrize( ("model", "custom_llm_provider", "expected"), ( From 30ed6e1dd6b854334966e7e93f7b5ca0de6b5346 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:29:18 -0700 Subject: [PATCH 02/18] fix(openrouter): bill responses, decisions and pass-through from usage.cost and always allow reasoning params (#45640) Backport of #45640 to rc/1.106.0. Cherry-picked from 465f20f2747c2b6f0e1965b7d45b4813d3e1e815 (main). --- litellm/decisions/main.py | 3 + .../llms/base_llm/decisions/transformation.py | 3 + .../llms/openrouter/chat/transformation.py | 13 +--- .../openrouter/decisions/transformation.py | 9 +++ .../openrouter/responses/transformation.py | 70 ++++++++++++++++++ .../typesafe_passthrough_logging_handler.py | 5 +- .../providers/test_decisions_wire.py | 35 ++++++--- tests/unit/decisions/test_main.py | 31 ++++++++ .../test_openrouter_chat_transformation.py | 17 +++-- ...est_openrouter_responses_transformation.py | 73 +++++++++++++++++++ ...st_typesafe_passthrough_logging_handler.py | 14 ++-- 11 files changed, 240 insertions(+), 33 deletions(-) diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index 2874327caccf..f1b83bb929a6 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -7,6 +7,7 @@ from typing_extensions import assert_never import litellm +from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( @@ -205,11 +206,13 @@ def _prepare_call( def _format_response(response: DecisionsIRResponse, call: _DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse: formatted: Final = _formatted_response(response, call) + provider_cost: Final = call.provider_config.provider_reported_cost(response) formatted.set_hidden_params( { "model": f"{call.custom_llm_provider}/{call.model}", "custom_llm_provider": call.custom_llm_provider, "provider_response_model": f"{call.custom_llm_provider}/{call.model}", + **({"additional_headers": {RESPONSE_COST_HEADER: provider_cost}} if provider_cost is not None else {}), } ) return formatted diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index 9d1ed188d7b1..c57c643983a5 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -308,6 +308,9 @@ class BaseDecisionsConfig(ABC): def get_default_api_base(self) -> str | None: return None + def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None: + return None + def missing_api_base_message(self, custom_llm_provider: str) -> str: return f"api_base is required for Decisions provider '{custom_llm_provider}'" diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index e97962a1c48f..c63863bbca48 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -12,7 +12,6 @@ import httpx -import litellm from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -41,18 +40,10 @@ class CacheControlSupportedModels(str, Enum): class OpenrouterConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> list: """ - Allow reasoning parameters for models flagged as reasoning-capable. + Include reasoning parameters regardless of model-map capabilities. """ supported_params: Final = super().get_supported_openai_params(model=model) - try: - if litellm.supports_reasoning(model=model, custom_llm_provider="openrouter") or litellm.supports_reasoning( - model=model - ): - supported_params.append("reasoning_effort") - supported_params.append("thinking") - except Exception: - pass - return list(dict.fromkeys(supported_params)) + return list(dict.fromkeys([*supported_params, "reasoning_effort", "thinking"])) def map_openai_params( self, diff --git a/litellm/llms/openrouter/decisions/transformation.py b/litellm/llms/openrouter/decisions/transformation.py index 53ddce8b4029..0fa5c7ca7cb1 100644 --- a/litellm/llms/openrouter/decisions/transformation.py +++ b/litellm/llms/openrouter/decisions/transformation.py @@ -1,4 +1,7 @@ +from typing import Final + from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig +from litellm.types.decisions import DecisionsIRResponse class OpenRouterDecisionsConfig(BaseDecisionsConfig): @@ -8,3 +11,9 @@ class OpenRouterDecisionsConfig(BaseDecisionsConfig): def get_default_api_base(self) -> str | None: return "https://openrouter.ai/api" + + def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None: + cost: Final = response.usage.extra.get("cost") + if not isinstance(cost, (int, float)): + return None + return float(cost) diff --git a/litellm/llms/openrouter/responses/transformation.py b/litellm/llms/openrouter/responses/transformation.py index 9e32356c2fa0..aeb5b09e14a8 100644 --- a/litellm/llms/openrouter/responses/transformation.py +++ b/litellm/llms/openrouter/responses/transformation.py @@ -8,11 +8,25 @@ Docs: https://openrouter.ai/docs/api/reference/responses/overview """ +from collections.abc import Mapping from typing import Final +import httpx +from pydantic import TypeAdapter + import litellm +from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER +from litellm.litellm_core_utils.hidden_params import set_hidden_params +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders @@ -33,6 +47,62 @@ class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.OPENROUTER + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIResponse: + response: Final = super().transform_response_api_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + self._set_provider_reported_cost(response) + return response + + def transform_streaming_response( + self, + model: str, + parsed_chunk: Mapping[str, object], + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIStreamingResponse: + response_event: Final = super().transform_streaming_response( + model=model, + parsed_chunk=dict(parsed_chunk), + logging_obj=logging_obj, + ) + terminal_response: Final = ( + response_event.response + if isinstance( + response_event, + (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent), + ) + else None + ) + if terminal_response is not None: + self._set_provider_reported_cost(terminal_response) + return response_event + + @staticmethod + def _set_provider_reported_cost(response: ResponsesAPIResponse) -> None: + cost: Final = response.usage.cost if response.usage is not None else None + if cost is None: + return + additional_headers: Final = TypeAdapter(dict[str, object]).validate_python( + response.hidden_params.get("additional_headers", {}) + ) + set_hidden_params( + response, + { + **response.hidden_params, + "additional_headers": { + **additional_headers, + RESPONSE_COST_HEADER: float(cost), + }, + }, + ) + def validate_environment( self, headers: dict, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/typesafe_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/typesafe_passthrough_logging_handler.py index 9f8318f9b744..c24915c7126d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/typesafe_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/typesafe_passthrough_logging_handler.py @@ -21,6 +21,7 @@ class _TypeSafeUsage(LiteLLMBaseModel): input_tokens: int = 0 output_tokens: int = 0 + cost: float | None = None class _TypeSafeResponse(LiteLLMBaseModel): @@ -88,7 +89,9 @@ def typesafe_passthrough_handler( ) pricing: Final = _pricing_for(candidate_model_keys) response_cost: Final = ( - input_tokens * pricing.input_cost_per_token + output_tokens * pricing.output_cost_per_token + usage.cost + if custom_llm_provider == "openrouter" and usage.cost is not None + else input_tokens * pricing.input_cost_per_token + output_tokens * pricing.output_cost_per_token ) usage_object: Final = Usage( prompt_tokens=input_tokens, diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index e61739262697..27a9783cf436 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -69,6 +69,7 @@ class _Provider: api_key: str | None wraps_result: bool cost_map_key: str | None + provider_reported_cost: float | None = None _PROVIDERS: Final = ( @@ -89,7 +90,8 @@ class _Provider: "typesafe/jev-1.13", _API_KEY, False, - "openrouter/typesafe/jev-1.13", + None, + 1.5834e-5, ), _Provider( "strands_decider", "strands_decider/systemone-decider", "/v1/systemone", "systemone-decider", None, False, None @@ -136,17 +138,32 @@ def _number(value: JsonValue) -> float: return float(value) -def _expected_spend(cost_map_key: str | None) -> float: - if cost_map_key is None: +def _expected_spend(provider: _Provider) -> float: + if provider.provider_reported_cost is not None: + return provider.provider_reported_cost + if provider.cost_map_key is None: return 0.0 - prices: Final = object_value(json.loads(Path("model_prices_and_context_window.json").read_text())[cost_map_key]) + prices: Final = object_value( + json.loads(Path("model_prices_and_context_window.json").read_text())[provider.cost_map_key] + ) return _number(_USAGE["input_tokens"]) * _number(prices["input_cost_per_token"]) + _number( _USAGE["output_tokens"] ) * _number(prices["output_cost_per_token"]) +def _usage(provider: _Provider) -> dict[str, JsonValue]: + return { + **_USAGE, + **({"cost": provider.provider_reported_cost} if provider.provider_reported_cost is not None else {}), + } + + def _answer_body(provider: _Provider) -> dict[str, JsonValue]: - answer: Final[dict[str, JsonValue]] = {"model": provider.body_model, "answers": _ANSWERS, "usage": _USAGE} + answer: Final[dict[str, JsonValue]] = { + "model": provider.body_model, + "answers": _ANSWERS, + "usage": _usage(provider), + } return {"result": answer, "success": True} if provider.wraps_result else answer @@ -229,16 +246,14 @@ def _pass_through_config(directory: Path, pass_through_target: str, native_api_b @pytest.mark.parametrize("provider", _PROVIDERS, ids=lambda provider: provider.name) -def test_each_provider_gets_its_own_path_key_and_body_and_is_billed_from_the_cost_map( - gateway: Gateway, provider: _Provider -) -> None: - expected_spend: Final = _expected_spend(provider.cost_map_key) +def test_each_provider_gets_its_own_path_key_and_body_and_is_billed(gateway: Gateway, provider: _Provider) -> None: + expected_spend: Final = _expected_spend(provider) with gateway.scenario() as scenario: handle: Final = _register(scenario, _answer_body(provider)) model: Final = _deployment(scenario, handle, provider) response: Final = _decide(gateway, model) assert response.status_code == 200, response.text - assert response.json() == {"model": provider.body_model, "answers": _ANSWERS, "usage": _USAGE} + assert response.json() == {"model": provider.body_model, "answers": _ANSWERS, "usage": _usage(provider)} assert response.headers["x-litellm-model-group"] == model assert math.isclose(float(response.headers.get("x-litellm-response-cost", "0")), expected_spend, rel_tol=1e-9) (call,) = _upstream_calls(gateway, handle) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index bf90364fbbb8..ec5375bc71eb 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -11,6 +11,7 @@ import respx import litellm +from litellm.cost_calculator import get_response_cost_from_hidden_params from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.types.decisions import ( @@ -1057,6 +1058,36 @@ async def test_requests_a_provider_cannot_serve_are_rejected_before_http( assert len(respx_mock.calls) == 0 +@pytest.mark.asyncio +async def test_openrouter_decisions_uses_provider_reported_cost_without_cost_map( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delitem(litellm.model_cost, "openrouter/typesafe/jev-1.13", raising=False) + cost: Final = 1.1802e-5 + route: Final = respx_mock.post("https://openrouter.ai/api/alpha/decisions").respond( + json={ + "model": "typesafe/jev-1.13", + "answers": _RESPONSE["answers"], + "usage": { + "input_tokens": _INPUT_TOKENS, + "output_tokens": _OUTPUT_TOKENS, + "cost": cost, + }, + } + ) + + response: Final = await litellm.adecisions( + model="openrouter/typesafe/jev-1.13", + state={"source": "unit-test"}, + questions=_QUESTIONS, + api_key="caller-key", + ) + + assert route.called + assert get_response_cost_from_hidden_params(response.hidden_params) == cost + + @pytest.mark.asyncio @pytest.mark.parametrize( "api_base", diff --git a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py index b177b80aed10..beffcf9b69c2 100644 --- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -96,6 +96,15 @@ def test_openrouter_extra_body_transformation(): ] +def test_openrouter_supported_params_include_reasoning_for_unmapped_model(): + supported_params = OpenrouterConfig().get_supported_openai_params( + "openrouter/unmapped-model" + ) + + assert "reasoning_effort" in supported_params + assert "thinking" in supported_params + + def test_openrouter_cache_control_flag_removal(): transformed_request = OpenrouterConfig().transform_request( model="openrouter/deepseek/deepseek-chat", @@ -537,17 +546,15 @@ def test_openrouter_reasoning_models_allow_reasoning_effort_param(): assert supported_params.count("reasoning_effort") == 1 -def test_openrouter_non_reasoning_models_do_not_add_reasoning_effort(): - """ - Models without reasoning support should not gain reasoning-specific params. - """ +def test_openrouter_non_reasoning_models_support_reasoning_params(): config = OpenrouterConfig() supported_params = config.get_supported_openai_params( model="openrouter/anthropic/claude-3-5-haiku" ) - assert "reasoning_effort" not in supported_params + assert "reasoning_effort" in supported_params + assert "thinking" in supported_params def test_openrouter_reasoning_effort_max_maps_to_xhigh(): diff --git a/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py index d3ea8d5b907b..19517e58b310 100644 --- a/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py +++ b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py @@ -9,12 +9,22 @@ Related issue: https://github.com/BerriAI/litellm/issues/22189 """ +from unittest.mock import MagicMock + +import httpx import pytest import litellm +from litellm.cost_calculator import get_response_cost_from_hidden_params +from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER from litellm.llms.openrouter.responses.transformation import ( OpenRouterResponsesAPIConfig, ) +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, +) from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -116,3 +126,66 @@ def test_openrouter_not_using_completion_bridge(self): # The URL should point to OpenRouter's responses endpoint url = config.get_complete_url(api_base=None, litellm_params={}) assert "/responses" in url + + +@pytest.mark.parametrize("cost", [1.23e-5, None]) +def test_openrouter_responses_returns_provider_reported_cost(cost): + usage: dict[str, object] = { + "input_tokens": 2, + "output_tokens": 3, + "total_tokens": 5, + } + if cost is not None: + usage["cost"] = cost + response = OpenRouterResponsesAPIConfig().transform_response_api_response( + model="openrouter/unmapped-model", + raw_response=httpx.Response( + 200, + json={ + "id": "resp_test", + "created_at": 1, + "model": "unmapped-model", + "object": "response", + "output": [], + "status": "completed", + "usage": usage, + }, + ), + logging_obj=MagicMock(), + ) + + assert get_response_cost_from_hidden_params(response.hidden_params) == cost + + +@pytest.mark.parametrize( + "event_type", + ["response.completed", "response.incomplete", "response.failed"], +) +def test_openrouter_terminal_stream_events_return_provider_reported_cost(event_type): + event = OpenRouterResponsesAPIConfig().transform_streaming_response( + model="openrouter/unmapped-model", + parsed_chunk={ + "type": event_type, + "response": { + "id": "resp_test", + "created_at": 1, + "model": "unmapped-model", + "object": "response", + "output": [], + "status": "completed", + "usage": { + "input_tokens": 2, + "output_tokens": 3, + "total_tokens": 5, + "cost": 1.23e-5, + }, + }, + }, + logging_obj=MagicMock(), + ) + assert isinstance( + event, + (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent), + ) + + assert event.response.hidden_params["additional_headers"][RESPONSE_COST_HEADER] == 1.23e-5 diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py index 7a411c95afae..07b52711dc35 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py @@ -14,7 +14,7 @@ @pytest.fixture(autouse=True) -def local_model_cost_map(monkeypatch: pytest.MonkeyPatch): +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) @@ -222,14 +222,17 @@ async def test_oss_gateway_accounts_for_checkpoint_usage_and_registered_cost( await budget_limiter.is_key_within_model_budget(auth, f"{provider}/{requested}") -def test_openrouter_decisions_response_is_priced_from_request_model_registry_row(): +def test_openrouter_decisions_response_uses_provider_cost_without_cost_map_row( + monkeypatch: pytest.MonkeyPatch, +): logging_obj = _logging_obj() - model_cost = litellm.model_cost["openrouter/typesafe/jev-1.13"] + cost: Final = 1.5834e-5 + monkeypatch.delitem(litellm.model_cost, "openrouter/typesafe/jev-1.13", raising=False) response = TypeSafePassthroughLoggingHandler.typesafe_passthrough_handler( httpx_response=_response(), response_body={ "model": "typesafe/jev-1.13-20260917", - "usage": {"input_tokens": 282, "output_tokens": 20}, + "usage": {"input_tokens": 282, "output_tokens": 20, "cost": cost}, }, logging_obj=logging_obj, url_route="https://openrouter.ai/api/alpha/decisions", @@ -241,10 +244,9 @@ def test_openrouter_decisions_response_is_priced_from_request_model_registry_row custom_llm_provider="openrouter", ) - expected_cost = 282 * model_cost["input_cost_per_token"] + 20 * model_cost["output_cost_per_token"] assert response["kwargs"]["model"] == "openrouter/typesafe/jev-1.13-20260917" assert response["kwargs"]["custom_llm_provider"] == "openrouter" - assert response["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert response["kwargs"]["response_cost"] == cost assert response["kwargs"]["combined_usage_object"].prompt_tokens == 282 assert response["kwargs"]["combined_usage_object"].completion_tokens == 20 assert response["kwargs"]["combined_usage_object"].total_tokens == 302 From ddd4ee61c35ef8cfac7cadaa5836cf48bdfa2e92 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:29:18 -0700 Subject: [PATCH 03/18] test(decisions): add hosted_vllm /v1/decisions and provider-400 translation cases (#45661) Backport of #45661 to rc/1.106.0. Cherry-picked from 097d018dbf962318a582ee805994f3a32b3e66d9 (main). --- tests/integration/translation/case.py | 4 +- .../decisions/bases/hosted_vllm.py | 108 ++++++++++++++++++ .../basic/test_decisions_basic_hosted_vllm.py | 12 +- tests/integration/translation/runner.py | 2 +- 4 files changed, 122 insertions(+), 4 deletions(-) diff --git a/tests/integration/translation/case.py b/tests/integration/translation/case.py index ac8f43439cee..b64edea21a1c 100644 --- a/tests/integration/translation/case.py +++ b/tests/integration/translation/case.py @@ -7,7 +7,8 @@ @dataclass(frozen=True, slots=True, kw_only=True) class TranslationTestCase: """One request through the proxy to a deployment in `proxy_config.yaml`: what the test sends to LiteLLM, the - exact request the provider must receive, the fake provider's reply, and the exact response LiteLLM must return.""" + exact request the provider must receive, the fake provider's status and reply, and the exact response LiteLLM must + return.""" scenario: str litellm_endpoint: str @@ -15,6 +16,7 @@ class TranslationTestCase: expected_provider_endpoint: str expected_provider_headers: Mapping[str, str] expected_provider_request: Mapping[str, JsonValue] + mock_provider_status_code: int = 200 mock_provider_response: Mapping[str, JsonValue] expected_litellm_status_code: int = 200 expected_litellm_response: Mapping[str, JsonValue] diff --git a/tests/integration/translation/decisions/bases/hosted_vllm.py b/tests/integration/translation/decisions/bases/hosted_vllm.py index 2da37f83a861..ad7759b97174 100644 --- a/tests/integration/translation/decisions/bases/hosted_vllm.py +++ b/tests/integration/translation/decisions/bases/hosted_vllm.py @@ -1,3 +1,5 @@ +import json +from dataclasses import replace from typing import Final from integration.translation.case import TranslationTestCase @@ -65,3 +67,109 @@ "diagnostics": {"is_urgent": {"label_mass": 0.9979320910923947, "argmax_is_label": True}}, }, ) + +"""OpenAI Decisions client shape in, vLLM /v1/systemone body out. Mock reply captured live on 2026-10-09 from the same +vLLM build. +""" +QWEN3_0_6B_DECISIONS_TEST_CASE: Final = replace( + QWEN3_0_6B_TEST_CASE, + scenario="decisions", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "hosted_vllm/Qwen/Qwen3-0.6B", + "input": "The checkout page crashed twice today", + "questions": [ + { + "type": "choice", + "name": "is_defect", + "instructions": "Is this a defect?", + "choices": [{"value": "yes"}, {"value": "no"}], + } + ], + "cache": {"no-cache": True}, + }, + expected_provider_request={ + "model": "Qwen/Qwen3-0.6B", + "state": "The checkout page crashed twice today", + "questions": { + "is_defect": {"type": "choice", "instructions": "Is this a defect?", "criteria": {"yes": None, "no": None}} + }, + }, + mock_provider_response={ + "id": "decision-91ce41627dbf016f", + "object": "structured_decision", + "created": 1791577835, + "model": "Qwen/Qwen3-0.6B", + "answers": { + "is_defect": { + "type": "choice", + "choice": "yes", + "probabilities": {"yes": 0.8670357477770336, "no": 0.13296425222296632}, + "confidence": 0.8669148648509282, + } + }, + "usage": {"input_tokens": 43, "output_tokens": 1}, + "diagnostics": {"is_defect": {"label_mass": 0.9998605790748358, "argmax_is_label": True}}, + }, + expected_litellm_response={ + "model": "Qwen/Qwen3-0.6B", + "answers": [ + { + "type": "choice", + "name": "is_defect", + "choice": "yes", + "probabilities": [ + {"value": "yes", "probability": 0.8670357477770336}, + {"value": "no", "probability": 0.13296425222296632}, + ], + "confidence": 0.8669148648509282, + } + ], + "usage": { + "input_tokens": 43, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 1, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 44, + }, + }, +) + +"""vLLM only scores `choice` questions, so an OpenAI `predicate` goes out as `noul` and comes back as vLLM's 400, +which LiteLLM passes through as a 400. +""" +VLLM_NOUL_REJECTED: Final = { + "error": { + "message": "unknown question type 'noul'; supported: ['choice']", + "type": "BadRequestError", + "param": None, + "code": 400, + } +} +QWEN3_0_6B_PREDICATE_REJECTED_TEST_CASE: Final = replace( + QWEN3_0_6B_DECISIONS_TEST_CASE, + scenario="predicate_rejected", + litellm_request={ + "model": "hosted_vllm/Qwen/Qwen3-0.6B", + "input": "The checkout page crashed twice today", + "questions": [{"type": "predicate", "name": "is_bug", "instructions": "Is this a bug?"}], + "cache": {"no-cache": True}, + }, + expected_provider_request={ + "model": "Qwen/Qwen3-0.6B", + "state": "The checkout page crashed twice today", + "questions": {"is_bug": {"type": "noul", "instructions": "Is this a bug?"}}, + }, + mock_provider_status_code=400, + mock_provider_response=VLLM_NOUL_REJECTED, + expected_litellm_status_code=400, + expected_litellm_response={ + "error": { + "message": f"litellm.BadRequestError: Hosted_vllmException - {json.dumps(VLLM_NOUL_REJECTED)}\n\n" + "LiteLLM: model group 'hosted_vllm/Qwen/Qwen3-0.6B' failed with the error above. No fallback was attempted.", + "type": "invalid_request_error", + "param": None, + "code": "400", + } + }, +) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_hosted_vllm.py b/tests/integration/translation/decisions/basic/test_decisions_basic_hosted_vllm.py index f52d9147273d..c53f4277de51 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_hosted_vllm.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_hosted_vllm.py @@ -2,10 +2,18 @@ from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.hosted_vllm import QWEN3_0_6B_TEST_CASE +from integration.translation.decisions.bases.hosted_vllm import ( + QWEN3_0_6B_DECISIONS_TEST_CASE, + QWEN3_0_6B_PREDICATE_REJECTED_TEST_CASE, + QWEN3_0_6B_TEST_CASE, +) from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [QWEN3_0_6B_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize( + "case", + [QWEN3_0_6B_TEST_CASE, QWEN3_0_6B_DECISIONS_TEST_CASE, QWEN3_0_6B_PREDICATE_REJECTED_TEST_CASE], + ids=lambda case: case.id, +) def test_decisions_basic_hosted_vllm(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/runner.py b/tests/integration/translation/runner.py index eeb7e116c50f..427c55e9cc87 100644 --- a/tests/integration/translation/runner.py +++ b/tests/integration/translation/runner.py @@ -10,7 +10,7 @@ def assert_translation(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: - provider.expect(Reply(body=json.dumps(case.mock_provider_response).encode())) + provider.expect(Reply(status=case.mock_provider_status_code, body=json.dumps(case.mock_provider_response).encode())) response: Final = gateway.request("POST", case.litellm_endpoint, case.litellm_request) received: Final = provider.received() assert [(request.method, request.target) for request in received] == [("POST", case.expected_provider_endpoint)] From 3c42f96a9d7e68b63e07eefbdd93b99c7b48083c Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:29:32 -0700 Subject: [PATCH 04/18] feat(decisions): add Databricks ai_decide as a /v1/decisions provider and auto-router decider (#45200) Backport of #45200 to rc/1.106.0. Cherry-picked from 883e3210ee4075a556d0b26a03866a2c704809e8 (main). --- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 2 + litellm/decisions/main.py | 5 +- .../databricks/decisions/transformation.py | 57 + ...odel_prices_and_context_window_backup.json | 12 + .../common_utils/registry_read_through.py | 34 +- .../auto_router_permissions.py | 2 +- .../complexity_router/complexity_router.py | 9 + .../complexity_router/config.py | 23 +- .../complexity_router/jev_classifier.py | 36 +- litellm/utils.py | 2 + model_prices_and_context_window.json | 12 + ...uto_router_databricks_classifier_writes.py | 334 ++++++ .../management/test_model_health_check.py | 40 + .../test_databricks_decisions_wire.py | 519 +++++++++ .../providers/test_decisions_wire.py | 9 + ...complexity_router_databricks_classifier.py | 1023 +++++++++++++++++ tests/unit/decisions/test_main.py | 174 +++ .../llms/databricks/decisions/__init__.py | 0 .../decisions/test_transformation.py | 56 + .../test_registry_read_through.py | 66 +- .../decisions_endpoints/test_endpoints.py | 45 + .../test_auto_router_permissions.py | 7 +- .../complexity_router/test_jev_classifier.py | 104 +- .../add_model/AutoRouterClassifierTabs.tsx | 18 +- .../add_model/ClassifierTypeRadios.tsx | 4 +- .../JevClassifierConfig.integration.test.tsx | 13 +- .../add_model/JevClassifierConfig.tsx | 15 +- ...d_auto_router_routing_test_request.test.ts | 11 +- .../build_complexity_router_config.test.ts | 6 + .../add_model/jev_classifier_config.test.ts | 44 + .../add_model/jev_classifier_config.ts | 48 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 33 files changed, 2675 insertions(+), 60 deletions(-) create mode 100644 litellm/llms/databricks/decisions/transformation.py create mode 100644 tests/integration/management/test_auto_router_databricks_classifier_writes.py create mode 100644 tests/integration/providers/test_databricks_decisions_wire.py create mode 100644 tests/integration/routing/test_complexity_router_databricks_classifier.py create mode 100644 tests/unit/llms/databricks/decisions/__init__.py create mode 100644 tests/unit/llms/databricks/decisions/test_transformation.py create mode 100644 ui/litellm-dashboard/src/components/add_model/jev_classifier_config.test.ts diff --git a/litellm/__init__.py b/litellm/__init__.py index ebaef292e2ba..8da0771cc99a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1681,6 +1681,9 @@ def set_global_gitlab_config(config: Mapping[str, object]) -> None: from .llms.strands_decider.decisions.transformation import ( StrandsDeciderDecisionsConfig as StrandsDeciderDecisionsConfig, ) + from .llms.databricks.decisions.transformation import ( + DatabricksDecisionsConfig as DatabricksDecisionsConfig, + ) from .llms.hosted_vllm.decisions.transformation import ( HostedVLLMDecisionsConfig as HostedVLLMDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 9545ed16851c..9caba6a060a9 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -162,6 +162,7 @@ "OpenRouterDecisionsConfig", "CloudflareDecisionsConfig", "StrandsDeciderDecisionsConfig", + "DatabricksDecisionsConfig", "HostedVLLMDecisionsConfig", "OpenAIDecisionsConfig", "AzureAIDecisionsConfig", @@ -725,6 +726,7 @@ ".llms.strands_decider.decisions.transformation", "StrandsDeciderDecisionsConfig", ), + "DatabricksDecisionsConfig": (".llms.databricks.decisions.transformation", "DatabricksDecisionsConfig"), "HostedVLLMDecisionsConfig": ( ".llms.hosted_vllm.decisions.transformation", "HostedVLLMDecisionsConfig", diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index f1b83bb929a6..c47ba0335c7c 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -127,7 +127,6 @@ def _prepare_call( api_key=api_key, ) provider_config: Final = _provider_config(upstream_model, provider) - canonical_model: Final = provider_config.canonical_model(upstream_model) if not upstream_model: raise litellm.BadRequestError( message="A model name is required for the Decisions API", @@ -140,6 +139,10 @@ def _prepare_call( model=model, llm_provider=provider, ) + try: + canonical_model: Final = provider_config.canonical_model(upstream_model) + except ValueError as error: + raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error try: request: Final = _validate_request( state=state, questions=questions, decision_input=decision_input, safety_identifier=safety_identifier diff --git a/litellm/llms/databricks/decisions/transformation.py b/litellm/llms/databricks/decisions/transformation.py new file mode 100644 index 000000000000..a454da83e50d --- /dev/null +++ b/litellm/llms/databricks/decisions/transformation.py @@ -0,0 +1,57 @@ +import re +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Final + +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + +_SERVING_ENDPOINT_NAME: Final = re.compile(r"[A-Za-z0-9_-][A-Za-z0-9._-]*") + + +def validate_serving_endpoint_name(name: str) -> str: + if _SERVING_ENDPOINT_NAME.fullmatch(name) is None: + raise ValueError( + f"Databricks serving endpoint name {name!r} must be the bare endpoint name (letters, digits, '-', '_' " + "and '.', with no '/', '?', '#', spaces, or a leading '.'), e.g. databricks-openjev-qwen35-4b" + ) + return name + + +@dataclass(frozen=True, slots=True) +class DatabricksDecisionsConnection: + api_base: str + api_key: str = field(repr=False) + + +class DatabricksDecisionsConfig(BaseDecisionsConfig): + api_key_env = ("DATABRICKS_API_KEY", "DATABRICKS_TOKEN") + api_base_env = ("DATABRICKS_API_BASE",) + + def missing_api_base_message(self, custom_llm_provider: str) -> str: + return ( + f"api_base is required for Decisions provider '{custom_llm_provider}': set DATABRICKS_API_BASE to " + "https:///serving-endpoints" + ) + + def canonical_model(self, model: str) -> str: + return validate_serving_endpoint_name(model) + + def get_complete_url(self, api_base: str, model: str) -> str: + return f"{api_base.rstrip('/')}/{validate_serving_endpoint_name(model)}/invocations" + + def classifier_response(self, body: Mapping[str, object], requested_model: str) -> Mapping[str, object]: + return MappingProxyType({**body, "model": requested_model}) + + def connection(self, api_base: str | None, api_key: str | None) -> DatabricksDecisionsConnection: + base: Final = self.resolve_api_base(api_base) + key: Final = api_key if api_base is not None else self.resolve_api_key(api_key) + if not base or not key: + raise ValueError( + "Databricks requires api_key or DATABRICKS_API_KEY (or DATABRICKS_TOKEN) and api_base or " + "DATABRICKS_API_BASE pointing to https:///serving-endpoints" + ) + return DatabricksDecisionsConnection(api_base=base.rstrip("/"), api_key=key) + + +DATABRICKS_DECISIONS_CONFIG: Final[DatabricksDecisionsConfig] = DatabricksDecisionsConfig() diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 28a920eb4f73..525c4199ddee 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21165,6 +21165,18 @@ "output_vector_size": 1024, "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, + "databricks/databricks-openjev-qwen35-4b": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "databricks", + "metadata": { + "notes": "Pay-per-token Foundation Model API decision endpoint (OpenJev, Qwen 3.5 4B) queried through the System One API. Databricks published no DBU rate for it as of 2026-10-07, so a deployment's model_info prices set the billed cost." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" + }, "dataforseo/search": { "input_cost_per_query": 0.003, "litellm_provider": "dataforseo", diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 0a47de96645a..8f64072b4d0a 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -7,18 +7,24 @@ On a registry miss, callers here fetch the missing row from the DB and load it into the local registry before giving up. A short negative-result TTL per key plus a global resync budget per window bound the DB load from lookups of -genuinely unknown names. +genuinely unknown names. An auto router's tiers are deployment rows of their +own, so a miss on an auto router reconciles every row instead of just its own, +or the first routed request on this replica would pick a tier it has not loaded. """ import asyncio import time -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper if TYPE_CHECKING: + from prisma import models as prisma_models from prisma.types import ( LiteLLM_AgentsTableInclude, LiteLLM_AgentsTableWhereUniqueInput, @@ -129,6 +135,11 @@ async def _resync_model_deployments(model_name: str) -> bool: prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj ) return proxy_server.llm_router is not None + if _holds_an_auto_router(rows): + await proxy_server.proxy_config.add_deployment( + prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj + ) + return _model_is_loaded(model_name) async with proxy_server.MODEL_RECONCILE_LOCK: await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client) proxy_server.proxy_config._add_deployment(db_models=rows) @@ -136,6 +147,25 @@ async def _resync_model_deployments(model_name: str) -> bool: return True +def _holds_an_auto_router(rows: Sequence["prisma_models.LiteLLM_ProxyModelTable"]) -> bool: + return any(_names_an_auto_router(row.litellm_params) for row in rows) + + +_STORED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + +def _names_an_auto_router(stored_litellm_params: object) -> bool: + try: + params: Final = _STORED_LITELLM_PARAMS.validate_python(stored_litellm_params) + except ValidationError: + return False + stored_model: Final = params.get("model") + if not isinstance(stored_model, str): + return False + model: Final = decrypt_value_helper(value=stored_model, key="model", return_original_value=True) + return isinstance(model, str) and model.startswith("auto_router/") + + async def _resync_guardrails(guardrail_name: str) -> bool: from litellm.proxy import proxy_server from litellm.proxy.guardrails.guardrail_registry import ( diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py index 16f6b512e0d8..e74578688c3e 100644 --- a/litellm/proxy/management_helpers/auto_router_permissions.py +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -76,7 +76,7 @@ class _MemberOpenSourceClassifierConfig(LiteLLMBaseModel): model_config = ConfigDict(extra="forbid") - provider: Literal["jev", "laya", "bespoke"] = "jev" + provider: Literal["jev", "laya", "bespoke", "databricks"] = "jev" model: str api_key: None = None api_base: None = None diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index acc9770ab41a..a5f00aa7e254 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -59,6 +59,7 @@ from litellm.llms.anthropic.common_utils import is_claude_code_user_agent from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.context_compaction import compaction_pending from litellm.router_strategy.complexity_router.tier_predictor import ( @@ -1342,6 +1343,14 @@ def _build_jev_client(config: OpenSourceClassifierConfig) -> JevClassifierClient http_client=get_async_httpx_client(httpxSpecialProvider.PassThroughEndpoint), provider=config.provider, ) + if config.provider == "databricks": + databricks: Final = DATABRICKS_DECISIONS_CONFIG.connection(config.api_base, config.api_key) + return HttpJevClassifierClient( + api_key=databricks.api_key, + api_base=databricks.api_base, + http_client=get_async_httpx_client(httpxSpecialProvider.PassThroughEndpoint), + provider="databricks", + ) api_key: Final = config.api_key or get_secret_str("TYPESAFE_API_KEY") if not api_key: raise ValueError( diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 797db460422d..36c19e37b8de 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -705,10 +705,18 @@ def normalize_classifier_config_aliases(config: Mapping[str, object]) -> Mapping return normalized +_ENVIRONMENT_KEY_SCOPE: Final = MappingProxyType( + { + "jev": "TYPESAFE_API_KEY is only sent to TYPESAFE_API_BASE or https://api.typesafe.ai", + "databricks": "DATABRICKS_API_KEY or DATABRICKS_TOKEN is only sent to DATABRICKS_API_BASE", + } +) + + class OpenSourceClassifierConfig(LiteLLMBaseModel): model_config = ConfigDict(extra="forbid", frozen=True) - provider: Literal["jev", "laya", "bespoke"] = "jev" + provider: Literal["jev", "laya", "bespoke", "databricks"] = "jev" model: str = "jev-latest" api_key: str | None = Field(default=None, description="Provider API key; optional for self-hosted providers") api_base: str | None = Field( @@ -753,10 +761,19 @@ def _keep_the_environment_key_on_the_environment_base(self) -> "OpenSourceClassi if self.api_base is not None: _ = validate_oss_api_base(self.provider, self.api_base) return self + if self.provider == "databricks": + from litellm.llms.databricks.decisions.transformation import validate_serving_endpoint_name + + if "model" not in self.model_fields_set: + raise ValueError( + "opensource_classifier_config.model is required for provider 'databricks': the serving endpoint " + "name, e.g. databricks-openjev-qwen35-4b" + ) + _ = validate_serving_endpoint_name(self.model) if self.api_base is not None and self.api_key is None: raise ValueError( - "opensource_classifier_config.api_base requires opensource_classifier_config.api_key: TYPESAFE_API_KEY is only sent " - "to TYPESAFE_API_BASE or https://api.typesafe.ai" + "opensource_classifier_config.api_base requires opensource_classifier_config.api_key: " + f"{_ENVIRONMENT_KEY_SCOPE[self.provider]}" ) return self diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index 93c8b0c1bd37..f68462af8e03 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG from litellm.llms.laya.common_utils import laya_response_model from litellm.proxy.pass_through_endpoints.llm_provider_handlers.typesafe_passthrough_logging_handler import ( TypeSafePassthroughLoggingHandler, @@ -27,6 +28,7 @@ from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN JevProbability: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)] +ClassifierProvider: TypeAlias = Literal["typesafe", "laya", "bespoke", "databricks"] DEFAULT_JEV_INSTRUCTIONS: Final = _DEFAULT_JEV_INSTRUCTIONS @@ -85,12 +87,12 @@ def __init__( api_key: str | None, api_base: str, http_client: AsyncHTTPHandler, - provider: Literal["typesafe", "laya", "bespoke"] = "typesafe", + provider: ClassifierProvider = "typesafe", ) -> None: self._api_key = api_key self._api_base = api_base.rstrip("/") self._http_client = http_client - self._provider = provider + self._provider: ClassifierProvider = provider async def evaluate( self, @@ -103,33 +105,43 @@ async def evaluate( MappingProxyType({"Authorization": f"Bearer {self._api_key}"}) if self._api_key else MappingProxyType({}) ) response: Final = await self._http_client.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler has a dynamic post signature - f"{self._api_base}/v1/systemone", + self._request_url(request.model), json=request.model_dump(mode="json"), headers=MappingProxyType({**authorization, "Content-Type": "application/json"}), # pyright: ignore[reportArgumentType] # HTTP headers are not mutated by AsyncHTTPHandler timeout=timeout_s, ) response.raise_for_status() body: Final = TypeAdapter(dict[str, object]).validate_json(response.content) - normalized_body: Final = ( - MappingProxyType({**body, "model": laya_response_model(body, request.model)}) - if self._provider == "laya" - else body - ) + normalized_body: Final = self._normalized_body(body, request.model) try: - self._log_response(request, response, request_kwargs, start_time) + self._log_response(request, response, normalized_body, request_kwargs, start_time) except Exception as exc: # noqa: BLE001 # logging integrations must not discard a provider verdict verbose_router_logger.warning("JEV response logging failed (%s)", type(exc).__name__) return TypeAdapter(JevSystemOneResponse).validate_python(normalized_body) + def _request_url(self, model: str) -> str: + if self._provider == "databricks": + return DATABRICKS_DECISIONS_CONFIG.get_complete_url(self._api_base, model) + return f"{self._api_base}/v1/systemone" + + def _normalized_body(self, body: Mapping[str, object], requested_model: str) -> Mapping[str, object]: + match self._provider: + case "laya": + return MappingProxyType({**body, "model": laya_response_model(body, requested_model)}) + case "databricks": + return DATABRICKS_DECISIONS_CONFIG.classifier_response(body, requested_model) + case "typesafe" | "bespoke": + return body + def _log_response( self, request: JevSystemOneRequest, response: httpx.Response, + body: Mapping[str, object], request_kwargs: Mapping[str, object] | None, start_time: datetime, ) -> None: try: - body: Final = TypeAdapter(dict[str, object]).validate_json(response.content) _ = TypeAdapter(JevUsage | None).validate_python(body.get("usage")) except ValidationError: return @@ -202,7 +214,7 @@ class JevVerdict(NamedTuple): confidence: float model: str cost: float | None - provider: Literal["typesafe", "laya", "bespoke"] = "typesafe" + provider: ClassifierProvider = "typesafe" class _RegistryPricing(LiteLLMBaseModel): @@ -226,7 +238,7 @@ def build_jev_request( def jev_classifier_cost( - response: JevSystemOneResponse, configured_model: str, provider: Literal["typesafe", "laya", "bespoke"] = "typesafe" + response: JevSystemOneResponse, configured_model: str, provider: ClassifierProvider = "typesafe" ) -> float | None: usage: Final = response.usage if usage is None: diff --git a/litellm/utils.py b/litellm/utils.py index 35bb8df37d5f..488c57ef2feb 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8989,6 +8989,8 @@ def get_provider_decisions_config(model: str, provider: LlmProviders) -> BaseDec return litellm.CloudflareDecisionsConfig() if provider == LlmProviders.STRANDS_DECIDER: return litellm.StrandsDeciderDecisionsConfig() + if provider == LlmProviders.DATABRICKS: + return litellm.DatabricksDecisionsConfig() if provider == LlmProviders.HOSTED_VLLM: return litellm.HostedVLLMDecisionsConfig() if provider == LlmProviders.OPENAI: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 28a920eb4f73..525c4199ddee 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21165,6 +21165,18 @@ "output_vector_size": 1024, "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, + "databricks/databricks-openjev-qwen35-4b": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "databricks", + "metadata": { + "notes": "Pay-per-token Foundation Model API decision endpoint (OpenJev, Qwen 3.5 4B) queried through the System One API. Databricks published no DBU rate for it as of 2026-10-07, so a deployment's model_info prices set the billed cost." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" + }, "dataforseo/search": { "input_cost_per_query": 0.003, "litellm_provider": "dataforseo", diff --git a/tests/integration/management/test_auto_router_databricks_classifier_writes.py b/tests/integration/management/test_auto_router_databricks_classifier_writes.py new file mode 100644 index 000000000000..6f541f5e6220 --- /dev/null +++ b/tests/integration/management/test_auto_router_databricks_classifier_writes.py @@ -0,0 +1,334 @@ +import asyncio +import json +import uuid +from collections.abc import Callable, Mapping +from typing import Final + +import httpx +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery, chat_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +_ENDPOINT: Final = "databricks-openjev-qwen35-4b" +_SECOND_ENDPOINT: Final = "databricks-openjev-qwen35-4b-canary" +_API_KEY: Final = "synthetic-databricks-key" +_BARE_NAME_RULE: Final = ( + "must be the bare endpoint name (letters, digits, '-', '_' and '.', with no '/', '?', '#', spaces, or a " + "leading '.'), e.g. databricks-openjev-qwen35-4b" +) +_REJECTED_AT_LOAD: Final = "The router would drop this deployment at load time, so the write is rejected instead." +_MODEL_REQUIRED: Final = ( + "opensource_classifier_config.model is required for provider 'databricks': the serving endpoint name, " + "e.g. databricks-openjev-qwen35-4b" +) +_BASE_NEEDS_KEY: Final = ( + "opensource_classifier_config.api_base requires opensource_classifier_config.api_key: DATABRICKS_API_KEY or " + "DATABRICKS_TOKEN is only sent to DATABRICKS_API_BASE" +) +_BLANK_KEY: Final = ( + "opensource_classifier_config.api_key must be non-empty; omit it to use the provider environment key" +) +_CLASSIFIER_ANSWER: Final[dict[str, JsonValue]] = { + "answers": { + "tier": { + "type": "choice", + "choice": "COMPLEX", + "confidence": 0.91, + "probabilities": {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91}, + } + }, + "usage": {"input_tokens": 12, "output_tokens": 1}, +} +_ROWS_QUERY: Final = ( + "SELECT request_id, status, metadata->>'internal_call_origin' AS origin " + 'FROM "LiteLLM_SpendLogs" WHERE model_group = %s' +) + + +def _classifier(request: Request) -> Reply: + return Reply(body=json.dumps(_CLASSIFIER_ANSWER).encode()) + + +def _tier(text: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + body: Final = json.loads(request.body) + return chat_reply(f"{text}-{uuid.uuid4().hex[:8]}", "tier-model", text, stream=bool(body.get("stream"))) + + return answering_model_discovery(respond) + + +def _classifier_config(api_base: str, **overrides: JsonValue) -> dict[str, JsonValue]: + return { + "provider": "databricks", + "model": _ENDPOINT, + "api_base": api_base, + "api_key": _API_KEY, + "timeout_ms": 20000, + "circuit_breaker_enabled": False, + **overrides, + } + + +def _without(config: Mapping[str, JsonValue], field: str) -> dict[str, JsonValue]: + return {name: value for name, value in config.items() if name != field} + + +def _router_config(classifier: Mapping[str, JsonValue], simple: str, complex_: str) -> dict[str, JsonValue]: + return { + "classifier_type": "oss_classifier", + "opensource_classifier_config": dict(classifier), + "tiers": {"SIMPLE": simple, "MEDIUM": simple, "COMPLEX": complex_, "REASONING": complex_}, + } + + +def _validate(gateway: Gateway, config: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + response: Final = gateway.request( + "POST", "/auto_router/validate_complexity_router_config", {"complexity_router_config": dict(config)} + ) + assert response.status_code == 200, response.text + return object_value(response.json()) + + +def _create_router(scenario: Scenario, config: Mapping[str, JsonValue]) -> tuple[str, str]: + name: Final = f"router-{uuid.uuid4().hex[:10]}" + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": dict(config)}, + "model_info": {}, + }, + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return name, identity + + +def _stored_classifier(gateway: Gateway, identity: str) -> dict[str, JsonValue]: + stored: Final = gateway.request("GET", "/model/info", params={"litellm_model_id": identity}) + assert stored.status_code == 200, stored.text + (entry,) = stored.json()["data"] + router_config: Final = object_value(object_value(object_value(entry)["litellm_params"])["complexity_router_config"]) + return object_value(router_config["opensource_classifier_config"]) + + +def _judge_targets(judge: Wire) -> tuple[str, ...]: + return tuple(request.target for request in judge.drain() if request.method == "POST") + + +def _chat_body(model: str) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": f"design a distributed cache {uuid.uuid4().hex}"}]} + + +async def _burst(base_url: str, key: str, model: str, count: int) -> tuple[httpx.Response, ...]: + async with httpx.AsyncClient( + base_url=base_url, timeout=60, trust_env=False, headers={"Authorization": f"Bearer {key}"} + ) as client: + return tuple( + await asyncio.gather(*(client.post("/v1/chat/completions", json=_chat_body(model)) for _ in range(count))) + ) + + +def _rows(router: str, *, want: int) -> tuple[dict[str, JsonValue], ...]: + return tuple(eventually(lambda: read_rows(_ROWS_QUERY, (router,)), lambda found: len(found) >= want, seconds=70)) + + +def test_validate_accepts_a_databricks_classifier_with_a_bare_endpoint_name_of_any_length(gateway: Gateway) -> None: + config: Final = _router_config( + _classifier_config("https://dbc-example.cloud.databricks.com/serving-endpoints"), "simple", "complex" + ) + assert _validate(gateway, config) == {"valid": True, "error": None} + long_name: Final = _router_config( + _classifier_config("https://dbc-example.cloud.databricks.com/serving-endpoints", model="e" * 5000), + "simple", + "complex", + ) + assert _validate(gateway, long_name) == {"valid": True, "error": None} + + +def test_validate_names_each_databricks_classifier_misconfiguration(gateway: Gateway) -> None: + base: Final = _classifier_config("https://dbc-example.cloud.databricks.com/serving-endpoints") + cases: Final[tuple[tuple[str, dict[str, JsonValue], str], ...]] = ( + ("model omitted", _without(base, "model"), f"at opensource_classifier_config: Value error, {_MODEL_REQUIRED}"), + ( + "api_base without api_key", + _without(base, "api_key"), + f"at opensource_classifier_config: Value error, {_BASE_NEEDS_KEY}", + ), + ("slash in the name", {**base, "model": "a/b"}, f"Databricks serving endpoint name 'a/b' {_BARE_NAME_RULE}"), + ("space in the name", {**base, "model": "a b"}, f"Databricks serving endpoint name 'a b' {_BARE_NAME_RULE}"), + ("leading dot", {**base, "model": ".hidden"}, f"Databricks serving endpoint name '.hidden' {_BARE_NAME_RULE}"), + ("empty name", {**base, "model": ""}, f"Databricks serving endpoint name '' {_BARE_NAME_RULE}"), + ( + "blank api_key", + {**base, "api_key": ""}, + f"at opensource_classifier_config.api_key: Value error, {_BLANK_KEY}", + ), + ("int name", {**base, "model": 5}, "at opensource_classifier_config.model: Input should be a valid string"), + ( + "list name", + {**base, "model": ["a"]}, + "at opensource_classifier_config.model: Input should be a valid string", + ), + ("null name", {**base, "model": None}, "at opensource_classifier_config.model: Input should be a valid string"), + ) + for label, classifier, expected in cases: + verdict: Final = _validate(gateway, _router_config(classifier, "simple", "complex")) + assert verdict["valid"] is False, (label, verdict) + error: Final = string_value(verdict["error"]) + assert error.startswith("complexity_router_config is invalid "), (label, error) + assert expected in error, (label, error) + assert error.endswith(_REJECTED_AT_LOAD), (label, error) + + +def test_test_routing_classifies_through_the_databricks_endpoint_without_calling_the_tier_it_picked( + gateway: Gateway, +) -> None: + prompt: Final = f"design a distributed cache {uuid.uuid4().hex}" + with ( + wire_server(_classifier) as judge, + wire_server(_tier("simple answer")) as simple, + wire_server(_tier("complex answer")) as complex_tier, + gateway.scenario() as scenario, + ): + simple_model: Final = scenario.model(model="openai/simple-tier", api_base=simple.url) + complex_model: Final = scenario.model(model="openai/complex-tier", api_base=complex_tier.url) + config: Final = _router_config(_classifier_config(judge.url), simple_model, complex_model) + response: Final = gateway.request( + "POST", "/auto_router/test_routing", {"prompt": prompt, "complexity_router_config": config} + ) + assert response.status_code == 200, response.text + assert response.json() == { + "routed_model": complex_model, + "routed_model_configured": True, + "routing_decision": { + "router_model_name": "auto_router_routing_test", + "router_type": "complexity", + "routed_model": complex_model, + "cause": "jev_classifier", + "tier": "COMPLEX", + "signals": [ + "databricks-classifier:COMPLEX", + "databricks-confidence=0.910000", + "tier-probability:SIMPLE=0.030000", + "tier-probability:MEDIUM=0.060000", + "tier-probability:COMPLEX=0.910000", + ], + "classifier_model": f"databricks/{_ENDPOINT}", + "classifier_cost": 0.0, + "classifier_probabilities": {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91}, + "classifier_confidence": 0.91, + "conversation_continuing": False, + }, + }, response.text + (call,) = [request for request in judge.drain() if request.method == "POST"] + assert call.target == f"/{_ENDPOINT}/invocations", call.target + assert call.headers.get("authorization") == f"Bearer {_API_KEY}", call.headers + body: Final = json.loads(call.body) + assert body["model"] == _ENDPOINT, body + assert body["state"] == f"\nClassify this message:\n{prompt}", body + assert _judge_targets(simple) == () and _judge_targets(complex_tier) == () + + +def test_model_new_refuses_a_databricks_classifier_without_an_endpoint_name_and_stores_nothing( + gateway: Gateway, +) -> None: + name: Final = f"router-{uuid.uuid4().hex[:10]}" + config: Final = _router_config( + _without(_classifier_config("https://dbc-example.cloud.databricks.com/serving-endpoints"), "model"), + "simple", + "complex", + ) + response: Final = gateway.request( + "POST", + "/model/new", + { + "model_name": name, + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": config}, + "model_info": {}, + }, + ) + assert response.status_code == 400, response.text + error: Final = object_value(object_value(response.json())["error"]) + assert _MODEL_REQUIRED in string_value(error["message"]), error + assert (error["type"], error["param"]) == ("validation_error", "litellm_params.model"), error + listed: Final = gateway.get("/model/info")["data"] + assert isinstance(listed, list) + assert [entry for entry in map(object_value, listed) if entry["model_name"] == name] == [] + + +def test_model_update_refuses_dropping_the_endpoint_name_and_keeps_the_stored_classifier_serving( + gateway: Gateway, +) -> None: + with ( + wire_server(_classifier) as judge, + wire_server(_tier("simple answer")) as simple, + wire_server(_tier("complex answer")) as complex_tier, + gateway.scenario() as scenario, + ): + simple_model: Final = scenario.model(model="openai/simple-tier", api_base=simple.url) + complex_model: Final = scenario.model(model="openai/complex-tier", api_base=complex_tier.url) + classifier: Final = _classifier_config(judge.url) + name, identity = _create_router(scenario, _router_config(classifier, simple_model, complex_model)) + dropped: Final = _router_config(_without(classifier, "model"), simple_model, complex_model) + response: Final = gateway.request( + "PATCH", + f"/model/{identity}/update", + {"litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": dropped}}, + ) + assert response.status_code == 400, response.text + assert _MODEL_REQUIRED in string_value(object_value(object_value(response.json())["error"])["message"]), ( + response.text + ) + assert _stored_classifier(gateway, identity)["model"] == _ENDPOINT + chat: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(name)) + assert chat.status_code == 200, chat.text + assert chat.json()["choices"][0]["message"]["content"] == "complex answer", chat.text + assert chat.headers["x-litellm-complexity-router-cause"] == "jev_classifier", dict(chat.headers) + assert _judge_targets(judge) == (f"/{_ENDPOINT}/invocations",) + + +async def test_renaming_the_endpoint_under_a_burst_moves_the_classifier_and_logs_every_request_once( + gateway: Gateway, +) -> None: + base_url: Final = str(gateway.client.base_url) + with ( + wire_server(_classifier) as judge, + wire_server(_tier("simple answer")) as simple, + wire_server(_tier("complex answer")) as complex_tier, + gateway.scenario() as scenario, + ): + simple_model: Final = scenario.model(model="openai/simple-tier", api_base=simple.url) + complex_model: Final = scenario.model(model="openai/complex-tier", api_base=complex_tier.url) + classifier: Final = _classifier_config(judge.url) + name, identity = _create_router(scenario, _router_config(classifier, simple_model, complex_model)) + burst: Final = asyncio.create_task(_burst(base_url, gateway.key, name, 24)) + renamed: Final = _router_config({**classifier, "model": _SECOND_ENDPOINT}, simple_model, complex_model) + patched: Final = gateway.request( + "PATCH", + f"/model/{identity}/update", + {"litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": renamed}}, + ) + assert patched.status_code == 200, patched.text + during: Final = await burst + assert [response.status_code for response in during] == [200] * 24, [response.text for response in during] + assert set(_judge_targets(judge)) <= {f"/{_ENDPOINT}/invocations", f"/{_SECOND_ENDPOINT}/invocations"} + assert _stored_classifier(gateway, identity)["model"] == _SECOND_ENDPOINT + + def classifier_target() -> tuple[str, ...]: + chat: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(name)) + assert chat.status_code == 200, chat.text + return _judge_targets(judge) + + eventually(classifier_target, lambda targets: targets == (f"/{_SECOND_ENDPOINT}/invocations",), seconds=30) + after: Final = await _burst(base_url, gateway.key, name, 6) + assert [response.status_code for response in after] == [200] * 6, [response.text for response in after] + assert set(_judge_targets(judge)) == {f"/{_SECOND_ENDPOINT}/invocations"} + identities: Final = tuple(string_value(response.json()["id"]) for response in (*during, *after)) + assert len(set(identities)) == len(identities), identities + rows: Final = _rows(name, want=len(identities)) + request_rows: Final = [row for row in rows if row["origin"] != "autorouter_classifier"] + assert sorted(string_value(row["request_id"]) for row in request_rows) == sorted(identities), request_rows + assert all(row["status"] == "success" for row in rows), rows diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index 11b902b540a8..075c60513530 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -186,3 +186,43 @@ def test_evaluation_mode_health_check_of_hosted_vllm_sends_a_choice_probe(gatewa }, ) ] + + +_DATABRICKS_PROBE_REPLY: Final = JsonResponse( + content_type="application/json", + body={ + "model": "databricks-openjev-qwen35-4b", + "answers": {"reachable": {"type": "noul", "noul": 1.0}}, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, +) + + +def test_evaluation_mode_health_check_of_a_databricks_serving_endpoint_resolves_the_mode_from_the_cost_map( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + handle: Final = register_scenario(f"health-decisions-{uuid.uuid4().hex[:12]}", _DATABRICKS_PROBE_REPLY) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + model="databricks/databricks-openjev-qwen35-4b", + api_base=handle.api_base(), + api_key="synthetic-databricks-key", + ) + listed: Final = gateway.request("GET", "/v2/model/info", params={"model": model}) + assert listed.status_code == 200, listed.text + assert [object_value(object_value(entry)["model_info"])["mode"] for entry in listed.json()["data"]] == [ + "evaluation" + ], listed.text + report: Final = _health_report(gateway, model) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert _probes_sent_to(gateway, handle) == [ + ( + f"/{handle.scenario_id}/databricks-openjev-qwen35-4b/invocations", + { + "model": "databricks-openjev-qwen35-4b", + "state": os.environ.get("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"), + "questions": {"reachable": {"type": "noul", "instructions": "Is the service reachable?"}}, + }, + ) + ] diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py new file mode 100644 index 000000000000..5c3b7d1b1a1d --- /dev/null +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -0,0 +1,519 @@ +import threading +import time +import uuid +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse +from pydantic import JsonValue + +_ENDPOINT: Final = "databricks-openjev-qwen35-4b" +_API_KEY: Final = "synthetic-databricks-key" +_ENV_KEY: Final = "synthetic-databricks-env-key" +_ENV_TOKEN: Final = "synthetic-databricks-env-token" +_STATE: Final[dict[str, JsonValue]] = {"ticket": "The export job hangs at 99%", "component": "billing"} +_TIER_CRITERIA: Final[dict[str, JsonValue]] = {"SIMPLE": "a lookup", "MEDIUM": "some work", "COMPLEX": "deep work"} +_QUESTIONS: Final[dict[str, JsonValue]] = {"tier": {"type": "choice", "criteria": _TIER_CRITERIA}} +_OPENAI_INPUT: Final = "The export job hangs at 99% in billing" +_OPENAI_TIER_QUESTION: Final[dict[str, JsonValue]] = { + "type": "choice", + "name": "tier", + "instructions": "Which tier?", + "choices": [{"value": name, "description": description} for name, description in _TIER_CRITERIA.items()], +} +_ANSWERS: Final[dict[str, JsonValue]] = { + "tier": { + "type": "choice", + "choice": "COMPLEX", + "confidence": 0.91, + "probabilities": {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91}, + } +} +_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 367, "output_tokens": 3} +_ANSWER_BODY: Final[dict[str, JsonValue]] = {"model": _ENDPOINT, "answers": _ANSWERS, "usage": _USAGE} +_BARE_NAME_RULE: Final = ( + "must be the bare endpoint name (letters, digits, '-', '_' and '.', with no '/', '?', '#', spaces, or a " + "leading '.'), e.g. databricks-openjev-qwen35-4b" +) +_MISSING_KEY: Final = "Missing API key for Decisions provider 'databricks'" +_NON_BARE_NAMES: Final[tuple[tuple[str, str], ...]] = ( + ("slash", "foo/bar"), + ("space", "a b"), + ("leading-dot", ".hidden"), + ("question-mark", "a?b"), + ("hash", "a#b"), +) +_DATABRICKS_ERRORS: Final[tuple[tuple[int, dict[str, JsonValue]], ...]] = ( + (400, {"error_code": "BAD_REQUEST", "message": "Invalid input: questions must not be empty"}), + (404, {"error_code": "RESOURCE_DOES_NOT_EXIST", "message": f"Endpoint with name '{_ENDPOINT}' does not exist."}), +) +_SPEND_QUERY: Final = ( + 'SELECT spend, status, call_type, model_group, custom_llm_provider, api_base FROM "LiteLLM_SpendLogs" ' + "WHERE request_id = %s" +) +_OWNED_PROXY_BUDGET: Final = int(2 * graceful_stop_seconds() + 120) +_MEGABYTE_NAME: Final = "a" * 1_000_000 + "/" + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def _register(scenario: Scenario, body: dict[str, JsonValue], *, status: int = 200) -> ScenarioHandle: + handle: Final = register_scenario( + f"databricks-{uuid.uuid4().hex[:12]}", JsonResponse(content_type="application/json", body=body, status=status) + ) + scenario.cleanups.callback(delete_scenario, handle) + return handle + + +def _decide(gateway: Gateway, model: str, *, key: str | None = None) -> httpx.Response: + return gateway.request("POST", "/v1/systemone", {"model": model, "state": _STATE, "questions": _QUESTIONS}, key=key) + + +def _observed(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + return tuple(map(object_value, upstream.get("/__observations").json()["requests"])) + + +def _calls_to(requests: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> list[dict[str, JsonValue]]: + return [request for request in requests if string_value(request["path"]).startswith(f"/{handle.scenario_id}/")] + + +def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, JsonValue]]: + return _calls_to(_observed(gateway), handle) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _wildcard_deployment(scenario: Scenario, **parameters: JsonValue) -> str: + prefix: Final = f"databricks-{uuid.uuid4().hex[:8]}" + created: Final = scenario.gateway.post( + "/model/new", + {"model_name": f"{prefix}/*", "litellm_params": {"model": "databricks/*", **parameters}, "model_info": {}}, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return prefix + + +def _router_config(classifier: Mapping[str, JsonValue], simple: str, complex_: str) -> dict[str, JsonValue]: + return { + "classifier_type": "oss_classifier", + "opensource_classifier_config": dict(classifier), + "tiers": {"SIMPLE": simple, "MEDIUM": simple, "COMPLEX": complex_, "REASONING": complex_}, + } + + +def _environment_classifier() -> dict[str, JsonValue]: + return {"provider": "databricks", "model": _ENDPOINT, "timeout_ms": 20000, "circuit_breaker_enabled": False} + + +def _router(scenario: Scenario, classifier: Mapping[str, JsonValue]) -> str: + simple: Final = scenario.model(model="openai/simple-tier") + complex_: Final = scenario.model(model="openai/complex-tier") + name: Final = f"router-{uuid.uuid4().hex[:10]}" + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": _router_config(classifier, simple, complex_), + }, + "model_info": {}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _chat(gateway: Gateway, model: str, *, key: str | None = None) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"design a distributed cache {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _assert_classified_by_the_complex_tier(response: httpx.Response, router: str) -> None: + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"], response.text + assert ( + response.headers["x-litellm-model-group"], + response.headers["x-litellm-model-name"], + response.headers["x-litellm-complexity-router-tier"], + response.headers["x-litellm-complexity-router-cause"], + response.headers["x-litellm-classifier-cost"], + ) == (router, "openai/complex-tier", "COMPLEX", "jev_classifier", "0.0"), dict(response.headers) + + +def _assert_classifier_call( + call: Mapping[str, JsonValue], *, path: str, authorization: str, prompt_marker: str +) -> None: + assert call["path"] == path, call + assert call["authorization"] == authorization, call + body: Final = object_value(call["body"]) + assert body["model"] == _ENDPOINT, body + assert prompt_marker in string_value(body["state"]), body + assert sorted(object_value(object_value(body["questions"])["tier"])) == ["criteria", "instructions", "type"], body + + +def _assert_refused_before_any_upstream_call(gateway: Gateway, model: str, message: str) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _ANSWER_BODY) + deployment: Final = scenario.model(model=model, api_base=handle.api_base(), api_key=_API_KEY) + response: Final = _decide(gateway, deployment) + assert response.status_code == 400, response.text + assert message in response.text, response.text + assert _upstream_calls(gateway, handle) == [] + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], _number(row["spend"])) == ("failure", "adecisions", 0.0), row + + +@pytest.mark.parametrize("name", [name for _, name in _NON_BARE_NAMES], ids=[label for label, _ in _NON_BARE_NAMES]) +def test_a_deployment_naming_a_non_bare_endpoint_is_refused_before_any_upstream_call( + gateway: Gateway, name: str +) -> None: + _assert_refused_before_any_upstream_call( + gateway, f"databricks/{name}", f"Databricks serving endpoint name {name!r} {_BARE_NAME_RULE}" + ) + + +def test_a_deployment_naming_no_endpoint_is_refused_before_any_upstream_call(gateway: Gateway) -> None: + _assert_refused_before_any_upstream_call(gateway, "databricks/", "A model name is required for the Decisions API") + + +def test_an_openai_format_request_at_v1_decisions_reaches_the_serving_endpoint_as_system_one(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _ANSWER_BODY) + deployment: Final = scenario.model( + model=f"databricks/{_ENDPOINT}", api_base=handle.api_base(), api_key=_API_KEY + ) + response: Final = gateway.request( + "POST", + "/v1/decisions", + { + "model": deployment, + "input": _OPENAI_INPUT, + "questions": [_OPENAI_TIER_QUESTION], + "safety_identifier": "end-user-1", + }, + ) + assert response.status_code == 200, response.text + assert response.json() == { + "model": _ENDPOINT, + "answers": [ + { + "type": "choice", + "name": "tier", + "choice": "COMPLEX", + "probabilities": [ + {"value": "SIMPLE", "probability": 0.03}, + {"value": "MEDIUM", "probability": 0.06}, + {"value": "COMPLEX", "probability": 0.91}, + ], + "confidence": 0.91, + } + ], + "usage": { + "input_tokens": 367, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 3, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 370, + }, + }, response.text + assert response.headers["x-litellm-model-name"] == f"databricks/{_ENDPOINT}", dict(response.headers) + (call,) = _upstream_calls(gateway, handle) + assert call["path"] == f"/{handle.scenario_id}/{_ENDPOINT}/invocations", call + assert call["authorization"] == f"Bearer {_API_KEY}", call + assert call["body"] == { + "model": _ENDPOINT, + "state": _OPENAI_INPUT, + "questions": {"tier": {"type": "choice", "instructions": "Which tier?", "criteria": _TIER_CRITERIA}}, + }, call + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert ( + row["status"], + row["call_type"], + row["custom_llm_provider"], + row["model_group"], + row["api_base"], + _number(row["spend"]), + ) == ("success", "adecisions", "databricks", deployment, f"{handle.api_base()}/{_ENDPOINT}/invocations", 0.0), ( + row + ) + + +def test_a_wildcard_deployment_sends_each_request_to_the_named_endpoint_and_bills_each_once(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _ANSWER_BODY) + prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY) + model: Final = f"{prefix}/{_ENDPOINT}" + responses: Final = tuple(_decide(gateway, model) for _ in range(2)) + for response in responses: + assert response.status_code == 200, response.text + assert response.json() == _ANSWER_BODY, response.text + assert response.headers["x-litellm-model-group"] == model, dict(response.headers) + assert response.headers["x-litellm-model-name"] == f"databricks/{_ENDPOINT}", dict(response.headers) + call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses) + assert len(set(call_ids)) == 2, call_ids + calls: Final = _upstream_calls(gateway, handle) + assert len(calls) == 2, calls + for call in calls: + assert call["path"] == f"/{handle.scenario_id}/{_ENDPOINT}/invocations", call + assert call["authorization"] == f"Bearer {_API_KEY}", call + assert call["body"] == {"model": _ENDPOINT, "state": _STATE, "questions": _QUESTIONS}, call + for call_id in call_ids: + row: Final = _spend_row(call_id) + assert ( + row["status"], + row["call_type"], + row["custom_llm_provider"], + row["model_group"], + row["api_base"], + _number(row["spend"]), + ) == ("success", "adecisions", "databricks", model, f"{handle.api_base()}/{_ENDPOINT}/invocations", 0.0), ( + row + ) + + +@pytest.mark.parametrize("model", (5, ["a"]), ids=("int", "list")) +def test_a_non_string_model_is_refused_at_the_gateway_without_an_upstream_call( + gateway: Gateway, model: JsonValue +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _ANSWER_BODY) + _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY) + response: Final = gateway.request( + "POST", "/v1/systemone", {"model": model, "state": _STATE, "questions": _QUESTIONS} + ) + assert 400 <= response.status_code < 500, response.text + assert _upstream_calls(gateway, handle) == [] + + +def test_a_five_kilobyte_endpoint_name_reaches_the_upstream_and_its_not_found_answer_is_kept(gateway: Gateway) -> None: + name: Final = "n" * 5000 + with gateway.scenario() as scenario: + handle: Final = _register( + scenario, + {"error_code": "RESOURCE_DOES_NOT_EXIST", "message": f"Endpoint with name '{name}' does not exist."}, + status=404, + ) + prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY) + response: Final = _decide(gateway, f"{prefix}/{name}") + assert response.status_code == 404, response.text[:500] + assert "RESOURCE_DOES_NOT_EXIST" in response.text, response.text[:500] + (call,) = _upstream_calls(gateway, handle) + assert call["path"] == f"/{handle.scenario_id}/{name}/invocations", call["path"][:200] + assert call["authorization"] == f"Bearer {_API_KEY}", call["authorization"] + assert call["body"] == {"model": name, "state": _STATE, "questions": _QUESTIONS} + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == ( + "failure", + "adecisions", + f"{prefix}/{name}", + 0.0, + ), row + + +@pytest.mark.parametrize("status,body", _DATABRICKS_ERRORS, ids=("400", "404")) +def test_a_databricks_error_keeps_its_status_and_message_and_bills_nothing( + gateway: Gateway, status: int, body: dict[str, JsonValue] +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, body, status=status) + model: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=handle.api_base(), api_key=_API_KEY) + response: Final = _decide(gateway, model) + assert response.status_code == status, response.text + assert string_value(body["message"]) in response.text, response.text + (call,) = _upstream_calls(gateway, handle) + assert call["path"] == f"/{handle.scenario_id}/{_ENDPOINT}/invocations", call + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == ( + "failure", + "adecisions", + model, + 0.0, + ), row + + +def test_a_deployment_without_a_key_and_no_environment_key_is_refused_before_any_upstream_call( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _ANSWER_BODY) + model: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=handle.api_base(), api_key=None) + response: Final = _decide(gateway, model) + assert response.status_code == 401, response.text + assert _MISSING_KEY in response.text, response.text + assert _upstream_calls(gateway, handle) == [] + + +@pytest.mark.timeout(_OWNED_PROXY_BUDGET) +def test_the_environment_key_and_base_serve_a_deployment_a_classifier_and_a_team_member_that_leave_them_unset( + gateway: Gateway, tmp_path: Path +) -> None: + with gateway.scenario() as rig_scenario: + handle: Final = _register(rig_scenario, _ANSWER_BODY) + environment_base: Final = f"{handle.api_base()}/serving-endpoints" + environment_path: Final = f"/{handle.scenario_id}/serving-endpoints/{_ENDPOINT}/invocations" + environment: Final = { + "DATABRICKS_API_KEY": _ENV_KEY, + "DATABRICKS_TOKEN": _ENV_TOKEN, + "DATABRICKS_API_BASE": environment_base, + } + with owned_proxy_process(gateway, tmp_path, environment) as owned, owned.gateway.scenario() as scenario: + candidate: Final = owned.gateway + with_base: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=handle.api_base(), api_key=None) + bare: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=None, api_key=None) + decided: Final = tuple(_decide(candidate, model) for model in (with_base, bare)) + for response in decided: + assert response.status_code == 200, response.text + assert response.json() == _ANSWER_BODY, response.text + assert [(call["path"], call["authorization"]) for call in _upstream_calls(candidate, handle)] == [ + (f"/{handle.scenario_id}/{_ENDPOINT}/invocations", f"Bearer {_ENV_KEY}"), + (environment_path, f"Bearer {_ENV_KEY}"), + ] + for response, model in zip(decided, (with_base, bare)): + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["model_group"], _number(row["spend"])) == ("success", model, 0.0), row + assert ( + _spend_row(decided[1].headers["x-litellm-call-id"])["api_base"] + == f"{environment_base}/{_ENDPOINT}/invocations" + ) + + router: Final = _router(scenario, _environment_classifier()) + routed: Final = _chat(candidate, router) + _assert_classified_by_the_complex_tier(routed, router) + (judge_call,) = _upstream_calls(candidate, handle) + _assert_classifier_call( + judge_call, + path=environment_path, + authorization=f"Bearer {_ENV_KEY}", + prompt_marker="design a distributed cache", + ) + + team: Final = scenario.team(team_member_permissions=["/auto_router/manage"]) + member: Final = scenario.member(team) + member_key: Final = scenario.key(user_id=member, team_id=team) + simple: Final = scenario.model(model="openai/simple-tier") + complex_: Final = scenario.model(model="openai/complex-tier") + team_router: Final = f"team-router-{uuid.uuid4().hex[:8]}" + created: Final = candidate.request( + "POST", + "/model/new", + { + "model_name": team_router, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": _router_config(_environment_classifier(), simple, complex_), + }, + "model_info": {"team_id": team}, + }, + key=member_key, + ) + assert created.status_code == 200, created.text + team_router_id: Final = string_value(object_value(object_value(created.json())["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, team_router_id) + stored: Final = candidate.request("GET", "/model/info", params={"litellm_model_id": team_router_id}) + assert stored.status_code == 200, stored.text + (entry,) = stored.json()["data"] + stored_classifier: Final = object_value( + object_value(object_value(object_value(entry)["litellm_params"])["complexity_router_config"])[ + "opensource_classifier_config" + ] + ) + assert (stored_classifier["provider"], stored_classifier["model"]) == ("databricks", _ENDPOINT), entry + assert "api_key" not in stored_classifier and "api_base" not in stored_classifier, entry + member_chat: Final = _chat(candidate, team_router, key=member_key) + _assert_classified_by_the_complex_tier(member_chat, team_router) + (member_judge_call,) = _upstream_calls(candidate, handle) + _assert_classifier_call( + member_judge_call, + path=environment_path, + authorization=f"Bearer {_ENV_KEY}", + prompt_marker="design a distributed cache", + ) + + +@pytest.mark.timeout(_OWNED_PROXY_BUDGET) +def test_the_environment_token_serves_a_deployment_and_a_classifier_when_no_environment_key_is_set( + gateway: Gateway, tmp_path: Path +) -> None: + with gateway.scenario() as rig_scenario: + handle: Final = _register(rig_scenario, _ANSWER_BODY) + environment_path: Final = f"/{handle.scenario_id}/serving-endpoints/{_ENDPOINT}/invocations" + environment: Final = { + "DATABRICKS_TOKEN": _ENV_TOKEN, + "DATABRICKS_API_BASE": f"{handle.api_base()}/serving-endpoints", + } + with ( + owned_proxy_process(gateway, tmp_path, environment, remove_environment=("DATABRICKS_API_KEY",)) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + bare: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=None, api_key=None) + decided: Final = _decide(candidate, bare) + assert decided.status_code == 200, decided.text + assert decided.json() == _ANSWER_BODY, decided.text + (decision_call,) = _upstream_calls(candidate, handle) + assert (decision_call["path"], decision_call["authorization"]) == (environment_path, f"Bearer {_ENV_TOKEN}") + assert _spend_row(decided.headers["x-litellm-call-id"])["status"] == "success" + router: Final = _router(scenario, _environment_classifier()) + routed: Final = _chat(candidate, router) + _assert_classified_by_the_complex_tier(routed, router) + (judge_call,) = _upstream_calls(candidate, handle) + _assert_classifier_call( + judge_call, + path=environment_path, + authorization=f"Bearer {_ENV_TOKEN}", + prompt_marker="design a distributed cache", + ) + + +def test_a_megabyte_endpoint_name_is_refused_in_linear_time_while_liveliness_stays_fast(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _ANSWER_BODY) + prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY) + model: Final = f"{prefix}/{_MEGABYTE_NAME}" + anonymous: Final = gateway.client.post( + "/v1/systemone", json={"model": model, "state": _STATE, "questions": _QUESTIONS} + ) + assert anonymous.status_code == 401, anonymous.text[:300] + liveliness: Final[list[tuple[int, float]]] = [] + stop: Final = threading.Event() + + def poll() -> None: + while not stop.is_set(): + started: Final = time.perf_counter() + probe: Final = gateway.client.get("/health/liveliness") + liveliness.append((probe.status_code, time.perf_counter() - started)) + + poller: Final = threading.Thread(target=poll) + poller.start() + started: Final = time.perf_counter() + response: Final = _decide(gateway, model) + elapsed: Final = time.perf_counter() - started + stop.set() + poller.join() + assert elapsed < 3, elapsed + assert liveliness and all(status == 200 for status, _ in liveliness), liveliness + assert max(seconds for _, seconds in liveliness) < 1, liveliness + assert response.status_code == 400, response.text[:300] + assert _BARE_NAME_RULE in response.text, response.text[-400:] + assert _upstream_calls(gateway, handle) == [] + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], _number(row["spend"])) == ("failure", "adecisions", 0.0), row diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 27a9783cf436..65bf5c619dc6 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -111,6 +111,15 @@ class _Provider: _Provider( "azure_ai", "azure_ai/decision-1", "/providers/microsoft/v1/systemone", "decision-1", _API_KEY, False, None ), + _Provider( + "databricks", + "databricks/databricks-openjev-qwen35-4b", + "/databricks-openjev-qwen35-4b/invocations", + "databricks-openjev-qwen35-4b", + _API_KEY, + False, + None, + ), ) _PERPLEXITY: Final = _PROVIDERS[0] _OPENROUTER: Final = _PROVIDERS[2] diff --git a/tests/integration/routing/test_complexity_router_databricks_classifier.py b/tests/integration/routing/test_complexity_router_databricks_classifier.py new file mode 100644 index 000000000000..25611b4ec08a --- /dev/null +++ b/tests/integration/routing/test_complexity_router_databricks_classifier.py @@ -0,0 +1,1023 @@ +import asyncio +import json +import os +import re +import signal +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import anthropic +import httpx +import openai +import psutil +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery, chat_reply, responses_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with +from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS + +_ENDPOINT: Final = "databricks-openjev-qwen35-4b" +_API_KEY: Final = "synthetic-databricks-key" +_CLASSIFIER_MODEL: Final = f"databricks/{_ENDPOINT}" +_TYPESAFE_MODEL: Final = "jev-1.13.0" +_TYPESAFE_KEY: Final = "synthetic-typesafe-key" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_ANTHROPIC_VERSION: Final = MappingProxyType({"anthropic-version": "2023-06-01"}) +_SALT: Final = os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt") +_TIER_CRITERIA: Final[dict[str, JsonValue]] = { + "SIMPLE": "Greetings, chitchat, or short factual lookups with known answers", + "MEDIUM": "Everyday requests needing explanation, light reasoning, or minor technical work", + "COMPLEX": "Non-trivial code, architecture, multi-step work, or specialized domain depth", + "REASONING": "Open-ended analysis, proofs, tradeoffs, or tasks requiring careful thought", +} +_PROBABILITIES: Final[dict[str, JsonValue]] = {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91} +_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 12, "output_tokens": 1} +_ANSWER: Final[dict[str, JsonValue]] = { + "answers": {"tier": {"type": "choice", "choice": "COMPLEX", "confidence": 0.91, "probabilities": _PROBABILITIES}}, + "usage": _USAGE, +} +_ROWS_QUERY: Final = ( + "SELECT request_id, status, model, call_type, cache_hit, custom_llm_provider, spend, " + "metadata->>'internal_call_origin' AS origin " + 'FROM "LiteLLM_SpendLogs" WHERE model_group = %s ORDER BY "startTime"' +) +_SIMPLE_TEXT: Final = "simple answer" +_COMPLEX_TEXT: Final = "complex answer" +_CALL_TYPES: Final = MappingProxyType( + {"/v1/chat/completions": "acompletion", "/v1/responses": "aresponses", "/v1/messages": "anthropic_messages"} +) + + +@dataclass(frozen=True, slots=True) +class _Routed: + name: str + identity: str + simple: str + complex_: str + + +@dataclass(frozen=True, slots=True) +class _Served: + route: str + stream: bool + status: int + identity: str + text: str + cause: str + classifier_cost: str | None + + +@dataclass(frozen=True, slots=True) +class _Failure: + label: str + respond: Callable[[Request], Reply] + overrides: Mapping[str, JsonValue] + reason: str + error_type: str + classifier_rows: int + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def _prompt() -> str: + return f"design a distributed cache {uuid.uuid4().hex}" + + +def _short_prompt() -> str: + return f"hi {uuid.uuid4().hex[:8]}" + + +def _answering(body: Mapping[str, JsonValue]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + return Reply(body=json.dumps(body).encode()) + + return respond + + +def _failing(status: int) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + return Reply(status=status, body=json.dumps({"error_code": "FAILED", "message": f"scripted {status}"}).encode()) + + return respond + + +def _slow(seconds: float) -> Callable[[Request], Reply]: + answer: Final = _answering(_ANSWER) + + def respond(request: Request) -> Reply: + time.sleep(seconds) + return answer(request) + + return respond + + +def _tier(prefix: str, text: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) if request.body else {} + identity: Final = f"{prefix}-{uuid.uuid4().hex[:8]}" + if request.target == "/chat/completions": + return chat_reply(identity, "tier-model", text, stream=bool(body.get("stream"))) + if request.target == "/responses": + return responses_reply(identity, "tier-model", text, stream=bool(body.get("stream"))) + return Reply(status=404, body=json.dumps({"error": {"message": f"no route {request.target}"}}).encode()) + + return answering_model_discovery(respond) + + +@contextmanager +def _wires(classifier: Callable[[Request], Reply] | None = None) -> Iterator[tuple[Wire, Wire, Wire]]: + with ExitStack() as stack: + judge: Final = stack.enter_context(wire_server(classifier or _answering(_ANSWER))) + simple: Final = stack.enter_context(wire_server(_tier("simple", _SIMPLE_TEXT))) + complex_: Final = stack.enter_context(wire_server(_tier("complex", _COMPLEX_TEXT))) + yield judge, simple, complex_ + + +def _classifier_config(judge_url: str, **overrides: JsonValue) -> dict[str, JsonValue]: + return { + "provider": "databricks", + "model": _ENDPOINT, + "api_base": judge_url, + "api_key": _API_KEY, + "timeout_ms": 20000, + "circuit_breaker_enabled": False, + **overrides, + } + + +def _router_config(classifier: Mapping[str, JsonValue], simple: str, complex_: str) -> dict[str, JsonValue]: + return { + "classifier_type": "oss_classifier", + "opensource_classifier_config": dict(classifier), + "tiers": {"SIMPLE": simple, "MEDIUM": simple, "COMPLEX": complex_, "REASONING": complex_}, + } + + +def _deploy( + scenario: Scenario, judge: Wire, simple: Wire, complex_: Wire, classifier: Mapping[str, JsonValue] | None = None +) -> _Routed: + simple_model: Final = scenario.model(model="openai/simple-tier", api_base=simple.url, api_key="synthetic-tier-key") + complex_model: Final = scenario.model( + model="openai/complex-tier", api_base=complex_.url, api_key="synthetic-tier-key" + ) + name: Final = f"router-{uuid.uuid4().hex[:10]}" + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": _router_config( + classifier if classifier is not None else _classifier_config(judge.url), simple_model, complex_model + ), + }, + "model_info": {}, + }, + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return _Routed(name=name, identity=identity, simple=simple_model, complex_=complex_model) + + +def _posts(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if request.method == "POST") + + +def _state(prompt: str, system: str | None) -> str: + if system is None: + return f"\nClassify this message:\n{prompt}" + return f"\nCaller system prompt, quoted as task context:\n{system}\n\nClassify this message:\n{prompt}" + + +def _assert_classifier_call( + call: Request, *, prompt: str, system: str | None = None, endpoint: str = _ENDPOINT, key: str = _API_KEY +) -> None: + assert call.target == f"/{endpoint}/invocations", call.target + assert call.headers.get("authorization") == f"Bearer {key}", call.headers + assert json.loads(call.body) == { + "state": _state(prompt, system), + "model": endpoint, + "questions": {"tier": {"type": "choice", "instructions": DEFAULT_JEV_INSTRUCTIONS, "criteria": _TIER_CRITERIA}}, + }, call.body + + +def _assert_classified(headers: Mapping[str, str], *, router: str, cost: str | None = "0.0") -> None: + assert ( + headers.get("x-litellm-model-group"), + headers.get("x-litellm-model-name"), + headers.get("x-litellm-complexity-router-tier"), + headers.get("x-litellm-complexity-router-cause"), + headers.get("x-litellm-classifier-cost"), + ) == (router, "openai/complex-tier", "COMPLEX", "jev_classifier", cost), dict(headers) + + +def _assert_fell_back(headers: Mapping[str, str], *, router: str) -> None: + assert ( + headers.get("x-litellm-model-group"), + headers.get("x-litellm-model-name"), + headers.get("x-litellm-complexity-router-tier"), + headers.get("x-litellm-complexity-router-cause"), + headers.get("x-litellm-classifier-cost"), + ) == (router, "openai/simple-tier", "SIMPLE", "heuristic_scorer", None), dict(headers) + + +def _rows(router: str, *, want: int) -> tuple[dict[str, JsonValue], ...]: + return tuple(eventually(lambda: read_rows(_ROWS_QUERY, (router,)), lambda found: len(found) >= want, seconds=70)) + + +def _classifier_rows(rows: Sequence[Mapping[str, JsonValue]]) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(row for row in rows if row["origin"] == "autorouter_classifier") + + +def _request_rows(rows: Sequence[Mapping[str, JsonValue]]) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(row for row in rows if row["origin"] != "autorouter_classifier") + + +def _assert_classifier_row( + row: Mapping[str, JsonValue], *, model: str = _CLASSIFIER_MODEL, provider: str = "databricks", spend: float = 0.0 +) -> None: + assert (row["status"], row["call_type"], row["model"], row["custom_llm_provider"], row["cache_hit"]) == ( + "success", + "pass_through_endpoint", + model, + provider, + "False", + ), row + assert _number(row["spend"]) == pytest.approx(spend), row + + +def _issued_response_id(caller_id: str) -> str: + decrypted: Final = decrypt_if_encrypted_with(caller_id.removeprefix("resp_"), _SALT) + assert decrypted is not None, caller_id + return decrypted.split(";")[0].split("response_id:")[-1] + + +def _logged_ids(route: str, caller_id: str) -> frozenset[str]: + if route != "/v1/responses": + return frozenset({caller_id}) + return frozenset({caller_id, _issued_response_id(caller_id)}) + + +def _assert_logged_once(rows: Sequence[Mapping[str, JsonValue]], *, route: str, request_id: str) -> None: + (classifier_row,) = _classifier_rows(rows) + _assert_classifier_row(classifier_row) + (request_row,) = _request_rows(rows) + assert request_row["request_id"] in _logged_ids(route, request_id), (request_row, request_id) + assert (request_row["call_type"], request_row["status"]) == (_CALL_TYPES[route], "success"), request_row + + +def _data_frames(lines: Sequence[str]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data:").strip()) + for line in lines + if line.startswith("data:") and line.removeprefix("data:").strip() != "[DONE]" + ) + + +def _delta_texts(frames: Sequence[Mapping[str, JsonValue]]) -> Iterator[str]: + for frame in frames: + choices: Final = frame["choices"] + if not isinstance(choices, list): + continue + for choice in choices: + content: Final = object_value(object_value(choice)["delta"]).get("content") + if isinstance(content, str): + yield content + + +def _chat_stream_identity_and_text(frames: Sequence[Mapping[str, JsonValue]]) -> tuple[str, str]: + identities: Final = {string_value(frame["id"]) for frame in frames} + assert len(identities) == 1, identities + return identities.pop(), "".join(_delta_texts(frames)) + + +def _responses_stream_identity_and_text(frames: Sequence[Mapping[str, JsonValue]]) -> tuple[str, str]: + created: Final = [frame for frame in frames if frame["type"] == "response.created"] + assert len(created) == 1, frames + identity: Final = string_value(object_value(created[0]["response"])["id"]) + text: Final = "".join( + string_value(frame["delta"]) for frame in frames if frame["type"] == "response.output_text.delta" + ) + return identity, text + + +def _messages_stream_identity_and_text(frames: Sequence[Mapping[str, JsonValue]]) -> tuple[str, str]: + started: Final = [frame for frame in frames if frame["type"] == "message_start"] + assert len(started) == 1, frames + identity: Final = string_value(object_value(started[0]["message"])["id"]) + deltas: Final = [object_value(frame["delta"]) for frame in frames if frame["type"] == "content_block_delta"] + text: Final = "".join(string_value(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + return identity, text + + +def _body(route: str, model: str, prompt: str, *, stream: bool) -> dict[str, JsonValue]: + if route == "/v1/chat/completions": + return {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": stream} + if route == "/v1/responses": + return {"model": model, "input": prompt, "stream": stream} + return {"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": prompt}], "stream": stream} + + +def _identity_and_text(route: str, body: Mapping[str, JsonValue]) -> tuple[str, str]: + if route == "/v1/chat/completions": + (choice,) = body["choices"] if isinstance(body["choices"], list) else () + return string_value(body["id"]), string_value(object_value(object_value(choice)["message"])["content"]) + if route == "/v1/responses": + (item,) = body["output"] if isinstance(body["output"], list) else () + (content,) = object_value(item)["content"] if isinstance(object_value(item)["content"], list) else () + return string_value(body["id"]), string_value(object_value(content)["text"]) + (block,) = body["content"] if isinstance(body["content"], list) else () + return string_value(body["id"]), string_value(object_value(block)["text"]) + + +def _stream_identity_and_text(route: str, lines: Sequence[str]) -> tuple[str, str]: + frames: Final = _data_frames(lines) + if route == "/v1/chat/completions": + return _chat_stream_identity_and_text(frames) + if route == "/v1/responses": + return _responses_stream_identity_and_text(frames) + return _messages_stream_identity_and_text(frames) + + +async def _send(client: httpx.AsyncClient, route: str, model: str, prompt: str, *, stream: bool) -> _Served: + body: Final = _body(route, model, prompt, stream=stream) + headers: Final = dict(_ANTHROPIC_VERSION) if route == "/v1/messages" else {} + if not stream: + response: Final = await client.post(route, json=body, headers=headers) + assert response.status_code == 200, response.text + identity, text = _identity_and_text(route, _JSON_OBJECT.validate_json(response.content)) + return _Served( + route=route, + stream=False, + status=response.status_code, + identity=identity, + text=text, + cause=response.headers.get("x-litellm-complexity-router-cause", ""), + classifier_cost=response.headers.get("x-litellm-classifier-cost"), + ) + async with client.stream("POST", route, json=body, headers=headers) as streamed: + lines: Final = [line async for line in streamed.aiter_lines()] + assert streamed.status_code == 200, lines + identity, text = _stream_identity_and_text(route, lines) + return _Served( + route=route, + stream=True, + status=streamed.status_code, + identity=identity, + text=text, + cause=streamed.headers.get("x-litellm-complexity-router-cause", ""), + classifier_cost=streamed.headers.get("x-litellm-classifier-cost"), + ) + + +def _mixed_calls(count: int) -> tuple[tuple[str, bool], ...]: + routes: Final = tuple(_CALL_TYPES) + return tuple((routes[index % len(routes)], index % 2 == 1) for index in range(count)) + + +async def _burst( + base_url: str, + key: str, + model: str, + calls: Sequence[tuple[str, bool]], + prompts: Sequence[str], + *, + tolerate_transport_errors: bool = False, +) -> tuple[_Served, ...]: + async with httpx.AsyncClient( + base_url=base_url, timeout=90, trust_env=False, headers={"Authorization": f"Bearer {key}"} + ) as client: + results: Final = await asyncio.gather( + *(_send(client, route, model, prompt, stream=stream) for (route, stream), prompt in zip(calls, prompts)), + return_exceptions=tolerate_transport_errors, + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _served_by_logged_id(served: Sequence[_Served]) -> Iterator[tuple[str, _Served]]: + for item in served: + for identity in _logged_ids(item.route, item.identity): + yield identity, item + + +def _assert_each_logged_once(rows: Sequence[Mapping[str, JsonValue]], served: Sequence[_Served]) -> None: + request_rows: Final = _request_rows(rows) + accepted: Final = dict(_served_by_logged_id(served)) + assert all(string_value(row["request_id"]) in accepted for row in request_rows), (request_rows, accepted) + matched: Final = tuple(accepted[string_value(row["request_id"])] for row in request_rows) + assert sorted(item.identity for item in matched) == sorted(item.identity for item in served), (matched, served) + for row, item in zip(request_rows, matched, strict=True): + assert (row["call_type"], row["status"]) == (_CALL_TYPES[item.route], "success"), (item, row) + + +def _test_routing(gateway: Gateway, config: Mapping[str, JsonValue], prompt: str) -> dict[str, JsonValue]: + response: Final = gateway.request( + "POST", "/auto_router/test_routing", {"prompt": prompt, "complexity_router_config": dict(config)} + ) + assert response.status_code == 200, response.text + return object_value(object_value(response.json())["routing_decision"]) + + +def _sdk_base(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def test_chat_completion_over_httpx_is_classified_by_the_databricks_endpoint_with_the_system_prompt_quoted( + gateway: Gateway, +) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": routed.name, + "messages": [{"role": "system", "content": "be terse"}, {"role": "user", "content": prompt}], + }, + ) + assert response.status_code == 200, response.text + identity, text = _identity_and_text("/v1/chat/completions", _JSON_OBJECT.validate_json(response.content)) + assert text == _COMPLEX_TEXT, response.text + _assert_classified(response.headers, router=routed.name) + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt, system="be terse") + (tier_call,) = _posts(complex_) + assert tier_call.target == "/chat/completions", tier_call.target + tier_body: Final = _JSON_OBJECT.validate_json(tier_call.body) + assert tier_body["model"] == "complex-tier", tier_body + assert tier_body["messages"] == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": prompt}, + ], tier_body + assert _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/chat/completions", request_id=identity) + + +def test_chat_completion_stream_over_httpx_is_classified_and_logged_under_the_chunk_id(gateway: Gateway) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json=_body("/v1/chat/completions", routed.name, prompt, stream=True), + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as streamed: + lines: Final = list(streamed.iter_lines()) + assert streamed.status_code == 200, lines + _assert_classified(streamed.headers, router=routed.name) + identity, text = _stream_identity_and_text("/v1/chat/completions", lines) + assert text == _COMPLEX_TEXT, lines + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt) + (tier_call,) = _posts(complex_) + assert _JSON_OBJECT.validate_json(tier_call.body)["stream"] is True, tier_call.body + assert _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/chat/completions", request_id=identity) + + +def test_chat_completion_through_the_openai_sdk_is_classified_and_logged_once(gateway: Gateway) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + client: Final = openai.OpenAI(base_url=f"{_sdk_base(gateway)}/v1", api_key=gateway.key, max_retries=0) + raw: Final = client.chat.completions.with_raw_response.create( + model=routed.name, messages=[{"role": "user", "content": prompt}] + ) + completion: Final = raw.parse() + assert completion.choices[0].message.content == _COMPLEX_TEXT, completion + _assert_classified(raw.headers, router=routed.name) + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt) + assert len(_posts(complex_)) == 1 and _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/chat/completions", request_id=completion.id) + + +async def test_chat_completion_stream_through_the_async_openai_sdk_is_classified_and_logged_once( + gateway: Gateway, +) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + client: Final = openai.AsyncOpenAI(base_url=f"{_sdk_base(gateway)}/v1", api_key=gateway.key, max_retries=0) + raw: Final = await client.chat.completions.with_raw_response.create( + model=routed.name, messages=[{"role": "user", "content": prompt}], stream=True + ) + chunks: Final = [chunk async for chunk in raw.parse()] + identities: Final = {chunk.id for chunk in chunks} + assert len(identities) == 1, identities + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == _COMPLEX_TEXT + _assert_classified(raw.headers, router=routed.name) + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt) + assert len(_posts(complex_)) == 1 and _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/chat/completions", request_id=identities.pop()) + + +def test_responses_request_over_httpx_is_classified_and_forwarded_to_the_tier_responses_route(gateway: Gateway) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + response: Final = gateway.request("POST", "/v1/responses", {"model": routed.name, "input": prompt}) + assert response.status_code == 200, response.text + identity, text = _identity_and_text("/v1/responses", _JSON_OBJECT.validate_json(response.content)) + assert text == _COMPLEX_TEXT, response.text + _assert_classified(response.headers, router=routed.name) + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt) + (tier_call,) = _posts(complex_) + assert tier_call.target == "/responses", tier_call.target + tier_body: Final = _JSON_OBJECT.validate_json(tier_call.body) + assert (tier_body["model"], tier_body["input"]) == ("complex-tier", prompt), tier_body + assert _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/responses", request_id=identity) + + +async def test_responses_stream_through_the_async_openai_sdk_is_classified_and_logged_under_the_created_id( + gateway: Gateway, +) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + client: Final = openai.AsyncOpenAI(base_url=f"{_sdk_base(gateway)}/v1", api_key=gateway.key, max_retries=0) + raw: Final = await client.responses.with_raw_response.create(model=routed.name, input=prompt, stream=True) + events: Final = [event async for event in raw.parse()] + created: Final = [event for event in events if event.type == "response.created"] + assert len(created) == 1, [event.type for event in events] + text: Final = "".join(event.delta for event in events if event.type == "response.output_text.delta") + assert text == _COMPLEX_TEXT, [event.type for event in events] + _assert_classified(raw.headers, router=routed.name) + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt) + (tier_call,) = _posts(complex_) + assert tier_call.target == "/responses", tier_call.target + assert _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/responses", request_id=created[0].response.id) + + +def test_messages_request_through_the_anthropic_sdk_is_classified_and_bridged_to_the_tier_responses_route( + gateway: Gateway, +) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + client: Final = anthropic.Anthropic(base_url=_sdk_base(gateway), api_key=gateway.key, max_retries=0) + raw: Final = client.messages.with_raw_response.create( + model=routed.name, max_tokens=64, messages=[{"role": "user", "content": prompt}] + ) + message: Final = raw.parse() + assert [block.text for block in message.content if block.type == "text"] == [_COMPLEX_TEXT], message + _assert_classified(raw.headers, router=routed.name) + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt) + (tier_call,) = _posts(complex_) + assert tier_call.target == "/responses", tier_call.target + assert prompt in tier_call.body.decode(), tier_call.body + assert _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/messages", request_id=message.id) + + +def test_messages_stream_over_httpx_is_classified_and_logged_under_the_message_start_id(gateway: Gateway) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + with gateway.client.stream( + "POST", + "/v1/messages", + json=_body("/v1/messages", routed.name, prompt, stream=True), + headers={"Authorization": f"Bearer {gateway.key}", **_ANTHROPIC_VERSION}, + ) as streamed: + lines: Final = list(streamed.iter_lines()) + assert streamed.status_code == 200, lines + _assert_classified(streamed.headers, router=routed.name) + identity, text = _stream_identity_and_text("/v1/messages", lines) + assert text == _COMPLEX_TEXT, lines + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt) + (tier_call,) = _posts(complex_) + assert tier_call.target == "/responses", tier_call.target + assert _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/messages", request_id=identity) + + +def test_a_non_catalog_endpoint_name_is_classified_without_a_cost_header_and_logged_at_zero(gateway: Gateway) -> None: + prompt: Final = _prompt() + endpoint: Final = "my-jev-endpoint" + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_, _classifier_config(judge.url, model=endpoint)) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _body("/v1/chat/completions", routed.name, prompt, stream=False) + ) + assert response.status_code == 200, response.text + identity, text = _identity_and_text("/v1/chat/completions", _JSON_OBJECT.validate_json(response.content)) + assert text == _COMPLEX_TEXT, response.text + _assert_classified(response.headers, router=routed.name, cost=None) + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt, endpoint=endpoint) + assert len(_posts(complex_)) == 1 and _posts(simple) == () + rows: Final = _rows(routed.name, want=2) + (classifier_row,) = _classifier_rows(rows) + _assert_classifier_row(classifier_row, model=f"databricks/{endpoint}") + (request_row,) = _request_rows(rows) + assert (request_row["request_id"], request_row["status"]) == (identity, "success"), request_row + + +def test_a_response_cache_hit_still_runs_the_classifier_and_logs_the_hit_row_at_zero(gateway: Gateway) -> None: + prompt: Final = _prompt() + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + body: Final = _body("/v1/chat/completions", routed.name, prompt, stream=False) + first: Final = gateway.request("POST", "/v1/chat/completions", body) + second: Final = gateway.request("POST", "/v1/chat/completions", body) + assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) + first_identity, first_text = _identity_and_text( + "/v1/chat/completions", _JSON_OBJECT.validate_json(first.content) + ) + second_identity, second_text = _identity_and_text( + "/v1/chat/completions", _JSON_OBJECT.validate_json(second.content) + ) + assert (first_text, second_text) == (_COMPLEX_TEXT, _COMPLEX_TEXT) + assert second_identity == first_identity, (first_identity, second_identity) + assert "x-litellm-cache-key" not in first.headers, dict(first.headers) + assert "x-litellm-cache-key" in second.headers, dict(second.headers) + _assert_classified(first.headers, router=routed.name) + _assert_classified(second.headers, router=routed.name) + classifier_calls: Final = _posts(judge) + assert len(classifier_calls) == 2, classifier_calls + for call in classifier_calls: + _assert_classifier_call(call, prompt=prompt) + assert len(_posts(complex_)) == 1 and _posts(simple) == () + rows: Final = _rows(routed.name, want=4) + classifier_rows: Final = _classifier_rows(rows) + assert len(classifier_rows) == 2, rows + for row in classifier_rows: + _assert_classifier_row(row) + request_rows: Final = _request_rows(rows) + assert [row["request_id"] for row in request_rows if row["cache_hit"] != "True"] == [first_identity], rows + (hit_row,) = [row for row in request_rows if row["cache_hit"] == "True"] + assert string_value(hit_row["request_id"]).startswith(f"{first_identity}_cache_hit"), hit_row + assert (hit_row["status"], _number(hit_row["spend"])) == ("success", 0.0), hit_row + + +_FAILURES: Final = ( + _Failure("500", _failing(500), {}, "classifier_error", "MaskedHTTPStatusError", 0), + _Failure("401", _failing(401), {}, "classifier_error", "MaskedHTTPStatusError", 0), + _Failure( + "malformed", + _answering({"answers": {"tier": {"type": "choice", "choice": "COMPLEX", "confidence": 0.91}}, "usage": _USAGE}), + {}, + "invalid_response", + "ValidationError", + 1, + ), + _Failure( + "unknown-tier", + _answering( + { + "answers": { + "tier": { + "type": "choice", + "choice": "GALAXY", + "confidence": 0.91, + "probabilities": {"GALAXY": 0.91}, + } + }, + "usage": _USAGE, + } + ), + {}, + "classifier_error", + "ValueError", + 1, + ), + _Failure("timeout", _slow(3), {"timeout_ms": 1000}, "timeout", "TimeoutError", 0), +) + + +@pytest.mark.parametrize("failure", _FAILURES, ids=[failure.label for failure in _FAILURES]) +def test_a_failing_databricks_classifier_falls_back_to_the_heuristic_and_names_the_failure( + gateway: Gateway, failure: _Failure +) -> None: + prompt: Final = _short_prompt() + with _wires(failure.respond) as (judge, simple, complex_), gateway.scenario() as scenario: + classifier: Final = _classifier_config(judge.url, **failure.overrides) + routed: Final = _deploy(scenario, judge, simple, complex_, classifier) + decision: Final = _test_routing(gateway, _router_config(classifier, routed.simple, routed.complex_), prompt) + assert ( + decision["cause"], + decision["tier"], + decision["routed_model"], + decision["classifier_failure_reason"], + decision["classifier_error_type"], + ) == ("heuristic_scorer", "SIMPLE", routed.simple, failure.reason, failure.error_type), decision + assert "classifier_cost" not in decision, decision + response: Final = gateway.request( + "POST", "/v1/chat/completions", _body("/v1/chat/completions", routed.name, prompt, stream=False) + ) + assert response.status_code == 200, response.text + identity, text = _identity_and_text("/v1/chat/completions", _JSON_OBJECT.validate_json(response.content)) + assert text == _SIMPLE_TEXT, response.text + _assert_fell_back(response.headers, router=routed.name) + classifier_calls: Final = _posts(judge) + assert len(classifier_calls) == 2, classifier_calls + for call in classifier_calls: + _assert_classifier_call(call, prompt=prompt) + assert len(_posts(simple)) == 1 and _posts(complex_) == () + rows: Final = _rows(routed.name, want=1 + failure.classifier_rows) + (request_row,) = _request_rows(rows) + assert (request_row["request_id"], request_row["call_type"], request_row["status"]) == ( + identity, + "acompletion", + "success", + ), request_row + classifier_rows: Final = _classifier_rows(rows) + assert len(classifier_rows) == failure.classifier_rows, rows + for row in classifier_rows: + _assert_classifier_row(row) + + +def test_a_five_kilobyte_prompt_with_a_system_message_reaches_the_classifier_whole(gateway: Gateway) -> None: + prompt: Final = "p" * 5000 + uuid.uuid4().hex + with _wires() as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": routed.name, + "messages": [{"role": "system", "content": "answer in one line"}, {"role": "user", "content": prompt}], + }, + ) + assert response.status_code == 200, response.text[:400] + identity, text = _identity_and_text("/v1/chat/completions", _JSON_OBJECT.validate_json(response.content)) + assert text == _COMPLEX_TEXT, response.text[:400] + _assert_classified(response.headers, router=routed.name) + (classifier_call,) = _posts(judge) + _assert_classifier_call(classifier_call, prompt=prompt, system="answer in one line") + assert len(_posts(complex_)) == 1 and _posts(simple) == () + _assert_logged_once(_rows(routed.name, want=2), route="/v1/chat/completions", request_id=identity) + + +def _typesafe_cost() -> float: + prices: Final = object_value( + json.loads(Path("model_prices_and_context_window.json").read_text())[f"typesafe/{_TYPESAFE_MODEL}"] + ) + return _number(_USAGE["input_tokens"]) * _number(prices["input_cost_per_token"]) + _number( + _USAGE["output_tokens"] + ) * _number(prices["output_cost_per_token"]) + + +def test_the_default_jev_provider_still_posts_to_systemone_and_bills_from_the_cost_map(gateway: Gateway) -> None: + prompt: Final = _prompt() + expected_cost: Final = _typesafe_cost() + with ( + _wires(_answering({**_ANSWER, "model": _TYPESAFE_MODEL})) as (judge, simple, complex_), + gateway.scenario() as scenario, + ): + classifier: Final[dict[str, JsonValue]] = { + "model": _TYPESAFE_MODEL, + "api_base": judge.url, + "api_key": _TYPESAFE_KEY, + "timeout_ms": 20000, + "circuit_breaker_enabled": False, + } + routed: Final = _deploy(scenario, judge, simple, complex_, classifier) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _body("/v1/chat/completions", routed.name, prompt, stream=False) + ) + assert response.status_code == 200, response.text + identity, text = _identity_and_text("/v1/chat/completions", _JSON_OBJECT.validate_json(response.content)) + assert text == _COMPLEX_TEXT, response.text + assert ( + response.headers["x-litellm-model-name"], + response.headers["x-litellm-complexity-router-tier"], + response.headers["x-litellm-complexity-router-cause"], + ) == ("openai/complex-tier", "COMPLEX", "jev_classifier"), dict(response.headers) + assert float(response.headers["x-litellm-classifier-cost"]) == pytest.approx(expected_cost) + (classifier_call,) = _posts(judge) + assert classifier_call.target == "/v1/systemone", classifier_call.target + assert classifier_call.headers.get("authorization") == f"Bearer {_TYPESAFE_KEY}", classifier_call.headers + assert prompt in classifier_call.body.decode(), classifier_call.body + assert len(_posts(complex_)) == 1 and _posts(simple) == () + rows: Final = _rows(routed.name, want=2) + (classifier_row,) = _classifier_rows(rows) + _assert_classifier_row( + classifier_row, model=f"typesafe/{_TYPESAFE_MODEL}", provider="typesafe", spend=expected_cost + ) + (request_row,) = _request_rows(rows) + assert (request_row["request_id"], request_row["status"]) == (identity, "success"), request_row + + +async def test_a_classifier_outage_mid_traffic_falls_back_and_recovers_with_every_request_logged_once( + gateway: Gateway, +) -> None: + base_url: Final = _sdk_base(gateway) + calls: Final = _mixed_calls(12) + with ( + wire_server(_tier("simple", _SIMPLE_TEXT)) as simple, + wire_server(_tier("complex", _COMPLEX_TEXT)) as complex_, + gateway.scenario() as scenario, + ): + with wire_server(_answering(_ANSWER)) as judge: + routed: Final = _deploy(scenario, judge, simple, complex_) + config: Final = _router_config(_classifier_config(judge.url), routed.simple, routed.complex_) + port: Final = urlsplit(judge.url).port + assert port is not None + before_prompts: Final = tuple(_prompt() for _ in calls) + before: Final = await _burst(base_url, gateway.key, routed.name, calls, before_prompts) + assert [item.cause for item in before] == ["jev_classifier"] * 12, before + assert [item.classifier_cost for item in before] == ["0.0"] * 12, before + assert [item.text for item in before] == [_COMPLEX_TEXT] * 12, before + assert sorted(json.loads(call.body)["state"] for call in _posts(judge)) == sorted( + _state(prompt, None) for prompt in before_prompts + ) + during_prompts: Final = tuple(_short_prompt() for _ in calls) + during: Final = await _burst(base_url, gateway.key, routed.name, calls, during_prompts) + assert [item.cause for item in during] == ["heuristic_scorer"] * 12, during + assert [item.classifier_cost for item in during] == [None] * 12, during + assert [item.text for item in during] == [_SIMPLE_TEXT] * 12, during + decision: Final = _test_routing(gateway, config, _short_prompt()) + assert (decision["cause"], decision["classifier_failure_reason"]) == ( + "heuristic_scorer", + "classifier_error", + ), decision + with wire_server(_answering(_ANSWER), port=port) as recovered: + after_prompts: Final = tuple(_prompt() for _ in calls) + after: Final = await _burst(base_url, gateway.key, routed.name, calls, after_prompts) + assert [item.cause for item in after] == ["jev_classifier"] * 12, after + assert [item.text for item in after] == [_COMPLEX_TEXT] * 12, after + assert sorted(json.loads(call.body)["state"] for call in _posts(recovered)) == sorted( + _state(prompt, None) for prompt in after_prompts + ) + served: Final = (*before, *during, *after) + assert len({item.identity for item in served}) == 36, served + rows: Final = _rows(routed.name, want=60) + _assert_each_logged_once(rows, served) + classifier_rows: Final = _classifier_rows(rows) + assert len(classifier_rows) == 24, rows + for row in classifier_rows: + _assert_classifier_row(row) + + +async def test_a_slow_classifier_under_a_burst_times_out_every_call_into_the_heuristic_without_classifier_rows( + gateway: Gateway, +) -> None: + base_url: Final = _sdk_base(gateway) + calls: Final = tuple(("/v1/chat/completions", False) for _ in range(20)) + prompts: Final = tuple(_short_prompt() for _ in calls) + with _wires(_slow(3)) as (judge, simple, complex_), gateway.scenario() as scenario: + routed: Final = _deploy(scenario, judge, simple, complex_, _classifier_config(judge.url, timeout_ms=500)) + served: Final = await _burst(base_url, gateway.key, routed.name, calls, prompts) + assert [item.cause for item in served] == ["heuristic_scorer"] * 20, served + assert [item.text for item in served] == [_SIMPLE_TEXT] * 20, served + assert len({item.identity for item in served}) == 20, served + seen: Final[list[Request]] = [] + + def drained() -> int: + seen.extend(_posts(judge)) + return len(seen) + + eventually(drained, lambda count: count == 20, seconds=10) + assert sorted(json.loads(call.body)["state"] for call in seen) == sorted( + _state(prompt, None) for prompt in prompts + ) + assert len(_posts(simple)) == 20 and _posts(complex_) == () + rows: Final = _rows(routed.name, want=20) + _assert_each_logged_once(rows, served) + assert _classifier_rows(rows) == (), rows + + +def _worker_startups(log: Path) -> tuple[tuple[int, ...], int]: + text: Final = log.read_text() + return tuple(int(pid) for pid in _STARTED_WORKER.findall(text)), text.count("Application startup complete.") + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _chaos_config(judge: Wire, simple: Wire, complex_: Wire, tmp_path: Path, router: str) -> Path: + config: Final = { + **_JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())), + "model_list": [ + { + "model_name": "chaos-simple", + "litellm_params": { + "model": "openai/simple-tier", + "api_base": simple.url, + "api_key": "synthetic-tier-key", + }, + }, + { + "model_name": "chaos-complex", + "litellm_params": { + "model": "openai/complex-tier", + "api_base": complex_.url, + "api_key": "synthetic-tier-key", + }, + }, + { + "model_name": router, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": _router_config( + _classifier_config(judge.url, timeout_ms=60000), "chaos-simple", "chaos-complex" + ), + }, + }, + ], + } + path: Final = tmp_path / "databricks-classifier-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + 180)) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_classifying_through_the_databricks_endpoint( + gateway: Gateway, tmp_path: Path +) -> None: + router: Final = f"chaos-router-{uuid.uuid4().hex[:8]}" + calls: Final = tuple(("/v1/chat/completions", False) for _ in range(20)) + prompts: Final = tuple(_prompt() for _ in calls) + release: Final = threading.Event() + held_states: Final[SimpleQueue[str]] = SimpleQueue() + answer: Final = _answering(_ANSWER) + + def held(request: Request) -> Reply: + held_states.put(string_value(_JSON_OBJECT.validate_json(request.body)["state"])) + assert release.wait(timeout=120), "The burst was never released" + return answer(request) + + with ( + wire_server(held) as judge, + wire_server(_tier("simple", _SIMPLE_TEXT)) as simple, + wire_server(_tier("complex", _COMPLEX_TEXT)) as complex_, + ): + path: Final = _chaos_config(judge, simple, complex_, tmp_path, router) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + base_url: Final = str(candidate.client.base_url) + workers, _ = eventually( + lambda: _worker_startups(owned.log), lambda found: len(found[0]) == 2 and found[1] == 2, seconds=120 + ) + burst: Final = asyncio.create_task( + _burst(base_url, candidate.key, router, calls, prompts, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, held_states.qsize, lambda size: size == 20, 90) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, judge.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + victim: Final = psutil.Process(victim_pid) + victim.suspend() + victim.send_signal(signal.SIGKILL) + release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + assert [item.cause for item in served] == ["jev_classifier"] * len(served), served + follow_up_prompt: Final = _prompt() + (answered,) = await _burst( + base_url, candidate.key, router, (("/v1/chat/completions", False),), (follow_up_prompt,) + ) + assert (answered.cause, answered.text) == ("jev_classifier", _COMPLEX_TEXT), answered + await asyncio.to_thread( + eventually, lambda: _worker_startups(owned.log), lambda found: len(found[0]) == 3 and found[1] == 3, 180 + ) + everything: Final = (*served, answered) + assert len({item.identity for item in everything}) == len(everything), everything + rows: Final = _rows(router, want=2 * len(everything)) + _assert_each_logged_once(rows, everything) + classifier_rows: Final = _classifier_rows(rows) + assert len(classifier_rows) == len(everything), rows + for row in classifier_rows: + _assert_classifier_row(row) + states: Final[set[str]] = set() + while not held_states.empty(): + states.add(held_states.get()) + assert states == {_state(prompt, None) for prompt in (*prompts, follow_up_prompt)}, states diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index ec5375bc71eb..8f05a492119f 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -22,6 +22,8 @@ OpenAIDecisionResponse, ScoreAnswer, ) +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager _QUESTIONS: Final[Mapping[str, object]] = MappingProxyType( { @@ -68,6 +70,20 @@ "usage": {"input_tokens": 216, "output_tokens": 3}, "latency_ms": 3722.17, } +_DATABRICKS_ENDPOINT: Final = "databricks-openjev-qwen35-4b" +_DATABRICKS_RESPONSE: Final[Mapping[str, object]] = { + "model": "/mosaicml/local_model", + "answers": { + "is_defect": {"type": "noul", "noul": 0.9}, + "sentiment": { + "type": "choice", + "choice": "positive", + "confidence": 0.8, + "probabilities": {"positive": 0.8, "negative": 0.2}, + }, + }, + "usage": {"input_tokens": 214, "output_tokens": 0}, +} _PROVIDERS: Final[tuple[tuple[str, str, str, str], ...]] = ( ( "perplexity", @@ -867,6 +883,164 @@ async def test_hosted_vllm_choice_question_goes_out_as_a_jev_choice_with_the_ser } +@pytest.mark.asyncio +async def test_databricks_requires_api_base_before_http( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("DATABRICKS_API_BASE", raising=False) + monkeypatch.setenv("DATABRICKS_API_KEY", "dapi-key") + + with pytest.raises(litellm.BadRequestError, match="DATABRICKS_API_BASE"): + await litellm.adecisions( + model=f"databricks/{_DATABRICKS_ENDPOINT}", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + ) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "endpoint_name", ["serving-endpoints/openjev", "openjev?x=1", "openjev#frag", "..", "open jev"] +) +async def test_databricks_rejects_a_serving_endpoint_name_that_would_rewrite_the_url( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + endpoint_name: str, +) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.example/serving-endpoints") + monkeypatch.setenv("DATABRICKS_API_KEY", "dapi-key") + + with pytest.raises(litellm.BadRequestError, match="bare endpoint name"): + await litellm.adecisions( + model=f"databricks/{endpoint_name}", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + ) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +async def test_databricks_requires_a_key_before_http( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.example/serving-endpoints") + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.delenv("DATABRICKS_TOKEN", raising=False) + + with pytest.raises(litellm.AuthenticationError, match="Missing API key for Decisions provider 'databricks'"): + await litellm.adecisions( + model=f"databricks/{_DATABRICKS_ENDPOINT}", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + ) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "api_base", ["https://workspace.example/serving-endpoints", "https://workspace.example/serving-endpoints/"] +) +async def test_databricks_posts_the_jev_body_to_the_serving_endpoint_invocations_route( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + api_base: str, +) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", api_base) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.setenv("DATABRICKS_TOKEN", "dapi-token") + route: Final = respx_mock.post( + f"https://workspace.example/serving-endpoints/{_DATABRICKS_ENDPOINT}/invocations" + ).respond(json=_DATABRICKS_RESPONSE) + questions: Final = { + "is_defect": {"type": "noul", "instructions": "Is this a defect?"}, + "sentiment": {"type": "choice", "criteria": {"positive": None, "negative": "unhappy"}}, + } + + response: Final = await litellm.adecisions( + model=f"databricks/{_DATABRICKS_ENDPOINT}", + state={"ticket": "export hangs"}, + questions=questions, + ) + + assert route.called + sent: Final = respx_mock.calls[0].request + assert sent.headers["authorization"] == "Bearer dapi-token" + assert json.loads(sent.content) == { + "model": _DATABRICKS_ENDPOINT, + "state": {"ticket": "export hangs"}, + "questions": questions, + } + assert response.model == "/mosaicml/local_model" + assert isinstance(response.answers["is_defect"], NoulAnswer) + assert isinstance(response.answers["sentiment"], ChoiceAnswer) + assert response.hidden_params["custom_llm_provider"] == "databricks" + assert response.hidden_params["model"] == f"databricks/{_DATABRICKS_ENDPOINT}" + + +@pytest.mark.asyncio +async def test_databricks_router_deployment_sends_its_own_key_to_its_own_base( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", "https://other-workspace.example/serving-endpoints") + monkeypatch.setenv("DATABRICKS_API_KEY", "env-key-stays-home") + router: Final = litellm.Router( + model_list=[ + { + "model_name": "decider", + "litellm_params": { + "model": f"databricks/{_DATABRICKS_ENDPOINT}", + "api_base": "https://workspace.example/serving-endpoints", + "api_key": "deployment-key", + }, + } + ] + ) + route: Final = respx_mock.post( + f"https://workspace.example/serving-endpoints/{_DATABRICKS_ENDPOINT}/invocations" + ).respond(json=_DATABRICKS_RESPONSE) + + response: Final = await router.adecisions( + model="decider", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + ) + + assert route.called + assert respx_mock.calls[0].request.headers["authorization"] == "Bearer deployment-key" + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + + +_PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset( + {LlmProviders.OPENAI.value, LlmProviders.OPENROUTER.value, LlmProviders.HOSTED_VLLM.value} +) +_DEDICATED_DECIDER_PROVIDERS: Final[tuple[str, ...]] = tuple( + sorted( + provider.value + for provider in LlmProviders + if provider.value not in _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE + and ProviderConfigManager.get_provider_decisions_config(model="", provider=provider) is not None + ) +) + + +@pytest.mark.parametrize("provider", _DEDICATED_DECIDER_PROVIDERS) +def test_every_dedicated_decider_provider_ships_an_evaluation_mode_cost_map_entry(provider: str) -> None: + evaluation_entries: Final = tuple( + name + for name, info in litellm.model_cost.items() + if isinstance(info, Mapping) and info.get("litellm_provider") == provider and info.get("mode") == "evaluation" + ) + + assert evaluation_entries, f"a {provider} decider would be health-checked as chat without a mode: evaluation entry" + + @pytest.mark.asyncio @pytest.mark.parametrize( ("api_base", "url"), diff --git a/tests/unit/llms/databricks/decisions/__init__.py b/tests/unit/llms/databricks/decisions/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/unit/llms/databricks/decisions/test_transformation.py b/tests/unit/llms/databricks/decisions/test_transformation.py new file mode 100644 index 000000000000..e6de1d1a635a --- /dev/null +++ b/tests/unit/llms/databricks/decisions/test_transformation.py @@ -0,0 +1,56 @@ +from typing import Final + +import pytest + +from litellm.llms.databricks.decisions.transformation import DatabricksDecisionsConfig + +_BASE: Final = "https://workspace.example/serving-endpoints" +_CONFIG: Final = DatabricksDecisionsConfig() + + +@pytest.mark.parametrize("api_base", [_BASE, f"{_BASE}/"]) +def test_the_complete_url_is_the_serving_endpoint_invocations_route(api_base: str) -> None: + assert _CONFIG.get_complete_url(api_base, "openjev") == f"{_BASE}/openjev/invocations" + + +@pytest.mark.parametrize( + "name", + ["", "serving-endpoints/openjev", "openjev?x=1", "openjev#frag", ".", "..", ".openjev", "open jev", "openjev%2Fx"], +) +def test_a_name_that_would_rewrite_the_url_is_rejected_before_any_request(name: str) -> None: + with pytest.raises(ValueError, match="bare endpoint name"): + _ = _CONFIG.get_complete_url(_BASE, name) + with pytest.raises(ValueError, match="bare endpoint name"): + _ = _CONFIG.canonical_model(name) + + +def test_connection_prefers_the_configured_values_over_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", "https://env.example/serving-endpoints") + monkeypatch.setenv("DATABRICKS_API_KEY", "dapi-env") + connection: Final = _CONFIG.connection(f"{_BASE}/", "dapi-configured") + assert (connection.api_base, connection.api_key) == (_BASE, "dapi-configured") + + +def test_connection_falls_back_to_the_workspace_token_when_the_api_key_variable_is_unset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", _BASE) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.setenv("DATABRICKS_TOKEN", "dapi-token") + connection: Final = _CONFIG.connection(None, None) + assert (connection.api_base, connection.api_key) == (_BASE, "dapi-token") + + +def test_connection_never_sends_the_environment_key_to_a_configured_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", _BASE) + monkeypatch.setenv("DATABRICKS_API_KEY", "dapi-env") + with pytest.raises(ValueError, match="DATABRICKS_API_BASE"): + _ = _CONFIG.connection("https://collector.example/serving-endpoints", None) + + +def test_classifier_response_accounts_the_serving_endpoint_instead_of_the_container_model() -> None: + body: Final = {"model": "/mosaicml/local_model", "answers": {}, "usage": {"input_tokens": 3, "output_tokens": 0}} + normalized: Final = _CONFIG.classifier_response(body, "openjev") + assert normalized["model"] == "openjev" + assert normalized["usage"] == body["usage"] + assert body["model"] == "/mosaicml/local_model" diff --git a/tests/unit/proxy/common_utils/test_registry_read_through.py b/tests/unit/proxy/common_utils/test_registry_read_through.py index 35f448c4fcf2..dd0dca328ac4 100644 --- a/tests/unit/proxy/common_utils/test_registry_read_through.py +++ b/tests/unit/proxy/common_utils/test_registry_read_through.py @@ -203,9 +203,7 @@ def clean_agent_registry(): @pytest.mark.asyncio -async def test_get_agent_with_read_through_recovers_agent_created_on_sibling_replica( - clean_agent_registry, monkeypatch -): +async def test_get_agent_with_read_through_recovers_agent_created_on_sibling_replica(clean_agent_registry, monkeypatch): from unittest.mock import AsyncMock, MagicMock import litellm.proxy.proxy_server as proxy_server @@ -601,7 +599,9 @@ async def test_resync_guardrails_syncs_decrypted_litellm_params(monkeypatch): synced: list[dict] = [] monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) - monkeypatch.setattr(IN_MEMORY_GUARDRAIL_HANDLER, "sync_guardrail_from_db", lambda guardrail: synced.append(guardrail)) + monkeypatch.setattr( + IN_MEMORY_GUARDRAIL_HANDLER, "sync_guardrail_from_db", lambda guardrail: synced.append(guardrail) + ) monkeypatch.setattr(read_through_module, "_initialized_guardrail", lambda guardrail_name: MagicMock()) assert await _resync_guardrails("enc-guardrail") is True @@ -610,15 +610,21 @@ async def test_resync_guardrails_syncs_decrypted_litellm_params(monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("lookup", ["agent-id", "Agent name"]) -async def test_agent_read_through_hydrates_identity_binding(lookup, clean_agent_registry, fresh_agent_read_through, monkeypatch): +async def test_agent_read_through_hydrates_identity_binding( + lookup, clean_agent_registry, fresh_agent_read_through, monkeypatch +): from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through binding = { - "agent_id": "agent-id", "provider": "microsoft_entra", "tenant_id": "tenant", "client_id": "client", - "issuer": "https://login.microsoftonline.com/tenant/v2.0", "revision": "revision", + "agent_id": "agent-id", + "provider": "microsoft_entra", + "tenant_id": "tenant", + "client_id": "client", + "issuer": "https://login.microsoftonline.com/tenant/v2.0", + "revision": "revision", } async def load_row(*, where, include): @@ -733,3 +739,49 @@ async def test_agent_read_through_answers_a_loaded_agent_without_reading_the_db( assert await agent_registry_read_through.attempt("wired-loaded-agent-id") is True assert await agent_registry_read_through.attempt("wired-loaded-agent") is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("router_holds_the_name_after_reconcile", [True, False]) +async def test_resync_model_deployments_reconciles_every_db_row_when_the_missed_model_is_an_auto_router( + monkeypatch: pytest.MonkeyPatch, router_holds_the_name_after_reconcile: bool +) -> None: + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-auto-router-read-through") + router_name: Final = "auto-router-created-on-a-sibling-replica" + row: Final = MagicMock() + row.model_name = router_name + row.litellm_params = { + "model": encrypt_value_helper("auto_router/complexity-router"), + "auto_router_config": {"model_tiers": {"SIMPLE": "simple-tier", "COMPLEX": "complex-tier"}}, + } + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[row]) + router: Final = MagicMock() + router.model_names = [] + router.has_model_id.return_value = False + reconciled: Final = MagicMock() + + async def reconcile_every_row(prisma_client: object, proxy_logging_obj: object) -> None: + reconciled(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + router.model_names = ( + [router_name, "simple-tier", "complex-tier"] if router_holds_the_name_after_reconcile else [] + ) + + def reject_single_row_path(db_models: object) -> None: + raise AssertionError(f"an auto router row took the single-row path: {db_models}") + + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + monkeypatch.setattr(proxy_server.proxy_config, "add_deployment", reconcile_every_row) + monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", reject_single_row_path) + + assert await _resync_model_deployments(router_name) is router_holds_the_name_after_reconcile + reconciled.assert_called_once_with(prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj) diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index ccfa12e10d62..6ce15636f1e2 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -257,6 +257,51 @@ def test_proxy_decisions_dispatches_strands_decider( assert "authorization" not in upstream.calls[0].request.headers +def test_proxy_systemone_dispatches_databricks_serving_endpoint( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("DATABRICKS_API_BASE", raising=False) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.delenv("DATABRICKS_TOKEN", raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "databricks-decider", + "litellm_params": { + "model": "databricks/databricks-openjev-qwen35-4b", + "api_base": "https://workspace.example/serving-endpoints", + "api_key": "dapi-deployment", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + upstream: Final = respx_mock.post( + "https://workspace.example/serving-endpoints/databricks-openjev-qwen35-4b/invocations" + ).respond(json={**_RESPONSE, "model": "/mosaicml/local_model"}) + + response: Final = client.post( + "/v1/systemone", + json={ + "model": "databricks-decider", + "state": {"source": "proxy-test"}, + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + }, + ) + + assert response.status_code == 200, response.text + assert response.json()["answers"] == _RESPONSE["answers"] + assert upstream.called + assert json.loads(upstream.calls[0].request.content) == { + "model": "databricks-openjev-qwen35-4b", + "state": {"source": "proxy-test"}, + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert upstream.calls[0].request.headers["authorization"] == "Bearer dapi-deployment" + + def test_proxy_decisions_without_model_uses_the_proxy_default_model( client: TestClient, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/management_helpers/test_auto_router_permissions.py b/tests/unit/proxy/management_helpers/test_auto_router_permissions.py index 1ccfbab7b1fa..b19445c95d9a 100644 --- a/tests/unit/proxy/management_helpers/test_auto_router_permissions.py +++ b/tests/unit/proxy/management_helpers/test_auto_router_permissions.py @@ -146,6 +146,8 @@ def test_tier_config_is_normalized_and_unknown_router_extras_are_rejected() -> N ({"provider": "laya", "model": "english", "api_key": "sk-member"}, "api_key"), ({"provider": "bespoke", "model": "nimble-latest", "api_base": "https://collector.invalid"}, "api_base"), ({"provider": "bespoke", "model": "nimble-latest", "api_key": "sk-member"}, "api_key"), + ({"provider": "databricks", "model": "my-openjev", "api_base": "https://collector.invalid"}, "opensource_classifier_config"), + ({"provider": "databricks", "model": "my-openjev", "api_key": "dapi-member"}, "api_key"), ], ) @pytest.mark.parametrize("legacy", [False, True]) @@ -164,7 +166,10 @@ def test_members_cannot_move_the_jev_classifier_off_the_proxys_typesafe_account( assert denied.value.detail == f"Invalid member auto-router configuration at {rejected_at}." -@pytest.mark.parametrize(("provider", "model"), [("typesafe", "jev-preview"), ("laya", "english"), ("bespoke", "nimble-latest")]) +@pytest.mark.parametrize( + ("provider", "model"), + [("typesafe", "jev-preview"), ("laya", "english"), ("bespoke", "nimble-latest"), ("databricks", "my-openjev")], +) @pytest.mark.parametrize("legacy", [False, True]) def test_members_can_still_tune_the_jev_classifier(provider: str, model: str, legacy: bool) -> None: validated: Final = validate_member_auto_router_config( diff --git a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py index 418bf522b6aa..fea283cbe9cf 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -451,7 +451,14 @@ def test_jev_config_requires_classifier_config() -> None: ) @pytest.mark.parametrize( ("provider", "model", "canonical_provider"), - [(None, "jev-latest", "jev"), ("typesafe", "jev-latest", "jev"), ("jev", "jev-latest", "jev"), ("laya", "english", "laya"), ("bespoke", "nimble-latest", "bespoke")], + [ + (None, "jev-latest", "jev"), + ("typesafe", "jev-latest", "jev"), + ("jev", "jev-latest", "jev"), + ("laya", "english", "laya"), + ("bespoke", "nimble-latest", "bespoke"), + ("databricks", "databricks-openjev-qwen35-4b", "databricks"), + ], ) def test_classifier_aliases_load_and_serialize_one_canonical_config( classifier_type: str, config_key: str, provider: str | None, model: str, canonical_provider: str @@ -533,6 +540,85 @@ async def test_oss_routes_with_its_own_credentials_and_accounts_the_checkpoint( assert recorder.calls[0]["response_cost"] == pytest.approx(0.31) +def test_databricks_requires_the_bare_serving_endpoint_name_as_the_model() -> None: + with pytest.raises(ValueError, match="model is required for provider 'databricks'"): + JevClassifierConfig.model_validate({"provider": "databricks"}) + with pytest.raises(ValueError, match="bare endpoint name"): + JevClassifierConfig.model_validate({"provider": "databricks", "model": "serving-endpoints/my-openjev"}) + with pytest.raises(ValueError, match="bare endpoint name"): + JevClassifierConfig.model_validate({"provider": "databricks", "model": ".."}) + assert JevClassifierConfig.model_validate({"provider": "databricks", "model": "my-openjev"}).model == "my-openjev" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("configured_key", [False, True]) +@pytest.mark.parametrize("legacy", [False, True]) +async def test_databricks_routes_through_the_serving_endpoint_and_accounts_the_endpoint_model( + monkeypatch: pytest.MonkeyPatch, configured_key: bool, legacy: bool +) -> None: + endpoint: Final = "databricks-openjev-qwen35-4b" + monkeypatch.setenv("TYPESAFE_API_KEY", "never-send-typesafe-key") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.test/serving-endpoints") + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.setenv("DATABRICKS_TOKEN", "dapi-env-token") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(litellm.model_cost, f"databricks/{endpoint}", {"input_cost_per_token": 0.01}) + recorder: Final = _UsageRecorder(f"databricks/{endpoint}") + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + router: Final = ComplexityRouter( + "databricks-route", + litellm.Router(model_list=[]), + { + "classifier_type": "jev" if legacy else "oss_classifier", + "jev_classifier_config" if legacy else "opensource_classifier_config": { + "provider": "databricks", + "model": endpoint, + **({"api_key": "dapi-configured", "api_base": "https://workspace.test/serving-endpoints"} if configured_key else {}), + }, + "tiers": {"SIMPLE": "cheap"}, + }, + derive_savings_baseline=False, + ) + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post(f"https://workspace.test/serving-endpoints/{endpoint}/invocations").respond( + 200, + json={ + "model": "/mosaicml/local_model", + "answers": {"tier": _answer().model_dump()}, + "usage": {"input_tokens": 31, "output_tokens": 0}, + }, + ) + outcome: Final = await router.aclassify("choose a tier") + await GLOBAL_LOGGING_WORKER.flush() + + assert outcome.cause == "jev_classifier" + assert outcome.jev_verdict is not None + assert (outcome.jev_verdict.provider, outcome.jev_verdict.model) == ("databricks", endpoint) + assert outcome.classifier_cost == pytest.approx(0.31) + sent: Final = route.calls.last.request + assert sent.headers.get("authorization") == ("Bearer dapi-configured" if configured_key else "Bearer dapi-env-token") + assert json.loads(sent.content)["model"] == endpoint + assert len(recorder.calls) == 1 + assert recorder.calls[0]["response_cost"] == pytest.approx(0.31) + + +def test_databricks_without_a_base_or_key_fails_at_client_build(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DATABRICKS_API_BASE", raising=False) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.delenv("DATABRICKS_TOKEN", raising=False) + with pytest.raises(ValueError, match="DATABRICKS_API_BASE"): + ComplexityRouter( + "databricks-route", + litellm.Router(model_list=[]), + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"provider": "databricks", "model": "databricks-openjev-qwen35-4b"}, + "tiers": {"SIMPLE": "cheap"}, + }, + derive_savings_baseline=False, + ) + + def test_jev_config_is_rejected_for_other_classifier_types() -> None: with pytest.raises(ValueError, match="has no effect"): ComplexityRouterConfig.model_validate( @@ -570,6 +656,22 @@ def test_jev_api_base_without_its_own_key_is_rejected_so_the_environment_key_sta assert JevClassifierConfig(api_key="sk-own").api_base is None +def test_databricks_api_base_without_its_own_key_is_rejected_so_the_workspace_token_stays_home() -> None: + with pytest.raises(ValueError, match="DATABRICKS_API_KEY or DATABRICKS_TOKEN is only sent to DATABRICKS_API_BASE"): + JevClassifierConfig.model_validate( + {"provider": "databricks", "model": "my-openjev", "api_base": "https://collector.invalid/serving-endpoints"} + ) + paired: Final = JevClassifierConfig.model_validate( + { + "provider": "databricks", + "model": "my-openjev", + "api_base": "https://workspace.invalid/serving-endpoints", + "api_key": "dapi-own", + } + ) + assert (paired.api_base, paired.api_key) == ("https://workspace.invalid/serving-endpoints", "dapi-own") + + @pytest.mark.parametrize( ("probabilities", "confidence"), [ diff --git a/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx b/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx index a8d492a6ced8..f21dc3d795aa 100644 --- a/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AutoRouterClassifierTabs.tsx @@ -16,7 +16,11 @@ import { type ClassifierType, type ComplexityRouterConfigValue, } from "./ComplexityRouterConfig"; -import { defaultJevClassifierConfig, normalizeJevClassifierConfig } from "./jev_classifier_config"; +import { + defaultJevClassifierConfig, + isOssClassifierProvider, + normalizeJevClassifierConfig, +} from "./jev_classifier_config"; import { transitionClassifierType } from "./classifier_type_transition"; import { isForecastClassifier } from "./forecast_classifier_config"; import { @@ -150,7 +154,7 @@ const AutoRouterClassifierTabs: React.FC = ({ val if (next === "jev") changeType("jev"); }; const changeProvider = (provider: unknown) => { - if (provider !== "jev" && provider !== "laya" && provider !== "bespoke") return; + if (!isOssClassifierProvider(provider)) return; const defaults = defaultJevClassifierConfig(provider); onChange({ ...value, @@ -173,7 +177,11 @@ const AutoRouterClassifierTabs: React.FC = ({ val {[ { value: "heuristics", label: "Heuristics", description: "Classify locally, with no API call" }, { value: "llm", label: "LLM", description: "Use a judge model to choose a solver" }, - { value: "jev", label: "OSS Classifier", description: "Use Jev, Laya, or Bespoke Nimble to choose a tier" }, + { + value: "jev", + label: "OSS Classifier", + description: "Use Jev, Laya, Bespoke Nimble, or Databricks to choose a tier", + }, ].map((option) => ( + )} diff --git a/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx b/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx index 2eda2c63f641..9deb53753bda 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx @@ -55,7 +55,9 @@ const ClassifierTypeRadios: React.FC = ({ value, clas OSS Classifier{" "} - uses Jev or Laya to decide the tier + + uses Jev, Laya, Bespoke, or a Databricks endpoint to decide the tier + diff --git a/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.integration.test.tsx index 99726a8f071b..15a736de7b29 100644 --- a/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.integration.test.tsx @@ -101,9 +101,11 @@ describe("JEV classifier editor", () => { ["jev", "Jev", "jev-test"], ["laya", "Laya", "multilingual"], ["bespoke", "Bespoke Nimble", "bespokelabs/Bespoke-Nimble-9B"], + ["databricks", "Databricks", "databricks-openjev-qwen35-4b"], ] as const)( "preserves %s, custom tiers and context through save, reload and probe", async (provider, label, model) => { + const typesTheModel = provider === "jev" || provider === "databricks"; renderWithProviders(
); expect(screen.getByLabelText("Judge model")).toBeInTheDocument(); expect(screen.getByText("Reasoning Effort")).toBeInTheDocument(); @@ -123,10 +125,11 @@ describe("JEV classifier editor", () => { expect(screen.getByLabelText("Classifier Model")).toHaveValue("jev-latest"); fireEvent.click(screen.getByRole("radio", { name: label })); if (provider === "bespoke") expect(screen.getByLabelText("Classifier Model")).toHaveTextContent("nimble-latest"); - if (provider !== "jev") { - await chooseSelectOption(userEvent, screen.getByLabelText("Classifier Model"), model); - } else { + if (provider === "databricks") expect(screen.getByLabelText("Classifier Model")).toHaveValue(""); + if (typesTheModel) { fireEvent.change(screen.getByLabelText("Classifier Model"), { target: { value: model } }); + } else { + await chooseSelectOption(userEvent, screen.getByLabelText("Classifier Model"), model); } fireEvent.change(screen.getByLabelText("Classifier Timeout (ms)"), { target: { value: "4200" } }); fireEvent.change(screen.getByLabelText("Context Window Size"), { target: { value: "6" } }); @@ -136,8 +139,8 @@ describe("JEV classifier editor", () => { fireEvent.click(screen.getByRole("button", { name: "Save and reload" })); expect(screen.getByRole("radio", { name: /^OSS Classifier$/ })).toBeChecked(); expect(screen.getByRole("radio", { name: label })).toBeChecked(); - if (provider !== "jev") expect(screen.getByLabelText("Classifier Model")).toHaveTextContent(model); - else expect(screen.getByLabelText("Classifier Model")).toHaveValue(model); + if (typesTheModel) expect(screen.getByLabelText("Classifier Model")).toHaveValue(model); + else expect(screen.getByLabelText("Classifier Model")).toHaveTextContent(model); expect(screen.getByLabelText("Classifier Timeout (ms)")).toHaveValue(4200); expect(screen.getByLabelText("Context Window Size")).toHaveValue("6"); expect(screen.getByRole("switch", { name: "Classifier circuit breaker" })).not.toBeChecked(); diff --git a/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx index 3818e3ae9c2b..86d2d0f83483 100644 --- a/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx @@ -7,13 +7,15 @@ import { Textarea } from "@/components/ui/textarea"; import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig"; import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { defaultJevClassifierConfig, OSS_CLASSIFIER_MODELS } from "./jev_classifier_config"; +import { defaultJevClassifierConfig, fixedClassifierModels } from "./jev_classifier_config"; const providerDescriptions = { jev: "Uses TypeSafe System One Choice evaluation with your configured tiers", laya: "Uses Laya with your configured tiers. Set LAYA_API_BASE on the gateway to connect your Laya server.", bespoke: "Uses Bespoke Nimble with your configured tiers. Set BESPOKE_API_BASE on the gateway to connect your Nimble server.", + databricks: + "Uses a Databricks ai_decide serving endpoint with your configured tiers. Set DATABRICKS_API_BASE and DATABRICKS_API_KEY on the gateway, and enter the serving endpoint name as the classifier model.", }; export default function JevClassifierConfig({ @@ -25,7 +27,7 @@ export default function JevClassifierConfig({ }) { const id = useId(); const config = value.jev_classifier_config ?? defaultJevClassifierConfig(); - const models = config.provider && config.provider !== "jev" ? OSS_CLASSIFIER_MODELS[config.provider] : undefined; + const models = fixedClassifierModels(config.provider); const update = (patch: Partial) => onChange({ ...value, jev_classifier_config: { ...config, ...patch } }); @@ -48,7 +50,14 @@ export default function JevClassifierConfig({ ) : ( - update({ model: event.target.value })} /> + update({ model: event.target.value })} + /> )}
diff --git a/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts b/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts index 4e3ef490079e..9ce99d395641 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts @@ -52,6 +52,8 @@ describe("buildAutoRouterRoutingTestRequest", () => { ["bespoke", "nimble-latest", "json"], ["bespoke", "bespokelabs/Bespoke-Nimble-9B", "object"], ["bespoke", "bespokelabs/Bespoke-Nimble-9B", "json"], + ["databricks", "databricks-openjev-qwen35-4b", "object"], + ["databricks", "databricks-openjev-qwen35-4b", "json"], ])("probes saved %s/%s %s configuration with custom tiers and team context", (provider, model, format) => { const config = { classifier_type: "oss_classifier", @@ -71,10 +73,15 @@ describe("buildAutoRouterRoutingTestRequest", () => { buildSavedJevConnectionTestRequest(format === "json" ? JSON.stringify(config) : config, "saved-id", "team-1"), ).toEqual(expectedRequest); }); - it.each(["laya", "bespoke"])("does not probe unsupported %s models", (provider) => { + it.each([ + ["laya", "unsupported"], + ["bespoke", "unsupported"], + ["databricks", "serving-endpoints/openjev"], + ["databricks", ".."], + ])("does not probe unsupported %s model %s", (provider, model) => { const config = { classifier_type: "oss_classifier", - opensource_classifier_config: { provider, model: "unsupported" }, + opensource_classifier_config: { provider, model }, tiers: CONFIG.tiers, }; expect(buildSavedJevConnectionTestRequest(config, "saved-id")).toBeUndefined(); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index e1153e34bc8d..23dedbbc8603 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -141,6 +141,9 @@ describe("buildComplexityRouterConfig", () => { { model: " " }, { provider: "laya" as const, model: "unsupported" }, { provider: "bespoke" as const, model: "unsupported" }, + { provider: "databricks" as const, model: "" }, + { provider: "databricks" as const, model: "serving-endpoints/openjev" }, + { provider: "databricks" as const, model: ".." }, { timeout_ms: 0 }, { timeout_ms: 1.5 }, { timeout_ms: Number.NaN }, @@ -159,6 +162,7 @@ describe("buildComplexityRouterConfig", () => { ["bespoke", "nimble-latest"], ["bespoke", "nimble"], ["bespoke", "bespokelabs/Bespoke-Nimble-9B"], + ["databricks", "databricks-openjev-qwen35-4b"], ["jev", "custom-jev-model"], [undefined, "custom-jev-model"], ] as const)("accepts %s model %s before saving or testing", (provider, model) => { @@ -177,6 +181,8 @@ describe("buildComplexityRouterConfig", () => { ["laya", "english", true], ["bespoke", "nimble-latest", false], ["bespoke", "nimble-latest", true], + ["databricks", "databricks-openjev-qwen35-4b", false], + ["databricks", "databricks-openjev-qwen35-4b", true], ] as const)("serializes %s/%s with shared context and no LLM config, custom tiers: %s", (provider, model, custom) => { const params: BuildComplexityRouterConfigParams = { ...baseParams, diff --git a/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.test.ts b/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.test.ts new file mode 100644 index 000000000000..05a10dd4927c --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from "vitest"; +import { defaultJevClassifierConfig, hydrateOssClassifier, jevClassifierConfigSchema } from "./jev_classifier_config"; + +describe("hydrateOssClassifier", () => { + it("keeps a stored Databricks classifier instead of falling back to Jev", () => { + const stored = { + provider: "databricks", + model: "databricks-openjev-qwen35-4b", + timeout_ms: 4200, + circuit_breaker_enabled: true, + circuit_breaker_cooldown_seconds: 50, + }; + + const hydrated = hydrateOssClassifier({ classifier_type: "oss_classifier", opensource_classifier_config: stored }); + + expect(hydrated).toEqual({ classifier_type: "jev", jev_classifier_config: stored }); + }); + + it("falls back to Jev only when the stored classifier cannot be parsed", () => { + const hydrated = hydrateOssClassifier({ + classifier_type: "oss_classifier", + opensource_classifier_config: { provider: "databricks", model: "serving-endpoints/openjev" }, + }); + + expect(hydrated.jev_classifier_config).toEqual(defaultJevClassifierConfig()); + }); +}); + +describe("jevClassifierConfigSchema", () => { + it("starts a Databricks classifier with an empty endpoint name that does not pass until typed", () => { + const fresh = defaultJevClassifierConfig("databricks"); + + expect(fresh).toEqual({ provider: "databricks", model: "", timeout_ms: 3000 }); + expect(jevClassifierConfigSchema.safeParse(fresh).success).toBe(false); + expect(jevClassifierConfigSchema.safeParse({ ...fresh, model: "my-openjev" }).success).toBe(true); + }); + + it.each(["serving-endpoints/openjev", "openjev?x=1", "openjev#frag", ".", "..", ".openjev", "open jev"])( + "rejects the Databricks endpoint name %j because it would change the request URL", + (model) => { + expect(jevClassifierConfigSchema.safeParse({ provider: "databricks", model }).success).toBe(false); + }, + ); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.ts b/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.ts index a9a7cc18711a..02cd8bcbfe97 100644 --- a/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.ts @@ -6,13 +6,33 @@ export const OSS_CLASSIFIER_MODELS = { bespoke: ["nimble-latest", "nimble", "bespokelabs/Bespoke-Nimble-9B"], } as const; +const OSS_CLASSIFIER_PROVIDERS = ["jev", "laya", "bespoke", "databricks"] as const; +export type OssClassifierProvider = (typeof OSS_CLASSIFIER_PROVIDERS)[number]; +export const isOssClassifierProvider = (value: unknown): value is OssClassifierProvider => + OSS_CLASSIFIER_PROVIDERS.some((provider) => provider === value); + +const DEFAULT_TIMEOUT_MS = 3000; +const SERVING_ENDPOINT_NAME = /^[A-Za-z0-9_-][A-Za-z0-9._-]*$/; + +export const fixedClassifierModels = (provider: OssClassifierProvider | undefined): readonly string[] | undefined => + provider === "laya" || provider === "bespoke" ? OSS_CLASSIFIER_MODELS[provider] : undefined; + +const defaultClassifierModel = (provider: OssClassifierProvider | undefined): string => + fixedClassifierModels(provider)?.[0] ?? (provider === "databricks" ? "" : "jev-latest"); + +const isSupportedClassifierModel = (provider: OssClassifierProvider | undefined, model: string): boolean => { + if (provider === "databricks") return SERVING_ENDPOINT_NAME.test(model); + const fixed = fixedClassifierModels(provider); + return fixed === undefined || fixed.some((candidate) => candidate === model); +}; + const jevClassifierConfigFields = { provider: z.preprocess( (value) => (value === "typesafe" ? "jev" : value), - z.enum(["jev", "laya", "bespoke"]).optional(), + z.enum(["jev", "laya", "bespoke", "databricks"]).optional(), ), model: z.string().trim().min(1).optional(), - timeout_ms: z.number().int().positive().default(3000), + timeout_ms: z.number().int().positive().default(DEFAULT_TIMEOUT_MS), instructions: z .string() .nullish() @@ -24,24 +44,18 @@ const jevClassifierConfigFields = { export const jevClassifierConfigSchema = z .object(jevClassifierConfigFields) - .transform((config) => ({ - ...config, - model: - config.model ?? - (config.provider && config.provider !== "jev" ? OSS_CLASSIFIER_MODELS[config.provider][0] : "jev-latest"), - })) - .refine( - (config) => - !config.provider || - config.provider === "jev" || - OSS_CLASSIFIER_MODELS[config.provider].some((model) => model === config.model), - { error: "Select a supported classifier model", path: ["model"] }, - ); + .transform((config) => ({ ...config, model: config.model ?? defaultClassifierModel(config.provider) })) + .refine((config) => isSupportedClassifierModel(config.provider, config.model), { + error: "Select a supported classifier model, or enter the bare Databricks serving endpoint name", + path: ["model"], + }); export type JevClassifierConfig = z.infer; -export const defaultJevClassifierConfig = (provider: JevClassifierConfig["provider"] = "jev"): JevClassifierConfig => - jevClassifierConfigSchema.parse({ provider }); +export const defaultJevClassifierConfig = (provider: OssClassifierProvider = "jev"): JevClassifierConfig => + provider === "databricks" + ? { provider, model: "", timeout_ms: DEFAULT_TIMEOUT_MS } + : jevClassifierConfigSchema.parse({ provider }); export const hydrateOssClassifier = (config: { classifier_type?: ClassifierType | "oss_classifier"; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 85219633c6c4..95a60bd9aede 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -40472,7 +40472,7 @@ export interface components { * @default jev * @enum {string} */ - provider: "jev" | "laya" | "bespoke"; + provider: "jev" | "laya" | "bespoke" | "databricks"; /** * Timeout Ms * @default 3000 From 3880fc285b1bffebc1278362d11a8fe087f1e164 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:29:32 -0700 Subject: [PATCH 05/18] fix(decisions)!: refuse safety_identifier on providers that cannot take it unless drop_params drops it (#44955) Backport of #44955 to rc/1.106.0. Cherry-picked from 546402c98c7ecea5f32d263fadd3c95bd9b934ff (main). --- litellm/decisions/main.py | 66 +- .../llms/base_llm/decisions/transformation.py | 1 + .../llms/openai/decisions/transformation.py | 1 + .../proxy/decisions_endpoints/endpoints.py | 7 +- .../test_databricks_decisions_wire.py | 1 - .../test_decisions_openai_format_wire.py | 586 ++++++++++++++++++ .../translation/decisions/bases/cloudflare.py | 97 +++ .../translation/decisions/bases/openrouter.py | 94 +++ .../translation/decisions/bases/perplexity.py | 95 +++ .../decisions/bases/strands_decider.py | 34 + .../translation/decisions/bases/typesafe.py | 92 +++ .../basic/test_decisions_basic_cloudflare.py | 4 +- .../basic/test_decisions_basic_openrouter.py | 9 +- .../basic/test_decisions_basic_perplexity.py | 9 +- .../test_decisions_basic_strands_decider.py | 11 +- .../basic/test_decisions_basic_typesafe.py | 4 +- tests/unit/decisions/test_main.py | 135 ++++ .../decisions_endpoints/test_endpoints.py | 62 ++ 18 files changed, 1289 insertions(+), 19 deletions(-) create mode 100644 tests/integration/providers/test_decisions_openai_format_wire.py diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index c47ba0335c7c..f463bd58e3c7 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -1,13 +1,13 @@ from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Final, TypeAlias import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError from typing_extensions import assert_never import litellm -from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER +from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( @@ -40,6 +40,9 @@ _SYSTEMONE_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody) _OPENAI_REQUEST_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(OpenAIDecisionRequestBody) +_SAFETY_IDENTIFIER_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter( + str | None, config=ConfigDict(title="safety_identifier") +) _HANDLER: Final = BaseLLMHTTPHandler() @@ -96,16 +99,46 @@ def _validate_request( ) -def _ir_request(request: DecisionsRequestFormat) -> DecisionsIRRequest: +def _ir_request(request: DecisionsRequestFormat, safety_identifier: str | None) -> DecisionsIRRequest: match request: case DecisionsRequestBody(): - return systemone_request_to_ir(request) + return replace(systemone_request_to_ir(request), safety_identifier=safety_identifier) case OpenAIDecisionRequestBody(): return openai_request_to_ir(request) case _: assert_never(request) +def _drops_params(kwargs: Mapping[str, object]) -> bool: + return litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True + + +def _request_safety_identifier(safety_identifier: object, kwargs: Mapping[str, object]) -> str | None: + if isinstance(safety_identifier, str) or not _drops_params(kwargs): + return _SAFETY_IDENTIFIER_ADAPTER.validate_python(safety_identifier) + return None + + +@dataclass(frozen=True, slots=True) +class _UnsupportedSafetyIdentifier: + pass + + +def _provider_ir_request( + request: DecisionsRequestFormat, + *, + safety_identifier: str | None, + provider_config: BaseDecisionsConfig, + kwargs: Mapping[str, object], +) -> DecisionsIRRequest | _UnsupportedSafetyIdentifier: + ir_request: Final = _ir_request(request, safety_identifier) + if ir_request.safety_identifier is None or provider_config.supports_safety_identifier: + return ir_request + if _drops_params(kwargs): + return replace(ir_request, safety_identifier=None) + return _UnsupportedSafetyIdentifier() + + def _prepare_call( *, model: str, @@ -144,8 +177,12 @@ def _prepare_call( except ValueError as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error try: + request_safety_identifier: Final = _request_safety_identifier(safety_identifier, kwargs) request: Final = _validate_request( - state=state, questions=questions, decision_input=decision_input, safety_identifier=safety_identifier + state=state, + questions=questions, + decision_input=decision_input, + safety_identifier=request_safety_identifier, ) except ValidationError as error: raise litellm.BadRequestError( @@ -169,7 +206,22 @@ def _prepare_call( llm_provider=provider, ) - ir_request: Final = _ir_request(request) + ir_request: Final = _provider_ir_request( + request, + safety_identifier=request_safety_identifier, + provider_config=provider_config, + kwargs=kwargs, + ) + if isinstance(ir_request, _UnsupportedSafetyIdentifier): + raise litellm.UnsupportedParamsError( + message=( + f"{provider} does not support parameters: ['safety_identifier'], for model={model}. " + "To drop these, set `litellm.drop_params=True` or for proxy:\n\n" + "`litellm_settings:\n drop_params: true`\n" + ), + model=model, + llm_provider=provider, + ) body: Final = provider_config.transform_decisions_request( model=canonical_model, request=ir_request, custom_llm_provider=provider ) diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index c57c643983a5..baeeb3f04492 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -301,6 +301,7 @@ class BaseDecisionsConfig(ABC): api_key_env: tuple[str, ...] = () api_base_env: tuple[str, ...] = () api_key_required: bool = True + supports_safety_identifier: bool = False health_check_questions: Mapping[str, Mapping[str, object]] = MappingProxyType( {"reachable": MappingProxyType({"type": "noul", "instructions": "Is the service reachable?"})} ) diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 0afbe1a434e6..317e5abeeb5b 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -292,6 +292,7 @@ def ir_to_openai_response( class OpenAIDecisionsConfig(BaseDecisionsConfig): path = "/v1/decisions" + supports_safety_identifier = True def get_default_api_base(self) -> str | None: return "https://api.openai.com" diff --git a/litellm/proxy/decisions_endpoints/endpoints.py b/litellm/proxy/decisions_endpoints/endpoints.py index f2891c184bf6..4f4209208aca 100644 --- a/litellm/proxy/decisions_endpoints/endpoints.py +++ b/litellm/proxy/decisions_endpoints/endpoints.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from types import MappingProxyType from typing import Annotated, Final from fastapi import APIRouter, Depends, Request, Response @@ -38,6 +39,10 @@ async def _invalid_request( ) +def _fields_checked_before_routing(data: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in data.items() if key != "safety_identifier"}) + + async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: body: Final = await request.body() try: @@ -75,7 +80,7 @@ async def _process_decisions( data: Final = await _request_data(request, user_api_key_dict) try: - body_adapter.validate_python(data) + body_adapter.validate_python(_fields_checked_before_routing(data)) except ValidationError as error: raise await _invalid_request(raw_data=data, error=error, user_api_key_dict=user_api_key_dict) general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(proxy_general_settings) diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py index 5c3b7d1b1a1d..93dc2ea04faf 100644 --- a/tests/integration/providers/test_databricks_decisions_wire.py +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -208,7 +208,6 @@ def test_an_openai_format_request_at_v1_decisions_reaches_the_serving_endpoint_a "model": deployment, "input": _OPENAI_INPUT, "questions": [_OPENAI_TIER_QUESTION], - "safety_identifier": "end-user-1", }, ) assert response.status_code == 200, response.text diff --git a/tests/integration/providers/test_decisions_openai_format_wire.py b/tests/integration/providers/test_decisions_openai_format_wire.py new file mode 100644 index 000000000000..f0e7163e264b --- /dev/null +++ b/tests/integration/providers/test_decisions_openai_format_wire.py @@ -0,0 +1,586 @@ +import math +import uuid +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse +from pydantic import JsonValue, TypeAdapter + +import litellm + +_API_KEY: Final = "synthetic-decisions-key" +_SAFETY_IDENTIFIER: Final = "end-user-7" +_DROPPING_DEPLOYMENT: Final = "decisions-under-litellm-settings-drop-params" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_INPUT: Final = "Ticket (billing): The export job hangs at 99%" +_FOLLOW_UP: Final = "Customer: still stuck after retrying" +_QUESTIONS: Final[list[JsonValue]] = [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [{"value": "low", "description": "cosmetic"}, {"value": "high", "description": "blocks users"}], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, +] +_SYSTEM_ONE_QUESTIONS: Final[dict[str, JsonValue]] = { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, +} +_SYSTEM_ONE_ANSWERS: Final[dict[str, JsonValue]] = { + "defect": {"type": "noul", "noul": 0.93}, + "severity": {"type": "choice", "choice": "high", "confidence": 0.8, "probabilities": {"low": 0.2, "high": 0.8}}, + "confidence": { + "type": "score", + "score": 1.0, + "confidence": 0.7, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.3, "1": 0.7}, + }, +} +_ANSWERS: Final[list[JsonValue]] = [ + {"type": "predicate", "name": "defect", "probability": 0.93}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.2}, {"value": "high", "probability": 0.8}], + "confidence": 0.8, + }, + { + "type": "score", + "name": "confidence", + "score": 1.0, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.3}, + {"value": 1, "label": "sure", "probability": 0.7}, + ], + "confidence": 0.7, + }, +] +_INPUT_TOKENS: Final = 367 +_OUTPUT_TOKENS: Final = 3 +_SYSTEM_ONE_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": _INPUT_TOKENS, "output_tokens": _OUTPUT_TOKENS} +_USAGE: Final[dict[str, JsonValue]] = { + "input_tokens": _INPUT_TOKENS, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": _OUTPUT_TOKENS, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": _INPUT_TOKENS + _OUTPUT_TOKENS, +} +_SDK_QUESTIONS: Final = TypeAdapter(list[dict[str, object]]).validate_python(_QUESTIONS) +_SPEND_QUERY: Final = ( + "SELECT spend, status, call_type, model_group, custom_llm_provider, api_base, prompt_tokens, completion_tokens, " + 'request_tags FROM "LiteLLM_SpendLogs" WHERE request_id = %s' +) + + +@dataclass(frozen=True, slots=True) +class _Provider: + name: str + model: str + path: str + body_model: str + api_key: str | None + wraps_result: bool + cost_map_key: str | None + speaks_openai: bool = False + + def upstream_body(self) -> dict[str, JsonValue]: + if self.speaks_openai: + return {"model": self.body_model, "input": _INPUT, "questions": _QUESTIONS} + return {"model": self.body_model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS} + + def upstream_reply(self) -> dict[str, JsonValue]: + if self.speaks_openai: + return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE} + answer: Final[dict[str, JsonValue]] = { + "model": self.body_model, + "answers": _SYSTEM_ONE_ANSWERS, + "usage": _SYSTEM_ONE_USAGE, + } + return {"result": answer, "success": True} if self.wraps_result else answer + + def litellm_response(self) -> dict[str, JsonValue]: + return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE} + + +_PROVIDERS: Final = ( + _Provider( + "perplexity", + "perplexity/pplx-decider-v1-27b", + "/v1/decisions", + "pplx-decider-v1-27b", + _API_KEY, + False, + "perplexity/pplx-decider-v1-27b", + ), + _Provider("typesafe", "typesafe/jev-1.13.0", "/v1/systemone", "jev-1.13.0", _API_KEY, False, "typesafe/jev-1.13.0"), + _Provider( + "openrouter", + "openrouter/typesafe/jev-1.13", + "/alpha/decisions", + "typesafe/jev-1.13", + _API_KEY, + False, + "openrouter/typesafe/jev-1.13", + ), + _Provider( + "strands_decider", "strands_decider/systemone-decider", "/v1/systemone", "systemone-decider", None, False, None + ), + _Provider( + "cloudflare", + "cloudflare/clef", + "/ai/run/@cf/cloudflare/clef", + "clef", + _API_KEY, + True, + "cloudflare/@cf/cloudflare/clef", + ), + _Provider("hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None), + _Provider( + "databricks", + "databricks/databricks-openjev-qwen35-4b", + "/databricks-openjev-qwen35-4b/invocations", + "databricks-openjev-qwen35-4b", + _API_KEY, + False, + None, + ), + _Provider( + "azure_ai", "azure_ai/decision-1", "/providers/microsoft/v1/systemone", "decision-1", _API_KEY, False, None + ), + _Provider("openai", "openai/gpt-6-luna", "/v1/decisions", "gpt-6-luna", _API_KEY, False, "gpt-6-luna", True), +) +_SYSTEM_ONE_PROVIDERS: Final = tuple(provider for provider in _PROVIDERS if not provider.speaks_openai) +_PERPLEXITY: Final = _PROVIDERS[0] +_TYPESAFE: Final = _PROVIDERS[1] +_OPENAI: Final = next(provider for provider in _PROVIDERS if provider.speaks_openai) +_PREDICATE: Final[dict[str, JsonValue]] = {"type": "predicate", "name": "q", "instructions": "Is it?"} +_INVALID_BODIES: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ("missing questions", {"input": _INPUT}), + ("missing input", {"questions": _QUESTIONS}), + ("numeric input", {"input": 5, "questions": _QUESTIONS}), + ("assistant message input", {"input": [{"role": "assistant", "content": "hi"}], "questions": _QUESTIONS}), + ("empty questions", {"input": _INPUT, "questions": []}), + ("questions as a map", {"input": _INPUT, "questions": {"q": _PREDICATE}}), + ("predicate without instructions", {"input": _INPUT, "questions": [{"type": "predicate", "name": "q"}]}), + ( + "choice with one choice", + {"input": _INPUT, "questions": [{**_PREDICATE, "type": "choice", "choices": [{"value": "only"}]}]}, + ), + ( + "score with one level", + {"input": _INPUT, "questions": [{**_PREDICATE, "type": "score", "levels": [{"label": "only"}]}]}, + ), + ("unknown question type", {"input": _INPUT, "questions": [{**_PREDICATE, "type": "ranking"}]}), + ("numeric safety_identifier", {"input": _INPUT, "questions": _QUESTIONS, "safety_identifier": 7}), + ("list safety_identifier", {"input": _INPUT, "questions": _QUESTIONS, "safety_identifier": [_SAFETY_IDENTIFIER]}), +) +_IMAGE_INPUT: Final[list[JsonValue]] = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": _INPUT}, + {"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo=", "detail": "auto"}, + ], + } +] +_OPENAI_IMAGE_MESSAGES: Final[list[JsonValue]] = [ + { + "role": "user", + "type": "message", + "content": [ + {"type": "input_text", "text": _INPUT}, + {"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo=", "detail": "auto"}, + ], + } +] +_MESSAGE_LIST_INPUT: Final[list[JsonValue]] = [ + {"role": "user", "content": _INPUT}, + {"role": "user", "content": [{"type": "input_text", "text": _FOLLOW_UP}]}, +] + + +def _response_cost(response: httpx.Response) -> float: + return float(response.headers["x-litellm-response-cost"]) if "x-litellm-response-cost" in response.headers else 0.0 + + +def _provider_id(provider: _Provider) -> str: + return provider.name + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def _expected_spend(cost_map_key: str | None) -> float: + if cost_map_key is None: + return 0.0 + cost_map: Final = _JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes()) + prices: Final = object_value(cost_map[cost_map_key]) + return _INPUT_TOKENS * _number(prices["input_cost_per_token"]) + _OUTPUT_TOKENS * _number( + prices["output_cost_per_token"] + ) + + +def _register(scenario: Scenario, body: dict[str, JsonValue], *, status: int = 200) -> ScenarioHandle: + handle: Final = register_scenario( + f"decisions-{uuid.uuid4().hex[:12]}", JsonResponse(content_type="application/json", body=body, status=status) + ) + scenario.cleanups.callback(delete_scenario, handle) + return handle + + +def _deployment(scenario: Scenario, handle: ScenarioHandle, provider: _Provider, *, drop_params: bool = False) -> str: + if drop_params: + return scenario.model( + model=provider.model, api_base=handle.api_base(), api_key=provider.api_key, drop_params=True + ) + return scenario.model(model=provider.model, api_base=handle.api_base(), api_key=provider.api_key) + + +def _decide(gateway: Gateway, model: str, *, route: str = "/v1/decisions", **extra: JsonValue) -> httpx.Response: + return gateway.request("POST", route, {"model": model, "input": _INPUT, "questions": _QUESTIONS, **extra}) + + +def _decide_in_system_one_format(gateway: Gateway, model: str, **extra: JsonValue) -> httpx.Response: + return gateway.request( + "POST", "/v1/systemone", {"model": model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS, **extra} + ) + + +def _observed_requests(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observations: Final = _JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"] + assert isinstance(observations, list), observations + return tuple(map(object_value, observations)) + + +def _calls_to(requests: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> list[dict[str, JsonValue]]: + return [request for request in requests if string_value(request["path"]).startswith(f"/{handle.scenario_id}/")] + + +def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, JsonValue]]: + return _calls_to(_observed_requests(gateway), handle) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _assert_refused_as_invalid(gateway: Gateway, model: str, label: str, body: dict[str, JsonValue]) -> None: + response: Final = gateway.request("POST", "/v1/decisions", {"model": model, **body}) + assert response.status_code == 400, (label, response.text) + assert "Invalid Decisions request" in response.text, (label, response.text) + + +def _assert_refused_for_safety_identifier(response: httpx.Response) -> None: + assert response.status_code == 400, response.text + assert "safety_identifier" in response.text and "drop_params" in response.text, response.text + + +def _drop_params_config(directory: Path, api_base: str) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "litellm_settings": {**object_value(base["litellm_settings"]), "drop_params": True}, + "model_list": [ + { + "model_name": _DROPPING_DEPLOYMENT, + "litellm_params": {"model": _PERPLEXITY.model, "api_base": api_base, "api_key": _API_KEY}, + } + ], + } + path: Final = directory / "decisions-drop-params.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.mark.parametrize("provider", _PROVIDERS, ids=_provider_id) +def test_each_provider_gets_its_own_path_key_and_body_and_is_billed_from_the_cost_map( + gateway: Gateway, provider: _Provider +) -> None: + expected_spend: Final = _expected_spend(provider.cost_map_key) + with gateway.scenario() as scenario: + handle: Final = _register(scenario, provider.upstream_reply()) + model: Final = _deployment(scenario, handle, provider) + response: Final = _decide(gateway, model) + assert response.status_code == 200, response.text + assert response.json() == provider.litellm_response() + assert response.headers["x-litellm-model-group"] == model + assert math.isclose(_response_cost(response), expected_spend, rel_tol=1e-9) + (call,) = _upstream_calls(gateway, handle) + assert call["path"] == f"/{handle.scenario_id}{provider.path}" + assert call["authorization"] == (f"Bearer {provider.api_key}" if provider.api_key else "") + assert call["body"] == provider.upstream_body() + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert ( + row["status"], + row["call_type"], + row["custom_llm_provider"], + row["model_group"], + row["api_base"], + row["prompt_tokens"], + row["completion_tokens"], + ) == ( + "success", + "adecisions", + provider.name, + model, + f"{handle.api_base()}{provider.path}", + _INPUT_TOKENS, + _OUTPUT_TOKENS, + ) + assert math.isclose(_number(row["spend"]), expected_spend, rel_tol=1e-9), row + + +def test_the_unversioned_alias_serves_the_same_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + response: Final = _decide(gateway, model, route="/decisions") + assert response.status_code == 200, response.text + assert response.json() == _PERPLEXITY.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + + +def test_repeated_identical_requests_each_reach_the_upstream_and_are_each_billed(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + responses: Final = tuple(_decide(gateway, model) for _ in range(2)) + assert [response.status_code for response in responses] == [200, 200], [r.text for r in responses] + call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses) + assert len(set(call_ids)) == 2, call_ids + assert len(_upstream_calls(gateway, handle)) == 2 + for call_id in call_ids: + assert _spend_row(call_id)["status"] == "success" + + +async def test_sdk_sync_and_async_clients_send_the_same_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _TYPESAFE.upstream_reply()) + synchronous: Final = litellm.decisions( + model=_TYPESAFE.model, input=_INPUT, questions=_SDK_QUESTIONS, api_base=handle.api_base(), api_key=_API_KEY + ) + asynchronous: Final = await litellm.adecisions( + model=_TYPESAFE.model, input=_INPUT, questions=_SDK_QUESTIONS, api_base=handle.api_base(), api_key=_API_KEY + ) + for response in (synchronous, asynchronous): + assert response.model_dump(mode="json") == _TYPESAFE.litellm_response() + calls: Final = _upstream_calls(gateway, handle) + assert len(calls) == 2, calls + for call in calls: + assert call["path"] == f"/{handle.scenario_id}{_TYPESAFE.path}" + assert call["authorization"] == f"Bearer {_API_KEY}" + assert call["body"] == _TYPESAFE.upstream_body() + + +def test_gateway_only_fields_stay_at_the_gateway_and_tags_reach_the_spend_log(gateway: Gateway) -> None: + tag: Final = f"decisions-openai-format-{uuid.uuid4().hex[:8]}" + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + response: Final = _decide( + gateway, model, user="auditor", num_retries=0, temperature=0.2, metadata={"tags": [tag]} + ) + assert response.status_code == 200, response.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + tags: Final = row["request_tags"] + assert isinstance(tags, list) and tag in tags, row + + +def test_invalid_bodies_are_refused_at_the_gateway_without_an_upstream_call(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + for label, body in _INVALID_BODIES: + _assert_refused_as_invalid(gateway, model, label, body) + assert _upstream_calls(gateway, handle) == [] + + +@pytest.mark.parametrize("provider", _SYSTEM_ONE_PROVIDERS, ids=_provider_id) +def test_system_one_providers_refuse_images_at_the_gateway_without_an_upstream_call( + gateway: Gateway, provider: _Provider +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, provider.upstream_reply()) + model: Final = _deployment(scenario, handle, provider) + response: Final = _decide(gateway, model, input=_IMAGE_INPUT) + assert response.status_code == 400, response.text + assert "input_image" in response.text + assert _upstream_calls(gateway, handle) == [] + + +def test_openai_forwards_image_input_in_its_own_message_shape(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _OPENAI.upstream_reply()) + model: Final = _deployment(scenario, handle, _OPENAI) + response: Final = _decide(gateway, model, input=_IMAGE_INPUT) + assert response.status_code == 200, response.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == {"model": _OPENAI.body_model, "input": _OPENAI_IMAGE_MESSAGES, "questions": _QUESTIONS} + + +def test_a_message_list_input_reaches_a_system_one_provider_as_its_flattened_text(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + model: Final = _deployment(scenario, handle, _PERPLEXITY) + response: Final = _decide(gateway, model, input=_MESSAGE_LIST_INPUT) + assert response.status_code == 200, response.text + assert response.json() == _PERPLEXITY.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == { + "model": _PERPLEXITY.body_model, + "state": f"{_INPUT}\n\n{_FOLLOW_UP}", + "questions": _SYSTEM_ONE_QUESTIONS, + } + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["prompt_tokens"], row["completion_tokens"]) == ( + "success", + "adecisions", + _INPUT_TOKENS, + _OUTPUT_TOKENS, + ) + + +@pytest.mark.parametrize("provider", _SYSTEM_ONE_PROVIDERS, ids=_provider_id) +def test_safety_identifier_is_refused_by_system_one_providers_unless_the_deployment_drops_params( + gateway: Gateway, provider: _Provider +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, provider.upstream_reply()) + strict: Final = _deployment(scenario, handle, provider) + refused: Final = _decide(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER) + _assert_refused_for_safety_identifier(refused) + assert _upstream_calls(gateway, handle) == [] + + dropping: Final = _deployment(scenario, handle, provider, drop_params=True) + accepted: Final = _decide(gateway, dropping, safety_identifier=_SAFETY_IDENTIFIER) + assert accepted.status_code == 200, accepted.text + assert accepted.json() == provider.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == provider.upstream_body() + row: Final = _spend_row(accepted.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["model_group"]) == ("success", "adecisions", dropping) + + +@pytest.mark.parametrize("value", ("", "x" * 5000), ids=("empty", "5kb")) +def test_every_string_safety_identifier_is_refused_or_dropped_like_the_usual_one(gateway: Gateway, value: str) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + strict: Final = _deployment(scenario, handle, _PERPLEXITY) + _assert_refused_for_safety_identifier(_decide(gateway, strict, safety_identifier=value)) + dropping: Final = _deployment(scenario, handle, _PERPLEXITY, drop_params=True) + accepted: Final = _decide(gateway, dropping, safety_identifier=value) + assert accepted.status_code == 200, accepted.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + + +def test_a_request_body_drop_params_drops_the_safety_identifier_like_chat(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + strict: Final = _deployment(scenario, handle, _PERPLEXITY) + response: Final = _decide(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER, drop_params=True) + assert response.status_code == 200, response.text + assert response.json() == _PERPLEXITY.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + + +def test_openai_keeps_the_safety_identifier_on_the_wire_without_drop_params(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _OPENAI.upstream_reply()) + model: Final = _deployment(scenario, handle, _OPENAI) + response: Final = _decide(gateway, model, safety_identifier=_SAFETY_IDENTIFIER) + assert response.status_code == 200, response.text + assert response.json() == _OPENAI.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == {**_OPENAI.upstream_body(), "safety_identifier": _SAFETY_IDENTIFIER} + + +def test_a_system_one_format_safety_identifier_is_refused_unless_the_deployment_drops_params(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + strict: Final = _deployment(scenario, handle, _PERPLEXITY) + refused: Final = _decide_in_system_one_format(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER) + _assert_refused_for_safety_identifier(refused) + assert _upstream_calls(gateway, handle) == [] + + dropping: Final = _deployment(scenario, handle, _PERPLEXITY, drop_params=True) + accepted: Final = _decide_in_system_one_format(gateway, dropping, safety_identifier=_SAFETY_IDENTIFIER) + assert accepted.status_code == 200, accepted.text + assert accepted.json()["answers"] == _SYSTEM_ONE_ANSWERS + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() + + +def test_openai_keeps_a_system_one_format_safety_identifier_on_the_wire(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _OPENAI.upstream_reply()) + model: Final = _deployment(scenario, handle, _OPENAI) + response: Final = _decide_in_system_one_format(gateway, model, safety_identifier=_SAFETY_IDENTIFIER) + assert response.status_code == 200, response.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == {**_OPENAI.upstream_body(), "safety_identifier": _SAFETY_IDENTIFIER} + + +@pytest.mark.parametrize("decide", (_decide, _decide_in_system_one_format), ids=("openai_format", "system_one_format")) +@pytest.mark.parametrize("value", (7, [_SAFETY_IDENTIFIER]), ids=("numeric", "list")) +def test_a_non_string_safety_identifier_is_refused_as_invalid_unless_the_deployment_drops_params( + gateway: Gateway, value: JsonValue, decide: Callable[..., httpx.Response] +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _OPENAI.upstream_reply()) + strict: Final = _deployment(scenario, handle, _OPENAI) + refused: Final = decide(gateway, strict, safety_identifier=value) + assert refused.status_code == 400, refused.text + assert "Invalid Decisions request" in refused.text, refused.text + assert _upstream_calls(gateway, handle) == [] + + dropping: Final = _deployment(scenario, handle, _OPENAI, drop_params=True) + accepted: Final = decide(gateway, dropping, safety_identifier=value) + assert accepted.status_code == 200, accepted.text + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _OPENAI.upstream_body() + + +def test_litellm_settings_drop_params_drops_the_safety_identifier_for_a_strict_deployment( + gateway: Gateway, tmp_path: Path +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _PERPLEXITY.upstream_reply()) + config: Final = _drop_params_config(tmp_path, handle.api_base()) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + response: Final = _decide(owned.gateway, _DROPPING_DEPLOYMENT, safety_identifier=_SAFETY_IDENTIFIER) + assert response.status_code == 200, response.text + assert response.json() == _PERPLEXITY.litellm_response() + (call,) = _upstream_calls(gateway, handle) + assert call["body"] == _PERPLEXITY.upstream_body() diff --git a/tests/integration/translation/decisions/bases/cloudflare.py b/tests/integration/translation/decisions/bases/cloudflare.py index ae86a2ccd874..3efc96eeb878 100644 --- a/tests/integration/translation/decisions/bases/cloudflare.py +++ b/tests/integration/translation/decisions/bases/cloudflare.py @@ -6,6 +6,103 @@ """ CLEF_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "cloudflare/@cf/cloudflare/clef", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/ai/run/@cf/cloudflare/clef", + expected_provider_headers={"authorization": "Bearer synthetic-cloudflare-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "clef", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "result": { + "model": "clef", + "answers": { + "defect": {"type": "noul", "noul": 0.9345}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.8067, + "probabilities": {"low": 0.0509, "high": 0.9491}, + }, + "confidence": { + "type": "score", + "score": 0.9036, + "confidence": 0.6515, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.0964, "1": 0.9036}, + }, + }, + "usage": {"input_tokens": 290, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + }, + expected_litellm_response={ + "model": "clef", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.9345}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.0509}, {"value": "high", "probability": 0.9491}], + "confidence": 0.8067, + }, + { + "type": "score", + "name": "confidence", + "score": 0.9036, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.0964}, + {"value": 1, "label": "sure", "probability": 0.9036}, + ], + "confidence": 0.6515, + }, + ], + "usage": { + "input_tokens": 290, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 290, + }, + }, +) +CLEF_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "cloudflare/@cf/cloudflare/clef", diff --git a/tests/integration/translation/decisions/bases/openrouter.py b/tests/integration/translation/decisions/bases/openrouter.py index b65acff6b06e..ed39bc1a55be 100644 --- a/tests/integration/translation/decisions/bases/openrouter.py +++ b/tests/integration/translation/decisions/bases/openrouter.py @@ -6,6 +6,100 @@ """ TYPESAFE_JEV_1_13_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "openrouter/typesafe/jev-1.13", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/alpha/decisions", + expected_provider_headers={"authorization": "Bearer synthetic-openrouter-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "typesafe/jev-1.13", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "model": "typesafe/jev-1.13-20260917", + "answers": { + "defect": {"type": "noul", "noul": 0.81}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.99, + "probabilities": {"low": 0.01, "high": 0.99}, + }, + "confidence": { + "type": "score", + "score": 0.5, + "confidence": 0, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.5, "1": 0.5}, + }, + }, + "usage": {"input_tokens": 377, "output_tokens": 62, "cost": 1.5834e-05}, + "id": "gen-dec-1791323839-GA15kY0nt34oiJ7srfki", + "provider": "TypeSafe", + }, + expected_litellm_response={ + "model": "typesafe/jev-1.13-20260917", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.81}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.01}, {"value": "high", "probability": 0.99}], + "confidence": 0.99, + }, + { + "type": "score", + "name": "confidence", + "score": 0.5, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.5}, + {"value": 1, "label": "sure", "probability": 0.5}, + ], + "confidence": 0, + }, + ], + "usage": { + "input_tokens": 377, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 62, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 439, + }, + }, +) +TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "openrouter/typesafe/jev-1.13", diff --git a/tests/integration/translation/decisions/bases/perplexity.py b/tests/integration/translation/decisions/bases/perplexity.py index 02b1de51f242..9e29ceab8ccd 100644 --- a/tests/integration/translation/decisions/bases/perplexity.py +++ b/tests/integration/translation/decisions/bases/perplexity.py @@ -6,6 +6,101 @@ """ PPLX_DECIDER_V1_27B_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "perplexity/pplx-decider-v1-27b", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/decisions", + expected_provider_headers={"authorization": "Bearer synthetic-perplexity-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "pplx-decider-v1-27b", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "model": "pplx-decider-v1-27b", + "answers": { + "defect": {"type": "noul", "noul": 0.9989100737587077}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.9964631215356778, + "probabilities": {"low": 0.0017684392321610232, "high": 0.9982315607678389}, + }, + "confidence": { + "type": "score", + "score": 0.07367392327139817, + "confidence": 0.8526521534572037, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.9263260767286018, "1": 0.07367392327139817}, + }, + }, + "usage": {"input_tokens": 318, "output_tokens": 3}, + }, + expected_litellm_response={ + "model": "pplx-decider-v1-27b", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.9989100737587077}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [ + {"value": "low", "probability": 0.0017684392321610232}, + {"value": "high", "probability": 0.9982315607678389}, + ], + "confidence": 0.9964631215356778, + }, + { + "type": "score", + "name": "confidence", + "score": 0.07367392327139817, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.9263260767286018}, + {"value": 1, "label": "sure", "probability": 0.07367392327139817}, + ], + "confidence": 0.8526521534572037, + }, + ], + "usage": { + "input_tokens": 318, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 3, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 321, + }, + }, +) +PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "perplexity/pplx-decider-v1-27b", diff --git a/tests/integration/translation/decisions/bases/strands_decider.py b/tests/integration/translation/decisions/bases/strands_decider.py index 8a1a51f8c3e6..2950b0907377 100644 --- a/tests/integration/translation/decisions/bases/strands_decider.py +++ b/tests/integration/translation/decisions/bases/strands_decider.py @@ -6,6 +6,40 @@ """ STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "strands_decider/strands-decider-2B-hobson-v19", + "input": "Help! My payouts have been failing for 3 days!", + "questions": [{"type": "predicate", "name": "is_urgent", "instructions": "Does this convey urgency?"}], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/systemone", + expected_provider_headers={"content-type": "application/json"}, + expected_provider_request={ + "model": "strands-decider-2B-hobson-v19", + "state": "Help! My payouts have been failing for 3 days!", + "questions": {"is_urgent": {"type": "noul", "instructions": "Does this convey urgency?"}}, + }, + mock_provider_response={ + "model": "strands-decider-2B-hobson-v19", + "answers": {"is_urgent": {"type": "noul", "noul": 0.8277}}, + "usage": {"input_tokens": 86, "output_tokens": 1}, + "latency_ms": 140.03, + }, + expected_litellm_response={ + "model": "strands-decider-2B-hobson-v19", + "answers": [{"type": "predicate", "name": "is_urgent", "probability": 0.8277}], + "usage": { + "input_tokens": 86, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 1, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 87, + }, + }, +) +STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "strands_decider/strands-decider-2B-hobson-v19", diff --git a/tests/integration/translation/decisions/bases/typesafe.py b/tests/integration/translation/decisions/bases/typesafe.py index 5da5b05cb09d..7d47ce6b0827 100644 --- a/tests/integration/translation/decisions/bases/typesafe.py +++ b/tests/integration/translation/decisions/bases/typesafe.py @@ -6,6 +6,98 @@ """ JEV_1_13_0_TEST_CASE: Final = TranslationTestCase( scenario="basic", + litellm_endpoint="/v1/decisions", + litellm_request={ + "model": "typesafe/jev-1.13.0", + "input": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": [ + {"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}, + { + "type": "choice", + "name": "severity", + "instructions": "How severe is it?", + "choices": [ + {"value": "low", "description": "cosmetic"}, + {"value": "high", "description": "blocks users"}, + ], + }, + { + "type": "score", + "name": "confidence", + "instructions": "How sure are you?", + "levels": [{"label": "unsure"}, {"label": "sure"}], + }, + ], + "cache": {"no-cache": True}, + }, + expected_provider_endpoint="/v1/systemone", + expected_provider_headers={"authorization": "Bearer synthetic-typesafe-key", "content-type": "application/json"}, + expected_provider_request={ + "model": "jev-1.13.0", + "state": "Ticket (billing): The export job hangs at 99% and never finishes", + "questions": { + "defect": {"type": "noul", "instructions": "Is this a defect?"}, + "severity": { + "type": "choice", + "instructions": "How severe is it?", + "criteria": {"low": "cosmetic", "high": "blocks users"}, + }, + "confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]}, + }, + }, + mock_provider_response={ + "model": "jev-1.13.0", + "answers": { + "defect": {"type": "noul", "noul": 0.78}, + "severity": { + "type": "choice", + "choice": "high", + "confidence": 0.99, + "probabilities": {"low": 0.01, "high": 0.99}, + }, + "confidence": { + "type": "score", + "score": 0.51, + "confidence": 0.03, + "legend": {"0": "unsure", "1": "sure"}, + "probabilities": {"0": 0.49, "1": 0.51}, + }, + }, + "usage": {"input_tokens": 377, "output_tokens": 62}, + }, + expected_litellm_response={ + "model": "jev-1.13.0", + "answers": [ + {"type": "predicate", "name": "defect", "probability": 0.78}, + { + "type": "choice", + "name": "severity", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.01}, {"value": "high", "probability": 0.99}], + "confidence": 0.99, + }, + { + "type": "score", + "name": "confidence", + "score": 0.51, + "probabilities": [ + {"value": 0, "label": "unsure", "probability": 0.49}, + {"value": 1, "label": "sure", "probability": 0.51}, + ], + "confidence": 0.03, + }, + ], + "usage": { + "input_tokens": 377, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 62, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 439, + }, + }, +) +JEV_1_13_0_SYSTEMONE_TEST_CASE: Final = TranslationTestCase( + scenario="systemone", litellm_endpoint="/v1/systemone", litellm_request={ "model": "typesafe/jev-1.13.0", diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py b/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py index 79fc0574b3bb..21e0538897d5 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py @@ -2,10 +2,10 @@ from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.cloudflare import CLEF_TEST_CASE +from integration.translation.decisions.bases.cloudflare import CLEF_TEST_CASE, CLEF_SYSTEMONE_TEST_CASE from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [CLEF_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize("case", [CLEF_TEST_CASE, CLEF_SYSTEMONE_TEST_CASE], ids=lambda case: case.id) def test_decisions_basic_cloudflare(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py b/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py index a74cb23f6b43..973526d72729 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py @@ -2,10 +2,15 @@ from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.openrouter import TYPESAFE_JEV_1_13_TEST_CASE +from integration.translation.decisions.bases.openrouter import ( + TYPESAFE_JEV_1_13_TEST_CASE, + TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE, +) from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [TYPESAFE_JEV_1_13_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize( + "case", [TYPESAFE_JEV_1_13_TEST_CASE, TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE], ids=lambda case: case.id +) def test_decisions_basic_openrouter(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py b/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py index e341d24ba537..d9662ec4b9f8 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py @@ -2,10 +2,15 @@ from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.perplexity import PPLX_DECIDER_V1_27B_TEST_CASE +from integration.translation.decisions.bases.perplexity import ( + PPLX_DECIDER_V1_27B_TEST_CASE, + PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE, +) from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [PPLX_DECIDER_V1_27B_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize( + "case", [PPLX_DECIDER_V1_27B_TEST_CASE, PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE], ids=lambda case: case.id +) def test_decisions_basic_perplexity(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py b/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py index cae7ca70c780..2d2cd9caaf54 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py @@ -2,10 +2,17 @@ from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.strands_decider import STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE +from integration.translation.decisions.bases.strands_decider import ( + STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE, + STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE, +) from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize( + "case", + [STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE, STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE], + ids=lambda case: case.id, +) def test_decisions_basic_strands_decider(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py b/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py index 302636e6fb41..c1111cc6ba95 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py @@ -2,10 +2,10 @@ from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.decisions.bases.typesafe import JEV_1_13_0_TEST_CASE +from integration.translation.decisions.bases.typesafe import JEV_1_13_0_TEST_CASE, JEV_1_13_0_SYSTEMONE_TEST_CASE from integration.translation.runner import assert_translation -@pytest.mark.parametrize("case", [JEV_1_13_0_TEST_CASE], ids=lambda case: case.id) +@pytest.mark.parametrize("case", [JEV_1_13_0_TEST_CASE, JEV_1_13_0_SYSTEMONE_TEST_CASE], ids=lambda case: case.id) def test_decisions_basic_typesafe(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: assert_translation(case, gateway, provider) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 8f05a492119f..c1cac2d386f5 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -1262,6 +1262,141 @@ async def test_openrouter_decisions_uses_provider_reported_cost_without_cost_map assert get_response_cost_from_hidden_params(response.hidden_params) == cost +_SAFETY_IDENTIFIER: Final = "end-user-7" +_PREDICATE_QUESTIONS: Final[tuple[Mapping[str, object], ...]] = ( + {"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}, +) +_PREDICATE_REQUESTS: Final[tuple[Mapping[str, object], ...]] = ( + MappingProxyType({"input": "review", "questions": _PREDICATE_QUESTIONS}), + MappingProxyType( + {"state": "review", "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}} + ), +) +_PREDICATE_REQUEST_IDS: Final = ("input_format", "state_format") + + +@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +def test_safety_identifier_is_refused_before_http_when_the_provider_cannot_take_it( + request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) + + with pytest.raises(litellm.UnsupportedParamsError, match=r"safety_identifier.*drop_params") as caught: + litellm.decisions( + model="perplexity/pplx-decider-v1-27b", + safety_identifier=_SAFETY_IDENTIFIER, + api_key="caller-key", + **request_kwargs, + ) + + assert caught.value.status_code == 400 + assert not route.called + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ("call", "global")) +async def test_safety_identifier_is_dropped_from_the_wire_under_drop_params( + scope: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", scope == "global") + route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) + + response: Final = await litellm.adecisions( + model="perplexity/pplx-decider-v1-27b", + input="review", + questions=_PREDICATE_QUESTIONS, + safety_identifier=_SAFETY_IDENTIFIER, + api_key="caller-key", + **({"drop_params": True} if scope == "call" else {}), + ) + + assert route.called + assert json.loads(respx_mock.calls[0].request.content) == { + "model": "pplx-decider-v1-27b", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert isinstance(response, OpenAIDecisionResponse) + assert [answer.model_dump(mode="json") for answer in response.answers] == [ + {"type": "predicate", "name": "is_defect", "probability": 0.9} + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +@pytest.mark.parametrize("drop_params", (False, True), ids=("strict", "drop_params")) +async def test_safety_identifier_reaches_openai_whether_or_not_params_are_dropped( + drop_params: bool, + request_kwargs: Mapping[str, object], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_params) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + + await litellm.adecisions( + model="openai/gpt-6-luna", + safety_identifier=_SAFETY_IDENTIFIER, + api_key="caller-key", + **request_kwargs, + ) + + assert route.called + assert json.loads(respx_mock.calls[0].request.content) == { + "model": "gpt-6-luna", + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + "safety_identifier": _SAFETY_IDENTIFIER, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +async def test_non_string_safety_identifier_is_rejected_before_http( + request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + invalid_safety_identifier: Final[Mapping[str, object]] = MappingProxyType({"safety_identifier": 7}) + + with pytest.raises(litellm.BadRequestError, match=r"(?s)Invalid Decisions request.*safety_identifier"): + await litellm.adecisions( + model="openai/gpt-6-luna", api_key="caller-key", **request_kwargs, **invalid_safety_identifier + ) + + assert not route.called + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS) +async def test_non_string_safety_identifier_is_dropped_from_the_wire_under_drop_params( + request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + dropped_safety_identifier: Final[Mapping[str, object]] = MappingProxyType( + {"safety_identifier": 7, "drop_params": True} + ) + + await litellm.adecisions( + model="openai/gpt-6-luna", api_key="caller-key", **request_kwargs, **dropped_safety_identifier + ) + + assert route.called + assert json.loads(respx_mock.calls[0].request.content) == { + "model": "gpt-6-luna", + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + } + + @pytest.mark.asyncio @pytest.mark.parametrize( "api_base", diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index 6ce15636f1e2..ed78fc7063af 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -575,6 +575,68 @@ def test_openai_format_decisions_reach_an_openai_deployment_unchanged_including_ assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected_cost) +@pytest.mark.parametrize("safety_identifier", (7, ["end-user-1"]), ids=("numeric", "list")) +@pytest.mark.parametrize("deployment_drops_params", (False, True), ids=("strict", "drop_params")) +def test_a_non_string_safety_identifier_is_refused_unless_the_deployment_drops_params( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + safety_identifier: object, + deployment_drops_params: bool, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "decider", + "litellm_params": { + "model": "openai/gpt-6-luna", + "api_key": "k", + "drop_params": deployment_drops_params, + }, + } + ] + ), + ) + upstream: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond( + json={ + "model": "gpt-6-luna", + "answers": _OPENAI_FORMAT_ANSWERS[:1], + "usage": { + "input_tokens": _INPUT_TOKENS, + "output_tokens": _OUTPUT_TOKENS, + "total_tokens": _INPUT_TOKENS + _OUTPUT_TOKENS, + }, + } + ) + request_body: Final = { + "model": "decider", + "input": "The package arrived with a broken screen.", + "questions": _OPENAI_FORMAT_REQUEST["questions"][:1], + "safety_identifier": safety_identifier, + } + + response: Final = client.post("/v1/decisions", json=request_body) + + if not deployment_drops_params: + assert response.status_code == 400, response.text + assert "safety_identifier" in response.json()["error"]["message"] + assert not upstream.called + return + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == { + "model": "gpt-6-luna", + "input": "The package arrived with a broken screen.", + "questions": _OPENAI_FORMAT_REQUEST["questions"][:1], + } + + def _decisions_feature() -> LazyFeature: return next(feature for feature in LAZY_FEATURES if feature.name == "decisions") From a978acffb389eb924c71e6fc10f54dbbe7d2af3e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:29:33 -0700 Subject: [PATCH 06/18] feat(strands_decider): call a Strands Decider on an AgentCore runtime ARN (#45678) Backport of #45678 to rc/1.106.0. Cherry-picked from ed9c31cfa8574b9d7516443f43d1466de4053cf3 (main). --- .github/workflows/test-unit.yml | 1 + litellm/decisions/main.py | 4 + .../llms/base_llm/decisions/transformation.py | 14 + litellm/llms/custom_httpx/llm_http_handler.py | 38 +- .../decisions/transformation.py | 154 ++++ litellm/proxy/auth/auth_utils.py | 2 + tests/integration/_support/forward_proxy.py | 46 + .../test_strands_decider_agentcore_wire.py | 827 ++++++++++++++++++ tests/unit/decisions/test_main.py | 200 +++++ tests/unit/llms/strands_decider/__init__.py | 0 .../strands_decider/decisions/__init__.py | 0 ...trands_decider_decisions_transformation.py | 256 ++++++ tests/unit/proxy/auth/test_auth_utils.py | 16 + 13 files changed, 1555 insertions(+), 3 deletions(-) create mode 100644 tests/integration/providers/test_strands_decider_agentcore_wire.py create mode 100644 tests/unit/llms/strands_decider/__init__.py create mode 100644 tests/unit/llms/strands_decider/decisions/__init__.py create mode 100644 tests/unit/llms/strands_decider/decisions/test_strands_decider_decisions_transformation.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 6dee2bb605fc..803b5fba1324 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -343,6 +343,7 @@ jobs: tests/unit/llms/snowflake tests/unit/llms/soniox tests/unit/llms/stability + tests/unit/llms/strands_decider tests/unit/llms/tavily tests/unit/llms/tencent tests/unit/llms/tinyfish diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index f463bd58e3c7..5fd5a133eb29 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -60,6 +60,7 @@ class _DecisionsCall: logging_obj: LiteLLMLoggingObj | None headers: Mapping[str, str] timeout: float | httpx.Timeout | None + litellm_params: Mapping[str, object] = field(repr=False) def supported_decisions_providers() -> tuple[str, ...]: @@ -256,6 +257,7 @@ def _prepare_call( logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None, headers=extra_headers or {}, timeout=timeout, + litellm_params=kwargs, ) @@ -339,6 +341,7 @@ async def adecisions( api_key=call.api_key, headers=call.headers, timeout=call.timeout, + litellm_params=call.litellm_params, ) except Exception as error: raise _map_upstream_exception(error, call) from error @@ -384,6 +387,7 @@ def decisions( api_key=call.api_key, headers=call.headers, timeout=call.timeout, + litellm_params=call.litellm_params, ) except Exception as error: raise _map_upstream_exception(error, call) from error diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index baeeb3f04492..724ca23c2abd 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -341,6 +341,20 @@ def validate_environment(self, headers: Mapping[str, str], model: str, api_key: def get_complete_url(self, api_base: str, model: str) -> str: return f"{api_base.rstrip('/').removesuffix('/v1')}{self.path}" + def signs_with_aws(self, api_base: str) -> bool: + return False + + def sign_request( + self, + headers: Mapping[str, str], + url: str, + api_base: str, + body: Mapping[str, object], + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> tuple[Mapping[str, str], bytes | None]: + return headers, None + def transform_decisions_request( self, model: str, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a9b8fee2e8b..8fff7a7e98a5 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4,7 +4,7 @@ import ssl from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager -from functools import lru_cache +from functools import lru_cache, partial from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -1517,6 +1517,7 @@ def decisions( api_key: str | None, headers: Mapping[str, str], timeout: float | httpx.Timeout | None, + litellm_params: Mapping[str, object], client: HTTPHandler | None = None, ) -> DecisionsIRResponse: url, outbound_headers, data = self._prepare_decisions_request( @@ -1528,10 +1529,23 @@ def decisions( api_key=api_key, headers=headers, ) + signed_headers, signed_body = provider_config.sign_request( + headers=outbound_headers, + url=url, + api_base=api_base, + body=data, + api_key=api_key, + litellm_params=litellm_params, + ) sync_httpx_client: Final = client if client is not None else get_httpx_client() try: response: Final = sync_httpx_client.post( - url, json=data, headers=outbound_headers, timeout=timeout, logging_obj=logging_obj + url, + content=signed_body, + json=data if signed_body is None else None, + headers=dict(signed_headers), + timeout=timeout, + logging_obj=logging_obj, ) except httpx.HTTPError as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -1551,6 +1565,7 @@ async def adecisions( api_key: str | None, headers: Mapping[str, str], timeout: float | httpx.Timeout | None, + litellm_params: Mapping[str, object], client: AsyncHTTPHandler | None = None, ) -> DecisionsIRResponse: url, outbound_headers, data = self._prepare_decisions_request( @@ -1562,6 +1577,18 @@ async def adecisions( api_key=api_key, headers=headers, ) + sign: Final = partial( + provider_config.sign_request, + headers=outbound_headers, + url=url, + api_base=api_base, + body=data, + api_key=api_key, + litellm_params=litellm_params, + ) + signed_headers, signed_body = ( + await run_aws_signing(sign) if provider_config.signs_with_aws(api_base) else sign() + ) async_httpx_client: Final = ( client if client is not None @@ -1569,7 +1596,12 @@ async def adecisions( ) try: response: Final = await async_httpx_client.post( - url, json=data, headers=outbound_headers, timeout=timeout, logging_obj=logging_obj + url, + content=signed_body, + json=data if signed_body is None else None, + headers=dict(signed_headers), + timeout=timeout, + logging_obj=logging_obj, ) except httpx.HTTPError as e: raise self._handle_error(e=e, provider_config=provider_config) diff --git a/litellm/llms/strands_decider/decisions/transformation.py b/litellm/llms/strands_decider/decisions/transformation.py index 9ef966e18f2e..fd8a643be2a8 100644 --- a/litellm/llms/strands_decider/decisions/transformation.py +++ b/litellm/llms/strands_decider/decisions/transformation.py @@ -1,7 +1,161 @@ +import hashlib +import json +import re +from collections.abc import Mapping +from dataclasses import dataclass +from functools import partial +from types import MappingProxyType +from typing import Final +from urllib.parse import quote, unquote + +import httpx +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, sign_aws_json_post +from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse +from litellm.types.llms.bedrock import AwsAuthParams + +AGENTCORE_SESSION_HEADER: Final = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" +_RESERVED_HEADER_PREFIXES: Final = ("x-amzn-bedrock-agentcore-runtime-", "x-amz-") +_RUNTIME_ARN: Final = re.compile( + r"\Aarn:aws(?:-[a-z]+)*:bedrock-agentcore:(?P[a-z0-9-]+):[0-9]{12}:runtime/[A-Za-z0-9_-]+\Z" +) +_INVOCATIONS_PATH: Final = re.compile(r"\A/runtimes/(?P[^/]+)/invocations\Z") +_RUNTIME_ERROR_STATUS: Final = MappingProxyType({"bad_request": 400, "invalid_request": 400, "loading": 503}) + + +@dataclass(frozen=True, slots=True) +class AgentCoreRuntime: + arn: str + region: str + + @property + def invocations_url(self) -> str: + host: Final = f"bedrock-agentcore.{self.region}.{get_aws_dns_suffix(self.region)}" + return f"https://{host}/runtimes/{quote(self.arn, safe='')}/invocations" + + @property + def default_session_id(self) -> str: + return f"litellm-decider-{hashlib.sha256(self.arn.encode()).hexdigest()}" + + +def agentcore_runtime(api_base: str) -> AgentCoreRuntime | None: + match: Final = _RUNTIME_ARN.match(api_base) + return None if match is None else AgentCoreRuntime(arn=api_base, region=match["region"]) + + +def invoked_runtime(url: httpx.URL) -> AgentCoreRuntime | None: + match: Final = _INVOCATIONS_PATH.match(url.raw_path.decode("ascii")) + runtime: Final = None if match is None else agentcore_runtime(unquote(match["arn"])) + return runtime if runtime is not None and runtime.invocations_url == str(url) else None + + +class _RuntimeError(BaseModel): + model_config = ConfigDict(frozen=True) + + code: str + message: str = "" + + +class _RuntimeReply(BaseModel): + model_config = ConfigDict(frozen=True) + + error: _RuntimeError | None = None + answers: object = None + + +def _runtime_error(content: bytes) -> _RuntimeError | None: + try: + reply: Final = _RuntimeReply.model_validate_json(content) + except ValidationError: + return None + return reply.error if reply.answers is None else None + + +class _SessionParams(BaseModel): + model_config = ConfigDict(frozen=True) + + agentcore_runtime_session_id: str | None = None + + +def _is_reserved_header(name: str) -> bool: + lowered: Final = name.lower() + return lowered == "host" or lowered.startswith(_RESERVED_HEADER_PREFIXES) + + +def _runtime_headers( + headers: Mapping[str, str], runtime: AgentCoreRuntime, litellm_params: Mapping[str, object] +) -> Mapping[str, str]: + dropped: Final = sorted(name for name in headers if _is_reserved_header(name)) + if dropped: + verbose_logger.warning("Strands Decider: not forwarding reserved AgentCore header(s) %s", dropped) + configured: Final = _SessionParams.model_validate(litellm_params).agentcore_runtime_session_id + return { + **{name: value for name, value in headers.items() if not _is_reserved_header(name)}, + AGENTCORE_SESSION_HEADER: runtime.default_session_id if configured is None else configured, + } class StrandsDeciderDecisionsConfig(BaseDecisionsConfig): api_key_env = ("STRANDS_DECIDER_API_KEY",) api_base_env = ("STRANDS_DECIDER_API_BASE",) api_key_required = False + + def __init__(self, aws: BaseAWSLLM | None = None) -> None: + self._aws: Final = aws if aws is not None else BaseAWSLLM() + + def get_complete_url(self, api_base: str, model: str) -> str: + runtime: Final = agentcore_runtime(api_base) + return super().get_complete_url(api_base, model) if runtime is None else runtime.invocations_url + + def signs_with_aws(self, api_base: str) -> bool: + return agentcore_runtime(api_base) is not None + + def sign_request( + self, + headers: Mapping[str, str], + url: str, + api_base: str, + body: Mapping[str, object], + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> tuple[Mapping[str, str], bytes | None]: + runtime: Final = agentcore_runtime(api_base) + if runtime is None: + return headers, None + session_headers: Final = _runtime_headers(headers, runtime, litellm_params) + if api_key is not None: + return session_headers, None + payload: Final = json.dumps(body) + signed: Final = sign_aws_json_post( + get_credentials=partial( + self._aws.resolve_credentials, AwsAuthParams.model_validate(litellm_params), runtime.region + ), + service_name="bedrock-agentcore", + aws_region_name=runtime.region, + url=url, + body=payload, + headers=session_headers, + ) + return dict(signed.headers.items()), payload.encode() + + def transform_decisions_response( + self, + model: str, + custom_llm_provider: str, + raw_response: httpx.Response, + request: DecisionsIRRequest, + ) -> DecisionsIRResponse: + runtime_error: Final = ( + None if invoked_runtime(raw_response.request.url) is None else _runtime_error(raw_response.content) + ) + if runtime_error is not None: + raise BaseLLMException( + status_code=_RUNTIME_ERROR_STATUS.get(runtime_error.code, 500), + message=f"Strands Decider runtime error '{runtime_error.code}': {runtime_error.message}", + ) + return super().transform_decisions_response(model, custom_llm_provider, raw_response, request) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d31ce4dae034..28089ab9f77f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -411,6 +411,8 @@ def _build_banned_observability_params() -> frozenset[str]: "anthropic_workspace_id", "anthropic-workspace-id", "bedrock_tags", + # Deployment-pinned AgentCore session: every new caller-chosen value provisions another microVM. + "agentcore_runtime_session_id", # Provider-specific endpoint overrides that flow into the outbound # request via ``optional_params``. Same threat as ``api_base``: # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker diff --git a/tests/integration/_support/forward_proxy.py b/tests/integration/_support/forward_proxy.py index 9f5cff94277a..0d92c3d07194 100644 --- a/tests/integration/_support/forward_proxy.py +++ b/tests/integration/_support/forward_proxy.py @@ -1,5 +1,7 @@ from __future__ import annotations +import contextlib +import socket import threading from collections.abc import Generator, Mapping from contextlib import contextmanager @@ -69,3 +71,47 @@ def log_message(self, format: str, *args: object) -> None: thread.join(timeout=6) assert not thread.is_alive(), "Owned forward proxy survived cleanup" server.server_close() + + +def _pipe(source: socket.socket, sink: socket.socket) -> None: + with contextlib.suppress(OSError): + for chunk in iter(lambda: source.recv(65536), b""): + sink.sendall(chunk) + with contextlib.suppress(OSError): + sink.shutdown(socket.SHUT_WR) + + +@contextmanager +def tunnelling_forward_proxy(destination_port: int, port: int = 0) -> Generator[ForwardProxy, None, None]: + received: Final[SimpleQueue[Tunnel]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + rbufsize = 0 + timeout = 30 + request: socket.socket + + def do_CONNECT(self) -> None: + received.put(Tunnel(self.command, self.path, {name.lower(): value for name, value in self.headers.items()})) + self.close_connection = True + with socket.create_connection(("127.0.0.1", destination_port), timeout=30) as destination: + self.send_response(200, "Connection established") + self.end_headers() + outbound: Final = threading.Thread(target=_pipe, args=(self.request, destination)) + outbound.start() + _pipe(destination, self.request) + outbound.join(timeout=35) + + def log_message(self, format: str, *args: object) -> None: + pass + + with ThreadingHTTPServer(("127.0.0.1", port), Handler) as server: + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread.start() + try: + yield ForwardProxy(f"http://127.0.0.1:{server.server_port}", server.server_port, received) + finally: + server.shutdown() + thread.join(timeout=6) + assert not thread.is_alive(), "Owned forward proxy survived cleanup" + server.server_close() diff --git a/tests/integration/providers/test_strands_decider_agentcore_wire.py b/tests/integration/providers/test_strands_decider_agentcore_wire.py new file mode 100644 index 000000000000..bf5da4531378 --- /dev/null +++ b/tests/integration/providers/test_strands_decider_agentcore_wire.py @@ -0,0 +1,827 @@ +import hashlib +import json +import os +import re +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, replace +from functools import partial +from pathlib import Path +from typing import Final +from urllib.parse import quote, unquote + +import httpx +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.forward_proxy import ForwardProxy, tunnelling_forward_proxy +from integration._support.process import owned_proxy +from integration._support.sigv4 import encoded_path, signature +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.llms.strands_decider.decisions.transformation import StrandsDeciderDecisionsConfig + +_ACCOUNT: Final = "123456789012" +_ACCESS_KEY: Final = "AKIAAGENTCOREAUDIT01" +_SECRET_KEY: Final = "synthetic-agentcore-audit-secret" +_SIGNING: Final[Mapping[str, JsonValue]] = {"aws_access_key_id": _ACCESS_KEY, "aws_secret_access_key": _SECRET_KEY} +_JWT: Final = "synthetic-agentcore-jwt" +_ENV_JWT: Final = "synthetic-agentcore-env-jwt" +_US_EAST: Final = "us-east-1" +_EU_WEST: Final = "eu-west-1" +_MODEL: Final = "strands_decider/systemone-decider" +_BODY_MODEL: Final = "systemone-decider" +_SESSION_HEADER: Final = "x-amzn-bedrock-agentcore-runtime-session-id" +_CONFIGURED_SESSION: Final = "agentcore-audit-configured-session-0000000001" +_STATE: Final[dict[str, JsonValue]] = {"ticket": "Refund requested twice for one order", "queue": "payments"} +_QUESTIONS: Final[dict[str, JsonValue]] = { + "duplicate": {"type": "noul", "instructions": "Is this a duplicate refund request?"}, + "urgency": {"type": "choice", "criteria": {"low": "can wait", "high": "money at risk"}}, +} +_SDK_QUESTIONS: Final = TypeAdapter(dict[str, dict[str, object]]).validate_python(_QUESTIONS) +_ANSWERS: Final[dict[str, JsonValue]] = { + "duplicate": {"type": "noul", "noul": 0.88}, + "urgency": {"type": "choice", "choice": "high", "confidence": 0.75, "probabilities": {"low": 0.25, "high": 0.75}}, +} +_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 211, "output_tokens": 2} +_UPSTREAM_BODY: Final[dict[str, JsonValue]] = {"model": _BODY_MODEL, "state": _STATE, "questions": _QUESTIONS} +_INPUT: Final = "Refund requested twice for one order" +_OPENAI_QUESTIONS: Final[list[JsonValue]] = [ + {"type": "predicate", "name": "duplicate", "instructions": "Is this a duplicate refund request?"}, + { + "type": "choice", + "name": "urgency", + "instructions": "How urgent is it?", + "choices": [{"value": "low", "description": "can wait"}, {"value": "high", "description": "money at risk"}], + }, +] +_OPENAI_UPSTREAM_BODY: Final[dict[str, JsonValue]] = { + "model": _BODY_MODEL, + "state": _INPUT, + "questions": { + "duplicate": {"type": "noul", "instructions": "Is this a duplicate refund request?"}, + "urgency": { + "type": "choice", + "instructions": "How urgent is it?", + "criteria": {"low": "can wait", "high": "money at risk"}, + }, + }, +} +_OPENAI_ANSWER: Final[dict[str, JsonValue]] = { + "model": _BODY_MODEL, + "answers": [ + {"type": "predicate", "name": "duplicate", "probability": 0.88}, + { + "type": "choice", + "name": "urgency", + "choice": "high", + "probabilities": [{"value": "low", "probability": 0.25}, {"value": "high", "probability": 0.75}], + "confidence": 0.75, + }, + ], + "usage": { + "input_tokens": 211, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 2, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 213, + }, +} + + +@dataclass(frozen=True, slots=True) +class _Route: + path: str + body: Mapping[str, JsonValue] + upstream: Mapping[str, JsonValue] + answer: Mapping[str, JsonValue] + + +_SYSTEMONE: Final = _Route( + "/v1/systemone", + {"state": _STATE, "questions": _QUESTIONS}, + _UPSTREAM_BODY, + {"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE}, +) +_OPENAI: Final = _Route( + "/v1/decisions", {"input": _INPUT, "questions": _OPENAI_QUESTIONS}, _OPENAI_UPSTREAM_BODY, _OPENAI_ANSWER +) +_ROUTES: Final = (_SYSTEMONE, _OPENAI, replace(_SYSTEMONE, path="/systemone"), replace(_OPENAI, path="/decisions")) +_SPEND_QUERY: Final = ( + "SELECT status, call_type, custom_llm_provider, model_group, api_base, prompt_tokens, completion_tokens " + 'FROM "LiteLLM_SpendLogs" WHERE request_id = %s' +) +_RUNTIME_TARGET: Final = re.compile( + r"\A/runtimes/arn:aws:bedrock-agentcore:[a-z0-9-]+:[0-9]{12}:runtime/(?P[A-Za-z0-9_-]+)/invocations\Z" +) +_ERROR_STATUS: Final = {"bad_request": 400, "invalid_request": 400, "loading": 503, "inference_error": 500} +_HTTP_ERRORS: Final = { + 403: "AccessDeniedException", + 404: "ResourceNotFoundException", + 429: "ThrottlingException", +} +_MALFORMED_SESSIONS: Final[Mapping[str, JsonValue]] = {"int": 1234567, "list": ["a" * 40], "empty": ""} +_LONG_SESSION: Final = "s" * 5120 +_NEAR_MISS_ARNS: Final = { + "endpoint": f"arn:aws:bedrock-agentcore:us-east-1:{_ACCOUNT}:runtime/audit_qualified/runtime-endpoint/DEFAULT", + "account": "arn:aws:bedrock-agentcore:us-east-1:12345:runtime/audit_short_account", + "service": f"arn:aws:bedrock:us-east-1:{_ACCOUNT}:runtime/audit_other_service", + "host": f"arn:aws:bedrock-agentcore:us-east-1.example.com:{_ACCOUNT}:runtime/audit_crafted_host", +} +_DEPLOYMENT_HEADERS: Final[dict[str, JsonValue]] = { + "x-amz-target": "AmazonBedrockAgentCore.Forged", + "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "deployment-forged-session-00000000000001", + "Host": "bedrock-agentcore.us-west-2.amazonaws.com", + "x-audit-trace": "deployment-trace", +} +_SIGNED: Final = "agentcore-audit-signed" +_REGION: Final = "agentcore-audit-region" +_SESSION: Final = "agentcore-audit-session" +_LONG: Final = "agentcore-audit-session-long" +_BEARER: Final = "agentcore-audit-bearer" +_HEADERS: Final = "agentcore-audit-headers" +_OPT_IN: Final = "agentcore-audit-opt-in" +_HEALTH: Final = "agentcore-audit-health" +_DROPPED: Final = "agentcore-audit-dropped" +_OUTAGE: Final = "agentcore-audit-outage" +_ENV: Final = "agentcore-audit-env" +_UNSIGNED: Final = "agentcore-audit-unsigned" +_SCRUBBED_PREFIXES: Final = ("AWS_", "STRANDS_DECIDER_") +_SCRUBBED_NAMES: Final = frozenset( + {"HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY", "SSL_VERIFY", "SSL_CERT_FILE", "REQUESTS_CA_BUNDLE"} +) +_load_yaml: Final[Callable[[str], object]] = yaml.safe_load +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_OWNED_PROXY_CELL_SECONDS: Final = 2 * max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS") or 70)) + 120 +pytestmark: Final = pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) + + +def _host(region: str) -> str: + return f"bedrock-agentcore.{region}.amazonaws.com" + + +def _arn(runtime: str, region: str = _US_EAST) -> str: + return f"arn:aws:bedrock-agentcore:{region}:{_ACCOUNT}:runtime/{runtime}" + + +def _invocations_url(runtime: str, region: str = _US_EAST) -> str: + return f"https://{_host(region)}/runtimes/{quote(_arn(runtime, region), safe='')}/invocations" + + +def _default_session(runtime: str, region: str = _US_EAST) -> str: + return f"litellm-decider-{hashlib.sha256(_arn(runtime, region).encode()).hexdigest()}" + + +def _runtime_of(request: Request) -> str | None: + match: Final = _RUNTIME_TARGET.match(unquote(request.target)) + return None if match is None else match["runtime"] + + +def _json_reply(body: Mapping[str, JsonValue], status: int = 200, **headers: str) -> Reply: + return Reply(status=status, body=json.dumps(body).encode(), headers=headers) + + +def _respond(outage: threading.Event, request: Request) -> Reply: + runtime: Final = _runtime_of(request) or "" + if runtime == "audit_dropped" or (runtime == "audit_outage" and outage.is_set()): + return Reply(drop_connection=True) + if runtime.startswith("audit_error_"): + code: Final = runtime.removeprefix("audit_error_") + return _json_reply({"error": {"code": code, "message": f"runtime refused with {code}"}}) + if runtime.startswith("audit_status_"): + status: Final = int(runtime.removeprefix("audit_status_")) + return _json_reply( + {"message": f"scripted {_HTTP_ERRORS[status]}"}, status, **{"x-amzn-errortype": _HTTP_ERRORS[status]} + ) + if not runtime: + return _json_reply({"message": f"unknown target {request.target}"}, 404) + return _json_reply({"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE}) + + +@dataclass(frozen=True, slots=True) +class _Runtime: + wire: Wire + tunnel: ForwardProxy + cert: Path + outage: threading.Event + + def reset(self) -> None: + self.outage.clear() + self.wire.drain() + self.tunnel.drain() + + def calls(self) -> tuple[Request, ...]: + return self.wire.drain() + + +def _calls_to(calls: Sequence[Request], runtime: str) -> tuple[Request, ...]: + return tuple(call for call in calls if _runtime_of(call) == runtime) + + +def _deployment(name: str, **litellm_params: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": _MODEL, **litellm_params}, + "model_info": {"id": name, "mode": "evaluation"}, + } + + +def _signed_deployment(name: str, runtime: str, region: str = _US_EAST, **extra: JsonValue) -> dict[str, JsonValue]: + return _deployment(name, api_base=_arn(runtime, region), **_SIGNING, **extra) + + +def _deployments() -> list[dict[str, JsonValue]]: + return [ + _signed_deployment(_SIGNED, "audit_signed", aws_region_name=_US_EAST), + _signed_deployment(_REGION, "audit_region", _EU_WEST, aws_region_name=_US_EAST), + _signed_deployment(_SESSION, "audit_session", agentcore_runtime_session_id=_CONFIGURED_SESSION), + _signed_deployment(_LONG, "audit_session_long", agentcore_runtime_session_id=_LONG_SESSION), + _deployment(_BEARER, api_base=_arn("audit_bearer"), api_key=_JWT), + _signed_deployment(_HEADERS, "audit_headers", extra_headers=_DEPLOYMENT_HEADERS), + _signed_deployment( + _OPT_IN, "audit_opt_in", configurable_clientside_auth_params=["agentcore_runtime_session_id"] + ), + _signed_deployment(_HEALTH, "audit_health"), + _signed_deployment(_DROPPED, "audit_dropped"), + _signed_deployment(_OUTAGE, "audit_outage"), + _deployment(_ENV, **_SIGNING), + _deployment(_UNSIGNED, api_base=_arn("audit_unsigned")), + *(_signed_deployment(f"agentcore-audit-error-{code}", f"audit_error_{code}") for code in _ERROR_STATUS), + *(_signed_deployment(f"agentcore-audit-status-{status}", f"audit_status_{status}") for status in _HTTP_ERRORS), + *( + _signed_deployment( + f"agentcore-audit-session-{label}", f"audit_session_{label}", agentcore_runtime_session_id=value + ) + for label, value in _MALFORMED_SESSIONS.items() + ), + *( + _deployment(f"agentcore-audit-near-miss-{label}", api_base=value, **_SIGNING) + for label, value in _NEAR_MISS_ARNS.items() + ), + ] + + +def _write_config(directory: Path) -> Path: + base: Final = JSON_OBJECT.validate_python(_load_yaml(Path("tests/integration/proxy_config.yaml").read_text())) + configuration: Final = { + **base, + "model_list": _deployments(), + "router_settings": {"disable_cooldowns": True, "num_retries": 0}, + } + path: Final = directory / "agentcore.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +def _scrubbed() -> tuple[str, ...]: + return tuple( + name for name in os.environ if name.upper().startswith(_SCRUBBED_PREFIXES) or name.upper() in _SCRUBBED_NAMES + ) + + +def _isolated_aws(directory: Path) -> dict[str, str]: + return { + "AWS_EC2_METADATA_DISABLED": "true", + "AWS_SHARED_CREDENTIALS_FILE": str(directory / "absent-credentials"), + "AWS_CONFIG_FILE": str(directory / "absent-config"), + } + + +def _proxy_environment(runtime: _Runtime, directory: Path) -> dict[str, str]: + return { + "HTTPS_PROXY": runtime.tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(runtime.cert), + "STRANDS_DECIDER_API_BASE": _arn("audit_env"), + **_isolated_aws(directory), + } + + +@pytest.fixture(scope="module") +def module_gateway() -> Iterator[Gateway]: + with gateway_from_environment() as value: + yield value + + +@pytest.fixture(scope="module") +def runtime(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Runtime]: + cert, key = write_self_signed_cert(tmp_path_factory.mktemp("agentcore-cert"), (_host(_US_EAST), _host(_EU_WEST))) + outage: Final = threading.Event() + with wire_server(partial(_respond, outage), tls=server_context(cert, key)) as wire: + with tunnelling_forward_proxy(int(wire.url.rsplit(":", 1)[1])) as tunnel: + yield _Runtime(wire, tunnel, cert, outage) + + +@pytest.fixture(scope="module") +def agentcore( + module_gateway: Gateway, runtime: _Runtime, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("agentcore-proxy") + with owned_proxy( + module_gateway, + directory, + _proxy_environment(runtime, directory), + config=_write_config(directory), + remove_environment=_scrubbed(), + workers=2, + ) as candidate: + yield candidate + + +def _decide(gateway: Gateway, model: str, route: _Route = _SYSTEMONE, **extra: JsonValue) -> httpx.Response: + return gateway.request("POST", route.path, {"model": model, **route.body, **extra}) + + +def _authorization_fields(request: Request) -> dict[str, str]: + authorization: Final = request.headers.get("authorization", "") + assert authorization.startswith("AWS4-HMAC-SHA256 "), dict(request.headers) + return dict(part.split("=", 1) for part in authorization.removeprefix("AWS4-HMAC-SHA256 ").split(", ")) + + +def _assert_signed(request: Request, region: str = _US_EAST) -> None: + fields: Final = _authorization_fields(request) + access_key, scope = fields["Credential"].split("/", 1) + assert access_key == _ACCESS_KEY, fields + assert scope == f"{request.headers['x-amz-date'][:8]}/{region}/bedrock-agentcore/aws4_request", fields + signed: Final = fields["SignedHeaders"] + assert {"host", "x-amz-date", _SESSION_HEADER}.issubset(signed.split(";")), signed + expected: Final = signature( + "POST", encoded_path(request.target), request.headers, signed, request.body, _SECRET_KEY, scope + ) + assert fields["Signature"] == expected[1], (request.target, fields) + + +def _assert_answered(response: httpx.Response, model: str, route: _Route = _SYSTEMONE) -> None: + assert response.status_code == 200, response.text + assert response.json() == route.answer, response.text + assert response.headers["x-litellm-model-group"] == model, dict(response.headers) + + +def _assert_runtime_target(request: Request, runtime: str, region: str = _US_EAST) -> None: + assert request.method == "POST", request.method + assert request.target == f"/runtimes/{quote(_arn(runtime, region), safe='')}/invocations", request.target + assert request.headers["host"] == _host(region), dict(request.headers) + + +def _assert_runtime_call(request: Request, runtime: str, region: str = _US_EAST, route: _Route = _SYSTEMONE) -> None: + _assert_runtime_target(request, runtime, region) + assert json.loads(request.body) == route.upstream, request.body + + +def _canonical(body: JsonValue | Mapping[str, JsonValue]) -> str: + return json.dumps(body, sort_keys=True) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _spend_rows(call_ids: Sequence[str]) -> list[dict[str, JsonValue]]: + query: Final = ( + "SELECT request_id, status FROM \"LiteLLM_SpendLogs\" WHERE request_id = ANY(string_to_array(%s, ','))" + ) + return eventually( + lambda: read_rows(query, (",".join(call_ids),)), lambda found: len(found) >= len(call_ids), seconds=70 + ) + + +def _connect_targets(runtime: _Runtime) -> frozenset[str]: + return frozenset(runtime.tunnel.targets()) + + +def test_an_arn_deployment_answers_on_every_decisions_route_signed_for_bedrock_agentcore( + agentcore: Gateway, runtime: _Runtime +) -> None: + runtime.reset() + responses: Final = tuple(_decide(agentcore, _SIGNED, route) for route in _ROUTES) + for response, route in zip(responses, _ROUTES, strict=True): + _assert_answered(response, _SIGNED, route) + assert _connect_targets(runtime) == {f"{_host(_US_EAST)}:443"} + calls: Final = runtime.calls() + assert len(calls) == len(_ROUTES), calls + for call, route in zip(calls, _ROUTES, strict=True): + _assert_runtime_call(call, "audit_signed", route=route) + assert call.headers[_SESSION_HEADER] == _default_session("audit_signed"), dict(call.headers) + _assert_signed(call) + for response in responses: + row = _spend_row(response.headers["x-litellm-call-id"]) + assert row == { + "status": "success", + "call_type": "adecisions", + "custom_llm_provider": "strands_decider", + "model_group": _SIGNED, + "api_base": _invocations_url("audit_signed"), + "prompt_tokens": 211, + "completion_tokens": 2, + }, row + + +def test_the_arn_region_picks_the_host_and_signing_scope_over_aws_region_name( + agentcore: Gateway, runtime: _Runtime +) -> None: + runtime.reset() + _assert_answered(_decide(agentcore, _REGION), _REGION) + assert _connect_targets(runtime) == {f"{_host(_EU_WEST)}:443"} + (call,) = runtime.calls() + _assert_runtime_call(call, "audit_region", _EU_WEST) + _assert_signed(call, _EU_WEST) + + +def test_a_configured_session_id_is_sent_as_written_and_signed(agentcore: Gateway, runtime: _Runtime) -> None: + runtime.reset() + _assert_answered(_decide(agentcore, _SESSION), _SESSION) + _assert_answered(_decide(agentcore, _LONG), _LONG) + calls: Final = runtime.calls() + (configured,) = _calls_to(calls, "audit_session") + (long,) = _calls_to(calls, "audit_session_long") + assert configured.headers[_SESSION_HEADER] == _CONFIGURED_SESSION, dict(configured.headers) + assert long.headers[_SESSION_HEADER] == _LONG_SESSION + _assert_signed(configured) + _assert_signed(long) + + +def test_a_deployment_api_key_sends_bearer_and_the_session_without_sigv4(agentcore: Gateway, runtime: _Runtime) -> None: + runtime.reset() + _assert_answered(_decide(agentcore, _BEARER), _BEARER) + (call,) = runtime.calls() + _assert_runtime_call(call, "audit_bearer") + assert call.headers["authorization"] == f"Bearer {_JWT}", dict(call.headers) + assert call.headers[_SESSION_HEADER] == _default_session("audit_bearer"), dict(call.headers) + assert "x-amz-date" not in call.headers, dict(call.headers) + + +def test_reserved_extra_headers_are_dropped_and_the_rest_are_kept(agentcore: Gateway, runtime: _Runtime) -> None: + runtime.reset() + _assert_answered(_decide(agentcore, _HEADERS), _HEADERS) + caller_headers: Final[dict[str, JsonValue]] = { + "x-amz-security-token": "caller-forged-token", + "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "caller-forged-session-0000000000000000001", + "x-audit-caller": "caller-trace", + } + _assert_answered(_decide(agentcore, _HEADERS, extra_headers=caller_headers), _HEADERS) + deployment, caller = runtime.calls() + for call in (deployment, caller): + _assert_runtime_call(call, "audit_headers") + assert call.headers[_SESSION_HEADER] == _default_session("audit_headers"), dict(call.headers) + assert "x-amz-target" not in call.headers, dict(call.headers) + assert "x-amz-security-token" not in call.headers, dict(call.headers) + _assert_signed(call) + assert deployment.headers["x-audit-trace"] == "deployment-trace", dict(deployment.headers) + assert caller.headers["x-audit-caller"] == "caller-trace", dict(caller.headers) + + +def test_an_opted_in_caller_session_id_becomes_the_header_and_stays_out_of_the_body( + agentcore: Gateway, runtime: _Runtime +) -> None: + runtime.reset() + chosen: Final = f"agentcore-audit-caller-session-{uuid.uuid4().hex}" + _assert_answered(_decide(agentcore, _OPT_IN, agentcore_runtime_session_id=chosen), _OPT_IN) + _assert_answered(_decide(agentcore, _OPT_IN), _OPT_IN) + first, second = runtime.calls() + for call in (first, second): + _assert_runtime_call(call, "audit_opt_in") + _assert_signed(call) + assert first.headers[_SESSION_HEADER] == chosen, dict(first.headers) + assert second.headers[_SESSION_HEADER] == _default_session("audit_opt_in"), dict(second.headers) + + +def test_an_arn_deployment_without_api_base_takes_the_runtime_from_the_environment( + agentcore: Gateway, runtime: _Runtime +) -> None: + runtime.reset() + response: Final = _decide(agentcore, _ENV) + _assert_answered(response, _ENV) + (call,) = runtime.calls() + _assert_runtime_call(call, "audit_env") + _assert_signed(call) + assert _spend_row(response.headers["x-litellm-call-id"])["api_base"] == _invocations_url("audit_env") + + +def test_an_arn_deployment_added_through_the_management_api_answers_signed( + agentcore: Gateway, runtime: _Runtime +) -> None: + runtime.reset() + with agentcore.scenario() as scenario: + model: Final = scenario.model( + model_info=None, + model=_MODEL, + api_base=_arn("audit_stored"), + api_key=None, + aws_access_key_id=_ACCESS_KEY, + aws_secret_access_key=_SECRET_KEY, + ) + _assert_answered(_decide(agentcore, model), model) + (call,) = runtime.calls() + _assert_runtime_call(call, "audit_stored") + _assert_signed(call) + + +def test_health_and_test_connection_probe_the_runtime_signed(agentcore: Gateway, runtime: _Runtime) -> None: + runtime.reset() + health: Final = agentcore.request("GET", "/health", params={"model": _HEALTH}) + assert health.status_code == 200, health.text + assert (health.json()["healthy_count"], health.json()["unhealthy_count"]) == (1, 0), health.text + (probe,) = runtime.calls() + assert probe.target == f"/runtimes/{quote(_arn('audit_health'), safe='')}/invocations", probe.target + _assert_signed(probe) + connection: Final = agentcore.request( + "POST", + "/health/test_connection", + { + "litellm_params": {"model": _MODEL, "api_base": _arn("audit_test_connection"), **_SIGNING}, + "mode": "evaluation", + }, + ) + assert connection.status_code == 200, connection.text + assert connection.json()["status"] == "success", connection.text + (tested,) = runtime.calls() + assert _runtime_of(tested) == "audit_test_connection", tested.target + _assert_signed(tested) + + +@pytest.mark.parametrize(("code", "status"), tuple(_ERROR_STATUS.items()), ids=tuple(_ERROR_STATUS)) +def test_a_runtime_error_reply_maps_to_its_status_with_the_runtime_message( + agentcore: Gateway, runtime: _Runtime, code: str, status: int +) -> None: + runtime.reset() + model: Final = f"agentcore-audit-error-{code}" + response: Final = _decide(agentcore, model) + assert response.status_code == status, response.text + assert f"Strands Decider runtime error '{code}': runtime refused with {code}" in response.text, response.text + assert len(_calls_to(runtime.calls(), f"audit_error_{code}")) == 1 + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["model_group"]) == ("failure", model), row + + +@pytest.mark.parametrize("status", tuple(_HTTP_ERRORS), ids=tuple(str(status) for status in _HTTP_ERRORS)) +def test_a_runtime_http_error_keeps_its_status_and_message(agentcore: Gateway, runtime: _Runtime, status: int) -> None: + runtime.reset() + model: Final = f"agentcore-audit-status-{status}" + response: Final = _decide(agentcore, model) + assert response.status_code == status, response.text + assert f"scripted {_HTTP_ERRORS[status]}" in response.text, response.text + assert len(_calls_to(runtime.calls(), f"audit_status_{status}")) == 1 + assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" + + +@pytest.mark.parametrize("label", tuple(_MALFORMED_SESSIONS)) +def test_a_malformed_configured_session_id_fails_only_its_own_deployment( + agentcore: Gateway, runtime: _Runtime, label: str +) -> None: + runtime.reset() + model: Final = f"agentcore-audit-session-{label}" + response: Final = _decide(agentcore, model) + calls: Final = _calls_to(runtime.calls(), f"audit_session_{label}") + if label == "empty": + _assert_answered(response, model) + (call,) = calls + assert call.headers[_SESSION_HEADER] == "", dict(call.headers) + else: + assert response.status_code == 500, response.text + assert "agentcore_runtime_session_id" in response.text, response.text + assert calls == () + _assert_answered(_decide(agentcore, _SIGNED), _SIGNED) + + +def test_a_deployment_without_aws_credentials_fails_before_dialing_the_runtime( + agentcore: Gateway, runtime: _Runtime +) -> None: + runtime.reset() + response: Final = _decide(agentcore, _UNSIGNED) + assert response.status_code == 500, response.text + assert "credentials" in response.text.lower(), response.text + assert runtime.calls() == () + assert _connect_targets(runtime) == frozenset() + _assert_answered(_decide(agentcore, _SIGNED), _SIGNED) + + +@pytest.mark.parametrize("label", tuple(_NEAR_MISS_ARNS)) +def test_a_near_miss_arn_keeps_the_plain_api_base_path(agentcore: Gateway, runtime: _Runtime, label: str) -> None: + runtime.reset() + response: Final = _decide(agentcore, f"agentcore-audit-near-miss-{label}") + assert 500 <= response.status_code < 600, response.text + assert "litellm.APIConnectionError" in response.text, response.text + assert _connect_targets(runtime) == frozenset() + assert runtime.calls() == () + + +def test_a_dropped_connection_is_a_connection_error_after_one_attempt(agentcore: Gateway, runtime: _Runtime) -> None: + runtime.reset() + response: Final = _decide(agentcore, _DROPPED) + assert response.status_code == 500, response.text + assert "litellm.APIConnectionError" in response.text, response.text + assert len(_calls_to(runtime.calls(), "audit_dropped")) == 1 + assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" + + +def _burst_routes(size: int) -> tuple[_Route, ...]: + return tuple(_ROUTES[index % len(_ROUTES)] for index in range(size)) + + +def _burst(gateway: Gateway, model: str, size: int) -> tuple[httpx.Response, ...]: + with ThreadPoolExecutor(max_workers=size) as pool: + return tuple(pool.map(partial(_decide, gateway, model), _burst_routes(size))) + + +def _health_counts(gateway: Gateway, model: str) -> tuple[int, int, int]: + response: Final = gateway.request("GET", "/health", params={"model": model}) + body: Final = JSON_OBJECT.validate_json(response.content) + healthy: Final = body["healthy_count"] + unhealthy: Final = body["unhealthy_count"] + assert isinstance(healthy, int) and isinstance(unhealthy, int), body + return response.status_code, healthy, unhealthy + + +def test_a_runtime_outage_mid_burst_fails_each_call_once_and_recovers_signed( + agentcore: Gateway, runtime: _Runtime +) -> None: + runtime.reset() + runtime.outage.set() + failed: Final = _burst(agentcore, _OUTAGE, 10) + assert [response.status_code for response in failed] == [500] * 10, [response.text for response in failed] + assert _health_counts(agentcore, _OUTAGE) == (503, 0, 1) + runtime.outage.clear() + runtime.calls() + served: Final = _burst(agentcore, _OUTAGE, 20) + for response, route in zip(served, _burst_routes(20), strict=True): + _assert_answered(response, _OUTAGE, route) + calls: Final = runtime.calls() + assert len(calls) == 20, calls + for call in calls: + _assert_runtime_target(call, "audit_outage") + _assert_signed(call) + assert sorted(_canonical(_JSON.validate_json(call.body)) for call in calls) == sorted( + _canonical(route.upstream) for route in _burst_routes(20) + ) + assert _health_counts(agentcore, _OUTAGE) == (200, 1, 0) + expected: Final = { + **{response.headers["x-litellm-call-id"]: "failure" for response in failed}, + **{response.headers["x-litellm-call-id"]: "success" for response in served}, + } + assert len(expected) == 30, expected + rows: Final = _spend_rows(tuple(expected)) + assert sorted((str(row["request_id"]), row["status"]) for row in rows) == sorted(expected.items()), rows + + +async def test_sdk_sync_and_async_calls_reach_the_runtime_signed( + runtime: _Runtime, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + runtime.reset() + for name in _scrubbed(): + monkeypatch.delenv(name) + for name, value in { + "HTTPS_PROXY": runtime.tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(runtime.cert), + **_isolated_aws(tmp_path), + }.items(): + monkeypatch.setenv(name, value) + litellm.in_memory_llm_clients_cache.flush_cache() + try: + synchronous: Final = litellm.decisions( + model=_MODEL, + state=_STATE, + questions=_SDK_QUESTIONS, + api_base=_arn("audit_sdk"), + aws_access_key_id=_ACCESS_KEY, + aws_secret_access_key=_SECRET_KEY, + ) + asynchronous: Final = await litellm.adecisions( + model=_MODEL, + state=_STATE, + questions=_SDK_QUESTIONS, + api_base=_arn("audit_sdk"), + aws_access_key_id=_ACCESS_KEY, + aws_secret_access_key=_SECRET_KEY, + ) + monkeypatch.setenv("STRANDS_DECIDER_API_KEY", _ENV_JWT) + bearer: Final = await litellm.adecisions( + model=_MODEL, state=_STATE, questions=_SDK_QUESTIONS, api_base=_arn("audit_sdk_env_key") + ) + finally: + litellm.in_memory_llm_clients_cache.flush_cache() + for response in (synchronous, asynchronous, bearer): + assert response.model_dump(mode="json") == {"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE} + calls: Final = runtime.calls() + signed: Final = _calls_to(calls, "audit_sdk") + assert len(signed) == 2, calls + for call in signed: + _assert_runtime_call(call, "audit_sdk") + assert call.headers[_SESSION_HEADER] == _default_session("audit_sdk"), dict(call.headers) + _assert_signed(call) + (env_key,) = _calls_to(calls, "audit_sdk_env_key") + assert env_key.headers["authorization"] == f"Bearer {_ENV_JWT}", dict(env_key.headers) + assert "x-amz-date" not in env_key.headers, dict(env_key.headers) + + +def _answering(request: Request) -> Reply: + return _json_reply({"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE}) + + +def _envelope(request: Request) -> Reply: + return _json_reply({"error": {"code": "bad_request", "message": "plain server refused"}}) + + +def _plain_model(scenario: Scenario, wire: Wire, **extra: JsonValue) -> str: + return scenario.model(model_info=None, model=_MODEL, api_base=wire.url, api_key=None, **extra) + + +def test_a_request_body_session_id_is_refused_on_every_route_without_an_upstream_call(gateway: Gateway) -> None: + with wire_server(_answering) as wire, gateway.scenario() as scenario: + model: Final = _plain_model(scenario, wire) + responses: Final = tuple( + _decide(gateway, model, route, agentcore_runtime_session_id=f"caller-{uuid.uuid4().hex}") + for route in _ROUTES + ) + for response in responses: + assert response.status_code == 401, response.text + assert "agentcore_runtime_session_id is not allowed in request body" in response.text, response.text + assert wire.drain() == () + + +def test_a_plain_api_base_opted_into_the_session_param_sends_no_session_header_or_body_field( + gateway: Gateway, +) -> None: + with wire_server(_answering) as wire, gateway.scenario() as scenario: + model: Final = _plain_model( + scenario, wire, configurable_clientside_auth_params=["agentcore_runtime_session_id"] + ) + _assert_answered(_decide(gateway, model, agentcore_runtime_session_id=f"caller-{uuid.uuid4().hex}"), model) + (call,) = wire.drain() + assert call.target == "/v1/systemone", call.target + assert _SESSION_HEADER not in call.headers, dict(call.headers) + assert json.loads(call.body) == _UPSTREAM_BODY, call.body + + +def test_a_plain_api_base_error_envelope_keeps_the_unexpected_response_server_error(gateway: Gateway) -> None: + with wire_server(_envelope) as wire, gateway.scenario() as scenario: + model: Final = _plain_model(scenario, wire) + response: Final = _decide(gateway, model) + assert response.status_code == 500, response.text + assert "Strands Decider runtime error" not in response.text, response.text + assert len(wire.drain()) == 1 + + +def test_a_plain_api_base_forwards_amazon_shaped_extra_headers_unchanged(gateway: Gateway) -> None: + headers: Final[dict[str, JsonValue]] = { + "x-amz-target": "plain-target", + "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "plain-session-kept", + "x-audit-trace": "plain-trace", + } + with wire_server(_answering) as wire, gateway.scenario() as scenario: + model: Final = _plain_model(scenario, wire, extra_headers=headers) + _assert_answered(_decide(gateway, model), model) + (call,) = wire.drain() + assert (call.headers["x-amz-target"], call.headers[_SESSION_HEADER], call.headers["x-audit-trace"]) == ( + "plain-target", + "plain-session-kept", + "plain-trace", + ), dict(call.headers) + assert "authorization" not in call.headers, dict(call.headers) + assert json.loads(call.body) == _UPSTREAM_BODY, call.body + + +def _worst_case_api_base() -> str: + return "arn:aws" + "-a" * 131072 + "!" + + +def test_a_worst_case_runtime_arn_pattern_input_is_classified_in_linear_time( + gateway: Gateway, record_property: Callable[[str, object], None] +) -> None: + worst: Final = _worst_case_api_base() + started: Final = time.perf_counter() + url: Final = StrandsDeciderDecisionsConfig().get_complete_url(worst, _BODY_MODEL) + classified: Final = time.perf_counter() - started + record_property("classifier_seconds", f"{classified:.6f}") + assert url == f"{worst}/v1/systemone" + assert classified < 0.5, classified + with wire_server(_answering) as wire, gateway.scenario() as scenario: + model: Final = _plain_model(scenario, wire, configurable_clientside_auth_params=["api_base"]) + with ThreadPoolExecutor(max_workers=2) as pool: + timed: Final = pool.submit(_timed_decide, gateway, model, worst) + liveliness_started: Final = time.perf_counter() + liveliness: Final = gateway.request("GET", "/health/liveliness") + liveliness_seconds: Final = time.perf_counter() - liveliness_started + response, decide_seconds = timed.result() + record_property("decide_seconds", f"{decide_seconds:.6f}") + record_property("liveliness_seconds", f"{liveliness_seconds:.6f}") + assert liveliness.status_code == 200, liveliness.text + assert 500 <= response.status_code < 600, response.text + assert "litellm.APIConnectionError" in response.text, response.text[:2000] + assert decide_seconds < 5, decide_seconds + assert liveliness_seconds < 2, liveliness_seconds + assert wire.drain() == () + + +def _timed_decide(gateway: Gateway, model: str, api_base: str) -> tuple[httpx.Response, float]: + started: Final = time.perf_counter() + response: Final = _decide(gateway, model, api_base=api_base) + return response, time.perf_counter() - started diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index c1cac2d386f5..ebe0eedea7fa 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -2,6 +2,8 @@ import asyncio import json +import threading +from concurrent.futures import wait from collections.abc import Mapping from types import MappingProxyType from typing import Final @@ -9,11 +11,16 @@ import httpx import pytest import respx +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials import litellm from litellm.cost_calculator import get_response_cost_from_hidden_params +from litellm.constants import AWS_SIGNING_MAX_THREADS from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR from litellm.types.decisions import ( ChoiceAnswer, DecisionsResponse, @@ -698,6 +705,33 @@ async def test_strands_decider_without_key_preserves_response_extras( assert severity.legend == {"0": "none", "1": "low", "2": "high"} +@pytest.mark.asyncio +async def test_strands_decider_plain_api_base_answers_while_every_aws_signing_thread_is_busy( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False) + respx_mock.post("https://strands.example/v1/systemone").respond(json=_STRANDS_RESPONSE) + release: Final = threading.Event() + busy: Final = tuple(AWS_SIGNING_EXECUTOR.submit(release.wait) for _ in range(AWS_SIGNING_MAX_THREADS)) + + try: + response: Final = await asyncio.wait_for( + litellm.adecisions( + model="strands_decider/strands-decider-2B-hobson-v19", + state="review", + questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}}, + api_base="https://strands.example", + ), + timeout=5, + ) + finally: + release.set() + wait(busy) + + assert isinstance(response.answers["severity"], ScoreAnswer) + + @pytest.mark.asyncio async def test_strands_decider_uses_key_from_matching_environment_base( monkeypatch: pytest.MonkeyPatch, @@ -750,6 +784,172 @@ async def test_strands_decider_provider_resolution_and_router_dispatch( assert response.model == _STRANDS_RESPONSE["model"] +_RUNTIME_ARN: Final = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/qa_decider-AbC123xyz0" +_RUNTIME_URL: Final = ( + "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/" + "arn%3Aaws%3Abedrock-agentcore%3Aus-east-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0/invocations" +) +_RUNTIME_CREDENTIALS: Final = Credentials(access_key="AKIDEXAMPLE", secret_key="example-secret") +_SESSION_HEADER: Final = "x-amzn-bedrock-agentcore-runtime-session-id" + + +def _sigv4_verifies(request: httpx.Request) -> bool: + authorization: Final = request.headers["authorization"] + signed_names: Final = authorization.split("SignedHeaders=")[1].split(",")[0].split(";") + replayed: Final = AWSRequest( + method=request.method, + url=str(request.url), + data=request.content, + headers={name: request.headers[name] for name in signed_names}, + ) + replayed.context["timestamp"] = request.headers["x-amz-date"] + auth: Final = SigV4Auth(_RUNTIME_CREDENTIALS, "bedrock-agentcore", "us-east-1") + string_to_sign: Final = auth.string_to_sign(replayed, auth.canonical_request(replayed)) + return authorization.endswith(f"Signature={auth.signature(string_to_sign, replayed)}") + + +@pytest.mark.parametrize("use_async", [True, False]) +@pytest.mark.asyncio +async def test_strands_decider_runtime_arn_sends_a_sigv4_signed_invoke_agent_runtime_call( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + use_async: bool, +) -> None: + monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE) + call_args: Final[Mapping[str, object]] = { + "model": "strands_decider/strands-decider-2B-hobson-v19", + "state": "review", + "questions": {"severity": {"type": "score", "criteria": ["none", "low", "high"]}}, + "api_base": _RUNTIME_ARN, + "aws_access_key_id": _RUNTIME_CREDENTIALS.access_key, + "aws_secret_access_key": _RUNTIME_CREDENTIALS.secret_key, + } + + response: Final = await litellm.adecisions(**call_args) if use_async else litellm.decisions(**call_args) + + assert route.call_count == 1 + sent: Final = route.calls[0].request + assert _sigv4_verifies(sent) + assert json.loads(sent.content)["model"] == "strands-decider-2B-hobson-v19" + assert sent.headers[_SESSION_HEADER].startswith("litellm-decider-") + assert response.model == _STRANDS_RESPONSE["model"] + + +@pytest.mark.asyncio +async def test_strands_decider_runtime_arn_router_deployment_answers_openai_format_calls( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "strands", + "litellm_params": { + "model": "strands_decider/strands-decider-2B-hobson-v19", + "api_base": _RUNTIME_ARN, + "aws_access_key_id": _RUNTIME_CREDENTIALS.access_key, + "aws_secret_access_key": _RUNTIME_CREDENTIALS.secret_key, + "aws_region_name": "eu-west-1", + "agentcore_runtime_session_id": "pooled-session-0123456789abcdef0123456", + }, + } + ] + ) + route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE) + + response: Final = await router.adecisions( + model="strands", + input="review", + extra_headers={_SESSION_HEADER: "caller-session-0123456789abcdef01234567"}, + questions=[ + { + "type": "score", + "name": "severity", + "instructions": "How bad?", + "levels": [{"label": "none"}, {"label": "low"}, {"label": "high"}], + } + ], + ) + + sent: Final = route.calls[0].request + assert _sigv4_verifies(sent) + assert sent.headers[_SESSION_HEADER] == "pooled-session-0123456789abcdef0123456" + assert isinstance(response, OpenAIDecisionResponse) + + +@pytest.mark.asyncio +async def test_strands_decider_runtime_arn_with_api_key_sends_the_bearer_token_unsigned( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False) + route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE) + + await litellm.adecisions( + model="strands_decider/strands-decider-2B-hobson-v19", + state="review", + questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}}, + api_base=_RUNTIME_ARN, + api_key="runtime-jwt", + ) + + sent: Final = route.calls[0].request + assert sent.headers["authorization"] == "Bearer runtime-jwt" + assert "x-amz-date" not in sent.headers + assert _SESSION_HEADER in sent.headers + + +@pytest.mark.parametrize( + ("code", "error_class"), + [ + ("bad_request", litellm.BadRequestError), + ("loading", litellm.ServiceUnavailableError), + ("inference_error", litellm.InternalServerError), + ], +) +@pytest.mark.asyncio +async def test_strands_decider_runtime_error_envelope_raises_instead_of_answering( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + code: str, + error_class: type[Exception], +) -> None: + monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False) + respx_mock.post(_RUNTIME_URL).respond(json={"error": {"code": code, "message": "at most 16 questions"}}) + + with pytest.raises(error_class, match="at most 16 questions"): + await litellm.adecisions( + model="strands_decider/strands-decider-2B-hobson-v19", + state="review", + questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}}, + api_base=_RUNTIME_ARN, + aws_access_key_id=_RUNTIME_CREDENTIALS.access_key, + aws_secret_access_key=_RUNTIME_CREDENTIALS.secret_key, + ) + + +@pytest.mark.asyncio +async def test_strands_decider_plain_api_base_error_envelope_stays_an_unexpected_response( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False) + respx_mock.post("https://strands.example/v1/systemone").respond( + json={"error": {"code": "bad_request", "message": "at most 16 questions"}} + ) + + with pytest.raises(litellm.InternalServerError, match="returned an unexpected response"): + await litellm.adecisions( + model="strands_decider/strands-decider-2B-hobson-v19", + state="review", + questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}}, + api_base="https://strands.example", + ) + + _VLLM_RESPONSE: Final[Mapping[str, object]] = { "id": "systemone-4af3d2c1", "object": "structured_decision", diff --git a/tests/unit/llms/strands_decider/__init__.py b/tests/unit/llms/strands_decider/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/unit/llms/strands_decider/decisions/__init__.py b/tests/unit/llms/strands_decider/decisions/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/unit/llms/strands_decider/decisions/test_strands_decider_decisions_transformation.py b/tests/unit/llms/strands_decider/decisions/test_strands_decider_decisions_transformation.py new file mode 100644 index 000000000000..4d249e592335 --- /dev/null +++ b/tests/unit/llms/strands_decider/decisions/test_strands_decider_decisions_transformation.py @@ -0,0 +1,256 @@ +import json +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +from botocore.credentials import Credentials +from pydantic import TypeAdapter + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.decisions.transformation import systemone_request_to_ir +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.strands_decider.decisions.transformation import ( + AGENTCORE_SESSION_HEADER, + StrandsDeciderDecisionsConfig, +) +from litellm.types.decisions import DecisionsIRResponse, DecisionsRequestBody +from litellm.types.llms.bedrock import AwsAuthParams + +_ARN: Final = "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0" +_BODY: Final[Mapping[str, object]] = { + "model": "strands-decider-2B-hobson-v21", + "state": "My order arrived broken.", + "questions": {"is_complaint": {"type": "noul", "instructions": "Is this a complaint?"}}, +} +_IR_REQUEST: Final = systemone_request_to_ir( + TypeAdapter(DecisionsRequestBody).validate_python({"state": _BODY["state"], "questions": _BODY["questions"]}) +) + + +class _RecordingAWS(BaseAWSLLM): + def __init__(self) -> None: + super().__init__() + self.resolved: tuple[tuple[AwsAuthParams, str | None], ...] = () + + def resolve_credentials(self, auth_params: AwsAuthParams, aws_region_name: str | None) -> Credentials: + self.resolved = (*self.resolved, (auth_params, aws_region_name)) + return Credentials(access_key="AKIDEXAMPLE", secret_key="example-secret") + + +def _sign( + config: StrandsDeciderDecisionsConfig, + api_base: str, + headers: Mapping[str, str], + api_key: str | None = None, + litellm_params: Mapping[str, object] | None = None, +) -> tuple[Mapping[str, str], bytes | None]: + return config.sign_request( + headers=headers, + url=config.get_complete_url(api_base=api_base, model="strands-decider-2B-hobson-v21"), + api_base=api_base, + body=_BODY, + api_key=api_key, + litellm_params=litellm_params or {}, + ) + + +@pytest.mark.parametrize( + ("arn", "expected_url"), + [ + ( + _ARN, + "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/" + "arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0/invocations", + ), + ( + "arn:aws-cn:bedrock-agentcore:cn-north-1:123456789012:runtime/qa_decider-AbC123xyz0", + "https://bedrock-agentcore.cn-north-1.amazonaws.com.cn/runtimes/" + "arn%3Aaws-cn%3Abedrock-agentcore%3Acn-north-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0" + "/invocations", + ), + ( + "arn:aws-us-gov:bedrock-agentcore:us-gov-west-1:123456789012:runtime/qa_decider-AbC123xyz0", + "https://bedrock-agentcore.us-gov-west-1.amazonaws.com/runtimes/" + "arn%3Aaws-us-gov%3Abedrock-agentcore%3Aus-gov-west-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0" + "/invocations", + ), + ], +) +def test_runtime_arn_targets_the_regional_invoke_agent_runtime_url(arn: str, expected_url: str) -> None: + assert StrandsDeciderDecisionsConfig().get_complete_url(api_base=arn, model="any") == expected_url + + +@pytest.mark.parametrize( + "api_base", + [ + "https://strands.example", + "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0/runtime-endpoint/DEFAULT", + "arn:aws:bedrock-agentcore:evil.example/x:123456789012:runtime/qa_decider-AbC123xyz0", + "arn:aws:bedrock:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0", + ], +) +def test_other_api_bases_keep_the_plain_systemone_request(api_base: str) -> None: + aws: Final = _RecordingAWS() + config: Final = StrandsDeciderDecisionsConfig(aws=aws) + headers: Final = {"Content-Type": "application/json"} + + assert config.get_complete_url(api_base=api_base, model="any") == f"{api_base}/v1/systemone" + assert _sign(config, api_base, headers) == (headers, None) + assert aws.resolved == () + + +def test_sign_request_signs_the_body_for_the_arn_region_with_deployment_credentials() -> None: + aws: Final = _RecordingAWS() + config: Final = StrandsDeciderDecisionsConfig(aws=aws) + + signed_headers, signed_body = _sign( + config, + _ARN, + {"Content-Type": "application/json"}, + litellm_params={"aws_profile_name": "qa", "aws_role_name": "decider-role", "aws_region_name": "us-east-1"}, + ) + + assert signed_body == json.dumps(_BODY).encode() + assert len(aws.resolved) == 1 + auth_params, region = aws.resolved[0] + assert (auth_params.aws_profile_name, auth_params.aws_role_name, region) == ("qa", "decider-role", "us-west-2") + authorization: Final = signed_headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/") + assert "/us-west-2/bedrock-agentcore/aws4_request" in authorization + assert AGENTCORE_SESSION_HEADER.lower() in authorization + + +def test_default_session_id_is_stable_per_runtime_and_within_the_agentcore_length_bounds() -> None: + other_arn: Final = _ARN.replace("AbC123xyz0", "ZyX987cba1") + + first, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {}) + again, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {}) + other, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), other_arn, {}) + + session_id: Final = first[AGENTCORE_SESSION_HEADER] + assert session_id == again[AGENTCORE_SESSION_HEADER] + assert session_id != other[AGENTCORE_SESSION_HEADER] + assert 33 <= len(session_id) <= 256 + + +_CALLER_RESERVED_HEADERS: Final = { + AGENTCORE_SESSION_HEADER.lower(): "caller-session-0123456789abcdef0123456789", + "X-Amzn-Bedrock-AgentCore-Runtime-User-Id": "someone-else", + "X-Amz-Security-Token": "caller-token", + "Host": "evil.example", +} + + +@pytest.mark.parametrize("api_key", [None, "runtime-jwt"]) +def test_caller_reserved_agentcore_headers_never_reach_the_runtime(api_key: str | None) -> None: + signed_headers, _ = _sign( + StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), + _ARN, + {**_CALLER_RESERVED_HEADERS, "X-Request-Source": "batch"}, + api_key=api_key, + ) + default_headers, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {}, api_key=api_key) + + forwarded: Final = {name.lower(): value for name, value in signed_headers.items()} + assert forwarded[AGENTCORE_SESSION_HEADER.lower()] == default_headers[AGENTCORE_SESSION_HEADER] + assert not any(value in forwarded.values() for value in _CALLER_RESERVED_HEADERS.values()) + assert "x-amzn-bedrock-agentcore-runtime-user-id" not in forwarded + assert "host" not in forwarded + assert forwarded["x-request-source"] == "batch" + + +def test_deployment_session_param_picks_the_session_over_the_default_and_any_caller_header() -> None: + replica_session: Final = "strands-decider-replica-a-0000000000" + + signed_headers, _ = _sign( + StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), + _ARN, + _CALLER_RESERVED_HEADERS, + litellm_params={"agentcore_runtime_session_id": replica_session}, + ) + + session_headers: Final = [ + value for name, value in signed_headers.items() if name.lower() == "x-amzn-bedrock-agentcore-runtime-session-id" + ] + assert session_headers == [replica_session] + assert AGENTCORE_SESSION_HEADER.lower() in signed_headers["Authorization"] + + +def test_api_key_on_a_runtime_arn_sends_the_bearer_token_without_sigv4() -> None: + aws: Final = _RecordingAWS() + + signed_headers, signed_body = _sign( + StrandsDeciderDecisionsConfig(aws=aws), _ARN, {"Authorization": "Bearer runtime-jwt"}, api_key="runtime-jwt" + ) + + assert signed_body is None + assert signed_headers["Authorization"] == "Bearer runtime-jwt" + assert AGENTCORE_SESSION_HEADER in signed_headers + assert aws.resolved == () + + +def _response(url: str, payload: object) -> httpx.Response: + return httpx.Response(200, json=payload, request=httpx.Request("POST", url)) + + +def _transform(api_base: str, payload: object) -> DecisionsIRResponse: + config: Final = StrandsDeciderDecisionsConfig() + url: Final = config.get_complete_url(api_base=api_base, model="strands-decider-2B-hobson-v21") + return config.transform_decisions_response( + model="strands-decider-2B-hobson-v21", + custom_llm_provider="strands_decider", + raw_response=_response(url, payload), + request=_IR_REQUEST, + ) + + +@pytest.mark.parametrize( + ("code", "status_code"), + [("bad_request", 400), ("invalid_request", 400), ("loading", 503), ("inference_error", 500), ("new_code", 500)], +) +def test_runtime_error_envelope_raises_with_the_mapped_status(code: str, status_code: int) -> None: + with pytest.raises(BaseLLMException) as raised: + _transform(_ARN, {"error": {"code": code, "message": "at most 16 questions"}}) + + assert raised.value.status_code == status_code + assert code in raised.value.message + assert "at most 16 questions" in raised.value.message + + +@pytest.mark.parametrize("code", ["bad_request", "loading", "inference_error"]) +def test_plain_api_base_error_envelope_keeps_the_unexpected_response_error(code: str) -> None: + with pytest.raises(BaseLLMException) as raised: + _transform("https://strands.example", {"error": {"code": code, "message": "at most 16 questions"}}) + + assert raised.value.status_code == 500 + assert raised.value.message.startswith("Decisions provider 'strands_decider' returned an unexpected response") + + +@pytest.mark.parametrize( + "url", + [ + "https://strands.example/runtimes/arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A123456789012%3Aruntime" + "%2Fqa_decider-AbC123xyz0/invocations", + "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/not-an-arn/invocations", + ], +) +def test_error_envelope_from_a_url_that_is_not_the_runtime_keeps_the_unexpected_response_error(url: str) -> None: + with pytest.raises(BaseLLMException) as raised: + StrandsDeciderDecisionsConfig().transform_decisions_response( + model="strands-decider-2B-hobson-v21", + custom_llm_provider="strands_decider", + raw_response=_response(url, {"error": {"code": "bad_request", "message": "at most 16 questions"}}), + request=_IR_REQUEST, + ) + + assert raised.value.status_code == 500 + + +@pytest.mark.parametrize("payload", [["not", "an", "object"], {"error": "unstructured"}]) +def test_runtime_reply_that_is_not_an_error_envelope_keeps_the_unexpected_response_error(payload: object) -> None: + with pytest.raises(BaseLLMException) as raised: + _transform(_ARN, payload) + + assert raised.value.status_code == 500 + assert raised.value.message.startswith("Decisions provider 'strands_decider' returned an unexpected response") diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index e2cec1cc5a2e..92bd57343538 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -2854,6 +2854,22 @@ def test_admin_opt_in_proxy_wide_allows_project_id(self): ) +def test_agentcore_runtime_session_id_in_request_body_is_rejected_without_admin_opt_in(): + body: Final = {"model": "strands-decider", "agentcore_runtime_session_id": "caller-session-0123456789abcdef0123"} + + with pytest.raises(ValueError, match="agentcore_runtime_session_id"): + is_request_body_safe(request_body=body, general_settings={}, llm_router=None, model="strands-decider") + assert ( + is_request_body_safe( + request_body=body, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="strands-decider", + ) + is True + ) + + class TestIsRequestBodySafeBlocksClaudePlatformWorkspaceOverride: @pytest.mark.parametrize( "alias", ["workspace_id", "aws_workspace_id", "anthropic_workspace_id", "anthropic-workspace-id"] From 8dcac8ab2381fabd84cb4d4bc4540ad310426bce Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:29:33 -0700 Subject: [PATCH 07/18] feat(ui): rename the Playground System One tab to Decisions and give the JSON editor VS Code light and dark themes (#45729) Backport of #45729 to rc/1.106.0. Cherry-picked from 512692dd36e222758fcbd5e656bb6ea7dae56618 (main). --- .../components/systemOneUI/JsonEditor.tsx | 110 ++++++++++-------- .../systemOneUI/QuestionBreakdown.tsx | 9 +- .../systemOneUI/ResponseView.test.tsx | 14 +++ .../components/systemOneUI/ResponseView.tsx | 4 +- .../systemOneUI/lib/rootBlocks.test.ts | 53 --------- .../components/systemOneUI/lib/rootBlocks.ts | 65 ----------- .../app/(dashboard)/playground/page.test.tsx | 4 +- .../src/app/(dashboard)/playground/page.tsx | 2 +- 8 files changed, 85 insertions(+), 176 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/rootBlocks.test.ts delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/rootBlocks.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx index fce617b4c476..cc7bebb0bb82 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx @@ -1,11 +1,12 @@ import { Badge } from "@/components/ui/badge"; import { cn } from "@/lib/cva.config"; import { CircleAlert, CircleCheck, TriangleAlert } from "lucide-react"; -import { useId, useMemo, useRef } from "react"; +import { useTheme } from "next-themes"; +import { useId, useRef } from "react"; import { createElement, PrismLight as SyntaxHighlighter } from "react-syntax-highlighter"; import type { SyntaxHighlighterProps } from "react-syntax-highlighter"; import json from "react-syntax-highlighter/dist/esm/languages/prism/json"; -import { findRootBlocks, ROOT_BLOCK_STYLES, type RootBlock } from "./lib/rootBlocks"; +import { vs, vscDarkPlus } from "react-syntax-highlighter/dist/esm/styles/prism"; import type { SystemOnePayloadValidation } from "./lib/validatePayload"; SyntaxHighlighter.registerLanguage("json", json); @@ -14,16 +15,44 @@ const EDITOR_TEXT = "m-0 whitespace-pre-wrap wrap-anywhere py-3 font-mono text-x const GUTTER_WIDTH = "w-11"; type LineRendererProps = Parameters>[0]; const CONTENT_INSET = "pl-14 pr-3"; -const CODE_TAG_PROPS = { className: "language-json", style: { whiteSpace: "pre-wrap" } } as const; +type SyntaxTheme = NonNullable; +const THEME_BLOCKS = ['pre[class*="language-"]', 'code[class*="language-"]'] as const; -const TOKEN_COLORS = [ - "[&_.token.property]:text-sky-700 dark:[&_.token.property]:text-sky-300", - "[&_.token.string]:text-emerald-700 dark:[&_.token.string]:text-emerald-300", - "[&_.token.number]:text-amber-700 dark:[&_.token.number]:text-amber-300", - "[&_.token.boolean]:text-violet-700 dark:[&_.token.boolean]:text-violet-300", - "[&_.token.null]:text-violet-700 dark:[&_.token.null]:text-violet-300", - "[&_.token.punctuation]:text-muted-foreground [&_.token.operator]:text-muted-foreground", -].join(" "); +function withoutBlockBackground(theme: SyntaxTheme): SyntaxTheme { + const stripped = THEME_BLOCKS.map((block) => { + const { background: _background, backgroundColor: _backgroundColor, ...rest } = theme[block] ?? {}; + return [block, rest] as const; + }); + return { ...theme, ...Object.fromEntries(stripped) }; +} + +const VS_CODE_DARK_PLUS = withoutBlockBackground(vscDarkPlus); +const VS_CODE_LIGHT_PLUS = { + ...withoutBlockBackground(vs), + property: { color: "#0451a5" }, + string: { color: "#a31515" }, + number: { color: "#098658" }, + boolean: { color: "#0000ff" }, + punctuation: { color: "#000000" }, + operator: { color: "#000000" }, +} as const; +const INHERIT_TEXT = { + fontFamily: "inherit", + fontSize: "inherit", + lineHeight: "inherit", + whiteSpace: "pre-wrap", + overflowWrap: "anywhere", + textShadow: "none", +} as const; +const CODE_TAG_PROPS = { className: "language-json", style: INHERIT_TEXT } as const; +const TRANSPARENT_PRE = { + ...INHERIT_TEXT, + background: "transparent", + border: 0, + margin: 0, + padding: 0, + overflow: "visible", +} as const; interface JsonEditorProps { value: string; @@ -69,39 +98,24 @@ export function IssueList({ id, validation }: { id: string; validation: SystemOn ); } -function renderLines(rootBlocks: readonly RootBlock[]) { - return function LineRows({ rows, stylesheet, useInlineStyles }: LineRendererProps) { - return rows.map((row, line) => { - const lineElement = { node: row, stylesheet, useInlineStyles, key: line }; - const block = rootBlocks.find(({ startLine, endLine }) => line >= startLine && line <= endLine); - return ( -
- - {line + 1} - - - {createElement(lineElement)} - -
- ); - }); - }; +function LineRows({ rows, stylesheet, useInlineStyles }: LineRendererProps) { + return rows.map((row, line) => { + const lineElement = { node: row, stylesheet, useInlineStyles, key: line }; + return ( +
+ + {line + 1} + + {createElement(lineElement)} +
+ ); + }); } export default function JsonEditor({ value, onChange, validation }: JsonEditorProps) { const issuesId = useId(); const highlightRef = useRef(null); - const renderer = useMemo(() => renderLines(findRootBlocks(value)), [value]); + const syntaxTheme = useTheme().resolvedTheme === "dark" ? VS_CODE_DARK_PLUS : VS_CODE_LIGHT_PLUS; const lineCount = value.split("\n").length; const hasErrors = !validation.isValid; @@ -121,20 +135,22 @@ export default function JsonEditor({ value, onChange, validation }: JsonEditorPr {lineCount} {lineCount === 1 ? "line" : "lines"}
-
-