# Code generated by smithy-python-codegen DO NOT EDIT.

from dataclasses import dataclass
from typing import Any, Callable, TypeAlias, Union

from smithy_aws_core.aio.protocols import RestJsonClientProtocol
from smithy_aws_core.auth import SigV4AuthScheme
from smithy_aws_core.endpoints.standard_regional import (
    StandardRegionalEndpointsResolver as _RegionalResolver,
)
from smithy_aws_core.identity import AWSCredentialsIdentity, AWSIdentityProperties
from smithy_core.aio.interfaces import (
    ClientProtocol,
    ClientTransport,
    EndpointResolver as _EndpointResolver,
)
from smithy_core.aio.interfaces.auth import AuthScheme
from smithy_core.aio.interfaces.identity import IdentityResolver
from smithy_core.interceptors import Interceptor
from smithy_core.interfaces import URI
from smithy_core.interfaces.retries import RetryStrategy
from smithy_core.retries import RetryStrategyOptions
from smithy_core.shapes import ShapeID
from smithy_http.aio.crt import AWSCRTHTTPClient
from smithy_http.interfaces import HTTPRequestConfiguration

from ._private.schemas import (
    AMAZON_BEDROCK_FRONTEND_SERVICE as _SCHEMA_AMAZON_BEDROCK_FRONTEND_SERVICE,
)
from .auth import HTTPAuthSchemeResolver
from .models import (
    ApplyGuardrailInput,
    ApplyGuardrailOutput,
    ConverseInput,
    ConverseOperationOutput,
    ConverseStreamInput,
    ConverseStreamOperationOutput,
    CountTokensOperationInput,
    CountTokensOutput,
    GetAsyncInvokeInput,
    GetAsyncInvokeOutput,
    InvokeModelInput,
    InvokeModelOutput,
    InvokeModelWithBidirectionalStreamOperationInput,
    InvokeModelWithBidirectionalStreamOperationOutput,
    InvokeModelWithResponseStreamInput,
    InvokeModelWithResponseStreamOutput,
    ListAsyncInvokesInput,
    ListAsyncInvokesOutput,
    StartAsyncInvokeInput,
    StartAsyncInvokeOutput,
)


_ServiceInterceptor = Union[
    Interceptor[ApplyGuardrailInput, ApplyGuardrailOutput, Any, Any],
    Interceptor[ConverseInput, ConverseOperationOutput, Any, Any],
    Interceptor[ConverseStreamInput, ConverseStreamOperationOutput, Any, Any],
    Interceptor[CountTokensOperationInput, CountTokensOutput, Any, Any],
    Interceptor[GetAsyncInvokeInput, GetAsyncInvokeOutput, Any, Any],
    Interceptor[InvokeModelInput, InvokeModelOutput, Any, Any],
    Interceptor[
        InvokeModelWithBidirectionalStreamOperationInput,
        InvokeModelWithBidirectionalStreamOperationOutput,
        Any,
        Any,
    ],
    Interceptor[
        InvokeModelWithResponseStreamInput,
        InvokeModelWithResponseStreamOutput,
        Any,
        Any,
    ],
    Interceptor[ListAsyncInvokesInput, ListAsyncInvokesOutput, Any, Any],
    Interceptor[StartAsyncInvokeInput, StartAsyncInvokeOutput, Any, Any],
]


@dataclass(init=False)
class Config:
    """Configuration for Bedrock Runtime."""

    auth_scheme_resolver: HTTPAuthSchemeResolver
    """
    An auth scheme resolver that determines the auth scheme for each
    operation.
    """

    auth_schemes: dict[ShapeID, AuthScheme[Any, Any, Any, Any]]
    """A map of auth scheme ids to auth schemes."""

    aws_access_key_id: str | None
    """The identifier for a secret access key."""

    aws_credentials_identity_resolver: (
        IdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties] | None
    )
    """Resolves AWS Credentials. Required for operations that use Sigv4 Auth."""

    aws_secret_access_key: str | None
    """A secret access key that can be used to sign requests."""

    aws_session_token: str | None
    """An access key ID that identifies temporary security credentials."""

    endpoint_resolver: _EndpointResolver
    """
    The endpoint resolver used to resolve the final endpoint per-operation
    based on the configuration.
    """

    endpoint_uri: str | URI | None
    """A static URI to route requests to."""

    http_request_config: HTTPRequestConfiguration | None
    """Configuration for individual HTTP requests."""

    interceptors: list[_ServiceInterceptor]
    """
    The list of interceptors, which are hooks that are called during the
    execution of a request.
    """

    protocol: ClientProtocol[Any, Any] | None
    """The protocol to serialize and deserialize requests with."""

    region: str | None
    """
    The AWS region to connect to. The configured region is used to determine
    the service endpoint.
    """

    retry_strategy: RetryStrategy | RetryStrategyOptions | None
    """
    The retry strategy or options for configuring retry behavior. Can be
    either a configured RetryStrategy or RetryStrategyOptions to create one.
    """

    sdk_ua_app_id: str | None
    """
    A unique and opaque application ID that is appended to the User-Agent
    header.
    """

    transport: ClientTransport[Any, Any] | None
    """The transport to use to send requests (e.g. an HTTP client)."""

    user_agent_extra: str | None
    """Additional suffix to be added to the User-Agent header."""

    def __init__(
        self,
        *,
        auth_scheme_resolver: HTTPAuthSchemeResolver | None = None,
        auth_schemes: dict[ShapeID, AuthScheme[Any, Any, Any, Any]] | None = None,
        aws_access_key_id: str | None = None,
        aws_credentials_identity_resolver: IdentityResolver[
            AWSCredentialsIdentity, AWSIdentityProperties
        ]
        | None = None,
        aws_secret_access_key: str | None = None,
        aws_session_token: str | None = None,
        endpoint_resolver: _EndpointResolver | None = None,
        endpoint_uri: str | URI | None = None,
        http_request_config: HTTPRequestConfiguration | None = None,
        interceptors: list[_ServiceInterceptor] | None = None,
        protocol: ClientProtocol[Any, Any] | None = None,
        region: str | None = None,
        retry_strategy: RetryStrategy | RetryStrategyOptions | None = None,
        sdk_ua_app_id: str | None = None,
        transport: ClientTransport[Any, Any] | None = None,
        user_agent_extra: str | None = None,
    ):
        self.auth_scheme_resolver = auth_scheme_resolver or HTTPAuthSchemeResolver()
        self.auth_schemes = auth_schemes or {
            ShapeID("aws.auth#sigv4"): SigV4AuthScheme(service="bedrock")
        }
        self.aws_access_key_id = aws_access_key_id
        self.aws_credentials_identity_resolver = aws_credentials_identity_resolver
        self.aws_secret_access_key = aws_secret_access_key
        self.aws_session_token = aws_session_token
        self.endpoint_resolver = endpoint_resolver or _RegionalResolver(
            endpoint_prefix="bedrock-runtime"
        )
        self.endpoint_uri = endpoint_uri
        self.http_request_config = http_request_config
        self.interceptors = interceptors or []
        self.protocol = protocol or RestJsonClientProtocol(
            _SCHEMA_AMAZON_BEDROCK_FRONTEND_SERVICE
        )
        self.region = region
        self.retry_strategy = retry_strategy
        self.sdk_ua_app_id = sdk_ua_app_id
        self.transport = transport or AWSCRTHTTPClient()
        self.user_agent_extra = user_agent_extra

    def set_auth_scheme(self, scheme: AuthScheme[Any, Any, Any, Any]) -> None:
        """
        Sets the implementation of an auth scheme.

        Using this method ensures the correct key is used.

        Args:
            scheme:
                The auth scheme to add.
        """
        self.auth_schemes[scheme.scheme_id] = scheme


Plugin: TypeAlias = Callable[[Config], None]
"""A callable that allows customizing the config object on each request."""