Draft [In Progress]: draft to allow for multiple retrieval backends

chatnoir_pyterrier/retrieve.py

@@ -4,11 +4,12 @@
 from typing import Set, Optional, Iterable, Union, Any, Dict
 
 from chatnoir_api import Index, Result, Slop, ExplainedResult
+from chatnoir_api.model import SearchMethod
 from chatnoir_api.v1 import (
     search, search_phrases
 )
 from chatnoir_api.defaults import (
-    DEFAULT_INDEX, DEFAULT_SLOP, DEFAULT_RETRIES, DEFAULT_BACKOFF_SECONDS, DEFAULT_API_KEY
+    DEFAULT_INDEX, DEFAULT_SLOP, DEFAULT_RETRIES, DEFAULT_BACKOFF_SECONDS, DEFAULT_API_KEY, DEFAULT_SEARCH_METHOD
 )
 from pandas import DataFrame
 from pandas.core.groupby import DataFrameGroupBy
@@ -35,13 +36,14 @@ class ChatNoirRetrieve(Transformer):
     backoff_seconds: float = DEFAULT_BACKOFF_SECONDS
     verbose: bool = False
     api_key: str = DEFAULT_API_KEY
+    search_method: SearchMethod = DEFAULT_SEARCH_METHOD
 
     def _merge_result(
-            self,
-            row: Dict[str, Any],
-            result: Union[
-                Result, ExplainedResult,
-            ]
+        self,
+        row: Dict[str, Any],
+        result: Union[
+            Result, ExplainedResult,
+        ]
     ) -> Dict[str, Any]:
         row = {
             **row,
@@ -135,6 +137,7 @@ def _transform_query(self, topic: DataFrame) -> DataFrame:
                     retries=self.retries,
                     backoff_seconds=self.backoff_seconds,
                     api_key=self.api_key,
+                    search_method=self.search_method
                 ).results
             else:
                 results = search(
@@ -147,6 +150,7 @@ def _transform_query(self, topic: DataFrame) -> DataFrame:
                     retries=self.retries,
                     backoff_seconds=self.backoff_seconds,
                     api_key=self.api_key,
+                    search_method=self.search_method
                 ).results
         else:
             if explain:
@@ -161,6 +165,7 @@ def _transform_query(self, topic: DataFrame) -> DataFrame:
                     retries=self.retries,
                     backoff_seconds=self.backoff_seconds,
                     api_key=self.api_key,
+                    search_method=self.search_method
                 ).results
             else:
                 results = search_phrases(
@@ -174,6 +179,7 @@ def _transform_query(self, topic: DataFrame) -> DataFrame:
                     retries=self.retries,
                     backoff_seconds=self.backoff_seconds,
                     api_key=self.api_key,
+                    search_method=self.search_method
                 ).results
 
         if self.filter_unknown:

pyproject.toml

@@ -20,7 +20,7 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 dependencies = [
-    "chatnoir-api~=3.1.3",
+    "chatnoir-api~=3.2.0",
     "importlib-metadata~=8.5",
     "pandas~=2.0",
     "python-terrier>=0.11,<0.13",

tests/test_retrieve.py

@@ -5,7 +5,9 @@
 
 
 def test_retrieve_hash(api_key: str):
-    retrieve = ChatNoirRetrieve(api_key)
+    retrieve = ChatNoirRetrieve(
+        api_key=api_key,
+    )
     retrieve_hash = hash(retrieve)
     assert isinstance(retrieve_hash, int)
 
@@ -26,9 +28,9 @@ def test_retrieve_query(api_key: str, query: str, index: Index):
 
 
 def test_retrieve_feature(
-        api_key: str,
-        query: str,
-        feature: Feature,
+    api_key: str,
+    query: str,
+    feature: Feature,
 ):
     retrieve = ChatNoirRetrieve(
         api_key=api_key,