mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-28 02:53:30 +00:00
a79a083e39
This PR fixes the broken pgvector provider as well as wraps all cursor object creations with context manager to ensure that they get properly closed to avoid potential memory leaks. ``` > pytest llama_stack/providers/tests/vector_io/test_vector_io.py -m "pgvector" --env EMBEDDING_DIMENSION=384 --env PGVECTOR_PORT=7432 --env PGVECTOR_DB=db --env PGVECTOR_USER=user --env PGVECTOR_PASSWORD=pass -v -s --tb=short --disable-warnings llama_stack/providers/tests/vector_io/test_vector_io.py::TestVectorIO::test_banks_list[-pgvector] PASSED llama_stack/providers/tests/vector_io/test_vector_io.py::TestVectorIO::test_banks_register[-pgvector] PASSED llama_stack/providers/tests/vector_io/test_vector_io.py::TestVectorIO::test_query_documents[-pgvector] The scores are: [0.8168284974053789, 0.8080469278964486, 0.8050996198466661] PASSED ``` --------- Signed-off-by: Yuan Tang <terrytangyuan@gmail.com>
102 lines
2.9 KiB
Python
102 lines
2.9 KiB
Python
# 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.
|
|
|
|
import pytest
|
|
|
|
from ..conftest import (
|
|
get_provider_fixture_overrides,
|
|
get_provider_fixture_overrides_from_test_config,
|
|
get_test_config_for_api,
|
|
)
|
|
|
|
from ..inference.fixtures import INFERENCE_FIXTURES
|
|
from .fixtures import VECTOR_IO_FIXTURES
|
|
|
|
|
|
DEFAULT_PROVIDER_COMBINATIONS = [
|
|
pytest.param(
|
|
{
|
|
"inference": "sentence_transformers",
|
|
"vector_io": "faiss",
|
|
},
|
|
id="sentence_transformers",
|
|
marks=pytest.mark.sentence_transformers,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "ollama",
|
|
"vector_io": "pgvector",
|
|
},
|
|
id="pgvector",
|
|
marks=pytest.mark.pgvector,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "ollama",
|
|
"vector_io": "faiss",
|
|
},
|
|
id="ollama",
|
|
marks=pytest.mark.ollama,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "sentence_transformers",
|
|
"vector_io": "chroma",
|
|
},
|
|
id="chroma",
|
|
marks=pytest.mark.chroma,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "bedrock",
|
|
"vector_io": "qdrant",
|
|
},
|
|
id="qdrant",
|
|
marks=pytest.mark.qdrant,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "fireworks",
|
|
"vector_io": "weaviate",
|
|
},
|
|
id="weaviate",
|
|
marks=pytest.mark.weaviate,
|
|
),
|
|
]
|
|
|
|
|
|
def pytest_configure(config):
|
|
for fixture_name in VECTOR_IO_FIXTURES:
|
|
config.addinivalue_line(
|
|
"markers",
|
|
f"{fixture_name}: marks tests as {fixture_name} specific",
|
|
)
|
|
|
|
|
|
def pytest_generate_tests(metafunc):
|
|
test_config = get_test_config_for_api(metafunc.config, "vector_io")
|
|
if "embedding_model" in metafunc.fixturenames:
|
|
model = getattr(test_config, "embedding_model", None)
|
|
# Fall back to the default if not specified by the config file
|
|
model = model or metafunc.config.getoption("--embedding-model")
|
|
if model:
|
|
params = [pytest.param(model, id="")]
|
|
else:
|
|
params = [pytest.param("all-minilm:l6-v2", id="")]
|
|
|
|
metafunc.parametrize("embedding_model", params, indirect=True)
|
|
|
|
if "vector_io_stack" in metafunc.fixturenames:
|
|
available_fixtures = {
|
|
"inference": INFERENCE_FIXTURES,
|
|
"vector_io": VECTOR_IO_FIXTURES,
|
|
}
|
|
combinations = (
|
|
get_provider_fixture_overrides_from_test_config(metafunc.config, "vector_io", DEFAULT_PROVIDER_COMBINATIONS)
|
|
or get_provider_fixture_overrides(metafunc.config, available_fixtures)
|
|
or DEFAULT_PROVIDER_COMBINATIONS
|
|
)
|
|
metafunc.parametrize("vector_io_stack", combinations, indirect=True)
|