diff --git a/litellm/llms/databricks/decisions/transformation.py b/litellm/llms/databricks/decisions/transformation.py index d9bd4f9ab1f..c2ac1f220b8 100644 --- a/litellm/llms/databricks/decisions/transformation.py +++ b/litellm/llms/databricks/decisions/transformation.py @@ -1,3 +1,4 @@ +import re from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType @@ -9,15 +10,23 @@ 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 validate_ai_decide_model(model: str) -> str: - if model != DATABRICKS_AI_DECIDE_MODEL: - raise ValueError( - f"Databricks decisions model must be {DATABRICKS_AI_DECIDE_MODEL!r} (the ai_decide AI Function)" - ) - return model +def is_ai_decide_model(model: str) -> bool: + return model == DATABRICKS_AI_DECIDE_MODEL + + +def validate_databricks_decisions_model(model: str) -> str: + if is_ai_decide_model(model) or _SERVING_ENDPOINT_NAME.fullmatch(model) is not None: + return model + raise ValueError( + f"Databricks decisions model must be {DATABRICKS_AI_DECIDE_MODEL!r} (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" + ) def databricks_workspace_host(api_base: str) -> str: @@ -61,21 +70,22 @@ class DatabricksDecisionsConnection: class DatabricksDecisionsConfig(BaseDecisionsConfig): - path = "/api/2.0/ai-functions/ai-decide" 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://" + "https:///serving-endpoints" ) def canonical_model(self, model: str) -> str: - return validate_ai_decide_model(model) + return validate_databricks_decisions_model(model) def get_complete_url(self, api_base: str, model: str) -> str: - return f"{databricks_workspace_host(api_base)}{self.path}" + 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, @@ -84,7 +94,7 @@ class DatabricksDecisionsConfig(BaseDecisionsConfig): custom_llm_provider: str, ) -> Mapping[str, object] | UnsupportedDecisionsRequest: body: Final = super().transform_decisions_request(model, request, custom_llm_provider) - if isinstance(body, UnsupportedDecisionsRequest): + 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"}) @@ -101,9 +111,9 @@ class DatabricksDecisionsConfig(BaseDecisionsConfig): 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://" + "DATABRICKS_API_BASE pointing to https:///serving-endpoints" ) - return DatabricksDecisionsConnection(api_base=databricks_workspace_host(base), api_key=key) + 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 b5cad20d326..6d9a0f60bc9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21147,6 +21147,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" + }, "databricks/ai_decide": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 4bdeb035610..c0a87b95d7e 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -764,15 +764,15 @@ class OpenSourceClassifierConfig(LiteLLMBaseModel): if self.provider == "databricks": from litellm.llms.databricks.decisions.transformation import ( DATABRICKS_AI_DECIDE_MODEL, - validate_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': set it to " - f"{DATABRICKS_AI_DECIDE_MODEL!r}" + "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_ai_decide_model(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 a9a87d88f70..d746c360f06 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -18,7 +18,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, @@ -107,7 +107,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", exclude=_WITHOUT_MODEL if self._provider == "databricks" else None), + 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, ) @@ -120,6 +120,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 b5cad20d326..6d9a0f60bc9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21147,6 +21147,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" + }, "databricks/ai_decide": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, diff --git a/tests/integration/management/test_auto_router_databricks_classifier_writes.py b/tests/integration/management/test_auto_router_databricks_classifier_writes.py index 16b3eeca5d2..2b533b03c6e 100644 --- a/tests/integration/management/test_auto_router_databricks_classifier_writes.py +++ b/tests/integration/management/test_auto_router_databricks_classifier_writes.py @@ -5,6 +5,7 @@ from collections.abc import Callable, Mapping, Sequence 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.openai_wire import answering_model_discovery, chat_reply @@ -13,12 +14,18 @@ from pydantic import JsonValue _MODEL: Final = "ai_decide" _ROUTE: Final = "/api/2.0/ai-functions/ai-decide" +_ENDPOINT: Final = "databricks-openjev-qwen35-4b" _SECOND_WORKSPACE: Final = "/second-workspace" _API_KEY: Final = "synthetic-databricks-key" -_MODEL_RULE: Final = "Databricks decisions model must be 'ai_decide' (the ai_decide AI Function)" +_MODEL_RULE: Final = ( + "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" +) _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': set it to 'ai_decide'" + "opensource_classifier_config.model is required for provider 'databricks': 'ai_decide' or a 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 " @@ -27,18 +34,19 @@ _BASE_NEEDS_KEY: Final = ( _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]] = { - "response": { - "answers": { - "tier": { - "type": "choice", - "choice": "COMPLEX", - "confidence": 0.91, - "probabilities": {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91}, - } - } - }, - "metadata": {"version": "1.0"}, +_ANSWERS: Final[dict[str, JsonValue]] = { + "tier": { + "type": "choice", + "choice": "COMPLEX", + "confidence": 0.91, + "probabilities": {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91}, + } +} +_AI_DECIDE_ANSWER: Final[dict[str, JsonValue]] = {"response": {"answers": _ANSWERS}, "metadata": {"version": "1.0"}} +_ENDPOINT_ANSWER: Final[dict[str, JsonValue]] = { + "model": "/mosaicml/local_model", + "answers": _ANSWERS, + "usage": {"input_tokens": 12, "output_tokens": 1}, } _ROWS_QUERY: Final = ( "SELECT request_id, status, metadata->>'internal_call_origin' AS origin " @@ -47,7 +55,8 @@ _ROWS_QUERY: Final = ( def _classifier(request: Request) -> Reply: - return Reply(body=json.dumps(_CLASSIFIER_ANSWER).encode()) + answer: Final = _ENDPOINT_ANSWER if request.target.endswith("/invocations") else _AI_DECIDE_ANSWER + return Reply(body=json.dumps(answer).encode()) def _tier(text: str) -> Callable[[Request], Reply]: @@ -155,6 +164,16 @@ def test_validate_accepts_a_databricks_ai_decide_classifier_with_or_without_the_ assert _validate(gateway, config) == {"valid": True, "error": None}, api_base +def test_validate_accepts_a_databricks_classifier_with_a_bare_endpoint_name_of_any_length(gateway: Gateway) -> None: + for model in (_ENDPOINT, "e" * 5000): + config: Final = _router_config( + _classifier_config("https://dbc-example.cloud.databricks.com/serving-endpoints", model=model), + "simple", + "complex", + ) + assert _validate(gateway, config) == {"valid": True, "error": None}, model[:40] + + 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], ...]] = ( @@ -164,12 +183,9 @@ def test_validate_names_each_databricks_classifier_misconfiguration(gateway: Gat _without(base, "api_key"), f"at opensource_classifier_config: Value error, {_BASE_NEEDS_KEY}", ), - ( - "serving endpoint name", - {**base, "model": "databricks-openjev-qwen35-4b"}, - _MODEL_RULE, - ), - ("dash", {**base, "model": "ai-decide"}, _MODEL_RULE), + ("slash in the name", {**base, "model": "a/b"}, _MODEL_RULE), + ("space in the name", {**base, "model": "a b"}, _MODEL_RULE), + ("leading dot", {**base, "model": ".hidden"}, _MODEL_RULE), ("empty name", {**base, "model": ""}, _MODEL_RULE), ( "blank api_key", @@ -193,8 +209,13 @@ def test_validate_names_each_databricks_classifier_misconfiguration(gateway: Gat assert error.endswith(_REJECTED_AT_LOAD), (label, error) -def test_test_routing_classifies_through_databricks_ai_decide_without_calling_the_tier_it_picked( - gateway: Gateway, +@pytest.mark.parametrize( + ("model", "target", "cost"), + [(_MODEL, _ROUTE, {}), (_ENDPOINT, f"/{_ENDPOINT}/invocations", {"classifier_cost": 0.0})], + ids=["ai_decide", "serving_endpoint"], +) +def test_test_routing_classifies_through_databricks_without_calling_the_tier_it_picked( + gateway: Gateway, model: str, target: str, cost: dict[str, JsonValue] ) -> None: prompt: Final = f"design a distributed cache {uuid.uuid4().hex}" with ( @@ -205,7 +226,7 @@ def test_test_routing_classifies_through_databricks_ai_decide_without_calling_th ): 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) + config: Final = _router_config(_classifier_config(judge.url, model=model), simple_model, complex_model) response: Final = gateway.request( "POST", "/auto_router/test_routing", {"prompt": prompt, "complexity_router_config": config} ) @@ -226,17 +247,18 @@ def test_test_routing_classifies_through_databricks_ai_decide_without_calling_th "tier-probability:MEDIUM=0.060000", "tier-probability:COMPLEX=0.910000", ], - "classifier_model": f"databricks/{_MODEL}", + "classifier_model": f"databricks/{model}", + **cost, "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 == _ROUTE, call.target + assert call.target == target, call.target assert call.headers.get("authorization") == f"Bearer {_API_KEY}", call.headers body: Final = json.loads(call.body) - assert "model" not in body, body + assert body.get("model") == (None if model == _MODEL else model), body assert body["state"] == f"\nClassify this message:\n{prompt}", body assert _judge_targets(simple) == () and _judge_targets(complex_tier) == () diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index a59f7f59133..1a6ece07db5 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -188,7 +188,47 @@ def test_evaluation_mode_health_check_of_hosted_vllm_sends_a_choice_probe(gatewa ] -_DATABRICKS_PROBE_REPLY: Final = JsonResponse( +_DATABRICKS_ENDPOINT_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_ENDPOINT_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?"}}, + }, + ) + ] + + +_DATABRICKS_AI_DECIDE_PROBE_REPLY: Final = JsonResponse( content_type="application/json", body={ "response": {"answers": {"reachable": {"type": "noul", "probability": 1.0}}}, @@ -201,7 +241,7 @@ def test_evaluation_mode_health_check_of_databricks_ai_decide_resolves_the_mode_ gateway: Gateway, ) -> None: with gateway.scenario() as scenario: - handle: Final = register_scenario(f"health-decisions-{uuid.uuid4().hex[:12]}", _DATABRICKS_PROBE_REPLY) + handle: Final = register_scenario(f"health-decisions-{uuid.uuid4().hex[:12]}", _DATABRICKS_AI_DECIDE_PROBE_REPLY) scenario.cleanups.callback(delete_scenario, handle) model: Final = scenario.model( model="databricks/ai_decide", diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py index 346103a2793..2409bea5f5f 100644 --- a/tests/integration/providers/test_databricks_decisions_wire.py +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -16,6 +16,7 @@ from pydantic import JsonValue _MODEL: Final = "ai_decide" _ROUTE: Final = "/api/2.0/ai-functions/ai-decide" +_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" @@ -37,6 +38,8 @@ _ANSWERS: Final[dict[str, JsonValue]] = { "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} _METADATA: Final[dict[str, JsonValue]] = {"version": "1.0"} _AI_DECIDE_BODY: Final[dict[str, JsonValue]] = {"response": {"answers": _ANSWERS}, "metadata": _METADATA} _SYSTEM_ONE_BODY: Final[dict[str, JsonValue]] = { @@ -45,17 +48,29 @@ _SYSTEM_ONE_BODY: Final[dict[str, JsonValue]] = { "usage": {"input_tokens": 0, "output_tokens": 0}, "metadata": _METADATA, } -_MODEL_RULE: Final = "Databricks decisions model must be 'ai_decide' (the ai_decide AI Function)" +_MODEL_RULE: Final = ( + "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" +) _MISSING_KEY: Final = "Missing API key for Decisions provider 'databricks'" -_OTHER_MODELS: Final[tuple[tuple[str, str], ...]] = ( - ("serving-endpoint", "databricks-openjev-qwen35-4b"), - ("dash", "ai-decide"), - ("upper-case", "AI_DECIDE"), +_NON_BARE_NAMES: Final[tuple[tuple[str, str], ...]] = ( + ("slash", "foo/bar"), + ("space", "a b"), + ("leading-dot", ".hidden"), + ("question-mark", "a?b"), + ("hash", "a#b"), ("route-path", "serving-endpoints/ai_decide"), ) -_DATABRICKS_ERRORS: Final[tuple[tuple[int, dict[str, JsonValue]], ...]] = ( - (400, {"error_code": "BAD_REQUEST", "message": "Invalid input: questions must not be empty"}), - (403, {"error_code": "PERMISSION_DENIED", "message": "ai_decide is not enabled for this workspace"}), +_DATABRICKS_ERRORS: Final[tuple[tuple[str, str, int, dict[str, JsonValue]], ...]] = ( + (_MODEL, _ROUTE, 400, {"error_code": "BAD_REQUEST", "message": "Invalid input: questions must not be empty"}), + (_MODEL, _ROUTE, 403, {"error_code": "PERMISSION_DENIED", "message": "ai_decide is not enabled for this workspace"}), + ( + _ENDPOINT, + f"/{_ENDPOINT}/invocations", + 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" ' @@ -192,8 +207,8 @@ def _assert_refused_before_any_upstream_call(gateway: Gateway, model: str, messa assert (row["status"], row["call_type"], _number(row["spend"])) == ("failure", "adecisions", 0.0), row -@pytest.mark.parametrize("name", [name for _, name in _OTHER_MODELS], ids=[label for label, _ in _OTHER_MODELS]) -def test_a_deployment_naming_a_model_other_than_ai_decide_is_refused_before_any_upstream_call( +@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}", _MODEL_RULE) @@ -203,6 +218,67 @@ def test_a_deployment_naming_no_model_is_refused_before_any_upstream_call(gatewa _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], + }, + ) + 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_an_openai_format_request_at_v1_decisions_reaches_the_ai_decide_route_without_a_model_field( gateway: Gateway, ) -> None: @@ -278,6 +354,39 @@ def test_a_noul_probability_reaches_the_caller_as_the_system_one_noul_answer(gat assert call["body"] == {"state": _STATE, "questions": {"defect": {"type": "noul", "instructions": "Bug?"}}} +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 + ) + + def test_a_wildcard_deployment_sends_each_ai_decide_request_to_the_route_and_logs_each_once(gateway: Gateway) -> None: with gateway.scenario() as scenario: handle: Final = _register(scenario, _AI_DECIDE_BODY) @@ -323,18 +432,45 @@ def test_a_non_string_model_is_refused_at_the_gateway_without_an_upstream_call( assert _upstream_calls(gateway, handle) == [] -@pytest.mark.parametrize("status,body", _DATABRICKS_ERRORS, ids=("400", "403")) +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( + "name,path,status,body", _DATABRICKS_ERRORS, ids=("ai_decide-400", "ai_decide-403", "serving_endpoint-404") +) def test_a_databricks_error_keeps_its_status_and_message_and_bills_nothing( - gateway: Gateway, status: int, body: dict[str, JsonValue] + gateway: Gateway, name: str, path: str, 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/{_MODEL}", api_base=handle.api_base(), api_key=_API_KEY) + model: Final = scenario.model(model=f"databricks/{name}", 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}{_ROUTE}", call + assert call["path"] == f"/{handle.scenario_id}{path}", call row: Final = _spend_row(response.headers["x-litellm-call-id"]) assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == ( "failure", diff --git a/tests/integration/providers/test_decisions_openai_format_wire.py b/tests/integration/providers/test_decisions_openai_format_wire.py index 30d072c2761..cc8b218ecd2 100644 --- a/tests/integration/providers/test_decisions_openai_format_wire.py +++ b/tests/integration/providers/test_decisions_openai_format_wire.py @@ -189,6 +189,15 @@ _PROVIDERS: Final = ( None, ai_decide=True, ), + _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 ), @@ -249,7 +258,7 @@ def _response_cost(response: httpx.Response) -> float: def _provider_id(provider: _Provider) -> str: - return provider.name + return provider.model def _number(value: JsonValue) -> float: diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index df205a80c38..16c938e3995 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -111,6 +111,15 @@ _PROVIDERS: Final = ( _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 index c78f7c113be..419b786bfc1 100644 --- a/tests/integration/routing/test_complexity_router_databricks_classifier.py +++ b/tests/integration/routing/test_complexity_router_databricks_classifier.py @@ -211,11 +211,14 @@ def _state(prompt: str, system: str | None) -> str: 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, key: str = _API_KEY) -> None: - assert call.target == _ROUTE, call.target +def _assert_classifier_call( + call: Request, *, prompt: str, system: str | None = None, key: str = _API_KEY, endpoint: str | None = None +) -> None: + assert call.target == (_ROUTE if endpoint is None else f"/{endpoint}/invocations"), call.target assert call.headers.get("authorization") == f"Bearer {key}", call.headers assert json.loads(call.body) == { "state": _state(prompt, system), + **({} if endpoint is None else {"model": endpoint}), "questions": {"tier": {"type": "choice", "instructions": DEFAULT_JEV_INSTRUCTIONS, "criteria": _TIER_CRITERIA}}, }, call.body @@ -622,6 +625,37 @@ def test_messages_stream_over_httpx_is_classified_and_logged_under_the_message_s _assert_logged_once(_rows(routed.name, want=2), route="/v1/messages", request_id=identity) +def test_a_serving_endpoint_name_is_classified_through_its_invocations_route_and_logged_at_zero( + gateway: Gateway, +) -> None: + prompt: Final = _prompt() + endpoint: Final = "my-jev-endpoint" + with ( + _wires(_answering({"model": "/mosaicml/local_model", "answers": _ANSWERS, "usage": _USAGE})) 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) + (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: diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index ccb39d38b02..cb8043e5fab 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -70,6 +70,20 @@ _STRANDS_RESPONSE: Final[Mapping[str, object]] = { "usage": {"input_tokens": 216, "output_tokens": 3}, "latency_ms": 3722.17, } +_DATABRICKS_ENDPOINT: Final = "databricks-openjev-qwen35-4b" +_DATABRICKS_ENDPOINT_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}, +} _DATABRICKS_ROUTE: Final = "https://workspace.example/api/2.0/ai-functions/ai-decide" _DATABRICKS_RESPONSE: Final[Mapping[str, object]] = { "response": { @@ -903,18 +917,20 @@ async def test_databricks_requires_api_base_before_http( @pytest.mark.asyncio -@pytest.mark.parametrize("model", ["databricks-openjev-qwen35-4b", "serving-endpoints/ai_decide", "ai-decide"]) -async def test_databricks_rejects_a_model_other_than_ai_decide_before_http( +@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, - model: str, + endpoint_name: str, ) -> None: - monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.example") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.example/serving-endpoints") monkeypatch.setenv("DATABRICKS_API_KEY", "dapi-key") - with pytest.raises(litellm.BadRequestError, match="'ai_decide'"): + with pytest.raises(litellm.BadRequestError, match="bare serving endpoint name"): await litellm.adecisions( - model=f"databricks/{model}", + model=f"databricks/{endpoint_name}", state="review", questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, ) @@ -941,6 +957,47 @@ async def test_databricks_requires_a_key_before_http( 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_ENDPOINT_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", diff --git a/tests/unit/llms/databricks/decisions/test_transformation.py b/tests/unit/llms/databricks/decisions/test_transformation.py index f683747f316..f6ea3444db1 100644 --- a/tests/unit/llms/databricks/decisions/test_transformation.py +++ b/tests/unit/llms/databricks/decisions/test_transformation.py @@ -11,34 +11,55 @@ from litellm.types.decisions import ( ) _HOST: Final = "https://workspace.example" -_ROUTE: Final = f"{_HOST}/api/2.0/ai-functions/ai-decide" +_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", [_HOST, f"{_HOST}/", f"{_HOST}/serving-endpoints", f"{_HOST}/serving-endpoints/"]) -def test_the_complete_url_is_the_workspace_ai_decide_route(api_base: str) -> None: - assert _CONFIG.get_complete_url(api_base, "ai_decide") == _ROUTE +@pytest.mark.parametrize("api_base", [_BASE, f"{_BASE}/"]) +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("model", ["", "databricks-openjev-qwen35-4b", "ai-decide", "AI_DECIDE", "ai_decide/x"]) -def test_a_model_other_than_ai_decide_is_rejected_before_any_request(model: str) -> None: - with pytest.raises(ValueError, match="'ai_decide'"): - _ = _CONFIG.canonical_model(model) +@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 -def test_the_request_body_omits_model_because_the_route_rejects_unexpected_parameters() -> None: +@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 serving endpoint name"): + _ = _CONFIG.get_complete_url(_BASE, 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": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, - } + assert body == {"state": "export hangs", "questions": _QUESTIONS} -def test_the_response_envelope_is_unwrapped_and_noul_probability_becomes_the_noul_answer() -> None: +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"}, @@ -49,31 +70,50 @@ def test_the_response_envelope_is_unwrapped_and_noul_probability_becomes_the_nou 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") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://env.example/serving-endpoints") monkeypatch.setenv("DATABRICKS_API_KEY", "dapi-env") - connection: Final = _CONFIG.connection(f"{_HOST}/serving-endpoints/", "dapi-configured") - assert (connection.api_base, connection.api_key) == (_HOST, "dapi-configured") + 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", _HOST) + 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) == (_HOST, "dapi-token") + 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", _HOST) + 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", None) + _ = _CONFIG.connection("https://collector.example/serving-endpoints", None) -def test_classifier_response_unwraps_the_envelope_and_accounts_the_configured_model() -> 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" + + +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") diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index eb1bb3c55e9..9416eb85d8a 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_systemone_dispatches_databricks_ai_decide( 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 7286dcd4f7d..4416e0dfab1 100644 --- a/tests/unit/proxy/management_helpers/test_auto_router_permissions.py +++ b/tests/unit/proxy/management_helpers/test_auto_router_permissions.py @@ -168,7 +168,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", "ai_decide")], + [ + ("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 98f62be8ca9..256757a3b65 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -457,6 +457,7 @@ def test_jev_config_requires_classifier_config() -> None: ("jev", "jev-latest", "jev"), ("laya", "english", "laya"), ("bespoke", "nimble-latest", "bespoke"), + ("databricks", "databricks-openjev-qwen35-4b", "databricks"), ("databricks", "ai_decide", "databricks"), ], ) @@ -540,11 +541,70 @@ async def test_oss_routes_with_its_own_credentials_and_accounts_the_checkpoint( assert recorder.calls[0]["response_cost"] == pytest.approx(0.31) -@pytest.mark.parametrize("model", [None, "databricks-openjev-qwen35-4b", "serving-endpoints/ai_decide", "ai-decide"]) -def test_databricks_accepts_only_ai_decide_as_the_model(model: str | None) -> None: - with pytest.raises(ValueError, match="'ai_decide'"): - JevClassifierConfig.model_validate({"provider": "databricks", **({"model": model} if model is not None else {})}) - assert JevClassifierConfig.model_validate({"provider": "databricks", "model": "ai_decide"}).model == "ai_decide" +@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 serving endpoint name"): + JevClassifierConfig.model_validate({"provider": "databricks", "model": "serving-endpoints/my-openjev"}) + with pytest.raises(ValueError, match="bare serving endpoint name"): + JevClassifierConfig.model_validate({"provider": "databricks", "model": ".."}) + + +@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) @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx b/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx index 8181ffa380a..6ed46e2a27d 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassifierTypeRadios.tsx @@ -55,9 +55,7 @@ const ClassifierTypeRadios: React.FC = ({ value, clas OSS Classifier{" "} - - uses Jev, Laya, Bespoke, or Databricks ai_decide 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 bb2556d9c94..1252b0f24d9 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,10 +101,11 @@ describe("JEV classifier editor", () => { ["jev", "Jev", "jev-test"], ["laya", "Laya", "multilingual"], ["bespoke", "Bespoke Nimble", "bespokelabs/Bespoke-Nimble-9B"], - ["databricks", "Databricks", "ai_decide"], + ["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(); @@ -124,11 +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 === "databricks") expect(screen.getByLabelText("Classifier Model")).toHaveTextContent("ai_decide"); - if (provider !== "jev") { - await chooseSelectOption(userEvent, screen.getByLabelText("Classifier Model"), model); - } else { + if (provider === "databricks") expect(screen.getByLabelText("Classifier Model")).toHaveValue("ai_decide"); + 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" } }); @@ -138,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 897c02905cd..9252891d7b5 100644 --- a/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx @@ -7,7 +7,7 @@ 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", @@ -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 the Databricks ai_decide AI Function with your configured tiers. Set DATABRICKS_API_BASE to your workspace URL and DATABRICKS_API_KEY on the gateway.", + "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({ @@ -27,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 } }); @@ -50,7 +50,16 @@ 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 950c26011da..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 @@ -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"], ["databricks", "ai_decide", "object"], ["databricks", "ai_decide", "json"], ])("probes saved %s/%s %s configuration with custom tiers and team context", (provider, model, format) => { @@ -76,8 +78,8 @@ describe("buildAutoRouterRoutingTestRequest", () => { it.each([ ["laya", "unsupported"], ["bespoke", "unsupported"], - ["databricks", "databricks-openjev-qwen35-4b"], - ["databricks", "serving-endpoints/ai_decide"], + ["databricks", "serving-endpoints/openjev"], + ["databricks", ".."], ])("does not probe unsupported %s model %s", (provider, model) => { 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 1bd3d9f360e..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 @@ -141,8 +141,9 @@ describe("buildComplexityRouterConfig", () => { { model: " " }, { provider: "laya" as const, model: "unsupported" }, { provider: "bespoke" as const, model: "unsupported" }, - { provider: "databricks" as const, model: "databricks-openjev-qwen35-4b" }, - { provider: "databricks" as const, model: "serving-endpoints/ai_decide" }, + { 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 }, @@ -161,6 +162,7 @@ describe("buildComplexityRouterConfig", () => { ["bespoke", "nimble-latest"], ["bespoke", "nimble"], ["bespoke", "bespokelabs/Bespoke-Nimble-9B"], + ["databricks", "databricks-openjev-qwen35-4b"], ["databricks", "ai_decide"], ["jev", "custom-jev-model"], [undefined, "custom-jev-model"], @@ -180,8 +182,8 @@ describe("buildComplexityRouterConfig", () => { ["laya", "english", true], ["bespoke", "nimble-latest", false], ["bespoke", "nimble-latest", true], - ["databricks", "ai_decide", false], - ["databricks", "ai_decide", 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 index 20365e613e8..324a23d8c3a 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 @@ -5,7 +5,7 @@ describe("hydrateOssClassifier", () => { it("keeps a stored Databricks classifier instead of falling back to Jev", () => { const stored = { provider: "databricks", - model: "ai_decide", + model: "databricks-openjev-qwen35-4b", timeout_ms: 4200, circuit_breaker_enabled: true, circuit_breaker_cooldown_seconds: 50, @@ -19,7 +19,7 @@ describe("hydrateOssClassifier", () => { 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: "databricks-openjev-qwen35-4b" }, + opensource_classifier_config: { provider: "databricks", model: "serving-endpoints/openjev" }, }); expect(hydrated.jev_classifier_config).toEqual(defaultJevClassifierConfig()); @@ -27,15 +27,16 @@ describe("hydrateOssClassifier", () => { }); describe("jevClassifierConfigSchema", () => { - it("starts a Databricks classifier on ai_decide, which passes as is", () => { + 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: "ai_decide", timeout_ms: 3000 }); expect(jevClassifierConfigSchema.safeParse(fresh).success).toBe(true); + expect(jevClassifierConfigSchema.safeParse({ ...fresh, model: "my-openjev" }).success).toBe(true); }); - it.each(["databricks-openjev-qwen35-4b", "serving-endpoints/ai_decide", "ai-decide", "AI_DECIDE"])( - "rejects the Databricks model %j because the ai-decide route serves only ai_decide", + 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 d7ecf9a7aa7..6866bddc205 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 @@ -4,7 +4,6 @@ import type { ClassifierType } from "./classifier_types"; export const OSS_CLASSIFIER_MODELS = { laya: ["english", "multilingual", "typed-decisions"], bespoke: ["nimble-latest", "nimble", "bespokelabs/Bespoke-Nimble-9B"], - databricks: ["ai_decide"], } as const; const OSS_CLASSIFIER_PROVIDERS = ["jev", "laya", "bespoke", "databricks"] as const; @@ -12,13 +11,29 @@ 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._-]*$/; +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" ? DATABRICKS_AI_DECIDE_MODEL : "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", "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() @@ -30,19 +45,11 @@ 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 ai_decide or a bare Databricks serving endpoint name", + path: ["model"], + }); export type JevClassifierConfig = z.infer;