milvus-haystack 0.0.7__tar.gz → 0.0.8__tar.gz

{milvus_haystack-0.0.7 → milvus_haystack-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: milvus-haystack
-Version: 0.0.7
+Version: 0.0.8
 Project-URL: Documentation, https://github.com/milvus-io/milvus-haystack#readme
 Project-URL: Issues, https://github.com/milvus-io/milvus-haystack/issues
 Project-URL: Source, https://github.com/milvus-io/milvus-haystack

{milvus_haystack-0.0.7 → milvus_haystack-0.0.8}/src/milvus_haystack/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2023-present Tuana Celik <tuana.celik@deepset.ai>
 #
 # SPDX-License-Identifier: Apache-2.0
-__version__ = "0.0.7"
+__version__ = "0.0.8"

{milvus_haystack-0.0.7 → milvus_haystack-0.0.8}/src/milvus_haystack/milvus_embedding_retriever.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Optional
 
-from haystack import Document, component
+from haystack import DeserializationError, Document, component, default_from_dict, default_to_dict
 
 from milvus_haystack import MilvusDocumentStore
 
@@ -23,6 +23,35 @@ class MilvusEmbeddingRetriever:
         self.top_k = top_k
         self.document_store = document_store
 
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Returns a dictionary representation of the retriever component.
+
+        :returns:
+            A dictionary representation of the retriever component.
+        """
+        return default_to_dict(
+            self, document_store=self.document_store.to_dict(), filters=self.filters, top_k=self.top_k
+        )
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "MilvusEmbeddingRetriever":
+        """
+        Creates a new retriever from a dictionary.
+
+        :param data: The dictionary to use to create the retriever.
+        :return: A new retriever.
+        """
+        init_params = data.get("init_parameters", {})
+        if "document_store" not in init_params:
+            err_msg = "Missing 'document_store' in serialization data"
+            raise DeserializationError(err_msg)
+
+        docstore = MilvusDocumentStore.from_dict(init_params["document_store"])
+        data["init_parameters"]["document_store"] = docstore
+
+        return default_from_dict(cls, data)
+
     @component.output_types(documents=List[Document])
     def run(self, query_embedding: List[float]) -> Dict[str, List[Document]]:
         """

{milvus_haystack-0.0.7 → milvus_haystack-0.0.8}/tests/test_document_store.py

@@ -10,18 +10,17 @@ from src.milvus_haystack import MilvusDocumentStore
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_CONNECTION_ARGS = {
+    "uri": "http://localhost:19530",
+    # "uri": "./milvus_test.db",  # This uri works for Milvus Lite
+}
+
 
 class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsTest):
     @pytest.fixture
     def document_store(self) -> MilvusDocumentStore:
         return MilvusDocumentStore(
-            connection_args={
-                "host": "localhost",
-                "port": "19530",
-                "user": "",
-                "password": "",
-                "secure": False,
-            },
+            connection_args=DEFAULT_CONNECTION_ARGS,
             drop_old=True,
         )
 
@@ -61,7 +60,7 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
                 "collection_name": "HaystackCollection",
                 "collection_description": "",
                 "collection_properties": None,
-                "connection_args": {"host": "localhost", "port": "19530", "user": "", "password": "", "secure": False},
+                "connection_args": DEFAULT_CONNECTION_ARGS,
                 "consistency_level": "Session",
                 "index_params": None,
                 "search_params": None,
@@ -78,6 +77,6 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
         assert document_store_dict == expected_dict
         reconstructed_document_store = MilvusDocumentStore.from_dict(document_store_dict)
         for field in vars(reconstructed_document_store):
-            if field.startswith("__"):
+            if field.startswith("__") or field == "alias":
                 continue
             assert getattr(reconstructed_document_store, field) == getattr(document_store, field)

milvus_haystack-0.0.8/tests/test_embedding_retriever.py

@@ -0,0 +1,117 @@
+import logging
+
+import pytest
+from haystack import Document
+
+from src.milvus_haystack import MilvusDocumentStore
+from src.milvus_haystack.milvus_embedding_retriever import MilvusEmbeddingRetriever
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_CONNECTION_ARGS = {
+    "uri": "http://localhost:19530",
+    # "uri": "./milvus_test.db",  # This uri works for Milvus Lite
+}
+
+
+class TestMilvusEmbeddingTests:
+    @pytest.fixture
+    def document_store(self) -> MilvusDocumentStore:
+        return MilvusDocumentStore(
+            connection_args=DEFAULT_CONNECTION_ARGS,
+            drop_old=True,
+        )
+
+    def test_run(self, document_store: MilvusDocumentStore):
+        documents = []
+        doc = Document(
+            content="A Foo Document",
+            meta={
+                "name": "name_0",
+                "page": "100",
+                "chapter": "intro",
+                "number": 2,
+                "date": "1969-07-21T20:17:40",
+            },
+            embedding=[-10.0] * 128,
+        )
+        documents.append(doc)
+        document_store.write_documents(documents)
+        retriever = MilvusEmbeddingRetriever(
+            document_store,
+        )
+        query_embedding = [-10.0] * 128
+        res = retriever.run(query_embedding)
+        assert res["documents"] == documents
+
+    def test_to_dict(self, document_store: MilvusDocumentStore):
+        expected_dict = {
+            "type": "src.milvus_haystack.document_store.MilvusDocumentStore",
+            "init_parameters": {
+                "collection_name": "HaystackCollection",
+                "collection_description": "",
+                "collection_properties": None,
+                "connection_args": DEFAULT_CONNECTION_ARGS,
+                "consistency_level": "Session",
+                "index_params": None,
+                "search_params": None,
+                "drop_old": True,
+                "primary_field": "id",
+                "text_field": "text",
+                "vector_field": "vector",
+                "partition_key_field": None,
+                "partition_names": None,
+                "replica_number": 1,
+                "timeout": None,
+            },
+        }
+        retriever = MilvusEmbeddingRetriever(document_store)
+        result = retriever.to_dict()
+
+        assert result["type"] == "src.milvus_haystack.milvus_embedding_retriever.MilvusEmbeddingRetriever"
+        assert result["init_parameters"]["document_store"] == expected_dict
+
+    def test_from_dict(self, document_store: MilvusDocumentStore):
+        retriever_dict = {
+            "type": "src.milvus_haystack.milvus_embedding_retriever.MilvusEmbeddingRetriever",
+            "init_parameters": {
+                "document_store": {
+                    "type": "milvus_haystack.document_store.MilvusDocumentStore",
+                    "init_parameters": {
+                        "collection_name": "HaystackCollection",
+                        "collection_description": "",
+                        "collection_properties": None,
+                        "connection_args": DEFAULT_CONNECTION_ARGS,
+                        "consistency_level": "Session",
+                        "index_params": None,
+                        "search_params": None,
+                        "drop_old": True,
+                        "primary_field": "id",
+                        "text_field": "text",
+                        "vector_field": "vector",
+                        "partition_key_field": None,
+                        "partition_names": None,
+                        "replica_number": 1,
+                        "timeout": None,
+                    },
+                },
+                "filters": None,
+                "top_k": 10,
+            },
+        }
+
+        retriever = MilvusEmbeddingRetriever(document_store)
+
+        reconstructed_retriever = MilvusEmbeddingRetriever.from_dict(retriever_dict)
+        for field in vars(reconstructed_retriever):
+            if field.startswith("__"):
+                continue
+            elif field == "document_store":
+                for doc_store_field in vars(document_store):
+                    if doc_store_field.startswith("__") or doc_store_field == "alias":
+                        continue
+                    assert getattr(reconstructed_retriever.document_store, doc_store_field) == getattr(
+                        document_store, doc_store_field
+                    )
+            else:
+                assert getattr(reconstructed_retriever, field) == getattr(retriever, field)

milvus_haystack-0.0.7/tests/test_embedding_retriever.py

@@ -1,46 +0,0 @@
-import logging
-
-import pytest
-from haystack import Document
-
-from src.milvus_haystack import MilvusDocumentStore
-from src.milvus_haystack.milvus_embedding_retriever import MilvusEmbeddingRetriever
-
-logger = logging.getLogger(__name__)
-
-
-class TestMilvusEmbeddingTests:
-    @pytest.fixture
-    def document_store(self) -> MilvusDocumentStore:
-        return MilvusDocumentStore(
-            connection_args={
-                "host": "localhost",
-                "port": "19530",
-                "user": "",
-                "password": "",
-                "secure": False,
-            },
-            drop_old=True,
-        )
-
-    def test_run(self, document_store: MilvusDocumentStore):
-        documents = []
-        doc = Document(
-            content="A Foo Document",
-            meta={
-                "name": "name_0",
-                "page": "100",
-                "chapter": "intro",
-                "number": 2,
-                "date": "1969-07-21T20:17:40",
-            },
-            embedding=[-10.0] * 128,
-        )
-        documents.append(doc)
-        document_store.write_documents(documents)
-        retriever = MilvusEmbeddingRetriever(
-            document_store,
-        )
-        query_embedding = [-10.0] * 128
-        res = retriever.run(query_embedding)
-        assert res["documents"] == documents