From 1bb3434562a608485abd7596e76b03e16576aabf Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Mon, 13 Oct 2025 10:09:04 -0400 Subject: [PATCH] chore: Auto-detect Provider ID when only 1 Vector Store Provider available and Require Embedding Model Signed-off-by: Francisco Javier Arceo --- llama_stack/core/routers/vector_io.py | 57 +++++++-------- .../vector_io/test_openai_vector_stores.py | 70 ------------------- tests/integration/vector_io/test_vector_io.py | 10 --- 3 files changed, 27 insertions(+), 110 deletions(-) diff --git a/llama_stack/core/routers/vector_io.py b/llama_stack/core/routers/vector_io.py index 79789ef0a..8e00873e8 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, @@ -134,15 +122,24 @@ class VectorIORouter(VectorIO): 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") + + # Always extract embedding dimension from the model registry + embedding_dimension = await self._get_embedding_model_dimension(embedding_model) + + # Auto-select provider if not specified + if provider_id is None: + if len(self.routing_table.impls_by_provider_id) == 1: + provider_id = list(self.routing_table.impls_by_provider_id.keys())[0] + logger.info(f"No provider_id specified, using the only available vector_io provider: {provider_id}") + else: + 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}" + ) vector_db_id = f"vs_{uuid.uuid4()}" registered_vector_db = await self.routing_table.register_vector_db( 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..53d46e28c 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", }, )