AWS Bedrock on Snowflake (Talk to Claude, LLAMA)

Directly Linking AWS Bedrock and Snowflake Without the Hassle

Yes it is that simple to connect to AWS Bedrock directly from Snowflake without the need for snowflake external function or API Gateway.

When AI Sneaks into Data Warehousing: Expect the Unexpected

Overview of AWS Bedrock

Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs) from leading AI companies like AI21 Labs, Anthropic, Cohere, Meta, Stability AI, and Amazon with a single API, along with a broad set of capabilities you need to build generative AI applications, simplifying development while maintaining privacy and security.
Amazon Bedrock is serverless, you don't have to manage any infrastructure

What are the services required to run the project ?

• AWS Lambda
• Snowflake External Access Integration
• Snowflake Stored Procedure

Note: This process can be simplified further without the need of AWS Lambda by just using snowflake stored procedure if the anaconda snowflake channel has their boto3 package upgraded to the latest version (boto3>=1.28.57)

How does the whole process look like ?

Integration Steps

1. Setting up Snowflake Secrets

First we create two snowflake secrets to call aws services AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY.

CREATE OR REPLACE SECRET AWS_ACCESS_KEY_ID
    TYPE = GENERIC_STRING
    SECRET_STRING = '....';

CREATE OR REPLACE SECRET AWS_SECRET_ACCESS_KEY
    TYPE = GENERIC_STRING
    SECRET_STRING = '......';

2. Configuring Snowflake Network Rules

Now create snowflake network rule, to tell snowflake to allow egress on specific traffic.

CREATE OR REPLACE NETWORK RULE AWS_EGRESS_RULE
    MODE = EGRESS
    TYPE = HOST_PORT
    VALUE_LIST = (
        'lambda.us-east-1.amazonaws.com'
    );

3. Setting Up External Access Integration in Snowflake

CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION bedrock_client
    ALLOWED_NETWORK_RULES = (AWS_EGRESS_RULE)
    ALLOWED_AUTHENTICATION_SECRETS = (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)
    ENABLED = TRUE;
4. Creating the Stored Procedure
This procedure will call AWS Lambda.
CREATE OR REPLACE 
  PROCEDURE snow_claude(question STRING)
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION=3.9
EXTERNAL_ACCESS_INTEGRATIONS = (bedrock_client)
SECRETS = ('AWS_ACCESS_KEY_ID' = AWS_ACCESS_KEY_ID, 'AWS_SECRET_ACCESS_KEY' = AWS_SECRET_ACCESS_KEY)
PACKAGES=('snowflake-snowpark-python', 'boto3', 'botocore', 'requests')
HANDLER='handler'
AS
$$
import json

import _snowflake
import boto3
from snowflake.snowpark import Session

# Your AWS credentials
AWS_ACCESS_KEY_ID = _snowflake.get_generic_secret_string("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = _snowflake.get_generic_secret_string("AWS_SECRET_ACCESS_KEY")

#Lambda function name
FUNCTION_NAME = "snowbedrock"

# Initialize the boto3 client for Lambda
session = boto3.session.Session(region_name="us-east-1")
lambda_client = session.client(
    "lambda",
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
)

def invoke_lambda_function(payload):
    response = lambda_client.invoke(
        FunctionName=FUNCTION_NAME,
        InvocationType="RequestResponse",
        Payload=json.dumps(payload),
    )
    # Since the Lambda response is already parsed into a string
    return response["Payload"].read().decode("utf-8")

def handler(session: Session, question: str):
    body = {"question": question}
    try:
        response = invoke_lambda_function(body)
        parsed_response = json.loads(response)
        return parsed_response
    except Exception as err:
        print("An error occurred:", err)
        raise

$$
;

5. Lambda Function Setup

Note: The boto3 package is not upto date in AWS Lambda, make sure to zip the latest version package (boto3>=1.28.57) as a layer, the older version doesn't know about bedrock.

#snowbedrock

import json
import boto3
import os

def claude(prompt):
    bedrock_runtime = boto3.client(
        service_name="bedrock-runtime",
        region_name="us-east-1",
    )
    TEMPLATE = f'Human: "{prompt}"\nAssistant: '

    body = json.dumps(
        {
            "prompt": TEMPLATE,
            "max_tokens_to_sample": 256,
            "stop_sequences": [],
            "temperature": 0,
            "top_p": 0.9,
        }
    )
    modelId = "anthropic.claude-instant-v1"
    accept = "application/json"
    contentType = "application/json"
    response = bedrock_runtime.invoke_model(
        body=body, modelId=modelId, accept=accept, contentType=contentType
    )
    response_body = json.loads(response.get("body").read())
    print("Model Response:", response_body)

    return response_body

def lambda_handler(event, context) -> str:    
    question = event.get("question")
    if not question:
        return "No question provided"

    answer = claude(question)
    return answer["completion"]

Make sure the lambda has permission to trigger bedrock
{
 "Version": "2012–10–17",
 "Statement": [
 {
 "Effect": "Allow",
 "Action": "bedrock:*",
 "Resource": "*"
 }
 ]
}

---

6. Time to test

call snow_claude('who is this elon musk?');

Elon Musk is a South African-born American entrepreneur and businessman. Some key things to know about him: - He is the founder, CEO, and Chief Engineer of SpaceX…..

---

I hope this guide simplifies your AWS Bedrock and Snowflake integration process.
Follow for more insightful content and to explore the world of open-source with me!
Medium: Kaarthikandavar
X: Kaarthikcodes
LinkedIn: Kaarthik

Thanks for Reading! Just another day in the tech soap opera.