diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index f2cef9b80a5..8e247b5228d 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import ORJSONResponse, StreamingResponse +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -46,6 +47,15 @@ if TYPE_CHECKING: router: Final = APIRouter() +class _RAGIngestErrorDetail(TypedDict): + error: ReadOnly[str] + + +def _rag_ingest_bad_request(message: str) -> HTTPException: + detail: Final[_RAGIngestErrorDetail] = {"error": message} + return HTTPException(status_code=400, detail=detail) + + def _raise_vector_store_scan_depth_exceeded() -> None: raise HTTPException( status_code=400, @@ -430,14 +440,13 @@ async def parse_rag_ingest_request( ) provider: Final = vector_store_opts.get("custom_llm_provider", "openai") - if isinstance(provider, str): - try: - get_ingestion_class(provider) - except ValueError as error: - raise HTTPException( - status_code=400, - detail={"error": str(error)}, - ) from error + if not isinstance(provider, str): + raise _rag_ingest_bad_request("custom_llm_provider must be a string") + + try: + get_ingestion_class(provider) + except ValueError as error: + raise _rag_ingest_bad_request(str(error)) from error return ingest_options, secured_file_data, file_url, file_id diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 14473fecf00..ee81aae8b78 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -7,12 +7,11 @@ Covers: import io from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest from fastapi.testclient import TestClient - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app @@ -77,6 +76,33 @@ def test_rag_ingest_rejects_unsupported_provider_before_execution( mock_aingest.assert_not_awaited() +def test_rag_ingest_rejects_non_string_provider_before_execution( + client_internal_user: TestClient, +) -> None: + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + ) as mock_aingest, + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + response: Final = client_internal_user.post( + "/v1/rag/ingest", + json={ + "file_id": "file-test", + "ingest_options": { + "vector_store": {"custom_llm_provider": {"provider": "milvus"}} + }, + }, + ) + + assert response.status_code == 400 + assert response.json() == { + "detail": {"error": "custom_llm_provider must be a string"} + } + mock_aingest.assert_not_awaited() + + def test_internal_user_viewer_rag_ingest_without_vector_store_id_rejected( client_internal_user_viewer, ):