diff --git a/pyproject.toml b/pyproject.toml index 5f086bd9d..fed424003 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,19 +87,14 @@ unit = [ "pypdf", "mcp", "chardet", - "qdrant-client", "sqlalchemy", "sqlalchemy[asyncio]>=2.0.41", "blobfile", "faiss-cpu", - "pymilvus>=2.6.1", - "milvus-lite>=2.5.0", "litellm", "together", "coverage", - "chromadb>=1.0.15", "moto[s3]>=5.1.10", - "weaviate-client>=4.16.4", ] # These are the core dependencies required for running integration tests. They are shared across all # providers. If a provider requires additional dependencies, please add them to your environment diff --git a/tests/unit/providers/vector_io/conftest.py b/tests/unit/providers/vector_io/conftest.py index f5a72da0d..485164194 100644 --- a/tests/unit/providers/vector_io/conftest.py +++ b/tests/unit/providers/vector_io/conftest.py @@ -9,27 +9,22 @@ from unittest.mock import AsyncMock, MagicMock, patch import numpy as np import pytest -from chromadb import PersistentClient from llama_stack.apis.vector_dbs import VectorDB from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse -from llama_stack.providers.inline.vector_io.chroma.config import ChromaVectorIOConfig from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig from llama_stack.providers.inline.vector_io.faiss.faiss import FaissIndex, FaissVectorIOAdapter -from llama_stack.providers.inline.vector_io.milvus.config import SqliteKVStoreConfig -from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig from llama_stack.providers.inline.vector_io.sqlite_vec import SQLiteVectorIOConfig from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import SQLiteVecIndex, SQLiteVecVectorIOAdapter -from llama_stack.providers.remote.vector_io.chroma.chroma import ChromaIndex, ChromaVectorIOAdapter, maybe_await from llama_stack.providers.remote.vector_io.pgvector.config import PGVectorVectorIOConfig from llama_stack.providers.remote.vector_io.pgvector.pgvector import PGVectorIndex, PGVectorVectorIOAdapter -from llama_stack.providers.remote.vector_io.qdrant.qdrant import QdrantVectorIOAdapter +from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig EMBEDDING_DIMENSION = 384 COLLECTION_PREFIX = "test_collection" -@pytest.fixture(params=["sqlite_vec", "faiss", "chroma", "pgvector"]) +@pytest.fixture(params=["sqlite_vec", "faiss", "pgvector"]) def vector_provider(request): return request.param @@ -201,98 +196,6 @@ async def faiss_vec_adapter(unique_kvstore_config, mock_inference_api, embedding await adapter.shutdown() -@pytest.fixture -def chroma_vec_db_path(tmp_path_factory): - persist_dir = tmp_path_factory.mktemp(f"chroma_{np.random.randint(1e6)}") - return str(persist_dir) - - -@pytest.fixture -async def chroma_vec_index(chroma_vec_db_path, embedding_dimension): - client = PersistentClient(path=chroma_vec_db_path) - name = f"{COLLECTION_PREFIX}_{np.random.randint(1e6)}" - collection = await maybe_await(client.get_or_create_collection(name)) - index = ChromaIndex(client=client, collection=collection) - await index.initialize() - yield index - await index.delete() - - -@pytest.fixture -async def chroma_vec_adapter(chroma_vec_db_path, unique_kvstore_config, mock_inference_api, embedding_dimension): - config = ChromaVectorIOConfig( - db_path=chroma_vec_db_path, - kvstore=unique_kvstore_config, - ) - adapter = ChromaVectorIOAdapter( - config=config, - inference_api=mock_inference_api, - files_api=None, - ) - await adapter.initialize() - await adapter.register_vector_db( - VectorDB( - identifier=f"chroma_test_collection_{random.randint(1, 1_000_000)}", - provider_id="test_provider", - embedding_model="test_model", - embedding_dimension=embedding_dimension, - ) - ) - yield adapter - await adapter.shutdown() - - -@pytest.fixture -def qdrant_vec_db_path(tmp_path_factory): - import uuid - - db_path = str(tmp_path_factory.getbasetemp() / f"test_qdrant_{uuid.uuid4()}.db") - return db_path - - -@pytest.fixture -async def qdrant_vec_adapter(qdrant_vec_db_path, unique_kvstore_config, mock_inference_api, embedding_dimension): - import uuid - - config = QdrantVectorIOConfig( - db_path=qdrant_vec_db_path, - kvstore=unique_kvstore_config, - ) - adapter = QdrantVectorIOAdapter( - config=config, - inference_api=mock_inference_api, - files_api=None, - ) - collection_id = f"qdrant_test_collection_{uuid.uuid4()}" - await adapter.initialize() - await adapter.register_vector_db( - VectorDB( - identifier=collection_id, - provider_id="test_provider", - embedding_model="test_model", - embedding_dimension=embedding_dimension, - ) - ) - adapter.test_collection_id = collection_id - yield adapter - await adapter.shutdown() - - -@pytest.fixture -async def qdrant_vec_index(qdrant_vec_db_path, embedding_dimension): - import uuid - - from qdrant_client import AsyncQdrantClient - - from llama_stack.providers.remote.vector_io.qdrant.qdrant import QdrantIndex - - client = AsyncQdrantClient(path=qdrant_vec_db_path) - collection_name = f"qdrant_test_collection_{uuid.uuid4()}" - index = QdrantIndex(client, collection_name) - yield index - await index.delete() - - @pytest.fixture def mock_psycopg2_connection(): connection = MagicMock() @@ -410,8 +313,6 @@ def vector_io_adapter(vector_provider, request): vector_provider_dict = { "faiss": "faiss_vec_adapter", "sqlite_vec": "sqlite_vec_adapter", - "chroma": "chroma_vec_adapter", - "qdrant": "qdrant_vec_adapter", "pgvector": "pgvector_vec_adapter", } return request.getfixturevalue(vector_provider_dict[vector_provider]) diff --git a/tests/unit/providers/vector_io/test_qdrant.py b/tests/unit/providers/vector_io/test_qdrant.py deleted file mode 100644 index aab5b6f45..000000000 --- a/tests/unit/providers/vector_io/test_qdrant.py +++ /dev/null @@ -1,147 +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 asyncio -import os -from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from llama_stack.apis.inference import Inference -from llama_stack.apis.inference.inference import OpenAIEmbeddingData, OpenAIEmbeddingsResponse, OpenAIEmbeddingUsage -from llama_stack.apis.vector_io import ( - QueryChunksResponse, - VectorDB, - VectorDBStore, -) -from llama_stack.providers.inline.vector_io.qdrant.config import ( - QdrantVectorIOConfig as InlineQdrantVectorIOConfig, -) -from llama_stack.providers.remote.vector_io.qdrant.qdrant import ( - QdrantVectorIOAdapter, -) -from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig - -# This test is a unit test for the QdrantVectorIOAdapter class. This should only contain -# tests which are specific to this class. More general (API-level) tests should be placed in -# tests/integration/vector_io/ -# -# How to run this test: -# -# pytest tests/unit/providers/vector_io/test_qdrant.py \ -# -v -s --tb=short --disable-warnings --asyncio-mode=auto - - -@pytest.fixture -def qdrant_config(tmp_path) -> InlineQdrantVectorIOConfig: - kvstore_config = SqliteKVStoreConfig(db_name=os.path.join(tmp_path, "test_kvstore.db")) - return InlineQdrantVectorIOConfig(path=os.path.join(tmp_path, "qdrant.db"), kvstore=kvstore_config) - - -@pytest.fixture(scope="session") -def loop(): - return asyncio.new_event_loop() - - -@pytest.fixture -def mock_vector_db(vector_db_id) -> MagicMock: - mock_vector_db = MagicMock(spec=VectorDB) - mock_vector_db.embedding_model = "embedding_model" - mock_vector_db.identifier = vector_db_id - mock_vector_db.embedding_dimension = 384 - mock_vector_db.model_dump_json.return_value = ( - '{"identifier": "' - + vector_db_id - + '", "provider_id": "qdrant", "embedding_model": "embedding_model", "embedding_dimension": 384}' - ) - return mock_vector_db - - -@pytest.fixture -def mock_vector_db_store(mock_vector_db) -> MagicMock: - mock_store = MagicMock(spec=VectorDBStore) - mock_store.get_vector_db = AsyncMock(return_value=mock_vector_db) - return mock_store - - -@pytest.fixture -def mock_api_service(sample_embeddings): - mock_api_service = MagicMock(spec=Inference) - mock_api_service.openai_embeddings = AsyncMock( - return_value=OpenAIEmbeddingsResponse( - model="mock-embedding-model", - data=[OpenAIEmbeddingData(embedding=sample, index=i) for i, sample in enumerate(sample_embeddings)], - usage=OpenAIEmbeddingUsage(prompt_tokens=10, total_tokens=10), - ) - ) - return mock_api_service - - -@pytest.fixture -async def qdrant_adapter(qdrant_config, mock_vector_db_store, mock_api_service, loop) -> QdrantVectorIOAdapter: - adapter = QdrantVectorIOAdapter(config=qdrant_config, inference_api=mock_api_service, files_api=None) - adapter.vector_db_store = mock_vector_db_store - await adapter.initialize() - yield adapter - await adapter.shutdown() - - -__QUERY = "Sample query" - - -@pytest.mark.parametrize("max_query_chunks, expected_chunks", [(2, 2), (100, 60)]) -async def test_qdrant_adapter_returns_expected_chunks( - qdrant_adapter: QdrantVectorIOAdapter, - vector_db_id, - sample_chunks, - sample_embeddings, - max_query_chunks, - expected_chunks, -) -> None: - assert qdrant_adapter is not None - await qdrant_adapter.insert_chunks(vector_db_id, sample_chunks) - - index = await qdrant_adapter._get_and_cache_vector_db_index(vector_db_id=vector_db_id) - assert index is not None - - response = await qdrant_adapter.query_chunks( - query=__QUERY, - vector_db_id=vector_db_id, - params={"max_chunks": max_query_chunks, "mode": "vector"}, - ) - assert isinstance(response, QueryChunksResponse) - assert len(response.chunks) == expected_chunks - - -# To by-pass attempt to convert a Mock to JSON -def _prepare_for_json(value: Any) -> str: - return str(value) - - -@patch("llama_stack.providers.utils.telemetry.trace_protocol._prepare_for_json", new=_prepare_for_json) -async def test_qdrant_register_and_unregister_vector_db( - qdrant_adapter: QdrantVectorIOAdapter, - mock_vector_db, - sample_chunks, -) -> None: - # Initially, no collections - vector_db_id = mock_vector_db.identifier - assert len((await qdrant_adapter.client.get_collections()).collections) == 0 - - # Register does not create a collection - assert not (await qdrant_adapter.client.collection_exists(vector_db_id)) - await qdrant_adapter.register_vector_db(mock_vector_db) - assert not (await qdrant_adapter.client.collection_exists(vector_db_id)) - - # First insert creates the collection - await qdrant_adapter.insert_chunks(vector_db_id, sample_chunks) - assert await qdrant_adapter.client.collection_exists(vector_db_id) - - # Unregister deletes the collection - await qdrant_adapter.unregister_vector_db(vector_db_id) - assert not (await qdrant_adapter.client.collection_exists(vector_db_id)) - assert len((await qdrant_adapter.client.get_collections()).collections) == 0 diff --git a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py index 6cd60fe7e..ed0934224 100644 --- a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py +++ b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py @@ -104,12 +104,8 @@ async def test_register_and_unregister_vector_db(vector_io_adapter): async def test_query_unregistered_raises(vector_io_adapter, vector_provider): fake_emb = np.zeros(8, dtype=np.float32) - if vector_provider == "chroma": - with pytest.raises(AttributeError): - await vector_io_adapter.query_chunks("no_such_db", fake_emb) - else: - with pytest.raises(ValueError): - await vector_io_adapter.query_chunks("no_such_db", fake_emb) + with pytest.raises(ValueError): + await vector_io_adapter.query_chunks("no_such_db", fake_emb) async def test_insert_chunks_calls_underlying_index(vector_io_adapter):