mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
fix(databricks): preserve serving endpoint decisions
(cherry picked from commit cd081a2623)
This commit is contained in:
mateo-berri
mateo
parent
574dedf970
commit
0aa586115a
22 changed files with 644 additions and 138 deletions
|
|
@ -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://<workspace-host>"
|
||||
"https://<workspace-host>/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://<workspace-host>"
|
||||
"DATABRICKS_API_BASE pointing to https://<workspace-host>/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()
|
||||
|
|
|
|||
|
|
@ -21165,6 +21165,18 @@
|
|||
"output_vector_size": 1024,
|
||||
"source": "https://www.databricks.com/product/pricing/foundation-model-serving"
|
||||
},
|
||||
"databricks/databricks-openjev-qwen35-4b": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 0.0,
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "databricks",
|
||||
"metadata": {
|
||||
"notes": "Pay-per-token Foundation Model API decision endpoint (OpenJev, Qwen 3.5 4B) queried through the System One API. Databricks published no DBU rate for it as of 2026-10-07, so a deployment's model_info prices set the billed cost."
|
||||
},
|
||||
"mode": "evaluation",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models"
|
||||
},
|
||||
"databricks/ai_decide": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 0.0,
|
||||
|
|
|
|||
|
|
@ -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: "
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -21165,6 +21165,18 @@
|
|||
"output_vector_size": 1024,
|
||||
"source": "https://www.databricks.com/product/pricing/foundation-model-serving"
|
||||
},
|
||||
"databricks/databricks-openjev-qwen35-4b": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 0.0,
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "databricks",
|
||||
"metadata": {
|
||||
"notes": "Pay-per-token Foundation Model API decision endpoint (OpenJev, Qwen 3.5 4B) queried through the System One API. Databricks published no DBU rate for it as of 2026-10-07, so a deployment's model_info prices set the billed cost."
|
||||
},
|
||||
"mode": "evaluation",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models"
|
||||
},
|
||||
"databricks/ai_decide": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 0.0,
|
||||
|
|
|
|||
|
|
@ -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) == ()
|
||||
|
||||
|
|
|
|||
|
|
@ -147,7 +147,47 @@ def test_evaluation_mode_health_check_of_the_self_hosted_strands_model_resolves_
|
|||
},
|
||||
)
|
||||
]
|
||||
_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}}},
|
||||
|
|
@ -160,7 +200,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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
@ -279,6 +355,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)
|
||||
|
|
@ -324,18 +433,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",
|
||||
|
|
|
|||
|
|
@ -106,6 +106,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]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -69,6 +69,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": {
|
||||
|
|
@ -769,18 +783,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?"}},
|
||||
)
|
||||
|
|
@ -807,6 +823,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",
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -55,9 +55,7 @@ const ClassifierTypeRadios: React.FC<ClassifierTypeRadiosProps> = ({ value, clas
|
|||
<RadioGroupItem value="jev" className="mt-0.5" />
|
||||
<span>
|
||||
<strong className="font-semibold">OSS Classifier</strong>{" "}
|
||||
<span className="text-muted-foreground">
|
||||
uses Jev, Laya, Bespoke, or Databricks ai_decide to decide the tier
|
||||
</span>
|
||||
<span className="text-muted-foreground">uses Jev, Laya, Bespoke, or Databricks to decide the tier</span>
|
||||
</span>
|
||||
</Label>
|
||||
<SimpleTooltip content={scorerLockedReason}>
|
||||
|
|
|
|||
|
|
@ -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(<Form />);
|
||||
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();
|
||||
|
|
|
|||
|
|
@ -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://<workspace-host>/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<typeof config>) =>
|
||||
onChange({ ...value, jev_classifier_config: { ...config, ...patch } });
|
||||
|
||||
|
|
@ -50,7 +50,16 @@ export default function JevClassifierConfig({
|
|||
</SelectContent>
|
||||
</Select>
|
||||
) : (
|
||||
<Input id={`${id}-model`} value={config.model} onChange={(event) => update({ model: event.target.value })} />
|
||||
<Input
|
||||
id={`${id}-model`}
|
||||
value={config.model}
|
||||
placeholder={
|
||||
config.provider === "databricks"
|
||||
? "ai_decide, or a serving endpoint name such as databricks-openjev-qwen35-4b"
|
||||
: undefined
|
||||
}
|
||||
onChange={(event) => update({ model: event.target.value })}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
<div>
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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<typeof jevClassifierConfigSchema>;
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue