phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-12 12:06:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

chore: Auto-detect Provider ID when only 1 Vector Store Provider available and Require Embedding Model

Browse source 
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Javier Arceo 2025-10-13 10:09:04 -04:00
parent b95f095a54
commit 1bb3434562
3 changed files with 27 additions and 110 deletions
Download patch file Download diff file Expand all files Collapse all files

57
llama_stack/core/routers/vector_io.py
View file

@ -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(

70
tests/integration/vector_io/test_openai_vector_stores.py
View file

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

10
tests/integration/vector_io/test_vector_io.py
View file

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