about summary refs log tree commit diff
path: root/.venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py
diff options
context:
space:
mode:
authorS. Solomon Darnell2025-03-28 21:52:21 -0500
committerS. Solomon Darnell2025-03-28 21:52:21 -0500
commit4a52a71956a8d46fcb7294ac71734504bb09bcc2 (patch)
treeee3dc5af3b6313e921cd920906356f5d4febc4ed /.venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py
parentcc961e04ba734dd72309fb548a2f97d67d578813 (diff)
downloadgn-ai-master.tar.gz
two version of R2R are here HEAD master
Diffstat (limited to '.venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py')
-rw-r--r--.venv/lib/python3.12/site-packages/litellm/rerank_api/rerank_utils.py46
1 files changed, 46 insertions, 0 deletions
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,
+    )