IAM Role Auth with Snowpark External Access for Generative AI with Snowflake and AWS

Key Contributors: James Sun — Sr. Partner Engineer @ Snowflake and Bhanu Prakash — Sr. Product Manager @ Snowflake

Being at Snowflake during the Generative AI boom has been exhilarating. Keeping up with the latest models and technologies in the market while diving deep into Cortex features has kept me on my toes! While many announcements in the tech space are both flashy and useful, some of the most valuable features are those that quietly make our lives easier, more secure, and more enterprise-grade.

Snowpark External Access in Snowflake is widely adopted by our customers, offering several authentication methods to access external services or APIs. But now, joint Snowflake and AWS customers can take advantage of a new feature: using AWS IAM Roles to authenticate with AWS services, such as Bedrock, directly from Snowflake. This method is quick, secure, and aligns with common AWS practices.

This private preview functionality is available by contacting your Snowflake account team to request access as it progresses toward general availability. Additionally, the Snowflake quickstart for utilizing Bedrock from Streamlit has been updated at step 5 to include instructions for using IAM Roles for authorization to make secure calls to Bedrock. The steps are also outlined below.

Using Snowpark External Access with Bedrock via IAM Role Authorization

use role <ROLE>;
use database <DATABASE>;

-- create network rule
CREATE OR REPLACE NETWORK RULE bedrock_network_rule
MODE = EGRESS
TYPE = HOST_PORT
VALUE_LIST = ('places.geo.<AWS Region - e.g. us-west-2>.amazonaws.com','bedrock-runtime.<AWS Region - e.g. us-west-2>.amazonaws.com');

-- create security integration with aws role arn
CREATE OR REPLACE SECURITY INTEGRATION bedrock_security_integration
TYPE = API_AUTHENTICATION
AUTH_TYPE = AWS_IAM
ENABLED = TRUE
AWS_ROLE_ARN = '<aws role arn>';

-- view integration information for trust relationship
DESC SECURITY INTEGRATION bedrock_security_integration;

-- create secret
CREATE OR REPLACE SECRET aws_bedrock_access_token
TYPE = CLOUD_PROVIDER_TOKEN
API_AUTHENTICATION = bedrock_security_integration;

-- create external access object
CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION bedrock_external_access_integration
ALLOWED_NETWORK_RULES = (bedrock_network_rule)
ALLOWED_AUTHENTICATION_SECRETS=(aws_bedrock_access_token)
ENABLED=true ;

-- create function
CREATE OR REPLACE FUNCTION ask_bedrock(
instructions VARCHAR,
user_context VARCHAR,
model_id VARCHAR
)
RETURNS VARCHAR
LANGUAGE PYTHON
EXTERNAL_ACCESS_INTEGRATIONS = (bedrock_external_access_integration)
RUNTIME_VERSION = '3.11'
SECRETS = ('cred' = aws_bedrock_access_token)
PACKAGES = ('boto3')
HANDLER = 'bedrock_py'
AS
$$
import boto3
import json
import _snowflake
def bedrock_py(instructions, user_context, model_id):
# Get the previously created token as an object
cloud_provider_object = _snowflake.get_cloud_provider_token('cred')
cloud_provider_dictionary = {
"ACCESS_KEY_ID": cloud_provider_object.access_key_id,
"SECRET_ACCESS_KEY": cloud_provider_object.secret_access_key,
"TOKEN": cloud_provider_object.token
}
# Assign AWS credentials and choose a region
boto3_session_args = {
'aws_access_key_id': cloud_provider_dictionary["ACCESS_KEY_ID"],
'aws_secret_access_key': cloud_provider_dictionary["SECRET_ACCESS_KEY"],
'aws_session_token': cloud_provider_dictionary["TOKEN"],
'region_name': 'us-west-2'
}
session = boto3.Session(**boto3_session_args)
client = session.client('bedrock-runtime')
# Prepare the request body for the specified model
def prepare_request_body(model_id, instructions, user_context):
default_max_tokens = 512
default_temperature = 0.7
default_top_p = 1.0
if model_id == 'amazon.titan-text-express-v1':
body = {
"inputText": f"<SYSTEM>Follow these:{instructions}<END_SYSTEM>\n<USER_CONTEXT>Use this user context in your response:{user_context}<END_USER_CONTEXT>",
"textGenerationConfig": {
"maxTokenCount": default_max_tokens,
"stopSequences": [],
"temperature": default_temperature,
"topP": default_top_p
}
}
elif model_id == 'ai21.j2-ultra-v1':
body = {
"prompt": f"<SYSTEM>Follow these:{instructions}<END_SYSTEM>\n<USER_CONTEXT>Use this user context in your response:{user_context}<END_USER_CONTEXT>",
"temperature": default_temperature,
"topP": default_top_p,
"maxTokens": default_max_tokens
}
elif model_id == 'anthropic.claude-3-sonnet-20240229-v1:0':
body = {
"max_tokens": default_max_tokens,
"messages": [{"role": "user", "content": f"<SYSTEM>Follow these:{instructions}<END_SYSTEM>\n<USER_CONTEXT>Use this user context in your response:{user_context}<END_USER_CONTEXT>"}],
"anthropic_version": "bedrock-2023–05–31"
}
else:
raise ValueError("Unsupported model ID")
return json.dumps(body)
# Call Bedrock to get a completion
body = prepare_request_body(model_id, instructions, user_context)
response = client.invoke_model(modelId=model_id, body=body)
response_body = json.loads(response.get('body').read())
# Parse the API response based on the model
def get_completion_from_response(response_body, model_id):
if model_id == 'amazon.titan-text-express-v1':
output_text = response_body.get('results')[0].get('outputText')
elif model_id == 'ai21.j2-ultra-v1':
output_text = response_body.get('completions')[0].get('data').get('text')
elif model_id == 'anthropic.claude-3-sonnet-20240229-v1:0':
output_text = response_body.get('content')[0].get('text')
else:
raise ValueError("Unsupported model ID")
return output_text
# Get the generated text from Bedrock
output_text = get_completion_from_response(response_body, model_id)
return output_text
$$;

-- test function
SET DEFAULT_LLM_INSTRUCTIONS = 'Review the customer\'s most frequent retail purchases from last year. Write a personalized email explaining their shopper profile based on these habits. Add a tailored message suggesting products and brands for them to consider, from their purchase history.';

SET DEFAULT_MODEL = 'anthropic.claude-3-sonnet-20240229-v1:0';
select ask_bedrock($DEFAULT_LLM_INSTRUCTIONS, 'Home Decor, Furniture, Lighting', $DEFAULT_MODEL);

While Snowflake sees significant usage of Snowpark External Access with Bedrock, other AWS services, such as location services, can also be utilized effectively.

Using Snowpark External Access with Location Services via IAM Role Authorization

The steps for using the Location Service are explained below.

1.Use the same IAM Role created in previous example AWS for Bedrock, add “geo:*” permission to it.
2. Create a place index in AWS location service.
https://docs.aws.amazon.com/location/latest/developerguide/places-prerequisites.html#create-place-index-resource
3. Insert the place index in the line of code where you see ‘<place index name>’.
4. Run demo.sql to ask Bedrock what are the top 5-star Michelin restaurants in Las Vegas and convert their addresses to geo coordinates.

-- Begin geo_udf.sql file
CREATE OR REPLACE FUNCTION get_coordinate(ADDRESS STRING)
RETURNS VARIANT
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
HANDLER = 'get_location'
EXTERNAL_ACCESS_INTEGRATIONS = (aws_geo_ea_integration)
PACKAGES = ('requests','boto3')
SECRETS = ('cred' = aws_geo_access_token)
AS
$$
import _snowflake
from botocore.config import Config
import requests, json, boto3

region_name = '<AWS Region - e.g. us-west-2>'

def get_location(ADDRESS):
    # Get token object
    cloud_provider_object = _snowflake.get_cloud_provider_token('cred')

    # Boto3 configuration
    config = Config(
      retries=dict(total_max_attempts=9),
      connect_timeout=30,
      read_timeout=30,
      max_pool_connections=50
    )

    # Connect to location service using boto3
    location_client = boto3.client(
      'location',
      region_name=region_name,
      aws_access_key_id=cloud_provider_object.access_key_id,
      aws_secret_access_key=cloud_provider_object.secret_access_key,
      aws_session_token=cloud_provider_object.token,
      config=config
    )

    # Perform geocoding request
    r = location_client.search_place_index_for_text(
        #Create a place index:
        #https://docs.aws.amazon.com/location/latest/developerguide/places-prerequisites.html#create-place-index-resource
        IndexName='<place index name>',
        Text=ADDRESS
        )

    result={}
    # Extract latitude and longitude
    try:
        lon = r["Results"][0]["Place"]["Geometry"]["Point"][0]
        lat = r["Results"][0]["Place"]["Geometry"]["Point"][1]

        result['latitude'] = lat
        result['longitude'] = lon

    except (KeyError, IndexError) as e:
    # Error handleing
        result = {}
        
    return result

$$;

-- Begin demo.sql file
SET DEFAULT_MODEL = 'anthropic.claude-3-sonnet-20240229-v1:0';

SET DEFAULT_LLM_INSTRUCTIONS = 'You are a helpful agent only respond with json format, not anything else';

SET PROMPT = 'Give me a few 5-star michelin restaurants in Las Vegas using json format as {"places":[{"place_name","address","city_name","state_name" }]}';

with pjson as (
  select parse_json(ask_bedrock($DEFAULT_LLM_INSTRUCTIONS, $PROMPT, $DEFAULT_MODEL)) v
),
f as (
select  p.VALUE:"address"::text as addr,
        p.VALUE:"city_name"::text as city,
        p.VALUE:"place_name"::text as place,
        p.VALUE:"state_name"::text as state,
FROM pjson,
LATERAL FLATTEN (pjson.v:places) as p
),
t as (
select place,addr,city,state,get_coordinate(addr) as coord from f
)
select place,addr,city,state,coord:latitude as lat, coord:longitude as lon from t;

Conclusion

Efficiency, security, and ease of use are core principles of many Snowflake features, including the Cortex suite of Generative AI services. Enhancing Snowpark External Access with IAM Role authentication supports these principles, allowing users to easily and securely connect to AWS services like Bedrock and Location Services, enabling the creation of more robust workflows and applications on the Snowflake platform.

Please reach out to your Snowflake account team for access to this feature while it’s in private preview.