forked from phoenix-oss/llama-stack-mirror
[memory refactor][2/n] Update faiss and make it pass tests (#830)
See https://github.com/meta-llama/llama-stack/issues/827 for the broader design. Second part: - updates routing table / router code - updates the faiss implementation ## Test Plan ``` pytest -s -v -k sentence test_vector_io.py --env EMBEDDING_DIMENSION=384 ```
This commit is contained in:
Ashwin Bharambe
GitHub
parent
3ae8585b65
commit
78a481bb22
19 changed files with 343 additions and 353 deletions
|
|
@ -302,7 +302,7 @@ def pytest_collection_modifyitems(session, config, items):
|
|||
pytest_plugins = [
|
||||
"llama_stack.providers.tests.inference.fixtures",
|
||||
"llama_stack.providers.tests.safety.fixtures",
|
||||
"llama_stack.providers.tests.memory.fixtures",
|
||||
"llama_stack.providers.tests.vector_io.fixtures",
|
||||
"llama_stack.providers.tests.agents.fixtures",
|
||||
"llama_stack.providers.tests.datasetio.fixtures",
|
||||
"llama_stack.providers.tests.scoring.fixtures",
|
||||
|
|
|
|||
|
|
@ -1,192 +0,0 @@
|
|||
# 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 uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from llama_stack.apis.memory import MemoryBankDocument, QueryDocumentsResponse
|
||||
|
||||
from llama_stack.apis.memory_banks import (
|
||||
MemoryBank,
|
||||
MemoryBanks,
|
||||
VectorMemoryBankParams,
|
||||
)
|
||||
|
||||
# How to run this test:
|
||||
#
|
||||
# pytest llama_stack/providers/tests/memory/test_memory.py
|
||||
# -m "sentence_transformers" --env EMBEDDING_DIMENSION=384
|
||||
# -v -s --tb=short --disable-warnings
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_documents():
|
||||
return [
|
||||
MemoryBankDocument(
|
||||
document_id="doc1",
|
||||
content="Python is a high-level programming language.",
|
||||
metadata={"category": "programming", "difficulty": "beginner"},
|
||||
),
|
||||
MemoryBankDocument(
|
||||
document_id="doc2",
|
||||
content="Machine learning is a subset of artificial intelligence.",
|
||||
metadata={"category": "AI", "difficulty": "advanced"},
|
||||
),
|
||||
MemoryBankDocument(
|
||||
document_id="doc3",
|
||||
content="Data structures are fundamental to computer science.",
|
||||
metadata={"category": "computer science", "difficulty": "intermediate"},
|
||||
),
|
||||
MemoryBankDocument(
|
||||
document_id="doc4",
|
||||
content="Neural networks are inspired by biological neural networks.",
|
||||
metadata={"category": "AI", "difficulty": "advanced"},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
async def register_memory_bank(
|
||||
banks_impl: MemoryBanks, embedding_model: str
|
||||
) -> MemoryBank:
|
||||
bank_id = f"test_bank_{uuid.uuid4().hex}"
|
||||
return await banks_impl.register_memory_bank(
|
||||
memory_bank_id=bank_id,
|
||||
params=VectorMemoryBankParams(
|
||||
embedding_model=embedding_model,
|
||||
chunk_size_in_tokens=512,
|
||||
overlap_size_in_tokens=64,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestMemory:
|
||||
@pytest.mark.asyncio
|
||||
async def test_banks_list(self, memory_stack, embedding_model):
|
||||
_, banks_impl = memory_stack
|
||||
|
||||
# Register a test bank
|
||||
registered_bank = await register_memory_bank(banks_impl, embedding_model)
|
||||
|
||||
try:
|
||||
# Verify our bank shows up in list
|
||||
response = await banks_impl.list_memory_banks()
|
||||
assert isinstance(response, list)
|
||||
assert any(
|
||||
bank.memory_bank_id == registered_bank.memory_bank_id
|
||||
for bank in response
|
||||
)
|
||||
finally:
|
||||
# Clean up
|
||||
await banks_impl.unregister_memory_bank(registered_bank.memory_bank_id)
|
||||
|
||||
# Verify our bank was removed
|
||||
response = await banks_impl.list_memory_banks()
|
||||
assert all(
|
||||
bank.memory_bank_id != registered_bank.memory_bank_id for bank in response
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_banks_register(self, memory_stack, embedding_model):
|
||||
_, banks_impl = memory_stack
|
||||
|
||||
bank_id = f"test_bank_{uuid.uuid4().hex}"
|
||||
|
||||
try:
|
||||
# Register initial bank
|
||||
await banks_impl.register_memory_bank(
|
||||
memory_bank_id=bank_id,
|
||||
params=VectorMemoryBankParams(
|
||||
embedding_model=embedding_model,
|
||||
chunk_size_in_tokens=512,
|
||||
overlap_size_in_tokens=64,
|
||||
),
|
||||
)
|
||||
|
||||
# Verify our bank exists
|
||||
response = await banks_impl.list_memory_banks()
|
||||
assert isinstance(response, list)
|
||||
assert any(bank.memory_bank_id == bank_id for bank in response)
|
||||
|
||||
# Try registering same bank again
|
||||
await banks_impl.register_memory_bank(
|
||||
memory_bank_id=bank_id,
|
||||
params=VectorMemoryBankParams(
|
||||
embedding_model=embedding_model,
|
||||
chunk_size_in_tokens=512,
|
||||
overlap_size_in_tokens=64,
|
||||
),
|
||||
)
|
||||
|
||||
# Verify still only one instance of our bank
|
||||
response = await banks_impl.list_memory_banks()
|
||||
assert isinstance(response, list)
|
||||
assert (
|
||||
len([bank for bank in response if bank.memory_bank_id == bank_id]) == 1
|
||||
)
|
||||
finally:
|
||||
# Clean up
|
||||
await banks_impl.unregister_memory_bank(bank_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_documents(
|
||||
self, memory_stack, embedding_model, sample_documents
|
||||
):
|
||||
memory_impl, banks_impl = memory_stack
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
await memory_impl.insert_documents("test_bank", sample_documents)
|
||||
|
||||
registered_bank = await register_memory_bank(banks_impl, embedding_model)
|
||||
await memory_impl.insert_documents(
|
||||
registered_bank.memory_bank_id, sample_documents
|
||||
)
|
||||
|
||||
query1 = "programming language"
|
||||
response1 = await memory_impl.query_documents(
|
||||
registered_bank.memory_bank_id, query1
|
||||
)
|
||||
assert_valid_response(response1)
|
||||
assert any("Python" in chunk.content for chunk in response1.chunks)
|
||||
|
||||
# Test case 3: Query with semantic similarity
|
||||
query3 = "AI and brain-inspired computing"
|
||||
response3 = await memory_impl.query_documents(
|
||||
registered_bank.memory_bank_id, query3
|
||||
)
|
||||
assert_valid_response(response3)
|
||||
assert any(
|
||||
"neural networks" in chunk.content.lower() for chunk in response3.chunks
|
||||
)
|
||||
|
||||
# Test case 4: Query with limit on number of results
|
||||
query4 = "computer"
|
||||
params4 = {"max_chunks": 2}
|
||||
response4 = await memory_impl.query_documents(
|
||||
registered_bank.memory_bank_id, query4, params4
|
||||
)
|
||||
assert_valid_response(response4)
|
||||
assert len(response4.chunks) <= 2
|
||||
|
||||
# Test case 5: Query with threshold on similarity score
|
||||
query5 = "quantum computing" # Not directly related to any document
|
||||
params5 = {"score_threshold": 0.01}
|
||||
response5 = await memory_impl.query_documents(
|
||||
registered_bank.memory_bank_id, query5, params5
|
||||
)
|
||||
assert_valid_response(response5)
|
||||
print("The scores are:", response5.scores)
|
||||
assert all(score >= 0.01 for score in response5.scores)
|
||||
|
||||
|
||||
def assert_valid_response(response: QueryDocumentsResponse):
|
||||
assert isinstance(response, QueryDocumentsResponse)
|
||||
assert len(response.chunks) > 0
|
||||
assert len(response.scores) > 0
|
||||
assert len(response.chunks) == len(response.scores)
|
||||
for chunk in response.chunks:
|
||||
assert isinstance(chunk.content, str)
|
||||
assert chunk.document_id is not None
|
||||
|
|
@ -12,11 +12,11 @@ from pydantic import BaseModel
|
|||
|
||||
from llama_stack.apis.datasets import DatasetInput
|
||||
from llama_stack.apis.eval_tasks import EvalTaskInput
|
||||
from llama_stack.apis.memory_banks import MemoryBankInput
|
||||
from llama_stack.apis.models import ModelInput
|
||||
from llama_stack.apis.scoring_functions import ScoringFnInput
|
||||
from llama_stack.apis.shields import ShieldInput
|
||||
from llama_stack.apis.tools import ToolGroupInput
|
||||
from llama_stack.apis.vector_dbs import VectorDBInput
|
||||
from llama_stack.distribution.build import print_pip_install_help
|
||||
from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
|
||||
from llama_stack.distribution.datatypes import Provider, StackRunConfig
|
||||
|
|
@ -39,7 +39,7 @@ async def construct_stack_for_test(
|
|||
provider_data: Optional[Dict[str, Any]] = None,
|
||||
models: Optional[List[ModelInput]] = None,
|
||||
shields: Optional[List[ShieldInput]] = None,
|
||||
memory_banks: Optional[List[MemoryBankInput]] = None,
|
||||
vector_dbs: Optional[List[VectorDBInput]] = None,
|
||||
datasets: Optional[List[DatasetInput]] = None,
|
||||
scoring_fns: Optional[List[ScoringFnInput]] = None,
|
||||
eval_tasks: Optional[List[EvalTaskInput]] = None,
|
||||
|
|
@ -53,7 +53,7 @@ async def construct_stack_for_test(
|
|||
metadata_store=SqliteKVStoreConfig(db_path=sqlite_file.name),
|
||||
models=models or [],
|
||||
shields=shields or [],
|
||||
memory_banks=memory_banks or [],
|
||||
vector_dbs=vector_dbs or [],
|
||||
datasets=datasets or [],
|
||||
scoring_fns=scoring_fns or [],
|
||||
eval_tasks=eval_tasks or [],
|
||||
|
|
|
|||
|
|
@ -13,14 +13,14 @@ from ..conftest import (
|
|||
)
|
||||
|
||||
from ..inference.fixtures import INFERENCE_FIXTURES
|
||||
from .fixtures import MEMORY_FIXTURES
|
||||
from .fixtures import VECTOR_IO_FIXTURES
|
||||
|
||||
|
||||
DEFAULT_PROVIDER_COMBINATIONS = [
|
||||
pytest.param(
|
||||
{
|
||||
"inference": "sentence_transformers",
|
||||
"memory": "faiss",
|
||||
"vector_io": "faiss",
|
||||
},
|
||||
id="sentence_transformers",
|
||||
marks=pytest.mark.sentence_transformers,
|
||||
|
|
@ -28,7 +28,7 @@ DEFAULT_PROVIDER_COMBINATIONS = [
|
|||
pytest.param(
|
||||
{
|
||||
"inference": "ollama",
|
||||
"memory": "faiss",
|
||||
"vector_io": "faiss",
|
||||
},
|
||||
id="ollama",
|
||||
marks=pytest.mark.ollama,
|
||||
|
|
@ -36,7 +36,7 @@ DEFAULT_PROVIDER_COMBINATIONS = [
|
|||
pytest.param(
|
||||
{
|
||||
"inference": "sentence_transformers",
|
||||
"memory": "chroma",
|
||||
"vector_io": "chroma",
|
||||
},
|
||||
id="chroma",
|
||||
marks=pytest.mark.chroma,
|
||||
|
|
@ -44,7 +44,7 @@ DEFAULT_PROVIDER_COMBINATIONS = [
|
|||
pytest.param(
|
||||
{
|
||||
"inference": "bedrock",
|
||||
"memory": "qdrant",
|
||||
"vector_io": "qdrant",
|
||||
},
|
||||
id="qdrant",
|
||||
marks=pytest.mark.qdrant,
|
||||
|
|
@ -52,7 +52,7 @@ DEFAULT_PROVIDER_COMBINATIONS = [
|
|||
pytest.param(
|
||||
{
|
||||
"inference": "fireworks",
|
||||
"memory": "weaviate",
|
||||
"vector_io": "weaviate",
|
||||
},
|
||||
id="weaviate",
|
||||
marks=pytest.mark.weaviate,
|
||||
|
|
@ -61,7 +61,7 @@ DEFAULT_PROVIDER_COMBINATIONS = [
|
|||
|
||||
|
||||
def pytest_configure(config):
|
||||
for fixture_name in MEMORY_FIXTURES:
|
||||
for fixture_name in VECTOR_IO_FIXTURES:
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
f"{fixture_name}: marks tests as {fixture_name} specific",
|
||||
|
|
@ -69,7 +69,7 @@ def pytest_configure(config):
|
|||
|
||||
|
||||
def pytest_generate_tests(metafunc):
|
||||
test_config = get_test_config_for_api(metafunc.config, "memory")
|
||||
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
|
||||
|
|
@ -81,16 +81,16 @@ def pytest_generate_tests(metafunc):
|
|||
|
||||
metafunc.parametrize("embedding_model", params, indirect=True)
|
||||
|
||||
if "memory_stack" in metafunc.fixturenames:
|
||||
if "vector_io_stack" in metafunc.fixturenames:
|
||||
available_fixtures = {
|
||||
"inference": INFERENCE_FIXTURES,
|
||||
"memory": MEMORY_FIXTURES,
|
||||
"vector_io": VECTOR_IO_FIXTURES,
|
||||
}
|
||||
combinations = (
|
||||
get_provider_fixture_overrides_from_test_config(
|
||||
metafunc.config, "memory", DEFAULT_PROVIDER_COMBINATIONS
|
||||
metafunc.config, "vector_io", DEFAULT_PROVIDER_COMBINATIONS
|
||||
)
|
||||
or get_provider_fixture_overrides(metafunc.config, available_fixtures)
|
||||
or DEFAULT_PROVIDER_COMBINATIONS
|
||||
)
|
||||
metafunc.parametrize("memory_stack", combinations, indirect=True)
|
||||
metafunc.parametrize("vector_io_stack", combinations, indirect=True)
|
||||
|
|
@ -12,11 +12,12 @@ import pytest_asyncio
|
|||
|
||||
from llama_stack.apis.models import ModelInput, ModelType
|
||||
from llama_stack.distribution.datatypes import Api, Provider
|
||||
from llama_stack.providers.inline.memory.chroma import ChromaInlineImplConfig
|
||||
from llama_stack.providers.inline.memory.faiss import FaissImplConfig
|
||||
from llama_stack.providers.remote.memory.chroma import ChromaRemoteImplConfig
|
||||
from llama_stack.providers.remote.memory.pgvector import PGVectorConfig
|
||||
from llama_stack.providers.remote.memory.weaviate import WeaviateConfig
|
||||
|
||||
from llama_stack.providers.inline.vector_io.chroma import ChromaInlineImplConfig
|
||||
from llama_stack.providers.inline.vector_io.faiss import FaissImplConfig
|
||||
from llama_stack.providers.remote.vector_io.chroma import ChromaRemoteImplConfig
|
||||
from llama_stack.providers.remote.vector_io.pgvector import PGVectorConfig
|
||||
from llama_stack.providers.remote.vector_io.weaviate import WeaviateConfig
|
||||
from llama_stack.providers.tests.resolver import construct_stack_for_test
|
||||
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
|
||||
|
||||
|
|
@ -32,12 +33,12 @@ def embedding_model(request):
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def memory_remote() -> ProviderFixture:
|
||||
def vector_io_remote() -> ProviderFixture:
|
||||
return remote_stack_fixture()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def memory_faiss() -> ProviderFixture:
|
||||
def vector_io_faiss() -> ProviderFixture:
|
||||
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
|
||||
return ProviderFixture(
|
||||
providers=[
|
||||
|
|
@ -53,7 +54,7 @@ def memory_faiss() -> ProviderFixture:
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def memory_pgvector() -> ProviderFixture:
|
||||
def vector_io_pgvector() -> ProviderFixture:
|
||||
return ProviderFixture(
|
||||
providers=[
|
||||
Provider(
|
||||
|
|
@ -72,7 +73,7 @@ def memory_pgvector() -> ProviderFixture:
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def memory_weaviate() -> ProviderFixture:
|
||||
def vector_io_weaviate() -> ProviderFixture:
|
||||
return ProviderFixture(
|
||||
providers=[
|
||||
Provider(
|
||||
|
|
@ -89,7 +90,7 @@ def memory_weaviate() -> ProviderFixture:
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def memory_chroma() -> ProviderFixture:
|
||||
def vector_io_chroma() -> ProviderFixture:
|
||||
url = os.getenv("CHROMA_URL")
|
||||
if url:
|
||||
config = ChromaRemoteImplConfig(url=url)
|
||||
|
|
@ -110,23 +111,23 @@ def memory_chroma() -> ProviderFixture:
|
|||
)
|
||||
|
||||
|
||||
MEMORY_FIXTURES = ["faiss", "pgvector", "weaviate", "remote", "chroma"]
|
||||
VECTOR_IO_FIXTURES = ["faiss", "pgvector", "weaviate", "chroma"]
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="session")
|
||||
async def memory_stack(embedding_model, request):
|
||||
async def vector_io_stack(embedding_model, request):
|
||||
fixture_dict = request.param
|
||||
|
||||
providers = {}
|
||||
provider_data = {}
|
||||
for key in ["inference", "memory"]:
|
||||
for key in ["inference", "vector_io"]:
|
||||
fixture = request.getfixturevalue(f"{key}_{fixture_dict[key]}")
|
||||
providers[key] = fixture.providers
|
||||
if fixture.provider_data:
|
||||
provider_data.update(fixture.provider_data)
|
||||
|
||||
test_stack = await construct_stack_for_test(
|
||||
[Api.memory, Api.inference],
|
||||
[Api.vector_io, Api.inference],
|
||||
providers,
|
||||
provider_data,
|
||||
models=[
|
||||
|
|
@ -140,4 +141,4 @@ async def memory_stack(embedding_model, request):
|
|||
],
|
||||
)
|
||||
|
||||
return test_stack.impls[Api.memory], test_stack.impls[Api.memory_banks]
|
||||
return test_stack.impls[Api.vector_io], test_stack.impls[Api.vector_dbs]
|
||||
200
llama_stack/providers/tests/vector_io/test_vector_io.py
Normal file
200
llama_stack/providers/tests/vector_io/test_vector_io.py
Normal file
|
|
@ -0,0 +1,200 @@
|
|||
# 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 uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from llama_stack.apis.vector_dbs import ListVectorDBsResponse, VectorDB
|
||||
from llama_stack.apis.vector_io import QueryChunksResponse
|
||||
|
||||
from llama_stack.providers.utils.memory.vector_store import (
|
||||
make_overlapped_chunks,
|
||||
MemoryBankDocument,
|
||||
)
|
||||
|
||||
# How to run this test:
|
||||
#
|
||||
# pytest llama_stack/providers/tests/memory/test_memory.py
|
||||
# -m "sentence_transformers" --env EMBEDDING_DIMENSION=384
|
||||
# -v -s --tb=short --disable-warnings
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def sample_chunks():
|
||||
docs = [
|
||||
MemoryBankDocument(
|
||||
document_id="doc1",
|
||||
content="Python is a high-level programming language.",
|
||||
metadata={"category": "programming", "difficulty": "beginner"},
|
||||
),
|
||||
MemoryBankDocument(
|
||||
document_id="doc2",
|
||||
content="Machine learning is a subset of artificial intelligence.",
|
||||
metadata={"category": "AI", "difficulty": "advanced"},
|
||||
),
|
||||
MemoryBankDocument(
|
||||
document_id="doc3",
|
||||
content="Data structures are fundamental to computer science.",
|
||||
metadata={"category": "computer science", "difficulty": "intermediate"},
|
||||
),
|
||||
MemoryBankDocument(
|
||||
document_id="doc4",
|
||||
content="Neural networks are inspired by biological neural networks.",
|
||||
metadata={"category": "AI", "difficulty": "advanced"},
|
||||
),
|
||||
]
|
||||
chunks = []
|
||||
for doc in docs:
|
||||
chunks.extend(
|
||||
make_overlapped_chunks(
|
||||
doc.document_id, doc.content, window_len=512, overlap_len=64
|
||||
)
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
async def register_vector_db(vector_dbs_impl: VectorDB, embedding_model: str):
|
||||
vector_db_id = f"test_vector_db_{uuid.uuid4().hex}"
|
||||
return await vector_dbs_impl.register_vector_db(
|
||||
vector_db_id=vector_db_id,
|
||||
embedding_model=embedding_model,
|
||||
embedding_dimension=384,
|
||||
)
|
||||
|
||||
|
||||
class TestVectorIO:
|
||||
@pytest.mark.asyncio
|
||||
async def test_banks_list(self, vector_io_stack, embedding_model):
|
||||
_, vector_dbs_impl = vector_io_stack
|
||||
|
||||
# Register a test bank
|
||||
registered_vector_db = await register_vector_db(
|
||||
vector_dbs_impl, embedding_model
|
||||
)
|
||||
|
||||
try:
|
||||
# Verify our bank shows up in list
|
||||
response = await vector_dbs_impl.list_vector_dbs()
|
||||
assert isinstance(response, ListVectorDBsResponse)
|
||||
assert any(
|
||||
vector_db.vector_db_id == registered_vector_db.vector_db_id
|
||||
for vector_db in response.data
|
||||
)
|
||||
finally:
|
||||
# Clean up
|
||||
await vector_dbs_impl.unregister_vector_db(
|
||||
registered_vector_db.vector_db_id
|
||||
)
|
||||
|
||||
# Verify our bank was removed
|
||||
response = await vector_dbs_impl.list_vector_dbs()
|
||||
assert isinstance(response, ListVectorDBsResponse)
|
||||
assert all(
|
||||
vector_db.vector_db_id != registered_vector_db.vector_db_id
|
||||
for vector_db in response.data
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_banks_register(self, vector_io_stack, embedding_model):
|
||||
_, vector_dbs_impl = vector_io_stack
|
||||
|
||||
vector_db_id = f"test_vector_db_{uuid.uuid4().hex}"
|
||||
|
||||
try:
|
||||
# Register initial bank
|
||||
await vector_dbs_impl.register_vector_db(
|
||||
vector_db_id=vector_db_id,
|
||||
embedding_model=embedding_model,
|
||||
embedding_dimension=384,
|
||||
)
|
||||
|
||||
# Verify our bank exists
|
||||
response = await vector_dbs_impl.list_vector_dbs()
|
||||
assert isinstance(response, ListVectorDBsResponse)
|
||||
assert any(
|
||||
vector_db.vector_db_id == vector_db_id for vector_db in response.data
|
||||
)
|
||||
|
||||
# Try registering same bank again
|
||||
await vector_dbs_impl.register_vector_db(
|
||||
vector_db_id=vector_db_id,
|
||||
embedding_model=embedding_model,
|
||||
embedding_dimension=384,
|
||||
)
|
||||
|
||||
# Verify still only one instance of our bank
|
||||
response = await vector_dbs_impl.list_vector_dbs()
|
||||
assert isinstance(response, ListVectorDBsResponse)
|
||||
assert (
|
||||
len(
|
||||
[
|
||||
vector_db
|
||||
for vector_db in response.data
|
||||
if vector_db.vector_db_id == vector_db_id
|
||||
]
|
||||
)
|
||||
== 1
|
||||
)
|
||||
finally:
|
||||
# Clean up
|
||||
await vector_dbs_impl.unregister_vector_db(vector_db_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_documents(
|
||||
self, vector_io_stack, embedding_model, sample_chunks
|
||||
):
|
||||
vector_io_impl, vector_dbs_impl = vector_io_stack
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
await vector_io_impl.insert_chunks("test_vector_db", sample_chunks)
|
||||
|
||||
registered_db = await register_vector_db(vector_dbs_impl, embedding_model)
|
||||
await vector_io_impl.insert_chunks(registered_db.vector_db_id, sample_chunks)
|
||||
|
||||
query1 = "programming language"
|
||||
response1 = await vector_io_impl.query_chunks(
|
||||
registered_db.vector_db_id, query1
|
||||
)
|
||||
assert_valid_response(response1)
|
||||
assert any("Python" in chunk.content for chunk in response1.chunks)
|
||||
|
||||
# Test case 3: Query with semantic similarity
|
||||
query3 = "AI and brain-inspired computing"
|
||||
response3 = await vector_io_impl.query_chunks(
|
||||
registered_db.vector_db_id, query3
|
||||
)
|
||||
assert_valid_response(response3)
|
||||
assert any(
|
||||
"neural networks" in chunk.content.lower() for chunk in response3.chunks
|
||||
)
|
||||
|
||||
# Test case 4: Query with limit on number of results
|
||||
query4 = "computer"
|
||||
params4 = {"max_chunks": 2}
|
||||
response4 = await vector_io_impl.query_chunks(
|
||||
registered_db.vector_db_id, query4, params4
|
||||
)
|
||||
assert_valid_response(response4)
|
||||
assert len(response4.chunks) <= 2
|
||||
|
||||
# Test case 5: Query with threshold on similarity score
|
||||
query5 = "quantum computing" # Not directly related to any document
|
||||
params5 = {"score_threshold": 0.01}
|
||||
response5 = await vector_io_impl.query_chunks(
|
||||
registered_db.vector_db_id, query5, params5
|
||||
)
|
||||
assert_valid_response(response5)
|
||||
print("The scores are:", response5.scores)
|
||||
assert all(score >= 0.01 for score in response5.scores)
|
||||
|
||||
|
||||
def assert_valid_response(response: QueryChunksResponse):
|
||||
assert len(response.chunks) > 0
|
||||
assert len(response.scores) > 0
|
||||
assert len(response.chunks) == len(response.scores)
|
||||
for chunk in response.chunks:
|
||||
assert isinstance(chunk.content, str)
|
||||
|
|
@ -11,8 +11,11 @@ from pathlib import Path
|
|||
|
||||
import pytest
|
||||
|
||||
from llama_stack.apis.memory.memory import MemoryBankDocument, URL
|
||||
from llama_stack.providers.utils.memory.vector_store import content_from_doc
|
||||
from llama_stack.providers.utils.memory.vector_store import (
|
||||
content_from_doc,
|
||||
MemoryBankDocument,
|
||||
URL,
|
||||
)
|
||||
|
||||
DUMMY_PDF_PATH = Path(os.path.abspath(__file__)).parent / "fixtures" / "dummy.pdf"
|
||||
|
||||
Loading…
Add table
Add a link
Reference in a new issue