mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(proxy): guard managed file id parsing for non-output unified formats
Only treat decoded unified ids as already-wrapped output files when they contain llm_output_file_id. Skip other litellm_proxy id shapes instead of IndexError on split. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
6748607d91
commit
a7f0a5baa5
2 changed files with 50 additions and 3 deletions
|
|
@ -504,7 +504,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if retrieve_file_id
|
||||
else False
|
||||
)
|
||||
if potential_file_id:
|
||||
if potential_file_id and "llm_output_file_id," in potential_file_id:
|
||||
model_id = self.get_model_id_from_unified_file_id(potential_file_id)
|
||||
if model_id:
|
||||
data["model"] = model_id
|
||||
|
|
@ -1058,7 +1058,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
return file_id.split("llm_output_file_model_id,")[1].split(";")[0]
|
||||
|
||||
def get_output_file_id_from_unified_file_id(self, file_id: str) -> str:
|
||||
return file_id.split("llm_output_file_id,")[1].split(";")[0]
|
||||
marker = "llm_output_file_id,"
|
||||
if marker not in file_id:
|
||||
raise ValueError(
|
||||
f"Unified id does not contain {marker!r}: {file_id[:80]!r}"
|
||||
)
|
||||
return file_id.split(marker, 1)[1].split(";")[0]
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
|
||||
|
|
@ -1102,13 +1107,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
decoded_unified_file_id = _is_base64_encoded_unified_file_id(
|
||||
file_id_value
|
||||
)
|
||||
if decoded_unified_file_id:
|
||||
if (
|
||||
decoded_unified_file_id
|
||||
and "llm_output_file_id," in decoded_unified_file_id
|
||||
):
|
||||
provider_file_id = (
|
||||
self.get_output_file_id_from_unified_file_id(
|
||||
decoded_unified_file_id
|
||||
)
|
||||
)
|
||||
unified_file_id = file_id_value
|
||||
elif decoded_unified_file_id:
|
||||
verbose_logger.warning(
|
||||
f"Skipping {file_attr}={file_id_value!r}: "
|
||||
"unified id is not a managed file output id"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
provider_file_id = file_id_value
|
||||
unified_file_id = self.get_unified_output_file_id(
|
||||
|
|
|
|||
|
|
@ -222,3 +222,36 @@ async def test_should_not_double_wrap_already_unified_output_file_id():
|
|||
already_unified + "=" * (-len(already_unified) % 4)
|
||||
).decode()
|
||||
assert decoded.count(f"llm_output_file_id,{provider_file_id}") == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_non_file_unified_id_on_output_file_id():
|
||||
"""Batch-style unified ids lack llm_output_file_id; must not IndexError or re-wrap."""
|
||||
import base64
|
||||
|
||||
managed_files = _make_managed_files_instance()
|
||||
batch_unified = (
|
||||
base64.urlsafe_b64encode(
|
||||
b"litellm_proxy;model_id:openai/openai/gpt-5.5-batch;llm_batch_id:batch_abc"
|
||||
)
|
||||
.decode()
|
||||
.rstrip("=")
|
||||
)
|
||||
|
||||
batch_response = _make_batch_response(
|
||||
model_id="openai/openai/gpt-5.5-batch",
|
||||
model_name="openai/openai/gpt-5.5-batch",
|
||||
output_file_id=batch_unified,
|
||||
)
|
||||
user_api_key_dict = _make_user_api_key_dict()
|
||||
|
||||
with patch("litellm.afile_retrieve", AsyncMock()) as mock_afile_retrieve:
|
||||
await managed_files.async_post_call_success_hook(
|
||||
data={},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=batch_response,
|
||||
)
|
||||
|
||||
assert batch_response.output_file_id == batch_unified
|
||||
mock_afile_retrieve.assert_not_called()
|
||||
managed_files.store_unified_file_id.assert_not_awaited()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue