diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 1f130f84283..476e22fc088 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -78,6 +78,8 @@ class HuggingFaceRerankConfig(BaseRerankConfig): "documents", "top_n", "return_documents", + "truncate", + "truncation_direction", ] def map_cohere_rerank_params( @@ -104,6 +106,10 @@ class HuggingFaceRerankConfig(BaseRerankConfig): optional_rerank_params["return_text"] = v elif k == "top_n" and v is not None: optional_rerank_params["top_n"] = v + elif k == "truncate" and isinstance(v, bool): + optional_rerank_params["truncate"] = v + elif k == "truncation_direction" and v is not None: + optional_rerank_params["truncation_direction"] = v elif k == "documents" and v is not None: optional_rerank_params["texts"] = v elif k == "query" and v is not None: @@ -151,7 +157,6 @@ class HuggingFaceRerankConfig(BaseRerankConfig): # HuggingFace API expects return_text parameter, corresponding to our return_documents parameter request_body: Final = { "raw_scores": False, - "truncate": False, "truncation_direction": "Right", } diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 6d1ab85012b..5b1d7ef0a7b 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -48,6 +48,8 @@ class OptionalRerankParams(TypedDict, total=False): max_chunks_per_doc: int | None max_tokens_per_doc: int | None instruction: str | None + truncate: ReadOnly[bool] + truncation_direction: ReadOnly[str | None] truncate_prompt_tokens: ReadOnly[int | None] truncation_side: ReadOnly[Literal["left", "right"] | None] max_tokens_per_query: ReadOnly[int | None]