diff --git a/litellm/__init__.py b/litellm/__init__.py index 0b304919727..c19b125c7c6 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -441,6 +441,10 @@ autorouter_presets_url: str = os.getenv( "LITELLM_AUTOROUTER_PRESETS_URL", "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json", ) +whats_new_url: str = os.getenv( + "LITELLM_WHATS_NEW_URL", + "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/whats_new.json", +) suppress_debug_info: bool = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None diff --git a/litellm/llms/databricks/decisions/transformation.py b/litellm/llms/databricks/decisions/transformation.py index a454da83e50..c2ac1f220b8 100644 --- a/litellm/llms/databricks/decisions/transformation.py +++ b/litellm/llms/databricks/decisions/transformation.py @@ -4,18 +4,63 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter, ValidationError + from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig +from litellm.types.decisions import DecisionsIRRequest, UnsupportedDecisionsRequest +DATABRICKS_AI_DECIDE_MODEL: Final = "ai_decide" +_AI_DECIDE_PATH: Final = "/api/2.0/ai-functions/ai-decide" _SERVING_ENDPOINT_NAME: Final = re.compile(r"[A-Za-z0-9_-][A-Za-z0-9._-]*") +_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) -def validate_serving_endpoint_name(name: str) -> str: - if _SERVING_ENDPOINT_NAME.fullmatch(name) is None: - raise ValueError( - f"Databricks serving endpoint name {name!r} must be the bare endpoint name (letters, digits, '-', '_' " - "and '.', with no '/', '?', '#', spaces, or a leading '.'), e.g. databricks-openjev-qwen35-4b" - ) - return name +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: + return api_base.rstrip("/").removesuffix("/serving-endpoints") + + +def _mapping(value: object) -> Mapping[str, object] | None: + try: + return _MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _systemone_answer(answer: object) -> object: + fields: Final = _mapping(answer) + if fields is None or fields.get("type") != "noul" or "probability" not in fields: + return answer + return MappingProxyType( + {**{key: value for key, value in fields.items() if key != "probability"}, "noul": fields["probability"]} + ) + + +def ai_decide_response(payload: Mapping[str, object]) -> Mapping[str, object]: + response: Final = _mapping(payload.get("response")) + answers: Final = None if response is None else _mapping(response.get("answers")) + if response is None or answers is None: + return payload + return MappingProxyType( + { + **{key: value for key, value in payload.items() if key != "response"}, + **response, + "answers": MappingProxyType({key: _systemone_answer(answer) for key, answer in answers.items()}), + } + ) @dataclass(frozen=True, slots=True) @@ -35,13 +80,30 @@ class DatabricksDecisionsConfig(BaseDecisionsConfig): ) def canonical_model(self, model: str) -> str: - return validate_serving_endpoint_name(model) + return validate_databricks_decisions_model(model) def get_complete_url(self, api_base: str, model: str) -> str: - return f"{api_base.rstrip('/')}/{validate_serving_endpoint_name(model)}/invocations" + if is_ai_decide_model(model): + return f"{databricks_workspace_host(api_base)}{_AI_DECIDE_PATH}" + return f"{api_base.rstrip('/')}/{validate_databricks_decisions_model(model)}/invocations" + + def transform_decisions_request( + self, + model: str, + request: DecisionsIRRequest, + custom_llm_provider: str, + ) -> Mapping[str, object] | UnsupportedDecisionsRequest: + body: Final = super().transform_decisions_request(model, request, custom_llm_provider) + if isinstance(body, UnsupportedDecisionsRequest) or not is_ai_decide_model(model): + return body + return MappingProxyType({key: value for key, value in body.items() if key != "model"}) + + def unwrap_response(self, payload: object) -> object: + fields: Final = _mapping(payload) + return payload if fields is None else ai_decide_response(fields) def classifier_response(self, body: Mapping[str, object], requested_model: str) -> Mapping[str, object]: - return MappingProxyType({**body, "model": requested_model}) + return MappingProxyType({**ai_decide_response(body), "model": requested_model}) def connection(self, api_base: str | None, api_key: str | None) -> DatabricksDecisionsConnection: base: Final = self.resolve_api_base(api_base) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 49432d4baa2..a1eaa6d4a90 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21196,6 +21196,18 @@ "output_vector_size": 1024, "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, + "databricks/ai_decide": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "databricks", + "metadata": { + "notes": "The ai_decide AI Function (Beta), called through POST /api/2.0/ai-functions/ai-decide. Its response carries no token usage, so LiteLLM records no spend for it." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.databricks.com/api/ai-functions/v1/ai-decide" + }, "databricks/databricks-openjev-qwen35-4b": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, @@ -21208,6 +21220,18 @@ "output_cost_per_token": 0.0, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" }, + "databricks/decisions/ai_decide": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "databricks", + "metadata": { + "notes": "The ai_decide AI Function (Beta), called through POST /api/2.0/ai-functions/ai-decide. Its response carries no token usage, so LiteLLM records no spend for it." + }, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://docs.databricks.com/api/ai-functions/v1/ai-decide" + }, "databricks/decisions/databricks-openjev-qwen35-4b": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py index 603b36d9406..a4428eb3a85 100644 --- a/litellm/proxy/management_helpers/auto_router_permissions.py +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -76,6 +76,7 @@ class _MemberOpenSourceClassifierConfig(LiteLLMBaseModel): model_config = ConfigDict(extra="forbid") + deployment_name: str | None = None provider: Literal["jev", "laya", "bespoke", "databricks"] = "jev" model: str api_key: None = None diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 47caacdcfff..75e8a5af6a5 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -6,8 +6,9 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from importlib.resources import files from typing import TYPE_CHECKING, Final, Protocol +import httpx from fastapi import APIRouter, HTTPException, Request -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -37,6 +38,7 @@ from litellm.types.proxy.public_endpoints.public_endpoints import ( ProviderCreateInfo, PublicModelHubInfo, SupportedEndpointsResponse, + WhatsNewResponse, ) from litellm.types.utils import LlmProviders @@ -505,14 +507,18 @@ def _load_bundled_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord]: ) -async def _fetch_remote_autorouter_presets(url: str) -> Mapping[str, AutoRouterPresetRecord]: +async def _fetch_remote_bytes(url: str) -> bytes: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.UI) response: Final = await client.get(url, timeout=5.0) response.raise_for_status() - presets: Final = _AUTOROUTER_PRESETS_ADAPTER.validate_python(response.json()) + return response.content + + +async def _fetch_remote_autorouter_presets(url: str) -> Mapping[str, AutoRouterPresetRecord]: + presets: Final = _AUTOROUTER_PRESETS_ADAPTER.validate_json(await _fetch_remote_bytes(url)) if not presets: raise ValueError("remote auto-router preset catalog is empty") return presets @@ -575,6 +581,70 @@ async def get_public_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord return await get_autorouter_presets(url=litellm.autorouter_presets_url) +def _load_bundled_whats_new() -> WhatsNewResponse: + return WhatsNewResponse.model_validate_json( + files("litellm.proxy.public_endpoints").joinpath("whats_new.json").read_text(encoding="utf-8") + ) + + +async def _fetch_remote_whats_new(url: str) -> WhatsNewResponse: + return WhatsNewResponse.model_validate_json(await _fetch_remote_bytes(url)) + + +async def _resolve_whats_new(url: str, fetch: Callable[[str], Awaitable[WhatsNewResponse]]) -> WhatsNewResponse: + if os.getenv("LITELLM_LOCAL_WHATS_NEW", "").lower() == "true": + return _load_bundled_whats_new() + try: + return await fetch(url) + except (httpx.HTTPError, ValidationError, OSError) as e: + verbose_logger.warning( + "LiteLLM: failed to fetch What's new launches from %s: %s. Serving the bundled list for the life of this process.", + url, + str(e), + ) + return _load_bundled_whats_new() + + +class _WhatsNewCache: + launches: WhatsNewResponse | None = None + lock: asyncio.Lock | None = None + + +async def get_whats_new( + url: str, + fetch: Callable[[str], Awaitable[WhatsNewResponse]] = _fetch_remote_whats_new, +) -> WhatsNewResponse: + cached: Final = _WhatsNewCache.launches + if cached is not None: + return cached + if _WhatsNewCache.lock is None: + _WhatsNewCache.lock = asyncio.Lock() + async with _WhatsNewCache.lock: + held: Final = _WhatsNewCache.launches + if held is not None: + return held + resolved: Final = await _resolve_whats_new(url=url, fetch=fetch) + _WhatsNewCache.launches = resolved + return resolved + + +@router.get( + "/public/whats_new", + tags=["public"], + response_model=WhatsNewResponse, +) +async def get_public_whats_new() -> WhatsNewResponse: + """ + Return the launches the dashboard Home page shows under What's new. + + Resolved once per process: fetched from ``litellm.whats_new_url`` (override with ``LITELLM_WHATS_NEW_URL``) + on the first request, falling back to the list bundled with the package on any failure, so an airgapped + proxy serves the bundled list without retrying. Set ``LITELLM_LOCAL_WHATS_NEW=True`` to skip the fetch. + A restart picks up a newly published list. + """ + return await get_whats_new(url=litellm.whats_new_url) + + @router.get( "/public/endpoints", tags=["public"], diff --git a/litellm/proxy/public_endpoints/whats_new.json b/litellm/proxy/public_endpoints/whats_new.json new file mode 100644 index 00000000000..40affb71e98 --- /dev/null +++ b/litellm/proxy/public_endpoints/whats_new.json @@ -0,0 +1,25 @@ +{ + "launches": [ + { + "icon": "box", + "title": "Claude Haiku 5.5", + "description": "Anthropic's newest small model, with day 0 pricing on Anthropic, Bedrock and Vertex AI.", + "href": "https://docs.litellm.ai/blog/claude-haiku-5-5", + "published_on": "2026-10-07" + }, + { + "icon": "scale", + "title": "Decision Models", + "description": "Call decision models at /v1/decisions and try them in the Decisions playground.", + "href": "https://docs.litellm.ai/docs/decisions", + "published_on": "2026-10-07" + }, + { + "icon": "zap", + "title": "OpenAI Ultrafast", + "description": "service_tier: ultrafast on GPT-6 Astra and GPT-6.1 Sol, including /ultrafast in Codex.", + "href": "https://docs.litellm.ai/docs/providers/openai/ultrafast", + "published_on": "2026-10-06" + } + ] +} diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a5f00aa7e25..7ef31f7a0a2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -123,6 +123,7 @@ from .jev_classifier import ( HttpJevClassifierClient, JevClassifierClient, JevVerdict, + RouterJevClassifierClient, build_jev_request, jev_classifier_cost, ) @@ -1331,8 +1332,9 @@ class ComplexityRouter(CustomLogger): - Question complexity (multiple questions) """ - @staticmethod - def _build_jev_client(config: OpenSourceClassifierConfig) -> JevClassifierClient: + def _build_jev_client(self, config: OpenSourceClassifierConfig) -> JevClassifierClient: + if config.deployment_name is not None: + return RouterJevClassifierClient(self.litellm_router_instance, config.deployment_name) if config.provider in ("laya", "bespoke"): from litellm.llms.oss_decision import oss_connection @@ -2318,13 +2320,23 @@ class ComplexityRouter(CustomLogger): if not self._tier_pools().get(tier_name): raise ValueError(f"Jev classifier returned tier {tier_name!r}, which has no models configured") model: Final = response.model or config.model - accounting_provider: Final = "typesafe" if config.provider == "jev" else config.provider + accounting_provider: Final = ( + response.classifier_provider + if config.deployment_name and response.classifier_provider + else "typesafe" + if config.provider == "jev" + else config.provider + ) verdict: Final = JevVerdict( label=answer.choice, probabilities=answer.probabilities, confidence=answer.confidence, model=model, - cost=jev_classifier_cost(response, config.model, accounting_provider), + cost=response.classifier_cost + if config.deployment_name + else jev_classifier_cost( + response, config.model, "typesafe" if config.provider == "jev" else config.provider + ), provider=accounting_provider, ) if breaker is not None and permit is not None: diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 36c19e37b8d..4a2c56daf00 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -716,6 +716,7 @@ _ENVIRONMENT_KEY_SCOPE: Final = MappingProxyType( class OpenSourceClassifierConfig(LiteLLMBaseModel): model_config = ConfigDict(extra="forbid", frozen=True) + deployment_name: str | None = Field(default=None, min_length=1) provider: Literal["jev", "laya", "bespoke", "databricks"] = "jev" model: str = "jev-latest" api_key: str | None = Field(default=None, description="Provider API key; optional for self-hosted providers") @@ -754,6 +755,12 @@ class OpenSourceClassifierConfig(LiteLLMBaseModel): @model_validator(mode="after") def _keep_the_environment_key_on_the_environment_base(self) -> "OpenSourceClassifierConfig": + if self.deployment_name is not None: + if not self.deployment_name.strip(): + raise ValueError("deployment_name must name a configured decision model") + if self.api_key is not None or self.api_base is not None: + raise ValueError("deployment_name uses the deployment connection; omit api_key and api_base") + return self if self.provider in ("laya", "bespoke"): from litellm.llms.oss_decision import validate_oss_api_base, validate_oss_model @@ -762,14 +769,17 @@ class OpenSourceClassifierConfig(LiteLLMBaseModel): _ = validate_oss_api_base(self.provider, self.api_base) return self if self.provider == "databricks": - from litellm.llms.databricks.decisions.transformation import validate_serving_endpoint_name + from litellm.llms.databricks.decisions.transformation import ( + DATABRICKS_AI_DECIDE_MODEL, + validate_databricks_decisions_model, + ) if "model" not in self.model_fields_set: raise ValueError( - "opensource_classifier_config.model is required for provider 'databricks': the serving endpoint " - "name, e.g. databricks-openjev-qwen35-4b" + "opensource_classifier_config.model is required for provider 'databricks': " + f"{DATABRICKS_AI_DECIDE_MODEL!r} or a serving endpoint name, e.g. databricks-openjev-qwen35-4b" ) - _ = validate_serving_endpoint_name(self.model) + _ = validate_databricks_decisions_model(self.model) if self.api_base is not None and self.api_key is None: raise ValueError( "opensource_classifier_config.api_base requires opensource_classifier_config.api_key: " @@ -833,7 +843,10 @@ def _resolve_normalized_complexity_router_config_write( { key: value for key, value in existing.items() - if same_provider and key in ("api_key", "api_base") and (key != "api_key" or same_base) + if same_provider + and not classifier.get("deployment_name") + and key in ("api_key", "api_base") + and (key != "api_key" or same_base) } ) return ComplexityRouterConfigWrite( diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index 65483b0a631..87bd031490a 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -1,7 +1,7 @@ from collections.abc import Mapping from datetime import datetime, timezone from types import MappingProxyType -from typing import Annotated, Final, Literal, NamedTuple, Protocol, TypeAlias +from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple, Protocol, TypeAlias from uuid import uuid4 import httpx @@ -10,6 +10,7 @@ from pydantic import ConfigDict, Field, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_router_logger from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.classifier_logging import masked_originating_request from litellm.litellm_core_utils.internal_call_metadata import ( effective_turn_off_message_logging, forwarded_internal_call_metadata, @@ -18,16 +19,21 @@ 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.llms.pass_through.typesafe_logging_handler import TypeSafePassthroughLoggingHandler from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS as _DEFAULT_JEV_INSTRUCTIONS +from litellm.types.decisions import DecisionsResponse from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN +if TYPE_CHECKING: + from litellm.router import Router + JevProbability: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)] ClassifierProvider: TypeAlias = Literal["typesafe", "laya", "bespoke", "databricks"] DEFAULT_JEV_INSTRUCTIONS: Final = _DEFAULT_JEV_INSTRUCTIONS +_WITHOUT_MODEL: Final = MappingProxyType({"model": True}) class JevChoiceQuestion(LiteLLMBaseModel): @@ -67,6 +73,8 @@ class JevSystemOneResponse(LiteLLMBaseModel): model: str | None = None answers: Mapping[str, JevChoiceAnswer] + classifier_cost: float | None = None + classifier_provider: str | None = None usage: JevUsage | None = None @@ -79,6 +87,76 @@ class JevClassifierClient(Protocol): ) -> JevSystemOneResponse: ... +class _DecisionsRouter(Protocol): + async def adecisions( + self, + *, + model: str, + state: str, + questions: Mapping[str, object], + timeout: float, + num_retries: int, + disable_fallbacks: bool, + metadata: Mapping[str, object], + proxy_server_request: Mapping[str, object], + turn_off_message_logging: bool | None, + litellm_session_id: str = "", + litellm_trace_id: str = "", + ) -> object: ... + + +_PARENT_METADATA: Final = TypeAdapter[Mapping[str, object] | None](Mapping[str, object] | None) + + +class RouterJevClassifierClient: + def __init__(self, router: "Router", deployment_name: str) -> None: + self._router: _DecisionsRouter = router + self._deployment_name = deployment_name + + async def evaluate( + self, + request: JevSystemOneRequest, + timeout_s: float, + request_kwargs: Mapping[str, object] | None = None, + ) -> JevSystemOneResponse: + parent: Final = request_kwargs or MappingProxyType({}) + metadata: Final = { + **forwarded_internal_call_metadata( + _PARENT_METADATA.validate_python(parent.get("litellm_metadata") or parent.get("metadata")), + AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + ), + INTERNAL_CALL_ORIGIN_METADATA_KEY: AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + } + payload: Final = { + "state": request.state, + "questions": {key: value.model_dump() for key, value in request.questions.items()}, + } + response: Final = DecisionsResponse.model_validate( + await self._router.adecisions( + model=self._deployment_name, + state=request.state, + questions={key: value.model_dump() for key, value in request.questions.items()}, + timeout=timeout_s, + num_retries=0, + disable_fallbacks=True, + metadata=metadata, + proxy_server_request={ + "originating_request_masked": masked_originating_request(request_kwargs), + "body": {"model": self._deployment_name, **payload}, + }, + turn_off_message_logging=effective_turn_off_message_logging(request_kwargs), + **parent_session_kwargs(request_kwargs), + ) + ) + return JevSystemOneResponse.model_validate( + { + **response.model_dump(), + "classifier_cost": response.hidden_params.get("response_cost"), + "classifier_provider": response.hidden_params.get("custom_llm_provider"), + } + ) + + class HttpJevClassifierClient: def __init__( self, @@ -104,7 +182,7 @@ class HttpJevClassifierClient: ) response: Final = await self._http_client.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler has a dynamic post signature self._request_url(request.model), - json=request.model_dump(mode="json"), + json=request.model_dump(mode="json", exclude=self._excluded_fields(request.model)), headers=MappingProxyType({**authorization, "Content-Type": "application/json"}), # pyright: ignore[reportArgumentType] # HTTP headers are not mutated by AsyncHTTPHandler timeout=timeout_s, ) @@ -117,6 +195,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) @@ -212,7 +293,7 @@ class JevVerdict(NamedTuple): confidence: float model: str cost: float | None - provider: ClassifierProvider = "typesafe" + provider: str = "typesafe" class _RegistryPricing(LiteLLMBaseModel): diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index 7ace4094cf0..3d78af94cde 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -169,9 +169,13 @@ def strategy_router_dependencies( else () ) + ( - _named( - f"{accounting_provider}/{decision_classifier.get('model', 'jev-latest')}", - "evaluation", + ( + _named(decision_classifier.get("deployment_name"), "classifier") + if decision_classifier.get("deployment_name") + else _named( + f"{accounting_provider}/{decision_classifier.get('model', 'jev-latest')}", + "evaluation", + ) ) if complexity.get("classifier_type") == "oss_classifier" else () diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index c487ce8b6c5..6e4adb4bbe7 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -1,4 +1,5 @@ from collections.abc import Mapping, Sequence +from datetime import date from typing import Any, Literal from pydantic import ConfigDict @@ -113,6 +114,24 @@ class AutoRouterPresetRecord(LiteLLMBaseModel): complexity_router_config: AutoRouterPresetConfig +class WhatsNewLaunch(LiteLLMBaseModel): + """One launch card in the dashboard Home page's What's new section.""" + + model_config = ConfigDict(frozen=True) + + icon: str + title: str + description: str + href: str + published_on: date + + +class WhatsNewResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + launches: tuple[WhatsNewLaunch, ...] + + class ComplexityScorerDefaults(LiteLLMBaseModel): """The complexity router's shipped heuristic scorer defaults. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 49432d4baa2..a1eaa6d4a90 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21196,6 +21196,18 @@ "output_vector_size": 1024, "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, + "databricks/ai_decide": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "databricks", + "metadata": { + "notes": "The ai_decide AI Function (Beta), called through POST /api/2.0/ai-functions/ai-decide. Its response carries no token usage, so LiteLLM records no spend for it." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.databricks.com/api/ai-functions/v1/ai-decide" + }, "databricks/databricks-openjev-qwen35-4b": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, @@ -21208,6 +21220,18 @@ "output_cost_per_token": 0.0, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" }, + "databricks/decisions/ai_decide": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "databricks", + "metadata": { + "notes": "The ai_decide AI Function (Beta), called through POST /api/2.0/ai-functions/ai-decide. Its response carries no token usage, so LiteLLM records no spend for it." + }, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://docs.databricks.com/api/ai-functions/v1/ai-decide" + }, "databricks/decisions/databricks-openjev-qwen35-4b": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, diff --git a/pyproject.toml b/pyproject.toml index 133fc735df1..0ed8b8757ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -321,6 +321,7 @@ include = [ "litellm/router_strategy/complexity_router/artifacts/*.json", "litellm/router_strategy/complexity_router/fuse_presets.json", "litellm/proxy/model_insights_tasks.json", + "litellm/proxy/public_endpoints/whats_new.json", "litellm/proxy/common_utils/codex_base_instructions.md", "litellm/proxy/common_utils/codex_bundled_models_0.159.3.json", ] 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 e277eddf009..d954fd49d81 100644 --- a/tests/integration/management/test_auto_router_databricks_classifier_writes.py +++ b/tests/integration/management/test_auto_router_databricks_classifier_writes.py @@ -2,28 +2,34 @@ import asyncio import json import threading import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from queue import SimpleQueue 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 from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue +from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS + +_MODEL: Final = "ai_decide" +_ROUTE: Final = "/api/2.0/ai-functions/ai-decide" _ENDPOINT: Final = "databricks-openjev-qwen35-4b" -_SECOND_ENDPOINT: Final = "databricks-openjev-qwen35-4b-canary" +_SECOND_WORKSPACE: Final = "/second-workspace" _API_KEY: Final = "synthetic-databricks-key" -_BARE_NAME_RULE: Final = ( - "must be the bare endpoint name (letters, digits, '-', '_' and '.', with no '/', '?', '#', spaces, or a " - "leading '.'), e.g. databricks-openjev-qwen35-4b" +_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': the serving endpoint name, " - "e.g. databricks-openjev-qwen35-4b" + "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 " @@ -32,17 +38,26 @@ _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]] = { - "answers": { - "tier": { - "type": "choice", - "choice": "COMPLEX", - "confidence": 0.91, - "probabilities": {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91}, - } - }, +_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}, } +_TIER_CRITERIA: Final[dict[str, JsonValue]] = { + "SIMPLE": "Greetings, chitchat, or short factual lookups with known answers", + "MEDIUM": "Everyday requests needing explanation, light reasoning, or minor technical work", + "COMPLEX": "Non-trivial code, architecture, multi-step work, or specialized domain depth", + "REASONING": "Open-ended analysis, proofs, tradeoffs, or tasks requiring careful thought", +} _ROWS_QUERY: Final = ( "SELECT request_id, status, metadata->>'internal_call_origin' AS origin " 'FROM "LiteLLM_SpendLogs" WHERE model_group = %s' @@ -50,7 +65,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]: @@ -64,7 +80,7 @@ def _tier(text: str) -> Callable[[Request], Reply]: def _classifier_config(api_base: str, **overrides: JsonValue) -> dict[str, JsonValue]: return { "provider": "databricks", - "model": _ENDPOINT, + "model": _MODEL, "api_base": api_base, "api_key": _API_KEY, "timeout_ms": 20000, @@ -133,21 +149,39 @@ async def _burst(base_url: str, key: str, model: str, count: int) -> tuple[httpx ) -def _rows(router: str, *, want: int) -> tuple[dict[str, JsonValue], ...]: - return tuple(eventually(lambda: read_rows(_ROWS_QUERY, (router,)), lambda found: len(found) >= want, seconds=70)) +def _request_rows(rows: Sequence[dict[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: + return tuple(row for row in rows if row["origin"] != "autorouter_classifier") + + +def _rows(router: str, *, requests: int) -> tuple[dict[str, JsonValue], ...]: + return tuple( + eventually( + lambda: read_rows(_ROWS_QUERY, (router,)), + lambda found: len(_request_rows(found)) >= requests, + seconds=70, + ) + ) + + +def test_validate_accepts_a_databricks_ai_decide_classifier_with_or_without_the_serving_endpoints_suffix( + gateway: Gateway, +) -> None: + for api_base in ( + "https://dbc-example.cloud.databricks.com", + "https://dbc-example.cloud.databricks.com/serving-endpoints", + ): + config: Final = _router_config(_classifier_config(api_base), "simple", "complex") + 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: - config: Final = _router_config( - _classifier_config("https://dbc-example.cloud.databricks.com/serving-endpoints"), "simple", "complex" - ) - assert _validate(gateway, config) == {"valid": True, "error": None} - long_name: Final = _router_config( - _classifier_config("https://dbc-example.cloud.databricks.com/serving-endpoints", model="e" * 5000), - "simple", - "complex", - ) - assert _validate(gateway, long_name) == {"valid": True, "error": None} + 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: @@ -159,10 +193,10 @@ def test_validate_names_each_databricks_classifier_misconfiguration(gateway: Gat _without(base, "api_key"), f"at opensource_classifier_config: Value error, {_BASE_NEEDS_KEY}", ), - ("slash in the name", {**base, "model": "a/b"}, f"Databricks serving endpoint name 'a/b' {_BARE_NAME_RULE}"), - ("space in the name", {**base, "model": "a b"}, f"Databricks serving endpoint name 'a b' {_BARE_NAME_RULE}"), - ("leading dot", {**base, "model": ".hidden"}, f"Databricks serving endpoint name '.hidden' {_BARE_NAME_RULE}"), - ("empty name", {**base, "model": ""}, f"Databricks serving endpoint name '' {_BARE_NAME_RULE}"), + ("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", {**base, "api_key": ""}, @@ -185,8 +219,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_the_databricks_endpoint_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 ( @@ -197,7 +236,7 @@ def test_test_routing_classifies_through_the_databricks_endpoint_without_calling ): 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} ) @@ -218,23 +257,27 @@ def test_test_routing_classifies_through_the_databricks_endpoint_without_calling "tier-probability:MEDIUM=0.060000", "tier-probability:COMPLEX=0.910000", ], - "classifier_model": f"databricks/{_ENDPOINT}", - "classifier_cost": 0.0, + "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 == f"/{_ENDPOINT}/invocations", 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 body["model"] == _ENDPOINT, body - assert body["state"] == f"\nClassify this message:\n{prompt}", body + assert json.loads(call.body) == { + "state": f"\nClassify this message:\n{prompt}", + **({} if model == _MODEL else {"model": model}), + "questions": { + "tier": {"type": "choice", "instructions": DEFAULT_JEV_INSTRUCTIONS, "criteria": _TIER_CRITERIA} + }, + }, call.body assert _judge_targets(simple) == () and _judge_targets(complex_tier) == () -def test_model_new_refuses_a_databricks_classifier_without_an_endpoint_name_and_stores_nothing( +def test_model_new_refuses_a_databricks_classifier_without_a_model_and_stores_nothing( gateway: Gateway, ) -> None: name: Final = f"router-{uuid.uuid4().hex[:10]}" @@ -261,7 +304,7 @@ def test_model_new_refuses_a_databricks_classifier_without_an_endpoint_name_and_ assert [entry for entry in map(object_value, listed) if entry["model_name"] == name] == [] -def test_model_update_refuses_dropping_the_endpoint_name_and_keeps_the_stored_classifier_serving( +def test_model_update_refuses_dropping_the_model_and_keeps_the_stored_classifier_serving( gateway: Gateway, ) -> None: with ( @@ -284,15 +327,15 @@ def test_model_update_refuses_dropping_the_endpoint_name_and_keeps_the_stored_cl assert _MODEL_REQUIRED in string_value(object_value(object_value(response.json())["error"])["message"]), ( response.text ) - assert _stored_classifier(gateway, identity)["model"] == _ENDPOINT + assert _stored_classifier(gateway, identity)["model"] == _MODEL chat: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(name)) assert chat.status_code == 200, chat.text assert chat.json()["choices"][0]["message"]["content"] == "complex answer", chat.text assert chat.headers["x-litellm-complexity-router-cause"] == "jev_classifier", dict(chat.headers) - assert _judge_targets(judge) == (f"/{_ENDPOINT}/invocations",) + assert _judge_targets(judge) == (_ROUTE,) -async def test_renaming_the_endpoint_under_a_burst_moves_the_classifier_and_logs_every_request_once( +async def test_moving_the_workspace_under_a_burst_moves_the_classifier_and_logs_every_request_once( gateway: Gateway, ) -> None: base_url: Final = str(gateway.client.base_url) @@ -315,7 +358,9 @@ async def test_renaming_the_endpoint_under_a_burst_moves_the_classifier_and_logs complex_model: Final = scenario.model(model="openai/complex-tier", api_base=complex_tier.url) classifier: Final = _classifier_config(judge.url) name, identity = _create_router(scenario, _router_config(classifier, simple_model, complex_model)) - renamed: Final = _router_config({**classifier, "model": _SECOND_ENDPOINT}, simple_model, complex_model) + moved: Final = _router_config( + {**classifier, "api_base": f"{judge.url}{_SECOND_WORKSPACE}"}, simple_model, complex_model + ) async with asyncio.TaskGroup() as requests: burst: Final = requests.create_task(_burst(base_url, gateway.key, name, 24)) try: @@ -323,15 +368,15 @@ async def test_renaming_the_endpoint_under_a_burst_moves_the_classifier_and_logs patched: Final = gateway.request( "PATCH", f"/model/{identity}/update", - {"litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": renamed}}, + {"litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": moved}}, ) assert patched.status_code == 200, patched.text finally: release.set() during: Final = await burst assert [response.status_code for response in during] == [200] * 24, [response.text for response in during] - assert set(_judge_targets(judge)) <= {f"/{_ENDPOINT}/invocations", f"/{_SECOND_ENDPOINT}/invocations"} - assert _stored_classifier(gateway, identity)["model"] == _SECOND_ENDPOINT + assert set(_judge_targets(judge)) <= {_ROUTE, f"{_SECOND_WORKSPACE}{_ROUTE}"} + assert _stored_classifier(gateway, identity)["api_base"] == f"{judge.url}{_SECOND_WORKSPACE}" convergence_calls: Final[SimpleQueue[str]] = SimpleQueue() @@ -341,16 +386,16 @@ async def test_renaming_the_endpoint_under_a_burst_moves_the_classifier_and_logs convergence_calls.put(string_value(chat.json()["id"])) return _judge_targets(judge) - eventually(classifier_target, lambda targets: targets == (f"/{_SECOND_ENDPOINT}/invocations",), seconds=30) + eventually(classifier_target, lambda targets: targets == (f"{_SECOND_WORKSPACE}{_ROUTE}",), seconds=30) after: Final = await _burst(base_url, gateway.key, name, 6) assert [response.status_code for response in after] == [200] * 6, [response.text for response in after] - assert set(_judge_targets(judge)) == {f"/{_SECOND_ENDPOINT}/invocations"} + assert set(_judge_targets(judge)) == {f"{_SECOND_WORKSPACE}{_ROUTE}"} identities: Final = ( *tuple(string_value(response.json()["id"]) for response in (*during, *after)), *tuple(convergence_calls.get_nowait() for _ in range(convergence_calls.qsize())), ) assert len(set(identities)) == len(identities), identities - rows: Final = _rows(name, want=len(identities)) - request_rows: Final = [row for row in rows if row["origin"] != "autorouter_classifier"] + rows: Final = _rows(name, requests=len(identities)) + request_rows: Final = _request_rows(rows) assert sorted(string_value(row["request_id"]) for row in request_rows) == sorted(identities), request_rows assert all(row["status"] == "success" for row in rows), rows diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index ef8412a6ed1..8d77c7d1094 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -190,7 +190,7 @@ 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", @@ -204,7 +204,7 @@ def test_evaluation_mode_health_check_of_a_databricks_serving_endpoint_resolves_ 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_ENDPOINT_PROBE_REPLY) scenario.cleanups.callback(delete_scenario, handle) model: Final = scenario.model( model="databricks/databricks-openjev-qwen35-4b", @@ -228,3 +228,41 @@ def test_evaluation_mode_health_check_of_a_databricks_serving_endpoint_resolves_ }, ) ] + + +_DATABRICKS_AI_DECIDE_PROBE_REPLY: Final = JsonResponse( + content_type="application/json", + body={ + "response": {"answers": {"reachable": {"type": "noul", "probability": 1.0}}}, + "metadata": {"version": "1.0"}, + }, +) + + +def test_evaluation_mode_health_check_of_databricks_ai_decide_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_AI_DECIDE_PROBE_REPLY) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + model="databricks/ai_decide", + 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}/api/2.0/ai-functions/ai-decide", + { + "state": os.environ.get("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"), + "questions": {"reachable": {"type": "noul", "instructions": "Is the service reachable?"}}, + }, + ) + ] diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py index 9b9f54e815d..24db0153bcd 100644 --- a/tests/integration/providers/test_databricks_decisions_wire.py +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -14,6 +14,8 @@ from integration._support.upstream import ScenarioHandle, delete_scenario, regis from integration.cost_calculation.cost_tracking_case import JsonResponse 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" @@ -38,9 +40,18 @@ _ANSWERS: Final[dict[str, JsonValue]] = { } _USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 367, "output_tokens": 3} _ANSWER_BODY: Final[dict[str, JsonValue]] = {"model": _ENDPOINT, "answers": _ANSWERS, "usage": _USAGE} -_BARE_NAME_RULE: Final = ( - "must be the bare endpoint name (letters, digits, '-', '_' and '.', with no '/', '?', '#', spaces, or a " - "leading '.'), e.g. databricks-openjev-qwen35-4b" +_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]] = { + "model": None, + "answers": _ANSWERS, + "usage": {"input_tokens": 0, "output_tokens": 0}, + "metadata": _METADATA, +} +_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'" _NON_BARE_NAMES: Final[tuple[tuple[str, str], ...]] = ( @@ -49,10 +60,17 @@ _NON_BARE_NAMES: Final[tuple[tuple[str, str], ...]] = ( ("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"}), - (404, {"error_code": "RESOURCE_DOES_NOT_EXIST", "message": f"Endpoint with name '{_ENDPOINT}' does not exist."}), +_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" ' @@ -60,6 +78,7 @@ _SPEND_QUERY: Final = ( ) _OWNED_PROXY_BUDGET: Final = int(2 * graceful_stop_seconds() + 120) _MEGABYTE_NAME: Final = "a" * 1_000_000 + "/" +_TENTH_MEGABYTE_NAME: Final = "a" * 100_000 + "/" def _number(value: JsonValue) -> float: @@ -92,6 +111,12 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, return _calls_to(_observed(gateway), handle) +def _timed_decide(gateway: Gateway, model: str) -> tuple[httpx.Response, float]: + started: Final = time.perf_counter() + response: Final = _decide(gateway, model) + return response, time.perf_counter() - started + + def _spend_row(call_id: str) -> dict[str, JsonValue]: rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) return rows[0] @@ -116,7 +141,7 @@ def _router_config(classifier: Mapping[str, JsonValue], simple: str, complex_: s def _environment_classifier() -> dict[str, JsonValue]: - return {"provider": "databricks", "model": _ENDPOINT, "timeout_ms": 20000, "circuit_breaker_enabled": False} + return {"provider": "databricks", "model": _MODEL, "timeout_ms": 20000, "circuit_breaker_enabled": False} def _router(scenario: Scenario, classifier: Mapping[str, JsonValue]) -> str: @@ -155,8 +180,8 @@ def _assert_classified_by_the_complex_tier(response: httpx.Response, router: str response.headers["x-litellm-model-name"], response.headers["x-litellm-complexity-router-tier"], response.headers["x-litellm-complexity-router-cause"], - response.headers["x-litellm-classifier-cost"], - ) == (router, "openai/complex-tier", "COMPLEX", "jev_classifier", "0.0"), dict(response.headers) + ) == (router, "openai/complex-tier", "COMPLEX", "jev_classifier"), dict(response.headers) + assert "x-litellm-classifier-cost" not in response.headers, dict(response.headers) def _assert_classifier_call( @@ -165,14 +190,14 @@ def _assert_classifier_call( assert call["path"] == path, call assert call["authorization"] == authorization, call body: Final = object_value(call["body"]) - assert body["model"] == _ENDPOINT, body + assert "model" not in body, body assert prompt_marker in string_value(body["state"]), body assert sorted(object_value(object_value(body["questions"])["tier"])) == ["criteria", "instructions", "type"], body def _assert_refused_before_any_upstream_call(gateway: Gateway, model: str, message: str) -> None: with gateway.scenario() as scenario: - handle: Final = _register(scenario, _ANSWER_BODY) + handle: Final = _register(scenario, _AI_DECIDE_BODY) deployment: Final = scenario.model(model=model, api_base=handle.api_base(), api_key=_API_KEY) response: Final = _decide(gateway, deployment) assert response.status_code == 400, response.text @@ -186,12 +211,10 @@ def _assert_refused_before_any_upstream_call(gateway: Gateway, model: str, messa def test_a_deployment_naming_a_non_bare_endpoint_is_refused_before_any_upstream_call( gateway: Gateway, name: str ) -> None: - _assert_refused_before_any_upstream_call( - gateway, f"databricks/{name}", f"Databricks serving endpoint name {name!r} {_BARE_NAME_RULE}" - ) + _assert_refused_before_any_upstream_call(gateway, f"databricks/{name}", _MODEL_RULE) -def test_a_deployment_naming_no_endpoint_is_refused_before_any_upstream_call(gateway: Gateway) -> None: +def test_a_deployment_naming_no_model_is_refused_before_any_upstream_call(gateway: Gateway) -> None: _assert_refused_before_any_upstream_call(gateway, "databricks/", "A model name is required for the Decisions API") @@ -251,11 +274,86 @@ def test_an_openai_format_request_at_v1_decisions_reaches_the_serving_endpoint_a row["model_group"], row["api_base"], _number(row["spend"]), - ) == ("success", "asystemone", "databricks", deployment, f"{handle.api_base()}/{_ENDPOINT}/invocations", 0.0), ( + ) == ("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: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _AI_DECIDE_BODY) + deployment: Final = scenario.model(model=f"databricks/{_MODEL}", 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": f"databricks/{_MODEL}", + "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": 0, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 0, + }, + }, response.text + assert response.headers["x-litellm-model-name"] == f"databricks/{_MODEL}", dict(response.headers) + (call,) = _upstream_calls(gateway, handle) + assert call["path"] == f"/{handle.scenario_id}{_ROUTE}", call + assert call["authorization"] == f"Bearer {_API_KEY}", call + assert call["body"] == { + "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()}{_ROUTE}", 0.0), row + + +def test_a_noul_probability_reaches_the_caller_as_the_system_one_noul_answer(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _register( + scenario, {"response": {"answers": {"defect": {"type": "noul", "probability": 0.8}}}, "metadata": _METADATA} + ) + deployment: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=handle.api_base(), api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/systemone", + {"model": deployment, "state": _STATE, "questions": {"defect": {"type": "noul", "instructions": "Bug?"}}}, + ) + assert response.status_code == 200, response.text + assert response.json()["answers"] == {"defect": {"type": "noul", "noul": 0.8}}, response.text + (call,) = _upstream_calls(gateway, handle) + 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) @@ -289,12 +387,43 @@ def test_a_wildcard_deployment_sends_each_request_to_the_named_endpoint_and_bill ) +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) + prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY) + model: Final = f"{prefix}/{_MODEL}" + responses: Final = tuple(_decide(gateway, model) for _ in range(2)) + for response in responses: + assert response.status_code == 200, response.text + assert response.json() == _SYSTEM_ONE_BODY, response.text + assert response.headers["x-litellm-model-group"] == model, dict(response.headers) + assert response.headers["x-litellm-model-name"] == f"databricks/{_MODEL}", 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}{_ROUTE}", call + assert call["authorization"] == f"Bearer {_API_KEY}", call + assert call["body"] == {"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", "asystemone", "databricks", model, f"{handle.api_base()}{_ROUTE}", 0.0), row + + @pytest.mark.parametrize("model", (5, ["a"]), ids=("int", "list")) def test_a_non_string_model_is_refused_at_the_gateway_without_an_upstream_call( gateway: Gateway, model: JsonValue ) -> None: with gateway.scenario() as scenario: - handle: Final = _register(scenario, _ANSWER_BODY) + handle: Final = _register(scenario, _AI_DECIDE_BODY) _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY) response: Final = gateway.request( "POST", "/v1/systemone", {"model": model, "state": _STATE, "questions": _QUESTIONS} @@ -328,18 +457,20 @@ def test_a_five_kilobyte_endpoint_name_reaches_the_upstream_and_its_not_found_an ), row -@pytest.mark.parametrize("status,body", _DATABRICKS_ERRORS, ids=("400", "404")) +@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/{_ENDPOINT}", 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}/{_ENDPOINT}/invocations", 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", @@ -353,8 +484,8 @@ def test_a_deployment_without_a_key_and_no_environment_key_is_refused_before_any gateway: Gateway, ) -> None: with gateway.scenario() as scenario: - handle: Final = _register(scenario, _ANSWER_BODY) - model: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=handle.api_base(), api_key=None) + handle: Final = _register(scenario, _AI_DECIDE_BODY) + model: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=handle.api_base(), api_key=None) response: Final = _decide(gateway, model) assert response.status_code == 401, response.text assert _MISSING_KEY in response.text, response.text @@ -366,9 +497,9 @@ def test_the_environment_key_and_base_serve_a_deployment_a_classifier_and_a_team gateway: Gateway, tmp_path: Path ) -> None: with gateway.scenario() as rig_scenario: - handle: Final = _register(rig_scenario, _ANSWER_BODY) - environment_base: Final = f"{handle.api_base()}/serving-endpoints" - environment_path: Final = f"/{handle.scenario_id}/serving-endpoints/{_ENDPOINT}/invocations" + handle: Final = _register(rig_scenario, _AI_DECIDE_BODY) + environment_base: Final = f"{handle.api_base()}/environment-workspace/serving-endpoints" + environment_path: Final = f"/{handle.scenario_id}/environment-workspace{_ROUTE}" environment: Final = { "DATABRICKS_API_KEY": _ENV_KEY, "DATABRICKS_TOKEN": _ENV_TOKEN, @@ -376,14 +507,14 @@ def test_the_environment_key_and_base_serve_a_deployment_a_classifier_and_a_team } with owned_proxy_process(gateway, tmp_path, environment) as owned, owned.gateway.scenario() as scenario: candidate: Final = owned.gateway - with_base: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=handle.api_base(), api_key=None) - bare: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=None, api_key=None) + with_base: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=handle.api_base(), api_key=None) + bare: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=None, api_key=None) decided: Final = tuple(_decide(candidate, model) for model in (with_base, bare)) for response in decided: assert response.status_code == 200, response.text - assert response.json() == _ANSWER_BODY, response.text + assert response.json() == _SYSTEM_ONE_BODY, response.text assert [(call["path"], call["authorization"]) for call in _upstream_calls(candidate, handle)] == [ - (f"/{handle.scenario_id}/{_ENDPOINT}/invocations", f"Bearer {_ENV_KEY}"), + (f"/{handle.scenario_id}{_ROUTE}", f"Bearer {_ENV_KEY}"), (environment_path, f"Bearer {_ENV_KEY}"), ] for response, model in zip(decided, (with_base, bare)): @@ -391,7 +522,7 @@ def test_the_environment_key_and_base_serve_a_deployment_a_classifier_and_a_team assert (row["status"], row["model_group"], _number(row["spend"])) == ("success", model, 0.0), row assert ( _spend_row(decided[1].headers["x-litellm-call-id"])["api_base"] - == f"{environment_base}/{_ENDPOINT}/invocations" + == f"{handle.api_base()}/environment-workspace{_ROUTE}" ) router: Final = _router(scenario, _environment_classifier()) @@ -435,7 +566,7 @@ def test_the_environment_key_and_base_serve_a_deployment_a_classifier_and_a_team "opensource_classifier_config" ] ) - assert (stored_classifier["provider"], stored_classifier["model"]) == ("databricks", _ENDPOINT), entry + assert (stored_classifier["provider"], stored_classifier["model"]) == ("databricks", _MODEL), entry assert "api_key" not in stored_classifier and "api_base" not in stored_classifier, entry member_chat: Final = _chat(candidate, team_router, key=member_key) _assert_classified_by_the_complex_tier(member_chat, team_router) @@ -453,21 +584,21 @@ def test_the_environment_token_serves_a_deployment_and_a_classifier_when_no_envi gateway: Gateway, tmp_path: Path ) -> None: with gateway.scenario() as rig_scenario: - handle: Final = _register(rig_scenario, _ANSWER_BODY) - environment_path: Final = f"/{handle.scenario_id}/serving-endpoints/{_ENDPOINT}/invocations" + handle: Final = _register(rig_scenario, _AI_DECIDE_BODY) + environment_path: Final = f"/{handle.scenario_id}/environment-workspace{_ROUTE}" environment: Final = { "DATABRICKS_TOKEN": _ENV_TOKEN, - "DATABRICKS_API_BASE": f"{handle.api_base()}/serving-endpoints", + "DATABRICKS_API_BASE": f"{handle.api_base()}/environment-workspace/serving-endpoints", } with ( owned_proxy_process(gateway, tmp_path, environment, remove_environment=("DATABRICKS_API_KEY",)) as owned, owned.gateway.scenario() as scenario, ): candidate: Final = owned.gateway - bare: Final = scenario.model(model=f"databricks/{_ENDPOINT}", api_base=None, api_key=None) + bare: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=None, api_key=None) decided: Final = _decide(candidate, bare) assert decided.status_code == 200, decided.text - assert decided.json() == _ANSWER_BODY, decided.text + assert decided.json() == _SYSTEM_ONE_BODY, decided.text (decision_call,) = _upstream_calls(candidate, handle) assert (decision_call["path"], decision_call["authorization"]) == (environment_path, f"Bearer {_ENV_TOKEN}") assert _spend_row(decided.headers["x-litellm-call-id"])["status"] == "success" @@ -483,15 +614,16 @@ def test_the_environment_token_serves_a_deployment_and_a_classifier_when_no_envi ) -def test_a_megabyte_endpoint_name_is_refused_in_linear_time_while_liveliness_stays_fast(gateway: Gateway) -> None: +def test_a_megabyte_model_name_is_refused_in_linear_time_while_liveliness_stays_up(gateway: Gateway) -> None: with gateway.scenario() as scenario: - handle: Final = _register(scenario, _ANSWER_BODY) + handle: Final = _register(scenario, _AI_DECIDE_BODY) prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY) model: Final = f"{prefix}/{_MEGABYTE_NAME}" anonymous: Final = gateway.client.post( "/v1/systemone", json={"model": model, "state": _STATE, "questions": _QUESTIONS} ) assert anonymous.status_code == 401, anonymous.text[:300] + tenth_response, tenth_elapsed = _timed_decide(gateway, f"{prefix}/{_TENTH_MEGABYTE_NAME}") liveliness: Final[list[tuple[int, float]]] = [] stop: Final = threading.Event() @@ -503,16 +635,15 @@ def test_a_megabyte_endpoint_name_is_refused_in_linear_time_while_liveliness_sta poller: Final = threading.Thread(target=poll) poller.start() - started: Final = time.perf_counter() - response: Final = _decide(gateway, model) - elapsed: Final = time.perf_counter() - started + response, elapsed = _timed_decide(gateway, model) stop.set() poller.join() - assert elapsed < 3, elapsed + assert tenth_response.status_code == 400, tenth_response.text[:300] + assert elapsed < 25 * tenth_elapsed, (elapsed, tenth_elapsed) assert liveliness and all(status == 200 for status, _ in liveliness), liveliness - assert max(seconds for _, seconds in liveliness) < 1, liveliness + assert max(seconds for _, seconds in liveliness) < max(1.0, elapsed / 2), (elapsed, liveliness) assert response.status_code == 400, response.text[:300] - assert _BARE_NAME_RULE in response.text, response.text[-400:] + assert _MODEL_RULE in response.text, response.text[:400] assert _upstream_calls(gateway, handle) == [] row: Final = _spend_row(response.headers["x-litellm-call-id"]) assert (row["status"], row["call_type"], _number(row["spend"])) == ("failure", "asystemone", 0.0), row diff --git a/tests/integration/providers/test_decisions_openai_format_wire.py b/tests/integration/providers/test_decisions_openai_format_wire.py index 1104a6151db..64387b293ce 100644 --- a/tests/integration/providers/test_decisions_openai_format_wire.py +++ b/tests/integration/providers/test_decisions_openai_format_wire.py @@ -88,6 +88,17 @@ _USAGE: Final[dict[str, JsonValue]] = { "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": _INPUT_TOKENS + _OUTPUT_TOKENS, } +_NO_USAGE: Final[dict[str, JsonValue]] = { + "input_tokens": 0, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 0, +} +_AI_DECIDE_ANSWERS: Final[dict[str, JsonValue]] = { + **_SYSTEM_ONE_ANSWERS, + "defect": {"type": "noul", "probability": 0.93}, +} _SDK_QUESTIONS: Final = TypeAdapter(list[dict[str, object]]).validate_python(_QUESTIONS) _SPEND_QUERY: Final = ( "SELECT spend, status, call_type, model_group, custom_llm_provider, api_base, prompt_tokens, completion_tokens, " @@ -105,11 +116,14 @@ class _Provider: wraps_result: bool cost_map_key: str | None speaks_openai: bool = False + ai_decide: bool = False provider_reported_cost: float | None = None def upstream_body(self) -> dict[str, JsonValue]: if self.speaks_openai: return {"model": self.body_model, "input": _INPUT, "questions": _QUESTIONS} + if self.ai_decide: + return {"state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS} return {"model": self.body_model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS} def _with_reported_cost(self, usage: dict[str, JsonValue]) -> dict[str, JsonValue]: @@ -120,6 +134,8 @@ class _Provider: def upstream_reply(self) -> dict[str, JsonValue]: if self.speaks_openai: return {"model": self.body_model, "answers": _ANSWERS, "usage": self._with_reported_cost(_USAGE)} + if self.ai_decide: + return {"response": {"answers": _AI_DECIDE_ANSWERS}, "metadata": {"version": "1.0"}} answer: Final[dict[str, JsonValue]] = { "model": self.body_model, "answers": _SYSTEM_ONE_ANSWERS, @@ -128,8 +144,13 @@ class _Provider: return {"result": answer, "success": True} if self.wraps_result else answer def litellm_response(self) -> dict[str, JsonValue]: + if self.ai_decide: + return {"model": self.model, "answers": _ANSWERS, "usage": _NO_USAGE} return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE} + def billed_tokens(self) -> tuple[int, int]: + return (0, 0) if self.ai_decide else (_INPUT_TOKENS, _OUTPUT_TOKENS) + _PROVIDERS: Final = ( _Provider( @@ -165,6 +186,16 @@ _PROVIDERS: Final = ( "cloudflare/@cf/cloudflare/clef", ), _Provider("hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None), + _Provider( + "databricks", + "databricks/ai_decide", + "/api/2.0/ai-functions/ai-decide", + "ai_decide", + _API_KEY, + False, + None, + ai_decide=True, + ), _Provider( "databricks", "databricks/databricks-openjev-qwen35-4b", @@ -234,7 +265,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: @@ -358,8 +389,7 @@ def test_each_provider_gets_its_own_path_key_and_body_and_is_billed(gateway: Gat provider.name, model, f"{handle.api_base()}{provider.path}", - _INPUT_TOKENS, - _OUTPUT_TOKENS, + *provider.billed_tokens(), ) assert math.isclose(_number(row["spend"]), expected_spend, rel_tol=1e-9), row diff --git a/tests/integration/providers/test_strands_decider_agentcore_wire.py b/tests/integration/providers/test_strands_decider_agentcore_wire.py index 493d8c238fa..fc62849db9b 100644 --- a/tests/integration/providers/test_strands_decider_agentcore_wire.py +++ b/tests/integration/providers/test_strands_decider_agentcore_wire.py @@ -114,6 +114,12 @@ _OPENAI: Final = _Route( "/v1/decisions", {"input": _INPUT, "questions": _OPENAI_QUESTIONS}, _OPENAI_UPSTREAM_BODY, _OPENAI_ANSWER ) _ROUTES: Final = (_SYSTEMONE, _OPENAI, replace(_SYSTEMONE, path="/systemone"), replace(_OPENAI, path="/decisions")) +_ROUTE_CALL_TYPES: Final[Mapping[str, str]] = { + "/v1/systemone": "asystemone", + "/systemone": "asystemone", + "/v1/decisions": "adecisions", + "/decisions": "adecisions", +} _SPEND_QUERY: Final = ( "SELECT status, call_type, custom_llm_provider, model_group, api_base, prompt_tokens, completion_tokens " 'FROM "LiteLLM_SpendLogs" WHERE request_id = %s' @@ -411,11 +417,11 @@ def test_an_arn_deployment_answers_on_every_decisions_route_signed_for_bedrock_a _assert_runtime_call(call, "audit_signed", route=route) assert call.headers[_SESSION_HEADER] == _default_session("audit_signed"), dict(call.headers) _assert_signed(call) - for response in responses: + for response, route in zip(responses, _ROUTES, strict=True): row = _spend_row(response.headers["x-litellm-call-id"]) assert row == { "status": "success", - "call_type": "asystemone", + "call_type": _ROUTE_CALL_TYPES[route.path], "custom_llm_provider": "strands_decider", "model_group": _SIGNED, "api_base": _invocations_url("audit_signed"), diff --git a/tests/integration/routing/test_complexity_router_databricks_classifier.py b/tests/integration/routing/test_complexity_router_databricks_classifier.py index 25611b4ec08..419b786bfc1 100644 --- a/tests/integration/routing/test_complexity_router_databricks_classifier.py +++ b/tests/integration/routing/test_complexity_router_databricks_classifier.py @@ -31,9 +31,10 @@ from pydantic import JsonValue, TypeAdapter from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS -_ENDPOINT: Final = "databricks-openjev-qwen35-4b" +_MODEL: Final = "ai_decide" +_ROUTE: Final = "/api/2.0/ai-functions/ai-decide" _API_KEY: Final = "synthetic-databricks-key" -_CLASSIFIER_MODEL: Final = f"databricks/{_ENDPOINT}" +_CLASSIFIER_MODEL: Final = f"databricks/{_MODEL}" _TYPESAFE_MODEL: Final = "jev-1.13.0" _TYPESAFE_KEY: Final = "synthetic-typesafe-key" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -48,10 +49,12 @@ _TIER_CRITERIA: Final[dict[str, JsonValue]] = { } _PROBABILITIES: Final[dict[str, JsonValue]] = {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91} _USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 12, "output_tokens": 1} -_ANSWER: Final[dict[str, JsonValue]] = { - "answers": {"tier": {"type": "choice", "choice": "COMPLEX", "confidence": 0.91, "probabilities": _PROBABILITIES}}, - "usage": _USAGE, +_ANSWERS: Final[dict[str, JsonValue]] = { + "tier": {"type": "choice", "choice": "COMPLEX", "confidence": 0.91, "probabilities": _PROBABILITIES} } +_METADATA: Final[dict[str, JsonValue]] = {"version": "1.0"} +_ANSWER: Final[dict[str, JsonValue]] = {"response": {"answers": _ANSWERS}, "metadata": _METADATA} +_TYPESAFE_ANSWER: Final[dict[str, JsonValue]] = {"model": _TYPESAFE_MODEL, "answers": _ANSWERS, "usage": _USAGE} _ROWS_QUERY: Final = ( "SELECT request_id, status, model, call_type, cache_hit, custom_llm_provider, spend, " "metadata->>'internal_call_origin' AS origin " @@ -155,7 +158,7 @@ def _wires(classifier: Callable[[Request], Reply] | None = None) -> Iterator[tup def _classifier_config(judge_url: str, **overrides: JsonValue) -> dict[str, JsonValue]: return { "provider": "databricks", - "model": _ENDPOINT, + "model": _MODEL, "api_base": judge_url, "api_key": _API_KEY, "timeout_ms": 20000, @@ -209,25 +212,25 @@ def _state(prompt: str, system: str | None) -> str: def _assert_classifier_call( - call: Request, *, prompt: str, system: str | None = None, endpoint: str = _ENDPOINT, key: str = _API_KEY + call: Request, *, prompt: str, system: str | None = None, key: str = _API_KEY, endpoint: str | None = None ) -> None: - assert call.target == f"/{endpoint}/invocations", call.target + 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), - "model": endpoint, + **({} if endpoint is None else {"model": endpoint}), "questions": {"tier": {"type": "choice", "instructions": DEFAULT_JEV_INSTRUCTIONS, "criteria": _TIER_CRITERIA}}, }, call.body -def _assert_classified(headers: Mapping[str, str], *, router: str, cost: str | None = "0.0") -> None: +def _assert_classified(headers: Mapping[str, str], *, router: str) -> None: assert ( headers.get("x-litellm-model-group"), headers.get("x-litellm-model-name"), headers.get("x-litellm-complexity-router-tier"), headers.get("x-litellm-complexity-router-cause"), headers.get("x-litellm-classifier-cost"), - ) == (router, "openai/complex-tier", "COMPLEX", "jev_classifier", cost), dict(headers) + ) == (router, "openai/complex-tier", "COMPLEX", "jev_classifier", None), dict(headers) def _assert_fell_back(headers: Mapping[str, str], *, router: str) -> None: @@ -622,10 +625,19 @@ 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_non_catalog_endpoint_name_is_classified_without_a_cost_header_and_logged_at_zero(gateway: Gateway) -> None: +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() as (judge, simple, complex_), gateway.scenario() as scenario: + 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) @@ -633,7 +645,7 @@ def test_a_non_catalog_endpoint_name_is_classified_without_a_cost_header_and_log assert response.status_code == 200, response.text identity, text = _identity_and_text("/v1/chat/completions", _JSON_OBJECT.validate_json(response.content)) assert text == _COMPLEX_TEXT, response.text - _assert_classified(response.headers, router=routed.name, cost=None) + _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) == () @@ -686,7 +698,7 @@ _FAILURES: Final = ( _Failure("401", _failing(401), {}, "classifier_error", "MaskedHTTPStatusError", 0), _Failure( "malformed", - _answering({"answers": {"tier": {"type": "choice", "choice": "COMPLEX", "confidence": 0.91}}, "usage": _USAGE}), + _answering({"response": {"answers": {"tier": {"type": "choice", "choice": "COMPLEX", "confidence": 0.91}}}}), {}, "invalid_response", "ValidationError", @@ -696,15 +708,17 @@ _FAILURES: Final = ( "unknown-tier", _answering( { - "answers": { - "tier": { - "type": "choice", - "choice": "GALAXY", - "confidence": 0.91, - "probabilities": {"GALAXY": 0.91}, + "response": { + "answers": { + "tier": { + "type": "choice", + "choice": "GALAXY", + "confidence": 0.91, + "probabilities": {"GALAXY": 0.91}, + } } }, - "usage": _USAGE, + "metadata": _METADATA, } ), {}, @@ -793,7 +807,7 @@ def test_the_default_jev_provider_still_posts_to_systemone_and_bills_from_the_co prompt: Final = _prompt() expected_cost: Final = _typesafe_cost() with ( - _wires(_answering({**_ANSWER, "model": _TYPESAFE_MODEL})) as (judge, simple, complex_), + _wires(_answering(_TYPESAFE_ANSWER)) as (judge, simple, complex_), gateway.scenario() as scenario, ): classifier: Final[dict[str, JsonValue]] = { @@ -848,7 +862,7 @@ async def test_a_classifier_outage_mid_traffic_falls_back_and_recovers_with_ever before_prompts: Final = tuple(_prompt() for _ in calls) before: Final = await _burst(base_url, gateway.key, routed.name, calls, before_prompts) assert [item.cause for item in before] == ["jev_classifier"] * 12, before - assert [item.classifier_cost for item in before] == ["0.0"] * 12, before + assert [item.classifier_cost for item in before] == [None] * 12, before assert [item.text for item in before] == [_COMPLEX_TEXT] * 12, before assert sorted(json.loads(call.body)["state"] for call in _posts(judge)) == sorted( _state(prompt, None) for prompt in before_prompts diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 5ed96e9b616..06a22631730 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -79,7 +79,7 @@ _STRANDS_RESPONSE: Final[Mapping[str, object]] = { "latency_ms": 3722.17, } _DATABRICKS_ENDPOINT: Final = "databricks-openjev-qwen35-4b" -_DATABRICKS_RESPONSE: Final[Mapping[str, object]] = { +_DATABRICKS_ENDPOINT_RESPONSE: Final[Mapping[str, object]] = { "model": "/mosaicml/local_model", "answers": { "is_defect": {"type": "noul", "noul": 0.9}, @@ -92,6 +92,21 @@ _DATABRICKS_RESPONSE: Final[Mapping[str, object]] = { }, "usage": {"input_tokens": 214, "output_tokens": 0}, } +_DATABRICKS_ROUTE: Final = "https://workspace.example/api/2.0/ai-functions/ai-decide" +_DATABRICKS_RESPONSE: Final[Mapping[str, object]] = { + "response": { + "answers": { + "is_defect": {"type": "noul", "probability": 0.9}, + "sentiment": { + "type": "choice", + "choice": "positive", + "confidence": 0.8, + "probabilities": {"positive": 0.8, "negative": 0.2}, + }, + } + }, + "metadata": {"version": "1.0"}, +} _PROVIDERS: Final[tuple[tuple[str, str, str, str], ...]] = ( ( "perplexity", @@ -1170,7 +1185,7 @@ async def test_databricks_requires_api_base_before_http( with pytest.raises(litellm.BadRequestError, match="DATABRICKS_API_BASE"): await litellm.asystemone( - model=f"databricks/{_DATABRICKS_ENDPOINT}", + model="databricks/ai_decide", state="review", questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, ) @@ -1190,7 +1205,7 @@ async def test_databricks_rejects_a_serving_endpoint_name_that_would_rewrite_the monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.example/serving-endpoints") monkeypatch.setenv("DATABRICKS_API_KEY", "dapi-key") - with pytest.raises(litellm.BadRequestError, match="bare endpoint name"): + with pytest.raises(litellm.BadRequestError, match="bare serving endpoint name"): await litellm.asystemone( model=f"databricks/{endpoint_name}", state="review", @@ -1205,13 +1220,13 @@ async def test_databricks_requires_a_key_before_http( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, ) -> None: - monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.example/serving-endpoints") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://workspace.example") monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) monkeypatch.delenv("DATABRICKS_TOKEN", raising=False) with pytest.raises(litellm.AuthenticationError, match="Missing API key for Decisions provider 'databricks'"): await litellm.asystemone( - model=f"databricks/{_DATABRICKS_ENDPOINT}", + model="databricks/ai_decide", state="review", questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, ) @@ -1233,7 +1248,7 @@ async def test_databricks_posts_the_jev_body_to_the_serving_endpoint_invocations monkeypatch.setenv("DATABRICKS_TOKEN", "dapi-token") route: Final = respx_mock.post( f"https://workspace.example/serving-endpoints/{_DATABRICKS_ENDPOINT}/invocations" - ).respond(json=_DATABRICKS_RESPONSE) + ).respond(json=_DATABRICKS_ENDPOINT_RESPONSE) questions: Final = { "is_defect": {"type": "noul", "instructions": "Is this a defect?"}, "sentiment": {"type": "choice", "criteria": {"positive": None, "negative": "unhappy"}}, @@ -1260,28 +1275,68 @@ async def test_databricks_posts_the_jev_body_to_the_serving_endpoint_invocations assert response.hidden_params["model"] == f"databricks/{_DATABRICKS_ENDPOINT}" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "api_base", + [ + "https://workspace.example", + "https://workspace.example/", + "https://workspace.example/serving-endpoints", + "https://workspace.example/serving-endpoints/", + ], +) +async def test_databricks_posts_the_systemone_body_without_model_to_the_ai_decide_route( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + api_base: str, +) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", api_base) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.setenv("DATABRICKS_TOKEN", "dapi-token") + route: Final = respx_mock.post(_DATABRICKS_ROUTE).respond(json=_DATABRICKS_RESPONSE) + questions: Final = { + "is_defect": {"type": "noul", "instructions": "Is this a defect?"}, + "sentiment": {"type": "choice", "criteria": {"positive": None, "negative": "unhappy"}}, + } + + response: Final = await litellm.adecisions( + model="databricks/ai_decide", + state={"ticket": "export hangs"}, + questions=questions, + ) + + assert route.called + sent: Final = respx_mock.calls[0].request + assert sent.headers["authorization"] == "Bearer dapi-token" + assert json.loads(sent.content) == {"state": {"ticket": "export hangs"}, "questions": questions} + assert response.model is None + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert isinstance(response.answers["sentiment"], ChoiceAnswer) + assert response.model_extra == {"metadata": {"version": "1.0"}} + assert response.hidden_params["custom_llm_provider"] == "databricks" + assert response.hidden_params["model"] == "databricks/ai_decide" + + @pytest.mark.asyncio async def test_databricks_router_deployment_sends_its_own_key_to_its_own_base( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, ) -> None: - monkeypatch.setenv("DATABRICKS_API_BASE", "https://other-workspace.example/serving-endpoints") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://other-workspace.example") monkeypatch.setenv("DATABRICKS_API_KEY", "env-key-stays-home") router: Final = litellm.Router( model_list=[ { "model_name": "decider", "litellm_params": { - "model": f"databricks/{_DATABRICKS_ENDPOINT}", - "api_base": "https://workspace.example/serving-endpoints", + "model": "databricks/ai_decide", + "api_base": "https://workspace.example", "api_key": "deployment-key", }, } ] ) - route: Final = respx_mock.post( - f"https://workspace.example/serving-endpoints/{_DATABRICKS_ENDPOINT}/invocations" - ).respond(json=_DATABRICKS_RESPONSE) + route: Final = respx_mock.post(_DATABRICKS_ROUTE).respond(json=_DATABRICKS_RESPONSE) response: Final = await router.asystemone( model="decider", diff --git a/tests/unit/llms/databricks/decisions/test_transformation.py b/tests/unit/llms/databricks/decisions/test_transformation.py index e6de1d1a635..f6ea3444db1 100644 --- a/tests/unit/llms/databricks/decisions/test_transformation.py +++ b/tests/unit/llms/databricks/decisions/test_transformation.py @@ -3,27 +3,84 @@ from typing import Final import pytest from litellm.llms.databricks.decisions.transformation import DatabricksDecisionsConfig +from litellm.types.decisions import ( + DecisionsIRPredicateAnswer, + DecisionsIRPredicateQuestion, + DecisionsIRRequest, + DecisionsIRState, +) -_BASE: Final = "https://workspace.example/serving-endpoints" +_HOST: Final = "https://workspace.example" +_BASE: Final = f"{_HOST}/serving-endpoints" +_AI_DECIDE_ROUTE: Final = f"{_HOST}/api/2.0/ai-functions/ai-decide" _CONFIG: Final = DatabricksDecisionsConfig() +_REQUEST: Final = DecisionsIRRequest( + input=DecisionsIRState(state="export hangs"), + questions=(DecisionsIRPredicateQuestion(name="is_defect", instructions="Is this a defect?", criteria=None),), +) +_QUESTIONS: Final = {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}} @pytest.mark.parametrize("api_base", [_BASE, f"{_BASE}/"]) -def test_the_complete_url_is_the_serving_endpoint_invocations_route(api_base: str) -> None: +def test_a_serving_endpoint_name_is_called_on_its_invocations_route(api_base: str) -> None: assert _CONFIG.get_complete_url(api_base, "openjev") == f"{_BASE}/openjev/invocations" +@pytest.mark.parametrize("api_base", [_HOST, f"{_HOST}/", _BASE, f"{_BASE}/"]) +def test_ai_decide_is_called_on_the_workspace_ai_function_route(api_base: str) -> None: + assert _CONFIG.get_complete_url(api_base, "ai_decide") == _AI_DECIDE_ROUTE + + @pytest.mark.parametrize( "name", ["", "serving-endpoints/openjev", "openjev?x=1", "openjev#frag", ".", "..", ".openjev", "open jev", "openjev%2Fx"], ) def test_a_name_that_would_rewrite_the_url_is_rejected_before_any_request(name: str) -> None: - with pytest.raises(ValueError, match="bare endpoint name"): + with pytest.raises(ValueError, match="bare serving endpoint name"): _ = _CONFIG.get_complete_url(_BASE, name) - with pytest.raises(ValueError, match="bare endpoint name"): + with pytest.raises(ValueError, match="bare serving endpoint name"): _ = _CONFIG.canonical_model(name) +def test_the_rejection_does_not_echo_the_rejected_name() -> None: + name: Final = "openjev?token=" + "x" * 4096 + with pytest.raises(ValueError, match="bare serving endpoint name") as raised: + _ = _CONFIG.canonical_model(name) + assert "openjev?token=" not in str(raised.value) + + +def test_a_serving_endpoint_request_keeps_the_model_the_endpoint_expects() -> None: + body: Final = _CONFIG.transform_decisions_request("openjev", _REQUEST, "databricks") + assert body == {"model": "openjev", "state": "export hangs", "questions": _QUESTIONS} + + +def test_the_ai_decide_request_omits_model_because_the_route_rejects_unexpected_parameters() -> None: + body: Final = _CONFIG.transform_decisions_request("ai_decide", _REQUEST, "databricks") + assert body == {"state": "export hangs", "questions": _QUESTIONS} + + +def test_the_ai_decide_envelope_is_unwrapped_and_noul_probability_becomes_the_noul_answer() -> None: + payload: Final = { + "response": {"answers": {"is_defect": {"type": "noul", "probability": 0.9}}}, + "metadata": {"version": "1.0"}, + } + parsed: Final = _CONFIG.parse_response(payload, _REQUEST) + assert parsed.answers == (DecisionsIRPredicateAnswer(probability=0.9),) + assert parsed.model is None + assert dict(parsed.extra) == {"metadata": {"version": "1.0"}} + + +def test_a_serving_endpoint_answer_is_parsed_as_sent() -> None: + payload: Final = { + "model": "openjev", + "answers": {"is_defect": {"type": "noul", "noul": 0.25}}, + "usage": {"input_tokens": 12, "output_tokens": 0}, + } + parsed: Final = _CONFIG.parse_response(payload, _REQUEST) + assert parsed.answers == (DecisionsIRPredicateAnswer(probability=0.25),) + assert parsed.model == "openjev" + + def test_connection_prefers_the_configured_values_over_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DATABRICKS_API_BASE", "https://env.example/serving-endpoints") monkeypatch.setenv("DATABRICKS_API_KEY", "dapi-env") @@ -54,3 +111,10 @@ def test_classifier_response_accounts_the_serving_endpoint_instead_of_the_contai assert normalized["model"] == "openjev" assert normalized["usage"] == body["usage"] assert body["model"] == "/mosaicml/local_model" + + +def test_classifier_response_unwraps_the_ai_decide_envelope_and_accounts_the_configured_model() -> None: + choice: Final = {"type": "choice", "choice": "simple", "confidence": 0.9, "probabilities": {"simple": 0.9}} + body: Final = {"response": {"answers": {"tier": choice}}, "metadata": {"version": "1.0"}} + normalized: Final = _CONFIG.classifier_response(body, "ai_decide") + assert dict(normalized) == {"answers": {"tier": choice}, "metadata": {"version": "1.0"}, "model": "ai_decide"} diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 6254b1bc44c..078fff0266b 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6164,7 +6164,8 @@ class TestMCPDcrBridgeDelegateAdmission: policy gates re-check. ``team_blocked=True`` patches the centralized gate's ``get_team_object`` at the ``user_api_key_auth`` namespace it actually calls; ``owner`` patches the SCIM gate's ``get_user_object`` (``auth_checks`` namespace); ``project_object`` - patches the centralized gate's ``get_project_object``. Unpatched lookups hit the MagicMock + patches the centralized gate's ``get_project_object``. The team-membership lookup is + patched to "no membership row". Unpatched lookups hit the MagicMock prisma and are swallowed (``_safe_fetch`` / the SCIM gate's fail-open), so their checks skip. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was the reload key.""" @@ -6179,6 +6180,10 @@ class TestMCPDcrBridgeDelegateAdmission: ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_membership", + AsyncMock(return_value=None), + ), ] if team_blocked: patchers.append( @@ -6206,7 +6211,8 @@ class TestMCPDcrBridgeDelegateAdmission: def _patch_user_reload(*, return_value=None, side_effect=None): """Patch the user-subject reload path an interactively-minted envelope takes: the ``get_user_object`` lookup ``reload_admitted_user`` runs (which also drives the SCIM gate), - plus the ``prisma_client`` / ``user_api_key_cache`` globals. The centralized gate's own + plus the ``prisma_client`` / ``user_api_key_cache`` globals. The team-membership lookup is + patched to "no membership row". The centralized gate's own fetches fail-safe to None under the MagicMock prisma, so an unblocked user admits. Yields the ``get_user_object`` mock so a caller can assert the sealed user_id was the reload key.""" get_user_object = AsyncMock(return_value=return_value, side_effect=side_effect) @@ -6214,6 +6220,10 @@ class TestMCPDcrBridgeDelegateAdmission: patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_membership", + AsyncMock(return_value=None), + ), ): yield get_user_object diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index a7006a7c1c6..56f3ce5f192 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -305,6 +305,53 @@ def test_proxy_systemone_dispatches_databricks_serving_endpoint( assert upstream.calls[0].request.headers["authorization"] == "Bearer dapi-deployment" +def test_proxy_systemone_dispatches_databricks_ai_decide( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("DATABRICKS_API_BASE", raising=False) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.delenv("DATABRICKS_TOKEN", raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "databricks-decider", + "litellm_params": { + "model": "databricks/ai_decide", + "api_base": "https://workspace.example/serving-endpoints", + "api_key": "dapi-deployment", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + upstream: Final = respx_mock.post("https://workspace.example/api/2.0/ai-functions/ai-decide").respond( + json={ + "response": {"answers": {"is_defect": {"type": "noul", "probability": 0.9}}}, + "metadata": {"version": "1.0"}, + } + ) + + response: Final = client.post( + "/v1/systemone", + json={ + "model": "databricks-decider", + "state": {"source": "proxy-test"}, + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + }, + ) + + assert response.status_code == 200, response.text + assert response.json()["answers"] == _RESPONSE["answers"] + assert upstream.called + assert json.loads(upstream.calls[0].request.content) == { + "state": {"source": "proxy-test"}, + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert upstream.calls[0].request.headers["authorization"] == "Bearer dapi-deployment" + + def test_proxy_decisions_without_model_uses_the_proxy_default_model( client: TestClient, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/management_helpers/test_auto_router_permissions.py b/tests/unit/proxy/management_helpers/test_auto_router_permissions.py index 5ef1399b0cb..e161ea3ecaf 100644 --- a/tests/unit/proxy/management_helpers/test_auto_router_permissions.py +++ b/tests/unit/proxy/management_helpers/test_auto_router_permissions.py @@ -146,8 +146,8 @@ def test_tier_config_is_normalized_and_unknown_router_extras_are_rejected() -> N ({"provider": "laya", "model": "english", "api_key": "sk-member"}, "api_key"), ({"provider": "bespoke", "model": "nimble-latest", "api_base": "https://collector.invalid"}, "api_base"), ({"provider": "bespoke", "model": "nimble-latest", "api_key": "sk-member"}, "api_key"), - ({"provider": "databricks", "model": "my-openjev", "api_base": "https://collector.invalid"}, "opensource_classifier_config"), - ({"provider": "databricks", "model": "my-openjev", "api_key": "dapi-member"}, "api_key"), + ({"provider": "databricks", "model": "ai_decide", "api_base": "https://collector.invalid"}, "opensource_classifier_config"), + ({"provider": "databricks", "model": "ai_decide", "api_key": "dapi-member"}, "api_key"), ], ) @pytest.mark.parametrize("legacy", [False, True]) @@ -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", "my-openjev")], + [ + ("typesafe", "jev-preview"), + ("laya", "english"), + ("bespoke", "nimble-latest"), + ("databricks", "ai_decide"), + ("databricks", "my-openjev"), + ], ) @pytest.mark.parametrize("legacy", [False, True]) def test_members_can_still_tune_the_jev_classifier(provider: str, model: str, legacy: bool) -> None: @@ -427,3 +433,26 @@ async def test_jev_evaluation_obeys_each_containing_scope( return await operation assert not catalog.get_model_list(f"{provider}/{model}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("restricted", ["key", "team", None]) +async def test_deployed_decision_classifier_requires_key_and_team_access(catalog: Router, restricted: str | None) -> None: + permitted: Final = ["allowed", "other"] + operation: Final = authorize_member_auto_router_dependencies( + config=validate_member_auto_router_config({ + "tiers": {"SIMPLE": "allowed"}, + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"deployment_name": "other"}, + }), + default_model=None, + user_api_key_dict=_actor(models=["allowed"] if restricted == "key" else permitted), + team=_team(models=["allowed"] if restricted == "team" else permitted), + prisma_client=_Client(), + llm_router=catalog, + ) + if restricted is not None: + with pytest.raises(ProxyException, match="other"): + await operation + return + await operation diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index b720ba24c3e..0779f692b17 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -1,6 +1,8 @@ import json import re from datetime import datetime, timezone +from importlib.resources import files +from collections.abc import Iterator from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -17,7 +19,7 @@ from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_pres from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) -from litellm.types.proxy.public_endpoints.public_endpoints import ProviderCreateInfo +from litellm.types.proxy.public_endpoints.public_endpoints import ProviderCreateInfo, WhatsNewResponse from litellm.types.utils import LlmProviders @@ -67,13 +69,10 @@ def test_get_provider_create_fields(): assert isinstance(first_provider["credential_fields"], list) has_detailed_fields = any( - provider.get("credential_fields") - and len(provider.get("credential_fields", [])) > 0 + provider.get("credential_fields") and len(provider.get("credential_fields", [])) > 0 for provider in response_data ) - assert ( - has_detailed_fields - ), "Expected at least one provider to have detailed credential fields" + assert has_detailed_fields, "Expected at least one provider to have detailed credential fields" def test_get_litellm_model_cost_map_catalog_only_excludes_runtime_registered_entries( @@ -124,10 +123,7 @@ def test_get_litellm_model_cost_map_returns_cost_map(): sample_model_data = payload[sample_model] assert isinstance(sample_model_data, dict) # Check for common cost fields that should be present - assert ( - "input_cost_per_token" in sample_model_data - or "output_cost_per_token" in sample_model_data - ) + assert "input_cost_per_token" in sample_model_data or "output_cost_per_token" in sample_model_data def test_public_ai_hub_info_is_public_by_default(monkeypatch): @@ -218,9 +214,9 @@ def test_anthropic_provider_fields_support_byok(): "Anthropic api_key must be optional so admins can configure BYOK models " "without entering a key. See BYOK tutorial." ) - assert fields_by_key["api_key"].get( - "tooltip" - ), "Anthropic api_key must have a tooltip explaining the BYOK use case." + assert fields_by_key["api_key"].get("tooltip"), ( + "Anthropic api_key must have a tooltip explaining the BYOK use case." + ) assert "api_base" in fields_by_key, ( "Anthropic provider form must expose api_base so cloud customers " "can override the upstream URL without env var access." @@ -228,16 +224,14 @@ def test_anthropic_provider_fields_support_byok(): api_base_field = fields_by_key["api_base"] assert api_base_field["required"] is False assert api_base_field["field_type"] == "text" - assert api_base_field.get( - "tooltip" - ), "api_base should have a tooltip explaining it is optional." + assert api_base_field.get("tooltip"), "api_base should have a tooltip explaining it is optional." # UI forms render fields in credential_fields order; api_base should come first # so an admin sees the URL override before the key field. field_order = [f["key"] for f in anthropic["credential_fields"]] - assert field_order.index("api_base") < field_order.index( - "api_key" - ), "api_base must appear before api_key in credential_fields (matches AI21 and ANTHROPIC_TEXT convention)." + assert field_order.index("api_base") < field_order.index("api_key"), ( + "api_base must appear before api_key in credential_fields (matches AI21 and ANTHROPIC_TEXT convention)." + ) def test_bedrock_mantle_provider_fields(): @@ -554,9 +548,7 @@ def test_google_ai_studio_provider_fields_expose_api_base(): assert response.status_code == 200 providers = response.json() - google_ai = next( - (p for p in providers if p["provider"] == "Google_AI_Studio"), None - ) + google_ai = next((p for p in providers if p["provider"] == "Google_AI_Studio"), None) assert google_ai is not None, "Google_AI_Studio provider entry not found" assert google_ai["litellm_provider"] == "gemini" @@ -576,18 +568,15 @@ def test_google_ai_studio_provider_fields_expose_api_base(): # placeholder shows the canonical URL so users still get the visual hint. # (See greptileai threads on PR #30419.) assert api_base_field["default_value"] is None - assert ( - api_base_field["placeholder"] - == "https://generativelanguage.googleapis.com/v1beta" - ) + assert api_base_field["placeholder"] == "https://generativelanguage.googleapis.com/v1beta" # UI forms render fields in credential_fields order; api_base should come # first so an admin sees the URL override before the key field (matches # OpenAI and Anthropic conventions). field_order = [f["key"] for f in google_ai["credential_fields"]] - assert field_order.index("api_base") < field_order.index( - "api_key" - ), "api_base must appear before api_key in credential_fields." + assert field_order.index("api_base") < field_order.index("api_key"), ( + "api_base must appear before api_key in credential_fields." + ) def test_public_model_hub_with_healthy_model(): @@ -615,20 +604,15 @@ def test_public_model_hub_with_healthy_model(): mock_llm_router = MagicMock() mock_prisma = MagicMock() - mock_prisma.get_all_latest_health_checks = AsyncMock( - return_value=[mock_health_check] - ) + mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[mock_health_check]) with ( patch("litellm.public_model_groups", ["gpt-3.5-turbo"]), patch("litellm.proxy.proxy_server.get_model_group_info") as mock_get_info, patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.health_endpoints._health_endpoints.convert_health_check_to_dict" - ) as mock_convert, + patch("litellm.proxy.health_endpoints._health_endpoints.convert_health_check_to_dict") as mock_convert, ): - mock_get_info.return_value = [mock_model_group] mock_convert.return_value = { "status": "healthy", @@ -673,20 +657,15 @@ def test_public_model_hub_with_unhealthy_model(): mock_llm_router = MagicMock() mock_prisma = MagicMock() - mock_prisma.get_all_latest_health_checks = AsyncMock( - return_value=[mock_health_check] - ) + mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[mock_health_check]) with ( patch("litellm.public_model_groups", ["gpt-4"]), patch("litellm.proxy.proxy_server.get_model_group_info") as mock_get_info, patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.health_endpoints._health_endpoints.convert_health_check_to_dict" - ) as mock_convert, + patch("litellm.proxy.health_endpoints._health_endpoints.convert_health_check_to_dict") as mock_convert, ): - mock_get_info.return_value = [mock_model_group] mock_convert.return_value = { "status": "unhealthy", @@ -732,7 +711,6 @@ def test_public_model_hub_without_health_check(): patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), ): - mock_get_info.return_value = [mock_model_group] response = client.get( @@ -789,9 +767,7 @@ def test_public_model_hub_mixed_health_statuses(): mock_llm_router = MagicMock() mock_prisma = MagicMock() - mock_prisma.get_all_latest_health_checks = AsyncMock( - return_value=[healthy_check, unhealthy_check] - ) + mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[healthy_check, unhealthy_check]) def convert_side_effect(check): if check.model_name == "gpt-3.5-turbo": @@ -813,11 +789,8 @@ def test_public_model_hub_mixed_health_statuses(): patch("litellm.proxy.proxy_server.get_model_group_info") as mock_get_info, patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.health_endpoints._health_endpoints.convert_health_check_to_dict" - ) as mock_convert, + patch("litellm.proxy.health_endpoints._health_endpoints.convert_health_check_to_dict") as mock_convert, ): - mock_get_info.return_value = [ healthy_model, unhealthy_model, @@ -1066,9 +1039,7 @@ def test_get_supported_endpoints_provider_fields(reset_endpoints_cache): def test_get_supported_endpoints_paths_start_with_slash(reset_endpoints_cache): endpoints = _make_client().get("/public/endpoints").json()["endpoints"] for item in endpoints: - assert item["endpoint"].startswith( - "/" - ), f"Expected path starting with /, got: {item['endpoint']}" + assert item["endpoint"].startswith("/"), f"Expected path starting with /, got: {item['endpoint']}" def test_get_supported_endpoints_chat_completions_present(reset_endpoints_cache): @@ -1092,9 +1063,9 @@ def test_get_supported_endpoints_display_names_have_no_slug_suffix( endpoints = _make_client().get("/public/endpoints").json()["endpoints"] for item in endpoints: for provider in item["providers"]: - assert not suffix_re.search( - provider["display_name"] - ), f"display_name still contains slug suffix: {provider['display_name']!r}" + assert not suffix_re.search(provider["display_name"]), ( + f"display_name still contains slug suffix: {provider['display_name']!r}" + ) def test_get_supported_endpoints_is_cached(reset_endpoints_cache): @@ -1320,7 +1291,6 @@ def test_public_mcp_hub_does_not_expose_upstream_url(): app.dependency_overrides.clear() - @pytest.fixture def reset_autorouter_presets_cache(): from litellm.proxy.public_endpoints.public_endpoints import _AutoRouterPresetsCache @@ -1332,9 +1302,7 @@ def reset_autorouter_presets_cache(): _AutoRouterPresetsCache.lock = None -def test_get_autorouter_presets_local_mode_serves_bundled_catalog( - monkeypatch, reset_autorouter_presets_cache -): +def test_get_autorouter_presets_local_mode_serves_bundled_catalog(monkeypatch, reset_autorouter_presets_cache): monkeypatch.setenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "True") app = FastAPI() app.include_router(router) @@ -1362,9 +1330,7 @@ def test_get_autorouter_presets_local_mode_serves_bundled_catalog( @pytest.mark.asyncio -async def test_get_autorouter_presets_fetches_once_per_process( - monkeypatch, reset_autorouter_presets_cache -): +async def test_get_autorouter_presets_fetches_once_per_process(monkeypatch, reset_autorouter_presets_cache): from litellm.proxy.public_endpoints.public_endpoints import ( _AUTOROUTER_PRESETS_ADAPTER, get_autorouter_presets, @@ -1376,7 +1342,9 @@ async def test_get_autorouter_presets_fetches_once_per_process( "remote_only": { "label": "Remote Only", "description": "from the remote catalog", - "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + "complexity_router_config": { + "tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]} + }, } } ) @@ -1411,7 +1379,9 @@ async def test_get_autorouter_presets_single_flight_on_concurrent_cold_start( "remote_only": { "label": "Remote Only", "description": "from the remote catalog", - "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + "complexity_router_config": { + "tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]} + }, } } ) @@ -1507,9 +1477,7 @@ async def test_autorouter_presets_adapter_rejects_wrong_shapes(): ) -def test_get_autorouter_presets_passes_unknown_catalog_fields_through( - monkeypatch, reset_autorouter_presets_cache -): +def test_get_autorouter_presets_passes_unknown_catalog_fields_through(monkeypatch, reset_autorouter_presets_cache): from litellm.proxy.public_endpoints.public_endpoints import ( _AUTOROUTER_PRESETS_ADAPTER, _AutoRouterPresetsCache, @@ -1551,12 +1519,14 @@ async def test_fetch_remote_autorouter_presets_parses_and_rejects_empty(monkeypa "remote_only": { "label": "Remote Only", "description": "from the remote catalog", - "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + "complexity_router_config": { + "tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]} + }, } } response = MagicMock() response.raise_for_status = MagicMock() - response.json = MagicMock(return_value=catalog) + response.content = json.dumps(catalog).encode() client = MagicMock() client.get = AsyncMock(return_value=response) monkeypatch.setattr(http_handler_module, "get_async_httpx_client", lambda llm_provider: client) @@ -1565,6 +1535,121 @@ async def test_fetch_remote_autorouter_presets_parses_and_rejects_empty(monkeypa assert presets["remote_only"].label == "Remote Only" response.raise_for_status.assert_called_once() - response.json = MagicMock(return_value={}) + response.content = b"{}" with pytest.raises(ValueError, match="empty"): await _fetch_remote_autorouter_presets("https://example.test/presets.json") + + +@pytest.fixture +def reset_whats_new_cache() -> Iterator[None]: + from litellm.proxy.public_endpoints.public_endpoints import _WhatsNewCache + + _WhatsNewCache.launches = None + _WhatsNewCache.lock = None + yield + _WhatsNewCache.launches = None + _WhatsNewCache.lock = None + + +def _remote_whats_new() -> WhatsNewResponse: + return WhatsNewResponse.model_validate( + { + "launches": [ + { + "icon": "sparkles", + "title": "Remote Launch", + "description": "from the remote list", + "href": "https://docs.litellm.ai/blog/remote", + "published_on": "2026-10-11", + } + ] + } + ) + + +def test_get_whats_new_local_mode_serves_bundled_list( + monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None +) -> None: + monkeypatch.setenv("LITELLM_LOCAL_WHATS_NEW", "True") + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + response = client.get("/public/whats_new") + + assert response.status_code == 200 + served: Final = WhatsNewResponse.model_validate_json(response.content) + bundled: Final = WhatsNewResponse.model_validate_json( + files("litellm.proxy.public_endpoints").joinpath("whats_new.json").read_text(encoding="utf-8") + ) + assert served == bundled + assert len(served.launches) > 0 + + +@pytest.mark.asyncio +async def test_get_whats_new_fetches_once_per_process( + monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None +) -> None: + from litellm.proxy.public_endpoints.public_endpoints import get_whats_new + + monkeypatch.delenv("LITELLM_LOCAL_WHATS_NEW", raising=False) + remote = _remote_whats_new() + calls: Final[list[str]] = [] + + async def fake_fetch(url: str) -> WhatsNewResponse: + calls.append(url) + return remote + + first = await get_whats_new(url="https://example.test/whats_new.json", fetch=fake_fetch) + second = await get_whats_new(url="https://example.test/whats_new.json", fetch=fake_fetch) + + assert first == remote + assert second == remote + assert calls == ["https://example.test/whats_new.json"] + + +@pytest.mark.asyncio +async def test_get_whats_new_caches_bundled_fallback_when_remote_is_unreachable( + monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None +) -> None: + from litellm.proxy.public_endpoints.public_endpoints import _load_bundled_whats_new, get_whats_new + + monkeypatch.delenv("LITELLM_LOCAL_WHATS_NEW", raising=False) + calls: Final[list[str]] = [] + + async def unreachable_fetch(url: str) -> WhatsNewResponse: + calls.append(url) + raise OSError("network is unreachable") + + first = await get_whats_new(url="https://example.test/whats_new.json", fetch=unreachable_fetch) + second = await get_whats_new(url="https://example.test/whats_new.json", fetch=unreachable_fetch) + + assert first == _load_bundled_whats_new() + assert second == first + assert len(calls) == 1 + + +@pytest.mark.asyncio +async def test_fetch_remote_whats_new_rejects_a_malformed_list_so_the_bundled_one_is_served( + monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None +) -> None: + import httpx + + from litellm.proxy.public_endpoints import public_endpoints + + monkeypatch.delenv("LITELLM_LOCAL_WHATS_NEW", raising=False) + request = httpx.Request("GET", "https://example.test/whats_new.json") + client = MagicMock() + client.get = AsyncMock( + return_value=httpx.Response( + 200, + request=request, + content=b'{"launches": [{"icon": "box", "title": "No date", "description": "d", "href": "https://x"}]}', + ) + ) + + with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=client): + served = await public_endpoints.get_whats_new(url="https://example.test/whats_new.json") + + assert served == public_endpoints._load_bundled_whats_new() + client.get.assert_awaited_once_with("https://example.test/whats_new.json", timeout=5.0) 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 7ef58695670..34c40df0218 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -459,6 +459,7 @@ def test_jev_config_requires_classifier_config() -> None: ("laya", "english", "laya"), ("bespoke", "nimble-latest", "bespoke"), ("databricks", "databricks-openjev-qwen35-4b", "databricks"), + ("databricks", "ai_decide", "databricks"), ], ) def test_classifier_aliases_load_and_serialize_one_canonical_config( @@ -546,14 +547,18 @@ async def test_oss_routes_with_its_own_credentials_and_accounts_the_checkpoint( assert recorder.calls[0]["response_cost"] == pytest.approx(0.31) -def test_databricks_requires_the_bare_serving_endpoint_name_as_the_model() -> None: +@pytest.mark.parametrize("model", ["ai_decide", "my-openjev", "databricks-openjev-qwen35-4b"]) +def test_databricks_accepts_ai_decide_or_a_bare_serving_endpoint_name(model: str) -> None: + assert JevClassifierConfig.model_validate({"provider": "databricks", "model": model}).model == model + + +def test_databricks_requires_a_model_that_cannot_rewrite_the_url() -> None: with pytest.raises(ValueError, match="model is required for provider 'databricks'"): JevClassifierConfig.model_validate({"provider": "databricks"}) - with pytest.raises(ValueError, match="bare endpoint name"): + with pytest.raises(ValueError, match="bare serving endpoint name"): JevClassifierConfig.model_validate({"provider": "databricks", "model": "serving-endpoints/my-openjev"}) - with pytest.raises(ValueError, match="bare endpoint name"): + with pytest.raises(ValueError, match="bare serving endpoint name"): JevClassifierConfig.model_validate({"provider": "databricks", "model": ".."}) - assert JevClassifierConfig.model_validate({"provider": "databricks", "model": "my-openjev"}).model == "my-openjev" @pytest.mark.asyncio @@ -584,7 +589,11 @@ async def test_databricks_routes_through_the_serving_endpoint_and_accounts_the_e "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 {}), + **( + {"api_key": "dapi-configured", "api_base": "https://workspace.test/serving-endpoints"} + if configured_key + else {} + ), }, "tiers": {"SIMPLE": "cheap"}, }, @@ -607,12 +616,61 @@ async def test_databricks_routes_through_the_serving_endpoint_and_accounts_the_e 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 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 +@pytest.mark.parametrize("configured_key", [False, True]) +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("api_base", ["https://workspace.test", "https://workspace.test/serving-endpoints/"]) +async def test_databricks_routes_through_ai_decide_without_a_model_field_and_accounts_ai_decide( + monkeypatch: pytest.MonkeyPatch, configured_key: bool, legacy: bool, api_base: str +) -> None: + monkeypatch.setenv("TYPESAFE_API_KEY", "never-send-typesafe-key") + monkeypatch.setenv("DATABRICKS_API_BASE", api_base) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + monkeypatch.setenv("DATABRICKS_TOKEN", "dapi-env-token") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _UsageRecorder("databricks/ai_decide") + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + router: Final = ComplexityRouter( + "databricks-route", + litellm.Router(model_list=[]), + { + "classifier_type": "jev" if legacy else "oss_classifier", + "jev_classifier_config" if legacy else "opensource_classifier_config": { + "provider": "databricks", + "model": "ai_decide", + **({"api_key": "dapi-configured", "api_base": api_base} if configured_key else {}), + }, + "tiers": {"SIMPLE": "cheap"}, + }, + derive_savings_baseline=False, + ) + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("https://workspace.test/api/2.0/ai-functions/ai-decide").respond( + 200, + json={"response": {"answers": {"tier": _answer().model_dump()}}, "metadata": {"version": "1.0"}}, + ) + outcome: Final = await router.aclassify("choose a tier") + await GLOBAL_LOGGING_WORKER.flush() + + assert outcome.cause == "jev_classifier" + assert outcome.jev_verdict is not None + assert (outcome.jev_verdict.provider, outcome.jev_verdict.model) == ("databricks", "ai_decide") + assert outcome.jev_verdict.label == "SIMPLE" + assert outcome.classifier_cost is None + sent: Final = route.calls.last.request + assert sent.headers.get("authorization") == ("Bearer dapi-configured" if configured_key else "Bearer dapi-env-token") + assert set(json.loads(sent.content)) == {"state", "questions"} + assert len(recorder.calls) == 1 + + def test_databricks_without_a_base_or_key_fails_at_client_build(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("DATABRICKS_API_BASE", raising=False) monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) @@ -623,7 +681,7 @@ def test_databricks_without_a_base_or_key_fails_at_client_build(monkeypatch: pyt litellm.Router(model_list=[]), { "classifier_type": "oss_classifier", - "opensource_classifier_config": {"provider": "databricks", "model": "databricks-openjev-qwen35-4b"}, + "opensource_classifier_config": {"provider": "databricks", "model": "ai_decide"}, "tiers": {"SIMPLE": "cheap"}, }, derive_savings_baseline=False, @@ -670,17 +728,12 @@ def test_jev_api_base_without_its_own_key_is_rejected_so_the_environment_key_sta def test_databricks_api_base_without_its_own_key_is_rejected_so_the_workspace_token_stays_home() -> None: with pytest.raises(ValueError, match="DATABRICKS_API_KEY or DATABRICKS_TOKEN is only sent to DATABRICKS_API_BASE"): JevClassifierConfig.model_validate( - {"provider": "databricks", "model": "my-openjev", "api_base": "https://collector.invalid/serving-endpoints"} + {"provider": "databricks", "model": "ai_decide", "api_base": "https://collector.invalid"} ) paired: Final = JevClassifierConfig.model_validate( - { - "provider": "databricks", - "model": "my-openjev", - "api_base": "https://workspace.invalid/serving-endpoints", - "api_key": "dapi-own", - } + {"provider": "databricks", "model": "ai_decide", "api_base": "https://workspace.invalid", "api_key": "dapi-own"} ) - assert (paired.api_base, paired.api_key) == ("https://workspace.invalid/serving-endpoints", "dapi-own") + assert (paired.api_base, paired.api_key) == ("https://workspace.invalid", "dapi-own") @pytest.mark.parametrize( @@ -776,3 +829,143 @@ async def test_http_jev_classifier_client_posts_to_system_one() -> None: assert captured["content_type"] == "application/json" assert captured["body"] == request.model_dump(mode="json") assert response.model == "jev-1.13.0" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [200, 503]) +@pytest.mark.parametrize("provider,model", [("typesafe", "jev-latest"), ("databricks", "my-endpoint")]) +async def test_deployed_decision_classifier_uses_registered_connection( + status_code: int, provider: str, model: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(litellm.model_cost, f"{provider}/{model}", {"input_cost_per_token": 0.01}) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "decision-deployment", + "litellm_params": { + "model": f"{provider}/{model}", + "api_key": "deployment-key", + "api_base": "https://decision.test/serving-endpoints" + if provider == "databricks" + else "https://decision.test", + }, + } + ], + num_retries=0, + ) + classifier: Final = ComplexityRouter( + "deployed-classifier", + router, + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"deployment_name": "decision-deployment"}, + "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}, + }, + derive_savings_baseline=False, + ) + with respx.mock(assert_all_called=True) as requests: + route: Final = requests.post( + "https://decision.test/serving-endpoints/my-endpoint/invocations" + if provider == "databricks" + else "https://decision.test/v1/systemone" + ).mock( + return_value=httpx.Response( + status_code, + json={ + "model": model, + "answers": {"tier": _answer().model_dump()}, + "usage": {"input_tokens": 31, "output_tokens": 0}, + }, + ) + ) + outcome: Final = await classifier.aclassify("What is two plus two?") + assert route.call_count == 1 + assert route.calls[0].request.headers["authorization"] == "Bearer deployment-key" + assert json.loads(route.calls[0].request.content)["model"] == model + assert (outcome.cause == "jev_classifier") is (status_code == 200) + if status_code == 200: + assert outcome.jev_verdict is not None + assert outcome.jev_verdict.label == _answer().choice + + assert outcome.jev_verdict.provider == provider + assert outcome.classifier_cost == pytest.approx(0.31) + + +@pytest.mark.parametrize("deployment_name", ["", " "]) +def test_deployed_decision_classifier_rejects_blank_deployment_name(deployment_name: str) -> None: + with pytest.raises(ValueError, match="deployment_name"): + JevClassifierConfig.model_validate({"deployment_name": deployment_name}) + + +@pytest.mark.parametrize("connection", [{"api_key": "secret"}, {"api_base": "https://override.test"}]) +def test_deployed_decision_classifier_rejects_connection_overrides(connection: Mapping[str, str]) -> None: + with pytest.raises(ValueError, match="deployment connection"): + JevClassifierConfig.model_validate({"deployment_name": "deployed", **connection}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["timeout", "cancelled", "invalid"]) +async def test_deployed_decision_classifier_preserves_failure_handling( + failure: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "decision-deployment", + "litellm_params": { + "model": "typesafe/jev-latest", + "api_key": "test", + "api_base": "https://decision.test", + }, + } + ], + num_retries=0, + ) + classifier: Final = ComplexityRouter( + "decision-router", + router, + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"deployment_name": "decision-deployment"}, + "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}, + }, + derive_savings_baseline=False, + ) + cancellation_reached: Final = asyncio.Event() + with respx.mock(assert_all_called=failure != "cancelled") as upstream: + route: Final = upstream.post("https://decision.test/v1/systemone") + if failure == "timeout": + route.mock(side_effect=httpx.ReadTimeout("classifier timeout")) + elif failure == "cancelled": + + async def cancel_request(request: httpx.Request) -> httpx.Response: + cancellation_reached.set() + raise asyncio.CancelledError() + + route.mock(side_effect=cancel_request) + else: + route.respond(json={"answers": {"tier": _answer("not-a-tier").model_dump()}}) + if failure == "cancelled": + with pytest.raises(asyncio.CancelledError): + await classifier.aclassify("short question") + assert cancellation_reached.is_set() + return + outcome: Final = await classifier.aclassify("short question") + assert outcome.cause != "jev_classifier" + assert outcome.jev_verdict is None + + +def test_switching_to_deployed_classifier_does_not_copy_legacy_credentials() -> None: + from litellm.router_strategy.complexity_router.config import resolve_complexity_router_config_write + + write: Final = resolve_complexity_router_config_write( + {"classifier_type": "oss_classifier", "opensource_classifier_config": {"deployment_name": "registered"}}, + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"api_key": "legacy-secret", "api_base": "https://legacy.test"}, + }, + ) + assert write.effective is not None + assert write.effective["opensource_classifier_config"] == {"deployment_name": "registered"} diff --git a/tests/unit/router_utils/test_auto_router_model_naming.py b/tests/unit/router_utils/test_auto_router_model_naming.py index d7016cbc9ec..1cfa3aa5c91 100644 --- a/tests/unit/router_utils/test_auto_router_model_naming.py +++ b/tests/unit/router_utils/test_auto_router_model_naming.py @@ -755,3 +755,20 @@ def test_chained_heuristic_and_judge_prompt_claim_each_active_capability( assert count_capability_routers(({"litellm_params": params},), capability=capability) == ( 1 if capability.key in expected else 0 ) + + +def test_deployed_decision_classifier_is_a_model_access_dependency() -> None: + dependencies: Final = strategy_router_dependencies( + { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"deployment_name": "decision-deployment"}, + "tiers": {"SIMPLE": "cheap"}, + }, + } + ) + assert tuple((dependency.model_name, dependency.role) for dependency in dependencies) == ( + ("cheap", "tier"), + ("decision-deployment", "classifier"), + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx index ad3f80b7d78..a17a18beee5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx @@ -6,7 +6,6 @@ import type { DailyActivityAggregatedResponse } from "@/components/UsagePage/dai import { EMPTY_DAILY_ACTIVITY_METADATA } from "@/components/UsagePage/dailyActivityApi"; import { all_admin_roles } from "@/utils/roles"; import HomePage from "./HomePage"; -import { WHATS_NEW_ITEMS } from "./homeContent"; const auth = vi.hoisted(() => ({ accessToken: "sk-test", @@ -31,6 +30,28 @@ vi.mock("@/components/networking", async (importOriginal) => ({ const blogPosts = { posts: [{ title: "Post one", description: "d", date: "2026-10-09", url: "https://x/1" }] }; +const launch = (title: string, icon: string, published_on: string) => ({ + icon, + title, + description: `${title} description`, + href: `https://docs.litellm.ai/blog/${title}`, + published_on, +}); + +const whatsNew = vi.hoisted(() => ({ launches: [] as unknown[] })); + +const requestUrl = (input: RequestInfo | URL): string => { + if (typeof input === "string") return input; + return input instanceof URL ? input.href : input.url; +}; + +const fakeFetch = (input: RequestInfo | URL) => { + const body = requestUrl(input).endsWith("/public/whats_new") ? whatsNew : blogPosts; + return Promise.resolve( + new Response(JSON.stringify(body), { status: 200, headers: { "Content-Type": "application/json" } }), + ); +}; + const dayResult = (date: string, model: string, tokens: number) => { const metrics = { spend: 1, @@ -86,7 +107,8 @@ describe("HomePage", () => { network.aggregated.mockReset().mockResolvedValue(activity); network.gateway.mockReset().mockResolvedValue({ total_successful_requests: 0, total_failed_requests: 0 }); network.userInfo.mockReset().mockResolvedValue({ user_info: { max_budget: null } }); - vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(JSON.stringify(blogPosts), { status: 200 }))); + whatsNew.launches = [launch("mid", "box", "2026-10-06"), launch("new", "zap", "2026-10-11")]; + vi.stubGlobal("fetch", vi.fn(fakeFetch)); }); afterEach(() => { vi.unstubAllGlobals(); @@ -121,13 +143,32 @@ describe("HomePage", () => { expect(network.aggregated).toHaveBeenCalledTimes(2); }); - it("lists the curated launches newest first with their links", () => { + it("renders the launches served by /public/whats_new newest first, whatever order the JSON lists them in", async () => { + whatsNew.launches = [ + launch("mid", "box", "2026-10-06"), + launch("new", "an-icon-from-a-newer-list", "2026-10-11"), + launch("old", "scale", "2026-09-30"), + ]; renderHome(); - const section = screen.getByTestId("home-whats-new"); - const links = within(section).getAllByRole("link"); - expect(links.map((link) => link.getAttribute("href"))).toEqual(WHATS_NEW_ITEMS.map((item) => item.href)); - const dates = WHATS_NEW_ITEMS.map((item) => item.publishedOn); - expect(dates).toEqual([...dates].sort().reverse()); + const section = await screen.findByTestId("home-whats-new"); + expect(await within(section).findByText("new")).toBeInTheDocument(); + expect( + within(section) + .getAllByRole("link") + .map((link) => link.getAttribute("href")), + ).toEqual([ + "https://docs.litellm.ai/blog/new", + "https://docs.litellm.ai/blog/mid", + "https://docs.litellm.ai/blog/old", + ]); + expect(within(section).getByText("Oct 11")).toBeInTheDocument(); + }); + + it("hides What's new when /public/whats_new has no launches", async () => { + whatsNew.launches = []; + renderHome(); + expect(await screen.findByRole("link", { name: /post one/i })).toBeInTheDocument(); + expect(screen.queryByTestId("home-whats-new")).not.toBeInTheDocument(); }); it("ranks model groups by token share on the leaderboard", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx index f4519958e9d..bcae30795c6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx @@ -4,21 +4,24 @@ import { ArrowRight, Plus, Sparkles } from "lucide-react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; +import { Skeleton } from "@/components/ui/skeleton"; import { Page } from "@/components/shared/Page"; import { PageHeader, PageHeaderControls, PageHeaderTitle } from "@/components/shared/PageHeader"; import { uiHref } from "@/utils/uiHref"; import { Leaderboard, useBreakdown } from "@/app/(dashboard)/usage/_components/components/overview/BreakdownChart"; import { Panel } from "@/app/(dashboard)/usage/_components/components/overview/Primitives"; import UsageStatStrip from "@/app/(dashboard)/usage/_components/components/overview/UsageStatStrip"; -import { formatPublishedOn, WHATS_NEW_ITEMS } from "./homeContent"; +import { formatPublishedOn, launchIcon } from "./homeContent"; import UpdatesFeed from "./UpdatesFeed"; import { HOME_USAGE_DAYS, useHomeUsage } from "./useHomeUsage"; +import { useWhatsNew } from "./useWhatsNew"; const LEADERBOARD = { metric: "tokens", dimension: "model_groups" } as const; export default function HomePage() { const { isViewOnly } = useAuthorized(); const usage = useHomeUsage(); + const whatsNew = useWhatsNew(); const { series, ranking } = useBreakdown(usage.results, LEADERBOARD, 8); return ( @@ -56,34 +59,40 @@ export default function HomePage() {
{item.description}
-{item.description}
+