123 lines
4.1 KiB
Python
123 lines
4.1 KiB
Python
|
"""
|
||
|
|
Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank` format.
|
||
|
|
|
||
|
|
Why separate file? Make it easy to see how transformation works
|
||
|
|
"""
|
||
|
|
|
||
|
|
from litellm._uuid import uuid
|
||
|
|
from typing import List, Optional
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
import litellm
|
||
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||
|
|
from litellm.llms.cohere.rerank.transformation import CohereRerankConfig
|
||
|
|
from litellm.secret_managers.main import get_secret_str
|
||
|
|
from litellm.types.rerank import (
|
||
|
|
RerankBilledUnits,
|
||
|
|
RerankResponse,
|
||
|
|
RerankResponseDocument,
|
||
|
|
RerankResponseMeta,
|
||
|
|
RerankResponseResult,
|
||
|
|
RerankTokens,
|
||
|
|
)
|
||
|
|
|
||
|
|
from ..common_utils import InfinityError
|
||
|
|
|
||
|
|
|
||
|
|
class InfinityRerankConfig(CohereRerankConfig):
|
||
|
|
def get_complete_url(
|
||
|
|
self,
|
||
|
|
api_base: Optional[str],
|
||
|
|
model: str,
|
||
|
|
optional_params: Optional[dict] = None,
|
||
|
|
) -> str:
|
||
|
|
if api_base is None:
|
||
|
|
raise ValueError("api_base is required for Infinity rerank")
|
||
|
|
# Remove trailing slashes and ensure clean base URL
|
||
|
|
api_base = api_base.rstrip("/")
|
||
|
|
if not api_base.endswith("/rerank"):
|
||
|
|
api_base = f"{api_base}/rerank"
|
||
|
|
return api_base
|
||
|
|
|
||
|
|
def validate_environment(
|
||
|
|
self,
|
||
|
|
headers: dict,
|
||
|
|
model: str,
|
||
|
|
api_key: Optional[str] = None,
|
||
|
|
optional_params: Optional[dict] = None,
|
||
|
|
) -> dict:
|
||
|
|
if api_key is None:
|
||
|
|
api_key = (
|
||
|
|
get_secret_str("INFINITY_API_KEY")
|
||
|
|
or get_secret_str("INFINITY_API_KEY")
|
||
|
|
or litellm.infinity_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_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:
|
||
|
|
"""
|
||
|
|
Transform Infinity rerank response
|
||
|
|
|
||
|
|
No transformation required, Infinity follows Cohere API response format
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
raw_response_json = raw_response.json()
|
||
|
|
except Exception:
|
||
|
|
raise InfinityError(
|
||
|
|
message=raw_response.text, status_code=raw_response.status_code
|
||
|
|
)
|
||
|
|
|
||
|
|
_billed_units = RerankBilledUnits(**raw_response_json.get("usage", {}))
|
||
|
|
_tokens = RerankTokens(
|
||
|
|
input_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0),
|
||
|
|
output_tokens=(
|
||
|
|
raw_response_json.get("usage", {}).get("total_tokens", 0)
|
||
|
|
- raw_response_json.get("usage", {}).get("prompt_tokens", 0)
|
||
|
|
),
|
||
|
|
)
|
||
|
|
rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
|
||
|
|
|
||
|
|
cohere_results: List[RerankResponseResult] = []
|
||
|
|
if raw_response_json.get("results"):
|
||
|
|
for result in raw_response_json.get("results"):
|
||
|
|
_rerank_response = RerankResponseResult(
|
||
|
|
index=result.get("index"),
|
||
|
|
relevance_score=result.get("relevance_score"),
|
||
|
|
)
|
||
|
|
if result.get("document"):
|
||
|
|
_rerank_response["document"] = RerankResponseDocument(
|
||
|
|
text=result.get("document")
|
||
|
|
)
|
||
|
|
cohere_results.append(_rerank_response)
|
||
|
|
if cohere_results is None:
|
||
|
|
raise ValueError(f"No results found in the response={raw_response_json}")
|
||
|
|
|
||
|
|
return RerankResponse(
|
||
|
|
id=raw_response_json.get("id") or str(uuid.uuid4()),
|
||
|
|
results=cohere_results,
|
||
|
|
meta=rerank_meta,
|
||
|
|
) # Return response
|