phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add test for rerank on custom api base

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-27 18:25:51 -07:00
parent be1bc99b60
commit 647504b462
3 changed files with 60 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

5
litellm/rerank_api/main.py
View file

@ -27,7 +27,7 @@ async def arerank(
custom_llm_provider: Optional[Literal["cohere", "together_ai"]] = None,
top_n: Optional[int] = None,
rank_fields: Optional[List[str]] = None,
return_documents: Optional[bool] = True,
return_documents: Optional[bool] = None,
max_chunks_per_doc: Optional[int] = None,
**kwargs,
) -> Union[RerankResponse, Coroutine[Any, Any, RerankResponse]]:
@ -112,7 +112,7 @@ def rerank(
optional_params.api_base
or litellm.api_base
or get_secret("COHERE_API_BASE")
or "https://api.cohere.ai/v1/generate"
or "https://api.cohere.com/v1/rerank"
)
headers: Dict = litellm.headers or {}
@ -126,6 +126,7 @@ def rerank(
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
api_key=cohere_key,
api_base=api_base,
_is_async=_is_async,
)
pass