naas-abi-core 1.4.1__py3-none-any.whl

naas_abi_core/services/vector_store/IVectorStorePort_test.py

@@ -0,0 +1,189 @@
+import pytest
+import numpy as np
+from typing import List
+from abc import ABC
+from .IVectorStorePort import IVectorStorePort, VectorDocument
+
+
+class GenericVectorStoreAdapterTest(ABC):
+    @pytest.fixture
+    def adapter(self) -> IVectorStorePort:
+        raise NotImplementedError("Subclasses must provide an adapter fixture")
+
+    @pytest.fixture
+    def test_collection_name(self) -> str:
+        return "test_collection"
+
+    @pytest.fixture
+    def test_dimension(self) -> int:
+        return 128
+
+    @pytest.fixture
+    def sample_vectors(self) -> List[np.ndarray]:
+        np.random.seed(42)
+        return [
+            np.random.randn(128).astype(np.float32) for _ in range(5)
+        ]
+
+    @pytest.fixture
+    def sample_documents(self, sample_vectors) -> List[VectorDocument]:
+        return [
+            VectorDocument(
+                id=f"doc_{i}",
+                vector=vector,
+                metadata={"category": f"cat_{i % 2}", "index": i},
+                payload={"data": f"payload_{i}"}
+            )
+            for i, vector in enumerate(sample_vectors)
+        ]
+
+    def test_initialize(self, adapter):
+        adapter.initialize()
+        assert True
+
+    def test_create_and_list_collections(
+        self, adapter, test_collection_name, test_dimension
+    ):
+        adapter.initialize()
+        
+        adapter.create_collection(
+            test_collection_name, test_dimension, distance_metric="cosine"
+        )
+        
+        collections = adapter.list_collections()
+        assert test_collection_name in collections
+
+    def test_store_and_retrieve_vectors(
+        self, adapter, test_collection_name, test_dimension, sample_documents
+    ):
+        adapter.initialize()
+        
+        adapter.create_collection(test_collection_name, test_dimension)
+        
+        adapter.store_vectors(test_collection_name, sample_documents)
+        
+        retrieved = adapter.get_vector(
+            test_collection_name, "doc_0", include_vector=True
+        )
+        
+        assert retrieved is not None
+        assert retrieved.id == "doc_0"
+        assert retrieved.metadata["category"] == "cat_0"
+        assert retrieved.payload["data"] == "payload_0"
+        assert retrieved.vector is not None
+        np.testing.assert_array_almost_equal(
+            retrieved.vector, sample_documents[0].vector, decimal=5
+        )
+
+    def test_search_vectors(
+        self, adapter, test_collection_name, test_dimension, sample_documents
+    ):
+        adapter.initialize()
+        
+        adapter.create_collection(test_collection_name, test_dimension)
+        adapter.store_vectors(test_collection_name, sample_documents)
+        
+        query_vector = sample_documents[0].vector
+        results = adapter.search(
+            test_collection_name,
+            query_vector,
+            k=3,
+            include_metadata=True
+        )
+        
+        assert len(results) <= 3
+        assert results[0].id == "doc_0"
+        assert results[0].score >= 0.99
+
+    def test_search_with_filter(
+        self, adapter, test_collection_name, test_dimension, sample_documents
+    ):
+        adapter.initialize()
+        
+        adapter.create_collection(test_collection_name, test_dimension)
+        adapter.store_vectors(test_collection_name, sample_documents)
+        
+        query_vector = sample_documents[0].vector
+        results = adapter.search(
+            test_collection_name,
+            query_vector,
+            k=10,
+            filter={"category": "cat_0"},
+            include_metadata=True
+        )
+        
+        for result in results:
+            assert result.metadata["category"] == "cat_0"
+
+    def test_update_vector(
+        self, adapter, test_collection_name, test_dimension, sample_documents
+    ):
+        adapter.initialize()
+        
+        adapter.create_collection(test_collection_name, test_dimension)
+        adapter.store_vectors(test_collection_name, [sample_documents[0]])
+        
+        new_metadata = {"category": "updated", "new_field": "value"}
+        adapter.update_vector(
+            test_collection_name,
+            "doc_0",
+            metadata=new_metadata
+        )
+        
+        retrieved = adapter.get_vector(test_collection_name, "doc_0")
+        assert retrieved.metadata["category"] == "updated"
+        assert retrieved.metadata["new_field"] == "value"
+
+    def test_delete_vectors(
+        self, adapter, test_collection_name, test_dimension, sample_documents
+    ):
+        adapter.initialize()
+        
+        adapter.create_collection(test_collection_name, test_dimension)
+        adapter.store_vectors(test_collection_name, sample_documents)
+        
+        initial_count = adapter.count_vectors(test_collection_name)
+        assert initial_count == len(sample_documents)
+        
+        adapter.delete_vectors(test_collection_name, ["doc_0", "doc_1"])
+        
+        final_count = adapter.count_vectors(test_collection_name)
+        assert final_count == initial_count - 2
+        
+        deleted_doc = adapter.get_vector(test_collection_name, "doc_0")
+        assert deleted_doc is None
+
+    def test_count_vectors(
+        self, adapter, test_collection_name, test_dimension, sample_documents
+    ):
+        adapter.initialize()
+        
+        adapter.create_collection(test_collection_name, test_dimension)
+        
+        count_empty = adapter.count_vectors(test_collection_name)
+        assert count_empty == 0
+        
+        adapter.store_vectors(test_collection_name, sample_documents)
+        
+        count_filled = adapter.count_vectors(test_collection_name)
+        assert count_filled == len(sample_documents)
+
+    def test_delete_collection(
+        self, adapter, test_collection_name, test_dimension
+    ):
+        adapter.initialize()
+        
+        adapter.create_collection(test_collection_name, test_dimension)
+        
+        collections_before = adapter.list_collections()
+        assert test_collection_name in collections_before
+        
+        adapter.delete_collection(test_collection_name)
+        
+        collections_after = adapter.list_collections()
+        assert test_collection_name not in collections_after
+
+    def test_close(self, adapter):
+        adapter.initialize()
+        adapter.close()
+        assert True

naas_abi_core/services/vector_store/VectorStoreFactory.py

@@ -0,0 +1,47 @@
+import logging
+import os
+from typing import Optional
+
+from .IVectorStorePort import IVectorStorePort
+from .VectorStoreService import VectorStoreService
+
+logger = logging.getLogger(__name__)
+
+
+class VectorStoreFactory:
+    _instance: Optional[VectorStoreService] = None
+
+    @classmethod
+    def create_adapter(cls) -> IVectorStorePort:
+        adapter_type = os.getenv("VECTOR_STORE_ADAPTER", "qdrant").lower()
+
+        if adapter_type == "qdrant":
+            # Lazy import to avoid loading qdrant_client at module import time
+            from .adapters.QdrantAdapter import QdrantAdapter
+
+            host = os.getenv("QDRANT_HOST", "localhost")
+            port = int(os.getenv("QDRANT_PORT", "6333"))
+            api_key = os.getenv("QDRANT_API_KEY")
+            https = os.getenv("QDRANT_HTTPS", "false").lower() == "true"
+            timeout = int(os.getenv("QDRANT_TIMEOUT", "30"))
+
+            logger.info(f"Creating Qdrant adapter (host={host}, port={port})")
+            return QdrantAdapter(
+                host=host, port=port, api_key=api_key, https=https, timeout=timeout
+            )
+        else:
+            raise ValueError(f"Unknown vector store adapter: {adapter_type}")
+
+    @classmethod
+    def get_service(cls) -> VectorStoreService:
+        if cls._instance is None:
+            adapter = cls.create_adapter()
+            cls._instance = VectorStoreService(adapter)
+            logger.info("VectorStoreService singleton created")
+        return cls._instance
+
+    @classmethod
+    def reset(cls) -> None:
+        if cls._instance:
+            logger.info("Resetting VectorStoreService singleton")
+        cls._instance = None

naas_abi_core/services/vector_store/VectorStoreService.py

@@ -0,0 +1,171 @@
+import logging
+from typing import List, Dict, Any, Optional
+import numpy as np
+from .IVectorStorePort import IVectorStorePort, VectorDocument, SearchResult
+
+logger = logging.getLogger(__name__)
+
+
+class VectorStoreService:
+    def __init__(self, adapter: IVectorStorePort):
+        self.adapter = adapter
+        self._initialized = False
+
+    def initialize(self) -> None:
+        if not self._initialized:
+            self.adapter.initialize()
+            self._initialized = True
+            logger.info("VectorStoreService initialized successfully")
+
+    def ensure_collection(
+        self,
+        collection_name: str,
+        dimension: int,
+        distance_metric: str = "cosine",
+        recreate: bool = False,
+        **kwargs
+    ) -> None:
+        self.initialize()
+        
+        existing_collections = self.adapter.list_collections()
+        
+        if collection_name in existing_collections:
+            if recreate:
+                logger.info(f"Recreating collection: {collection_name}")
+                self.adapter.delete_collection(collection_name)
+                self.adapter.create_collection(
+                    collection_name, dimension, distance_metric, **kwargs
+                )
+            else:
+                logger.debug(f"Collection {collection_name} already exists")
+        else:
+            logger.info(f"Creating new collection: {collection_name}")
+            self.adapter.create_collection(
+                collection_name, dimension, distance_metric, **kwargs
+            )
+
+    def add_documents(
+        self,
+        collection_name: str,
+        ids: List[str],
+        vectors: List[np.ndarray],
+        metadata: Optional[List[Dict[str, Any]]] = None,
+        payloads: Optional[List[Dict[str, Any]]] = None
+    ) -> None:
+        self.initialize()
+        
+        if not ids or not vectors:
+            raise ValueError("IDs and vectors cannot be empty")
+        
+        if len(ids) != len(vectors):
+            raise ValueError("Number of IDs must match number of vectors")
+        
+        if metadata and len(metadata) != len(ids):
+            raise ValueError("Number of metadata entries must match number of IDs")
+        
+        if payloads and len(payloads) != len(ids):
+            raise ValueError("Number of payloads must match number of IDs")
+        
+        documents = []
+        for i, (doc_id, vector) in enumerate(zip(ids, vectors)):
+            doc = VectorDocument(
+                id=doc_id,
+                vector=vector,
+                metadata=metadata[i] if metadata else {},
+                payload=payloads[i] if payloads else None
+            )
+            documents.append(doc)
+        
+        self.adapter.store_vectors(collection_name, documents)
+        logger.debug(f"Added {len(documents)} documents to collection {collection_name}")
+
+    def search_similar(
+        self,
+        collection_name: str,
+        query_vector: np.ndarray,
+        k: int = 10,
+        filter: Optional[Dict[str, Any]] = None,
+        score_threshold: Optional[float] = None,
+        include_vectors: bool = False,
+        include_metadata: bool = True
+    ) -> List[SearchResult]:
+        self.initialize()
+        
+        if k <= 0:
+            raise ValueError("k must be a positive integer")
+        
+        results = self.adapter.search(
+            collection_name=collection_name,
+            query_vector=query_vector,
+            k=k,
+            filter=filter,
+            include_vectors=include_vectors,
+            include_metadata=include_metadata
+        )
+        
+        if score_threshold is not None:
+            results = [r for r in results if r.score >= score_threshold]
+        
+        logger.debug(f"Found {len(results)} similar vectors in {collection_name}")
+        return results
+
+    def get_document(
+        self,
+        collection_name: str,
+        document_id: str,
+        include_vector: bool = True
+    ) -> Optional[VectorDocument]:
+        self.initialize()
+        return self.adapter.get_vector(
+            collection_name, document_id, include_vector
+        )
+
+    def update_document(
+        self,
+        collection_name: str,
+        document_id: str,
+        vector: Optional[np.ndarray] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        payload: Optional[Dict[str, Any]] = None
+    ) -> None:
+        self.initialize()
+        
+        if vector is None and metadata is None and payload is None:
+            raise ValueError("At least one of vector, metadata, or payload must be provided")
+        
+        self.adapter.update_vector(
+            collection_name, document_id, vector, metadata, payload
+        )
+        logger.debug(f"Updated document {document_id} in collection {collection_name}")
+
+    def delete_documents(
+        self,
+        collection_name: str,
+        document_ids: List[str]
+    ) -> None:
+        self.initialize()
+        
+        if not document_ids:
+            raise ValueError("Document IDs cannot be empty")
+        
+        self.adapter.delete_vectors(collection_name, document_ids)
+        logger.info(f"Deleted {len(document_ids)} documents from collection {collection_name}")
+
+    def get_collection_size(self, collection_name: str) -> int:
+        self.initialize()
+        return self.adapter.count_vectors(collection_name)
+
+    def list_collections(self) -> List[str]:
+        self.initialize()
+        return self.adapter.list_collections()
+
+    def delete_collection(self, collection_name: str) -> None:
+        self.initialize()
+        self.adapter.delete_collection(collection_name)
+        logger.info(f"Deleted collection: {collection_name}")
+
+    def close(self) -> None:
+        if self._initialized:
+            self.adapter.close()
+            self._initialized = False
+            logger.info("VectorStoreService closed")

naas_abi_core/services/vector_store/VectorStoreService_test.py

@@ -0,0 +1,185 @@
+import pytest
+from unittest.mock import Mock
+import numpy as np
+from .VectorStoreService import VectorStoreService
+from .IVectorStorePort import IVectorStorePort, VectorDocument, SearchResult
+
+
+class TestVectorStoreService:
+    @pytest.fixture
+    def mock_adapter(self):
+        adapter = Mock(spec=IVectorStorePort)
+        adapter.initialize = Mock()
+        adapter.list_collections = Mock(return_value=[])
+        adapter.create_collection = Mock()
+        adapter.delete_collection = Mock()
+        adapter.store_vectors = Mock()
+        adapter.search = Mock()
+        adapter.get_vector = Mock()
+        adapter.update_vector = Mock()
+        adapter.delete_vectors = Mock()
+        adapter.count_vectors = Mock()
+        adapter.close = Mock()
+        return adapter
+
+    @pytest.fixture
+    def service(self, mock_adapter):
+        return VectorStoreService(mock_adapter)
+
+    def test_initialize(self, service, mock_adapter):
+        service.initialize()
+        mock_adapter.initialize.assert_called_once()
+        
+        service.initialize()
+        mock_adapter.initialize.assert_called_once()
+
+    def test_ensure_collection_create_new(self, service, mock_adapter):
+        mock_adapter.list_collections.return_value = []
+        
+        service.ensure_collection("test_collection", 128)
+        
+        mock_adapter.create_collection.assert_called_once_with(
+            "test_collection", 128, "cosine"
+        )
+
+    def test_ensure_collection_exists_no_recreate(self, service, mock_adapter):
+        mock_adapter.list_collections.return_value = ["test_collection"]
+        
+        service.ensure_collection("test_collection", 128, recreate=False)
+        
+        mock_adapter.create_collection.assert_not_called()
+        mock_adapter.delete_collection.assert_not_called()
+
+    def test_ensure_collection_recreate(self, service, mock_adapter):
+        mock_adapter.list_collections.return_value = ["test_collection"]
+        
+        service.ensure_collection("test_collection", 128, recreate=True)
+        
+        mock_adapter.delete_collection.assert_called_once_with("test_collection")
+        mock_adapter.create_collection.assert_called_once_with(
+            "test_collection", 128, "cosine"
+        )
+
+    def test_add_documents(self, service, mock_adapter):
+        ids = ["doc1", "doc2"]
+        vectors = [np.random.randn(128), np.random.randn(128)]
+        metadata = [{"key": "value1"}, {"key": "value2"}]
+        
+        service.add_documents("test_collection", ids, vectors, metadata)
+        
+        mock_adapter.store_vectors.assert_called_once()
+        call_args = mock_adapter.store_vectors.call_args
+        assert call_args[0][0] == "test_collection"
+        assert len(call_args[0][1]) == 2
+        assert call_args[0][1][0].id == "doc1"
+
+    def test_add_documents_validation(self, service, mock_adapter):
+        with pytest.raises(ValueError, match="IDs and vectors cannot be empty"):
+            service.add_documents("test_collection", [], [])
+        
+        with pytest.raises(ValueError, match="Number of IDs must match"):
+            service.add_documents(
+                "test_collection",
+                ["doc1"],
+                [np.random.randn(128), np.random.randn(128)]
+            )
+
+    def test_search_similar(self, service, mock_adapter):
+        mock_results = [
+            SearchResult(id="doc1", score=0.95, metadata={"key": "value1"}),
+            SearchResult(id="doc2", score=0.85, metadata={"key": "value2"}),
+        ]
+        mock_adapter.search.return_value = mock_results
+        
+        query_vector = np.random.randn(128)
+        results = service.search_similar(
+            "test_collection", query_vector, k=5
+        )
+        
+        assert len(results) == 2
+        assert results[0].id == "doc1"
+        mock_adapter.search.assert_called_once()
+
+    def test_search_similar_with_threshold(self, service, mock_adapter):
+        mock_results = [
+            SearchResult(id="doc1", score=0.95),
+            SearchResult(id="doc2", score=0.85),
+            SearchResult(id="doc3", score=0.75),
+        ]
+        mock_adapter.search.return_value = mock_results
+        
+        query_vector = np.random.randn(128)
+        results = service.search_similar(
+            "test_collection", query_vector, k=5, score_threshold=0.8
+        )
+        
+        assert len(results) == 2
+        assert all(r.score >= 0.8 for r in results)
+
+    def test_get_document(self, service, mock_adapter):
+        mock_doc = VectorDocument(
+            id="doc1",
+            vector=np.random.randn(128),
+            metadata={"key": "value"}
+        )
+        mock_adapter.get_vector.return_value = mock_doc
+        
+        doc = service.get_document("test_collection", "doc1")
+        
+        assert doc.id == "doc1"
+        mock_adapter.get_vector.assert_called_once_with(
+            "test_collection", "doc1", True
+        )
+
+    def test_update_document(self, service, mock_adapter):
+        new_metadata = {"key": "new_value"}
+        
+        service.update_document(
+            "test_collection", "doc1", metadata=new_metadata
+        )
+        
+        mock_adapter.update_vector.assert_called_once_with(
+            "test_collection", "doc1", None, new_metadata, None
+        )
+
+    def test_update_document_validation(self, service, mock_adapter):
+        with pytest.raises(ValueError, match="At least one of"):
+            service.update_document("test_collection", "doc1")
+
+    def test_delete_documents(self, service, mock_adapter):
+        service.delete_documents("test_collection", ["doc1", "doc2"])
+        
+        mock_adapter.delete_vectors.assert_called_once_with(
+            "test_collection", ["doc1", "doc2"]
+        )
+
+    def test_delete_documents_validation(self, service, mock_adapter):
+        with pytest.raises(ValueError, match="Document IDs cannot be empty"):
+            service.delete_documents("test_collection", [])
+
+    def test_get_collection_size(self, service, mock_adapter):
+        mock_adapter.count_vectors.return_value = 42
+        
+        count = service.get_collection_size("test_collection")
+        
+        assert count == 42
+        mock_adapter.count_vectors.assert_called_once_with("test_collection")
+
+    def test_list_collections(self, service, mock_adapter):
+        mock_adapter.list_collections.return_value = ["coll1", "coll2"]
+        
+        collections = service.list_collections()
+        
+        assert collections == ["coll1", "coll2"]
+
+    def test_delete_collection(self, service, mock_adapter):
+        service.delete_collection("test_collection")
+        
+        mock_adapter.delete_collection.assert_called_once_with("test_collection")
+
+    def test_close(self, service, mock_adapter):
+        service.initialize()
+        service.close()
+        
+        mock_adapter.close.assert_called_once()
+        assert not service._initialized

naas_abi_core/services/vector_store/__init__.py

@@ -0,0 +1,13 @@
+# from .IVectorStorePort import IVectorStorePort, VectorDocument, SearchResult
+# from .VectorStoreService import VectorStoreService
+# from .VectorStoreFactory import VectorStoreFactory
+# from .adapters.QdrantAdapter import QdrantAdapter
+
+# __all__ = [
+#     "IVectorStorePort",
+#     "VectorDocument",
+#     "SearchResult",
+#     "VectorStoreService",
+#     "VectorStoreFactory",
+#     "QdrantAdapter",
+# ]