Merge pull request #11 from ysekiy/dev/bedrock

Bedrock support

amplify/backend/api/fargate/custom-policies.json

@@ -35,5 +35,14 @@
     "Resource": [
       "arn:aws:s3:::*/*"
     ]
+  },
+  {
+    "Effect": "Allow",
+    "Action": [
+      "bedrock:InvokeModel"
+    ],
+    "Resource": [
+      "arn:aws:bedrock:*::foundation-model/*"
+    ]
   }
 ]

amplify/backend/api/fargate/src/docker-compose-template.yml

@@ -8,9 +8,10 @@ services:
       - "8080:8080"
     environment:
       - AWS_REGION=us-west-2
+      - AWS_BEDROCK_REGION=us-east-1
       - ALLOW_ORIGINS=*
       - SAGEMAKER_ENDPOINT_NAME=Rinna-Inference
-      - LLM=rinna
+      - LLM=claude_bedrock
     secrets:
       - KENDRA_INDEX_ID
       # - ANTHROPIC_API_KEY

amplify/backend/api/fargate/src/langchain/.env.development-template

@@ -1,8 +1,9 @@
 AWS_REGION=us-west-2
+AWS_BEDROCK_REGION=us-east-1
 ALLOW_ORIGINS=*
 SAGEMAKER_ENDPOINT_NAME=Rinna-Inference
 KENDRA_INDEX_ID=*******
 ANTHROPIC_API_KEY=
-LLM=rinna
+LLM=claude_bedrock
 AWS_ACCESS_KEY_ID=*******
 AWS_SECRET_ACCESS_KEY=*******

amplify/backend/api/fargate/src/langchain/README.md

@@ -28,13 +28,14 @@
 2. `cp .env.development-template .env`
 3. `.env` ファイルを編集
       - (MUST) `AWS_REGION` は AWS のリージョンを指定(Kendra や SageMaker が動いている)
+      - (MUST) `AWS_BEDROCK_REGION` には bedrock の利用リージョンを指定
       - (MUST) `KENDRA_INDEX_ID` を Kendra の Index ID に指定
       - (MUST) `AWS_ACCESS_KEY_ID` に Amplifyを操作できるユーザ のアクセスキーを設定
       - (MUST) `AWS_SECRET_ACCESS_KEY` に に Amplifyを操作できるユーザ のシークレットアクセスキーを設定
       - (WANT) `ALLOW_ORIGINS` は Access-Control-Allow-Origin の設定値 を指定
       - (WANT) `SAGEMAKER_ENDPOINT_NAME` は SageMaker エンドポイント名 を指定ください。
       - (WANT) `ANTHROPIC_API_KEY` Anthropicを利用する場合のみご指定ください。
-      - (WANT) `LLM` は、rinna か claude を指定可能。 Anthropic を利用する場合は claudeを指定ください。
+      - (WANT) `LLM` は、rinna, claude, claude_bedrock を指定可能。Anthropic を利用する場合は claudeを指定ください。
 4. `./build_and_run.sh`
 5. `http://localhost:8080` でサーバーが起動します。
 

amplify/backend/api/fargate/src/langchain/app/chain/__init__.py

@@ -1,2 +1,3 @@
 from chain.claude import build_claude_chain, build_claude_chain_without_doc, run_claude_chain
+from chain.claude_bedrock import build_claude_bedrock_chain, build_claude_bedrock_chain_without_doc, run_claude_bedrock_chain
 from chain.rinna import build_rinna_chain, run_rinna_chain

amplify/backend/api/fargate/src/langchain/app/chain/claude_bedrock.py

@@ -0,0 +1,83 @@
+# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Licensed under the MIT-0 License (https://github.com/aws/mit-0)
+"""claude on Bedrock を使った chain を定義するモジュール
+"""
+import os
+from typing import List
+
+from langchain.chains import LLMChain
+from langchain.llms.bedrock import Bedrock
+from langchain.prompts import PromptTemplate
+from schemas import KendraDocument, LLMWithDocReqBody
+
+model_id = "anthropic.claude-v1"
+bedrock_region = os.environ.get("AWS_BEDROCK_REGION", "us-east-1")
+
+
+def build_claude_bedrock_chain():
+    """claude を LLM として利用する場合の Chain の作成"""
+
+    inference_modifier = {
+            "max_tokens_to_sample": 500,
+            "temperature": 0.85, 
+            "top_k": 10, 
+            "top_p": 0.85, 
+            "stop_sequences": ["\n\nHuman:"] 
+           }
+
+    claude = Bedrock(
+        model_id=model_id,
+        region_name=bedrock_region,
+        model_kwargs=inference_modifier
+    )
+    prompt = PromptTemplate(
+        template="""Human: 資料:
+{context}
+上記の資料をもとに以下の質問に回答しなさい。[0]の形式で参考にした資料を示しなさい。また資料がないものは「わかりません」と答えなさい。\n質問: 
+{question}
+
+Assistant:""",
+        input_variables=["context", "question"],
+    )
+    return LLMChain(llm=claude, prompt=prompt)
+
+
+def build_claude_bedrock_chain_without_doc():
+    """context が与えられていない場合のプロンプトを使う Chain の作成"""
+
+    inference_modifier = {
+            "max_tokens_to_sample": 200,
+            "temperature": 0.85, 
+            "top_k": 10, 
+            "top_p": 0.85, 
+            "stop_sequences": ["\n\nHuman:"] 
+           }
+
+    claude = Bedrock(
+        model_id=model_id,
+        region_name=bedrock_region,
+        model_kwargs=inference_modifier
+    )
+    prompt = PromptTemplate(
+        template="""Human: {question}
+
+Assistant:""",
+        input_variables=["question"],
+    )
+    return LLMChain(llm=claude, prompt=prompt)
+
+
+def run_claude_bedrock_chain(chain: LLMChain, body: LLMWithDocReqBody):
+    """claude の Chain を実行する"""
+    return chain.run(
+        context=_make_context_for_claude_from_docs(body.documents),
+        question=body.userUtterance,
+    )
+
+
+def _make_context_for_claude_from_docs(documents: List[KendraDocument]):
+    """与えられた Document 情報から claude のプロンプトに埋め込むための context 情報を作成する"""
+    context: str = ""
+    for doc_id, doc in enumerate(documents):
+        context += f"[{doc_id}]{doc.title}\n{doc.excerpt}\n"
+    return context

amplify/backend/api/fargate/src/langchain/app/logics/llm_with_doc.py

@@ -9,8 +9,11 @@
     build_claude_chain,
     build_claude_chain_without_doc,
     build_rinna_chain,
+    build_claude_bedrock_chain,
+    build_claude_bedrock_chain_without_doc,
     run_claude_chain,
     run_rinna_chain,
+    run_claude_bedrock_chain
 )
 from langchain.prompts import PromptTemplate
 from schemas import LLMWithDocReqBody
@@ -31,7 +34,7 @@ def llm_with_doc(
     body: LLMWithDocReqBody,
     endpoint_name: str,
     aws_region: str,
-    llm_type: Literal["rinna", "claude"] = "rinna",
+    llm_type: Literal["rinna", "claude", "claude_bedrock"] = "claude_bedrock",
 ):
     """chain を使わずに与えられた情報をもとにプロンプトを作成して LLM に投げる"""
     if llm_type == "rinna":
@@ -43,4 +46,10 @@ def llm_with_doc(
         else:
             chain = build_claude_chain_without_doc()
         return run_claude_chain(chain, body)
+    elif llm_type == "claude_bedrock":
+        if body.documents:
+            chain = build_claude_bedrock_chain()
+        else:
+            chain = build_claude_bedrock_chain_without_doc()
+        return run_claude_bedrock_chain(chain, body)
     raise ValueError(f"unsupported LLM")

amplify/backend/api/fargate/src/langchain/app/main.py

@@ -25,7 +25,7 @@
 REGION = os.environ["AWS_REGION"]
 KENDRA_INDEX_ID: str = os.environ["KENDRA_INDEX_ID"]
 SAGEMAKER_ENDPOINT_NAME: str = os.environ.get("SAGEMAKER_ENDPOINT_NAME", None)
-LLM: Literal["rinna", "claude"] = os.environ.get("LLM", "rinna")
+LLM: Literal["rinna", "claude", "claude_bedrock"] = os.environ.get("LLM", "claude_bedrock")
 kendra_client = boto3.client("kendra", region_name=REGION)
 
 

amplify/backend/api/fargate/src/langchain/requirements.txt

@@ -1,5 +1,5 @@
-langchain==0.0.223
-boto3>=1.26.137
+langchain==0.0.310
+boto3>=1.28.62
 fastapi>=0.95.2
 uvicorn>=0.22.0
 anthropic==0.2.10

docs/DeveloperGuide.md

@@ -38,9 +38,10 @@ SageMaker エンドポイントを作成する。
    2. (Anthropic を使用する場合) `amplify/backend/api/fargate/secrets/.secret-anthropic` ファイルを作成し、Anthropic の API キーを入れる。また、`amplify/backend/api/fargate/src/docker-compose.yml` のコメントを解除する。ファイルがない/空の場合はエラーになるため使わない場合はコメントアウトする。
    4. `amplify/backend/api/fargate/src/docker-compose.yml` の環境変数を必要に応じて変更する。
       - (MUST) `AWS_REGION` を amplify を立ち上げるリージョンにする。
+      - (MUST) `AWS_BEDROCK_REGION` には bedrock の利用リージョンを指定
       - (WANT) `ALLOW_ORIGINS` は Access-Control-Allow-Origin の設定値です。
       - (WANT) `SAGEMAKER_ENDPOINT_NAME` は立ち上げた SageMaker エンドポイント名です。deploy_llm.sh で立ち上げた場合、変更の必要はありません。
-      - (WANT) `LLM` は、rinna か claude を指定可能。 Anthropicを利用する場合は claudeを指定する。
+      - (WANT) `LLM` は、rinna, claude, claude_bedrock を指定可能。 Anthropicを利用する場合は claudeを指定する。
 6.  `amplify publish` でデプロイ
    1. `? Are you sure you want to continue? (Y/n) ` は Y と入力
       - もし `You are not authorized to perform this operation` というエラーが発生した場合、ユーザー に `AdministratorAccess` ポリシー を付与して再試行お願いします。