diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 22782c1300..5371b9a4b6 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -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 diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py new file mode 100644 index 0000000000..419327d9d5 --- /dev/null +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -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, + ) \ No newline at end of file diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index cc80d357d1..678d4963ac 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -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.) diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py new file mode 100644 index 0000000000..3d4d0e7b70 --- /dev/null +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -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) \ No newline at end of file