DEV Community

Tsubasa Kanno
Tsubasa Kanno

Posted on

Calling Amazon Bedrock directly from Streamlit in Snowflake (SiS)

Introduction

SiS has finally added support for external network access! This allows SiS to directly access domains permitted by network rules.

In this post, we'll create a web app using SiS that can use both Cortex LLM and Amazon Bedrock by directly calling the Amazon Bedrock API.

For any unclear points in this procedure, please refer to the following Snowflake documentation:

External Network Access in Streamlit

Note: This post represents my personal views and not those of Snowflake.

Feature Overview

Goal

  • Call Amazon Bedrock directly from SiS without using UDF / UDTF

Actual Screenshot (Excerpt)

Prerequisites

  • Snowflake
    • Snowflake account with access to Cortex LLM (If you only want to use Amazon Bedrock from SiS, a Snowflake account without Cortex LLM access is also fine)
    • snowflake-ml-python 1.1.2 or later
    • boto3 1.28.64
  • AWS
    • AWS account with access to Amazon Bedrock (We're using the us-east-1 region in this procedure to use Claude 3.5 Sonnet)

Procedure

Enable Target Model in Amazon Bedrock

Connect to the Amazon Bedrock console from the AWS Management Console.

Click "Model access" in the left pane and enable the model you want to use. In this case, we're enabling Claude 3.5 Sonnet.

Wait for approval from AWS. It was immediate for me, but it might take some time depending on your account.

Create an AWS IAM User

Create an IAM user for SiS access.

Connect to the IAM console from the AWS Management Console.

Click "Users" in the left pane, then click "Create user" on the next screen.

Enter a user name (we'll use sis_bedrock_user here), uncheck "Provide user access to the AWS Management Console", and click "Next".

Check "Attach policies directly", check the AmazonBedrockFullAccess managed policy, and click "Next".
Note: AmazonBedrockFullAccess gives full permissions for Amazon Bedrock. Consider assigning a more restricted IAM policy as needed.

Click "Create user" to create an IAM user with Amazon Bedrock access permissions.

Search for the IAM user you just created on the IAM users screen and click on the user name.

Click the "Security credentials" tab, then click "Create access key".

Check "Third-party service", check "I understand the above recommendation and want to proceed to create an access key", and click "Next".

On the "Set description tag - optional" screen, click "Create access key" without entering anything.

An "Access key" and "Secret access key" will be generated. Store these securely. (You can also download them as a CSV file)
Note: If you don't save the "Access key" and "Secret access key" here, they cannot be reissued. If lost, create a new access key.

This completes the AWS setup.

Create a Network Rule in Snowflake

External access from Snowflake is closed by default, so we'll create a network rule to access Amazon Bedrock. Don't worry, just creating a network rule doesn't enable external access.

Open a worksheet from Snowflake's Snowsight and execute the following SQL:

-- Create a network rule for external access
CREATE OR REPLACE NETWORK RULE bedrock_network_rule
mode = EGRESS
type = HOST_PORT
value_list = ('bedrock-runtime.us-east-1.amazonaws.com','bedrock-runtime-fips.us-east-1.amazonaws.com');
Enter fullscreen mode Exit fullscreen mode

Create a Secret

Continue in the worksheet and execute the following SQL:

-- Create a secret
CREATE OR REPLACE SECRET bedrock_key
  TYPE = PASSWORD
  USERNAME = '<Access Key>'
  PASSWORD = '<Secret Access Key>';
Enter fullscreen mode Exit fullscreen mode

Note: The Access Key and Secret Access Key are the values you noted when creating the AWS IAM user earlier.

Create an External Access Integration

Continue in the worksheet and execute the following SQL:

-- Create an external access integration
CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION bedrock_access_integration
ALLOWED_NETWORK_RULES = (bedrock_network_rule)
ALLOWED_AUTHENTICATION_SECRETS = (bedrock_key)
ENABLED = TRUE;
Enter fullscreen mode Exit fullscreen mode

Create a SiS App

Click "Streamlit" in the left pane of Snowsight, click the "+ Streamlit" button to create a SiS app.

Once the SiS app starts, check the URL displayed in the browser. The end of the URL will be:

https://(omitted)/streamlit-apps/<DB name>.<schema name>.<Streamlit object name>
Enter fullscreen mode Exit fullscreen mode

Note down the Streamlit object name.
Note: This is different from the SiS app title.
Note: You can also check with the SHOW STREAMLITS command.

Link the External Access Integration to the Streamlit Object

Return to the worksheet and execute the following SQL:

-- Map Streamlit object and external access integration
ALTER STREAMLIT <DB name>.<schema name>.<Streamlit object name>
  SET EXTERNAL_ACCESS_INTEGRATIONS = (bedrock_access_integration)
  SECRETS = ('bedrock_key' = <DB name>.<schema name>.bedrock_key);
Enter fullscreen mode Exit fullscreen mode

Run Amazon Bedrock from the SiS App

Click "Streamlit" in the left pane of Snowsight and open the SiS app you created earlier.
In the edit screen, select and install the following from "Packages":

  • snowflake-ml-python 1.5.3
  • boto3 1.28.64

Paste and run the following source code:

from snowflake.snowpark.context import get_active_session
import streamlit as st
from snowflake.cortex import Complete as CompleteText
import json
import boto3
import _snowflake

# Get Snowflake session
snowflake_session = get_active_session()

# Application title
st.title("Cortex & Bedrock Chat App")

# Function to get AWS credentials
def get_aws_credentials():
    aws_key_object = _snowflake.get_username_password('bedrock_key')
    region = 'us-east-1'

    boto3_session_args = {
        'aws_access_key_id': aws_key_object.username,
        'aws_secret_access_key': aws_key_object.password,
        'region_name': region
    }

    return boto3_session_args, region

# Configure Bedrock client
boto3_session_args, region = get_aws_credentials()
boto3_session = boto3.Session(**boto3_session_args)
bedrock = boto3_session.client('bedrock-runtime', region_name=region)

# AI settings
st.sidebar.title("LLM Settings")
lang_model = st.sidebar.radio("Select the language model you want to use", 
                              ("Cortex AI LLM", "Bedrock Claude3.5 Sonnet"))

if lang_model == "Cortex AI LLM":
    lang_model = st.sidebar.radio("Select Cortex AI model",
                                  ("snowflake-arctic", "reka-flash", "reka-core", 
                                   "mistral-large2", "mistral-large", "mixtral-8x7b", "mistral-7b", 
                                   "llama3.1-405b", "llama3.1-70b", "llama3.1-8b", 
                                   "llama3-70b", "llama3-8b", "llama2-70b-chat", 
                                   "jamba-instruct", "gemma-7b")
    )
else:
    lang_model = "anthropic.claude-3-5-sonnet-20240620-v1:0"

# Function to call Bedrock API
def ask_claude3(prompt):
    body = json.dumps({
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 100000,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt
                    }
                ]
            }
        ]
    })

    response = bedrock.invoke_model(
        body=body,
        modelId=lang_model,
        accept='application/json',
        contentType='application/json'
    )

    response_body = json.loads(response.get('body').read())
    return response_body.get('content')[0].get('text')

st.header("Let's chat freely!")

# Initialize chat history
if 'messages' not in st.session_state:
    st.session_state.messages = []

# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# User input
if prompt := st.chat_input("Enter your message."):
    # Add and display user message
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generate AI response
    if lang_model.startswith("anthropic.claude-3"):
        response = ask_claude3(prompt)
    else:
        response = CompleteText(lang_model, prompt)

    # Add and display AI response
    st.session_state.messages.append({"role": "assistant", "content": response})
    with st.chat_message("assistant"):
        st.markdown(response)
Enter fullscreen mode Exit fullscreen mode

Conclusion

How was it? Being able to access externally from SiS without UDF / UDTF allows us to create more advanced web applications more easily! Also, perhaps because we're not using UDF / UDTF, the response seems to have improved, enabling an even better user experience.

Please try creating wonderful web apps using SiS's external network access feature!

Announcements

Snowflake What's New Updates on X

I'm sharing Snowflake's What's New updates on X. Please feel free to follow if you're interested!

English Version

Snowflake What's New Bot (English Version)
https://x.com/snow_new_en

Japanese Version

Snowflake What's New Bot (Japanese Version)
https://x.com/snow_new_jp

Change History

(20240914) Initial post

Original Japanese Article

https://zenn.dev/tsubasa_tech/articles/ea53b5e37705cb

Top comments (0)