Feat: Cohere reranker (MemMachine#761)

Author: vinares

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "aiosqlite>=0.20.0",
     "boto3>=1.40.40",
     "boto3-stubs==1.40.64",
+    "cohere>=5.20.0",
     "coverage>=7.11.0",
     "dotenv>=0.9.9",
     "fastapi>=0.116.1",

sample_configs/episodic_memory_config.cpu.sample

@@ -100,6 +100,11 @@ resources:
       provider: "identity"
     bm_ranker_id:
       provider: "bm25"
+    cohere_reranker_id:
+      provider: "cohere"
+      config:
+        cohere_key: <COHERE_API_KEY>
+        model: "rerank-english-v3.0"
     aws_reranker_id:
       provider: "amazon-bedrock"
       config:

sample_configs/episodic_memory_config.gpu.sample

@@ -101,6 +101,11 @@ resources:
       provider: "identity"
     bm_ranker_id:
       provider: "bm25"
+    cohere_reranker_id:
+      provider: "cohere"
+      config:
+        cohere_key: <COHERE_API_KEY>
+        model: "rerank-english-v3.0"
     ce_ranker_id:
       provider: "cross-encoder"
       config:

src/memmachine/common/configuration/reranker_conf.py

@@ -50,6 +50,19 @@ class AmazonBedrockRerankerConf(YamlSerializableMixin):
     )
 
 
+class CohereRerankerConf(YamlSerializableMixin):
+    """Parameters for CohereReranker."""
+
+    cohere_key: SecretStr | None = Field(
+        ...,
+        description="Cohere API key for authentication.",
+    )
+    model: str = Field(
+        default="rerank-english-v3.0",
+        description="Cohere rerank model",
+    )
+
+
 class CrossEncoderRerankerConf(YamlSerializableMixin):
     """Parameters for CrossEncoderReranker."""
 
@@ -89,6 +102,7 @@ class RerankersConf(BaseModel):
 
     bm25: dict[str, BM25RerankerConf] = {}
     amazon_bedrock: dict[str, AmazonBedrockRerankerConf] = {}
+    cohere: dict[str, CohereRerankerConf] = {}
     cross_encoder: dict[str, CrossEncoderRerankerConf] = {}
     embedder: dict[str, EmbedderRerankerConf] = {}
     identity: dict[str, IdentityRerankerConf] = {}
@@ -101,6 +115,7 @@ def contains_reranker(self, reranker_id: str) -> bool:
         return reranker_id in self._saved_reranker_ids
 
     BM25: ClassVar[str] = "bm25"
+    COHERE: ClassVar[str] = "cohere"
     CROSS_ENCODER: ClassVar[str] = "cross-encoder"
     EMBEDDER: ClassVar[str] = "embedder"
     IDENTITY: ClassVar[str] = "identity"
@@ -126,6 +141,9 @@ def add_reranker(name: str, provider: str, config: dict) -> None:
         for reranker_id, conf in self.amazon_bedrock.items():
             add_reranker(reranker_id, self.AMAZON_BEDROCK, conf.to_yaml_dict())
 
+        for reranker_id, conf in self.cohere.items():
+            add_reranker(reranker_id, self.COHERE, conf.to_yaml_dict())
+
         for reranker_id, conf in self.cross_encoder.items():
             add_reranker(reranker_id, self.CROSS_ENCODER, conf.to_yaml_dict())
 
@@ -154,6 +172,7 @@ def parse(cls, input_dict: dict) -> Self:
 
         bm25_dict = {}
         amazon_bedrock_dict = {}
+        cohere_dict = {}
         cross_encoder_dict = {}
         embedder_dict = {}
         identity_dict = {}
@@ -167,6 +186,8 @@ def parse(cls, input_dict: dict) -> Self:
                 bm25_dict[reranker_id] = BM25RerankerConf(**conf)
             elif provider == cls.AMAZON_BEDROCK:
                 amazon_bedrock_dict[reranker_id] = AmazonBedrockRerankerConf(**conf)
+            elif provider == cls.COHERE:
+                cohere_dict[reranker_id] = CohereRerankerConf(**conf)
             elif provider == cls.CROSS_ENCODER:
                 cross_encoder_dict[reranker_id] = CrossEncoderRerankerConf(**conf)
             elif provider == cls.EMBEDDER:
@@ -183,6 +204,7 @@ def parse(cls, input_dict: dict) -> Self:
         ret = cls(
             bm25=bm25_dict,
             amazon_bedrock=amazon_bedrock_dict,
+            cohere=cohere_dict,
             cross_encoder=cross_encoder_dict,
             embedder=embedder_dict,
             identity=identity_dict,

src/memmachine/common/reranker/cohere_reranker.py

@@ -0,0 +1,66 @@
+"""Cohere reranker implementation."""
+
+import asyncio
+import logging
+from typing import Any
+
+import cohere
+from pydantic import BaseModel, Field
+
+from memmachine.common.data_types import ExternalServiceAPIError
+
+from .reranker import Reranker
+
+logger = logging.getLogger(__name__)
+
+
+class CohereRerankerParams(BaseModel):
+    """Configuration parameters for CohereReranker."""
+
+    client: Any = Field(
+        ...,
+        description="Cohere client instance for making API calls",
+    )
+    model: str = Field(
+        "rerank-english-v3.0",
+        description="Cohere rerank model",
+    )
+
+
+class CohereReranker(Reranker):
+    """Reranker using Cohere's rerank API."""
+
+    def __init__(self, params: CohereRerankerParams) -> None:
+        """Initialize a CohereReranker with the provided parameters."""
+        super().__init__()
+
+        self._client = params.client
+        self._model = params.model
+
+    async def score(self, query: str, candidates: list[str]) -> list[float]:
+        """Score candidates using Cohere's rerank API."""
+
+        # Build request parameters
+        def _call_rerank() -> cohere.RerankResponse:
+            return self._client.rerank(
+                model=self._model,
+                query=query,
+                documents=candidates,
+            )
+
+        try:
+            response = await asyncio.to_thread(_call_rerank)
+        except Exception as e:
+            error_message = (
+                f"Failed to score candidates with Cohere model {self._model} "
+                f"due to {type(e).__name__}: {e}"
+            )
+            logger.exception(error_message)
+            raise ExternalServiceAPIError(error_message) from e
+
+        # Cohere returns ranked order — map scores back to original positions
+        scores = [0.0] * len(candidates)
+        for result in response.results:
+            scores[result.index] = result.relevance_score
+
+        return scores

src/memmachine/common/resource_manager/reranker_manager.py

@@ -47,6 +47,7 @@ async def build_all(self) -> dict[str, Reranker]:
             name
             for keys in [
                 self.conf.bm25.keys(),
+                self.conf.cohere.keys(),
                 self.conf.cross_encoder.keys(),
                 self.conf.amazon_bedrock.keys(),
                 self.conf.embedder.keys(),
@@ -80,6 +81,8 @@ async def _build_reranker(self, name: str) -> Reranker:
         """Create a reranker based on provider-specific configuration."""
         if name in self.conf.bm25:
             return await self._build_bm25_reranker(name)
+        if name in self.conf.cohere:
+            return await self._build_cohere_reranker(name)
         if name in self.conf.cross_encoder:
             return await self._build_cross_encoder_reranker(name)
         if name in self.conf.amazon_bedrock:
@@ -129,6 +132,25 @@ def _default_tokenize(text: str) -> list[str]:
         )
         return self.rerankers[name]
 
+    async def _build_cohere_reranker(self, name: str) -> Reranker:
+        from cohere import ClientV2
+
+        from memmachine.common.reranker.cohere_reranker import (
+            CohereReranker,
+            CohereRerankerParams,
+        )
+
+        conf = self.conf.cohere[name]
+
+        cohere_api_key = conf.cohere_key.get_secret_value() if conf.cohere_key else None
+        client = ClientV2(api_key=cohere_api_key)
+        params = CohereRerankerParams(
+            client=client,
+            model=conf.model,
+        )
+        self.rerankers[name] = CohereReranker(params)
+        return self.rerankers[name]
+
     async def _build_cross_encoder_reranker(self, name: str) -> Reranker:
         from sentence_transformers import CrossEncoder
 

tests/memmachine/common/resource_manager/test_reranker_manager.py

@@ -6,6 +6,7 @@
 from memmachine.common.configuration.reranker_conf import (
     AmazonBedrockRerankerConf,
     BM25RerankerConf,
+    CohereRerankerConf,
     CrossEncoderRerankerConf,
     RerankersConf,
     RRFHybridRerankerConf,
@@ -23,7 +24,13 @@ def mock_conf():
             ),
         },
         identity={"id_ranker_id": {}},
-        bm25={"bm_ranker_id": BM25RerankerConf(tokenize="simple")},
+        bm25={"bm_ranker_id": BM25RerankerConf(tokenizer="simple")},
+        cohere={
+            "cohere_reranker_id": CohereRerankerConf(
+                cohere_key=SecretStr("<COHERE_API_KEY>"),
+                model="rerank-english-v3.0",
+            ),
+        },
         cross_encoder={
             "ce_ranker_id": CrossEncoderRerankerConf(
                 model_name="cross-encoder/qnli-electra-base",
@@ -79,6 +86,15 @@ async def test_reranker_not_found(reranker_manager):
         await reranker_manager.get_reranker("unknown_reranker_id")
 
 
+@pytest.mark.asyncio
+async def test_build_cohere_rerankers(reranker_manager):
+    await reranker_manager.build_all()
+
+    assert "cohere_reranker_id" in reranker_manager.rerankers
+    reranker = reranker_manager.rerankers["cohere_reranker_id"]
+    assert reranker is not None
+
+
 @pytest.mark.asyncio
 async def test_build_cross_encoder_rerankers(reranker_manager):
     await reranker_manager.build_all()