diff --git a/litellm/llms/databricks/decisions/transformation.py b/litellm/llms/databricks/decisions/transformation.py index a454da83e50..f33b424cf0c 100644 --- a/litellm/llms/databricks/decisions/transformation.py +++ b/litellm/llms/databricks/decisions/transformation.py @@ -4,20 +4,71 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import Final -from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig +from pydantic import TypeAdapter, ValidationError +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig +from litellm.types.decisions import DecisionsIRRequest, UnsupportedDecisionsRequest + +DATABRICKS_AI_DECIDE_MODEL: Final = "ai_decide" +_AI_DECIDE_PATH: Final = "/api/2.0/ai-functions/ai-decide" _SERVING_ENDPOINT_NAME: Final = re.compile(r"[A-Za-z0-9_-][A-Za-z0-9._-]*") +_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def is_ai_decide_model(model: str) -> bool: + return model == DATABRICKS_AI_DECIDE_MODEL 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" + "Databricks decisions model must be 'ai_decide' (the ai_decide AI Function) or a bare serving endpoint " + "name (letters, digits, '-', '_' and '.', with no '/', '?', '#', spaces, or a leading '.'), e.g. " + "databricks-openjev-qwen35-4b" ) return name +def validate_databricks_decisions_model(model: str) -> str: + if is_ai_decide_model(model): + return model + return validate_serving_endpoint_name(model) + + +def databricks_workspace_host(api_base: str) -> str: + return api_base.rstrip("/").removesuffix("/serving-endpoints") + + +def _mapping(value: object) -> Mapping[str, object] | None: + try: + return _MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _systemone_answer(answer: object) -> object: + fields: Final = _mapping(answer) + if fields is None or fields.get("type") != "noul" or "probability" not in fields: + return answer + return MappingProxyType( + {**{key: value for key, value in fields.items() if key != "probability"}, "noul": fields["probability"]} + ) + + +def ai_decide_response(payload: Mapping[str, object]) -> Mapping[str, object]: + response: Final = _mapping(payload.get("response")) + answers: Final = None if response is None else _mapping(response.get("answers")) + if response is None or answers is None: + return payload + return MappingProxyType( + { + **{key: value for key, value in payload.items() if key != "response"}, + **response, + "answers": MappingProxyType({key: _systemone_answer(answer) for key, answer in answers.items()}), + } + ) + + @dataclass(frozen=True, slots=True) class DatabricksDecisionsConnection: api_base: str @@ -35,13 +86,30 @@ class DatabricksDecisionsConfig(BaseDecisionsConfig): ) def canonical_model(self, model: str) -> str: - return validate_serving_endpoint_name(model) + return validate_databricks_decisions_model(model) def get_complete_url(self, api_base: str, model: str) -> str: - return f"{api_base.rstrip('/')}/{validate_serving_endpoint_name(model)}/invocations" + if is_ai_decide_model(model): + return f"{databricks_workspace_host(api_base)}{_AI_DECIDE_PATH}" + return f"{api_base.rstrip('/')}/{validate_databricks_decisions_model(model)}/invocations" + + def transform_decisions_request( + self, + model: str, + request: DecisionsIRRequest, + custom_llm_provider: str, + ) -> Mapping[str, object] | UnsupportedDecisionsRequest: + body: Final = super().transform_decisions_request(model, request, custom_llm_provider) + if isinstance(body, UnsupportedDecisionsRequest) or not is_ai_decide_model(model): + return body + return MappingProxyType({key: value for key, value in body.items() if key != "model"}) + + def unwrap_response(self, payload: object) -> object: + fields: Final = _mapping(payload) + return payload if fields is None else ai_decide_response(fields) def classifier_response(self, body: Mapping[str, object], requested_model: str) -> Mapping[str, object]: - return MappingProxyType({**body, "model": requested_model}) + return MappingProxyType({**ai_decide_response(body), "model": requested_model}) def connection(self, api_base: str | None, api_key: str | None) -> DatabricksDecisionsConnection: base: Final = self.resolve_api_base(api_base) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8d0157f0b97..351139617f7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21203,6 +21203,18 @@ "output_cost_per_token": 0.0, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" }, + "databricks/ai_decide": { + "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": "The ai_decide AI Function (Beta), called through POST /api/2.0/ai-functions/ai-decide. Its response carries no token usage, so LiteLLM records no spend for it." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.databricks.com/api/ai-functions/v1/ai-decide" + }, "dataforseo/search": { "input_cost_per_query": 0.003, "litellm_provider": "dataforseo", diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 992c098bc80..4a2c56daf00 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -769,14 +769,17 @@ class OpenSourceClassifierConfig(LiteLLMBaseModel): _ = 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 + from litellm.llms.databricks.decisions.transformation import ( + DATABRICKS_AI_DECIDE_MODEL, + validate_databricks_decisions_model, + ) 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" + "opensource_classifier_config.model is required for provider 'databricks': " + f"{DATABRICKS_AI_DECIDE_MODEL!r} or a serving endpoint name, e.g. databricks-openjev-qwen35-4b" ) - _ = validate_serving_endpoint_name(self.model) + _ = validate_databricks_decisions_model(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: " diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index 9425b4d35a4..30af68bdf30 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -19,7 +19,7 @@ from litellm.litellm_core_utils.internal_call_metadata import ( 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.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG, is_ai_decide_model 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, @@ -35,6 +35,7 @@ if TYPE_CHECKING: 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 +_WITHOUT_MODEL: Final = MappingProxyType({"model": True}) class JevChoiceQuestion(LiteLLMBaseModel): @@ -183,7 +184,7 @@ class HttpJevClassifierClient: ) response: Final = await self._http_client.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler has a dynamic post signature self._request_url(request.model), - json=request.model_dump(mode="json"), + json=request.model_dump(mode="json", exclude=self._excluded_fields(request.model)), headers=MappingProxyType({**authorization, "Content-Type": "application/json"}), # pyright: ignore[reportArgumentType] # HTTP headers are not mutated by AsyncHTTPHandler timeout=timeout_s, ) @@ -196,6 +197,9 @@ class HttpJevClassifierClient: verbose_router_logger.warning("JEV response logging failed (%s)", type(exc).__name__) return TypeAdapter(JevSystemOneResponse).validate_python(normalized_body) + def _excluded_fields(self, model: str) -> Mapping[str, bool] | None: + return _WITHOUT_MODEL if self._provider == "databricks" and is_ai_decide_model(model) else None + def _request_url(self, model: str) -> str: if self._provider == "databricks": return DATABRICKS_DECISIONS_CONFIG.get_complete_url(self._api_base, model) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8d0157f0b97..351139617f7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21203,6 +21203,18 @@ "output_cost_per_token": 0.0, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" }, + "databricks/ai_decide": { + "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": "The ai_decide AI Function (Beta), called through POST /api/2.0/ai-functions/ai-decide. Its response carries no token usage, so LiteLLM records no spend for it." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.databricks.com/api/ai-functions/v1/ai-decide" + }, "dataforseo/search": { "input_cost_per_query": 0.003, "litellm_provider": "dataforseo", diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 91ab5652051..2e72b212089 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -95,6 +95,21 @@ _DATABRICKS_RESPONSE: Final[Mapping[str, object]] = { }, "usage": {"input_tokens": 214, "output_tokens": 0}, } +_DATABRICKS_ROUTE: Final = "https://workspace.example/api/2.0/ai-functions/ai-decide" +_DATABRICKS_AI_DECIDE_RESPONSE: Final[Mapping[str, object]] = { + "response": { + "answers": { + "is_defect": {"type": "noul", "probability": 0.9}, + "sentiment": { + "type": "choice", + "choice": "positive", + "confidence": 0.8, + "probabilities": {"positive": 0.8, "negative": 0.2}, + }, + } + }, + "metadata": {"version": "1.0"}, +} _PROVIDERS: Final[tuple[tuple[str, str, str, str], ...]] = ( ( "perplexity", @@ -1117,7 +1132,7 @@ async def test_databricks_rejects_a_serving_endpoint_name_that_would_rewrite_the 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"): + with pytest.raises(litellm.BadRequestError, match="bare serving endpoint name"): await litellm.asystemone( model=f"databricks/{endpoint_name}", state="review", @@ -1245,6 +1260,180 @@ def test_every_dedicated_decider_provider_ships_an_evaluation_mode_cost_map_entr assert evaluation_entries, f"a {provider} decider would be health-checked as chat without a mode: evaluation entry" +@pytest.mark.asyncio +async def test_databricks_ai_decide_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="databricks/ai_decide", + 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_adecisions_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 serving 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_ai_decide_requires_a_key_before_http( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.example") + 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="databricks/ai_decide", + 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_adecisions_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 +@pytest.mark.parametrize( + "api_base", + [ + "https://workspace.example", + "https://workspace.example/", + "https://workspace.example/serving-endpoints", + "https://workspace.example/serving-endpoints/", + ], +) +async def test_databricks_posts_the_systemone_body_without_model_to_the_ai_decide_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(_DATABRICKS_ROUTE).respond(json=_DATABRICKS_AI_DECIDE_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="databricks/ai_decide", + 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) == {"state": {"ticket": "export hangs"}, "questions": questions} + assert response.model is None + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert isinstance(response.answers["sentiment"], ChoiceAnswer) + assert response.model_extra == {"metadata": {"version": "1.0"}} + assert response.hidden_params["custom_llm_provider"] == "databricks" + assert response.hidden_params["model"] == "databricks/ai_decide" + + +@pytest.mark.asyncio +async def test_databricks_ai_decide_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") + monkeypatch.setenv("DATABRICKS_API_KEY", "env-key-stays-home") + router: Final = litellm.Router( + model_list=[ + { + "model_name": "decider", + "litellm_params": { + "model": "databricks/ai_decide", + "api_base": "https://workspace.example", + "api_key": "deployment-key", + }, + } + ] + ) + route: Final = respx_mock.post(_DATABRICKS_ROUTE).respond(json=_DATABRICKS_AI_DECIDE_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) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("api_base", "url"), diff --git a/tests/unit/llms/databricks/decisions/test_transformation.py b/tests/unit/llms/databricks/decisions/test_transformation.py index e6de1d1a635..f6ea3444db1 100644 --- a/tests/unit/llms/databricks/decisions/test_transformation.py +++ b/tests/unit/llms/databricks/decisions/test_transformation.py @@ -3,27 +3,84 @@ from typing import Final import pytest from litellm.llms.databricks.decisions.transformation import DatabricksDecisionsConfig +from litellm.types.decisions import ( + DecisionsIRPredicateAnswer, + DecisionsIRPredicateQuestion, + DecisionsIRRequest, + DecisionsIRState, +) -_BASE: Final = "https://workspace.example/serving-endpoints" +_HOST: Final = "https://workspace.example" +_BASE: Final = f"{_HOST}/serving-endpoints" +_AI_DECIDE_ROUTE: Final = f"{_HOST}/api/2.0/ai-functions/ai-decide" _CONFIG: Final = DatabricksDecisionsConfig() +_REQUEST: Final = DecisionsIRRequest( + input=DecisionsIRState(state="export hangs"), + questions=(DecisionsIRPredicateQuestion(name="is_defect", instructions="Is this a defect?", criteria=None),), +) +_QUESTIONS: Final = {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}} @pytest.mark.parametrize("api_base", [_BASE, f"{_BASE}/"]) -def test_the_complete_url_is_the_serving_endpoint_invocations_route(api_base: str) -> None: +def test_a_serving_endpoint_name_is_called_on_its_invocations_route(api_base: str) -> None: assert _CONFIG.get_complete_url(api_base, "openjev") == f"{_BASE}/openjev/invocations" +@pytest.mark.parametrize("api_base", [_HOST, f"{_HOST}/", _BASE, f"{_BASE}/"]) +def test_ai_decide_is_called_on_the_workspace_ai_function_route(api_base: str) -> None: + assert _CONFIG.get_complete_url(api_base, "ai_decide") == _AI_DECIDE_ROUTE + + @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"): + with pytest.raises(ValueError, match="bare serving endpoint name"): _ = _CONFIG.get_complete_url(_BASE, name) - with pytest.raises(ValueError, match="bare endpoint name"): + with pytest.raises(ValueError, match="bare serving endpoint name"): _ = _CONFIG.canonical_model(name) +def test_the_rejection_does_not_echo_the_rejected_name() -> None: + name: Final = "openjev?token=" + "x" * 4096 + with pytest.raises(ValueError, match="bare serving endpoint name") as raised: + _ = _CONFIG.canonical_model(name) + assert "openjev?token=" not in str(raised.value) + + +def test_a_serving_endpoint_request_keeps_the_model_the_endpoint_expects() -> None: + body: Final = _CONFIG.transform_decisions_request("openjev", _REQUEST, "databricks") + assert body == {"model": "openjev", "state": "export hangs", "questions": _QUESTIONS} + + +def test_the_ai_decide_request_omits_model_because_the_route_rejects_unexpected_parameters() -> None: + body: Final = _CONFIG.transform_decisions_request("ai_decide", _REQUEST, "databricks") + assert body == {"state": "export hangs", "questions": _QUESTIONS} + + +def test_the_ai_decide_envelope_is_unwrapped_and_noul_probability_becomes_the_noul_answer() -> None: + payload: Final = { + "response": {"answers": {"is_defect": {"type": "noul", "probability": 0.9}}}, + "metadata": {"version": "1.0"}, + } + parsed: Final = _CONFIG.parse_response(payload, _REQUEST) + assert parsed.answers == (DecisionsIRPredicateAnswer(probability=0.9),) + assert parsed.model is None + assert dict(parsed.extra) == {"metadata": {"version": "1.0"}} + + +def test_a_serving_endpoint_answer_is_parsed_as_sent() -> None: + payload: Final = { + "model": "openjev", + "answers": {"is_defect": {"type": "noul", "noul": 0.25}}, + "usage": {"input_tokens": 12, "output_tokens": 0}, + } + parsed: Final = _CONFIG.parse_response(payload, _REQUEST) + assert parsed.answers == (DecisionsIRPredicateAnswer(probability=0.25),) + assert parsed.model == "openjev" + + 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") @@ -54,3 +111,10 @@ def test_classifier_response_accounts_the_serving_endpoint_instead_of_the_contai assert normalized["model"] == "openjev" assert normalized["usage"] == body["usage"] assert body["model"] == "/mosaicml/local_model" + + +def test_classifier_response_unwraps_the_ai_decide_envelope_and_accounts_the_configured_model() -> None: + choice: Final = {"type": "choice", "choice": "simple", "confidence": 0.9, "probabilities": {"simple": 0.9}} + body: Final = {"response": {"answers": {"tier": choice}}, "metadata": {"version": "1.0"}} + normalized: Final = _CONFIG.classifier_response(body, "ai_decide") + assert dict(normalized) == {"answers": {"tier": choice}, "metadata": {"version": "1.0"}, "model": "ai_decide"} diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index a4a67c7d5d0..7275e1dcb63 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -304,6 +304,53 @@ def test_proxy_systemone_dispatches_databricks_serving_endpoint( assert upstream.calls[0].request.headers["authorization"] == "Bearer dapi-deployment" +def test_proxy_systemone_dispatches_databricks_ai_decide( + 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/ai_decide", + "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/api/2.0/ai-functions/ai-decide").respond( + json={ + "response": {"answers": {"is_defect": {"type": "noul", "probability": 0.9}}}, + "metadata": {"version": "1.0"}, + } + ) + + 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) == { + "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 ba312ce0d2d..f5a097f60c1 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": "ai_decide", "api_base": "https://collector.invalid"}, "opensource_classifier_config"), + ({"provider": "databricks", "model": "ai_decide", "api_key": "dapi-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"), ], @@ -168,7 +170,13 @@ def test_members_cannot_move_the_jev_classifier_off_the_proxys_typesafe_account( @pytest.mark.parametrize( ("provider", "model"), - [("typesafe", "jev-preview"), ("laya", "english"), ("bespoke", "nimble-latest"), ("databricks", "my-openjev")], + [ + ("typesafe", "jev-preview"), + ("laya", "english"), + ("bespoke", "nimble-latest"), + ("databricks", "ai_decide"), + ("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: 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 c1c96588004..f7fbac21304 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -458,6 +458,7 @@ def test_jev_config_requires_classifier_config() -> None: ("laya", "english", "laya"), ("bespoke", "nimble-latest", "bespoke"), ("databricks", "databricks-openjev-qwen35-4b", "databricks"), + ("databricks", "ai_decide", "databricks"), ], ) def test_classifier_aliases_load_and_serialize_one_canonical_config( @@ -540,14 +541,18 @@ 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: +@pytest.mark.parametrize("model", ["ai_decide", "my-openjev", "databricks-openjev-qwen35-4b"]) +def test_databricks_accepts_ai_decide_or_a_bare_serving_endpoint_name(model: str) -> None: + assert JevClassifierConfig.model_validate({"provider": "databricks", "model": model}).model == model + + +def test_databricks_requires_a_model_that_cannot_rewrite_the_url() -> 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"): + with pytest.raises(ValueError, match="bare serving endpoint name"): JevClassifierConfig.model_validate({"provider": "databricks", "model": "serving-endpoints/my-openjev"}) - with pytest.raises(ValueError, match="bare endpoint name"): + with pytest.raises(ValueError, match="bare serving endpoint name"): JevClassifierConfig.model_validate({"provider": "databricks", "model": ".."}) - assert JevClassifierConfig.model_validate({"provider": "databricks", "model": "my-openjev"}).model == "my-openjev" @pytest.mark.asyncio @@ -608,6 +613,53 @@ async def test_databricks_routes_through_the_serving_endpoint_and_accounts_the_e assert recorder.calls[0]["response_cost"] == pytest.approx(0.31) +@pytest.mark.asyncio +@pytest.mark.parametrize("configured_key", [False, True]) +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("api_base", ["https://workspace.test", "https://workspace.test/serving-endpoints/"]) +async def test_databricks_routes_through_ai_decide_without_a_model_field_and_accounts_ai_decide( + monkeypatch: pytest.MonkeyPatch, configured_key: bool, legacy: bool, api_base: str +) -> None: + monkeypatch.setenv("TYPESAFE_API_KEY", "never-send-typesafe-key") + monkeypatch.setenv("DATABRICKS_API_BASE", api_base) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.setenv("DATABRICKS_TOKEN", "dapi-env-token") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _UsageRecorder("databricks/ai_decide") + 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": "ai_decide", + **({"api_key": "dapi-configured", "api_base": api_base} if configured_key else {}), + }, + "tiers": {"SIMPLE": "cheap"}, + }, + derive_savings_baseline=False, + ) + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("https://workspace.test/api/2.0/ai-functions/ai-decide").respond( + 200, + json={"response": {"answers": {"tier": _answer().model_dump()}}, "metadata": {"version": "1.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", "ai_decide") + assert outcome.jev_verdict.label == "SIMPLE" + assert outcome.classifier_cost is None + sent: Final = route.calls.last.request + assert sent.headers.get("authorization") == ("Bearer dapi-configured" if configured_key else "Bearer dapi-env-token") + assert set(json.loads(sent.content)) == {"state", "questions"} + assert len(recorder.calls) == 1 + + 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) @@ -618,7 +670,7 @@ def test_databricks_without_a_base_or_key_fails_at_client_build(monkeypatch: pyt litellm.Router(model_list=[]), { "classifier_type": "oss_classifier", - "opensource_classifier_config": {"provider": "databricks", "model": "databricks-openjev-qwen35-4b"}, + "opensource_classifier_config": {"provider": "databricks", "model": "ai_decide"}, "tiers": {"SIMPLE": "cheap"}, }, derive_savings_baseline=False, @@ -665,17 +717,12 @@ def test_jev_api_base_without_its_own_key_is_rejected_so_the_environment_key_sta 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"} + {"provider": "databricks", "model": "ai_decide", "api_base": "https://collector.invalid"} ) paired: Final = JevClassifierConfig.model_validate( - { - "provider": "databricks", - "model": "my-openjev", - "api_base": "https://workspace.invalid/serving-endpoints", - "api_key": "dapi-own", - } + {"provider": "databricks", "model": "ai_decide", "api_base": "https://workspace.invalid", "api_key": "dapi-own"} ) - assert (paired.api_base, paired.api_key) == ("https://workspace.invalid/serving-endpoints", "dapi-own") + assert (paired.api_base, paired.api_key) == ("https://workspace.invalid", "dapi-own") @pytest.mark.parametrize( diff --git a/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx b/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx index 9deb53753bd..edbd4d50ee5 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx @@ -56,7 +56,7 @@ const ClassifierTypeRadios: React.FC = ({ value, clas OSS Classifier{" "} - uses Jev, Laya, Bespoke, or a Databricks endpoint to decide the tier + uses Jev, Laya, Bespoke, or Databricks 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 3480a8e922a..f829a2f2e41 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 @@ -125,7 +125,7 @@ describe("JEV classifier editor", () => { expect(screen.getByLabelText("Classifier Model")).toHaveValue("jev-latest"); await chooseSelectOption(userEvent, screen.getByRole("combobox", { name: "Decision model" }), label); if (provider === "bespoke") expect(screen.getByLabelText("Classifier Model")).toHaveTextContent("nimble-latest"); - if (provider === "databricks") expect(screen.getByLabelText("Classifier Model")).toHaveValue(""); + if (provider === "databricks") expect(screen.getByLabelText("Classifier Model")).toHaveValue("ai_decide"); if (typesTheModel) { fireEvent.change(screen.getByLabelText("Classifier Model"), { target: { value: model } }); } else { diff --git a/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx index 31f29dd9e9d..4b3cfa21f9d 100644 --- a/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx @@ -15,7 +15,7 @@ const providerDescriptions = { 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.", + "Uses the Databricks ai_decide AI Function, or a decision serving endpoint you name, with your configured tiers. Set DATABRICKS_API_BASE to https:///serving-endpoints and DATABRICKS_API_KEY on the gateway.", }; export default function JevClassifierConfig({ @@ -59,7 +59,7 @@ export default function JevClassifierConfig({ value={config.model} placeholder={ config.provider === "databricks" - ? "Serving endpoint name, e.g. databricks-openjev-qwen35-4b" + ? "ai_decide, or a serving endpoint name such as databricks-openjev-qwen35-4b" : undefined } onChange={(event) => 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 9ce99d39564..5cfedba7b1e 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 @@ -54,6 +54,8 @@ describe("buildAutoRouterRoutingTestRequest", () => { ["bespoke", "bespokelabs/Bespoke-Nimble-9B", "json"], ["databricks", "databricks-openjev-qwen35-4b", "object"], ["databricks", "databricks-openjev-qwen35-4b", "json"], + ["databricks", "ai_decide", "object"], + ["databricks", "ai_decide", "json"], ])("probes saved %s/%s %s configuration with custom tiers and team context", (provider, model, format) => { const config = { classifier_type: "oss_classifier", 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 23dedbbc860..89d419d0487 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 @@ -163,6 +163,7 @@ describe("buildComplexityRouterConfig", () => { ["bespoke", "nimble"], ["bespoke", "bespokelabs/Bespoke-Nimble-9B"], ["databricks", "databricks-openjev-qwen35-4b"], + ["databricks", "ai_decide"], ["jev", "custom-jev-model"], [undefined, "custom-jev-model"], ] as const)("accepts %s model %s before saving or testing", (provider, model) => { 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 index 109e780cc68..06298817073 100644 --- 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 @@ -27,11 +27,11 @@ describe("hydrateOssClassifier", () => { }); describe("jevClassifierConfigSchema", () => { - it("starts a Databricks classifier with an empty endpoint name that does not pass until typed", () => { + it("starts a Databricks classifier on ai_decide and still takes a typed serving endpoint name", () => { const fresh = defaultJevClassifierConfig("databricks"); - expect(fresh).toEqual({ provider: "databricks", model: "", timeout_ms: 3000 }); - expect(jevClassifierConfigSchema.safeParse(fresh).success).toBe(false); + expect(fresh).toEqual({ provider: "databricks", model: "ai_decide", timeout_ms: 3000 }); + expect(jevClassifierConfigSchema.safeParse(fresh).success).toBe(true); expect(jevClassifierConfigSchema.safeParse({ ...fresh, model: "my-openjev" }).success).toBe(true); }); 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 a74247f2060..b75b2dfc314 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 @@ -13,19 +13,19 @@ export const isOssClassifierProvider = (value: unknown): value is OssClassifierP const DEFAULT_TIMEOUT_MS = 3000; const SERVING_ENDPOINT_NAME = /^[A-Za-z0-9_-][A-Za-z0-9._-]*$/; +const DATABRICKS_AI_DECIDE_MODEL = "ai_decide"; 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"); + fixedClassifierModels(provider)?.[0] ?? (provider === "databricks" ? DATABRICKS_AI_DECIDE_MODEL : "jev-latest"); const isSupportedClassifierModel = (provider: OssClassifierProvider | undefined, model: string): boolean => { - if (provider === "databricks") return SERVING_ENDPOINT_NAME.test(model); + if (provider === "databricks") return model === DATABRICKS_AI_DECIDE_MODEL || SERVING_ENDPOINT_NAME.test(model); const fixed = fixedClassifierModels(provider); return fixed === undefined || fixed.some((candidate) => candidate === model); }; - const jevClassifierConfigFields = { deployment_name: z.preprocess((value) => value ?? undefined, z.string().trim().min(1).optional()), provider: z.preprocess( @@ -49,7 +49,7 @@ export const jevClassifierConfigSchema = z .refine( (config) => config.deployment_name !== undefined || isSupportedClassifierModel(config.provider, config.model), { - error: "Select a supported classifier model, or enter the bare Databricks serving endpoint name", + error: "Select a supported classifier model, or enter ai_decide or a bare Databricks serving endpoint name", path: ["model"], }, ); @@ -57,9 +57,7 @@ export const jevClassifierConfigSchema = z export type JevClassifierConfig = z.infer; export const defaultJevClassifierConfig = (provider: OssClassifierProvider = "jev"): JevClassifierConfig => - provider === "databricks" - ? { provider, model: "", timeout_ms: DEFAULT_TIMEOUT_MS } - : jevClassifierConfigSchema.parse({ provider }); + jevClassifierConfigSchema.parse({ provider }); export const hydrateOssClassifier = (config: { classifier_type?: ClassifierType | "oss_classifier";