From ddb29b306c210f883f0e06f673d460efd613cf2d Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Thu, 3 Jul 2025 22:02:06 -0400 Subject: [PATCH] chore: Migrating VectorDB to use OpenAI compatible identifier Signed-off-by: Francisco Javier Arceo --- docs/_static/llama-stack-spec.html | 10 +++++++- docs/_static/llama-stack-spec.yaml | 6 +++++ llama_stack/apis/vector_dbs/vector_dbs.py | 3 +++ llama_stack/apis/vector_io/vector_io.py | 2 +- llama_stack/distribution/routers/vector_io.py | 11 +++++---- .../distribution/routing_tables/vector_dbs.py | 6 +++-- .../remote/vector_io/chroma/chroma.py | 2 +- .../remote/vector_io/pgvector/pgvector.py | 2 +- .../utils/memory/openai_vector_store_mixin.py | 24 +++++++++---------- 9 files changed, 44 insertions(+), 22 deletions(-) diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index ae9ad5d4c..eb9af5737 100644 --- a/docs/_static/llama-stack-spec.html +++ b/docs/_static/llama-stack-spec.html @@ -11285,6 +11285,9 @@ }, "embedding_dimension": { "type": "integer" + }, + "provider_vector_db_name": { + "type": "string" } }, "additionalProperties": false, @@ -13543,7 +13546,8 @@ }, "additionalProperties": false, "required": [ - "name" + "name", + "provider_vector_db_id" ], "title": "OpenaiCreateVectorStoreRequest" }, @@ -15571,6 +15575,10 @@ "provider_vector_db_id": { "type": "string", "description": "The identifier of the vector database in the provider." + }, + "provider_vector_db_name": { + "type": "string", + "description": "The name of the vector database." } }, "additionalProperties": false, diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml index 48cefe12b..ce461237b 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/docs/_static/llama-stack-spec.yaml @@ -7944,6 +7944,8 @@ components: type: string embedding_dimension: type: integer + provider_vector_db_name: + type: string additionalProperties: false required: - identifier @@ -9461,6 +9463,7 @@ components: additionalProperties: false required: - name + - provider_vector_db_id title: OpenaiCreateVectorStoreRequest VectorStoreFileCounts: type: object @@ -10897,6 +10900,9 @@ components: type: string description: >- The identifier of the vector database in the provider. + provider_vector_db_name: + type: string + description: The name of the vector database. additionalProperties: false required: - vector_db_id diff --git a/llama_stack/apis/vector_dbs/vector_dbs.py b/llama_stack/apis/vector_dbs/vector_dbs.py index 405852476..31bbfaa0a 100644 --- a/llama_stack/apis/vector_dbs/vector_dbs.py +++ b/llama_stack/apis/vector_dbs/vector_dbs.py @@ -19,6 +19,7 @@ class VectorDB(Resource): embedding_model: str embedding_dimension: int + provider_vector_db_name: str | None = None @property def vector_db_id(self) -> str: @@ -71,6 +72,7 @@ class VectorDBs(Protocol): embedding_dimension: int | None = 384, provider_id: str | None = None, provider_vector_db_id: str | None = None, + provider_vector_db_name: str | None = None, ) -> VectorDB: """Register a vector database. @@ -79,6 +81,7 @@ class VectorDBs(Protocol): :param embedding_dimension: The dimension of the embedding model. :param provider_id: The identifier of the provider. :param provider_vector_db_id: The identifier of the vector database in the provider. + :param provider_vector_db_name: The name of the vector database. :returns: A VectorDB. """ ... diff --git a/llama_stack/apis/vector_io/vector_io.py b/llama_stack/apis/vector_io/vector_io.py index 2d4131315..f8318bed8 100644 --- a/llama_stack/apis/vector_io/vector_io.py +++ b/llama_stack/apis/vector_io/vector_io.py @@ -346,7 +346,7 @@ class VectorIO(Protocol): embedding_model: str | None = None, embedding_dimension: int | None = 384, provider_id: str | None = None, - provider_vector_db_id: str | None = None, + provider_vector_db_id: str = "", ) -> VectorStoreObject: """Creates a vector store. diff --git a/llama_stack/distribution/routers/vector_io.py b/llama_stack/distribution/routers/vector_io.py index 4bd5952dc..50336056e 100644 --- a/llama_stack/distribution/routers/vector_io.py +++ b/llama_stack/distribution/routers/vector_io.py @@ -5,6 +5,7 @@ # the root directory of this source tree. import asyncio +import uuid from typing import Any from llama_stack.apis.common.content_types import ( @@ -82,6 +83,7 @@ class VectorIORouter(VectorIO): embedding_dimension: int | None = 384, provider_id: str | None = None, provider_vector_db_id: str | None = None, + provider_vector_db_name: str | None = None, ) -> None: logger.debug(f"VectorIORouter.register_vector_db: {vector_db_id}, {embedding_model}") await self.routing_table.register_vector_db( @@ -90,6 +92,7 @@ class VectorIORouter(VectorIO): embedding_dimension, provider_id, provider_vector_db_id, + provider_vector_db_name, ) async def insert_chunks( @@ -123,7 +126,7 @@ 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, + provider_vector_db_id: str = "", ) -> VectorStoreObject: logger.debug(f"VectorIORouter.openai_create_vector_store: name={name}, provider_id={provider_id}") @@ -135,17 +138,17 @@ class VectorIORouter(VectorIO): embedding_model, embedding_dimension = embedding_model_info logger.info(f"No embedding model specified, using first available: {embedding_model}") - vector_db_id = name + vector_db_id = f"vs_{uuid.uuid4()}" registered_vector_db = await self.routing_table.register_vector_db( vector_db_id, embedding_model, embedding_dimension, provider_id, provider_vector_db_id, + name, ) - return await self.routing_table.get_provider_impl(registered_vector_db.identifier).openai_create_vector_store( - vector_db_id, + name, file_ids=file_ids, expires_after=expires_after, chunking_strategy=chunking_strategy, diff --git a/llama_stack/distribution/routing_tables/vector_dbs.py b/llama_stack/distribution/routing_tables/vector_dbs.py index 542e965f8..3c7fd258a 100644 --- a/llama_stack/distribution/routing_tables/vector_dbs.py +++ b/llama_stack/distribution/routing_tables/vector_dbs.py @@ -35,9 +35,10 @@ class VectorDBsRoutingTable(CommonRoutingTableImpl, VectorDBs): embedding_model: str, embedding_dimension: int | None = 384, provider_id: str | None = None, - provider_vector_db_id: str | None = None, + provider_vector_db_id: str = "", + provider_vector_db_name: str | None = None, ) -> VectorDB: - if provider_vector_db_id is None: + if provider_vector_db_id == "": provider_vector_db_id = vector_db_id if provider_id is None: if len(self.impls_by_provider_id) > 0: @@ -62,6 +63,7 @@ class VectorDBsRoutingTable(CommonRoutingTableImpl, VectorDBs): "provider_resource_id": provider_vector_db_id, "embedding_model": embedding_model, "embedding_dimension": model.metadata["embedding_dimension"], + "vector_db_name": provider_vector_db_name, } vector_db = TypeAdapter(VectorDBWithOwner).validate_python(vector_db_data) await self.register_object(vector_db) diff --git a/llama_stack/providers/remote/vector_io/chroma/chroma.py b/llama_stack/providers/remote/vector_io/chroma/chroma.py index 3bef39e9c..4c5d3b1b0 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/llama_stack/providers/remote/vector_io/chroma/chroma.py @@ -217,7 +217,7 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate): embedding_model: str | None = None, embedding_dimension: int | None = 384, provider_id: str | None = None, - provider_vector_db_id: str | None = None, + provider_vector_db_id: str = "", ) -> VectorStoreObject: raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma") diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index c3cdef9b8..b30c73767 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -247,7 +247,7 @@ 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, + provider_vector_db_id: str = "", ) -> 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 7c97ff7f6..d1a01948e 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -8,7 +8,6 @@ import asyncio import logging import mimetypes import time -import uuid from abc import ABC, abstractmethod from typing import Any @@ -144,11 +143,12 @@ class OpenAIVectorStoreMixin(ABC): embedding_model: str | None = None, embedding_dimension: int | None = 384, provider_id: str | None = None, - provider_vector_db_id: str | None = None, + provider_vector_db_id: str = "", ) -> VectorStoreObject: """Creates a vector store.""" - store_id = name or str(uuid.uuid4()) created_at = int(time.time()) + if provider_vector_db_id is None: + raise ValueError("Provider vector DB ID is required") if provider_id is None: raise ValueError("Provider ID is required") @@ -160,13 +160,13 @@ class OpenAIVectorStoreMixin(ABC): if embedding_dimension is None: raise ValueError("Embedding dimension is required") - provider_vector_db_id = provider_vector_db_id or store_id vector_db = VectorDB( - identifier=store_id, + identifier=provider_vector_db_id, embedding_dimension=embedding_dimension, embedding_model=embedding_model, provider_id=provider_id, provider_resource_id=provider_vector_db_id, + provider_vector_db_name=name, ) # Register the vector DB await self.register_vector_db(vector_db) @@ -182,11 +182,11 @@ class OpenAIVectorStoreMixin(ABC): in_progress=0, total=0, ) - store_info = { - "id": store_id, + store_info: dict[str, Any] = { + "id": provider_vector_db_id, "object": "vector_store", "created_at": created_at, - "name": store_id, + "name": name, "usage_bytes": 0, "file_counts": file_counts.model_dump(), "status": status, @@ -206,18 +206,18 @@ class OpenAIVectorStoreMixin(ABC): store_info["metadata"] = metadata # Save to persistent storage (provider-specific) - await self._save_openai_vector_store(store_id, store_info) + await self._save_openai_vector_store(provider_vector_db_id, store_info) # Store in memory cache - self.openai_vector_stores[store_id] = store_info + self.openai_vector_stores[provider_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(store_id, file_id) for file_id in file_ids] + tasks = [self.openai_attach_file_to_vector_store(provider_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[store_id] + store_info = self.openai_vector_stores[provider_vector_db_id] return VectorStoreObject.model_validate(store_info) async def openai_list_vector_stores(