Enable client side to control model id and region

aws-samples · May 24, 2024 · 65e6d2d · 65e6d2d
1 parent 2fdab56
commit 65e6d2d
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 6 deletions.
diff --git a/.env.development-template b/.env.development-template
@@ -1,2 +1,4 @@
 VITE_INDEX_ID=*****
-# VITE_SERVER_URL=http://localhost:8080
+# VITE_SERVER_URL=http://localhost:8080
+VITE_MODEL_ID=anthropic.claude-3-haiku-20240307-v1:0
+VITE_BEDROCK_REGION=us-west-2
diff --git a/amplify/backend/awscloudformation/override.ts b/amplify/backend/awscloudformation/override.ts
@@ -55,6 +55,23 @@ export function override(resources: AmplifyRootStackTemplate, amplifyProjectInfo
     envContent += `\nVITE_STREAM_FUNC_NAME=streamClaude3-${envName}`;
   }
 
+  // VITE_MODEL_ID の値を追加
+  const modelId = "anthropic.claude-3-haiku-20240307-v1:0"
+  const modelIdPattern = /^VITE_MODEL_ID=.*/gm;
+  if (modelIdPattern.test(envContent)) {
+    envContent = envContent.replace(modelIdPattern, `VITE_MODEL_ID=${modelId}`);
+  } else {
+    envContent += `\nVITE_MODEL_ID=${modelId}`;
+  }
+
+  // VITE_BEDROCK_REGION の値を追加
+  const bedrockRegionPattern = /^VITE_BEDROCK_REGION=.*/gm;
+  if (bedrockRegionPattern.test(envContent)) {
+    envContent = envContent.replace(bedrockRegionPattern, `VITE_BEDROCK_REGION=${region_name}`);
+  } else {
+    envContent += `\nVITE_BEDROCK_REGION=${region_name}`;
+  }
+
   // .env ファイルに書き込む
   fs.writeFileSync('.env', envContent);
 }
diff --git a/amplify/backend/function/streamClaude3/src/index.js b/amplify/backend/function/streamClaude3/src/index.js
@@ -8,8 +8,6 @@ const {
   ThrottlingException,
 } = require('@aws-sdk/client-bedrock-runtime');
 
-const client = new BedrockRuntimeClient();
-
 
 const extractOutputTextClaude3Message = (body) => {
   if (body.type === 'message') {
@@ -20,7 +18,9 @@ const extractOutputTextClaude3Message = (body) => {
   return '';
 };
 
-async function* invokeStream(input) {
+async function* invokeStream(region, input) {
+  const client = new BedrockRuntimeClient({ region });
+
   try {
 
     const command = new InvokeModelWithResponseStreamCommand(input);
@@ -63,7 +63,7 @@ exports.handler = awslambda.streamifyResponse(
   async (event, responseStream, context) => {
     context.callbackWaitsForEmptyEventLoop = false;
 
-    for await (const token of invokeStream?.(event.body) ?? []) {
+    for await (const token of invokeStream?.(event.body.bedrockRegion, event.body) ?? []) {
       responseStream.write(token);
     }
     responseStream.end();

diff --git a/src/utils/service.ts b/src/utils/service.ts
@@ -32,6 +32,16 @@ if (!import.meta.env.VITE_STREAM_FUNC_NAME) {
     "環境変数にSTREAM_FUNC_ARNがありません"
   );
 }
+if (!import.meta.env.VITE_MODEL_ID) {
+  _loadingErrors.push(
+    "環境変数にMODEL_IDがありません"
+  );
+}
+if (!import.meta.env.VITE_BEDROCK_REGION) {
+  _loadingErrors.push(
+    "環境変数にBEDROCK_REGIONがありません"
+  );
+}
 const hasErrors = _loadingErrors.length > 0;
 if (hasErrors) {
   console.error(JSON.stringify(_loadingErrors));
@@ -44,6 +54,8 @@ const stream_func_name: string = import.meta.env.VITE_STREAM_FUNC_NAME ?? ""
 const local_server = import.meta.env.VITE_SERVER_URL ?? ""
 const remote_server = awsconfig.aws_cloud_logic_custom[0].endpoint ?? ""
 export const serverUrl: string = local_server ? local_server : remote_server;
+const model_id: string = import.meta.env.VITE_MODEL_ID ?? ""
+const bedrock_region: string = import.meta.env.VITE_BEDROCK_REGION ?? ""
 let jwtToken = "";
 
 Amplify.configure({
@@ -284,7 +296,8 @@ export async function* infStreamClaude(user_prompt: string) {
   }
   const req = {
     "body": {
-      "modelId": "anthropic.claude-3-haiku-20240307-v1:0",
+      "bedrockRegion": bedrock_region,
+      "modelId": model_id,
       "accept": "application/json",
       "contentType": "application/json",
       "body": JSON.stringify(body)