scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

refactor(proxy): extract encode_batch_response_ids helper and fix list_batches encoding

Browse source 
Extract duplicated batch ID encoding logic into a shared helper
encode_batch_response_ids() in common_utils.py. Use it in create_batch,
retrieve_batch, and cancel_batch. Also add encoding to list_batches
when x-litellm-model is used.
This commit is contained in:
Chesars 2026-03-03 11:38:50 -03:00
parent 7506fd0426
commit 59bde4a81a
2 changed files with 27 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

77
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -23,6 +23,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
decode_model_from_file_id,
encode_batch_response_ids,
encode_file_id_with_model,
get_batch_from_database,
get_credentials_for_model,
@ -243,29 +244,7 @@ async def create_batch( # noqa: PLR0915
**_create_batch_data # type: ignore
)
# Encode response IDs with model info so retrieve_batch
# can route back to the correct provider/credentials.
if response and hasattr(response, "id") and response.id:
response.id = encode_file_id_with_model(
file_id=response.id,
model=model_param,
id_type="batch",
)
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_param
)
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_param
)
if hasattr(response, "input_file_id") and response.input_file_id:
response.input_file_id = encode_file_id_with_model(
file_id=response.input_file_id, model=model_param
)
encode_batch_response_ids(response, model=model_param)
verbose_proxy_logger.debug(f"Created batch using model: {model_param}")
else:
@ -465,28 +444,7 @@ async def retrieve_batch( # noqa: PLR0915
**data # type: ignore
)
# Re-encode response IDs so the client always sees encoded IDs.
# The provider returns raw IDs (e.g. output_file_id, error_file_id)
# which the client needs encoded to route future file downloads.
if response and hasattr(response, "id") and response.id:
response.id = encode_file_id_with_model(
file_id=response.id, model=model_from_id, id_type="batch",
)
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_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_id
)
if hasattr(response, "input_file_id") and response.input_file_id:
response.input_file_id = encode_file_id_with_model(
file_id=response.input_file_id, model=model_from_id
)
encode_batch_response_ids(response, model=model_from_id)
verbose_proxy_logger.debug(
f"Retrieved batch using model: {model_from_id}, original_id: {original_batch_id}"
@ -679,7 +637,13 @@ async def list_batches(
limit=limit,
**data # type: ignore
)
# Encode batch IDs in the list response so clients can use
# them for retrieve/cancel/file downloads through the proxy.
if response and hasattr(response, "data") and response.data:
for batch in response.data:
encode_batch_response_ids(batch, model=model_param)
verbose_proxy_logger.debug(f"Listed batches using model: {model_param}")
# SCENARIO 2 (alternative): target_model_names based routing
@ -856,26 +820,7 @@ async def cancel_batch(
**data # type: ignore
)
# Re-encode response IDs so the client always sees encoded IDs.
if response and hasattr(response, "id") and response.id:
response.id = encode_file_id_with_model(
file_id=response.id, model=model_from_id, id_type="batch",
)
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_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_id
)
if hasattr(response, "input_file_id") and response.input_file_id:
response.input_file_id = encode_file_id_with_model(
file_id=response.input_file_id, model=model_from_id
)
encode_batch_response_ids(response, model=model_from_id)
verbose_proxy_logger.debug(
f"Cancelled batch using model: {model_from_id}, original_id: {original_batch_id}"

16
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -129,6 +129,22 @@ def encode_file_id_with_model(
return f"{prefix}{encoded_b64}"
def encode_batch_response_ids(response, model: str) -> 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"
)
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),
)
def decode_model_from_file_id(encoded_id: str) -> Optional[str]:
"""
Extract model name from an encoded file/batch ID.