From 75eec146919ce7b2c6d6a1a28f13bc710e866e72 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 26 Jan 2026 15:12:18 -0800 Subject: [PATCH] feat: _save_vector_store_to_db_from_rag_ingest --- litellm/proxy/rag_endpoints/endpoints.py | 85 ++++++++ .../management_endpoints.py | 189 ++++++++++++------ 2 files changed, 216 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 79b4fd6873d..470372f1476 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -26,6 +26,75 @@ from litellm.proxy.common_utils.http_parsing_utils import ( router = APIRouter() +async def _save_vector_store_to_db_from_rag_ingest( + response: Any, + ingest_options: Dict[str, Any], + prisma_client, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """ + Helper function to save a newly created vector store from RAG ingest to the database. + + This function: + - Extracts vector store ID and config from the ingest response + - Checks if the vector store already exists in the database + - Creates a new database entry if it doesn't exist + - Adds the vector store to the registry + + Args: + response: The response from litellm.aingest() + ingest_options: The ingest options containing vector store config + prisma_client: The Prisma database client + user_api_key_dict: User API key authentication info + """ + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + create_vector_store_in_db, + ) + + vector_store_id = response.get("vector_store_id") + if vector_store_id is None or not isinstance(vector_store_id, str): + verbose_proxy_logger.warning( + "Vector store ID is None or not a string, skipping database save" + ) + return + + vector_store_config = ingest_options.get("vector_store", {}) + custom_llm_provider = vector_store_config.get("custom_llm_provider") + + try: + # Check if vector store already exists in database + existing_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store_id} + ) + ) + + # Only create if it doesn't exist + if existing_vector_store is None: + verbose_proxy_logger.info( + f"Saving newly created vector store {vector_store_id} to database" + ) + + await create_vector_store_in_db( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider or "openai", + prisma_client=prisma_client, + vector_store_name=f"RAG Vector Store - {vector_store_id[:8]}", + vector_store_description="Created via RAG ingest endpoint", + created_by=user_api_key_dict.user_id, + updated_by=user_api_key_dict.user_id, + ) + + verbose_proxy_logger.info( + f"Vector store {vector_store_id} saved to database successfully" + ) + except Exception as db_error: + # Log the error but don't fail the request since ingestion succeeded + verbose_proxy_logger.warning( + f"Failed to save vector store {vector_store_id} to database: {db_error}" + ) + + async def parse_rag_ingest_request( request: Request, ) -> Tuple[Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str]]: @@ -158,9 +227,11 @@ async def rag_ingest( add_litellm_data_to_request, general_settings, llm_router, + prisma_client, proxy_config, version, ) + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore try: # Parse request @@ -189,6 +260,20 @@ async def rag_ingest( **request_data, ) + # Save vector store to database if it was newly created and prisma_client is available + if ( + prisma_client is not None + and response is not None + and isinstance(response, dict) + and response.get("vector_store_id") + ): + await _save_vector_store_to_db_from_rag_ingest( + response=response, + ingest_options=ingest_options, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + ) + return response except HTTPException: diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index bc61a60fe5a..47adb8c03ba 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -133,6 +133,115 @@ async def _resolve_embedding_config_from_db( return None +######################################################## +# Helper Functions +######################################################## +async def create_vector_store_in_db( + vector_store_id: str, + custom_llm_provider: str, + prisma_client, + vector_store_name: Optional[str] = None, + vector_store_description: Optional[str] = None, + vector_store_metadata: Optional[Dict] = None, + litellm_params: Optional[Dict] = None, + created_by: Optional[str] = None, + updated_by: Optional[str] = None, +) -> LiteLLM_ManagedVectorStore: + """ + Helper function to create a vector store in the database. + + This function handles: + - Checking if vector store already exists + - Creating the vector store in the database + - Adding it to the vector store registry + + Returns: + LiteLLM_ManagedVectorStore: The created vector store object + + Raises: + HTTPException: If vector store already exists or database error occurs + """ + from litellm.types.router import GenericLiteLLMParams + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + # Check if vector store already exists + existing_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store_id} + ) + ) + if existing_vector_store is not None: + raise HTTPException( + status_code=400, + detail=f"Vector store with ID {vector_store_id} already exists", + ) + + # Prepare data for database + data_to_create: Dict[str, Any] = { + "vector_store_id": vector_store_id, + "custom_llm_provider": custom_llm_provider, + } + + if vector_store_name is not None: + data_to_create["vector_store_name"] = vector_store_name + if vector_store_description is not None: + data_to_create["vector_store_description"] = vector_store_description + if vector_store_metadata is not None: + data_to_create["vector_store_metadata"] = safe_dumps(vector_store_metadata) + if created_by is not None: + data_to_create["created_by"] = created_by + if updated_by is not None: + data_to_create["updated_by"] = updated_by + + # Handle litellm_params + litellm_params_json: Optional[str] = None + if litellm_params: + # Auto-resolve embedding config if embedding model is provided but config is not + embedding_model = litellm_params.get("litellm_embedding_model") + if embedding_model and not litellm_params.get("litellm_embedding_config"): + resolved_config = await _resolve_embedding_config_from_db( + embedding_model=embedding_model, + prisma_client=prisma_client + ) + if resolved_config: + litellm_params["litellm_embedding_config"] = resolved_config + verbose_proxy_logger.info( + f"Auto-resolved embedding config for model {embedding_model}" + ) + + litellm_params_dict = GenericLiteLLMParams( + **litellm_params + ).model_dump(exclude_none=True) + litellm_params_json = safe_dumps(litellm_params_dict) + + data_to_create["litellm_params"] = litellm_params_json + + # Create in database + _new_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.create( + data=data_to_create + ) + ) + + new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( + **_new_vector_store.model_dump() + ) + + # Add vector store to registry + if litellm.vector_store_registry is not None: + litellm.vector_store_registry.add_vector_store_to_registry( + vector_store=new_vector_store + ) + + verbose_proxy_logger.info( + f"Vector store {vector_store_id} created in database successfully" + ) + + return new_vector_store + + ######################################################## # Management Endpoints ######################################################## @@ -156,71 +265,35 @@ async def new_vector_store( - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store """ from litellm.proxy.proxy_server import prisma_client - from litellm.types.router import GenericLiteLLMParams - - if prisma_client is None: - raise HTTPException(status_code=500, detail="Database not connected") try: - # Check if vector store already exists - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store.get("vector_store_id")} - ) - ) - if existing_vector_store is not None: + vector_store_id = vector_store.get("vector_store_id") + custom_llm_provider = vector_store.get("custom_llm_provider") + + if not vector_store_id or not custom_llm_provider: raise HTTPException( status_code=400, - detail=f"Vector store with ID {vector_store.get('vector_store_id')} already exists", - ) - - if vector_store.get("vector_store_metadata") is not None: - vector_store["vector_store_metadata"] = safe_dumps( - vector_store.get("vector_store_metadata") - ) - - # Safely handle JSON serialization of litellm_params - litellm_params_json: Optional[str] = None - _input_litellm_params: dict = vector_store.get("litellm_params", {}) or {} - if _input_litellm_params is not None: - # Auto-resolve embedding config if embedding model is provided but config is not - embedding_model = _input_litellm_params.get("litellm_embedding_model") - if embedding_model and not _input_litellm_params.get("litellm_embedding_config"): - resolved_config = await _resolve_embedding_config_from_db( - embedding_model=embedding_model, - prisma_client=prisma_client - ) - if resolved_config: - _input_litellm_params["litellm_embedding_config"] = resolved_config - verbose_proxy_logger.info( - f"Auto-resolved embedding config for model {embedding_model}" - ) - - litellm_params_dict = GenericLiteLLMParams( - **_input_litellm_params - ).model_dump(exclude_none=True) - litellm_params_json = safe_dumps(litellm_params_dict) - del vector_store["litellm_params"] - - _new_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.create( - data={ - **vector_store, - "litellm_params": litellm_params_json, - } + detail="vector_store_id and custom_llm_provider are required" ) + + # Extract and validate metadata + metadata = vector_store.get("vector_store_metadata") + validated_metadata: Optional[Dict] = None + if metadata is not None and isinstance(metadata, dict): + validated_metadata = metadata + + new_vector_store = await create_vector_store_in_db( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider, + prisma_client=prisma_client, + vector_store_name=vector_store.get("vector_store_name"), + vector_store_description=vector_store.get("vector_store_description"), + vector_store_metadata=validated_metadata, + litellm_params=vector_store.get("litellm_params"), + created_by=user_api_key_dict.user_id, + updated_by=user_api_key_dict.user_id, ) - new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( - **_new_vector_store.model_dump() - ) - - # Add vector store to registry - if litellm.vector_store_registry is not None: - litellm.vector_store_registry.add_vector_store_to_registry( - vector_store=new_vector_store - ) - return { "status": "success", "message": f"Vector store {vector_store.get('vector_store_id')} created successfully",