From eb2c8bb2132dea1465d2e4c9b12eda4ebcd12ca9 Mon Sep 17 00:00:00 2001 From: sujal011 Date: Wed, 29 Apr 2026 17:56:17 +0530 Subject: [PATCH] feat: implement Litellm-pgvector support OpenAI-specific RAG ingestion and add public ingestion API registry --- litellm/rag/ingestion/openai_ingestion.py | 10 +++++----- litellm/rag/main.py | 1 + litellm/types/utils.py | 1 + litellm/utils.py | 2 +- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/rag/ingestion/openai_ingestion.py b/litellm/rag/ingestion/openai_ingestion.py index 891e3d0e914..4e47f4e88fb 100644 --- a/litellm/rag/ingestion/openai_ingestion.py +++ b/litellm/rag/ingestion/openai_ingestion.py @@ -89,7 +89,7 @@ class OpenAIRAGIngestion(BaseRAGIngestion): ) create_response = await vector_store_acreate( name=self.ingest_name or "litellm-rag-ingest", - custom_llm_provider="openai", + custom_llm_provider=self.custom_llm_provider, expires_after=expires_after, api_key=api_key, api_base=api_base, @@ -99,7 +99,7 @@ class OpenAIRAGIngestion(BaseRAGIngestion): # Upload file and attach to vector store result_file_id = None if file_content and filename and vector_store_id: - # Upload file to OpenAI + # Upload file to OpenAI (or compatible provider) file_response = await litellm.acreate_file( file=( filename, @@ -107,17 +107,17 @@ class OpenAIRAGIngestion(BaseRAGIngestion): content_type or "application/octet-stream", ), purpose="assistants", - custom_llm_provider="openai", + custom_llm_provider=self.custom_llm_provider, api_key=api_key, api_base=api_base, ) result_file_id = file_response.id - # Attach file to vector store (OpenAI handles chunking/embedding) + # Attach file to vector store (Provider handles chunking/embedding) await vector_store_file_acreate( vector_store_id=vector_store_id, file_id=result_file_id, - custom_llm_provider="openai", + custom_llm_provider=self.custom_llm_provider, chunking_strategy=cast( Optional[Dict[str, Any]], self.chunking_strategy ), diff --git a/litellm/rag/main.py b/litellm/rag/main.py index e3d354b6c33..e05982efe17 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -52,6 +52,7 @@ INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = { "gemini": GeminiRAGIngestion, "s3_vectors": S3VectorsRAGIngestion, "vertex_ai": VertexAIRAGIngestion, + "pg_vector": OpenAIRAGIngestion, } diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ed29d49fc29..d25fedf6e19 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3340,6 +3340,7 @@ LlmProvidersSet = {provider.value for provider in LlmProviders} OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = { LlmProviders.OPENAI.value, LlmProviders.HOSTED_VLLM.value, + LlmProviders.PG_VECTOR.value, } ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "vertex_ai"] diff --git a/litellm/utils.py b/litellm/utils.py index 027c9fedced..5e6c24fd094 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8861,7 +8861,7 @@ class ProviderConfigManager: def get_provider_vector_store_files_config( provider: LlmProviders, ) -> Optional[BaseVectorStoreFilesConfig]: - if litellm.LlmProviders.OPENAI == provider: + if provider in (litellm.LlmProviders.OPENAI, litellm.LlmProviders.PG_VECTOR): from litellm.llms.openai.vector_store_files.transformation import ( OpenAIVectorStoreFilesConfig, )