litellm/litellm/llms/bedrock/rerank/handler.py

import json
from typing import TYPE_CHECKING, Any, Final, cast

import httpx

import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.llms.custom_httpx.http_handler import (
    AsyncHTTPHandler,
    HTTPHandler,
    _get_httpx_client,
    get_async_httpx_client,
)
from litellm.types.llms.bedrock import BedrockPreparedRequest
from litellm.types.rerank import RerankRequest
from litellm.types.utils import RerankResponse

from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError
from .transformation import BedrockRerankConfig

if TYPE_CHECKING:
    from botocore.awsrequest import AWSPreparedRequest
else:
    AWSPreparedRequest = Any


class BedrockRerankHandler(BaseAWSLLM):
    async def arerank(
        self,
        prepared_request: BedrockPreparedRequest,
        logging_obj: LitellmLogging,
        timeout: float | httpx.Timeout | None = None,
        client: AsyncHTTPHandler | None = None,
    ):
        if client is None:
            client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
        try:
            response: Final = await client.post(
                url=prepared_request["endpoint_url"],
                headers=dict(prepared_request["prepped"].headers),
                data=prepared_request["body"],
                timeout=timeout,
                logging_obj=logging_obj,
            )
            response.raise_for_status()
        except httpx.HTTPStatusError as err:
            error_code: Final = err.response.status_code
            raise BedrockError(status_code=error_code, message=err.response.text)
        except httpx.TimeoutException:
            raise BedrockError(status_code=408, message="Timeout error occurred.")

        return BedrockRerankConfig()._transform_response(response.json())

    def rerank(
        self,
        model: str,
        query: str,
        documents: list[str | dict[str, Any]],
        optional_params: dict,
        logging_obj: LitellmLogging,
        top_n: int | None = None,
        rank_fields: list[str] | None = None,
        return_documents: bool | None = True,
        max_chunks_per_doc: int | None = None,
        _is_async: bool | None = False,
        timeout: float | httpx.Timeout | None = None,
        api_base: str | None = None,
        extra_headers: dict | None = None,
        client: HTTPHandler | AsyncHTTPHandler | None = None,
    ) -> RerankResponse:
        request_data: Final = RerankRequest(
            model=model,
            query=query,
            documents=documents,
            top_n=top_n,
            rank_fields=rank_fields,
            return_documents=return_documents,
        )
        data: Final = BedrockRerankConfig()._transform_request(request_data)

        prepared_request: Final = self._prepare_request(
            model=model,
            optional_params=optional_params,
            api_base=api_base,
            extra_headers=extra_headers,
            data=cast(dict, data),
        )

        logging_obj.pre_call(
            input=data,
            api_key="",
            additional_args={
                "complete_input_dict": data,
                "api_base": prepared_request["endpoint_url"],
                "headers": dict(prepared_request["prepped"].headers),
            },
        )

        if _is_async:
            return self.arerank(
                prepared_request,
                logging_obj=logging_obj,
                timeout=timeout,
                client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None,
            )

        if client is None or not isinstance(client, HTTPHandler):
            client = _get_httpx_client()
        try:
            response: Final = client.post(
                url=prepared_request["endpoint_url"],
                headers=dict(prepared_request["prepped"].headers),
                data=prepared_request["body"],
                timeout=timeout,
            )
            response.raise_for_status()
        except httpx.HTTPStatusError as err:
            error_code: Final = err.response.status_code
            raise BedrockError(status_code=error_code, message=err.response.text)
        except httpx.TimeoutException:
            raise BedrockError(status_code=408, message="Timeout error occurred.")

        logging_obj.post_call(
            original_response=response.text,
            api_key="",
        )

        response_json: Final = response.json()

        return BedrockRerankConfig()._transform_response(response_json)

    def _prepare_request(
        self,
        model: str,
        api_base: str | None,
        extra_headers: dict | None,
        data: dict,
        optional_params: dict,
    ) -> BedrockPreparedRequest:
        try:
            from botocore.auth import SigV4Auth
            from botocore.awsrequest import AWSRequest
        except ImportError:
            raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
        boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model)

        ### SET RUNTIME ENDPOINT ###
        _, proxy_endpoint_url = self.get_runtime_endpoint(
            api_base=api_base,
            aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint,
            aws_region_name=boto3_credentials_info.aws_region_name,
        )
        proxy_endpoint_url = proxy_endpoint_url.replace("bedrock-runtime", "bedrock-agent-runtime")
        proxy_endpoint_url = f"{proxy_endpoint_url}/rerank"
        sigv4: Final = SigV4Auth(
            boto3_credentials_info.credentials,
            "bedrock",
            boto3_credentials_info.aws_region_name,
        )
        # Make POST Request
        body: Final = json.dumps(data).encode("utf-8")

        headers = {"Content-Type": "application/json"}
        if extra_headers is not None:
            headers = {"Content-Type": "application/json", **extra_headers}
        request: Final = AWSRequest(method="POST", url=proxy_endpoint_url, data=body, headers=headers)
        sigv4.add_auth(request)
        if (
            extra_headers is not None and "Authorization" in extra_headers
        ):  # prevent sigv4 from overwriting the auth header
            request.headers["Authorization"] = extra_headers["Authorization"]
        prepped: Final = request.prepare()

        return BedrockPreparedRequest(
            endpoint_url=proxy_endpoint_url,
            prepped=prepped,
            body=body,
            data=data,
        )