mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(batches): route header and model-param batch uploads and creates to the group's batch deployment
A model group that mixes an Azure chat deployment with an Azure batch deployment sent batch file uploads, batch creates, and batch lists named by the group to the first deployment, which cannot run batches. Batch-purpose uploads and model-param batch calls now resolve credentials for the group's batch-ready deployment, and the ids they return carry that deployment so later retrieve, cancel, content, and create-from-file calls reach the same resource. Ids without the hint keep resolving to the group's first deployment
This commit is contained in:
mateo-berri
parent
67ae3fb527
commit
fbfa364b4a
9 changed files with 499 additions and 88 deletions
|
|
@ -52,9 +52,10 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 #
|
|||
apply_team_provider_credentials,
|
||||
authorize_model_for_key,
|
||||
batch_cost_poller_is_active,
|
||||
batch_deployment_id,
|
||||
decode_model_from_file_id,
|
||||
encode_batch_response_ids,
|
||||
encode_file_id_with_model,
|
||||
encode_file_id_with_model, # pyright: ignore[reportUnusedImport] # backwards-compatible package export
|
||||
ensure_batch_response_managed_file_ids,
|
||||
get_authorized_credentials_for_model,
|
||||
get_batch_from_database,
|
||||
|
|
@ -62,6 +63,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 #
|
|||
get_model_id_from_unified_batch_id,
|
||||
get_models_from_unified_file_id,
|
||||
get_original_file_id,
|
||||
hinted_batch_deployment_id,
|
||||
is_base64_encoded_unified_file_id,
|
||||
is_litellm_executed_batch,
|
||||
prepare_data_with_credentials,
|
||||
|
|
@ -360,11 +362,13 @@ async def create_batch(
|
|||
|
||||
# SCENARIO 1: File ID is encoded with model info
|
||||
if model_from_file_id is not None and input_file_id:
|
||||
input_file_deployment_id: Final = hinted_batch_deployment_id(llm_router, model_from_file_id, input_file_id)
|
||||
credentials = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_from_file_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch creation (file created with model)",
|
||||
deployment_id=input_file_deployment_id,
|
||||
)
|
||||
|
||||
original_file_id: Final = get_original_file_id(input_file_id)
|
||||
|
|
@ -380,35 +384,11 @@ async def create_batch(
|
|||
**_create_batch_data,
|
||||
)
|
||||
|
||||
# Encode the batch ID and related file IDs with model information
|
||||
if response and hasattr(response, "id") and response.id:
|
||||
original_batch_id: Final = response.id
|
||||
encoded_batch_id: Final = encode_file_id_with_model(
|
||||
file_id=original_batch_id,
|
||||
model=model_from_file_id,
|
||||
id_type="batch",
|
||||
)
|
||||
response.id = encoded_batch_id
|
||||
|
||||
if hasattr(response, "output_file_id") and response.output_file_id:
|
||||
response.output_file_id = encode_file_id_with_model(
|
||||
file_id=response.output_file_id, model=model_from_file_id
|
||||
)
|
||||
|
||||
if hasattr(response, "error_file_id") and response.error_file_id:
|
||||
response.error_file_id = encode_file_id_with_model(
|
||||
file_id=response.error_file_id, model=model_from_file_id
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Created batch using model: %s, original_batch_id: %s, encoded: %s",
|
||||
model_from_file_id,
|
||||
original_batch_id,
|
||||
encoded_batch_id,
|
||||
)
|
||||
|
||||
encode_batch_response_ids(response, model=model_from_file_id, deployment_id=input_file_deployment_id)
|
||||
response.input_file_id = input_file_id
|
||||
|
||||
verbose_proxy_logger.debug("Created batch using model: %s", model_from_file_id)
|
||||
|
||||
elif litellm.enable_loadbalancing_on_batch_endpoints is True and is_router_model and router_model is not None:
|
||||
if llm_router is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -465,11 +445,13 @@ async def create_batch(
|
|||
# SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback
|
||||
if model_param:
|
||||
# SCENARIO 2: Use model-based routing from header/query/body
|
||||
model_param_deployment_id: Final = batch_deployment_id(llm_router, model_param)
|
||||
credentials = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_param,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch creation",
|
||||
deployment_id=model_param_deployment_id,
|
||||
)
|
||||
await _raise_when_input_file_must_be_managed(model_param, credentials)
|
||||
|
||||
|
|
@ -484,7 +466,7 @@ async def create_batch(
|
|||
**_create_batch_data,
|
||||
)
|
||||
|
||||
encode_batch_response_ids(response, model=model_param)
|
||||
encode_batch_response_ids(response, model=model_param, deployment_id=model_param_deployment_id)
|
||||
|
||||
verbose_proxy_logger.debug("Created batch using model: %s", model_param)
|
||||
else:
|
||||
|
|
@ -709,11 +691,13 @@ async def retrieve_batch(
|
|||
# Retrieve from provider (for non-terminal states or if DB lookup failed)
|
||||
# SCENARIO 1: Batch ID is encoded with model info
|
||||
if model_from_id is not None:
|
||||
retrieve_deployment_id: Final = hinted_batch_deployment_id(llm_router, model_from_id, batch_id)
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_from_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch retrieval (batch created with model)",
|
||||
deployment_id=retrieve_deployment_id,
|
||||
)
|
||||
|
||||
original_batch_id: Final = get_original_file_id(batch_id)
|
||||
|
|
@ -728,7 +712,9 @@ async def retrieve_batch(
|
|||
# so litellm.aretrieve_batch can load BedrockBatchesConfig. Without
|
||||
# it the call falls into the legacy provider switch and 400s.
|
||||
data["model"] = model_from_id
|
||||
add_deployment_model_info(data=data, llm_router=llm_router, model_id=model_from_id)
|
||||
add_deployment_model_info(
|
||||
data=data, llm_router=llm_router, model_id=retrieve_deployment_id or model_from_id
|
||||
)
|
||||
_litellm_metadata_of(data).setdefault("model_group", model_from_id)
|
||||
|
||||
# Retrieve batch using model credentials
|
||||
|
|
@ -737,7 +723,7 @@ async def retrieve_batch(
|
|||
**data,
|
||||
)
|
||||
|
||||
encode_batch_response_ids(response, model=model_from_id)
|
||||
encode_batch_response_ids(response, model=model_from_id, deployment_id=retrieve_deployment_id)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Retrieved batch using model: %s, original_id: %s", model_from_id, original_batch_id
|
||||
|
|
@ -939,11 +925,15 @@ async def list_batches(
|
|||
data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model")
|
||||
):
|
||||
# SCENARIO 2: Use model-based routing from header/query/body
|
||||
list_deployment_id: Final = (
|
||||
batch_deployment_id(llm_router, model_param) if isinstance(model_param, str) else None
|
||||
)
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_param,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch listing",
|
||||
deployment_id=list_deployment_id,
|
||||
)
|
||||
|
||||
prepare_data_with_credentials(data=data, credentials=credentials)
|
||||
|
|
@ -960,7 +950,7 @@ async def list_batches(
|
|||
response_data: Final = getattr(response, "data", None)
|
||||
if response_data:
|
||||
for batch in response_data:
|
||||
encode_batch_response_ids(batch, model=model_param)
|
||||
encode_batch_response_ids(batch, model=model_param, deployment_id=list_deployment_id)
|
||||
|
||||
verbose_proxy_logger.debug("Listed batches using model: %s", model_param)
|
||||
|
||||
|
|
@ -1140,11 +1130,13 @@ async def cancel_batch(
|
|||
|
||||
# SCENARIO 1: Batch ID is encoded with model info
|
||||
if model_from_id is not None:
|
||||
cancel_deployment_id: Final = hinted_batch_deployment_id(llm_router, model_from_id, batch_id)
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_from_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch cancellation (batch created with model)",
|
||||
deployment_id=cancel_deployment_id,
|
||||
)
|
||||
|
||||
original_batch_id: Final = get_original_file_id(batch_id)
|
||||
|
|
@ -1162,7 +1154,7 @@ async def cancel_batch(
|
|||
**data,
|
||||
)
|
||||
|
||||
encode_batch_response_ids(response, model=model_from_id)
|
||||
encode_batch_response_ids(response, model=model_from_id, deployment_id=cancel_deployment_id)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Cancelled batch using model: %s, original_id: %s", model_from_id, original_batch_id
|
||||
|
|
|
|||
|
|
@ -229,11 +229,13 @@ def is_litellm_executed_batch(decoded_unified_batch_id: str) -> bool:
|
|||
return bool(marker) and batch_id.startswith(LITELLM_EXECUTED_BATCH_ID_PREFIX)
|
||||
|
||||
|
||||
def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file", "batch"] = "file") -> str:
|
||||
def encode_file_id_with_model(
|
||||
file_id: str, model: str, id_type: Literal["file", "batch"] = "file", deployment_id: str | None = None
|
||||
) -> str:
|
||||
"""
|
||||
Encode a file/batch ID with model routing information.
|
||||
|
||||
Format: <prefix><base64(litellm:<original_id>;model,<model_name>)>
|
||||
Format: <prefix><base64(litellm:<original_id>;model,<model_name>[;deployment,<deployment_id>])>
|
||||
The result preserves the original prefix (file-, batch_, etc.) for OpenAI compliance.
|
||||
|
||||
Args:
|
||||
|
|
@ -242,6 +244,8 @@ def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file",
|
|||
id_type: Type of ID being encoded. Used to determine the correct prefix when
|
||||
the raw ID lacks a recognizable prefix (e.g., Vertex AI numeric IDs).
|
||||
Defaults to "file" for backward compatibility.
|
||||
deployment_id: The group deployment that holds the resource, set only when the group's
|
||||
first deployment cannot serve it (a mixed Azure chat and batch group).
|
||||
|
||||
Returns:
|
||||
Encoded ID starting with appropriate prefix and containing routing information
|
||||
|
|
@ -256,7 +260,8 @@ def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file",
|
|||
encode_file_id_with_model("3814889423749775360", "gemini-2.5-pro", id_type="batch")
|
||||
-> "batch_bGl0ZWxsbTozODE0ODg5NDIzNzQ5Nzc1MzYwO21vZGVsLGdlbWluaS0yLjUtcHJv"
|
||||
"""
|
||||
encoded_str: Final = f"litellm:{file_id};model,{model}"
|
||||
deployment_segment: Final = f";deployment,{deployment_id}" if deployment_id else ""
|
||||
encoded_str: Final = f"litellm:{file_id};model,{model}{deployment_segment}"
|
||||
encoded_bytes: Final = base64.urlsafe_b64encode(encoded_str.encode())
|
||||
encoded_b64: Final = encoded_bytes.decode().rstrip("=")
|
||||
|
||||
|
|
@ -273,47 +278,46 @@ def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file",
|
|||
return f"{prefix}{encoded_b64}"
|
||||
|
||||
|
||||
def encode_batch_response_ids(response, model: str) -> None:
|
||||
def encode_batch_response_ids(response, model: str, deployment_id: str | None = None) -> None:
|
||||
"""Encode all IDs in a batch response with model routing info (in-place)."""
|
||||
if not response or not hasattr(response, "id") or not response.id:
|
||||
return
|
||||
response.id = encode_file_id_with_model(file_id=response.id, model=model, id_type="batch")
|
||||
response.id = encode_file_id_with_model(
|
||||
file_id=response.id, model=model, id_type="batch", deployment_id=deployment_id
|
||||
)
|
||||
for attr in ("output_file_id", "error_file_id", "input_file_id"):
|
||||
if hasattr(response, attr) and getattr(response, attr):
|
||||
setattr(
|
||||
response,
|
||||
attr,
|
||||
encode_file_id_with_model(file_id=getattr(response, attr), model=model),
|
||||
encode_file_id_with_model(file_id=getattr(response, attr), model=model, deployment_id=deployment_id),
|
||||
)
|
||||
|
||||
|
||||
def decode_model_from_file_id(encoded_id: str) -> str | None:
|
||||
def _model_embedded_payload(encoded_id: object) -> str | None:
|
||||
if not isinstance(encoded_id, str):
|
||||
return None
|
||||
prefix: Final = next((prefix for prefix in ("file-", "batch_") if encoded_id.startswith(prefix)), "")
|
||||
b64_part: Final = encoded_id[len(prefix) :]
|
||||
try:
|
||||
decoded: Final = base64.urlsafe_b64decode(b64_part + "=" * (-len(b64_part) % 4)).decode()
|
||||
except ValueError:
|
||||
return None
|
||||
return decoded if decoded.startswith("litellm:") and ";model," in decoded else None
|
||||
|
||||
|
||||
def decode_model_from_file_id(encoded_id: object) -> str | None:
|
||||
"""
|
||||
Extract model name from an encoded file/batch ID.
|
||||
Handles IDs that start with "file-" or "batch_" prefix.
|
||||
"""
|
||||
try:
|
||||
if not isinstance(encoded_id, str):
|
||||
return None
|
||||
match: Final = re.search(r";model,([^;]+)", _model_embedded_payload(encoded_id) or "")
|
||||
return match.group(1).strip() if match else None
|
||||
|
||||
# Remove prefix if present (file-, batch_, etc.)
|
||||
if encoded_id.startswith("file-"):
|
||||
b64_part = encoded_id[5:] # Remove "file-"
|
||||
elif encoded_id.startswith("batch_"):
|
||||
b64_part = encoded_id[6:] # Remove "batch_"
|
||||
else:
|
||||
b64_part = encoded_id
|
||||
|
||||
padded: Final = b64_part + "=" * (-len(b64_part) % 4)
|
||||
decoded: Final = base64.urlsafe_b64decode(padded).decode()
|
||||
if decoded.startswith("litellm:") and ";model," in decoded:
|
||||
match: Final = re.search(r";model,([^;]+)", decoded)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
def decode_deployment_from_file_id(encoded_id: object) -> str | None:
|
||||
match: Final = re.search(r";deployment,([^;]+)", _model_embedded_payload(encoded_id) or "")
|
||||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def get_original_file_id(encoded_id: str) -> str:
|
||||
|
|
@ -321,29 +325,8 @@ def get_original_file_id(encoded_id: str) -> str:
|
|||
Extract the original provider file/batch ID from an encoded ID.
|
||||
Handles IDs that start with "file-" or "batch_" prefix.
|
||||
"""
|
||||
try:
|
||||
if not isinstance(encoded_id, str):
|
||||
return encoded_id
|
||||
|
||||
# Remove prefix if present (file-, batch_, etc.)
|
||||
if encoded_id.startswith("file-"):
|
||||
b64_part = encoded_id[5:] # Remove "file-"
|
||||
elif encoded_id.startswith("batch_"):
|
||||
b64_part = encoded_id[6:] # Remove "batch_"
|
||||
else:
|
||||
b64_part = encoded_id
|
||||
|
||||
padded: Final = b64_part + "=" * (-len(b64_part) % 4)
|
||||
decoded: Final = base64.urlsafe_b64decode(padded).decode()
|
||||
|
||||
if decoded.startswith("litellm:") and ";model," in decoded:
|
||||
match: Final = re.search(r"litellm:([^;]+);model,", decoded)
|
||||
if match:
|
||||
return match.group(1)
|
||||
|
||||
return encoded_id
|
||||
except Exception:
|
||||
return encoded_id
|
||||
match: Final = re.search(r"litellm:([^;]+);model,", _model_embedded_payload(encoded_id) or "")
|
||||
return match.group(1) if match else encoded_id
|
||||
|
||||
|
||||
def is_model_embedded_id(file_id: str) -> bool:
|
||||
|
|
@ -467,16 +450,32 @@ async def get_authorized_credentials_for_model(
|
|||
model_id: str,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
operation_context: str = "file operation",
|
||||
deployment_id: str | None = None,
|
||||
) -> dict: # mutable-ok: same contract as get_credentials_for_model, callers merge it into request data
|
||||
"""``get_credentials_for_model`` gated by ``authorize_model_for_key``."""
|
||||
await authorize_model_for_key(model_id=model_id, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
|
||||
return get_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_id,
|
||||
model_id=deployment_id or model_id,
|
||||
operation_context=operation_context,
|
||||
)
|
||||
|
||||
|
||||
def batch_deployment_id(
|
||||
llm_router: Optional["Router"], model: str, preferred_deployment_id: str | None = None
|
||||
) -> str | None:
|
||||
return (
|
||||
llm_router.get_batch_deployment_id(model_group_name=model, preferred_deployment_id=preferred_deployment_id)
|
||||
if llm_router is not None
|
||||
else None
|
||||
)
|
||||
|
||||
|
||||
def hinted_batch_deployment_id(llm_router: Optional["Router"], model: str, encoded_id: str) -> str | None:
|
||||
hint: Final = decode_deployment_from_file_id(encoded_id)
|
||||
return batch_deployment_id(llm_router, model, preferred_deployment_id=hint) if hint else None
|
||||
|
||||
|
||||
def deployment_drop_params(
|
||||
llm_router: Optional["Router"], model_id: str, request: Mapping[str, object]
|
||||
) -> Mapping[str, bool | str]:
|
||||
|
|
@ -801,7 +800,7 @@ def prepare_data_with_credentials(
|
|||
async def handle_model_based_routing(
|
||||
file_id: str,
|
||||
request, # FastAPI Request object
|
||||
llm_router, # Router instance
|
||||
llm_router: Optional["Router"],
|
||||
data: dict,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
check_file_id_encoding: bool = True,
|
||||
|
|
@ -845,6 +844,7 @@ async def handle_model_based_routing(
|
|||
model_id=model_from_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="file operation (file created with model)",
|
||||
deployment_id=hinted_batch_deployment_id(llm_router, model_from_id, file_id),
|
||||
)
|
||||
original_file_id: Final = get_original_file_id(file_id)
|
||||
return True, model_from_id, original_file_id, credentials
|
||||
|
|
|
|||
|
|
@ -80,6 +80,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 #
|
|||
add_internal_model_credentials,
|
||||
apply_team_provider_credentials,
|
||||
authorize_model_for_key,
|
||||
batch_deployment_id,
|
||||
deployment_drop_params,
|
||||
encode_file_id_with_model,
|
||||
extract_file_creation_params,
|
||||
|
|
@ -454,18 +455,23 @@ async def route_create_file(
|
|||
|
||||
# NEW: Handle model-based routing (no DB required)
|
||||
if model is not None:
|
||||
deployment_id: Final = batch_deployment_id(llm_router, model) if purpose == "batch" else None
|
||||
# Get credentials from model_list via router
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="file upload",
|
||||
deployment_id=deployment_id,
|
||||
)
|
||||
|
||||
# Merge credentials into the request
|
||||
prepare_data_with_credentials(
|
||||
data=_create_file_request,
|
||||
credentials={**credentials, **deployment_drop_params(llm_router, model, _create_file_request)},
|
||||
credentials={
|
||||
**credentials,
|
||||
**deployment_drop_params(llm_router, deployment_id or model, _create_file_request),
|
||||
},
|
||||
)
|
||||
|
||||
# Create the file with model credentials
|
||||
|
|
@ -477,7 +483,7 @@ async def route_create_file(
|
|||
# Encode the file ID with model information
|
||||
if response and hasattr(response, "id") and response.id:
|
||||
original_id: Final = response.id
|
||||
encoded_id: Final = encode_file_id_with_model(file_id=original_id, model=model)
|
||||
encoded_id: Final = encode_file_id_with_model(file_id=original_id, model=model, deployment_id=deployment_id)
|
||||
response.id = encoded_id
|
||||
verbose_proxy_logger.debug("Encoded file ID: %s -> %s (model: %s)", original_id, encoded_id, model)
|
||||
|
||||
|
|
|
|||
|
|
@ -724,6 +724,14 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj
|
|||
return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {})
|
||||
|
||||
|
||||
def _unpaused_deployment_id(deployment: Mapping[str, object]) -> str | None:
|
||||
model_info: Final = _configured_model_info(deployment)
|
||||
deployment_id: Final = model_info.get("id")
|
||||
if model_info.get("blocked") is True or not isinstance(deployment_id, str):
|
||||
return None
|
||||
return deployment_id
|
||||
|
||||
|
||||
def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None:
|
||||
from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request
|
||||
|
||||
|
|
@ -12021,6 +12029,26 @@ class Router:
|
|||
return None
|
||||
return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable
|
||||
|
||||
def get_batch_deployment_id(self, model_group_name: str, preferred_deployment_id: str | None = None) -> str | None:
|
||||
"""
|
||||
The shared deployment of ``model_group_name`` that file and batch calls
|
||||
should reach when the group mixes batch and non-batch deployments, or
|
||||
None when the batch rule drops nothing and the group name resolves as before.
|
||||
"""
|
||||
group: Final = tuple(
|
||||
deployment
|
||||
for deployment in self.deployments_named(model_group_name) or ()
|
||||
if self.deployment_usable_by_team(deployment, None)
|
||||
)
|
||||
with batch_create_routing():
|
||||
batch_ready: Final = without_deployments_that_cannot_batch(group, caller_deployments=lambda: group)
|
||||
if len(batch_ready) == len(group):
|
||||
return None
|
||||
ready_ids: Final = tuple(
|
||||
deployment_id for deployment_id in map(_unpaused_deployment_id, batch_ready) if deployment_id is not None
|
||||
)
|
||||
return preferred_deployment_id if preferred_deployment_id in ready_ids else next(iter(ready_ids), None)
|
||||
|
||||
async def arefresh_model_info(self, *, client: AsyncHTTPHandler | None = None) -> None:
|
||||
"""Refresh token limits advertised by configured OpenAI-compatible deployments."""
|
||||
deployments: Final = iter(tuple(self.model_list))
|
||||
|
|
|
|||
|
|
@ -95,6 +95,7 @@ ignored_function_names = [
|
|||
"arm_routing_read_prefetch", # Tested in tests/unit/caching/test_request_redis_batch_pre_call.py (file lacks "router" in name)
|
||||
"_deployment_params_base_model", # Tested through log_retry in test_router_helper_utils.py
|
||||
"_configured_model_info", # Tested through get_configured_service_tiers in test_router.py
|
||||
"_unpaused_deployment_id", # Tested through the paused batch deployment test of get_batch_deployment_id in test_router.py
|
||||
"_every_deployment_refused", # Tested through the unsupported transcription retry tests in test_router.py
|
||||
"_routable_deployments", # Tested through get_configured_service_tiers and get_routable_upstream_model in test_router.py
|
||||
"_async_get_available_deployment", # Body of the `route {model}` phase wrapper, exercised through async_get_available_deployment in test_router.py
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ cannot drift without a test failure.
|
|||
import base64
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
|
@ -50,7 +51,10 @@ import litellm.proxy.proxy_server as proxy_server
|
|||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
decode_deployment_from_file_id,
|
||||
decode_model_from_file_id,
|
||||
encode_file_id_with_model,
|
||||
get_original_file_id,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
|
@ -198,6 +202,7 @@ def harness(monkeypatch: pytest.MonkeyPatch):
|
|||
router.model_list = []
|
||||
router.acreate_batch = AsyncMock(return_value=make_batch())
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
||||
router.get_batch_deployment_id = MagicMock(return_value=None)
|
||||
|
||||
read_body = AsyncMock(side_effect=lambda request: body_holder["body"])
|
||||
pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock()))
|
||||
|
|
@ -1457,6 +1462,7 @@ def retrieve_harness():
|
|||
router.model_list = []
|
||||
router.aretrieve_batch = AsyncMock(return_value=make_batch())
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
||||
router.get_batch_deployment_id = MagicMock(return_value=None)
|
||||
router.get_credential_deployment = MagicMock(return_value=None)
|
||||
|
||||
pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock()))
|
||||
|
|
@ -2040,6 +2046,7 @@ def list_harness():
|
|||
router.model_list = []
|
||||
router.alist_batches = AsyncMock(return_value=FakeListPage([]))
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
||||
router.get_batch_deployment_id = MagicMock(return_value=None)
|
||||
|
||||
read_body = AsyncMock(side_effect=lambda request: body_holder["body"])
|
||||
pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock()))
|
||||
|
|
@ -2458,6 +2465,7 @@ def cancel_harness():
|
|||
router.model_list = []
|
||||
router.acancel_batch = AsyncMock(return_value=make_batch())
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
||||
router.get_batch_deployment_id = MagicMock(return_value=None)
|
||||
|
||||
pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock()))
|
||||
# add_litellm_data_to_request is a passthrough that returns the data it got.
|
||||
|
|
@ -3482,3 +3490,166 @@ async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_har
|
|||
factory.assert_not_called()
|
||||
runner.cancel.assert_not_called()
|
||||
cancel_harness.router_acancel.assert_not_called()
|
||||
|
||||
|
||||
def _mixed_azure_group_router() -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": model_group,
|
||||
"litellm_params": {
|
||||
"model": f"azure/{sku}",
|
||||
"api_base": f"https://{sku}.openai.azure.com",
|
||||
"api_key": f"{sku}-key",
|
||||
"api_version": "2025-04-01-preview",
|
||||
},
|
||||
"model_info": {"id": deployment_id, "mode": mode},
|
||||
}
|
||||
for model_group, sku, mode, deployment_id in (
|
||||
("gpt-mixed", "chat-sku", "chat", "chat-id"),
|
||||
("gpt-mixed", "batch-sku", "batch", "batch-id"),
|
||||
("gpt-other", "other-sku", "batch", "other-batch-id"),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _decoded(encoded_id: str) -> tuple[str | None, str, str | None]:
|
||||
return (
|
||||
decode_model_from_file_id(encoded_id),
|
||||
get_original_file_id(encoded_id),
|
||||
decode_deployment_from_file_id(encoded_id),
|
||||
)
|
||||
|
||||
|
||||
def _reached(provider_kwargs: Mapping[str, object]) -> tuple[object, object]:
|
||||
return provider_kwargs["api_base"], provider_kwargs["api_key"]
|
||||
|
||||
|
||||
_BATCH_SKU: Final = ("https://batch-sku.openai.azure.com", "batch-sku-key")
|
||||
_CHAT_SKU: Final = ("https://chat-sku.openai.azure.com", "chat-sku-key")
|
||||
_MIXED_GROUP_HINTS: Final = pytest.mark.parametrize(
|
||||
("deployment_hint", "reached", "kept_hint"),
|
||||
[
|
||||
("batch-id", _BATCH_SKU, "batch-id"),
|
||||
("other-batch-id", _BATCH_SKU, "batch-id"),
|
||||
(None, _CHAT_SKU, None),
|
||||
],
|
||||
ids=["hinted", "hint-naming-another-group", "group-only-id-from-before-the-upgrade"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@_MIXED_GROUP_HINTS
|
||||
async def test_create__mixed_group_input_file_reaches_the_deployment_its_hint_names(
|
||||
harness, deployment_hint: str | None, reached: tuple[str, str], kept_hint: str | None
|
||||
):
|
||||
input_file_id: Final = encode_file_id_with_model("file-original123", "gpt-mixed", deployment_id=deployment_hint)
|
||||
set_body(harness, {"input_file_id": input_file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"})
|
||||
provider_create: Final = AsyncMock(return_value=make_batch(output_file_id="file-out", error_file_id="file-err"))
|
||||
|
||||
with (
|
||||
patch.object(proxy_server, "llm_router", _mixed_azure_group_router()),
|
||||
patch.object(litellm, "acreate_batch", provider_create),
|
||||
):
|
||||
resp: Final = await call_create(harness)
|
||||
|
||||
assert _reached(provider_create.call_args.kwargs) == reached
|
||||
assert provider_create.call_args.kwargs["input_file_id"] == "file-original123"
|
||||
assert tuple(map(_decoded, (resp.id, resp.output_file_id, resp.error_file_id))) == (
|
||||
("gpt-mixed", "batch-provider-id", kept_hint),
|
||||
("gpt-mixed", "file-out", kept_hint),
|
||||
("gpt-mixed", "file-err", kept_hint),
|
||||
)
|
||||
assert resp.input_file_id == input_file_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create__model_header_on_a_mixed_group_reaches_its_batch_deployment(harness):
|
||||
set_body(harness, {"input_file_id": "file-raw", "endpoint": "/v1/chat/completions", "completion_window": "24h"})
|
||||
|
||||
with patch.object(proxy_server, "llm_router", _mixed_azure_group_router()):
|
||||
resp: Final = await call_create(harness, headers={"x-litellm-model": "gpt-mixed"})
|
||||
|
||||
assert _reached(harness.acreate_kwargs()) == _BATCH_SKU
|
||||
assert _decoded(resp.id) == ("gpt-mixed", "batch-provider-id", "batch-id")
|
||||
assert _decoded(resp.input_file_id) == ("gpt-mixed", "file-provider-input", "batch-id")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list__model_header_on_a_mixed_group_reaches_its_batch_deployment(list_harness):
|
||||
provider_list: Final = AsyncMock(return_value=FakeListPage([make_batch(id="batch-1")]))
|
||||
|
||||
with (
|
||||
patch.object(proxy_server, "llm_router", _mixed_azure_group_router()),
|
||||
patch.object(litellm, "alist_batches", provider_list),
|
||||
):
|
||||
resp: Final = await call_list(list_harness, headers={"x-litellm-model": "gpt-mixed"})
|
||||
|
||||
assert _reached(provider_list.call_args.kwargs) == _BATCH_SKU
|
||||
assert _decoded(resp.data[0].id) == ("gpt-mixed", "batch-1", "batch-id")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@_MIXED_GROUP_HINTS
|
||||
async def test_retrieve__mixed_group_batch_reaches_the_deployment_its_hint_names(
|
||||
retrieve_harness, deployment_hint: str | None, reached: tuple[str, str], kept_hint: str | None
|
||||
):
|
||||
batch_id: Final = encode_file_id_with_model(
|
||||
"batch-provider-id", "gpt-mixed", id_type="batch", deployment_id=deployment_hint
|
||||
)
|
||||
provider_retrieve: Final = AsyncMock(return_value=make_batch(output_file_id="file-out"))
|
||||
|
||||
with (
|
||||
patch.object(proxy_server, "llm_router", _mixed_azure_group_router()),
|
||||
patch.object(litellm, "aretrieve_batch", provider_retrieve),
|
||||
):
|
||||
resp: Final = await call_retrieve(retrieve_harness, batch_id)
|
||||
|
||||
assert _reached(provider_retrieve.call_args.kwargs) == reached
|
||||
assert provider_retrieve.call_args.kwargs["litellm_metadata"]["model_info"]["id"] == (kept_hint or "chat-id")
|
||||
assert _decoded(resp.id) == ("gpt-mixed", "batch-provider-id", kept_hint)
|
||||
assert _decoded(resp.output_file_id) == ("gpt-mixed", "file-out", kept_hint)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@_MIXED_GROUP_HINTS
|
||||
async def test_cancel__mixed_group_batch_reaches_the_deployment_its_hint_names(
|
||||
cancel_harness, deployment_hint: str | None, reached: tuple[str, str], kept_hint: str | None
|
||||
):
|
||||
batch_id: Final = encode_file_id_with_model(
|
||||
"batch-provider-id", "gpt-mixed", id_type="batch", deployment_id=deployment_hint
|
||||
)
|
||||
|
||||
with patch.object(proxy_server, "llm_router", _mixed_azure_group_router()):
|
||||
resp: Final = await call_cancel(cancel_harness, batch_id)
|
||||
|
||||
assert _reached(cancel_harness.acancel_kwargs()) == reached
|
||||
assert _decoded(resp.id) == ("gpt-mixed", "batch-provider-id", kept_hint)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve__a_key_granted_only_the_group_may_use_its_hinted_batch_id(retrieve_harness):
|
||||
batch_id: Final = encode_file_id_with_model(
|
||||
"batch-provider-id", "gpt-mixed", id_type="batch", deployment_id="batch-id"
|
||||
)
|
||||
|
||||
with patch.object(proxy_server, "llm_router", _mixed_azure_group_router()):
|
||||
resp: Final = await call_retrieve(retrieve_harness, batch_id, user=_key_restricted_to("gpt-mixed"))
|
||||
|
||||
assert _reached(retrieve_harness.aretrieve_kwargs()) == _BATCH_SKU
|
||||
assert _decoded(resp.id) == ("gpt-mixed", "batch-provider-id", "batch-id")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve__a_hint_naming_a_granted_deployment_never_authorizes_a_group_the_key_lacks(retrieve_harness):
|
||||
batch_id: Final = encode_file_id_with_model(
|
||||
"batch-provider-id", "gpt-mixed", id_type="batch", deployment_id="other-batch-id"
|
||||
)
|
||||
|
||||
with patch.object(proxy_server, "llm_router", _mixed_azure_group_router()):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_retrieve(retrieve_harness, batch_id, user=_key_restricted_to("gpt-other"))
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
retrieve_harness.litellm_aretrieve.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import base64
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -13,7 +14,11 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
apply_team_provider_credentials,
|
||||
apply_unified_file_ids,
|
||||
decode_deployment_from_file_id,
|
||||
decode_model_from_file_id,
|
||||
encode_file_id_with_model,
|
||||
get_credentials_for_model,
|
||||
get_original_file_id,
|
||||
get_team_provider_credentials,
|
||||
is_litellm_executed_batch,
|
||||
map_raw_file_ids_to_unified,
|
||||
|
|
@ -851,3 +856,37 @@ def test_a_team_keeps_its_first_config_file_openai_deployment_over_older_db_ones
|
|||
|
||||
assert credentials is not None
|
||||
assert credentials["api_key"] == "deployment-z-key"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("raw_id", "id_type"), [("file-abc123", "file"), ("batch_abc123", "batch")])
|
||||
def test_a_deployment_hinted_id_decodes_to_its_group_original_id_and_deployment(
|
||||
raw_id: str, id_type: Literal["file", "batch"]
|
||||
):
|
||||
encoded_id: Final = encode_file_id_with_model(raw_id, "gpt-mixed", id_type=id_type, deployment_id="batch-id")
|
||||
|
||||
assert (
|
||||
decode_model_from_file_id(encoded_id),
|
||||
get_original_file_id(encoded_id),
|
||||
decode_deployment_from_file_id(encoded_id),
|
||||
) == ("gpt-mixed", raw_id, "batch-id")
|
||||
|
||||
|
||||
def test_an_unhinted_id_keeps_the_encoding_ids_created_before_the_hint_existed():
|
||||
encoded_before_the_hint: Final = "file-" + base64.urlsafe_b64encode(
|
||||
b"litellm:file-abc123;model,gpt-mixed"
|
||||
).decode().rstrip("=")
|
||||
|
||||
assert encode_file_id_with_model("file-abc123", "gpt-mixed") == encoded_before_the_hint
|
||||
assert decode_deployment_from_file_id(encoded_before_the_hint) is None
|
||||
assert get_original_file_id(encoded_before_the_hint) == "file-abc123"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw_id", ["file-abc123", "batch_abc123", "file-é", ""])
|
||||
def test_a_provider_id_carries_no_deployment_hint(raw_id: str):
|
||||
assert decode_deployment_from_file_id(raw_id) is None
|
||||
assert get_original_file_id(raw_id) == raw_id
|
||||
|
||||
|
||||
@pytest.mark.parametrize("unvalidated_id", [None, 123, ["file-abc123"]])
|
||||
def test_a_non_string_id_from_a_request_body_decodes_to_no_model_or_deployment(unvalidated_id: object):
|
||||
assert (decode_model_from_file_id(unvalidated_id), decode_deployment_from_file_id(unvalidated_id)) == (None, None)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import itertools
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, List, Optional
|
||||
from urllib.parse import quote
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock
|
||||
|
|
@ -18,6 +19,12 @@ from litellm.files.types import FileContentStreamingResult
|
|||
from litellm.proxy._types import LiteLLM_UserTableFiltered, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.hooks import get_proxy_hook
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import ui_view_users
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
decode_deployment_from_file_id,
|
||||
decode_model_from_file_id,
|
||||
encode_file_id_with_model,
|
||||
get_original_file_id,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import (
|
||||
FileContentStreamingHandler,
|
||||
)
|
||||
|
|
@ -984,6 +991,108 @@ def test_batch_upload_keeps_the_callers_drop_params_over_the_deployment(batch_up
|
|||
assert provider_upload.call_args.kwargs["drop_params"] is True
|
||||
|
||||
|
||||
_MIXED_GROUP_SKUS: Final = ("chat-sku", "batch-sku", "other-sku")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mixed_azure_group(monkeypatch):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": model_group,
|
||||
"litellm_params": {
|
||||
"model": f"azure/{sku}",
|
||||
"api_base": f"https://{sku}.openai.azure.com",
|
||||
"api_key": f"{sku}-key",
|
||||
"api_version": "2025-04-01-preview",
|
||||
},
|
||||
"model_info": {"id": deployment_id, "mode": mode},
|
||||
}
|
||||
for model_group, sku, mode, deployment_id in (
|
||||
("gpt-mixed", "chat-sku", "chat", "chat-id"),
|
||||
("gpt-mixed", "batch-sku", "batch", "batch-id"),
|
||||
("gpt-other", "other-sku", "batch", "other-batch-id"),
|
||||
)
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
|
||||
setup_proxy_logging_object(monkeypatch, router)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user"
|
||||
)
|
||||
try:
|
||||
with respx.mock(assert_all_called=False) as azure:
|
||||
yield azure
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
def _azure_upload_routes(azure: respx.MockRouter) -> Mapping[str, respx.Route]:
|
||||
uploaded: Final = {
|
||||
"id": "file-up",
|
||||
"object": "file",
|
||||
"bytes": len(BATCH_JSONL_LINE),
|
||||
"created_at": 0,
|
||||
"filename": "batch.jsonl",
|
||||
"purpose": "batch",
|
||||
"status": "processed",
|
||||
}
|
||||
return {
|
||||
sku: azure.post(host=f"{sku}.openai.azure.com").mock(return_value=httpx.Response(200, json=uploaded))
|
||||
for sku in _MIXED_GROUP_SKUS
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("purpose", "reached_sku", "deployment_hint"),
|
||||
[("batch", "batch-sku", "batch-id"), ("assistants", "chat-sku", None)],
|
||||
ids=["batch-file", "other-purpose"],
|
||||
)
|
||||
def test_upload_naming_a_mixed_azure_group_reaches_its_batch_deployment_only_for_batch_files(
|
||||
mixed_azure_group: respx.MockRouter, purpose: str, reached_sku: str, deployment_hint: str | None
|
||||
):
|
||||
upload_routes: Final = _azure_upload_routes(mixed_azure_group)
|
||||
|
||||
response: Final = _upload_batch_file({"x-litellm-model": "gpt-mixed"}, {"purpose": purpose})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert {sku: route.call_count for sku, route in upload_routes.items()} == {
|
||||
sku: int(sku == reached_sku) for sku in _MIXED_GROUP_SKUS
|
||||
}
|
||||
assert upload_routes[reached_sku].calls.last.request.headers["api-key"] == f"{reached_sku}-key"
|
||||
encoded_id: Final = response.json()["id"]
|
||||
assert (
|
||||
decode_model_from_file_id(encoded_id),
|
||||
get_original_file_id(encoded_id),
|
||||
decode_deployment_from_file_id(encoded_id),
|
||||
) == ("gpt-mixed", "file-up", deployment_hint)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deployment_hint", "reached_sku"),
|
||||
[("batch-id", "batch-sku"), ("other-batch-id", "batch-sku"), (None, "chat-sku")],
|
||||
ids=["hinted", "hint-naming-another-group", "group-only-id-from-before-the-upgrade"],
|
||||
)
|
||||
def test_file_content_of_a_mixed_azure_group_id_comes_from_the_deployment_its_hint_names(
|
||||
mixed_azure_group: respx.MockRouter, deployment_hint: str | None, reached_sku: str
|
||||
):
|
||||
for sku in _MIXED_GROUP_SKUS:
|
||||
mixed_azure_group.get(host=f"{sku}.openai.azure.com", path__regex=r"/files/file-out/content$").mock(
|
||||
return_value=httpx.Response(200, content=f"rows from {sku}".encode())
|
||||
)
|
||||
output_file_id: Final = encode_file_id_with_model("file-out", "gpt-mixed", deployment_id=deployment_hint)
|
||||
|
||||
response: Final = client.get(f"/v1/files/{output_file_id}/content", headers={"Authorization": "Bearer test-key"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.content == f"rows from {reached_sku}".encode()
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="mock respx fails on ci/cd - unclear why")
|
||||
def test_create_file_and_call_chat_completion_e2e(
|
||||
mocker: MockerFixture, monkeypatch, llm_router: Router
|
||||
|
|
|
|||
|
|
@ -14168,6 +14168,71 @@ def test_credentials_for_a_mixed_group_name_keep_resolving_to_its_first_deployme
|
|||
assert credentials is not None and credentials["api_base"] == "https://chat-sku.openai.azure.com"
|
||||
|
||||
|
||||
def test_batch_deployment_of_a_mixed_group_resolves_to_its_batch_deployment_credentials():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
_azure_sku_deployment("gpt-mixed", "chat-sku", "chat", id="chat-id"),
|
||||
_azure_sku_deployment("gpt-mixed", "batch-sku", "batch", id="batch-id"),
|
||||
]
|
||||
)
|
||||
|
||||
deployment_id: Final = router.get_batch_deployment_id("gpt-mixed")
|
||||
credentials: Final = router.get_deployment_credentials_with_provider(model_id=deployment_id or "gpt-mixed")
|
||||
|
||||
assert deployment_id == "batch-id"
|
||||
assert credentials is not None and credentials["api_base"] == "https://batch-sku.openai.azure.com"
|
||||
|
||||
|
||||
def test_batch_deployment_is_none_for_a_group_without_a_batch_deployment():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
_azure_sku_deployment("gpt-mixed", "chat-sku", "chat", id="chat-id"),
|
||||
_azure_sku_deployment("gpt-mixed", "other-chat-sku", "chat", id="other-chat-id"),
|
||||
]
|
||||
)
|
||||
|
||||
assert router.get_batch_deployment_id("gpt-mixed") is None
|
||||
|
||||
|
||||
def test_batch_deployment_is_none_when_the_batch_rule_drops_no_deployment_of_the_group():
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{"model_name": "gpt-mixed", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "dummy-key"}},
|
||||
_azure_sku_deployment("gpt-mixed", "batch-sku", "batch", id="batch-id"),
|
||||
]
|
||||
)
|
||||
|
||||
assert router.get_batch_deployment_id("gpt-mixed") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("preferred_deployment_id", "expected"),
|
||||
[
|
||||
("second-batch-id", "second-batch-id"),
|
||||
(None, "first-batch-id"),
|
||||
("removed-id", "first-batch-id"),
|
||||
("other-group-batch-id", "first-batch-id"),
|
||||
("chat-id", "first-batch-id"),
|
||||
("paused-batch-id", "first-batch-id"),
|
||||
],
|
||||
ids=["still-batch-ready", "no-hint", "removed", "other-group", "chat-deployment", "paused"],
|
||||
)
|
||||
def test_batch_deployment_keeps_a_hinted_deployment_only_while_it_is_batch_ready_in_the_group(
|
||||
preferred_deployment_id: str | None, expected: str
|
||||
):
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
_azure_sku_deployment("gpt-mixed", "chat-sku", "chat", id="chat-id"),
|
||||
_azure_sku_deployment("gpt-mixed", "paused-sku", "batch", id="paused-batch-id", blocked=True),
|
||||
_azure_sku_deployment("gpt-mixed", "first-batch-sku", "batch", id="first-batch-id"),
|
||||
_azure_sku_deployment("gpt-mixed", "second-batch-sku", "batch", id="second-batch-id"),
|
||||
_azure_sku_deployment("gpt-other", "other-batch-sku", "batch", id="other-group-batch-id"),
|
||||
]
|
||||
)
|
||||
|
||||
assert router.get_batch_deployment_id("gpt-mixed", preferred_deployment_id=preferred_deployment_id) == expected
|
||||
|
||||
|
||||
_AZURE_CHAT_SKU_REJECTS_BATCH: Final = httpx.Response(
|
||||
400,
|
||||
json={"error": {"code": "invalid_deployment_type", "message": "The current deployment is not supported for batch jobs"}},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue