# Code generated by smithy-python-codegen DO NOT EDIT. from copy import deepcopy import logging from smithy_core.aio.client import ClientCall, RequestPipeline from smithy_core.aio.eventstream import DuplexEventStream, OutputEventStream from smithy_core.exceptions import ExpectationNotMetError from smithy_core.interceptors import InterceptorChain from smithy_core.interfaces.retries import RetryStrategy from smithy_core.retries import RetryStrategyOptions, RetryStrategyResolver from smithy_core.types import TypedProperties from smithy_http.plugins import user_agent_plugin from .config import Config, Plugin from .models import ( APPLY_GUARDRAIL, ApplyGuardrailInput, ApplyGuardrailOutput, CONVERSE, CONVERSE_STREAM, COUNT_TOKENS, ConverseInput, ConverseOperationOutput, ConverseStreamInput, ConverseStreamOperationOutput, ConverseStreamOutput, CountTokensOperationInput, CountTokensOutput, GET_ASYNC_INVOKE, GetAsyncInvokeInput, GetAsyncInvokeOutput, INVOKE_MODEL, INVOKE_MODEL_WITH_BIDIRECTIONAL_STREAM, INVOKE_MODEL_WITH_RESPONSE_STREAM, InvokeModelInput, InvokeModelOutput, InvokeModelWithBidirectionalStreamInput, InvokeModelWithBidirectionalStreamOperationInput, InvokeModelWithBidirectionalStreamOperationOutput, InvokeModelWithBidirectionalStreamOutput, InvokeModelWithResponseStreamInput, InvokeModelWithResponseStreamOutput, LIST_ASYNC_INVOKES, ListAsyncInvokesInput, ListAsyncInvokesOutput, ResponseStream, START_ASYNC_INVOKE, StartAsyncInvokeInput, StartAsyncInvokeOutput, _ConverseStreamOutputDeserializer, _InvokeModelWithBidirectionalStreamOutputDeserializer, _ResponseStreamDeserializer, ) from .user_agent import aws_user_agent_plugin logger = logging.getLogger(__name__) class BedrockRuntimeClient: """ Describes the API operations for running inference using Amazon Bedrock models. :param config: Optional configuration for the client. Here you can set things like the endpoint for HTTP services or auth credentials. :param plugins: A list of callables that modify the configuration dynamically. These can be used to set defaults, for example. """ def __init__( self, config: Config | None = None, plugins: list[Plugin] | None = None ): self._config = config or Config() client_plugins: list[Plugin] = [aws_user_agent_plugin, user_agent_plugin] if plugins: client_plugins.extend(plugins) for plugin in client_plugins: plugin(self._config) self._retry_strategy_resolver = RetryStrategyResolver() async def _resolve_retry_strategy( self, retry_strategy: RetryStrategy | RetryStrategyOptions | None ) -> RetryStrategy: if isinstance(retry_strategy, RetryStrategy): return retry_strategy elif isinstance(retry_strategy, RetryStrategyOptions): return await self._retry_strategy_resolver.resolve_retry_strategy( options=retry_strategy ) elif retry_strategy is None: return await self._retry_strategy_resolver.resolve_retry_strategy( options=RetryStrategyOptions() ) else: raise TypeError( f"retry_strategy must be RetryStrategy, RetryStrategyOptions, or None, " f"got {type(retry_strategy).__name__}" ) async def apply_guardrail( self, input: ApplyGuardrailInput, plugins: list[Plugin] | None = None ) -> ApplyGuardrailOutput: """ The action to apply a guardrail. For troubleshooting some of the common errors you might encounter when using the ``ApplyGuardrail`` API, see `Troubleshooting Amazon Bedrock API Error Codes `_ in the Amazon Bedrock User Guide :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=APPLY_GUARDRAIL, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline(call) async def converse( self, input: ConverseInput, plugins: list[Plugin] | None = None ) -> ConverseOperationOutput: """ Sends messages to the specified Amazon Bedrock model. ``Converse`` provides a consistent interface that works with all models that support messages. This allows you to write code once and use it with different models. If a model has unique inference parameters, you can also pass those unique parameters to the model. Amazon Bedrock doesn't store any text, images, or documents that you provide as content. The data is only used to generate the response. You can submit a prompt by including it in the ``messages`` field, specifying the ``modelId`` of a foundation model or inference profile to run inference on it, and including any other fields that are relevant to your use case. You can also submit a prompt from Prompt management by specifying the ARN of the prompt version and including a map of variables to values in the ``promptVariables`` field. You can append more messages to the prompt by using the ``messages`` field. If you use a prompt from Prompt management, you can't include the following fields in the request: ``additionalModelRequestFields``, ``inferenceConfig``, ``system``, or ``toolConfig``. Instead, these fields must be defined through Prompt management. For more information, see `Use a prompt from Prompt management `_ . For information about the Converse API, see *Use the Converse API* in the *Amazon Bedrock User Guide*. To use a guardrail, see *Use a guardrail with the Converse API* in the *Amazon Bedrock User Guide*. To use a tool with a model, see *Tool use (Function calling)* in the *Amazon Bedrock User Guide* For example code, see *Converse API examples* in the *Amazon Bedrock User Guide*. This operation requires permission for the ``bedrock:InvokeModel`` action. .. important:: To deny all inference access to resources that you specify in the modelId field, you need to deny access to the ``bedrock:InvokeModel`` and ``bedrock:InvokeModelWithResponseStream`` actions. Doing this also denies access to the resource through the base inference actions (`InvokeModel `_ and `InvokeModelWithResponseStream `_). For more information see `Deny access for inference on specific models `_ . For troubleshooting some of the common errors you might encounter when using the ``Converse`` API, see `Troubleshooting Amazon Bedrock API Error Codes `_ in the Amazon Bedrock User Guide :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=CONVERSE, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline(call) async def converse_stream( self, input: ConverseStreamInput, plugins: list[Plugin] | None = None ) -> OutputEventStream[ConverseStreamOutput, ConverseStreamOperationOutput]: """ Sends messages to the specified Amazon Bedrock model and returns the response in a stream. ``ConverseStream`` provides a consistent API that works with all Amazon Bedrock models that support messages. This allows you to write code once and use it with different models. Should a model have unique inference parameters, you can also pass those unique parameters to the model. To find out if a model supports streaming, call `GetFoundationModel `_ and check the ``responseStreamingSupported`` field in the response. .. note:: The CLI doesn't support streaming operations in Amazon Bedrock, including ``ConverseStream``. Amazon Bedrock doesn't store any text, images, or documents that you provide as content. The data is only used to generate the response. You can submit a prompt by including it in the ``messages`` field, specifying the ``modelId`` of a foundation model or inference profile to run inference on it, and including any other fields that are relevant to your use case. You can also submit a prompt from Prompt management by specifying the ARN of the prompt version and including a map of variables to values in the ``promptVariables`` field. You can append more messages to the prompt by using the ``messages`` field. If you use a prompt from Prompt management, you can't include the following fields in the request: ``additionalModelRequestFields``, ``inferenceConfig``, ``system``, or ``toolConfig``. Instead, these fields must be defined through Prompt management. For more information, see `Use a prompt from Prompt management `_ . For information about the Converse API, see *Use the Converse API* in the *Amazon Bedrock User Guide*. To use a guardrail, see *Use a guardrail with the Converse API* in the *Amazon Bedrock User Guide*. To use a tool with a model, see *Tool use (Function calling)* in the *Amazon Bedrock User Guide* For example code, see *Conversation streaming example* in the *Amazon Bedrock User Guide*. This operation requires permission for the ``bedrock:InvokeModelWithResponseStream`` action. .. important:: To deny all inference access to resources that you specify in the modelId field, you need to deny access to the ``bedrock:InvokeModel`` and ``bedrock:InvokeModelWithResponseStream`` actions. Doing this also denies access to the resource through the base inference actions (`InvokeModel `_ and `InvokeModelWithResponseStream `_). For more information see `Deny access for inference on specific models `_ . For troubleshooting some of the common errors you might encounter when using the ``ConverseStream`` API, see `Troubleshooting Amazon Bedrock API Error Codes `_ in the Amazon Bedrock User Guide :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=CONVERSE_STREAM, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline.output_stream( call, ConverseStreamOutput, _ConverseStreamOutputDeserializer().deserialize ) async def count_tokens( self, input: CountTokensOperationInput, plugins: list[Plugin] | None = None ) -> CountTokensOutput: """ Returns the token count for a given inference request. This operation helps you estimate token usage before sending requests to foundation models by returning the token count that would be used if the same input were sent to the model in an inference request. Token counting is model-specific because different models use different tokenization strategies. The token count returned by this operation will match the token count that would be charged if the same input were sent to the model in an ``InvokeModel`` or ``Converse`` request. You can use this operation to: * Estimate costs before sending inference requests. * Optimize prompts to fit within token limits. * Plan for token usage in your applications. This operation accepts the same input formats as ``InvokeModel`` and ``Converse``, allowing you to count tokens for both raw text inputs and structured conversation formats. The following operations are related to ``CountTokens``: * `InvokeModel `_ - Sends inference requests to foundation models * `Converse `_ - Sends conversation-based inference requests to foundation models :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=COUNT_TOKENS, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline(call) async def get_async_invoke( self, input: GetAsyncInvokeInput, plugins: list[Plugin] | None = None ) -> GetAsyncInvokeOutput: """ Retrieve information about an asynchronous invocation. :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=GET_ASYNC_INVOKE, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline(call) async def invoke_model( self, input: InvokeModelInput, plugins: list[Plugin] | None = None ) -> InvokeModelOutput: """ Invokes the specified Amazon Bedrock model to run inference using the prompt and inference parameters provided in the request body. You use model inference to generate text, images, and embeddings. For example code, see *Invoke model code examples* in the *Amazon Bedrock User Guide*. This operation requires permission for the ``bedrock:InvokeModel`` action. .. important:: To deny all inference access to resources that you specify in the modelId field, you need to deny access to the ``bedrock:InvokeModel`` and ``bedrock:InvokeModelWithResponseStream`` actions. Doing this also denies access to the resource through the Converse API actions (`Converse `_ and `ConverseStream `_). For more information see `Deny access for inference on specific models `_ . For troubleshooting some of the common errors you might encounter when using the ``InvokeModel`` API, see `Troubleshooting Amazon Bedrock API Error Codes `_ in the Amazon Bedrock User Guide :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=INVOKE_MODEL, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline(call) async def invoke_model_with_bidirectional_stream( self, input: InvokeModelWithBidirectionalStreamOperationInput, plugins: list[Plugin] | None = None, ) -> DuplexEventStream[ InvokeModelWithBidirectionalStreamInput, InvokeModelWithBidirectionalStreamOutput, InvokeModelWithBidirectionalStreamOperationOutput, ]: """ Invoke the specified Amazon Bedrock model to run inference using the bidirectional stream. The response is returned in a stream that remains open for 8 minutes. A single session can contain multiple prompts and responses from the model. The prompts to the model are provided as audio files and the model's responses are spoken back to the user and transcribed. It is possible for users to interrupt the model's response with a new prompt, which will halt the response speech. The model will retain contextual awareness of the conversation while pivoting to respond to the new prompt. :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=INVOKE_MODEL_WITH_BIDIRECTIONAL_STREAM, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline.duplex_stream( call, InvokeModelWithBidirectionalStreamInput, InvokeModelWithBidirectionalStreamOutput, _InvokeModelWithBidirectionalStreamOutputDeserializer().deserialize, ) async def invoke_model_with_response_stream( self, input: InvokeModelWithResponseStreamInput, plugins: list[Plugin] | None = None, ) -> OutputEventStream[ResponseStream, InvokeModelWithResponseStreamOutput]: """ Invoke the specified Amazon Bedrock model to run inference using the prompt and inference parameters provided in the request body. The response is returned in a stream. To see if a model supports streaming, call `GetFoundationModel `_ and check the ``responseStreamingSupported`` field in the response. .. note:: The CLI doesn't support streaming operations in Amazon Bedrock, including ``InvokeModelWithResponseStream``. For example code, see *Invoke model with streaming code example* in the *Amazon Bedrock User Guide*. This operation requires permissions to perform the ``bedrock:InvokeModelWithResponseStream`` action. .. important:: To deny all inference access to resources that you specify in the modelId field, you need to deny access to the ``bedrock:InvokeModel`` and ``bedrock:InvokeModelWithResponseStream`` actions. Doing this also denies access to the resource through the Converse API actions (`Converse `_ and `ConverseStream `_). For more information see `Deny access for inference on specific models `_ . For troubleshooting some of the common errors you might encounter when using the ``InvokeModelWithResponseStream`` API, see `Troubleshooting Amazon Bedrock API Error Codes `_ in the Amazon Bedrock User Guide :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=INVOKE_MODEL_WITH_RESPONSE_STREAM, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline.output_stream( call, ResponseStream, _ResponseStreamDeserializer().deserialize ) async def list_async_invokes( self, input: ListAsyncInvokesInput, plugins: list[Plugin] | None = None ) -> ListAsyncInvokesOutput: """ Lists asynchronous invocations. :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=LIST_ASYNC_INVOKES, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline(call) async def start_async_invoke( self, input: StartAsyncInvokeInput, plugins: list[Plugin] | None = None ) -> StartAsyncInvokeOutput: """ Starts an asynchronous invocation. This operation requires permission for the ``bedrock:InvokeModel`` action. .. important:: To deny all inference access to resources that you specify in the modelId field, you need to deny access to the ``bedrock:InvokeModel`` and ``bedrock:InvokeModelWithResponseStream`` actions. Doing this also denies access to the resource through the Converse API actions (`Converse `_ and `ConverseStream `_). For more information see `Deny access for inference on specific models `_ . :param input: The operation's input. :param plugins: A list of callables that modify the configuration dynamically. Changes made by these plugins only apply for the duration of the operation execution and will not affect any other operation invocations. """ operation_plugins: list[Plugin] = [] if plugins: operation_plugins.extend(plugins) config = deepcopy(self._config) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: raise ExpectationNotMetError( "protocol and transport MUST be set on the config to make calls." ) retry_strategy = await self._resolve_retry_strategy(config.retry_strategy) pipeline = RequestPipeline(protocol=config.protocol, transport=config.transport) call = ClientCall( input=input, operation=START_ASYNC_INVOKE, context=TypedProperties({"config": config}), interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, endpoint_resolver=config.endpoint_resolver, retry_strategy=retry_strategy, ) return await pipeline(call)