elasticsearch-haystack 3.0.1__py3-none-any.whl → 3.1.0__py3-none-any.whl

{elasticsearch_haystack-3.0.1.dist-info → elasticsearch_haystack-3.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: elasticsearch-haystack
-Version: 3.0.1
+Version: 3.1.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
@@ -50,7 +50,7 @@ docker-compose up
 Then run tests:
 
 ```console
-hatch run test
+hatch run test:all
 ```
 
 ## License

elasticsearch_haystack-3.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+haystack_integrations/components/retrievers/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+haystack_integrations/components/retrievers/elasticsearch/__init__.py,sha256=cSJBsYjz_T4kK-M-auAHVUnYIcgUqqwwQe_hsF0_IG4,307
+haystack_integrations/components/retrievers/elasticsearch/bm25_retriever.py,sha256=HsR42EolOBEIuLNwQ_8FeSmrHMJ6WscYulcoXSHoaYQ,7098
+haystack_integrations/components/retrievers/elasticsearch/embedding_retriever.py,sha256=-6eIHW5cU4k8-jAsUsCb15hJRalpkUhzy_dNxr5HUZo,7404
+haystack_integrations/document_stores/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+haystack_integrations/document_stores/elasticsearch/__init__.py,sha256=YTfu94dtVUBogbJFr1aJrKuaI6-Bw9VuHfPoyU7M8os,207
+haystack_integrations/document_stores/elasticsearch/document_store.py,sha256=d_u49ySnhQzK_jGGThAYCWKPGDdVcpmCGQ-CWgCaO58,27852
+haystack_integrations/document_stores/elasticsearch/filters.py,sha256=Umip-PP4uFjuWeB1JWkKhaKClQ0VpiykoDlDu99wIV0,9759
+elasticsearch_haystack-3.1.0.dist-info/METADATA,sha256=O1bfELb0DpiXMSLvZuq4upfSo-1So67b058LXqt7N4E,2261
+elasticsearch_haystack-3.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+elasticsearch_haystack-3.1.0.dist-info/licenses/LICENSE,sha256=_M2kulivnaiTHiW-5CRlZrPmH47tt04pBgAgeDvfYi4,11342
+elasticsearch_haystack-3.1.0.dist-info/RECORD,,

haystack_integrations/components/retrievers/elasticsearch/bm25_retriever.py

@@ -116,7 +116,9 @@ class ElasticsearchBM25Retriever:
         return default_from_dict(cls, data)
 
     @component.output_types(documents=List[Document])
-    def run(self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None):
+    def run(
+        self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None
+    ) -> Dict[str, List[Document]]:
         """
         Retrieve documents using the BM25 keyword-based algorithm.
 
@@ -139,7 +141,9 @@ class ElasticsearchBM25Retriever:
         return {"documents": docs}
 
     @component.output_types(documents=List[Document])
-    async def run_async(self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None):
+    async def run_async(
+        self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None
+    ) -> Dict[str, List[Document]]:
         """
         Asynchronously retrieve documents using the BM25 keyword-based algorithm.
 

haystack_integrations/components/retrievers/elasticsearch/embedding_retriever.py

@@ -114,7 +114,9 @@ class ElasticsearchEmbeddingRetriever:
         return default_from_dict(cls, data)
 
     @component.output_types(documents=List[Document])
-    def run(self, query_embedding: List[float], filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None):
+    def run(
+        self, query_embedding: List[float], filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None
+    ) -> Dict[str, List[Document]]:
         """
         Retrieve documents using a vector similarity metric.
 
@@ -139,7 +141,7 @@ class ElasticsearchEmbeddingRetriever:
     @component.output_types(documents=List[Document])
     async def run_async(
         self, query_embedding: List[float], filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None
-    ):
+    ) -> Dict[str, List[Document]]:
         """
         Asynchronously retrieve documents using a vector similarity metric.
 

haystack_integrations/document_stores/elasticsearch/document_store.py

@@ -5,16 +5,14 @@ from collections.abc import Mapping
 from typing import Any, Dict, List, Literal, Optional, Union
 
 import numpy as np
-
-# There are no import stubs for elastic_transport and elasticsearch so mypy fails
-from elastic_transport import NodeConfig  # type: ignore[import-not-found]
+from elastic_transport import NodeConfig
 from haystack import default_from_dict, default_to_dict, logging
 from haystack.dataclasses import Document
 from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
 from haystack.document_stores.types import DuplicatePolicy
 from haystack.version import __version__ as haystack_version
 
-from elasticsearch import AsyncElasticsearch, Elasticsearch, helpers  # type: ignore[import-not-found]
+from elasticsearch import AsyncElasticsearch, Elasticsearch, helpers
 
 from .filters import _normalize_filters
 
@@ -66,7 +64,7 @@ class ElasticsearchDocumentStore:
         custom_mapping: Optional[Dict[str, Any]] = None,
         index: str = "default",
         embedding_similarity_function: Literal["cosine", "dot_product", "l2_norm", "max_inner_product"] = "cosine",
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Creates a new ElasticsearchDocumentStore instance.
@@ -93,8 +91,8 @@ class ElasticsearchDocumentStore:
         :param **kwargs: Optional arguments that `Elasticsearch` takes.
         """
         self._hosts = hosts
-        self._client = None
-        self._async_client = None
+        self._client: Optional[Elasticsearch] = None
+        self._async_client: Optional[AsyncElasticsearch] = None
         self._index = index
         self._embedding_similarity_function = embedding_similarity_function
         self._custom_mapping = custom_mapping
@@ -166,6 +164,7 @@ class ElasticsearchDocumentStore:
         Returns the synchronous Elasticsearch client, initializing it if necessary.
         """
         self._ensure_initialized()
+        assert self._client is not None  # noqa: S101
         return self._client
 
     @property
@@ -174,6 +173,7 @@ class ElasticsearchDocumentStore:
         Returns the asynchronous Elasticsearch client, initializing it if necessary.
         """
         self._ensure_initialized()
+        assert self._async_client is not None  # noqa: S101
         return self._async_client
 
     def to_dict(self) -> Dict[str, Any]:
@@ -226,7 +226,7 @@ class ElasticsearchDocumentStore:
         result = await self._async_client.count(index=self._index)  # type: ignore
         return result["count"]
 
-    def _search_documents(self, **kwargs) -> List[Document]:
+    def _search_documents(self, **kwargs: Any) -> List[Document]:
         """
         Calls the Elasticsearch client's search method and handles pagination.
         """
@@ -253,7 +253,7 @@ class ElasticsearchDocumentStore:
                 break
         return documents
 
-    async def _search_documents_async(self, **kwargs) -> List[Document]:
+    async def _search_documents_async(self, **kwargs: Any) -> List[Document]:
         """
         Asynchronously calls the Elasticsearch client's search method and handles pagination.
         """
@@ -379,9 +379,12 @@ class ElasticsearchDocumentStore:
             refresh="wait_for",
             index=self._index,
             raise_on_error=False,
+            stats_only=False,
         )
 
         if errors:
+            # with stats_only=False, errors is guaranteed to be a list of dicts
+            assert isinstance(errors, list)  # noqa: S101
             duplicate_errors_ids = []
             other_errors = []
             for e in errors:
@@ -451,13 +454,16 @@ class ElasticsearchDocumentStore:
 
         try:
             success, failed = await helpers.async_bulk(
-                client=self._async_client,
+                client=self.async_client,
                 actions=actions,
                 index=self._index,
                 refresh=True,
                 raise_on_error=False,
+                stats_only=False,
             )
             if failed:
+                # with stats_only=False, failed is guaranteed to be a list of dicts
+                assert isinstance(failed, list)  # noqa: S101
                 if policy == DuplicatePolicy.FAIL:
                     for error in failed:
                         if "create" in error and error["create"]["status"] == DOC_ALREADY_EXISTS:
@@ -494,7 +500,7 @@ class ElasticsearchDocumentStore:
 
         try:
             await helpers.async_bulk(
-                client=self._async_client,
+                client=self.async_client,
                 actions=({"_op_type": "delete", "_id": id_} for id_ in document_ids),
                 index=self._index,
                 refresh=True,
@@ -551,6 +557,8 @@ class ElasticsearchDocumentStore:
 
         if scale_score:
             for doc in documents:
+                if doc.score is None:
+                    continue
                 doc.score = float(1 / (1 + np.exp(-np.asarray(doc.score / BM25_SCALING_FACTOR))))
 
         return documents

elasticsearch_haystack-3.0.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-haystack_integrations/components/retrievers/elasticsearch/__init__.py,sha256=cSJBsYjz_T4kK-M-auAHVUnYIcgUqqwwQe_hsF0_IG4,307
-haystack_integrations/components/retrievers/elasticsearch/bm25_retriever.py,sha256=ISHc6elYXoDXDvC62_3bMMCk_Dv67jvZIgQBCZ1ZHdw,7012
-haystack_integrations/components/retrievers/elasticsearch/embedding_retriever.py,sha256=jHDLMeecpf-DhvbRM1AAq2kIJn7xMNTR9vkm-FhHH7k,7332
-haystack_integrations/document_stores/elasticsearch/__init__.py,sha256=YTfu94dtVUBogbJFr1aJrKuaI6-Bw9VuHfPoyU7M8os,207
-haystack_integrations/document_stores/elasticsearch/document_store.py,sha256=pZ0pPyOCPTCKNYD4q5YbLrslSGTIbVPj60U18-BImX8,27406
-haystack_integrations/document_stores/elasticsearch/filters.py,sha256=Umip-PP4uFjuWeB1JWkKhaKClQ0VpiykoDlDu99wIV0,9759
-elasticsearch_haystack-3.0.1.dist-info/METADATA,sha256=defq5KselqQwwMLUa3jfrmXjjVxH15cLSn7iKlDKgJk,2257
-elasticsearch_haystack-3.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-elasticsearch_haystack-3.0.1.dist-info/licenses/LICENSE,sha256=_M2kulivnaiTHiW-5CRlZrPmH47tt04pBgAgeDvfYi4,11342
-elasticsearch_haystack-3.0.1.dist-info/RECORD,,