Public/kotaemon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: added base_url parameter to CochereReranking (#743)

Browse Source 
Co-authored-by: Mauro Gattari <mauro.gattari@infn.it>
This commit is contained in:
mginfn 2025-06-05 11:35:00 +02:00 committed by GitHub
parent 833982ac81
commit ffe766f24d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
libs/kotaemon/kotaemon/rerankings/cohere.py
View File

@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import os
from decouple import config from decouple import config
from kotaemon.base import Document, Param from kotaemon.base import Document, Param
@ -23,6 +24,9 @@ class CohereReranking(BaseReranking):
help="Cohere API key", help="Cohere API key",
required=True, 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]: def run(self, documents: list[Document], query: str) -> list[Document]:
"""Use Cohere Reranker model to re-order documents """Use Cohere Reranker model to re-order documents
@ -38,7 +42,7 @@ class CohereReranking(BaseReranking):
print("Cohere API key not found. Skipping rerankings.") print("Cohere API key not found. Skipping rerankings.")
return documents 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] = [] compressed_docs: list[Document] = []
if not documents: # to avoid empty api call if not documents: # to avoid empty api call