feat: implement Litellm-pgvector support OpenAI-specific RAG ingestion and add public ingestion API registry

This commit is contained in:
sujal011 2026-04-29 17:56:17 +05:30
parent 6e6b2ca2d8
commit eb2c8bb213
4 changed files with 8 additions and 6 deletions

View file

@ -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
),

View file

@ -52,6 +52,7 @@ INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = {
"gemini": GeminiRAGIngestion,
"s3_vectors": S3VectorsRAGIngestion,
"vertex_ai": VertexAIRAGIngestion,
"pg_vector": OpenAIRAGIngestion,
}

View file

@ -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"]

View file

@ -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,
)