diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index fc1f77bb684..e0b7e4bddcc 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -21,6 +21,9 @@ from litellm.proxy.container_endpoints.ownership import ( get_container_forwarding_params, record_container_owner, ) +from litellm.proxy.container_endpoints.pagination import ( + parse_container_list_query_params, +) router = APIRouter() @@ -206,9 +209,7 @@ async def list_containers( version, ) - # Read query parameters - query_params = dict(request.query_params) - data: Dict[str, Any] = {"query_params": query_params} + data: Dict[str, Any] = dict(parse_container_list_query_params(request.query_params)) # Extract custom_llm_provider using priority chain custom_llm_provider = ( diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index bc871479356..d8398de59fa 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -5,12 +5,12 @@ This module reads the endpoints.json config and dynamically creates FastAPI route handlers for ALL container file endpoints. """ -import json from pathlib import Path -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import ORJSONResponse +from pydantic import BaseModel from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -23,19 +23,37 @@ from litellm.proxy.container_endpoints.ownership import ( assert_user_can_access_container, get_container_forwarding_params, ) +from litellm.proxy.container_endpoints.pagination import ( + parse_container_list_query_params, +) -def _load_endpoints_config() -> Dict: +class ContainerEndpointConfig(BaseModel): + name: str + async_name: str + path: str + method: str + response_type: str + path_params: Tuple[str, ...] = () + query_params: Tuple[str, ...] = () + returns_binary: bool = False + is_multipart: bool = False + + +class ContainerEndpointsConfig(BaseModel): + endpoints: Tuple[ContainerEndpointConfig, ...] = () + + +def _load_endpoint_configs() -> Tuple[ContainerEndpointConfig, ...]: """Load the endpoints configuration from JSON file.""" config_path = Path(__file__).parent.parent.parent / "containers" / "endpoints.json" with open(config_path) as f: - return json.load(f) + return ContainerEndpointsConfig.model_validate_json(f.read()).endpoints def get_all_route_types() -> List[str]: """Get all async route types for registration in route_llm_request.py""" - config = _load_endpoints_config() - return [endpoint["async_name"] for endpoint in config["endpoints"]] + return [endpoint.async_name for endpoint in _load_endpoint_configs()] def _get_container_provider_config(custom_llm_provider: str): @@ -52,16 +70,17 @@ def _get_container_provider_config(custom_llm_provider: str): def _create_handler_for_path_params( - path_params: List[str], + path_params: Tuple[str, ...], route_type: str, returns_binary: bool = False, is_multipart: bool = False, + query_params: Tuple[str, ...] = (), ): """ Dynamically create a handler with the correct path parameter signature. """ # For binary content endpoints, use a different handler - if returns_binary and path_params == ["container_id", "file_id"]: + if returns_binary and path_params == ("container_id", "file_id"): async def handler_binary_content( request: Request, @@ -100,7 +119,7 @@ def _create_handler_for_path_params( return handler_multipart_upload # Create handlers for different path parameter combinations - if path_params == ["container_id"]: + if path_params == ("container_id",): async def handler_container_id( request: Request, @@ -114,11 +133,12 @@ def _create_handler_for_path_params( user_api_key_dict=user_api_key_dict, route_type=route_type, path_params={"container_id": container_id}, + query_param_names=query_params, ) return handler_container_id - elif path_params == ["container_id", "file_id"]: + elif path_params == ("container_id", "file_id"): async def handler_container_file( request: Request, @@ -133,6 +153,7 @@ def _create_handler_for_path_params( user_api_key_dict=user_api_key_dict, route_type=route_type, path_params={"container_id": container_id, "file_id": file_id}, + query_param_names=query_params, ) return handler_container_file @@ -150,6 +171,7 @@ def _create_handler_for_path_params( user_api_key_dict=user_api_key_dict, route_type=route_type, path_params={}, + query_param_names=query_params, ) return handler_no_params @@ -357,6 +379,7 @@ async def _process_request( user_api_key_dict: UserAPIKeyAuth, route_type: str, path_params: Dict[str, str], + query_param_names: Tuple[str, ...] = (), ): """Common request processing logic.""" from litellm.proxy.proxy_server import ( @@ -373,9 +396,8 @@ async def _process_request( version, ) - query_params = dict(request.query_params) data: Dict[str, Any] = { - "query_params": query_params, + **parse_container_list_query_params(request.query_params, supported_params=query_param_names), **path_params, } @@ -441,21 +463,21 @@ def register_container_file_endpoints(router: APIRouter) -> None: This single function registers all endpoints defined in endpoints.json, eliminating the need for manual endpoint definitions. """ - config = _load_endpoints_config() - - for endpoint_config in config["endpoints"]: - path = endpoint_config["path"] - method = endpoint_config["method"].lower() - path_params = endpoint_config.get("path_params", []) - route_type = endpoint_config["async_name"] - returns_binary = endpoint_config.get("returns_binary", False) - is_multipart = endpoint_config.get("is_multipart", False) + for endpoint_config in _load_endpoint_configs(): + path = endpoint_config.path + returns_binary = endpoint_config.returns_binary # Create handler with correct signature for path params - handler = _create_handler_for_path_params(path_params, route_type, returns_binary, is_multipart) + handler = _create_handler_for_path_params( + endpoint_config.path_params, + endpoint_config.async_name, + returns_binary, + endpoint_config.is_multipart, + query_params=endpoint_config.query_params, + ) # Register routes - route_method = getattr(router, method) + route_method = getattr(router, endpoint_config.method.lower()) # For binary endpoints, don't use ORJSONResponse if returns_binary: diff --git a/litellm/proxy/container_endpoints/pagination.py b/litellm/proxy/container_endpoints/pagination.py new file mode 100644 index 00000000000..76362eb297b --- /dev/null +++ b/litellm/proxy/container_endpoints/pagination.py @@ -0,0 +1,35 @@ +from typing import Dict, Literal, Mapping, Optional, Tuple, Union + +from fastapi import HTTPException +from pydantic import BaseModel, Field, ValidationError + +CONTAINER_LIST_QUERY_PARAMS: Tuple[str, ...] = ("after", "limit", "order") + + +class ContainerListPaginationParams(BaseModel): + after: Optional[str] = None + limit: Optional[int] = Field(default=None, ge=1) + order: Optional[Literal["asc", "desc"]] = None + + +def parse_container_list_query_params( + query_params: Mapping[str, str], + supported_params: Tuple[str, ...] = CONTAINER_LIST_QUERY_PARAMS, +) -> Dict[str, Union[str, int]]: + """ + Validate the pagination query params of a container list route and return them as + top-level SDK arguments (`after`, `limit`, `order`), which is what the container + SDK functions and the provider APIs expect. + """ + requested = {key: value for key, value in query_params.items() if key in supported_params and value != ""} + try: + parsed = ContainerListPaginationParams.model_validate(requested) + except ValidationError as e: + violations = "; ".join( + f"{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors() + ) + raise HTTPException( + status_code=400, + detail={"error": f"Invalid container list query parameters: {violations}"}, + ) + return parsed.model_dump(exclude_none=True) diff --git a/tests/test_litellm/containers/test_container_proxy_pagination.py b/tests/test_litellm/containers/test_container_proxy_pagination.py new file mode 100644 index 00000000000..e629100cfd9 --- /dev/null +++ b/tests/test_litellm/containers/test_container_proxy_pagination.py @@ -0,0 +1,124 @@ +from typing import Any, Dict, List + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.container_endpoints.pagination import ( + parse_container_list_query_params, +) + + +@pytest.fixture +def captured_data(monkeypatch) -> List[Dict[str, Any]]: + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + captured: List[Dict[str, Any]] = [] + + async def fake_base_process_llm_request(self, **kwargs): + captured.append(self.data) + return {"object": "list", "data": []} + + monkeypatch.setattr( + ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + fake_base_process_llm_request, + ) + return captured + + +@pytest.fixture +def client(monkeypatch) -> TestClient: + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.container_endpoints import endpoints, ownership + + async def fake_auth(): + return UserAPIKeyAuth(user_id="user-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + async def fake_assert_user_can_access_container(container_id, user_api_key_dict, custom_llm_provider): + return container_id, custom_llm_provider + + async def fake_get_container_forwarding_params(container_id, original_container_id, custom_llm_provider): + return {"container_id": container_id, "custom_llm_provider": custom_llm_provider} + + async def fake_filter_container_list_response(response, user_api_key_dict, custom_llm_provider): + return response + + for module in (endpoints, ownership): + monkeypatch.setattr( + module, + "assert_user_can_access_container", + fake_assert_user_can_access_container, + raising=False, + ) + monkeypatch.setattr( + module, + "get_container_forwarding_params", + fake_get_container_forwarding_params, + raising=False, + ) + monkeypatch.setattr( + endpoints, + "filter_container_list_response", + fake_filter_container_list_response, + ) + + from litellm.proxy.container_endpoints import handler_factory + + monkeypatch.setattr( + handler_factory, + "assert_user_can_access_container", + fake_assert_user_can_access_container, + ) + monkeypatch.setattr( + handler_factory, + "get_container_forwarding_params", + fake_get_container_forwarding_params, + ) + + app = FastAPI() + app.include_router(endpoints.router) + app.dependency_overrides[user_api_key_auth] = fake_auth + return TestClient(app) + + +def test_should_forward_container_list_pagination_as_top_level_params(client, captured_data): + response = client.get("/v1/containers?after=cntr_cursor&limit=20&order=desc") + + assert response.status_code == 200 + assert captured_data[0]["after"] == "cntr_cursor" + assert captured_data[0]["limit"] == 20 + assert captured_data[0]["order"] == "desc" + assert "query_params" not in captured_data[0] + + +def test_should_forward_container_file_list_pagination_as_top_level_params(client, captured_data): + response = client.get("/v1/containers/cntr_123/files?after=cfile_cursor&limit=5&order=asc") + + assert response.status_code == 200 + assert captured_data[0]["after"] == "cfile_cursor" + assert captured_data[0]["limit"] == 5 + assert captured_data[0]["order"] == "asc" + assert captured_data[0]["container_id"] == "cntr_123" + assert "query_params" not in captured_data[0] + + +def test_should_not_forward_pagination_params_to_non_list_container_routes(client, captured_data): + response = client.get("/v1/containers/cntr_123/files/cfile_1?limit=5") + + assert response.status_code == 200 + assert "limit" not in captured_data[0] + + +def test_should_reject_invalid_container_list_pagination_params(client, captured_data): + assert client.get("/v1/containers?order=sideways").status_code == 400 + assert client.get("/v1/containers?limit=abc").status_code == 400 + assert client.get("/v1/containers?limit=0").status_code == 400 + assert captured_data == [] + + +def test_should_ignore_non_pagination_query_params(): + parsed = parse_container_list_query_params({"custom_llm_provider": "azure", "limit": "3"}) + + assert parsed == {"limit": 3}