mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-17 08:25:03 +00:00
Add Hosted VLLM rerank provider integration (#12738)
* Vllm rerank (#12737) * Add Hosted VLLM rerank provider integration This commit implements the Hosted VLLM rerank provider integration for LiteLLM. The integration includes: Adding Hosted VLLM as a supported rerank provider in the main rerank function Implementing the HostedVLLMRerank handler class for making API requests Creating a transformation class to convert Hosted VLLM responses to LiteLLM's standardized format The integration supports both synchronous and asynchronous rerank operations. API credentials can be provided directly or through environment variables (HOSTED_VLLM_API_KEY and HOSTED_VLLM_API_BASE). Notable features: Proper error handling for missing credentials Standard response transformation Support for common rerank parameters (top_n, return_documents, etc.) Proper token usage tracking This expands LiteLLM's rerank provider ecosystem to include Hosted VLLM alongside existing providers like Cohere, Together AI, Azure AI, and Bedrock. * refactor(rerank): use base_llm_http_handler for hosted_vllm rerank - Replace custom HostedVLLMRerank handler with base_llm_http_handler - Implement proper HostedVLLMRerankConfig inheriting from BaseRerankConfig - Follow Cohere-compatible implementation pattern - Clean up unnecessary comments * Fix lint errors in hosted_vllm rerank transformer: remove unused imports * Fix linting errors in rerank transformation modules * fix: resolve type errors in Hosted VLLM rerank module --------- Co-authored-by: Philip D'Souza <philip.dsouza@macro4.com> Co-authored-by: Philip D'Souza <philip.a.dsouza@gmail.com> * added a few tests --------- Co-authored-by: Philip D'Souza <philip.dsouza@macro4.com> Co-authored-by: Philip D'Souza <philip.a.dsouza@gmail.com>
This commit is contained in:
Jugal D. Bhatt
co-authored by
Philip D'Souza
Philip D'Souza
Philip D'Souza
Philip D'Souza
parent
56eacd0a38
commit
7c49197f29
@@ -1,14 +1,13 @@
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
|
||||
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 OptionalRerankParams, RerankRequest
|
||||
from litellm.types.utils import RerankResponse
|
||||
from litellm.types.rerank import OptionalRerankParams, RerankRequest, RerankResponse
|
||||
|
||||
from ..common_utils import CohereError
|
||||
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
"""
|
||||
Transformation logic for Hosted VLLM rerank
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from litellm.types.rerank import (
|
||||
RerankBilledUnits,
|
||||
RerankResponse,
|
||||
RerankResponseDocument,
|
||||
RerankResponseMeta,
|
||||
RerankResponseResult,
|
||||
RerankTokens,
|
||||
OptionalRerankParams,
|
||||
RerankRequest,
|
||||
)
|
||||
|
||||
import httpx
|
||||
|
||||
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
|
||||
|
||||
|
||||
class HostedVLLMRerankError(BaseLLMException):
|
||||
def __init__(
|
||||
self,
|
||||
status_code: int,
|
||||
message: str,
|
||||
headers: Optional[Union[dict, httpx.Headers]] = None,
|
||||
):
|
||||
super().__init__(status_code=status_code, message=message, headers=headers)
|
||||
|
||||
|
||||
class HostedVLLMRerankConfig(BaseRerankConfig):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def get_complete_url(self, api_base: Optional[str], model: str) -> str:
|
||||
if api_base:
|
||||
# Remove trailing slashes and ensure clean base URL
|
||||
api_base = api_base.rstrip("/")
|
||||
if not api_base.endswith("/v1/rerank"):
|
||||
api_base = f"{api_base}/v1/rerank"
|
||||
return api_base
|
||||
raise ValueError("api_base must be provided for Hosted VLLM rerank")
|
||||
|
||||
def get_supported_cohere_rerank_params(self, model: str) -> list:
|
||||
return [
|
||||
"query",
|
||||
"documents",
|
||||
"top_n",
|
||||
"rank_fields",
|
||||
"return_documents",
|
||||
]
|
||||
|
||||
def map_cohere_rerank_params(
|
||||
self,
|
||||
non_default_params: Optional[dict],
|
||||
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,
|
||||
) -> OptionalRerankParams:
|
||||
"""
|
||||
Map parameters for Hosted VLLM rerank
|
||||
"""
|
||||
if max_chunks_per_doc is not None:
|
||||
raise ValueError("Hosted VLLM does not support max_chunks_per_doc")
|
||||
|
||||
return OptionalRerankParams(
|
||||
query=query,
|
||||
documents=documents,
|
||||
top_n=top_n,
|
||||
rank_fields=rank_fields,
|
||||
return_documents=return_documents,
|
||||
)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
) -> dict:
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key"
|
||||
|
||||
default_headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"accept": "application/json",
|
||||
"content-type": "application/json",
|
||||
}
|
||||
|
||||
# If 'Authorization' is provided in headers, it overrides the default.
|
||||
if "Authorization" in headers:
|
||||
default_headers["Authorization"] = headers["Authorization"]
|
||||
|
||||
# Merge other headers, overriding any default ones except Authorization
|
||||
return {**default_headers, **headers}
|
||||
|
||||
def transform_rerank_request(
|
||||
self,
|
||||
model: str,
|
||||
optional_rerank_params: OptionalRerankParams,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
if "query" not in optional_rerank_params:
|
||||
raise ValueError("query is required for Hosted VLLM rerank")
|
||||
if "documents" not in optional_rerank_params:
|
||||
raise ValueError("documents is required for Hosted VLLM rerank")
|
||||
|
||||
rerank_request = RerankRequest(
|
||||
model=model,
|
||||
query=optional_rerank_params["query"],
|
||||
documents=optional_rerank_params["documents"],
|
||||
top_n=optional_rerank_params.get("top_n", None),
|
||||
rank_fields=optional_rerank_params.get("rank_fields", None),
|
||||
return_documents=optional_rerank_params.get("return_documents", None),
|
||||
)
|
||||
return rerank_request.model_dump(exclude_none=True)
|
||||
|
||||
def transform_rerank_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: RerankResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
api_key: Optional[str] = None,
|
||||
request_data: dict = {},
|
||||
optional_params: dict = {},
|
||||
litellm_params: dict = {},
|
||||
) -> RerankResponse:
|
||||
"""
|
||||
Process response from Hosted VLLM rerank API
|
||||
"""
|
||||
try:
|
||||
raw_response_json = raw_response.json()
|
||||
except Exception:
|
||||
raise ValueError(
|
||||
f"Error parsing response: {raw_response.text}, status_code={raw_response.status_code}"
|
||||
)
|
||||
|
||||
return RerankResponse(**raw_response_json)
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
|
||||
) -> BaseLLMException:
|
||||
return HostedVLLMRerankError(message=error_message, status_code=status_code, headers=headers)
|
||||
|
||||
def _transform_response(self, response: dict) -> RerankResponse:
|
||||
# Extract usage information
|
||||
usage_data = response.get("usage", {})
|
||||
_billed_units = RerankBilledUnits(total_tokens=usage_data.get("total_tokens", 0))
|
||||
_tokens = RerankTokens(input_tokens=usage_data.get("total_tokens", 0))
|
||||
rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
|
||||
|
||||
# Extract results
|
||||
_results: Optional[List[dict]] = response.get("results")
|
||||
|
||||
if _results is None:
|
||||
raise ValueError(f"No results found in the response={response}")
|
||||
|
||||
rerank_results: List[RerankResponseResult] = []
|
||||
|
||||
for result in _results:
|
||||
# Validate required fields exist
|
||||
if not all(key in result for key in ["index", "relevance_score"]):
|
||||
raise ValueError(f"Missing required fields in the result={result}")
|
||||
|
||||
# Get document data if it exists
|
||||
document_data = result.get("document", {})
|
||||
document = (
|
||||
RerankResponseDocument(text=str(document_data.get("text", "")))
|
||||
if document_data
|
||||
else None
|
||||
)
|
||||
|
||||
# Create typed result
|
||||
rerank_result = RerankResponseResult(
|
||||
index=int(result["index"]),
|
||||
relevance_score=float(result["relevance_score"]),
|
||||
)
|
||||
|
||||
# Only add document if it exists
|
||||
if document:
|
||||
rerank_result["document"] = document
|
||||
|
||||
rerank_results.append(rerank_result)
|
||||
|
||||
return RerankResponse(
|
||||
id=response.get("id") or str(uuid.uuid4()),
|
||||
results=rerank_results,
|
||||
meta=rerank_meta,
|
||||
)
|
||||
@@ -75,7 +75,7 @@ def rerank( # noqa: PLR0915
|
||||
query: str,
|
||||
documents: List[Union[str, Dict[str, Any]]],
|
||||
custom_llm_provider: Optional[
|
||||
Literal["cohere", "together_ai", "azure_ai", "infinity", "litellm_proxy"]
|
||||
Literal["cohere", "together_ai", "azure_ai", "infinity", "litellm_proxy", "hosted_vllm"]
|
||||
] = None,
|
||||
top_n: Optional[int] = None,
|
||||
rank_fields: Optional[List[str]] = None,
|
||||
@@ -323,6 +323,39 @@ def rerank( # noqa: PLR0915
|
||||
logging_obj=litellm_logging_obj,
|
||||
client=client,
|
||||
)
|
||||
elif _custom_llm_provider == "hosted_vllm":
|
||||
# Implement Hosted VLLM rerank logic
|
||||
api_key = (
|
||||
dynamic_api_key
|
||||
or optional_params.api_key
|
||||
or get_secret_str("HOSTED_VLLM_API_KEY")
|
||||
)
|
||||
|
||||
api_base = (
|
||||
dynamic_api_base
|
||||
or optional_params.api_base
|
||||
or get_secret_str("HOSTED_VLLM_API_BASE")
|
||||
)
|
||||
|
||||
if api_base is None:
|
||||
raise ValueError(
|
||||
"api_base must be provided for Hosted VLLM rerank. Set in call or via HOSTED_VLLM_API_BASE env var."
|
||||
)
|
||||
|
||||
response = base_llm_http_handler.rerank(
|
||||
model=model,
|
||||
custom_llm_provider=_custom_llm_provider,
|
||||
provider_config=rerank_provider_config,
|
||||
optional_rerank_params=optional_rerank_params,
|
||||
logging_obj=litellm_logging_obj,
|
||||
timeout=optional_params.timeout,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
_is_async=_is_async,
|
||||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
)
|
||||
else:
|
||||
# Generic handler for all providers that use base_llm_http_handler
|
||||
# Provider-specific logic (API key validation, URL generation, etc.)
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
import sys
|
||||
import os
|
||||
import pytest
|
||||
from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig
|
||||
from litellm.types.rerank import OptionalRerankParams, RerankResponse, RerankResponseResult, RerankResponseMeta, RerankBilledUnits, RerankTokens, RerankResponseDocument
|
||||
|
||||
class TestHostedVLLMRerankTransform:
|
||||
def setup_method(self):
|
||||
self.config = HostedVLLMRerankConfig()
|
||||
self.model = "hosted-vllm-model"
|
||||
|
||||
def test_map_cohere_rerank_params_basic(self):
|
||||
params = self.config.map_cohere_rerank_params(
|
||||
non_default_params=None,
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
query="test query",
|
||||
documents=["doc1", "doc2"],
|
||||
top_n=2,
|
||||
rank_fields=["field1"],
|
||||
return_documents=True,
|
||||
)
|
||||
assert params["query"] == "test query"
|
||||
assert params["documents"] == ["doc1", "doc2"]
|
||||
assert params["top_n"] == 2
|
||||
assert params["rank_fields"] == ["field1"]
|
||||
assert params["return_documents"] is True
|
||||
|
||||
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"):
|
||||
self.config.map_cohere_rerank_params(
|
||||
non_default_params=None,
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
query="test query",
|
||||
documents=["doc1"],
|
||||
max_chunks_per_doc=5
|
||||
)
|
||||
|
||||
def test_get_complete_url(self):
|
||||
base = "https://api.example.com"
|
||||
url = self.config.get_complete_url(base, self.model)
|
||||
assert url == "https://api.example.com/v1/rerank"
|
||||
# Already ends with /v1/rerank
|
||||
url2 = self.config.get_complete_url("https://api.example.com/v1/rerank", self.model)
|
||||
assert url2 == "https://api.example.com/v1/rerank"
|
||||
# Raises if api_base is None
|
||||
with pytest.raises(ValueError):
|
||||
self.config.get_complete_url(None, self.model)
|
||||
|
||||
def test_transform_response(self):
|
||||
response_dict = {
|
||||
"id": "abc123",
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.9, "document": {"text": "doc1 text"}},
|
||||
{"index": 1, "relevance_score": 0.7, "document": {"text": "doc2 text"}},
|
||||
],
|
||||
"usage": {"total_tokens": 42}
|
||||
}
|
||||
result = self.config._transform_response(response_dict)
|
||||
assert result.id == "abc123"
|
||||
assert len(result.results) == 2
|
||||
assert result.results[0]["index"] == 0
|
||||
assert result.results[0]["relevance_score"] == 0.9
|
||||
assert result.results[0]["document"]["text"] == "doc1 text"
|
||||
assert result.meta["billed_units"]["total_tokens"] == 42
|
||||
assert result.meta["tokens"]["input_tokens"] == 42
|
||||
|
||||
def test_transform_response_missing_results(self):
|
||||
response_dict = {"id": "abc123", "usage": {"total_tokens": 10}}
|
||||
with pytest.raises(ValueError, match="No results found in the response="):
|
||||
self.config._transform_response(response_dict)
|
||||
|
||||
def test_transform_response_missing_required_fields(self):
|
||||
response_dict = {
|
||||
"id": "abc123",
|
||||
"results": [{"relevance_score": 0.5}],
|
||||
"usage": {"total_tokens": 10}
|
||||
}
|
||||
with pytest.raises(ValueError, match="Missing required fields in the result="):
|
||||
self.config._transform_response(response_dict)
|
||||
Reference in New Issue
Block a user