mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-24 00:28:32 +00:00
[Fix] Populate _hidden_params.model_id in batch terminal-state shortcut path
The terminal-state DB shortcut in retrieve_batch returned a LiteLLMBatch with empty _hidden_params, causing the managed_files hook to skip encoding output_file_id into a unified ID. This adds the same model_id extraction from unified_batch_id that the non-terminal path already has. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
5b755db329
commit
ca2033036f
@@ -412,6 +412,16 @@ async def retrieve_batch( # noqa: PLR0915
|
||||
"cancelled",
|
||||
"expired",
|
||||
]:
|
||||
# Populate _hidden_params.model_id so managed_files hook can
|
||||
# encode output_file_id / error_file_id into unified IDs.
|
||||
if unified_batch_id:
|
||||
response._hidden_params["unified_batch_id"] = unified_batch_id
|
||||
model_id_from_batch = get_model_id_from_unified_batch_id(
|
||||
unified_batch_id
|
||||
)
|
||||
if model_id_from_batch:
|
||||
response._hidden_params["model_id"] = model_id_from_batch
|
||||
|
||||
# Call hooks and return
|
||||
response = await proxy_logging_obj.post_call_success_hook(
|
||||
data=data, user_api_key_dict=user_api_key_dict, response=response
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
"""
|
||||
Test that the terminal-state shortcut path in retrieve_batch populates
|
||||
_hidden_params.model_id from the unified batch ID.
|
||||
|
||||
Regression: When a batch is in a terminal state (completed/failed/cancelled/expired),
|
||||
the DB shortcut path returned a LiteLLMBatch with empty _hidden_params, causing the
|
||||
managed_files hook to skip encoding output_file_id into a unified ID.
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
get_model_id_from_unified_batch_id,
|
||||
)
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
|
||||
MODEL_ID = "model-xyz-123"
|
||||
RAW_BATCH_ID = "batch-provider-456"
|
||||
DECODED_UNIFIED_BATCH_ID = (
|
||||
f"litellm_proxy;model_id:{MODEL_ID};llm_batch_id:{RAW_BATCH_ID}"
|
||||
)
|
||||
B64_UNIFIED_BATCH_ID = (
|
||||
base64.urlsafe_b64encode(DECODED_UNIFIED_BATCH_ID.encode()).decode().rstrip("=")
|
||||
)
|
||||
|
||||
|
||||
def test_get_model_id_from_unified_batch_id():
|
||||
"""model_id is correctly extracted from a unified batch ID."""
|
||||
assert get_model_id_from_unified_batch_id(DECODED_UNIFIED_BATCH_ID) == MODEL_ID
|
||||
|
||||
|
||||
def test_get_model_id_returns_none_for_invalid():
|
||||
"""Returns None for non-unified IDs."""
|
||||
assert get_model_id_from_unified_batch_id("plain-batch-id") is None
|
||||
assert get_model_id_from_unified_batch_id("") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", ["completed", "failed", "cancelled", "expired"])
|
||||
def test_terminal_batch_hidden_params_population(status):
|
||||
"""
|
||||
Simulate the terminal-state shortcut path logic: when a batch from the DB
|
||||
has a terminal status and we have a unified_batch_id, _hidden_params should
|
||||
get model_id and unified_batch_id set.
|
||||
|
||||
This mirrors the code added to endpoints.py lines 415-423.
|
||||
"""
|
||||
# Create a batch as it would come from the database (empty _hidden_params)
|
||||
response = LiteLLMBatch(
|
||||
id=B64_UNIFIED_BATCH_ID,
|
||||
completion_window="24h",
|
||||
created_at=1700000000,
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="file-input-abc",
|
||||
object="batch",
|
||||
status=status,
|
||||
output_file_id="file-output-raw",
|
||||
)
|
||||
|
||||
assert response._hidden_params.get("model_id") is None, "precondition: no model_id"
|
||||
|
||||
unified_batch_id = DECODED_UNIFIED_BATCH_ID
|
||||
|
||||
# This is the exact logic from the terminal-state shortcut path
|
||||
if unified_batch_id:
|
||||
response._hidden_params["unified_batch_id"] = unified_batch_id
|
||||
model_id_from_batch = get_model_id_from_unified_batch_id(unified_batch_id)
|
||||
if model_id_from_batch:
|
||||
response._hidden_params["model_id"] = model_id_from_batch
|
||||
|
||||
assert response._hidden_params["model_id"] == MODEL_ID
|
||||
assert response._hidden_params["unified_batch_id"] == DECODED_UNIFIED_BATCH_ID
|
||||
|
||||
|
||||
def test_terminal_batch_no_unified_id_leaves_hidden_params_empty():
|
||||
"""
|
||||
When there is no unified_batch_id (non-managed batch), _hidden_params
|
||||
should remain unchanged.
|
||||
"""
|
||||
response = LiteLLMBatch(
|
||||
id="batch-plain-id",
|
||||
completion_window="24h",
|
||||
created_at=1700000000,
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="file-input-abc",
|
||||
object="batch",
|
||||
status="completed",
|
||||
)
|
||||
|
||||
unified_batch_id = None
|
||||
|
||||
if unified_batch_id:
|
||||
response._hidden_params["unified_batch_id"] = unified_batch_id
|
||||
model_id_from_batch = get_model_id_from_unified_batch_id(unified_batch_id)
|
||||
if model_id_from_batch:
|
||||
response._hidden_params["model_id"] = model_id_from_batch
|
||||
|
||||
assert response._hidden_params.get("model_id") is None
|
||||
assert response._hidden_params.get("unified_batch_id") is None
|
||||
Reference in New Issue
Block a user