mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-15 06:37:58 +00:00
chore: Auto-detect Provider ID when only 1 Vector Store Provider avai… (#3802)
Some checks failed
SqlStore Integration Tests / test-postgres (3.12) (push) Failing after 0s
SqlStore Integration Tests / test-postgres (3.13) (push) Failing after 0s
Integration Auth Tests / test-matrix (oauth2_token) (push) Failing after 1s
Test External Providers Installed via Module / test-external-providers-from-module (venv) (push) Has been skipped
Python Package Build Test / build (3.13) (push) Failing after 1s
Python Package Build Test / build (3.12) (push) Failing after 1s
Integration Tests (Replay) / Integration Tests (, , , client=, ) (push) Failing after 4s
Test External API and Providers / test-external (venv) (push) Failing after 4s
Unit Tests / unit-tests (3.12) (push) Failing after 4s
Unit Tests / unit-tests (3.13) (push) Failing after 4s
Vector IO Integration Tests / test-matrix (push) Failing after 8s
API Conformance Tests / check-schema-compatibility (push) Successful in 18s
UI Tests / ui-tests (22) (push) Successful in 29s
Pre-commit / pre-commit (push) Successful in 1m24s
Some checks failed
SqlStore Integration Tests / test-postgres (3.12) (push) Failing after 0s
SqlStore Integration Tests / test-postgres (3.13) (push) Failing after 0s
Integration Auth Tests / test-matrix (oauth2_token) (push) Failing after 1s
Test External Providers Installed via Module / test-external-providers-from-module (venv) (push) Has been skipped
Python Package Build Test / build (3.13) (push) Failing after 1s
Python Package Build Test / build (3.12) (push) Failing after 1s
Integration Tests (Replay) / Integration Tests (, , , client=, ) (push) Failing after 4s
Test External API and Providers / test-external (venv) (push) Failing after 4s
Unit Tests / unit-tests (3.12) (push) Failing after 4s
Unit Tests / unit-tests (3.13) (push) Failing after 4s
Vector IO Integration Tests / test-matrix (push) Failing after 8s
API Conformance Tests / check-schema-compatibility (push) Successful in 18s
UI Tests / ui-tests (22) (push) Successful in 29s
Pre-commit / pre-commit (push) Successful in 1m24s
# What does this PR do? 2 main changes: 1. Remove `provider_id` requirement in call to vector stores and 2. Removes "register first embedding model" logic - Now forces embedding model id as required on Vector Store creation Simplifies the UX for OpenAI to: ```python vs = client.vector_stores.create( name="my_citations_db", extra_body={ "embedding_model": "ollama/nomic-embed-text:latest", } ) ``` <!-- If resolving an issue, uncomment and update the line below --> <!-- Closes #[issue-number] --> ## Test Plan <!-- Describe the tests you ran to verify your changes with result summaries. *Provide clear instructions so the plan can be easily re-executed.* --> --------- Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
parent
642126e13b
commit
968c364a3e
5 changed files with 123 additions and 115 deletions
|
@ -55,30 +55,18 @@ class VectorIORouter(VectorIO):
|
||||||
logger.debug("VectorIORouter.shutdown")
|
logger.debug("VectorIORouter.shutdown")
|
||||||
pass
|
pass
|
||||||
|
|
||||||
async def _get_first_embedding_model(self) -> tuple[str, int] | None:
|
async def _get_embedding_model_dimension(self, embedding_model_id: str) -> int:
|
||||||
"""Get the first available embedding model identifier."""
|
"""Get the embedding dimension for a specific embedding model."""
|
||||||
try:
|
all_models = await self.routing_table.get_all_with_type("model")
|
||||||
# Get all models from the routing table
|
|
||||||
all_models = await self.routing_table.get_all_with_type("model")
|
|
||||||
|
|
||||||
# Filter for embedding models
|
for model in all_models:
|
||||||
embedding_models = [
|
if model.identifier == embedding_model_id and model.model_type == ModelType.embedding:
|
||||||
model
|
dimension = model.metadata.get("embedding_dimension")
|
||||||
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)
|
|
||||||
if dimension is None:
|
if dimension is None:
|
||||||
raise ValueError(f"Embedding model {embedding_models[0].identifier} has no embedding dimension")
|
raise ValueError(f"Embedding model '{embedding_model_id}' has no embedding_dimension in metadata")
|
||||||
return embedding_models[0].identifier, dimension
|
return int(dimension)
|
||||||
else:
|
|
||||||
logger.warning("No embedding models found in the routing table")
|
raise ValueError(f"Embedding model '{embedding_model_id}' not found or not an embedding model")
|
||||||
return None
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Error getting embedding models: {e}")
|
|
||||||
return None
|
|
||||||
|
|
||||||
async def register_vector_db(
|
async def register_vector_db(
|
||||||
self,
|
self,
|
||||||
|
@ -129,20 +117,30 @@ class VectorIORouter(VectorIO):
|
||||||
# Extract llama-stack-specific parameters from extra_body
|
# Extract llama-stack-specific parameters from extra_body
|
||||||
extra = params.model_extra or {}
|
extra = params.model_extra or {}
|
||||||
embedding_model = extra.get("embedding_model")
|
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")
|
provider_id = extra.get("provider_id")
|
||||||
|
|
||||||
logger.debug(f"VectorIORouter.openai_create_vector_store: name={params.name}, provider_id={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
|
# Require explicit embedding model specification
|
||||||
# TODO: this branch will soon be deleted so you _must_ provide the embedding_model when
|
|
||||||
# creating a vector store
|
|
||||||
if embedding_model is None:
|
if embedding_model is None:
|
||||||
embedding_model_info = await self._get_first_embedding_model()
|
raise ValueError("embedding_model is required in extra_body when creating a vector store")
|
||||||
if embedding_model_info is None:
|
|
||||||
raise ValueError("No embedding model provided and no embedding models available in the system")
|
if embedding_dimension is None:
|
||||||
embedding_model, embedding_dimension = embedding_model_info
|
embedding_dimension = await self._get_embedding_model_dimension(embedding_model)
|
||||||
logger.info(f"No embedding model specified, using first available: {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()}"
|
vector_db_id = f"vs_{uuid.uuid4()}"
|
||||||
registered_vector_db = await self.routing_table.register_vector_db(
|
registered_vector_db = await self.routing_table.register_vector_db(
|
||||||
|
|
|
@ -353,14 +353,12 @@ class OpenAIVectorStoreMixin(ABC):
|
||||||
provider_vector_db_id = extra.get("provider_vector_db_id")
|
provider_vector_db_id = extra.get("provider_vector_db_id")
|
||||||
embedding_model = extra.get("embedding_model")
|
embedding_model = extra.get("embedding_model")
|
||||||
embedding_dimension = extra.get("embedding_dimension", 384)
|
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)
|
# 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()}")
|
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:
|
if embedding_model is None:
|
||||||
raise ValueError("Embedding model is required")
|
raise ValueError("Embedding model is required")
|
||||||
|
|
||||||
|
@ -369,6 +367,9 @@ class OpenAIVectorStoreMixin(ABC):
|
||||||
raise ValueError("Embedding dimension is required")
|
raise ValueError("Embedding dimension is required")
|
||||||
|
|
||||||
# Register the VectorDB backing this vector store
|
# 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(
|
vector_db = VectorDB(
|
||||||
identifier=vector_db_id,
|
identifier=vector_db_id,
|
||||||
embedding_dimension=embedding_dimension,
|
embedding_dimension=embedding_dimension,
|
||||||
|
|
|
@ -146,8 +146,6 @@ def test_openai_create_vector_store(
|
||||||
metadata={"purpose": "testing", "environment": "integration"},
|
metadata={"purpose": "testing", "environment": "integration"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"type": "test"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"embedding_model": embedding_model_id,
|
||||||
"embedding_dimension": embedding_dimension,
|
|
||||||
"provider_id": "my_provider",
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
store2 = client.vector_stores.create(
|
store2 = client.vector_stores.create(
|
||||||
|
@ -184,8 +180,6 @@ def test_openai_list_vector_stores(
|
||||||
metadata={"type": "test"},
|
metadata={"type": "test"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"purpose": "retrieval_test"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"version": "1.0"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"embedding_model": embedding_model_id,
|
||||||
"embedding_dimension": embedding_dimension,
|
|
||||||
"provider_id": "my_provider",
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
time.sleep(1)
|
time.sleep(1)
|
||||||
|
@ -282,8 +272,6 @@ def test_openai_delete_vector_store(
|
||||||
metadata={"purpose": "deletion_test"},
|
metadata={"purpose": "deletion_test"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"purpose": "search_testing"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"purpose": "chunks_testing"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"purpose": "relevance_testing"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"purpose": "ranking_testing"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"purpose": "high_score_filtering"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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"},
|
metadata={"purpose": "max_num_results_testing"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
file_ids=file_ids,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
attributes=attributes,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="test_store_with_files",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"embedding_model": embedding_model_id,
|
||||||
"embedding_dimension": embedding_dimension,
|
|
||||||
"provider_id": "my_provider",
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
assert vector_store.file_counts.completed == 0
|
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",
|
name="test_store_with_files",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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],
|
file_id=file_ids[0],
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"embedding_model": embedding_model_id,
|
||||||
"embedding_dimension": embedding_dimension,
|
|
||||||
"provider_id": "my_provider",
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
assert created_file.status == "completed"
|
assert created_file.status == "completed"
|
||||||
|
@ -1065,8 +1019,6 @@ def test_create_vector_store_files_duplicate_vector_store_name(
|
||||||
file_id=file_ids[1],
|
file_id=file_ids[1],
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"embedding_model": embedding_model_id,
|
||||||
"embedding_dimension": embedding_dimension,
|
|
||||||
"provider_id": "my_provider",
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
assert created_file_from_non_deleted_vector_store.status == "completed"
|
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"},
|
metadata={"purpose": "search_mode_testing"},
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="batch_test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
file_ids=file_ids,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="batch_list_test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
file_ids=file_ids,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="batch_cancel_test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
file_ids=file_ids,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="batch_contents_test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
file_ids=file_ids,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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",
|
name="batch_error_test_store",
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
file_ids=file_ids,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"embedding_model": embedding_model_id,
|
||||||
"embedding_dimension": embedding_dimension,
|
|
||||||
"provider_id": "my_provider",
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
@ -52,8 +52,6 @@ def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id, embe
|
||||||
name=vector_db_name,
|
name=vector_db_name,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
name=vector_db_name,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
name=vector_db_name,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
name=vector_db_name,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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,
|
name=vector_db_name,
|
||||||
extra_body={
|
extra_body={
|
||||||
"embedding_model": embedding_model_id,
|
"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 len(response.chunks) > 0
|
||||||
assert response.chunks[0].metadata["document_id"] == "doc1"
|
assert response.chunks[0].metadata["document_id"] == "doc1"
|
||||||
assert response.chunks[0].metadata["source"] == "precomputed"
|
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
|
||||||
|
|
57
tests/unit/core/routers/test_vector_io.py
Normal file
57
tests/unit/core/routers/test_vector_io.py
Normal file
|
@ -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)
|
Loading…
Add table
Add a link
Reference in a new issue