feat: support base_url for Cohere rerankers

Add optional `base_url` to `CohereRerankerConf` and forward it when
building `cohere.ClientV2`, allowing MemMachine to use self-hosted
Cohere-compatible rerank APIs without adding a new provider.

Also add unit tests for config parsing/validation and client kwargs
forwarding, and update the three episodic sample configs with
commented `base_url` examples.

Reported-by: Kwangjin Ko <kwangjin.ko@sk.com>
Signed-off-by: Hyeongtak Ji <hyeongtak.ji@sk.com>

Author: hyeongtakji

packages/server/server_tests/memmachine_server/common/configuration/test_reranker_conf.py

@@ -7,6 +7,7 @@
 from memmachine_server.common.configuration.reranker_conf import (
     AmazonBedrockRerankerConf,
     BM25RerankerConf,
+    CohereRerankerConf,
     CrossEncoderRerankerConf,
     EmbedderRerankerConf,
     IdentityRerankerConf,
@@ -47,6 +48,18 @@ def amazon_bedrock_reranker_conf() -> dict:
     }
 
 
+@pytest.fixture
+def cohere_reranker_conf() -> dict:
+    return {
+        "provider": "cohere",
+        "config": {
+            "cohere_key": "test-cohere-key",
+            "model": "rerank-english-v3.0",
+            "base_url": "http://localhost:8000",
+        },
+    }
+
+
 @pytest.fixture
 def cross_encoder_reranker_conf() -> dict:
     return {
@@ -83,6 +96,7 @@ def full_reranker_input(
     bm25_reranker_conf,
     identity_reranker_conf,
     amazon_bedrock_reranker_conf,
+    cohere_reranker_conf,
     cross_encoder_reranker_conf,
     embedder_reranker_conf,
     rrf_hybrid_reranker_conf,
@@ -93,6 +107,7 @@ def full_reranker_input(
             "id_ranker_id": identity_reranker_conf,
             "bm_ranker_id": bm25_reranker_conf,
             "aws_reranker_id": amazon_bedrock_reranker_conf,
+            "cohere_reranker_id": cohere_reranker_conf,
             "cross_encoder_id": cross_encoder_reranker_conf,
             "embedder_id": embedder_reranker_conf,
         },
@@ -121,6 +136,13 @@ def test_valid_amazon_bedrock_reranker_conf(amazon_bedrock_reranker_conf):
     assert conf.model_id == "amazon.rerank-v1:0"
 
 
+def test_valid_cohere_reranker_conf(cohere_reranker_conf):
+    conf = CohereRerankerConf(**cohere_reranker_conf["config"])
+    assert conf.cohere_key == SecretStr("test-cohere-key")
+    assert conf.model == "rerank-english-v3.0"
+    assert conf.base_url == "http://localhost:8000"
+
+
 def test_valid_cross_encoder_reranker_conf(cross_encoder_reranker_conf):
     conf = CrossEncoderRerankerConf(**cross_encoder_reranker_conf["config"])
     assert conf.model_name == "cross-encoder/qnli-electra-base"
@@ -150,6 +172,8 @@ def test_full_reranker_conf(full_reranker_input):
 
     assert "aws_reranker_id" in conf.amazon_bedrock
     assert conf.amazon_bedrock["aws_reranker_id"].region == "us-west-2"
+    assert "cohere_reranker_id" in conf.cohere
+    assert conf.cohere["cohere_reranker_id"].base_url == "http://localhost:8000"
 
     assert "cross_encoder_id" in conf.cross_encoder
     assert (
@@ -182,3 +206,12 @@ def test_missing_required_field_in_bedrock_reranker():
         AmazonBedrockRerankerConf(**config)
     assert "missing" in str(exc_info.value)
     assert "aws_secret_access_key" in str(exc_info.value)
+
+
+def test_invalid_cohere_base_url():
+    with pytest.raises(ValidationError, match="Invalid base URL"):
+        CohereRerankerConf(
+            cohere_key=SecretStr("test-cohere-key"),
+            model="rerank-english-v3.0",
+            base_url="localhost:8000",
+        )

packages/server/server_tests/memmachine_server/common/resource_manager/test_reranker_manager.py

@@ -34,8 +34,9 @@ def mock_conf():
         bm25={"bm_ranker_id": BM25RerankerConf(tokenizer="simple")},
         cohere={
             "cohere_reranker_id": CohereRerankerConf(
-                cohere_key=SecretStr("<COHERE_API_KEY>"),
+                cohere_key=SecretStr("test-cohere-key"),
                 model="rerank-english-v3.0",
+                base_url="http://localhost:8000",
             ),
         },
         cross_encoder={
@@ -132,6 +133,72 @@ async def test_build_cohere_rerankers(reranker_manager):
     assert reranker is not None
 
 
+@pytest.mark.asyncio
+async def test_build_cohere_reranker_passes_base_url(monkeypatch):
+    captured_kwargs = {}
+
+    class FakeCohereClient:
+        pass
+
+    def fake_client_v2(**kwargs):
+        captured_kwargs.update(kwargs)
+        return FakeCohereClient()
+
+    monkeypatch.setattr("cohere.ClientV2", fake_client_v2)
+
+    conf = RerankersConf(
+        cohere={
+            "cohere_reranker_id": CohereRerankerConf(
+                cohere_key=SecretStr("test-cohere-key"),
+                model="rerank-english-v3.0",
+                base_url="http://localhost:8000",
+            ),
+        },
+    )
+    reranker_manager = RerankerManager(
+        conf=conf,
+        embedder_factory=cast(EmbedderFactory, FakeEmbedderFactory()),
+    )
+
+    await reranker_manager.get_reranker("cohere_reranker_id")
+
+    assert captured_kwargs == {
+        "api_key": "test-cohere-key",
+        "base_url": "http://localhost:8000",
+    }
+
+
+@pytest.mark.asyncio
+async def test_build_cohere_reranker_omits_base_url_when_not_configured(monkeypatch):
+    captured_kwargs = {}
+
+    class FakeCohereClient:
+        pass
+
+    def fake_client_v2(**kwargs):
+        captured_kwargs.update(kwargs)
+        return FakeCohereClient()
+
+    monkeypatch.setattr("cohere.ClientV2", fake_client_v2)
+
+    conf = RerankersConf(
+        cohere={
+            "cohere_reranker_id": CohereRerankerConf(
+                cohere_key=SecretStr("test-cohere-key"),
+                model="rerank-english-v3.0",
+            ),
+        },
+    )
+    reranker_manager = RerankerManager(
+        conf=conf,
+        embedder_factory=cast(EmbedderFactory, FakeEmbedderFactory()),
+    )
+
+    await reranker_manager.get_reranker("cohere_reranker_id")
+
+    assert captured_kwargs == {"api_key": "test-cohere-key"}
+
+
 @requires_sentence_transformers
 @pytest.mark.asyncio
 async def test_build_cross_encoder_rerankers(reranker_manager):

packages/server/src/memmachine_server/common/configuration/reranker_conf.py

@@ -1,9 +1,10 @@
 """Reranker configuration models."""
 
 from typing import ClassVar, Self
+from urllib.parse import urlparse
 
 import yaml
-from pydantic import BaseModel, Field, PrivateAttr, SecretStr
+from pydantic import BaseModel, Field, PrivateAttr, SecretStr, field_validator
 
 from memmachine_server.common.configuration.mixin_confs import (
     MetricsFactoryIdMixin,
@@ -64,6 +65,20 @@ class CohereRerankerConf(YamlSerializableMixin):
         default="rerank-english-v3.0",
         description="Cohere rerank model",
     )
+    base_url: str | None = Field(
+        default=None,
+        description="Cohere-compatible Rerank API base URL",
+    )
+
+    @field_validator("base_url")
+    @classmethod
+    def validate_base_url(cls, v: str | None) -> str | None:
+        """Ensure the base URL includes a scheme and host."""
+        if v is not None:
+            parsed_url = urlparse(v)
+            if not parsed_url.scheme or not parsed_url.netloc:
+                raise ValueError(f"Invalid base URL: base_url={v}")
+        return v
 
 
 class CrossEncoderRerankerConf(YamlSerializableMixin):

packages/server/src/memmachine_server/common/resource_manager/reranker_manager.py

@@ -178,7 +178,10 @@ async def _build_cohere_reranker(self, name: str) -> Reranker:
         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)
+        if conf.base_url is not None:
+            client = ClientV2(api_key=cohere_api_key, base_url=conf.base_url)
+        else:
+            client = ClientV2(api_key=cohere_api_key)
         params = CohereRerankerParams(
             client=client,
             model=conf.model,

sample_configs/episodic_memory_config.cpu.sample

@@ -138,6 +138,9 @@ resources:
       config:
         cohere_key: <COHERE_API_KEY>
         model: "rerank-english-v3.0"
+      # Optional for self-hosted Cohere-compatible rerank APIs.
+      # Examples:
+      # base_url: "http://localhost:8000"
     aws_reranker_id:
       provider: "amazon-bedrock"
       config:

sample_configs/episodic_memory_config.gpu.sample

@@ -139,6 +139,9 @@ resources:
       config:
         cohere_key: <COHERE_API_KEY>
         model: "rerank-english-v3.0"
+      # Optional for self-hosted Cohere-compatible rerank APIs.
+      # Examples:
+      # base_url: "http://localhost:8000"
     ce_ranker_id:
       provider: "cross-encoder"
       config:

sample_configs/episodic_memory_config.nebula.sample

@@ -163,6 +163,9 @@ resources:
       config:
         cohere_key: <COHERE_API_KEY>
         model: "rerank-english-v3.0"
+      # Optional for self-hosted Cohere-compatible rerank APIs.
+      # Examples:
+      # base_url: "http://localhost:8000"
     aws_reranker_id:
       provider: "amazon-bedrock"
       config: