diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 8e247b5228d..e3c3e322567 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -46,6 +46,23 @@ if TYPE_CHECKING: router: Final = APIRouter() +_BLOCKED_VECTOR_STORE_CREDENTIAL_PARAMS: Final = { + "vertex_credentials", + "vertex_ai_credentials", + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_web_identity_token", + "aws_role_name", + "aws_session_name", + "aws_profile_name", + "aws_sts_endpoint", + "aws_external_id", + "azure_ad_token", + "api_key", + "api_base", +} + class _RAGIngestErrorDetail(TypedDict): error: ReadOnly[str] @@ -222,6 +239,7 @@ async def _save_vector_store_to_db_from_rag_ingest( vector_store_config: Final = ingest_options.get("vector_store", {}) custom_llm_provider: Final = vector_store_config.get("custom_llm_provider") + credential_name: Final = vector_store_config.get("litellm_credential_name") # Extract litellm_vector_store_params for custom name and description litellm_vector_store_params: Final = ingest_options.get("litellm_vector_store_params", {}) @@ -231,7 +249,12 @@ async def _save_vector_store_to_db_from_rag_ingest( # Extract provider-specific params from vector_store_config to save as litellm_params # This ensures params like aws_region_name, embedding_model, etc. are available for search provider_specific_params: Final = {} - excluded_keys: Final = {"custom_llm_provider", "vector_store_id"} + excluded_keys: Final = { + "custom_llm_provider", + "vector_store_id", + "litellm_credential_name", + *_BLOCKED_VECTOR_STORE_CREDENTIAL_PARAMS, + } for key, value in vector_store_config.items(): if key not in excluded_keys and value is not None: provider_specific_params[key] = value @@ -268,6 +291,7 @@ async def _save_vector_store_to_db_from_rag_ingest( vector_store_description=vector_store_description, vector_store_metadata=initial_metadata, litellm_params=(provider_specific_params if provider_specific_params else None), + litellm_credential_name=credential_name, team_id=user_api_key_dict.team_id, user_id=user_api_key_dict.user_id, ) @@ -411,22 +435,6 @@ async def parse_rag_ingest_request( # google-auth's identity_pool credential refresh. # api_base is also blocked: a user-controlled base URL causes the server # to send its configured provider credentials to an attacker endpoint. - _BLOCKED_VECTOR_STORE_CREDENTIAL_PARAMS: Final = { - "vertex_credentials", - "vertex_ai_credentials", - "aws_access_key_id", - "aws_secret_access_key", - "aws_session_token", - "aws_web_identity_token", - "aws_role_name", - "aws_session_name", - "aws_profile_name", - "aws_sts_endpoint", - "aws_external_id", - "azure_ad_token", - "api_key", - "api_base", - } vector_store_opts: Final[object] = ingest_options.get("vector_store", {}) if isinstance(vector_store_opts, dict): for field in _BLOCKED_VECTOR_STORE_CREDENTIAL_PARAMS: diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index da1bc0a1feb..ae937074269 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -14,6 +14,7 @@ from __future__ import annotations import base64 from abc import ABC, abstractmethod +from copy import deepcopy from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -60,7 +61,9 @@ class BaseRAGIngestion(ABC): ingest_options.get("chunking_strategy") or {"type": "auto"}, ) self.embedding_config = ingest_options.get("embedding") - self.vector_store_config: dict[str, Any] = cast(dict[str, Any], ingest_options.get("vector_store") or {}) + self.vector_store_config: dict[str, Any] = deepcopy( + cast(dict[str, Any], ingest_options.get("vector_store") or {}) + ) self.ingest_name = ingest_options.get("name") # Load credentials from litellm_credential_name if provided in vector_store config 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 e4474db556b..73d9fd4031d 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -6,6 +6,7 @@ Covers: """ import io +import json from copy import deepcopy from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -107,7 +108,9 @@ def test_rag_ingest_rejects_non_string_provider_before_execution( @pytest.mark.asyncio -async def test_rag_ingest_named_credentials_are_request_local_and_not_persisted() -> None: +async def test_rag_ingest_named_credentials_are_request_local_and_not_persisted( + monkeypatch: pytest.MonkeyPatch, +) -> None: from litellm.proxy.rag_endpoints.endpoints import _save_vector_store_to_db_from_rag_ingest from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion @@ -133,32 +136,40 @@ async def test_rag_ingest_named_credentials_are_request_local_and_not_persisted( }, ) prisma_client: Final = MagicMock() - prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) + vector_store_table: Final = prisma_client.db.litellm_managedvectorstorestable + vector_store_table.find_unique = AsyncMock(return_value=None) + created_vector_store: Final = MagicMock() + created_vector_store.model_dump.return_value = { + "vector_store_id": "vs_named_credential", + "custom_llm_provider": "openai", + "litellm_credential_name": "rag-openai", + "litellm_params": '{"ttl_days": 7}', + } + vector_store_table.create = AsyncMock(return_value=created_vector_store) + monkeypatch.setattr(litellm, "credential_list", [credential]) + monkeypatch.setattr(litellm, "vector_store_registry", None) - with ( - patch.object(litellm, "credential_list", [credential]), - patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.create_vector_store_in_db", - new_callable=AsyncMock, - ) as create_vector_store, - ): - OpenAIRAGIngestion(ingest_options=ingest_options) - await _save_vector_store_to_db_from_rag_ingest( - response={"vector_store_id": "vs_named_credential"}, - ingest_options=ingest_options, - prisma_client=prisma_client, - user_api_key_dict=UserAPIKeyAuth(user_id="user-123", team_id="team-123"), - ) + OpenAIRAGIngestion(ingest_options=ingest_options) + await _save_vector_store_to_db_from_rag_ingest( + response={"vector_store_id": "vs_named_credential"}, + ingest_options=ingest_options, + prisma_client=prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_id="user-123", team_id="team-123"), + ) - persistence_args: Final = create_vector_store.await_args.kwargs + persistence_data: Final = vector_store_table.create.await_args.kwargs["data"] + persisted_params: Final = json.loads(persistence_data["litellm_params"]) + credential_keys: Final = {*credential.credential_values, "litellm_credential_name"} assert { "caller_request": ingest_options, - "credential_column": persistence_args.get("litellm_credential_name"), - "litellm_params": persistence_args["litellm_params"], + "credential_column": persistence_data.get("litellm_credential_name"), + "provider_setting": persisted_params.get("ttl_days"), + "persisted_credential_keys": credential_keys.intersection(persisted_params), } == { "caller_request": original_ingest_options, "credential_column": "rag-openai", - "litellm_params": {"ttl_days": 7}, + "provider_setting": 7, + "persisted_credential_keys": set(), }