diff --git a/llama_stack/core/routers/vector_io.py b/llama_stack/core/routers/vector_io.py index 79789ef0a..dc7b3a694 100644 --- a/llama_stack/core/routers/vector_io.py +++ b/llama_stack/core/routers/vector_io.py @@ -55,30 +55,18 @@ class VectorIORouter(VectorIO): logger.debug("VectorIORouter.shutdown") pass - async def _get_first_embedding_model(self) -> tuple[str, int] | None: - """Get the first available embedding model identifier.""" - try: - # Get all models from the routing table - all_models = await self.routing_table.get_all_with_type("model") + async def _get_embedding_model_dimension(self, embedding_model_id: str) -> int: + """Get the embedding dimension for a specific embedding model.""" + all_models = await self.routing_table.get_all_with_type("model") - # Filter for embedding models - embedding_models = [ - model - for model in all_models - if hasattr(model, "model_type") and model.model_type == ModelType.embedding - ] - - if embedding_models: - dimension = embedding_models[0].metadata.get("embedding_dimension", None) + for model in all_models: + if model.identifier == embedding_model_id and model.model_type == ModelType.embedding: + dimension = model.metadata.get("embedding_dimension") if dimension is None: - raise ValueError(f"Embedding model {embedding_models[0].identifier} has no embedding dimension") - return embedding_models[0].identifier, dimension - else: - logger.warning("No embedding models found in the routing table") - return None - except Exception as e: - logger.error(f"Error getting embedding models: {e}") - return None + raise ValueError(f"Embedding model '{embedding_model_id}' has no embedding_dimension in metadata") + return int(dimension) + + raise ValueError(f"Embedding model '{embedding_model_id}' not found or not an embedding model") async def register_vector_db( self, @@ -129,20 +117,30 @@ class VectorIORouter(VectorIO): # Extract llama-stack-specific parameters from extra_body extra = params.model_extra or {} embedding_model = extra.get("embedding_model") - embedding_dimension = extra.get("embedding_dimension", 384) + embedding_dimension = extra.get("embedding_dimension") provider_id = extra.get("provider_id") logger.debug(f"VectorIORouter.openai_create_vector_store: name={params.name}, provider_id={provider_id}") - # If no embedding model is provided, use the first available one - # TODO: this branch will soon be deleted so you _must_ provide the embedding_model when - # creating a vector store + # Require explicit embedding model specification if embedding_model is None: - embedding_model_info = await self._get_first_embedding_model() - if embedding_model_info is None: - raise ValueError("No embedding model provided and no embedding models available in the system") - embedding_model, embedding_dimension = embedding_model_info - logger.info(f"No embedding model specified, using first available: {embedding_model}") + raise ValueError("embedding_model is required in extra_body when creating a vector store") + + if embedding_dimension is None: + embedding_dimension = await self._get_embedding_model_dimension(embedding_model) + + # Auto-select provider if not specified + if provider_id is None: + num_providers = len(self.routing_table.impls_by_provider_id) + if num_providers == 0: + raise ValueError("No vector_io providers available") + if num_providers > 1: + available_providers = list(self.routing_table.impls_by_provider_id.keys()) + raise ValueError( + f"Multiple vector_io providers available. Please specify provider_id in extra_body. " + f"Available providers: {available_providers}" + ) + provider_id = list(self.routing_table.impls_by_provider_id.keys())[0] vector_db_id = f"vs_{uuid.uuid4()}" registered_vector_db = await self.routing_table.register_vector_db( 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 70bcbba32..02c3d9730 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -353,14 +353,12 @@ class OpenAIVectorStoreMixin(ABC): provider_vector_db_id = extra.get("provider_vector_db_id") embedding_model = extra.get("embedding_model") embedding_dimension = extra.get("embedding_dimension", 384) - provider_id = extra.get("provider_id") + # use provider_id set by router; fallback to provider's own ID when used directly via --stack-config + provider_id = extra.get("provider_id") or getattr(self, "__provider_id__", None) # Derive the canonical vector_db_id (allow override, else generate) vector_db_id = provider_vector_db_id or generate_object_id("vector_store", lambda: f"vs_{uuid.uuid4()}") - if provider_id is None: - raise ValueError("Provider ID is required") - if embedding_model is None: raise ValueError("Embedding model is required") @@ -369,6 +367,9 @@ class OpenAIVectorStoreMixin(ABC): raise ValueError("Embedding dimension is required") # Register the VectorDB backing this vector store + if provider_id is None: + raise ValueError("Provider ID is required but was not provided") + vector_db = VectorDB( identifier=vector_db_id, embedding_dimension=embedding_dimension, diff --git a/tests/integration/vector_io/test_openai_vector_stores.py b/tests/integration/vector_io/test_openai_vector_stores.py index 347b43145..904e382e1 100644 --- a/tests/integration/vector_io/test_openai_vector_stores.py +++ b/tests/integration/vector_io/test_openai_vector_stores.py @@ -146,8 +146,6 @@ def test_openai_create_vector_store( metadata={"purpose": "testing", "environment": "integration"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -175,8 +173,6 @@ def test_openai_list_vector_stores( metadata={"type": "test"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) store2 = client.vector_stores.create( @@ -184,8 +180,6 @@ def test_openai_list_vector_stores( metadata={"type": "test"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -220,8 +214,6 @@ def test_openai_retrieve_vector_store( metadata={"purpose": "retrieval_test"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -249,8 +241,6 @@ def test_openai_update_vector_store( metadata={"version": "1.0"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) time.sleep(1) @@ -282,8 +272,6 @@ def test_openai_delete_vector_store( metadata={"purpose": "deletion_test"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -314,8 +302,6 @@ def test_openai_vector_store_search_empty( metadata={"purpose": "search_testing"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -346,8 +332,6 @@ def test_openai_vector_store_with_chunks( metadata={"purpose": "chunks_testing"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -412,8 +396,6 @@ def test_openai_vector_store_search_relevance( metadata={"purpose": "relevance_testing"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -457,8 +439,6 @@ def test_openai_vector_store_search_with_ranking_options( metadata={"purpose": "ranking_testing"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -500,8 +480,6 @@ def test_openai_vector_store_search_with_high_score_filter( metadata={"purpose": "high_score_filtering"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -561,8 +539,6 @@ def test_openai_vector_store_search_with_max_num_results( metadata={"purpose": "max_num_results_testing"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -596,8 +572,6 @@ def test_openai_vector_store_attach_file( name="test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -666,8 +640,6 @@ def test_openai_vector_store_attach_files_on_creation( file_ids=file_ids, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -713,8 +685,6 @@ def test_openai_vector_store_list_files( name="test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -799,8 +769,6 @@ def test_openai_vector_store_retrieve_file_contents( name="test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -819,8 +787,6 @@ def test_openai_vector_store_retrieve_file_contents( attributes=attributes, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -857,8 +823,6 @@ def test_openai_vector_store_delete_file( name="test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -918,8 +882,6 @@ def test_openai_vector_store_delete_file_removes_from_vector_store( name="test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -965,8 +927,6 @@ def test_openai_vector_store_update_file( name="test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1026,8 +986,6 @@ def test_create_vector_store_files_duplicate_vector_store_name( name="test_store_with_files", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) assert vector_store.file_counts.completed == 0 @@ -1040,8 +998,6 @@ def test_create_vector_store_files_duplicate_vector_store_name( name="test_store_with_files", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1053,8 +1009,6 @@ def test_create_vector_store_files_duplicate_vector_store_name( file_id=file_ids[0], extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) assert created_file.status == "completed" @@ -1065,8 +1019,6 @@ def test_create_vector_store_files_duplicate_vector_store_name( file_id=file_ids[1], extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) assert created_file_from_non_deleted_vector_store.status == "completed" @@ -1087,8 +1039,6 @@ def test_openai_vector_store_search_modes( metadata={"purpose": "search_mode_testing"}, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1120,8 +1070,6 @@ def test_openai_vector_store_file_batch_create_and_retrieve( name="batch_test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1139,8 +1087,6 @@ def test_openai_vector_store_file_batch_create_and_retrieve( file_ids=file_ids, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1187,8 +1133,6 @@ def test_openai_vector_store_file_batch_list_files( name="batch_list_test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1206,8 +1150,6 @@ def test_openai_vector_store_file_batch_list_files( file_ids=file_ids, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1284,8 +1226,6 @@ def test_openai_vector_store_file_batch_cancel( name="batch_cancel_test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1303,8 +1243,6 @@ def test_openai_vector_store_file_batch_cancel( file_ids=file_ids, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1343,8 +1281,6 @@ def test_openai_vector_store_file_batch_retrieve_contents( name="batch_contents_test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1367,8 +1303,6 @@ def test_openai_vector_store_file_batch_retrieve_contents( file_ids=file_ids, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1420,8 +1354,6 @@ def test_openai_vector_store_file_batch_error_handling( name="batch_error_test_store", extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -1433,8 +1365,6 @@ def test_openai_vector_store_file_batch_error_handling( file_ids=file_ids, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) diff --git a/tests/integration/vector_io/test_vector_io.py b/tests/integration/vector_io/test_vector_io.py index f2205ed0a..653299338 100644 --- a/tests/integration/vector_io/test_vector_io.py +++ b/tests/integration/vector_io/test_vector_io.py @@ -52,8 +52,6 @@ def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id, embe name=vector_db_name, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -73,8 +71,6 @@ def test_vector_db_register(client_with_empty_registry, embedding_model_id, embe name=vector_db_name, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -110,8 +106,6 @@ def test_insert_chunks(client_with_empty_registry, embedding_model_id, embedding name=vector_db_name, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -152,8 +146,6 @@ def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, e name=vector_db_name, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -202,8 +194,6 @@ def test_query_returns_valid_object_when_identical_to_embedding_in_vdb( name=vector_db_name, extra_body={ "embedding_model": embedding_model_id, - "embedding_dimension": embedding_dimension, - "provider_id": "my_provider", }, ) @@ -234,3 +224,35 @@ def test_query_returns_valid_object_when_identical_to_embedding_in_vdb( assert len(response.chunks) > 0 assert response.chunks[0].metadata["document_id"] == "doc1" assert response.chunks[0].metadata["source"] == "precomputed" + + +def test_auto_extract_embedding_dimension(client_with_empty_registry, embedding_model_id): + vs = client_with_empty_registry.vector_stores.create( + name="test_auto_extract", extra_body={"embedding_model": embedding_model_id} + ) + assert vs.id is not None + + +def test_provider_auto_selection_single_provider(client_with_empty_registry, embedding_model_id): + providers = [p for p in client_with_empty_registry.providers.list() if p.api == "vector_io"] + if len(providers) != 1: + pytest.skip(f"Test requires exactly one vector_io provider, found {len(providers)}") + + vs = client_with_empty_registry.vector_stores.create( + name="test_auto_provider", extra_body={"embedding_model": embedding_model_id} + ) + assert vs.id is not None + + +def test_provider_id_override(client_with_empty_registry, embedding_model_id): + providers = [p for p in client_with_empty_registry.providers.list() if p.api == "vector_io"] + if len(providers) != 1: + pytest.skip(f"Test requires exactly one vector_io provider, found {len(providers)}") + + provider_id = providers[0].provider_id + + vs = client_with_empty_registry.vector_stores.create( + name="test_provider_override", extra_body={"embedding_model": embedding_model_id, "provider_id": provider_id} + ) + assert vs.id is not None + assert vs.metadata.get("provider_id") == provider_id diff --git a/tests/unit/core/routers/test_vector_io.py b/tests/unit/core/routers/test_vector_io.py new file mode 100644 index 000000000..997df0d78 --- /dev/null +++ b/tests/unit/core/routers/test_vector_io.py @@ -0,0 +1,57 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +from unittest.mock import AsyncMock, Mock + +import pytest + +from llama_stack.apis.vector_io import OpenAICreateVectorStoreRequestWithExtraBody +from llama_stack.core.routers.vector_io import VectorIORouter + + +async def test_single_provider_auto_selection(): + # provider_id automatically selected during vector store create() when only one provider available + mock_routing_table = Mock() + mock_routing_table.impls_by_provider_id = {"inline::faiss": "mock_provider"} + mock_routing_table.get_all_with_type = AsyncMock( + return_value=[ + Mock(identifier="all-MiniLM-L6-v2", model_type="embedding", metadata={"embedding_dimension": 384}) + ] + ) + mock_routing_table.register_vector_db = AsyncMock( + return_value=Mock(identifier="vs_123", provider_id="inline::faiss", provider_resource_id="vs_123") + ) + mock_routing_table.get_provider_impl = AsyncMock( + return_value=Mock(openai_create_vector_store=AsyncMock(return_value=Mock(id="vs_123"))) + ) + router = VectorIORouter(mock_routing_table) + request = OpenAICreateVectorStoreRequestWithExtraBody.model_validate( + {"name": "test_store", "embedding_model": "all-MiniLM-L6-v2"} + ) + + result = await router.openai_create_vector_store(request) + assert result.id == "vs_123" + + +async def test_create_vector_stores_multiple_providers_missing_provider_id_error(): + # if multiple providers are available, vector store create will error without provider_id + mock_routing_table = Mock() + mock_routing_table.impls_by_provider_id = { + "inline::faiss": "mock_provider_1", + "inline::sqlite-vec": "mock_provider_2", + } + mock_routing_table.get_all_with_type = AsyncMock( + return_value=[ + Mock(identifier="all-MiniLM-L6-v2", model_type="embedding", metadata={"embedding_dimension": 384}) + ] + ) + router = VectorIORouter(mock_routing_table) + request = OpenAICreateVectorStoreRequestWithExtraBody.model_validate( + {"name": "test_store", "embedding_model": "all-MiniLM-L6-v2"} + ) + + with pytest.raises(ValueError, match="Multiple vector_io providers available"): + await router.openai_create_vector_store(request)