From a296276ecdfc79272f1f3b29ae6fc9da1cd6cfe0 Mon Sep 17 00:00:00 2001
From: nalbam <me@nalbam.com>
Date: Tue, 4 Jun 2024 15:46:15 +0900
Subject: [PATCH] add bedrock:InvokeModel

---
 .env.example               | 11 +++--------
 .github/workflows/push.yml |  4 ++++
 handler.py                 |  6 ++----
 serverless.yml             |  6 ++++++
 4 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/.env.example b/.env.example
index 90dc58f..20db494 100644
--- a/.env.example
+++ b/.env.example
@@ -5,14 +5,9 @@ SLACK_SIGNING_SECRET="xxxx"
 
 DYNAMODB_TABLE_NAME="gureumi-ai-bot-context"
 
-OPENAI_ORG_ID="org-xxxx"
-OPENAI_API_KEY="sk-xxxx"
-OPENAI_MODEL="gpt-4o"
-
-IMAGE_MODEL="dall-e-3"
-IMAGE_SIZE="1024x1024"
-IMAGE_QUALITY="standard"
+TEXT_MODEL_ID="anthropic.claude-3-sonnet-20240229-v1:0"
+IMAGE_MODEL_ID="stability.stable-diffusion-xl-v1"
 
 SYSTEM_MESSAGE="너는 최대한 정확하고 신뢰할 수 있는 정보를 알려줘. 너는 항상 사용자를 존중해."
 
-TEMPERATURE="0.5"
+TEMPERATURE="0"
diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml
index fd0f225..95a5f8d 100644
--- a/.github/workflows/push.yml
+++ b/.github/workflows/push.yml
@@ -14,6 +14,8 @@ env:
   SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
   SLACK_SIGNING_SECRET: ${{ secrets.SLACK_SIGNING_SECRET }}
   DYNAMODB_TABLE_NAME: ${{ secrets.DYNAMODB_TABLE_NAME }}
+  TEXT_MODEL_ID: ${{ secrets.TEXT_MODEL_ID }}
+  IMAGE_MODEL_ID: ${{ secrets.IMAGE_MODEL_ID }}
   SYSTEM_MESSAGE: ${{ secrets.SYSTEM_MESSAGE }}
 
 jobs:
@@ -49,6 +51,8 @@ jobs:
           echo "SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN}" >> .env
           echo "SLACK_SIGNING_SECRET=${SLACK_SIGNING_SECRET}" >> .env
           echo "DYNAMODB_TABLE_NAME=${DYNAMODB_TABLE_NAME}" >> .env
+          echo "TEXT_MODEL_ID=${TEXT_MODEL_ID}" >> .env
+          echo "IMAGE_MODEL_ID=${IMAGE_MODEL_ID}" >> .env
           echo "SYSTEM_MESSAGE=${SYSTEM_MESSAGE}" >> .env
 
       - name: Deploy to AWS Lambda 🚀
diff --git a/handler.py b/handler.py
index c4c3afb..e4193eb 100644
--- a/handler.py
+++ b/handler.py
@@ -19,10 +19,8 @@
 DYNAMODB_TABLE_NAME = os.environ.get("DYNAMODB_TABLE_NAME", "gureumi-ai-bot-context")
 
 # Amazon Bedrock Model ID
-TEXT_MODEL_ID = os.environ.get(
-    "TEXT_MODEL_ID", "anthropic.claude-3-sonnet-20240229-v1:0"
-)
-IMAGE_MODEL_ID = os.environ.get("IMAGE_MODEL_ID", "stability.stable-diffusion-xl-v1")
+TEXT_MODEL_ID = os.environ.get("TEXT_MODEL_ID", "anthropic.claude-3")
+IMAGE_MODEL_ID = os.environ.get("IMAGE_MODEL_ID", "stability.stable-diffusion-xl")
 
 ANTHROPIC_VERSION = os.environ.get("ANTHROPIC_VERSION", "bedrock-2023-05-31")
 ANTHROPIC_TOKENS = int(os.environ.get("ANTHROPIC_TOKENS", 1024))
diff --git a/serverless.yml b/serverless.yml
index a509e88..7e1a420 100644
--- a/serverless.yml
+++ b/serverless.yml
@@ -14,6 +14,12 @@ provider:
         - dynamodb:*
       Resource:
         - "arn:aws:dynamodb:*:*:table/${self:provider.environment.DYNAMODB_TABLE_NAME}"
+    - Effect: Allow
+      Action:
+        - bedrock:InvokeModel
+      Resource:
+        - "arn:aws:bedrock:*::foundation-model/anthropic.claude-3-*"
+        - "arn:aws:bedrock:*::foundation-model/stability.stable-diffusion-xl-*"
 
 functions:
   mention: