From 9ceaa2cbb0989ba4213baeca7ae03552240cd7a3 Mon Sep 17 00:00:00 2001 From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Date: Wed, 18 Feb 2026 16:58:38 -0800 Subject: [PATCH] tests and route permissions (#21508) --- litellm/proxy/auth/route_checks.py | 3 + litellm/proxy/rag_endpoints/endpoints.py | 15 +- tests/proxy_unit_tests/test_proxy_routes.py | 5 + .../proxy/auth/test_route_checks.py | 40 ++++++ .../proxy/auth/test_user_api_key_auth.py | 10 ++ .../proxy/rag_endpoints/__init__.py | 0 .../proxy/rag_endpoints/test_rag_endpoints.py | 130 ++++++++++++++++++ 7 files changed, 202 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/rag_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index fa63ff01b3..a00401008f 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -307,6 +307,9 @@ class RouteChecks: ): return True + if route in LiteLLMRoutes.litellm_native_routes.value: + return True + # fuzzy match routes like "/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" # Check for routes with placeholders or wildcard patterns for openai_route in LiteLLMRoutes.openai_routes.value: diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 39909df8e9..408a017967 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -10,7 +10,7 @@ import base64 from typing import Any, Dict, Optional, Tuple import orjson -from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import ORJSONResponse import litellm @@ -357,6 +357,19 @@ async def rag_ingest( # Parse request ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(request) + # INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores only + if ( + user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + and not ingest_options.get("vector_store", {}).get("vector_store_id") + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "internal_user_viewer role can only ingest files to an existing vector store. " + "Provide 'vector_store_id' in ingest_options.vector_store." + }, + ) + # Add litellm data request_data: Dict[str, Any] = {} request_data = await add_litellm_data_to_request( diff --git a/tests/proxy_unit_tests/test_proxy_routes.py b/tests/proxy_unit_tests/test_proxy_routes.py index 6d704a6267..6f77456047 100644 --- a/tests/proxy_unit_tests/test_proxy_routes.py +++ b/tests/proxy_unit_tests/test_proxy_routes.py @@ -91,6 +91,11 @@ def test_routes_on_litellm_proxy(): # Bedrock Pass Through Routes ("/bedrock/model/cohere.command-r-v1:0/converse", True), ("/vertex-ai/model/text-embedding-004/embeddings", True), + # LiteLLM native RAG routes + ("/rag/ingest", True), + ("/v1/rag/ingest", True), + ("/rag/query", True), + ("/v1/rag/query", True), ], ) def test_is_llm_api_route(route: str, expected: bool): diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index c339ed2b76..ec1a13b8ab 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -790,6 +790,46 @@ def test_containers_routes_are_llm_api_routes(route): assert RouteChecks.is_llm_api_route(route) is True +@pytest.mark.parametrize( + "route", + [ + "/rag/ingest", + "/v1/rag/ingest", + "/rag/query", + "/v1/rag/query", + ], +) +def test_rag_routes_are_llm_api_routes(route): + """Test that RAG routes are recognized as LLM API routes (internal_user_viewer can access)""" + + assert RouteChecks.is_llm_api_route(route) is True + + +def test_rag_routes_accessible_to_internal_user_viewer(): + """ + Test that internal_user_viewer can access RAG routes (/rag/ingest, /rag/query). + + internal_user_viewer should be able to call RAG endpoints like chat/completions + since they are LLM API routes. For /rag/ingest, they can only add to existing + vector stores (enforced in the endpoint). + """ + + valid_token = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ) + + for route in ["/rag/ingest", "/v1/rag/ingest", "/rag/query", "/v1/rag/query"]: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + route=route, + request=MagicMock(spec=Request), + valid_token=valid_token, + request_data={}, + ) + + def test_videos_route_accessible_to_internal_users(): """ Test that internal users can access the videos routes. diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 68abf2a1c9..79c2ed4158 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -181,6 +181,16 @@ def test_route_checks_is_llm_api_route(): for route in mcp_routes: assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" + # Test LiteLLM native RAG routes + rag_routes = [ + "/rag/ingest", + "/v1/rag/ingest", + "/rag/query", + "/v1/rag/query", + ] + for route in rag_routes: + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" + # Test routes with placeholders placeholder_routes = [ "/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ", diff --git a/tests/test_litellm/proxy/rag_endpoints/__init__.py b/tests/test_litellm/proxy/rag_endpoints/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py new file mode 100644 index 0000000000..945afd886c --- /dev/null +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -0,0 +1,130 @@ +""" +Tests for RAG proxy endpoints. + +Covers: +- internal_user_viewer restriction: can only ingest to existing vector stores (must provide vector_store_id) +""" + +import io +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +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 + + +@pytest.fixture +def client_internal_user_viewer(): + """Test client with internal_user_viewer auth.""" + mock_auth = UserAPIKeyAuth( + user_id="test_viewer_user", + user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ) + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + try: + yield TestClient(app) + finally: + app.dependency_overrides = original_overrides + + +@pytest.fixture +def client_internal_user(): + """Test client with internal_user auth (can create new vector stores).""" + mock_auth = UserAPIKeyAuth( + user_id="test_internal_user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + try: + yield TestClient(app) + finally: + app.dependency_overrides = original_overrides + + +def test_internal_user_viewer_rag_ingest_without_vector_store_id_rejected( + client_internal_user_viewer, +): + """ + internal_user_viewer cannot create new vector stores - must provide vector_store_id. + """ + # Form upload without vector_store_id (would create new store) + response = client_internal_user_viewer.post( + "/v1/rag/ingest", + files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, + data={ + "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' + }, + ) + + assert response.status_code == 403 + detail = response.json() + assert "detail" in detail + error_msg = ( + detail["detail"]["error"] + if isinstance(detail["detail"], dict) + else str(detail["detail"]) + ) + assert "internal_user_viewer" in error_msg + assert "vector_store_id" in error_msg + + +def test_internal_user_viewer_rag_ingest_with_vector_store_id_passes_check( + client_internal_user_viewer, +): + """ + internal_user_viewer with vector_store_id passes the role check. + (Actual ingest may fail due to missing API keys, but we get past 403.) + """ + with patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "vs_existing", "file_id": "file_123"}, + ): + response = client_internal_user_viewer.post( + "/v1/rag/ingest", + files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, + data={ + "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai","vector_store_id":"vs_699651f6b6688191b0a210c00a686d20"}}}' + }, + ) + + # Should not be 403 (role check passed) + assert response.status_code != 403, ( + f"internal_user_viewer with vector_store_id should pass role check. " + f"Response: {response.json()}" + ) + + +def test_internal_user_rag_ingest_without_vector_store_id_allowed(client_internal_user): + """ + internal_user can create new vector stores (no vector_store_id required). + """ + with patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "vs_new", "file_id": "file_123"}, + ): + response = client_internal_user.post( + "/v1/rag/ingest", + files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, + data={ + "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' + }, + ) + + # Should not be 403 + assert response.status_code != 403, ( + f"internal_user should be allowed to create new vector stores. " + f"Response: {response.json()}" + )