elasticsearch-haystack 0.4.0__tar.gz → 0.6.0__tar.gz

{elasticsearch_haystack-0.4.0 → elasticsearch_haystack-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: elasticsearch-haystack
-Version: 0.4.0
+Version: 0.6.0
 Summary: Haystack 2.x Document Store for ElasticSearch
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/elasticsearch#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
@@ -9,6 +9,7 @@ Author-email: Silvano Cerza <silvanocerza@gmail.com>
 License-Expression: Apache-2.0
 License-File: LICENSE
 Classifier: Development Status :: 4 - Beta
+Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9

{elasticsearch_haystack-0.4.0 → elasticsearch_haystack-0.6.0}/pydoc/config.yml

@@ -17,7 +17,7 @@ processors:
   - type: smart
   - type: crossref
 renderer:
-  type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer
+  type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer
   excerpt: Elasticsearch integration for Haystack
   category_slug: integrations-api
   title: Elasticsearch

{elasticsearch_haystack-0.4.0 → elasticsearch_haystack-0.6.0}/pyproject.toml

@@ -14,6 +14,7 @@ authors = [
   { name = "Silvano Cerza", email = "silvanocerza@gmail.com" },
 ]
 classifiers = [
+  "License :: OSI Approved :: Apache Software License",
   "Development Status :: 4 - Beta",
   "Programming Language :: Python",
   "Programming Language :: Python :: 3.8",

{elasticsearch_haystack-0.4.0 → elasticsearch_haystack-0.6.0}/src/haystack_integrations/document_stores/elasticsearch/document_store.py

@@ -63,6 +63,7 @@ class ElasticsearchDocumentStore:
         self,
         *,
         hosts: Optional[Hosts] = None,
+        custom_mapping: Optional[Dict[str, Any]] = None,
         index: str = "default",
         embedding_similarity_function: Literal["cosine", "dot_product", "l2_norm", "max_inner_product"] = "cosine",
         **kwargs,
@@ -82,6 +83,7 @@ class ElasticsearchDocumentStore:
         [reference](https://elasticsearch-py.readthedocs.io/en/stable/api.html#module-elasticsearch)
 
         :param hosts: List of hosts running the Elasticsearch client.
+        :param custom_mapping: Custom mapping for the index. If not provided, a default mapping will be used.
         :param index: Name of index in Elasticsearch.
         :param embedding_similarity_function: The similarity function used to compare Documents embeddings.
             This parameter only takes effect if the index does not yet exist and is created.
@@ -91,40 +93,60 @@ class ElasticsearchDocumentStore:
         :param **kwargs: Optional arguments that `Elasticsearch` takes.
         """
         self._hosts = hosts
-        self._client = Elasticsearch(
-            hosts,
-            headers={"user-agent": f"haystack-py-ds/{haystack_version}"},
-            **kwargs,
-        )
+        self._client = None
         self._index = index
         self._embedding_similarity_function = embedding_similarity_function
+        self._custom_mapping = custom_mapping
         self._kwargs = kwargs
 
-        # Check client connection, this will raise if not connected
-        self._client.info()
+        if self._custom_mapping and not isinstance(self._custom_mapping, Dict):
+            msg = "custom_mapping must be a dictionary"
+            raise ValueError(msg)
 
-        # configure mapping for the embedding field
-        mappings = {
-            "properties": {
-                "embedding": {"type": "dense_vector", "index": True, "similarity": embedding_similarity_function},
-                "content": {"type": "text"},
-            },
-            "dynamic_templates": [
-                {
-                    "strings": {
-                        "path_match": "*",
-                        "match_mapping_type": "string",
-                        "mapping": {
-                            "type": "keyword",
+    @property
+    def client(self) -> Elasticsearch:
+        if self._client is None:
+            client = Elasticsearch(
+                self._hosts,
+                headers={"user-agent": f"haystack-py-ds/{haystack_version}"},
+                **self._kwargs,
+            )
+            # Check client connection, this will raise if not connected
+            client.info()
+
+            if self._custom_mapping:
+                mappings = self._custom_mapping
+            else:
+                # Configure mapping for the embedding field if none is provided
+                mappings = {
+                    "properties": {
+                        "embedding": {
+                            "type": "dense_vector",
+                            "index": True,
+                            "similarity": self._embedding_similarity_function,
                         },
-                    }
+                        "content": {"type": "text"},
+                    },
+                    "dynamic_templates": [
+                        {
+                            "strings": {
+                                "path_match": "*",
+                                "match_mapping_type": "string",
+                                "mapping": {
+                                    "type": "keyword",
+                                },
+                            }
+                        }
+                    ],
                 }
-            ],
-        }
 
-        # Create the index if it doesn't exist
-        if not self._client.indices.exists(index=index):
-            self._client.indices.create(index=index, mappings=mappings)
+            # Create the index if it doesn't exist
+            if not client.indices.exists(index=self._index):
+                client.indices.create(index=self._index, mappings=mappings)
+
+            self._client = client
+
+        return self._client
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -139,6 +161,7 @@ class ElasticsearchDocumentStore:
         return default_to_dict(
             self,
             hosts=self._hosts,
+            custom_mapping=self._custom_mapping,
             index=self._index,
             embedding_similarity_function=self._embedding_similarity_function,
             **self._kwargs,
@@ -161,7 +184,7 @@ class ElasticsearchDocumentStore:
         Returns how many documents are present in the document store.
         :returns: Number of documents in the document store.
         """
-        return self._client.count(index=self._index)["count"]
+        return self.client.count(index=self._index)["count"]
 
     def _search_documents(self, **kwargs) -> List[Document]:
         """
@@ -176,7 +199,7 @@ class ElasticsearchDocumentStore:
         from_ = 0
         # Handle pagination
         while True:
-            res = self._client.search(
+            res = self.client.search(
                 index=self._index,
                 from_=from_,
                 **kwargs,
@@ -250,7 +273,7 @@ class ElasticsearchDocumentStore:
             )
 
         documents_written, errors = helpers.bulk(
-            client=self._client,
+            client=self.client,
             actions=elasticsearch_actions,
             refresh="wait_for",
             index=self._index,
@@ -306,7 +329,7 @@ class ElasticsearchDocumentStore:
         """
 
         helpers.bulk(
-            client=self._client,
+            client=self.client,
             actions=({"_op_type": "delete", "_id": id_} for id_ in document_ids),
             refresh="wait_for",
             index=self._index,

{elasticsearch_haystack-0.4.0 → elasticsearch_haystack-0.6.0}/tests/test_bm25_retriever.py

@@ -28,6 +28,7 @@ def test_to_dict(_mock_elasticsearch_client):
             "document_store": {
                 "init_parameters": {
                     "hosts": "some fake host",
+                    "custom_mapping": None,
                     "index": "default",
                     "embedding_similarity_function": "cosine",
                 },

{elasticsearch_haystack-0.4.0 → elasticsearch_haystack-0.6.0}/tests/test_document_store.py

@@ -4,7 +4,7 @@
 
 import random
 from typing import List
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 from elasticsearch.exceptions import BadRequestError  # type: ignore[import-not-found]
@@ -15,6 +15,12 @@ from haystack.testing.document_store import DocumentStoreBaseTests
 from haystack_integrations.document_stores.elasticsearch import ElasticsearchDocumentStore
 
 
+@patch("haystack_integrations.document_stores.elasticsearch.document_store.Elasticsearch")
+def test_init_is_lazy(_mock_es_client):
+    ElasticsearchDocumentStore(hosts="testhost")
+    _mock_es_client.assert_not_called()
+
+
 @patch("haystack_integrations.document_stores.elasticsearch.document_store.Elasticsearch")
 def test_to_dict(_mock_elasticsearch_client):
     document_store = ElasticsearchDocumentStore(hosts="some hosts")
@@ -23,6 +29,7 @@ def test_to_dict(_mock_elasticsearch_client):
         "type": "haystack_integrations.document_stores.elasticsearch.document_store.ElasticsearchDocumentStore",
         "init_parameters": {
             "hosts": "some hosts",
+            "custom_mapping": None,
             "index": "default",
             "embedding_similarity_function": "cosine",
         },
@@ -35,6 +42,7 @@ def test_from_dict(_mock_elasticsearch_client):
         "type": "haystack_integrations.document_stores.elasticsearch.document_store.ElasticsearchDocumentStore",
         "init_parameters": {
             "hosts": "some hosts",
+            "custom_mapping": None,
             "index": "default",
             "embedding_similarity_function": "cosine",
         },
@@ -42,6 +50,7 @@ def test_from_dict(_mock_elasticsearch_client):
     document_store = ElasticsearchDocumentStore.from_dict(data)
     assert document_store._hosts == "some hosts"
     assert document_store._index == "default"
+    assert document_store._custom_mapping is None
     assert document_store._embedding_similarity_function == "cosine"
 
 
@@ -70,7 +79,7 @@ class TestDocumentStore(DocumentStoreBaseTests):
             hosts=hosts, index=index, embedding_similarity_function=embedding_similarity_function
         )
         yield store
-        store._client.options(ignore_status=[400, 404]).indices.delete(index=index)
+        store.client.options(ignore_status=[400, 404]).indices.delete(index=index)
 
     def assert_documents_are_equal(self, received: List[Document], expected: List[Document]):
         """
@@ -98,7 +107,7 @@ class TestDocumentStore(DocumentStoreBaseTests):
         super().assert_documents_are_equal(received, expected)
 
     def test_user_agent_header(self, document_store: ElasticsearchDocumentStore):
-        assert document_store._client._headers["user-agent"].startswith("haystack-py-ds/")
+        assert document_store.client._headers["user-agent"].startswith("haystack-py-ds/")
 
     def test_write_documents(self, document_store: ElasticsearchDocumentStore):
         docs = [Document(id="1")]
@@ -280,3 +289,33 @@ class TestDocumentStore(DocumentStoreBaseTests):
 
         with pytest.raises(DocumentStoreError):
             document_store.write_documents(docs)
+
+    @patch("haystack_integrations.document_stores.elasticsearch.document_store.Elasticsearch")
+    def test_init_with_custom_mapping(self, mock_elasticsearch):
+        custom_mapping = {
+            "properties": {
+                "embedding": {"type": "dense_vector", "index": True, "similarity": "dot_product"},
+                "content": {"type": "text"},
+            },
+            "dynamic_templates": [
+                {
+                    "strings": {
+                        "path_match": "*",
+                        "match_mapping_type": "string",
+                        "mapping": {
+                            "type": "keyword",
+                        },
+                    }
+                }
+            ],
+        }
+        mock_client = Mock(
+            indices=Mock(create=Mock(), exists=Mock(return_value=False)),
+        )
+        mock_elasticsearch.return_value = mock_client
+
+        _ = ElasticsearchDocumentStore(hosts="some hosts", custom_mapping=custom_mapping).client
+        mock_client.indices.create.assert_called_once_with(
+            index="default",
+            mappings=custom_mapping,
+        )

{elasticsearch_haystack-0.4.0 → elasticsearch_haystack-0.6.0}/tests/test_embedding_retriever.py

@@ -29,6 +29,7 @@ def test_to_dict(_mock_elasticsearch_client):
             "document_store": {
                 "init_parameters": {
                     "hosts": "some fake host",
+                    "custom_mapping": None,
                     "index": "default",
                     "embedding_similarity_function": "cosine",
                 },