Skip to content

Commit

Permalink
Fix linting issues
Browse files Browse the repository at this point in the history
  • Loading branch information
silvanocerza committed Oct 6, 2023
1 parent 72e7309 commit 01daad8
Show file tree
Hide file tree
Showing 4 changed files with 16 additions and 13 deletions.
5 changes: 3 additions & 2 deletions src/elasticsearch_haystack/bm25_retriever.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
# SPDX-FileCopyrightText: 2023-present Silvano Cerza <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from typing import Optional, Dict, List, Any
from typing import Any, Dict, List, Optional

from haystack.preview import component, default_to_dict, default_from_dict
from haystack.preview import component, default_from_dict, default_to_dict
from haystack.preview.dataclasses import Document

from elasticsearch_haystack.document_store import ElasticsearchDocumentStore
Expand All @@ -13,6 +13,7 @@
class ElasticsearchBM25Retriever:
def __init__(
self,
*,
document_store: ElasticsearchDocumentStore,
filters: Optional[Dict[str, Any]] = None,
top_k: int = 10,
Expand Down
12 changes: 7 additions & 5 deletions src/elasticsearch_haystack/document_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import numpy as np
from elastic_transport import NodeConfig
from elasticsearch import Elasticsearch, helpers
from haystack.preview import default_to_dict, default_from_dict
from haystack.preview import default_from_dict, default_to_dict
from haystack.preview.dataclasses import Document
from haystack.preview.document_stores.decorator import document_store
from haystack.preview.document_stores.errors import DuplicateDocumentError
Expand Down Expand Up @@ -262,6 +262,7 @@ def delete_documents(self, document_ids: List[str]) -> None:
def _bm25_retrieval(
self,
query: str,
*,
filters: Optional[Dict[str, Any]] = None,
top_k: int = 10,
scale_score: bool = True,
Expand All @@ -278,17 +279,19 @@ def _bm25_retrieval(
`query` must be a non empty string, otherwise a `ValueError` will be raised.
:param query: String to search in saved Documents' text.
:param filters: Filters applied to the retrieved Documents, for more info see `ElasticsearchDocumentStore.filter_documents`, defaults to None
:param filters: Filters applied to the retrieved Documents, for more info
see `ElasticsearchDocumentStore.filter_documents`, defaults to None
:param top_k: Maximum number of Documents to return, defaults to 10
:param scale_score: If `True` scales the Document`s scores between 0 and 1, defaults to True
:raises ValueError: If `query` is an empty string
:return: List of Document that match `query`
"""

if not query:
raise ValueError("query must be a non empty string")
msg = "query must be a non empty string"
raise ValueError(msg)

body = {
body: Dict[str, Any] = {
"size": top_k,
"query": {
"bool": {
Expand All @@ -304,7 +307,6 @@ def _bm25_retrieval(
}
},
}
{"query": {"bool": {"must": [{"multi_match": {"query": "PHP", "type": "most_fields", "operator": "AND"}}]}}}

if filters:
body["query"]["bool"]["filter"] = _normalize_filters(filters)
Expand Down
8 changes: 4 additions & 4 deletions tests/test_bm25_retriever.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
# SPDX-FileCopyrightText: 2023-present Silvano Cerza <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import patch, Mock
from unittest.mock import Mock, patch

from haystack.preview.dataclasses import Document

from elasticsearch_haystack.document_store import ElasticsearchDocumentStore
from elasticsearch_haystack.bm25_retriever import ElasticsearchBM25Retriever
from elasticsearch_haystack.document_store import ElasticsearchDocumentStore


def test_init_default():
Expand All @@ -19,7 +19,7 @@ def test_init_default():


@patch("elasticsearch_haystack.document_store.Elasticsearch")
def test_to_dict(mock_elasticsearch_client):
def test_to_dict(_mock_elasticsearch_client):
document_store = ElasticsearchDocumentStore(hosts="some fake host")
retriever = ElasticsearchBM25Retriever(document_store=document_store)
res = retriever.to_dict()
Expand All @@ -38,7 +38,7 @@ def test_to_dict(mock_elasticsearch_client):


@patch("elasticsearch_haystack.document_store.Elasticsearch")
def test_from_dict(mock_elasticsearch_client):
def test_from_dict(_mock_elasticsearch_client):
data = {
"type": "ElasticsearchBM25Retriever",
"init_parameters": {
Expand Down
4 changes: 2 additions & 2 deletions tests/test_document_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ def docstore(self, request):
store._client.options(ignore_status=[400, 404]).indices.delete(index=index)

@patch("elasticsearch_haystack.document_store.Elasticsearch")
def test_to_dict(self, mock_elasticsearch_client):
def test_to_dict(self, _mock_elasticsearch_client):
document_store = ElasticsearchDocumentStore(hosts="some hosts")
res = document_store.to_dict()
assert res == {
Expand All @@ -45,7 +45,7 @@ def test_to_dict(self, mock_elasticsearch_client):
}

@patch("elasticsearch_haystack.document_store.Elasticsearch")
def test_from_dict(self, mock_elasticsearch_client):
def test_from_dict(self, _mock_elasticsearch_client):
data = {
"type": "ElasticsearchDocumentStore",
"init_parameters": {
Expand Down

0 comments on commit 01daad8

Please sign in to comment.