bedrock.py

from __future__ import annotations

import os
import re
import inspect
from typing import Literal, Callable, Awaitable, cast
from dataclasses import field, dataclass

import httpx2

from .._types import NOT_GIVEN, NotGiven
from .._utils import asyncify
from .._httpx2 import normalize_httpx_url, request_not_read_exceptions
from .._models import FinalRequestOptions
from .._provider import _Provider, _create_provider, _ProviderRuntime
from .._exceptions import OpenAIError
from ..lib._bedrock_auth import (
    BedrockAwsAuth,
    BedrockAwsAuthConfig,
    AwsCredentialsProvider,
)

BedrockTokenProvider = Callable[[], "str | Awaitable[str]"]
BedrockEndpoint = Literal["mantle", "runtime"]

_AWS_SIGNING_HEADERS = ("authorization", "x-amz-content-sha256", "x-amz-date", "x-amz-security-token")
_AWS_REGION = re.compile(r"^[a-z]{2,8}(?:-[a-z0-9]+)+-\d+$")


def _normalize_optional_string(value: str | None) -> str | None:
    if value is None:
        return None

    normalized = value.strip()
    return normalized or None


def _normalize_base_url(base_url: str | httpx2.URL) -> httpx2.URL:
    url = normalize_httpx_url(base_url)
    path = url.path.rstrip("/")
    responses_match = re.search(r"/responses(?:/.*)?$", path)
    if responses_match is not None:
        path = path[: responses_match.start()]

    return url.copy_with(path=path or "/")


def _runtime_dns_suffixes(region: str) -> tuple[str, str]:
    if region.startswith("cn-"):
        return "amazonaws.com.cn", "api.amazonwebservices.com.cn"
    if region.startswith("eusc-"):
        return "amazonaws.eu", "api.amazonwebservices.eu"
    if region.startswith("us-iso-"):
        return "c2s.ic.gov", "api.aws.ic.gov"
    if region.startswith("us-isob-"):
        return "sc2s.sgov.gov", "api.aws.scloud"
    if region.startswith("eu-isoe-"):
        return "cloud.adc-e.uk", "api.cloud-aws.adc-e.uk"
    if region.startswith("us-isof-"):
        return "csp.hci.ic.gov", "api.aws.hci.ic.gov"
    return "amazonaws.com", "api.aws"


def _parse_bedrock_endpoint_hostname(hostname: str) -> tuple[BedrockEndpoint, str] | None:
    service, separator, remainder = hostname.removesuffix(".").lower().partition(".")
    region, region_separator, suffix = remainder.partition(".")
    if not separator or not region_separator or _AWS_REGION.fullmatch(region) is None:
        return None

    if service == "bedrock-mantle" and suffix == "api.aws":
        return "mantle", region
    if service in {"bedrock-runtime", "bedrock-runtime-fips"} and suffix in _runtime_dns_suffixes(region):
        return "runtime", region
    return None


def _validate_bedrock_region(region: str | None) -> None:
    if region is not None and _AWS_REGION.fullmatch(region) is None:
        raise OpenAIError("The Bedrock AWS `region` is invalid. Use a standard AWS region such as `us-east-1`.")


def _validate_canonical_bedrock_endpoint(
    base_url: httpx2.URL, *, endpoint: BedrockEndpoint, region: str | None
) -> None:
    canonical_endpoint = _parse_bedrock_endpoint_hostname(base_url.host)
    if canonical_endpoint is None:
        return

    canonical_family, canonical_region = canonical_endpoint
    if base_url.scheme != "https":
        raise OpenAIError("Canonical Amazon Bedrock endpoints require HTTPS.")
    if canonical_family != endpoint:
        raise OpenAIError(
            f"The Bedrock {canonical_family} hostname does not match the selected `{endpoint}` endpoint. "
            f"Set `endpoint='{canonical_family}'` to use this hostname."
        )
    if region is not None and canonical_region != region:
        raise OpenAIError(
            f"The Bedrock endpoint region `{canonical_region}` does not match the configured AWS region `{region}`."
        )


def _default_bedrock_base_url(endpoint: BedrockEndpoint, region: str) -> httpx2.URL:
    hostname = (
        f"bedrock-runtime.{region}.{_runtime_dns_suffixes(region)[0]}"
        if endpoint == "runtime"
        else f"bedrock-mantle.{region}.api.aws"
    )
    return _normalize_base_url(f"https://{hostname}/openai/v1")


def _same_origin(left: httpx2.URL, right: httpx2.URL) -> bool:
    return (left.scheme, left.host, left.port) == (right.scheme, right.host, right.port)


def _body_for_signing(request: httpx2.Request) -> bytes:
    try:
        return request.content
    except request_not_read_exceptions() as exc:
        raise OpenAIError(
            "Bedrock SigV4 authentication requires a replayable request body. "
            "Buffer the body before sending or use bearer authentication."
        ) from exc


def _assert_provider_owns_authorization(request: httpx2.Request) -> None:
    if "Authorization" in request.headers:
        raise OpenAIError("Bedrock provider authentication cannot be combined with a custom `Authorization` header.")


def _without_redirects(options: FinalRequestOptions) -> FinalRequestOptions:
    if options.follow_redirects:
        raise OpenAIError(
            "Bedrock SigV4 authentication does not support automatic redirects. "
            "Send a new request to the redirect target so it can be signed again."
        )
    options.follow_redirects = False
    return options


class _BedrockBearerAuth:
    def __init__(self, token_provider: BedrockTokenProvider, *, base_url: httpx2.URL) -> None:
        self._token_provider = token_provider
        self._base_url = base_url

    def _validate_request(self, request: httpx2.Request) -> None:
        _assert_provider_owns_authorization(request)
        if not _same_origin(request.url, self._base_url):
            raise OpenAIError(
                "Refusing to authenticate a Bedrock request for an origin other than the configured provider URL."
            )

    def _resolve_token(self) -> str:
        try:
            token = cast(object, self._token_provider())
        except OpenAIError:
            raise
        except Exception as exc:
            raise OpenAIError("Failed to resolve a bearer credential for Bedrock.") from exc

        if inspect.isawaitable(token):
            close = getattr(token, "close", None)
            if callable(close):
                close()
            raise OpenAIError("An async Bedrock token provider requires `AsyncOpenAI`.")
        if not isinstance(token, str) or not token.strip():
            raise OpenAIError("The Bedrock bearer credential provider must return a non-empty string.")
        return token

    async def _resolve_token_async(self) -> str:
        try:
            token = cast(object, self._token_provider())
            if inspect.isawaitable(token):
                token = await token
        except OpenAIError:
            raise
        except Exception as exc:
            raise OpenAIError("Failed to resolve a bearer credential for Bedrock.") from exc

        if not isinstance(token, str) or not token.strip():
            raise OpenAIError("The Bedrock bearer credential provider must return a non-empty string.")
        return token

    def prepare_request(self, request: httpx2.Request) -> None:
        self._validate_request(request)
        request.headers["Authorization"] = f"Bearer {self._resolve_token()}"

    async def prepare_async_request(self, request: httpx2.Request) -> None:
        self._validate_request(request)
        request.headers["Authorization"] = f"Bearer {await self._resolve_token_async()}"


class _BedrockSigV4Auth:
    def __init__(
        self,
        *,
        config: BedrockAwsAuthConfig,
        base_url: httpx2.URL,
        auth: BedrockAwsAuth | None = None,
    ) -> None:
        self._config = config
        self._base_url = base_url
        self._auth = auth

    def _validate_request(self, request: httpx2.Request) -> bytes:
        _assert_provider_owns_authorization(request)
        if not _same_origin(request.url, self._base_url):
            raise OpenAIError(
                "Refusing to sign a Bedrock request for an origin other than the configured provider URL."
            )

        canonical_endpoint = _parse_bedrock_endpoint_hostname(request.url.host)
        if canonical_endpoint is not None:
            endpoint, region = canonical_endpoint
            expected_endpoint = "runtime" if self._config.service == "bedrock" else "mantle"
            if endpoint != expected_endpoint:
                raise OpenAIError(
                    f"The Bedrock {endpoint} hostname does not match the selected `{expected_endpoint}` endpoint."
                )
            if region != self._config.region:
                raise OpenAIError(
                    f"The Bedrock endpoint region `{region}` does not match the SigV4 region `{self._config.region}`."
                )

        return _body_for_signing(request)

    def _sign(self, request: httpx2.Request, *, auth: BedrockAwsAuth, body: bytes) -> None:
        for header in _AWS_SIGNING_HEADERS:
            request.headers.pop(header, None)

        signed_headers = auth.sign(
            method=request.method,
            url=str(request.url),
            headers=dict(request.headers),
            body=body,
        )
        request.headers.clear()
        request.headers.update(signed_headers)

    def prepare_request(self, request: httpx2.Request) -> None:
        body = self._validate_request(request)
        if self._auth is None:
            self._auth = BedrockAwsAuth(self._config)
        self._sign(request, auth=self._auth, body=body)

    async def prepare_async_request(self, request: httpx2.Request) -> None:
        body = self._validate_request(request)
        if self._auth is None:
            self._auth = await asyncify(BedrockAwsAuth)(self._config)

        signed_headers = await asyncify(self._auth.sign)(
            method=request.method,
            url=str(request.url),
            headers={
                name: value for name, value in request.headers.items() if name.lower() not in _AWS_SIGNING_HEADERS
            },
            body=body,
        )
        request.headers.clear()
        request.headers.update(signed_headers)


@dataclass
class _BedrockProviderRuntime(_ProviderRuntime):
    region: str | None = None


@dataclass(frozen=True)
class _BedrockProviderDefinition:
    endpoint: BedrockEndpoint
    configured_region: str | None
    region_source: Literal["explicit", "environment"] | None
    configured_base_url: httpx2.URL | None
    api_key: str | None = field(default=None, repr=False)
    token_provider: BedrockTokenProvider | None = field(default=None, repr=False, compare=False)
    use_environment_bearer: bool = False
    profile: str | None = None
    access_key_id: str | None = field(default=None, repr=False)
    secret_access_key: str | None = field(default=None, repr=False)
    session_token: str | None = field(default=None, repr=False)
    credential_provider: AwsCredentialsProvider | None = field(default=None, repr=False, compare=False)
    name: str = field(default="bedrock", init=False)

    def _aws_source(self) -> Literal["static", "profile", "provider", "default"]:
        if self.access_key_id is not None:
            return "static"
        if self.profile is not None:
            return "profile"
        if self.credential_provider is not None:
            return "provider"
        return "default"

    def _resolve_aws_auth(self) -> tuple[BedrockAwsAuthConfig, BedrockAwsAuth | None]:
        if self.configured_region is not None:
            return (
                BedrockAwsAuthConfig(
                    region=self.configured_region,
                    source=self._aws_source(),
                    service="bedrock" if self.endpoint == "runtime" else "bedrock-mantle",
                    region_source=self.region_source or "explicit",
                    profile=self.profile,
                    access_key_id=self.access_key_id,
                    secret_access_key=self.secret_access_key,
                    session_token=self.session_token,
                    credentials_provider=self.credential_provider,
                ),
                None,
            )

        auth = BedrockAwsAuth.resolve(
            region=None,
            profile=self.profile,
            access_key_id=self.access_key_id,
            secret_access_key=self.secret_access_key,
            session_token=self.session_token,
            credentials_provider=self.credential_provider,
            service="bedrock" if self.endpoint == "runtime" else "bedrock-mantle",
        )
        return auth.config, auth

    def configure(self) -> _ProviderRuntime:
        def environment_token() -> str:
            token = os.environ.get("AWS_BEARER_TOKEN_BEDROCK")
            if not token:
                raise OpenAIError(
                    "Could not find credentials for Bedrock. Pass a bearer credential or AWS credentials to "
                    "`bedrock(...)`, set `AWS_BEARER_TOKEN_BEDROCK`, or configure the default AWS credential chain."
                )
            return token

        auth: _BedrockBearerAuth | _BedrockSigV4Auth | None = None
        bearer_provider: BedrockTokenProvider | None = None
        if self.api_key is not None:
            bearer_provider = lambda: self.api_key or ""
            region = self.configured_region
        elif self.token_provider is not None:
            bearer_provider = self.token_provider
            region = self.configured_region
        elif self.use_environment_bearer:
            bearer_provider = environment_token
            region = self.configured_region
        else:
            aws_config, aws_auth = self._resolve_aws_auth()
            region = aws_config.region
            _validate_bedrock_region(region)
            base_url = self.configured_base_url or _default_bedrock_base_url(self.endpoint, region)
            _validate_canonical_bedrock_endpoint(base_url, endpoint=self.endpoint, region=region)
            auth = _BedrockSigV4Auth(config=aws_config, base_url=base_url, auth=aws_auth)

        if self.configured_base_url is not None:
            base_url = self.configured_base_url
        elif region is not None:
            base_url = _default_bedrock_base_url(self.endpoint, region)
        else:
            raise OpenAIError(
                "Bedrock requires an AWS region. Pass `region` to `bedrock(...)`, or set `AWS_REGION` or "
                "`AWS_DEFAULT_REGION`."
            )

        if bearer_provider is not None:
            auth = _BedrockBearerAuth(bearer_provider, base_url=base_url)

        assert auth is not None
        if isinstance(auth, _BedrockSigV4Auth):
            return _BedrockProviderRuntime(
                name=self.name,
                base_url=base_url,
                region=region,
                transform_request=_without_redirects,
                prepare_request=auth.prepare_request,
                prepare_async_request=auth.prepare_async_request,
            )

        return _BedrockProviderRuntime(
            name=self.name,
            base_url=base_url,
            region=region,
            prepare_request=auth.prepare_request,
            prepare_async_request=auth.prepare_async_request,
        )


def bedrock(
    *,
    endpoint: BedrockEndpoint | None = None,
    region: str | None = None,
    base_url: str | httpx2.URL | None | NotGiven = NOT_GIVEN,
    api_key: str | None | NotGiven = NOT_GIVEN,
    token_provider: BedrockTokenProvider | None = None,
    access_key_id: str | None = None,
    secret_access_key: str | None = None,
    session_token: str | None = None,
    profile: str | None = None,
    credential_provider: AwsCredentialsProvider | None = None,
) -> _Provider:
    """Configure the standard OpenAI client for Amazon Bedrock Mantle or Runtime."""

    if endpoint is not None and endpoint not in {"mantle", "runtime"}:
        raise OpenAIError("The Bedrock `endpoint` must be either `mantle` or `runtime`.")

    normalized_region = _normalize_optional_string(region)
    if region is not None and normalized_region is None:
        raise OpenAIError("The Bedrock AWS `region` must not be empty.")
    _validate_bedrock_region(normalized_region)

    region_source: Literal["explicit", "environment"] | None = "explicit" if normalized_region is not None else None

    configured_base_url: httpx2.URL | None
    if isinstance(base_url, NotGiven):
        environment_base_url = _normalize_optional_string(os.environ.get("AWS_BEDROCK_BASE_URL"))
        configured_base_url = _normalize_base_url(environment_base_url) if environment_base_url else None
    elif base_url is None:
        configured_base_url = None
    else:
        if isinstance(base_url, str) and not base_url.strip():
            raise OpenAIError("The Bedrock `base_url` must not be empty.")
        configured_base_url = _normalize_base_url(base_url)

    canonical_endpoint = (
        _parse_bedrock_endpoint_hostname(configured_base_url.host) if configured_base_url is not None else None
    )
    resolved_endpoint: BedrockEndpoint = endpoint or (
        canonical_endpoint[0] if canonical_endpoint is not None else "mantle"
    )

    normalized_profile = _normalize_optional_string(profile)
    if profile is not None and normalized_profile is None:
        raise OpenAIError("The Bedrock AWS `profile` must not be empty.")

    if (access_key_id is None) != (secret_access_key is None) or (session_token is not None and access_key_id is None):
        raise OpenAIError(
            "Static AWS credentials require both `access_key_id` and `secret_access_key`. "
            "A `session_token` may only be used with both."
        )
    if access_key_id is not None and (not access_key_id.strip() or not cast(str, secret_access_key).strip()):
        raise OpenAIError("Static AWS credentials require non-empty `access_key_id` and `secret_access_key` values.")
    if session_token is not None and not session_token.strip():
        raise OpenAIError("A static AWS `session_token` must not be empty when provided.")

    explicit_api_key = not isinstance(api_key, NotGiven) and api_key is not None
    if explicit_api_key and (not isinstance(api_key, str) or not api_key.strip()):
        raise OpenAIError("The Bedrock bearer credential must not be empty.")
    if explicit_api_key and token_provider is not None:
        raise OpenAIError("The `api_key` and `token_provider` options are mutually exclusive. Configure only one.")

    explicit_bearer = explicit_api_key or token_provider is not None
    aws_modes = sum(
        (
            access_key_id is not None,
            normalized_profile is not None,
            credential_provider is not None,
        )
    )
    if aws_modes > 1:
        raise OpenAIError(
            "Bedrock authentication is ambiguous. Configure exactly one explicit AWS mode: static credentials, "
            "profile, or credential provider."
        )
    if explicit_bearer and aws_modes:
        raise OpenAIError(
            "Bedrock authentication is ambiguous. Configure exactly one explicit mode: bearer credential, "
            "static AWS credentials, profile, or credential provider."
        )

    skip_environment_bearer = not isinstance(api_key, NotGiven) and api_key is None
    use_environment_bearer = (
        not explicit_bearer
        and not aws_modes
        and not skip_environment_bearer
        and bool(os.environ.get("AWS_BEARER_TOKEN_BEDROCK"))
    )

    if normalized_region is None and (configured_base_url is None or not (explicit_bearer or use_environment_bearer)):
        normalized_region = _normalize_optional_string(
            os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
        )
        _validate_bedrock_region(normalized_region)
        if normalized_region is not None:
            region_source = "environment"

    if configured_base_url is not None:
        _validate_canonical_bedrock_endpoint(configured_base_url, endpoint=resolved_endpoint, region=normalized_region)

    return _create_provider(
        _BedrockProviderDefinition(
            endpoint=resolved_endpoint,
            configured_region=normalized_region,
            region_source=region_source,
            configured_base_url=configured_base_url,
            api_key=cast("str | None", api_key) if explicit_api_key else None,
            token_provider=token_provider,
            use_environment_bearer=use_environment_bearer,
            profile=normalized_profile,
            access_key_id=access_key_id,
            secret_access_key=secret_access_key,
            session_token=session_token,
            credential_provider=credential_provider,
        )
    )


__all__ = ["bedrock", "BedrockEndpoint", "BedrockTokenProvider", "AwsCredentialsProvider"]