From ffe766f24d4ef8a91f8c61871d2b5a1930aa204e Mon Sep 17 00:00:00 2001 From: mginfn <116359611+mginfn@users.noreply.github.com> Date: Thu, 5 Jun 2025 11:35:00 +0200 Subject: [PATCH] chore: added base_url parameter to CochereReranking (#743) Co-authored-by: Mauro Gattari --- libs/kotaemon/kotaemon/rerankings/cohere.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/libs/kotaemon/kotaemon/rerankings/cohere.py b/libs/kotaemon/kotaemon/rerankings/cohere.py index 6e68fd49..18785f04 100644 --- a/libs/kotaemon/kotaemon/rerankings/cohere.py +++ b/libs/kotaemon/kotaemon/rerankings/cohere.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from decouple import config from kotaemon.base import Document, Param @@ -23,6 +24,9 @@ class CohereReranking(BaseReranking): help="Cohere API key", required=True, ) + base_url: str = Param( + None, help="Rerank API base url. Default is https://api.cohere.com", required=False + ) def run(self, documents: list[Document], query: str) -> list[Document]: """Use Cohere Reranker model to re-order documents @@ -38,7 +42,7 @@ class CohereReranking(BaseReranking): print("Cohere API key not found. Skipping rerankings.") return documents - cohere_client = cohere.Client(self.cohere_api_key) + cohere_client = cohere.Client(self.cohere_api_key, base_url=self.base_url or os.getenv("CO_API_URL")) compressed_docs: list[Document] = [] if not documents: # to avoid empty api call