From 4a52a71956a8d46fcb7294ac71734504bb09bcc2 Mon Sep 17 00:00:00 2001 From: S. Solomon Darnell Date: Fri, 28 Mar 2025 21:52:21 -0500 Subject: two version of R2R are here --- .../litellm/rerank_api/rerank_utils.py | 46 ++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 .venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py (limited to '.venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py') diff --git a/.venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py b/.venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py new file mode 100644 index 00000000..f70ec015 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py @@ -0,0 +1,46 @@ +from typing import Any, Dict, List, Optional, Union + +from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.types.rerank import OptionalRerankParams + + +def get_optional_rerank_params( + rerank_provider_config: BaseRerankConfig, + model: str, + drop_params: bool, + query: str, + documents: List[Union[str, Dict[str, Any]]], + custom_llm_provider: Optional[str] = None, + top_n: Optional[int] = None, + rank_fields: Optional[List[str]] = None, + return_documents: Optional[bool] = True, + max_chunks_per_doc: Optional[int] = None, + max_tokens_per_doc: Optional[int] = None, + non_default_params: Optional[dict] = None, +) -> OptionalRerankParams: + all_non_default_params = non_default_params or {} + if query is not None: + all_non_default_params["query"] = query + if top_n is not None: + all_non_default_params["top_n"] = top_n + if documents is not None: + all_non_default_params["documents"] = documents + if return_documents is not None: + all_non_default_params["return_documents"] = return_documents + if max_chunks_per_doc is not None: + all_non_default_params["max_chunks_per_doc"] = max_chunks_per_doc + if max_tokens_per_doc is not None: + all_non_default_params["max_tokens_per_doc"] = max_tokens_per_doc + return rerank_provider_config.map_cohere_rerank_params( + model=model, + drop_params=drop_params, + query=query, + documents=documents, + custom_llm_provider=custom_llm_provider, + top_n=top_n, + rank_fields=rank_fields, + return_documents=return_documents, + max_chunks_per_doc=max_chunks_per_doc, + max_tokens_per_doc=max_tokens_per_doc, + non_default_params=all_non_default_params, + ) -- cgit v1.2.3