From b01c72a7de397994cbd6ae81e6ffebf0d20f39c7 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Mon, 14 Jul 2025 23:31:34 -0400 Subject: [PATCH] removing provider_vector_db from VectorIORouter.openai_create_vector_store, VectorIO.openai_create_vector_store, and PGVector Signed-off-by: Francisco Javier Arceo --- llama_stack/distribution/routers/vector_io.py | 3 +-- .../remote/vector_io/pgvector/pgvector.py | 1 - .../utils/memory/openai_vector_store_mixin.py | 22 +++++++++---------- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/llama_stack/distribution/routers/vector_io.py b/llama_stack/distribution/routers/vector_io.py index d35d5fa05..eba292bfc 100644 --- a/llama_stack/distribution/routers/vector_io.py +++ b/llama_stack/distribution/routers/vector_io.py @@ -126,7 +126,6 @@ class VectorIORouter(VectorIO): embedding_model: str | None = None, embedding_dimension: int | None = None, provider_id: str | None = None, - provider_vector_db_id: str | None = None, ) -> VectorStoreObject: logger.debug(f"VectorIORouter.openai_create_vector_store: name={name}, provider_id={provider_id}") @@ -144,7 +143,7 @@ class VectorIORouter(VectorIO): embedding_model, embedding_dimension, provider_id, - provider_vector_db_id, + vector_db_id, name, ) return await self.routing_table.get_provider_impl(registered_vector_db.identifier).openai_create_vector_store( diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index c3cdef9b8..38ec3279b 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -247,7 +247,6 @@ class PGVectorVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate): embedding_model: str | None = None, embedding_dimension: int | None = 384, provider_id: str | None = None, - provider_vector_db_id: str | None = None, ) -> VectorStoreObject: raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector") diff --git a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py index 19147c43f..7a2bdddcd 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -148,8 +148,8 @@ class OpenAIVectorStoreMixin(ABC): ) -> VectorStoreObject: """Creates a vector store.""" created_at = int(time.time()) - if provider_vector_db_id is None: - provider_vector_db_id = f"vs_{uuid.uuid4()}" + # Derive the canonical vector_db_id (allow override, else generate) + vector_db_id = provider_vector_db_id or f"vs_{uuid.uuid4()}" if provider_id is None: raise ValueError("Provider ID is required") @@ -157,19 +157,19 @@ class OpenAIVectorStoreMixin(ABC): if embedding_model is None: raise ValueError("Embedding model is required") - # Use provided embedding dimension or default to 384 + # Embedding dimension is required (defaulted to 384 if not provided) if embedding_dimension is None: raise ValueError("Embedding dimension is required") + # Register the VectorDB backing this vector store vector_db = VectorDB( - identifier=provider_vector_db_id, + identifier=vector_db_id, embedding_dimension=embedding_dimension, embedding_model=embedding_model, provider_id=provider_id, - provider_resource_id=provider_vector_db_id, + provider_resource_id=vector_db_id, vector_db_name=name, ) - # Register the vector DB await self.register_vector_db(vector_db) # Create OpenAI vector store metadata @@ -184,7 +184,7 @@ class OpenAIVectorStoreMixin(ABC): total=0, ) store_info: dict[str, Any] = { - "id": provider_vector_db_id, + "id": vector_db_id, "object": "vector_store", "created_at": created_at, "name": name, @@ -207,18 +207,18 @@ class OpenAIVectorStoreMixin(ABC): store_info["metadata"] = metadata # Save to persistent storage (provider-specific) - await self._save_openai_vector_store(provider_vector_db_id, store_info) + await self._save_openai_vector_store(vector_db_id, store_info) # Store in memory cache - self.openai_vector_stores[provider_vector_db_id] = store_info + self.openai_vector_stores[vector_db_id] = store_info # Now that our vector store is created, attach any files that were provided file_ids = file_ids or [] - tasks = [self.openai_attach_file_to_vector_store(provider_vector_db_id, file_id) for file_id in file_ids] + tasks = [self.openai_attach_file_to_vector_store(vector_db_id, file_id) for file_id in file_ids] await asyncio.gather(*tasks) # Get the updated store info and return it - store_info = self.openai_vector_stores[provider_vector_db_id] + store_info = self.openai_vector_stores[vector_db_id] return VectorStoreObject.model_validate(store_info) async def openai_list_vector_stores(