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

updated vector io tests

Browse source 
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Javier Arceo 2025-10-10 16:49:11 -04:00
parent b94b701317
commit c730078b27
1 changed files with 46 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -34,9 +34,9 @@ def sample_chunks():
@pytest.fixture(scope="function")
def client_with_empty_registry(client_with_models):
def clear_registry():
vector_dbs = [vector_db.identifier for vector_db in client_with_models.vector_dbs.list()]
for vector_db_id in vector_dbs:
client_with_models.vector_dbs.unregister(vector_db_id=vector_db_id)
vector_stores = client_with_models.vector_stores.list()
for vector_store in vector_stores.data:
client_with_models.vector_stores.delete(vector_store_id=vector_store.id)
clear_registry()
yield client_with_models
@ -48,47 +48,48 @@ def client_with_empty_registry(client_with_models):
def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id, embedding_dimension):
vector_db_name = "test_vector_db"
register_response = client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_name,
embedding_model=embedding_model_id,
embedding_dimension=embedding_dimension,
create_response = client_with_empty_registry.vector_stores.create(
name=vector_db_name,
extra_body={
"embedding_model": embedding_model_id,
"embedding_dimension": embedding_dimension,
},
)
actual_vector_db_id = register_response.identifier
actual_vector_db_id = create_response.id
# Retrieve the memory bank and validate its properties
response = client_with_empty_registry.vector_dbs.retrieve(vector_db_id=actual_vector_db_id)
# Retrieve the vector store and validate its properties
response = client_with_empty_registry.vector_stores.retrieve(vector_store_id=actual_vector_db_id)
assert response is not None
assert response.identifier == actual_vector_db_id
assert response.embedding_model == embedding_model_id
assert response.identifier.startswith("vs_")
assert response.id == actual_vector_db_id
assert response.name == vector_db_name
assert response.id.startswith("vs_")
def test_vector_db_register(client_with_empty_registry, embedding_model_id, embedding_dimension):
vector_db_name = "test_vector_db"
response = client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_name,
embedding_model=embedding_model_id,
embedding_dimension=embedding_dimension,
response = client_with_empty_registry.vector_stores.create(
name=vector_db_name,
extra_body={
"embedding_model": embedding_model_id,
"embedding_dimension": embedding_dimension,
},
)
actual_vector_db_id = response.identifier
actual_vector_db_id = response.id
assert actual_vector_db_id.startswith("vs_")
assert actual_vector_db_id != vector_db_name
vector_dbs_after_register = [vector_db.identifier for vector_db in client_with_empty_registry.vector_dbs.list()]
assert vector_dbs_after_register == [actual_vector_db_id]
vector_stores = client_with_empty_registry.vector_stores.list()
assert len(vector_stores.data) == 1
vector_store = vector_stores.data[0]
assert vector_store.id == actual_vector_db_id
assert vector_store.name == vector_db_name
client_with_empty_registry.vector_dbs.unregister(vector_db_id=actual_vector_db_id)
client_with_empty_registry.vector_stores.delete(vector_store_id=actual_vector_db_id)
vector_dbs = [vector_db.identifier for vector_db in client_with_empty_registry.vector_dbs.list()]
assert len(vector_dbs) == 0
vector_stores = client_with_empty_registry.vector_stores.list()
assert len(vector_stores.data) == 0
@pytest.mark.parametrize(
@ -103,13 +104,15 @@ def test_vector_db_register(client_with_empty_registry, embedding_model_id, embe
)
def test_insert_chunks(client_with_empty_registry, embedding_model_id, embedding_dimension, sample_chunks, test_case):
vector_db_name = "test_vector_db"
register_response = client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_name,
embedding_model=embedding_model_id,
embedding_dimension=embedding_dimension,
create_response = client_with_empty_registry.vector_stores.create(
name=vector_db_name,
extra_body={
"embedding_model": embedding_model_id,
"embedding_dimension": embedding_dimension,
},
)
actual_vector_db_id = register_response.identifier
actual_vector_db_id = create_response.id
client_with_empty_registry.vector_io.insert(
vector_db_id=actual_vector_db_id,
@ -142,13 +145,15 @@ def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, e
"remote::qdrant": {"score_threshold": -1.0},
}
vector_db_name = "test_precomputed_embeddings_db"
register_response = client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_name,
embedding_model=embedding_model_id,
embedding_dimension=embedding_dimension,
register_response = client_with_empty_registry.vector_stores.create(
name=vector_db_name,
extra_body={
"embedding_model": embedding_model_id,
"embedding_dimension": embedding_dimension,
},
)
actual_vector_db_id = register_response.identifier
actual_vector_db_id = register_response.id
chunks_with_embeddings = [
Chunk(
@ -189,13 +194,15 @@ def test_query_returns_valid_object_when_identical_to_embedding_in_vdb(
"inline::qdrant": {"score_threshold": 0.0},
}
vector_db_name = "test_precomputed_embeddings_db"
register_response = client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_name,
embedding_model=embedding_model_id,
embedding_dimension=embedding_dimension,
register_response = client_with_empty_registry.vector_stores.create(
name=vector_db_name,
extra_body={
"embedding_model": embedding_model_id,
"embedding_dimension": embedding_dimension,
},
)
actual_vector_db_id = register_response.identifier
actual_vector_db_id = register_response.id
chunks_with_embeddings = [
Chunk(