ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(proxy): forward decoded container ids

Browse Source
This commit is contained in:
user 2026-04-30 19:13:20 -07:00
parent 9376b30bca
commit 6aac4552f9
4 changed files with 121 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
litellm/proxy/container_endpoints/endpoints.py
View File

@ -17,6 +17,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
from litellm.proxy.container_endpoints.ownership import (
assert_user_can_access_container,
filter_container_list_response,
get_container_forwarding_params,
record_container_owner,
)
@ -295,15 +296,18 @@ async def retrieve_container(
)
# Add custom_llm_provider to data
container_access = await assert_user_can_access_container(
original_container_id, custom_llm_provider = await assert_user_can_access_container(
container_id=container_id,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = container_access[1]
# Keep the managed id in request data so downstream container utilities can
# preserve encoded routing metadata while decoding before the provider call.
data["custom_llm_provider"] = custom_llm_provider
data.update(
get_container_forwarding_params(
container_id,
original_container_id,
custom_llm_provider,
)
)
# Process request using ProxyBaseLLMRequestProcessing
processor = ProxyBaseLLMRequestProcessing(data=data)
@ -397,15 +401,18 @@ async def delete_container(
)
# Add custom_llm_provider to data
container_access = await assert_user_can_access_container(
original_container_id, custom_llm_provider = await assert_user_can_access_container(
container_id=container_id,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = container_access[1]
# Keep the managed id in request data so downstream container utilities can
# preserve encoded routing metadata while decoding before the provider call.
data["custom_llm_provider"] = custom_llm_provider
data.update(
get_container_forwarding_params(
container_id,
original_container_id,
custom_llm_provider,
)
)
# Process request using ProxyBaseLLMRequestProcessing
processor = ProxyBaseLLMRequestProcessing(data=data)

46
litellm/proxy/container_endpoints/handler_factory.py
View File

@ -21,6 +21,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
)
from litellm.proxy.container_endpoints.ownership import (
assert_user_can_access_container,
get_container_forwarding_params,
)
@ -183,6 +184,13 @@ async def _process_binary_request(
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
forwarding_params = get_container_forwarding_params(
container_id,
original_container_id,
custom_llm_provider,
)
if "model_id" in forwarding_params:
litellm_params["model_id"] = forwarding_params["model_id"]
# Get the provider config
container_provider_config = _get_container_provider_config(custom_llm_provider)
@ -283,18 +291,19 @@ async def _process_multipart_upload_request(
or "openai"
)
container_access = await assert_user_can_access_container(
original_container_id, custom_llm_provider = await assert_user_can_access_container(
container_id=container_id,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = container_access[1]
# Keep the managed container id in the forwarded request. The container API
# layer decodes it before the upstream provider call and uses embedded
# routing metadata to preserve model/deployment affinity.
data["container_id"] = container_id
data["custom_llm_provider"] = custom_llm_provider
data.update(
get_container_forwarding_params(
container_id,
original_container_id,
custom_llm_provider,
)
)
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
@ -361,15 +370,22 @@ async def _process_request(
# Validate container_id ownership if present in path_params.
if "container_id" in path_params:
container_access = await assert_user_can_access_container(
container_id=path_params["container_id"],
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
original_container_id, custom_llm_provider = (
await assert_user_can_access_container(
container_id=path_params["container_id"],
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
)
custom_llm_provider = container_access[1]
# Preserve the managed id for downstream container decoding/routing.
data["custom_llm_provider"] = custom_llm_provider
data.update(
get_container_forwarding_params(
path_params["container_id"],
original_container_id,
custom_llm_provider,
)
)
else:
data["custom_llm_provider"] = custom_llm_provider
processor = ProxyBaseLLMRequestProcessing(data=data)
try:

35
litellm/proxy/container_endpoints/ownership.py
View File

@ -45,6 +45,22 @@ def decode_container_id_for_ownership(
return original_container_id, custom_llm_provider
def get_container_forwarding_params(
container_id: str,
original_container_id: str,
custom_llm_provider: str,
) -> Dict[str, str]:
params = {
"container_id": original_container_id,
"custom_llm_provider": custom_llm_provider,
}
decoded = ResponsesAPIRequestUtils._decode_container_id(container_id)
model_id = decoded.get("model_id")
if isinstance(model_id, str) and model_id:
params["model_id"] = model_id
return params
def _get_response_id(response: Any) -> Optional[str]:
if response is None:
return None
@ -139,11 +155,12 @@ async def record_container_owner(
raise
except Exception as e:
verbose_proxy_logger.warning(
"Failed to record container ownership for container_id=%s: %s",
"Failed to persist container ownership for container_id=%s; "
"falling back to in-process tracking: %s",
model_object_id,
e,
)
raise HTTPException(status_code=500, detail="Unable to track container")
_IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner
return response
@ -210,7 +227,9 @@ def _get_container_list_data(response: Any) -> Optional[List[Any]]:
return data if isinstance(data, list) else None
def _set_container_list_data(response: Any, data: List[Any]) -> Any:
def _set_container_list_data(
response: Any, data: List[Any], removed_filtered_items: bool = False
) -> Any:
if isinstance(response, dict):
response["data"] = data
if data:
@ -220,6 +239,8 @@ def _set_container_list_data(response: Any, data: List[Any]) -> Any:
response["first_id"] = None
response["last_id"] = None
response["has_more"] = False
if removed_filtered_items:
response["has_more"] = False
return response
response.data = data
@ -227,6 +248,8 @@ def _set_container_list_data(response: Any, data: List[Any]) -> Any:
response.last_id = _get_response_id(data[-1]) if data else None
if not data and hasattr(response, "has_more"):
response.has_more = False
if removed_filtered_items and hasattr(response, "has_more"):
response.has_more = False
return response
@ -291,4 +314,8 @@ async def filter_container_list_response(
):
filtered.append(item)
return _set_container_list_data(response, filtered)
return _set_container_list_data(
response,
filtered,
removed_filtered_items=len(filtered) != len(data),
)

51
tests/test_litellm/containers/test_container_proxy_ownership.py
View File

@ -82,6 +82,34 @@ async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch):
assert data["updated_by"] == "team:team-1"
@pytest.mark.asyncio
async def test_should_fallback_to_memory_when_persistent_owner_record_fails(
monkeypatch,
):
table = AsyncMock()
table.find_unique.side_effect = Exception("db unavailable")
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=prisma_client),
)
auth = UserAPIKeyAuth(user_id="user-1")
await ownership.record_container_owner(
response=_container("cntr_provider"),
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert (
ownership._IN_MEMORY_CONTAINER_OWNERS["container:openai:cntr_provider"]
== "user-1"
)
@pytest.mark.asyncio
async def test_should_track_container_owner_in_memory_without_prisma(monkeypatch):
monkeypatch.setattr(
@ -216,7 +244,7 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch):
response = ContainerListResponse(
object="list",
data=[_container("cntr_owned"), _container("cntr_other")],
has_more=False,
has_more=True,
)
filtered = await ownership.filter_container_list_response(
@ -228,6 +256,7 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch):
assert [item.id for item in filtered.data] == ["cntr_owned"]
assert filtered.first_id == "cntr_owned"
assert filtered.last_id == "cntr_owned"
assert filtered.has_more is False
where = table.find_many.await_args.kwargs["where"]
assert where["file_purpose"] == ownership.CONTAINER_OBJECT_PURPOSE
assert where["created_by"]["in"] == ["user-1", "user:user-1"]
@ -334,7 +363,7 @@ async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch
@pytest.mark.asyncio
async def test_should_preserve_managed_container_id_for_proxy_forwarding(monkeypatch):
async def test_should_forward_decoded_container_id_for_proxy_forwarding(monkeypatch):
from litellm.proxy.container_endpoints import handler_factory
proxy_server_stub = SimpleNamespace(
@ -388,12 +417,13 @@ async def test_should_preserve_managed_container_id_for_proxy_forwarding(monkeyp
path_params={"container_id": encoded_id},
)
assert result["container_id"] == encoded_id
assert result["container_id"] == "cntr_provider"
assert result["custom_llm_provider"] == "azure"
assert result["model_id"] == "router-gpt"
@pytest.mark.asyncio
async def test_should_preserve_managed_container_id_for_multipart_upload(monkeypatch):
async def test_should_forward_decoded_container_id_for_multipart_upload(monkeypatch):
from litellm.proxy.common_utils import http_parsing_utils
from litellm.proxy.container_endpoints import handler_factory
@ -458,13 +488,14 @@ async def test_should_preserve_managed_container_id_for_multipart_upload(monkeyp
container_id=encoded_id,
)
assert result["container_id"] == encoded_id
assert result["container_id"] == "cntr_provider"
assert result["custom_llm_provider"] == "azure"
assert result["model_id"] == "router-gpt"
assert result["file"] == "file-data"
@pytest.mark.asyncio
async def test_should_preserve_managed_container_id_for_proxy_retrieve(monkeypatch):
async def test_should_forward_decoded_container_id_for_proxy_retrieve(monkeypatch):
from litellm.proxy.container_endpoints import endpoints
proxy_server_stub = SimpleNamespace(
@ -513,12 +544,13 @@ async def test_should_preserve_managed_container_id_for_proxy_retrieve(monkeypat
user_api_key_dict=UserAPIKeyAuth(user_id="user-1"),
)
assert result["container_id"] == encoded_id
assert result["container_id"] == "cntr_provider"
assert result["custom_llm_provider"] == "azure"
assert result["model_id"] == "router-gpt"
@pytest.mark.asyncio
async def test_should_preserve_managed_container_id_for_proxy_delete(monkeypatch):
async def test_should_forward_decoded_container_id_for_proxy_delete(monkeypatch):
from litellm.proxy.container_endpoints import endpoints
proxy_server_stub = SimpleNamespace(
@ -567,5 +599,6 @@ async def test_should_preserve_managed_container_id_for_proxy_delete(monkeypatch
user_api_key_dict=UserAPIKeyAuth(user_id="user-1"),
)
assert result["container_id"] == encoded_id
assert result["container_id"] == "cntr_provider"
assert result["custom_llm_provider"] == "azure"
assert result["model_id"] == "router-gpt"