mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-17 11:19:47 +00:00
4f8b73b9e1
# What does this PR do? Moves all the memory providers to use the inference API and improved the memory tests to setup the inference stack correctly and use the embedding models ## Test Plan torchrun $CONDA_PREFIX/bin/pytest -v -s -k "meta_reference" --inference-model="Llama3.2-3B-Instruct" --embedding-model="sentence-transformers/all-MiniLM-L6-v2" llama_stack/providers/tests/inference/test_embeddings.py --env EMBEDDING_DIMENSION=384 pytest -v -s llama_stack/providers/tests/memory/test_memory.py --providers="inference=together,memory=weaviate" --embedding-model="togethercomputer/m2-bert-80M-2k-retrieval" --env EMBEDDING_DIMENSION=768 --env TOGETHER_API_KEY=<API-KEY> --env WEAVIATE_API_KEY=foo --env WEAVIATE_CLUSTER_URL=bar pytest -v -s llama_stack/providers/tests/memory/test_memory.py --providers="inference=together,memory=chroma" --embedding-model="togethercomputer/m2-bert-80M-2k-retrieval" --env EMBEDDING_DIMENSION=768 --env TOGETHER_API_KEY=<API-KEY>--env CHROMA_HOST=localhost --env CHROMA_PORT=8000 pytest -v -s llama_stack/providers/tests/memory/test_memory.py --providers="inference=together,memory=pgvector" --embedding-model="togethercomputer/m2-bert-80M-2k-retrieval" --env PGVECTOR_DB=postgres --env PGVECTOR_USER=postgres --env PGVECTOR_PASSWORD=mysecretpassword --env PGVECTOR_HOST=0.0.0.0 --env EMBEDDING_DIMENSION=768 --env TOGETHER_API_KEY=<API-KEY> pytest -v -s llama_stack/providers/tests/memory/test_memory.py --providers="inference=together,memory=faiss" --embedding-model="togethercomputer/m2-bert-80M-2k-retrieval" --env EMBEDDING_DIMENSION=768 --env TOGETHER_API_KEY=<API-KEY>
95 lines
2.5 KiB
Python
95 lines
2.5 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
|
|
|
|
from ..inference.fixtures import INFERENCE_FIXTURES
|
|
from .fixtures import MEMORY_FIXTURES
|
|
|
|
|
|
DEFAULT_PROVIDER_COMBINATIONS = [
|
|
pytest.param(
|
|
{
|
|
"inference": "meta_reference",
|
|
"memory": "faiss",
|
|
},
|
|
id="meta_reference",
|
|
marks=pytest.mark.meta_reference,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "ollama",
|
|
"memory": "pgvector",
|
|
},
|
|
id="ollama",
|
|
marks=pytest.mark.ollama,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "together",
|
|
"memory": "chroma",
|
|
},
|
|
id="chroma",
|
|
marks=pytest.mark.chroma,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "bedrock",
|
|
"memory": "qdrant",
|
|
},
|
|
id="qdrant",
|
|
marks=pytest.mark.qdrant,
|
|
),
|
|
pytest.param(
|
|
{
|
|
"inference": "fireworks",
|
|
"memory": "weaviate",
|
|
},
|
|
id="weaviate",
|
|
marks=pytest.mark.weaviate,
|
|
),
|
|
]
|
|
|
|
|
|
def pytest_addoption(parser):
|
|
parser.addoption(
|
|
"--embedding-model",
|
|
action="store",
|
|
default=None,
|
|
help="Specify the embedding model to use for testing",
|
|
)
|
|
|
|
|
|
def pytest_configure(config):
|
|
for fixture_name in MEMORY_FIXTURES:
|
|
config.addinivalue_line(
|
|
"markers",
|
|
f"{fixture_name}: marks tests as {fixture_name} specific",
|
|
)
|
|
|
|
|
|
def pytest_generate_tests(metafunc):
|
|
if "embedding_model" in metafunc.fixturenames:
|
|
model = metafunc.config.getoption("--embedding-model")
|
|
if not model:
|
|
raise ValueError(
|
|
"No embedding model specified. Please provide a valid embedding model."
|
|
)
|
|
params = [pytest.param(model, id="")]
|
|
|
|
metafunc.parametrize("embedding_model", params, indirect=True)
|
|
if "memory_stack" in metafunc.fixturenames:
|
|
available_fixtures = {
|
|
"inference": INFERENCE_FIXTURES,
|
|
"memory": MEMORY_FIXTURES,
|
|
}
|
|
combinations = (
|
|
get_provider_fixture_overrides(metafunc.config, available_fixtures)
|
|
or DEFAULT_PROVIDER_COMBINATIONS
|
|
)
|
|
metafunc.parametrize("memory_stack", combinations, indirect=True)
|