AWS Bedrock

Amazon Bedrock has an OpenAI-compatible API, so a Flyte task or app can call it with the OpenAI SDK by changing the base URL.

Setting	Value
Base URL	`https://bedrock-runtime.<region>.amazonaws.com/openai/v1`
API key	Bedrock console, API keys
Environment variable	`AWS_BEARER_TOKEN_BEDROCK`

Store your Bedrock API key

Run this command and paste the key at the prompt:

        
    
flyte create secret amazon-bedrock-api-key --project <your-project> --domain development

Call Bedrock from a task

call_amazon_bedrock.py
                
                    
                
            
                
            

                
                    
                
            
import os

import flyte

env = flyte.TaskEnvironment(
    name="amazon-bedrock-inference",
    image=flyte.Image.from_debian_base().with_pip_packages("openai"),
    secrets=flyte.Secret(key="amazon-bedrock-api-key", as_env_var="AWS_BEARER_TOKEN_BEDROCK"),
)

@env.task
async def summarize(text: str) -> str:
    from openai import AsyncOpenAI

    client = AsyncOpenAI(
        base_url="https://bedrock-runtime.<region>.amazonaws.com/openai/v1",
        api_key=os.environ["AWS_BEARER_TOKEN_BEDROCK"],
    )
    response = await client.chat.completions.create(
        model="<model-id>",
        messages=[{"role": "user", "content": f"Summarize in two sentences:\n\n{text}"}],
    )
    return response.choices[0].message.content

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(summarize, text="First document...")
    print(run.url)

Replace <model-id> with a model from the Bedrock docs.

Call Bedrock from an app

An AppEnvironment takes the same secrets argument, and the app reads the same environment variable. See Call the provider from an app.

Use a Bedrock API key

Pass a Bedrock API key to the OpenAI SDK as api_key. AWS recommends short-term keys for production, which you’d need to refresh in the secret before they expire.