Add Cohere Reranker

pyproject.toml

@@ -74,7 +74,8 @@ module = [
     'pinecone_text.*',
     'pinecone_datasets',
     'pinecone',
-    'transformers.*'
+    'transformers.*',
+    'cohere.*',
 ]
 ignore_missing_imports = true
 

src/canopy/knowledge_base/reranker/__init__.py

@@ -1 +1,3 @@
-from .reranker import TransparentReranker, Reranker
+from .reranker import Reranker
+from .transparent import TransparentReranker
+from .cohere import CohereReranker

src/canopy/knowledge_base/reranker/cohere.py

@@ -0,0 +1,69 @@
+from typing import List, Optional
+
+from canopy.knowledge_base.models import KBQueryResult
+from canopy.knowledge_base.reranker import Reranker
+
+try:
+    import cohere
+except (OSError, ImportError, ModuleNotFoundError):
+    _cohere_installed = False
+else:
+    _cohere_installed = True
+
+
+class CohereReranker(Reranker):
+    """
+    Reranker that uses Cohere's text embedding to rerank documents.
+
+    For each query and documents returned for that query, returns a list
+    of documents ordered by their relevance to the provided query.
+    """
+
+    def __init__(self,
+                 model_name: str = 'rerank-english-v2.0',
+                 *,
+                 top_n: int = 10,
+                 api_key: Optional[str] = None):
+        """
+            Initializes the Cohere reranker.
+
+            Args:
+                model_name: The identifier of the model to use, one of :
+                    rerank-english-v2.0, rerank-multilingual-v2.0
+                top_n: The number of most relevant documents return, defaults to 10
+                api_key: API key for Cohere. If not passed `CO_API_KEY` environment
+                    variable will be used.
+        """
+
+        if not _cohere_installed:
+            raise ImportError(
+                "Failed to import cohere. Make sure you install cohere extra "
+                "dependencies by running: "
+                "pip install canopy-sdk[cohere]"
+            )
+        self._client = cohere.Client(api_key=api_key)
+        self._model_name = model_name
+        self._top_n = top_n
+
+    def rerank(self, results: List[KBQueryResult]) -> List[KBQueryResult]:
+        reranked_query_results: List[KBQueryResult] = []
+        for result in results:
+            texts = [doc.text for doc in result.documents]
+            response = self._client.rerank(query=result.query,
+                                           documents=texts,
+                                           top_n=self._top_n,
+                                           model=self._model_name)
+            reranked_docs = []
+            for rerank_result in response:
+                doc = result.documents[rerank_result.index].copy(
+                    deep=True,
+                    update=dict(score=rerank_result.relevance_score)
+                )
+                reranked_docs.append(doc)
+
+            reranked_query_results.append(KBQueryResult(query=result.query,
+                                                        documents=reranked_docs))
+        return reranked_query_results
+
+    async def arerank(self, results: List[KBQueryResult]) -> List[KBQueryResult]:
+        raise NotImplementedError()

src/canopy/knowledge_base/reranker/reranker.py

@@ -19,24 +19,3 @@ def rerank(self, results: List[KBQueryResult]) -> List[KBQueryResult]:
     @abstractmethod
     async def arerank(self, results: List[KBQueryResult]) -> List[KBQueryResult]:
         pass
-
-
-class TransparentReranker(Reranker):
-    """
-    Transparent reranker that does nothing, it just returns the results as is. This is the default reranker.
-    The TransparentReranker is used as a placeholder for future development "forcing" every result set to be reranked.
-    """  # noqa: E501
-    def rerank(self, results: List[KBQueryResult]) -> List[KBQueryResult]:
-        """
-        Returns the results as is.
-
-        Args:
-            results: A list of KBQueryResult to rerank.
-
-        Returns:
-            results: A list of KBQueryResult, same as the input.
-        """  # noqa: E501
-        return results
-
-    async def arerank(self, results: List[KBQueryResult]) -> List[KBQueryResult]:
-        return results

src/canopy/knowledge_base/reranker/transparent.py

@@ -0,0 +1,26 @@
+from typing import List
+
+from canopy.knowledge_base.models import KBQueryResult
+from canopy.knowledge_base.reranker import Reranker
+
+
+class TransparentReranker(Reranker):
+    """
+    Transparent reranker that does nothing, it just returns the results as is. This is the default reranker.
+    The TransparentReranker is used as a placeholder for future development "forcing" every result set to be reranked.
+    """  # noqa: E501
+
+    def rerank(self, results: List[KBQueryResult]) -> List[KBQueryResult]:
+        """
+        Returns the results as is.
+
+        Args:
+            results: A list of KBQueryResult to rerank.
+
+        Returns:
+            results: A list of KBQueryResult, same as the input.
+        """  # noqa: E501
+        return results
+
+    async def arerank(self, results: List[KBQueryResult]) -> List[KBQueryResult]:
+        return results

tests/system/reranker/test_cohere_reranker.py

@@ -0,0 +1,62 @@
+import os
+
+import pytest
+
+from canopy.knowledge_base.models import KBQueryResult, KBDocChunkWithScore
+from canopy.knowledge_base.reranker import CohereReranker
+
+
+@pytest.fixture
+def should_run_test():
+    if os.getenv("CO_API_KEY") is None:
+        pytest.skip(
+            "Couldn't find Cohere API key. Skipping Cohere tests."
+        )
+
+
+@pytest.fixture
+def cohere_reranker(should_run_test):
+    return CohereReranker()
+
+
+@pytest.fixture
+def documents():
+    return [
+        KBDocChunkWithScore(
+            id=f"doc_1_{i}",
+            text=f"Sample chunk {i}",
+            document_id="doc_1",
+            source="doc_1",
+            score=0.1 * i
+        ) for i in range(4)
+    ]
+
+
+@pytest.fixture
+def query_result(documents):
+    return KBQueryResult(query="Sample query 1",
+                         documents=documents)
+
+
+def test_rerank_empty(cohere_reranker):
+    results = cohere_reranker.rerank([])
+    assert results == []
+
+
+def test_rerank(cohere_reranker, query_result, documents):
+    ranked_result = next(iter(cohere_reranker.rerank([query_result])))
+    scores = [doc.score for doc in ranked_result.documents]
+
+    assert len(ranked_result.documents) == len(documents)
+    assert scores == sorted(scores, reverse=True)
+
+
+def test_bad_api_key(should_run_test, query_result):
+    from cohere import CohereAPIError
+    with pytest.raises(CohereAPIError, match="invalid api token"):
+        CohereReranker(api_key="bad key").rerank([query_result])
+
+
+def test_top_n(should_run_test, query_result):
+    results = CohereReranker(top_n=1).rerank([query_result])
+    assert len(results[0].documents) == 1

tests/system/reranker/test_transparent_reranker.py

@@ -0,0 +1,27 @@
+import pytest
+
+from canopy.knowledge_base.models import KBDocChunkWithScore, KBQueryResult
+from canopy.knowledge_base.reranker import TransparentReranker
+
+
+@pytest.fixture
+def documents():
+    return [
+        KBDocChunkWithScore(
+            id=f"doc_1_{i}",
+            text=f"Sample chunk {i}",
+            document_id="doc_1",
+            source="doc_1",
+            score=0.1 * i
+        ) for i in range(1)
+    ]
+
+
+@pytest.fixture
+def query_result(documents):
+    return KBQueryResult(query="Sample query 1",
+                         documents=documents)
+
+
+def test_rerank(query_result):
+    assert TransparentReranker().rerank([query_result]) == [query_result]