mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge pull request #39363 from BerriAI/litellm_hosted_vllm_rerank_truncate_params
fix(hosted_vllm): forward truncate_prompt_tokens on rerank requests
This commit is contained in:
Mateo Wang
GitHub
commit
cd9ed8bad5
3 changed files with 175 additions and 10 deletions
|
|
@ -3,16 +3,20 @@ Transformation logic for Hosted VLLM rerank
|
|||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.rerank import (
|
||||
HostedVLLMRerankTruncationParams,
|
||||
OptionalRerankParams,
|
||||
RerankBilledUnits,
|
||||
RerankRequest,
|
||||
|
|
@ -34,6 +38,13 @@ class HostedVLLMRerankError(BaseLLMException):
|
|||
super().__init__(status_code=status_code, message=message, headers=headers)
|
||||
|
||||
|
||||
def validated_truncation_params(non_default_params: Mapping[str, object] | None) -> HostedVLLMRerankTruncationParams:
|
||||
try:
|
||||
return HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({}))
|
||||
except ValidationError as error:
|
||||
raise UnsupportedParamsError(status_code=400, message=f"hosted_vllm rerank: {error}") from error
|
||||
|
||||
|
||||
class HostedVLLMRerankConfig(BaseRerankConfig):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
|
@ -62,7 +73,11 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
|
|||
"top_n",
|
||||
"rank_fields",
|
||||
"return_documents",
|
||||
"max_tokens_per_doc",
|
||||
"instruction",
|
||||
"truncate_prompt_tokens",
|
||||
"truncation_side",
|
||||
"max_tokens_per_query",
|
||||
]
|
||||
|
||||
def map_cohere_rerank_params(
|
||||
|
|
@ -100,7 +115,15 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
|
|||
if instruction is not None:
|
||||
mapped_params["instruction"] = instruction
|
||||
|
||||
return dict(mapped_params)
|
||||
truncation: Final = validated_truncation_params(non_default_params)
|
||||
forwarded: Final[OptionalRerankParams] = {
|
||||
**mapped_params,
|
||||
"max_tokens_per_doc": max_tokens_per_doc,
|
||||
"truncate_prompt_tokens": truncation.truncate_prompt_tokens,
|
||||
"truncation_side": truncation.truncation_side,
|
||||
"max_tokens_per_query": truncation.max_tokens_per_query,
|
||||
}
|
||||
return dict(forwarded)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
|
|
@ -138,6 +161,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
|
|||
if "documents" not in optional_rerank_params:
|
||||
raise ValueError("documents is required for Hosted VLLM rerank")
|
||||
|
||||
truncation: Final = HostedVLLMRerankTruncationParams.model_validate(optional_rerank_params)
|
||||
rerank_request: Final = RerankRequest(
|
||||
model=model,
|
||||
query=optional_rerank_params["query"],
|
||||
|
|
@ -146,6 +170,10 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
|
|||
rank_fields=optional_rerank_params.get("rank_fields", None),
|
||||
return_documents=optional_rerank_params.get("return_documents", None),
|
||||
instruction=optional_rerank_params.get("instruction", None),
|
||||
max_tokens_per_doc=truncation.max_tokens_per_doc,
|
||||
truncate_prompt_tokens=truncation.truncate_prompt_tokens,
|
||||
truncation_side=truncation.truncation_side,
|
||||
max_tokens_per_query=truncation.max_tokens_per_query,
|
||||
)
|
||||
return rerank_request.model_dump(exclude_none=True)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,8 +4,10 @@ https://docs.cohere.com/reference/rerank
|
|||
|
||||
"""
|
||||
|
||||
from pydantic import BaseModel, PrivateAttr
|
||||
from typing_extensions import Required, TypedDict
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, PrivateAttr
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
|
||||
class RerankRequest(BaseModel):
|
||||
|
|
@ -21,6 +23,18 @@ class RerankRequest(BaseModel):
|
|||
# (e.g. hosted vLLM / Qwen3-Reranker, DeepInfra). Omitted from the outgoing
|
||||
# request when None, so this is fully backward-compatible.
|
||||
instruction: str | None = None
|
||||
truncate_prompt_tokens: int | None = None
|
||||
truncation_side: Literal["left", "right"] | None = None
|
||||
max_tokens_per_query: int | None = None
|
||||
|
||||
|
||||
class HostedVLLMRerankTruncationParams(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
truncate_prompt_tokens: int | None = None
|
||||
truncation_side: Literal["left", "right"] | None = None
|
||||
max_tokens_per_query: int | None = None
|
||||
max_tokens_per_doc: int | None = None
|
||||
|
||||
|
||||
class OptionalRerankParams(TypedDict, total=False):
|
||||
|
|
@ -32,6 +46,9 @@ class OptionalRerankParams(TypedDict, total=False):
|
|||
max_chunks_per_doc: int | None
|
||||
max_tokens_per_doc: int | None
|
||||
instruction: str | None
|
||||
truncate_prompt_tokens: ReadOnly[int | None]
|
||||
truncation_side: ReadOnly[Literal["left", "right"] | None]
|
||||
max_tokens_per_query: ReadOnly[int | None]
|
||||
|
||||
|
||||
class RerankBilledUnits(TypedDict, total=False):
|
||||
|
|
|
|||
|
|
@ -1,8 +1,14 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig
|
||||
from litellm.rerank_api.rerank_utils import get_optional_rerank_params
|
||||
from litellm.types.rerank import (
|
||||
|
|
@ -87,9 +93,7 @@ class TestHostedVLLMRerankTransform:
|
|||
assert "instruction" not in body
|
||||
|
||||
def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self):
|
||||
with pytest.raises(
|
||||
ValueError, match="Hosted VLLM does not support max_chunks_per_doc"
|
||||
):
|
||||
with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"):
|
||||
self.config.map_cohere_rerank_params(
|
||||
non_default_params=None,
|
||||
model=self.model,
|
||||
|
|
@ -104,12 +108,10 @@ class TestHostedVLLMRerankTransform:
|
|||
url = self.config.get_complete_url(base, self.model)
|
||||
assert url == "https://api.example.com/rerank"
|
||||
# Already ends with /rerank
|
||||
url2 = self.config.get_complete_url(
|
||||
"https://api.example.com/rerank", self.model
|
||||
)
|
||||
url2 = self.config.get_complete_url("https://api.example.com/rerank", self.model)
|
||||
assert url2 == "https://api.example.com/rerank"
|
||||
# Raises if api_base is None
|
||||
with pytest.raises(ValueError, match='api_base must be provided for Hosted VLLM rerank'):
|
||||
with pytest.raises(ValueError, match="api_base must be provided for Hosted VLLM rerank"):
|
||||
self.config.get_complete_url(None, self.model)
|
||||
|
||||
def test_transform_response(self):
|
||||
|
|
@ -173,3 +175,121 @@ class TestGetOptionalRerankParamsInstruction:
|
|||
documents=["doc1", "doc2"],
|
||||
)
|
||||
assert "instruction" not in params
|
||||
|
||||
|
||||
class TestHostedVLLMRerankTruncationParams:
|
||||
def setup_method(self):
|
||||
self.config = HostedVLLMRerankConfig()
|
||||
self.model = "hosted-vllm-model"
|
||||
|
||||
def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self):
|
||||
params: Final = self.config.map_cohere_rerank_params(
|
||||
non_default_params={
|
||||
"truncate_prompt_tokens": 512,
|
||||
"truncation_side": "left",
|
||||
"max_tokens_per_query": 64,
|
||||
"metadata": {"user_api_key": "sk-test"},
|
||||
},
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
query="test query",
|
||||
documents=["doc1", "doc2"],
|
||||
max_tokens_per_doc=128,
|
||||
)
|
||||
assert params["truncate_prompt_tokens"] == 512
|
||||
assert params["truncation_side"] == "left"
|
||||
assert params["max_tokens_per_query"] == 64
|
||||
assert params["max_tokens_per_doc"] == 128
|
||||
assert "metadata" not in params
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad_params",
|
||||
[{"truncation_side": "middle"}, {"truncate_prompt_tokens": "lots"}, {"max_tokens_per_query": -1.5}],
|
||||
)
|
||||
def test_map_cohere_rerank_params_rejects_invalid_truncation_params_as_400(self, bad_params: dict[str, object]):
|
||||
with pytest.raises(litellm.UnsupportedParamsError) as raised:
|
||||
self.config.map_cohere_rerank_params(
|
||||
non_default_params=dict(bad_params),
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
query="test query",
|
||||
documents=["doc1", "doc2"],
|
||||
)
|
||||
assert raised.value.status_code == 400
|
||||
assert next(iter(bad_params)) in str(raised.value)
|
||||
|
||||
def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self):
|
||||
params: Final = self.config.map_cohere_rerank_params(
|
||||
non_default_params={"metadata": {"user_api_key": "sk-test"}},
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
query="test query",
|
||||
documents=["doc1", "doc2"],
|
||||
)
|
||||
body: Final = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={})
|
||||
truncation_keys: Final = {
|
||||
"truncate_prompt_tokens",
|
||||
"truncation_side",
|
||||
"max_tokens_per_query",
|
||||
"max_tokens_per_doc",
|
||||
}
|
||||
assert not truncation_keys & body.keys()
|
||||
assert body == {
|
||||
"model": self.model,
|
||||
"query": "test query",
|
||||
"documents": ["doc1", "doc2"],
|
||||
"return_documents": True,
|
||||
}
|
||||
|
||||
def test_transform_request_forwards_truncation_params(self):
|
||||
body: Final = self.config.transform_rerank_request(
|
||||
model=self.model,
|
||||
optional_rerank_params={
|
||||
"query": "test query",
|
||||
"documents": ["doc1", "doc2"],
|
||||
"truncate_prompt_tokens": 512,
|
||||
"truncation_side": "left",
|
||||
"max_tokens_per_query": 64,
|
||||
"max_tokens_per_doc": 128,
|
||||
},
|
||||
headers={},
|
||||
)
|
||||
assert body["truncate_prompt_tokens"] == 512
|
||||
assert body["truncation_side"] == "left"
|
||||
assert body["max_tokens_per_query"] == 64
|
||||
assert body["max_tokens_per_doc"] == 128
|
||||
|
||||
def test_transform_request_omits_truncation_params_when_absent(self):
|
||||
body: Final = self.config.transform_rerank_request(
|
||||
model=self.model,
|
||||
optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]},
|
||||
headers={},
|
||||
)
|
||||
assert "truncate_prompt_tokens" not in body
|
||||
assert "truncation_side" not in body
|
||||
assert "max_tokens_per_query" not in body
|
||||
assert "max_tokens_per_doc" not in body
|
||||
|
||||
def test_rerank_sends_truncate_prompt_tokens_to_vllm(self):
|
||||
client: Final = HTTPHandler()
|
||||
mock_response: Final = MagicMock(spec=httpx.Response)
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"id": "score-1",
|
||||
"results": [{"index": 0, "relevance_score": 0.5}],
|
||||
"usage": {"total_tokens": 512},
|
||||
}
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
litellm.rerank(
|
||||
model="hosted_vllm/BAAI/bge-reranker-base",
|
||||
api_base="http://vllm.local:8000",
|
||||
query="List all the unique case ids",
|
||||
documents=["a document longer than the reranker context window"],
|
||||
truncate_prompt_tokens=512,
|
||||
truncation_side="left",
|
||||
client=client,
|
||||
)
|
||||
sent_body: Final = json.loads(mock_post.call_args.kwargs["data"])
|
||||
assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank"
|
||||
assert sent_body["truncate_prompt_tokens"] == 512
|
||||
assert sent_body["truncation_side"] == "left"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue