fix: allow running vector tests with embedding dimension (#2467)
Some checks failed
Integration Auth Tests / test-matrix (oauth2_token) (push) Failing after 8s
Integration Tests / test-matrix (http, 3.11, datasets) (push) Failing after 8s
Integration Tests / test-matrix (http, 3.11, providers) (push) Failing after 6s
Integration Tests / test-matrix (http, 3.11, agents) (push) Failing after 9s
Integration Tests / test-matrix (http, 3.11, post_training) (push) Failing after 9s
Integration Tests / test-matrix (http, 3.11, vector_io) (push) Failing after 7s
Integration Tests / test-matrix (http, 3.12, inspect) (push) Failing after 6s
Integration Tests / test-matrix (http, 3.12, tool_runtime) (push) Failing after 6s
Integration Tests / test-matrix (http, 3.11, inspect) (push) Failing after 13s
Integration Tests / test-matrix (library, 3.11, datasets) (push) Failing after 5s
Integration Tests / test-matrix (http, 3.11, tool_runtime) (push) Failing after 13s
Integration Tests / test-matrix (http, 3.12, inference) (push) Failing after 11s
Integration Tests / test-matrix (http, 3.12, scoring) (push) Failing after 12s
Integration Tests / test-matrix (library, 3.11, agents) (push) Failing after 10s
Integration Tests / test-matrix (library, 3.11, scoring) (push) Failing after 5s
Integration Tests / test-matrix (library, 3.11, vector_io) (push) Failing after 5s
Integration Tests / test-matrix (library, 3.12, datasets) (push) Failing after 5s
Integration Tests / test-matrix (library, 3.11, inspect) (push) Failing after 13s
Integration Tests / test-matrix (library, 3.11, providers) (push) Failing after 11s
Integration Tests / test-matrix (library, 3.11, tool_runtime) (push) Failing after 10s
Integration Tests / test-matrix (library, 3.12, inference) (push) Failing after 6s
Integration Tests / test-matrix (library, 3.12, agents) (push) Failing after 8s
Integration Tests / test-matrix (library, 3.12, inspect) (push) Failing after 6s
Integration Tests / test-matrix (library, 3.12, post_training) (push) Failing after 5s
Integration Tests / test-matrix (http, 3.11, scoring) (push) Failing after 28s
Integration Tests / test-matrix (http, 3.12, providers) (push) Failing after 24s
Integration Tests / test-matrix (http, 3.12, datasets) (push) Failing after 26s
Integration Tests / test-matrix (http, 3.11, inference) (push) Failing after 30s
Integration Tests / test-matrix (http, 3.12, agents) (push) Failing after 28s
Integration Tests / test-matrix (http, 3.12, post_training) (push) Failing after 26s
Integration Tests / test-matrix (http, 3.12, vector_io) (push) Failing after 23s
Test Llama Stack Build / generate-matrix (push) Successful in 5s
Integration Tests / test-matrix (library, 3.12, scoring) (push) Failing after 5s
Test Llama Stack Build / build-custom-container-distribution (push) Failing after 5s
Test External Providers / test-external-providers (venv) (push) Failing after 5s
Integration Tests / test-matrix (library, 3.11, post_training) (push) Failing after 20s
Integration Tests / test-matrix (library, 3.12, providers) (push) Failing after 7s
Unit Tests / unit-tests (3.11) (push) Failing after 7s
Update ReadTheDocs / update-readthedocs (push) Failing after 6s
Integration Tests / test-matrix (library, 3.12, vector_io) (push) Failing after 8s
Integration Tests / test-matrix (library, 3.11, inference) (push) Failing after 22s
Test Llama Stack Build / build (push) Failing after 17s
Unit Tests / unit-tests (3.13) (push) Failing after 37s
Integration Tests / test-matrix (library, 3.12, tool_runtime) (push) Failing after 1m7s
Test Llama Stack Build / build-single-provider (push) Failing after 1m15s
Test Llama Stack Build / build-ubi9-container-distribution (push) Failing after 1m17s
Unit Tests / unit-tests (3.12) (push) Failing after 1m32s
Pre-commit / pre-commit (push) Failing after 2m14s

# What does this PR do?

Do not force 384 for the embedding dimension, use the one provided by
the test run.

## Test Plan

```
 pytest -s -vvv tests/integration/vector_io/test_vector_io.py --stack-config=http://localhost:8321 \
    -k "not(builtin_tool or safety_with_image or code_interpreter or test_rag)" \
    --text-model="meta-llama/Llama-3.2-3B-Instruct" \
    --embedding-model=granite-embedding-125m --embedding-dimension=768
Uninstalled 1 package in 16ms
Installed 1 package in 11ms
INFO     2025-06-18 10:52:03,314 tests.integration.conftest:59 tests: Setting DISABLE_CODE_SANDBOX=1 for macOS
/Users/leseb/Documents/AI/llama-stack/.venv/lib/python3.10/site-packages/pytest_asyncio/plugin.py:207: PytestDeprecationWarning: The configuration option "asyncio_default_fixture_loop_scope" is unset.
The event loop scope for asynchronous fixtures will default to the fixture caching scope. Future versions of pytest-asyncio will default the loop scope for asynchronous fixtures to function scope. Set the default fixture loop scope explicitly in order to avoid unexpected behavior in the future. Valid fixture loop scopes are: "function", "class", "module", "package", "session"

  warnings.warn(PytestDeprecationWarning(_DEFAULT_FIXTURE_LOOP_SCOPE_UNSET))
================================================= test session starts =================================================
platform darwin -- Python 3.10.16, pytest-8.3.4, pluggy-1.5.0 -- /Users/leseb/Documents/AI/llama-stack/.venv/bin/python
cachedir: .pytest_cache
metadata: {'Python': '3.10.16', 'Platform': 'macOS-15.5-arm64-arm-64bit', 'Packages': {'pytest': '8.3.4', 'pluggy': '1.5.0'}, 'Plugins': {'cov': '6.0.0', 'html': '4.1.1', 'json-report': '1.5.0', 'timeout': '2.4.0', 'metadata': '3.1.1', 'asyncio': '0.25.3', 'anyio': '4.8.0', 'nbval': '0.11.0'}}
rootdir: /Users/leseb/Documents/AI/llama-stack
configfile: pyproject.toml
plugins: cov-6.0.0, html-4.1.1, json-report-1.5.0, timeout-2.4.0, metadata-3.1.1, asyncio-0.25.3, anyio-4.8.0, nbval-0.11.0
asyncio: mode=strict, asyncio_default_fixture_loop_scope=None
collected 8 items

tests/integration/vector_io/test_vector_io.py::test_vector_db_retrieve[emb=granite-embedding-125m:dim=768] PASSED
tests/integration/vector_io/test_vector_io.py::test_vector_db_register[emb=granite-embedding-125m:dim=768] PASSED
tests/integration/vector_io/test_vector_io.py::test_insert_chunks[emb=granite-embedding-125m:dim=768-test_case0] PASSED
tests/integration/vector_io/test_vector_io.py::test_insert_chunks[emb=granite-embedding-125m:dim=768-test_case1] PASSED
tests/integration/vector_io/test_vector_io.py::test_insert_chunks[emb=granite-embedding-125m:dim=768-test_case2] PASSED
tests/integration/vector_io/test_vector_io.py::test_insert_chunks[emb=granite-embedding-125m:dim=768-test_case3] PASSED
tests/integration/vector_io/test_vector_io.py::test_insert_chunks[emb=granite-embedding-125m:dim=768-test_case4] PASSED
tests/integration/vector_io/test_vector_io.py::test_insert_chunks_with_precomputed_embeddings[emb=granite-embedding-125m:dim=768] PASSED

================================================== 8 passed in 5.50s ==================================================
```

Signed-off-by: Sébastien Han <seb@redhat.com>
This commit is contained in:
Sébastien Han 2025-06-19 09:59:04 +02:00 committed by GitHub
parent d12f195f56
commit 6039d922c0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 25 additions and 18 deletions

View file

@ -95,6 +95,7 @@ def pytest_addoption(parser):
parser.addoption( parser.addoption(
"--embedding-dimension", "--embedding-dimension",
type=int, type=int,
default=384,
help="Output dimensionality of the embedding model to use for testing. Default: 384", help="Output dimensionality of the embedding model to use for testing. Default: 384",
) )
parser.addoption( parser.addoption(

View file

@ -63,12 +63,14 @@ def assert_valid_text_response(response):
assert all(isinstance(chunk.text, str) for chunk in response.content) assert all(isinstance(chunk.text, str) for chunk in response.content)
def test_vector_db_insert_inline_and_query(client_with_empty_registry, sample_documents, embedding_model_id): def test_vector_db_insert_inline_and_query(
client_with_empty_registry, sample_documents, embedding_model_id, embedding_dimension
):
vector_db_id = "test_vector_db" vector_db_id = "test_vector_db"
client_with_empty_registry.vector_dbs.register( client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_id, vector_db_id=vector_db_id,
embedding_model=embedding_model_id, embedding_model=embedding_model_id,
embedding_dimension=384, embedding_dimension=embedding_dimension,
) )
client_with_empty_registry.tool_runtime.rag_tool.insert( client_with_empty_registry.tool_runtime.rag_tool.insert(
@ -116,7 +118,9 @@ def test_vector_db_insert_inline_and_query(client_with_empty_registry, sample_do
assert all(score >= 0.01 for score in response4.scores) assert all(score >= 0.01 for score in response4.scores)
def test_vector_db_insert_from_url_and_query(client_with_empty_registry, sample_documents, embedding_model_id): def test_vector_db_insert_from_url_and_query(
client_with_empty_registry, sample_documents, embedding_model_id, embedding_dimension
):
providers = [p for p in client_with_empty_registry.providers.list() if p.api == "vector_io"] providers = [p for p in client_with_empty_registry.providers.list() if p.api == "vector_io"]
assert len(providers) > 0 assert len(providers) > 0
@ -125,7 +129,7 @@ def test_vector_db_insert_from_url_and_query(client_with_empty_registry, sample_
client_with_empty_registry.vector_dbs.register( client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_id, vector_db_id=vector_db_id,
embedding_model=embedding_model_id, embedding_model=embedding_model_id,
embedding_dimension=384, embedding_dimension=embedding_dimension,
) )
# list to check memory bank is successfully registered # list to check memory bank is successfully registered
@ -170,7 +174,7 @@ def test_vector_db_insert_from_url_and_query(client_with_empty_registry, sample_
assert any("llama2" in chunk.content.lower() for chunk in response2.chunks) assert any("llama2" in chunk.content.lower() for chunk in response2.chunks)
def test_rag_tool_insert_and_query(client_with_empty_registry, embedding_model_id): def test_rag_tool_insert_and_query(client_with_empty_registry, embedding_model_id, embedding_dimension):
providers = [p for p in client_with_empty_registry.providers.list() if p.api == "vector_io"] providers = [p for p in client_with_empty_registry.providers.list() if p.api == "vector_io"]
assert len(providers) > 0 assert len(providers) > 0
@ -179,7 +183,7 @@ def test_rag_tool_insert_and_query(client_with_empty_registry, embedding_model_i
client_with_empty_registry.vector_dbs.register( client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_id, vector_db_id=vector_db_id,
embedding_model=embedding_model_id, embedding_model=embedding_model_id,
embedding_dimension=384, embedding_dimension=embedding_dimension,
) )
available_vector_dbs = [vector_db.identifier for vector_db in client_with_empty_registry.vector_dbs.list()] available_vector_dbs = [vector_db.identifier for vector_db in client_with_empty_registry.vector_dbs.list()]

View file

@ -46,13 +46,13 @@ def client_with_empty_registry(client_with_models):
clear_registry() clear_registry()
def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id): def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id, embedding_dimension):
# Register a memory bank first # Register a memory bank first
vector_db_id = "test_vector_db" vector_db_id = "test_vector_db"
client_with_empty_registry.vector_dbs.register( client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_id, vector_db_id=vector_db_id,
embedding_model=embedding_model_id, embedding_model=embedding_model_id,
embedding_dimension=384, embedding_dimension=embedding_dimension,
) )
# Retrieve the memory bank and validate its properties # Retrieve the memory bank and validate its properties
@ -63,12 +63,12 @@ def test_vector_db_retrieve(client_with_empty_registry, embedding_model_id):
assert response.provider_resource_id == vector_db_id assert response.provider_resource_id == vector_db_id
def test_vector_db_register(client_with_empty_registry, embedding_model_id): def test_vector_db_register(client_with_empty_registry, embedding_model_id, embedding_dimension):
vector_db_id = "test_vector_db" vector_db_id = "test_vector_db"
client_with_empty_registry.vector_dbs.register( client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_id, vector_db_id=vector_db_id,
embedding_model=embedding_model_id, embedding_model=embedding_model_id,
embedding_dimension=384, embedding_dimension=embedding_dimension,
) )
vector_dbs_after_register = [vector_db.identifier for vector_db in client_with_empty_registry.vector_dbs.list()] vector_dbs_after_register = [vector_db.identifier for vector_db in client_with_empty_registry.vector_dbs.list()]
@ -90,12 +90,12 @@ def test_vector_db_register(client_with_empty_registry, embedding_model_id):
("How does machine learning improve over time?", "doc2"), ("How does machine learning improve over time?", "doc2"),
], ],
) )
def test_insert_chunks(client_with_empty_registry, embedding_model_id, sample_chunks, test_case): def test_insert_chunks(client_with_empty_registry, embedding_model_id, embedding_dimension, sample_chunks, test_case):
vector_db_id = "test_vector_db" vector_db_id = "test_vector_db"
client_with_empty_registry.vector_dbs.register( client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_id, vector_db_id=vector_db_id,
embedding_model=embedding_model_id, embedding_model=embedding_model_id,
embedding_dimension=384, embedding_dimension=embedding_dimension,
) )
client_with_empty_registry.vector_io.insert( client_with_empty_registry.vector_io.insert(
@ -122,19 +122,19 @@ def test_insert_chunks(client_with_empty_registry, embedding_model_id, sample_ch
assert top_match.metadata["document_id"] == expected_doc_id, f"Query '{query}' should match {expected_doc_id}" 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): def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, embedding_model_id, embedding_dimension):
vector_db_id = "test_precomputed_embeddings_db" vector_db_id = "test_precomputed_embeddings_db"
client_with_empty_registry.vector_dbs.register( client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_id, vector_db_id=vector_db_id,
embedding_model=embedding_model_id, embedding_model=embedding_model_id,
embedding_dimension=384, embedding_dimension=embedding_dimension,
) )
chunks_with_embeddings = [ chunks_with_embeddings = [
Chunk( Chunk(
content="This is a test chunk with precomputed embedding.", content="This is a test chunk with precomputed embedding.",
metadata={"document_id": "doc1", "source": "precomputed"}, metadata={"document_id": "doc1", "source": "precomputed"},
embedding=[0.1] * 384, embedding=[0.1] * int(embedding_dimension),
), ),
] ]
@ -156,19 +156,21 @@ def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, e
assert response.chunks[0].metadata["source"] == "precomputed" assert response.chunks[0].metadata["source"] == "precomputed"
def test_query_returns_valid_object_when_identical_to_embedding_in_vdb(client_with_empty_registry, embedding_model_id): def test_query_returns_valid_object_when_identical_to_embedding_in_vdb(
client_with_empty_registry, embedding_model_id, embedding_dimension
):
vector_db_id = "test_precomputed_embeddings_db" vector_db_id = "test_precomputed_embeddings_db"
client_with_empty_registry.vector_dbs.register( client_with_empty_registry.vector_dbs.register(
vector_db_id=vector_db_id, vector_db_id=vector_db_id,
embedding_model=embedding_model_id, embedding_model=embedding_model_id,
embedding_dimension=384, embedding_dimension=embedding_dimension,
) )
chunks_with_embeddings = [ chunks_with_embeddings = [
Chunk( Chunk(
content="duplicate", content="duplicate",
metadata={"document_id": "doc1", "source": "precomputed"}, metadata={"document_id": "doc1", "source": "precomputed"},
embedding=[0.1] * 384, embedding=[0.1] * int(embedding_dimension),
), ),
] ]