mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-22 16:23:08 +00:00
chore: Updating how default embedding model is set in stack (#3818)
# What does this PR do?
Refactor setting default vector store provider and embedding model to
use an optional `vector_stores` config in the `StackRunConfig` and clean
up code to do so (had to add back in some pieces of VectorDB). Also
added remote Qdrant and Weaviate to starter distro (based on other PR
where inference providers were added for UX).
New config is simply (default for Starter distro):
```yaml
vector_stores:
default_provider_id: faiss
default_embedding_model:
provider_id: sentence-transformers
model_id: nomic-ai/nomic-embed-text-v1.5
```
## Test Plan
CI and Unit tests.
---------
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
Co-authored-by: Ashwin Bharambe <ashwin.bharambe@gmail.com>
This commit is contained in:
Francisco Arceo
GitHub
parent
2c43285e22
commit
48581bf651
48 changed files with 973 additions and 818 deletions
|
|
@ -317,3 +317,72 @@ def pytest_ignore_collect(path: str, config: pytest.Config) -> bool:
|
|||
if p.is_relative_to(rp):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_vector_io_provider_ids(client):
|
||||
"""Get all available vector_io provider IDs."""
|
||||
providers = [p for p in client.providers.list() if p.api == "vector_io"]
|
||||
return [p.provider_id for p in providers]
|
||||
|
||||
|
||||
def vector_provider_wrapper(func):
|
||||
"""Decorator to run a test against all available vector_io providers."""
|
||||
import functools
|
||||
import os
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
# Get the vector_io_provider_id from the test arguments
|
||||
import inspect
|
||||
|
||||
sig = inspect.signature(func)
|
||||
bound_args = sig.bind(*args, **kwargs)
|
||||
bound_args.apply_defaults()
|
||||
|
||||
vector_io_provider_id = bound_args.arguments.get("vector_io_provider_id")
|
||||
if not vector_io_provider_id:
|
||||
pytest.skip("No vector_io_provider_id provided")
|
||||
|
||||
# Get client_with_models to check available providers
|
||||
client_with_models = bound_args.arguments.get("client_with_models")
|
||||
if client_with_models:
|
||||
available_providers = get_vector_io_provider_ids(client_with_models)
|
||||
if vector_io_provider_id not in available_providers:
|
||||
pytest.skip(f"Provider '{vector_io_provider_id}' not available. Available: {available_providers}")
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
# For replay tests, only use providers that are available in ci-tests environment
|
||||
if os.environ.get("LLAMA_STACK_TEST_INFERENCE_MODE") == "replay":
|
||||
all_providers = ["faiss", "sqlite-vec"]
|
||||
else:
|
||||
# For live tests, try all providers (they'll skip if not available)
|
||||
all_providers = [
|
||||
"faiss",
|
||||
"sqlite-vec",
|
||||
"milvus",
|
||||
"chromadb",
|
||||
"pgvector",
|
||||
"weaviate",
|
||||
"qdrant",
|
||||
]
|
||||
|
||||
return pytest.mark.parametrize("vector_io_provider_id", all_providers)(wrapper)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def vector_io_provider_id(request, client_with_models):
|
||||
"""Fixture that provides a specific vector_io provider ID, skipping if not available."""
|
||||
if hasattr(request, "param"):
|
||||
requested_provider = request.param
|
||||
available_providers = get_vector_io_provider_ids(client_with_models)
|
||||
|
||||
if requested_provider not in available_providers:
|
||||
pytest.skip(f"Provider '{requested_provider}' not available. Available: {available_providers}")
|
||||
|
||||
return requested_provider
|
||||
else:
|
||||
provider_ids = get_vector_io_provider_ids(client_with_models)
|
||||
if not provider_ids:
|
||||
pytest.skip("No vector_io providers available")
|
||||
return provider_ids[0]
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from llama_stack_client import LlamaStackClient
|
|||
from openai import OpenAI
|
||||
|
||||
from llama_stack import LlamaStackAsLibraryClient
|
||||
from llama_stack.core.datatypes import VectorStoresConfig
|
||||
from llama_stack.core.stack import run_config_from_adhoc_config_spec
|
||||
from llama_stack.env import get_env_or_fail
|
||||
|
||||
|
|
@ -236,6 +237,13 @@ def instantiate_llama_stack_client(session):
|
|||
|
||||
if "=" in config:
|
||||
run_config = run_config_from_adhoc_config_spec(config)
|
||||
|
||||
# --stack-config bypasses template so need this to set default embedding model
|
||||
if "vector_io" in config and "inference" in config:
|
||||
run_config.vector_stores = VectorStoresConfig(
|
||||
embedding_model_id="inline::sentence-transformers/nomic-ai/nomic-embed-text-v1.5"
|
||||
)
|
||||
|
||||
run_config_file = tempfile.NamedTemporaryFile(delete=False, suffix=".yaml")
|
||||
with open(run_config_file.name, "w") as f:
|
||||
yaml.dump(run_config.model_dump(mode="json"), f)
|
||||
|
|
|
|||
|
|
@ -8,14 +8,15 @@ import time
|
|||
from io import BytesIO
|
||||
|
||||
import pytest
|
||||
from llama_stack_client import BadRequestError, NotFoundError
|
||||
from llama_stack_client import BadRequestError
|
||||
from openai import BadRequestError as OpenAIBadRequestError
|
||||
from openai import NotFoundError as OpenAINotFoundError
|
||||
|
||||
from llama_stack.apis.vector_io import Chunk
|
||||
from llama_stack.core.library_client import LlamaStackAsLibraryClient
|
||||
from llama_stack.log import get_logger
|
||||
|
||||
from ..conftest import vector_provider_wrapper
|
||||
|
||||
logger = get_logger(name=__name__, category="vector_io")
|
||||
|
||||
|
||||
|
|
@ -133,8 +134,9 @@ def compat_client_with_empty_stores(compat_client):
|
|||
clear_files()
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_create_vector_store(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test creating a vector store using OpenAI API."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -146,6 +148,7 @@ def test_openai_create_vector_store(
|
|||
metadata={"purpose": "testing", "environment": "integration"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -159,14 +162,18 @@ def test_openai_create_vector_store(
|
|||
assert hasattr(vector_store, "created_at")
|
||||
|
||||
|
||||
def test_openai_create_vector_store_default(compat_client_with_empty_stores, client_with_models):
|
||||
@vector_provider_wrapper
|
||||
def test_openai_create_vector_store_default(compat_client_with_empty_stores, client_with_models, vector_io_provider_id):
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
vector_store = compat_client_with_empty_stores.vector_stores.create()
|
||||
vector_store = compat_client_with_empty_stores.vector_stores.create(
|
||||
extra_body={"provider_id": vector_io_provider_id}
|
||||
)
|
||||
assert vector_store.id
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_list_vector_stores(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test listing vector stores using OpenAI API."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -179,6 +186,7 @@ def test_openai_list_vector_stores(
|
|||
metadata={"type": "test"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
store2 = client.vector_stores.create(
|
||||
|
|
@ -186,6 +194,7 @@ def test_openai_list_vector_stores(
|
|||
metadata={"type": "test"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -206,8 +215,9 @@ def test_openai_list_vector_stores(
|
|||
assert len(limited_response.data) == 1
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_retrieve_vector_store(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test retrieving a specific vector store using OpenAI API."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -220,6 +230,7 @@ def test_openai_retrieve_vector_store(
|
|||
metadata={"purpose": "retrieval_test"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -233,8 +244,9 @@ def test_openai_retrieve_vector_store(
|
|||
assert retrieved_store.object == "vector_store"
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_update_vector_store(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test modifying a vector store using OpenAI API."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -247,6 +259,7 @@ def test_openai_update_vector_store(
|
|||
metadata={"version": "1.0"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
time.sleep(1)
|
||||
|
|
@ -264,8 +277,9 @@ def test_openai_update_vector_store(
|
|||
assert modified_store.last_active_at > created_store.last_active_at
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_delete_vector_store(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test deleting a vector store using OpenAI API."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -278,6 +292,7 @@ def test_openai_delete_vector_store(
|
|||
metadata={"purpose": "deletion_test"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -294,8 +309,9 @@ def test_openai_delete_vector_store(
|
|||
client.vector_stores.retrieve(vector_store_id=created_store.id)
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_search_empty(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test searching an empty vector store using OpenAI API."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -308,6 +324,7 @@ def test_openai_vector_store_search_empty(
|
|||
metadata={"purpose": "search_testing"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -323,8 +340,14 @@ def test_openai_vector_store_search_empty(
|
|||
assert search_response.has_more is False
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_with_chunks(
|
||||
compat_client_with_empty_stores, client_with_models, sample_chunks, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores,
|
||||
client_with_models,
|
||||
sample_chunks,
|
||||
embedding_model_id,
|
||||
embedding_dimension,
|
||||
vector_io_provider_id,
|
||||
):
|
||||
"""Test vector store functionality with actual chunks using both OpenAI and native APIs."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -338,6 +361,7 @@ def test_openai_vector_store_with_chunks(
|
|||
metadata={"purpose": "chunks_testing"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -380,6 +404,7 @@ def test_openai_vector_store_with_chunks(
|
|||
("What inspires neural networks?", "doc4", "ai"),
|
||||
],
|
||||
)
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_search_relevance(
|
||||
compat_client_with_empty_stores,
|
||||
client_with_models,
|
||||
|
|
@ -387,6 +412,7 @@ def test_openai_vector_store_search_relevance(
|
|||
test_case,
|
||||
embedding_model_id,
|
||||
embedding_dimension,
|
||||
vector_io_provider_id,
|
||||
):
|
||||
"""Test that OpenAI vector store search returns relevant results for different queries."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -402,6 +428,7 @@ def test_openai_vector_store_search_relevance(
|
|||
metadata={"purpose": "relevance_testing"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -430,8 +457,14 @@ def test_openai_vector_store_search_relevance(
|
|||
assert top_result.score > 0
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_search_with_ranking_options(
|
||||
compat_client_with_empty_stores, client_with_models, sample_chunks, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores,
|
||||
client_with_models,
|
||||
sample_chunks,
|
||||
embedding_model_id,
|
||||
embedding_dimension,
|
||||
vector_io_provider_id,
|
||||
):
|
||||
"""Test OpenAI vector store search with ranking options."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -445,6 +478,7 @@ def test_openai_vector_store_search_with_ranking_options(
|
|||
metadata={"purpose": "ranking_testing"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -483,8 +517,14 @@ def test_openai_vector_store_search_with_ranking_options(
|
|||
assert result.score >= threshold
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_search_with_high_score_filter(
|
||||
compat_client_with_empty_stores, client_with_models, sample_chunks, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores,
|
||||
client_with_models,
|
||||
sample_chunks,
|
||||
embedding_model_id,
|
||||
embedding_dimension,
|
||||
vector_io_provider_id,
|
||||
):
|
||||
"""Test that searching with text very similar to a document and high score threshold returns only that document."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -498,6 +538,7 @@ def test_openai_vector_store_search_with_high_score_filter(
|
|||
metadata={"purpose": "high_score_filtering"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -542,8 +583,14 @@ def test_openai_vector_store_search_with_high_score_filter(
|
|||
assert "python" in top_content.lower() or "programming" in top_content.lower()
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_search_with_max_num_results(
|
||||
compat_client_with_empty_stores, client_with_models, sample_chunks, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores,
|
||||
client_with_models,
|
||||
sample_chunks,
|
||||
embedding_model_id,
|
||||
embedding_dimension,
|
||||
vector_io_provider_id,
|
||||
):
|
||||
"""Test OpenAI vector store search with max_num_results."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -557,6 +604,7 @@ def test_openai_vector_store_search_with_max_num_results(
|
|||
metadata={"purpose": "max_num_results_testing"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -577,8 +625,9 @@ def test_openai_vector_store_search_with_max_num_results(
|
|||
assert len(search_response.data) == 2
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_attach_file(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test OpenAI vector store attach file."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -591,6 +640,7 @@ def test_openai_vector_store_attach_file(
|
|||
name="test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -637,8 +687,9 @@ def test_openai_vector_store_attach_file(
|
|||
assert "foobazbar" in top_content.lower()
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_attach_files_on_creation(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test OpenAI vector store attach files on creation."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -668,6 +719,7 @@ def test_openai_vector_store_attach_files_on_creation(
|
|||
file_ids=file_ids,
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -700,8 +752,9 @@ def test_openai_vector_store_attach_files_on_creation(
|
|||
assert updated_vector_store.file_counts.failed == 0
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_list_files(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test OpenAI vector store list files."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -714,6 +767,7 @@ def test_openai_vector_store_list_files(
|
|||
name="test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -773,8 +827,9 @@ def test_openai_vector_store_list_files(
|
|||
assert updated_vector_store.file_counts.in_progress == 0
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_list_files_invalid_vector_store(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test OpenAI vector store list files with invalid vector store ID."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -783,14 +838,15 @@ def test_openai_vector_store_list_files_invalid_vector_store(
|
|||
if isinstance(compat_client, LlamaStackAsLibraryClient):
|
||||
errors = ValueError
|
||||
else:
|
||||
errors = (NotFoundError, OpenAINotFoundError)
|
||||
errors = (BadRequestError, OpenAIBadRequestError)
|
||||
|
||||
with pytest.raises(errors):
|
||||
compat_client.vector_stores.files.list(vector_store_id="abc123")
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_retrieve_file_contents(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test OpenAI vector store retrieve file contents."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -803,6 +859,7 @@ def test_openai_vector_store_retrieve_file_contents(
|
|||
name="test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -848,8 +905,9 @@ def test_openai_vector_store_retrieve_file_contents(
|
|||
assert file_contents.attributes == attributes
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_delete_file(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test OpenAI vector store delete file."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -862,6 +920,7 @@ def test_openai_vector_store_delete_file(
|
|||
name="test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -912,8 +971,9 @@ def test_openai_vector_store_delete_file(
|
|||
assert updated_vector_store.file_counts.in_progress == 0
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_delete_file_removes_from_vector_store(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test OpenAI vector store delete file removes from vector store."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -926,6 +986,7 @@ def test_openai_vector_store_delete_file_removes_from_vector_store(
|
|||
name="test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -962,8 +1023,9 @@ def test_openai_vector_store_delete_file_removes_from_vector_store(
|
|||
assert not search_response.data
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_update_file(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test OpenAI vector store update file."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -976,6 +1038,7 @@ def test_openai_vector_store_update_file(
|
|||
name="test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1017,8 +1080,9 @@ def test_openai_vector_store_update_file(
|
|||
assert retrieved_file.attributes["foo"] == "baz"
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_create_vector_store_files_duplicate_vector_store_name(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""
|
||||
This test confirms that client.vector_stores.create() creates a unique ID
|
||||
|
|
@ -1044,6 +1108,7 @@ def test_create_vector_store_files_duplicate_vector_store_name(
|
|||
name="test_store_with_files",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
assert vector_store.file_counts.completed == 0
|
||||
|
|
@ -1056,6 +1121,7 @@ def test_create_vector_store_files_duplicate_vector_store_name(
|
|||
name="test_store_with_files",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1086,8 +1152,15 @@ def test_create_vector_store_files_duplicate_vector_store_name(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("search_mode", ["vector", "keyword", "hybrid"])
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_search_modes(
|
||||
llama_stack_client, client_with_models, sample_chunks, search_mode, embedding_model_id, embedding_dimension
|
||||
llama_stack_client,
|
||||
client_with_models,
|
||||
sample_chunks,
|
||||
search_mode,
|
||||
embedding_model_id,
|
||||
embedding_dimension,
|
||||
vector_io_provider_id,
|
||||
):
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
skip_if_provider_doesnt_support_openai_vector_stores_search(client_with_models, search_mode)
|
||||
|
|
@ -1097,6 +1170,7 @@ def test_openai_vector_store_search_modes(
|
|||
metadata={"purpose": "search_mode_testing"},
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1115,8 +1189,9 @@ def test_openai_vector_store_search_modes(
|
|||
assert search_response is not None
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_file_batch_create_and_retrieve(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test creating and retrieving a vector store file batch."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -1128,6 +1203,7 @@ def test_openai_vector_store_file_batch_create_and_retrieve(
|
|||
name="batch_test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1178,8 +1254,9 @@ def test_openai_vector_store_file_batch_create_and_retrieve(
|
|||
assert retrieved_batch.status == "completed" # Should be completed after processing
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_file_batch_list_files(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test listing files in a vector store file batch."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -1191,6 +1268,7 @@ def test_openai_vector_store_file_batch_list_files(
|
|||
name="batch_list_test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1271,8 +1349,9 @@ def test_openai_vector_store_file_batch_list_files(
|
|||
assert first_page_ids.isdisjoint(second_page_ids)
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_file_batch_cancel(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test cancelling a vector store file batch."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -1284,6 +1363,7 @@ def test_openai_vector_store_file_batch_cancel(
|
|||
name="batch_cancel_test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1326,8 +1406,9 @@ def test_openai_vector_store_file_batch_cancel(
|
|||
assert final_batch.status in ["completed", "cancelled"]
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_file_batch_retrieve_contents(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test retrieving file contents after file batch processing."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -1339,6 +1420,7 @@ def test_openai_vector_store_file_batch_retrieve_contents(
|
|||
name="batch_contents_test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1399,8 +1481,9 @@ def test_openai_vector_store_file_batch_retrieve_contents(
|
|||
assert file_data[i][1].decode("utf-8") in content_text
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_file_batch_error_handling(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test error handling for file batch operations."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -1412,6 +1495,7 @@ def test_openai_vector_store_file_batch_error_handling(
|
|||
name="batch_error_test_store",
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1443,11 +1527,11 @@ def test_openai_vector_store_file_batch_error_handling(
|
|||
batch_id="non_existent_batch_id",
|
||||
)
|
||||
|
||||
# Test operations on non-existent vector store (returns NotFoundError)
|
||||
# Test operations on non-existent vector store (returns BadRequestError)
|
||||
if isinstance(compat_client, LlamaStackAsLibraryClient):
|
||||
vector_store_errors = ValueError
|
||||
else:
|
||||
vector_store_errors = (NotFoundError, OpenAINotFoundError)
|
||||
vector_store_errors = (BadRequestError, OpenAIBadRequestError)
|
||||
|
||||
with pytest.raises(vector_store_errors): # Should raise an error for non-existent vector store
|
||||
compat_client.vector_stores.file_batches.create(
|
||||
|
|
@ -1456,8 +1540,9 @@ def test_openai_vector_store_file_batch_error_handling(
|
|||
)
|
||||
|
||||
|
||||
@vector_provider_wrapper
|
||||
def test_openai_vector_store_embedding_config_from_metadata(
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension
|
||||
compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
"""Test that embedding configuration works from metadata source."""
|
||||
skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
|
||||
|
|
@ -1471,6 +1556,9 @@ def test_openai_vector_store_embedding_config_from_metadata(
|
|||
"embedding_dimension": str(embedding_dimension),
|
||||
"test_source": "metadata",
|
||||
},
|
||||
extra_body={
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
assert vector_store_metadata is not None
|
||||
|
|
@ -1489,6 +1577,7 @@ def test_openai_vector_store_embedding_config_from_metadata(
|
|||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"embedding_dimension": int(embedding_dimension), # Ensure same type/value
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ import pytest
|
|||
|
||||
from llama_stack.apis.vector_io import Chunk
|
||||
|
||||
from ..conftest import vector_provider_wrapper
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def sample_chunks():
|
||||
|
|
@ -46,12 +48,13 @@ def client_with_empty_registry(client_with_models):
|
|||
clear_registry()
|
||||
|
||||
|
||||
def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id, embedding_dimension):
|
||||
@vector_provider_wrapper
|
||||
def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id):
|
||||
vector_db_name = "test_vector_db"
|
||||
create_response = client_with_empty_registry.vector_stores.create(
|
||||
name=vector_db_name,
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -65,12 +68,13 @@ def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id, embe
|
|||
assert response.id.startswith("vs_")
|
||||
|
||||
|
||||
def test_vector_db_register(client_with_empty_registry, embedding_model_id, embedding_dimension):
|
||||
@vector_provider_wrapper
|
||||
def test_vector_db_register(client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id):
|
||||
vector_db_name = "test_vector_db"
|
||||
response = client_with_empty_registry.vector_stores.create(
|
||||
name=vector_db_name,
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -100,12 +104,15 @@ def test_vector_db_register(client_with_empty_registry, embedding_model_id, embe
|
|||
("How does machine learning improve over time?", "doc2"),
|
||||
],
|
||||
)
|
||||
def test_insert_chunks(client_with_empty_registry, embedding_model_id, embedding_dimension, sample_chunks, test_case):
|
||||
@vector_provider_wrapper
|
||||
def test_insert_chunks(
|
||||
client_with_empty_registry, embedding_model_id, embedding_dimension, sample_chunks, test_case, vector_io_provider_id
|
||||
):
|
||||
vector_db_name = "test_vector_db"
|
||||
create_response = client_with_empty_registry.vector_stores.create(
|
||||
name=vector_db_name,
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -135,7 +142,10 @@ def test_insert_chunks(client_with_empty_registry, embedding_model_id, embedding
|
|||
assert top_match.metadata["document_id"] == expected_doc_id, f"Query '{query}' should match {expected_doc_id}"
|
||||
|
||||
|
||||
def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, embedding_model_id, embedding_dimension):
|
||||
@vector_provider_wrapper
|
||||
def test_insert_chunks_with_precomputed_embeddings(
|
||||
client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
vector_io_provider_params_dict = {
|
||||
"inline::milvus": {"score_threshold": -1.0},
|
||||
"inline::qdrant": {"score_threshold": -1.0},
|
||||
|
|
@ -145,7 +155,7 @@ def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, e
|
|||
register_response = client_with_empty_registry.vector_stores.create(
|
||||
name=vector_db_name,
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -181,8 +191,9 @@ def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, e
|
|||
|
||||
|
||||
# expect this test to fail
|
||||
@vector_provider_wrapper
|
||||
def test_query_returns_valid_object_when_identical_to_embedding_in_vdb(
|
||||
client_with_empty_registry, embedding_model_id, embedding_dimension
|
||||
client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
vector_io_provider_params_dict = {
|
||||
"inline::milvus": {"score_threshold": 0.0},
|
||||
|
|
@ -194,6 +205,7 @@ def test_query_returns_valid_object_when_identical_to_embedding_in_vdb(
|
|||
name=vector_db_name,
|
||||
extra_body={
|
||||
"embedding_model": embedding_model_id,
|
||||
"provider_id": vector_io_provider_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -226,33 +238,44 @@ def test_query_returns_valid_object_when_identical_to_embedding_in_vdb(
|
|||
assert response.chunks[0].metadata["source"] == "precomputed"
|
||||
|
||||
|
||||
def test_auto_extract_embedding_dimension(client_with_empty_registry, embedding_model_id):
|
||||
@vector_provider_wrapper
|
||||
def test_auto_extract_embedding_dimension(
|
||||
client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id
|
||||
):
|
||||
# This test specifically tests embedding model override, so we keep embedding_model
|
||||
vs = client_with_empty_registry.vector_stores.create(
|
||||
name="test_auto_extract", extra_body={"embedding_model": embedding_model_id}
|
||||
name="test_auto_extract",
|
||||
extra_body={"embedding_model": embedding_model_id, "provider_id": vector_io_provider_id},
|
||||
)
|
||||
assert vs.id is not None
|
||||
|
||||
|
||||
def test_provider_auto_selection_single_provider(client_with_empty_registry, embedding_model_id):
|
||||
@vector_provider_wrapper
|
||||
def test_provider_auto_selection_single_provider(
|
||||
client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_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}
|
||||
)
|
||||
# Test that when only one provider is available, it's auto-selected (no provider_id needed)
|
||||
vs = client_with_empty_registry.vector_stores.create(name="test_auto_provider")
|
||||
assert vs.id is not None
|
||||
|
||||
|
||||
def test_provider_id_override(client_with_empty_registry, embedding_model_id):
|
||||
@vector_provider_wrapper
|
||||
def test_provider_id_override(
|
||||
client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_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
|
||||
|
||||
# Test explicit provider_id specification (using default embedding model)
|
||||
vs = client_with_empty_registry.vector_stores.create(
|
||||
name="test_provider_override", extra_body={"embedding_model": embedding_model_id, "provider_id": provider_id}
|
||||
name="test_provider_override", extra_body={"provider_id": provider_id}
|
||||
)
|
||||
assert vs.id is not None
|
||||
assert vs.metadata.get("provider_id") == provider_id
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue