LangChainからクロスアカウントのAmazon Kendraサービスに接続する【LLMアプリ】

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": "arn:aws:sts::{{ メインAWSアカウントID }}:assumed-role/{{ lambda-function-name }}"
            },
            "Action": "sts:AssumeRole",
            "Condition": {}
        }
    ]
}

from aws_cdk import (
  Duration,
  Size,
  Stack,
  aws_ecr as ecr,
  aws_lambda as _lambda,
  aws_iam as iam,
  aws_logs as logs,
)

from constructs import Construct
import os
from dotenv import load_dotenv
load_dotenv()
class LangchainAiAgentStack(Stack):
  def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
    super().__init__(scope, construct_id, **kwargs)

    # LambdaイメージのECRリポジトリ 
    ## langchainのパッケージが大きいためimage経由とします。
    lambda_ecr_repository = ecr.Repository.from_repository_name(
      self, 'kendra-sts-lambda-image-ecr-repository',
      '{{ ##ecr-name }}'
    )
	
    # lambda_Roleの設定
    lambda_role = iam.Role(
      self, 'kendra-sts-lambda-role',
      assumed_by=iam.ServicePrincipal("lambda.amazonaws.com"),
      inline_policies={
        "lambda_policy": iam.PolicyDocument(
          statements=[iam.PolicyStatement(
            actions=[
              "logs:CreateLogGroup",
              "logs:CreateLogStream",
              "logs:PutLogEvents",
              "lambda:InvokeFunction",
              "kendra:*",
            ],
            resources=["*"],
          )]
        ),
        ## kendraアカウントを使用するためのクロスアカウントポリシー
        "kendra_account_access_policy": iam.PolicyDocument(
          statements=[iam.PolicyStatement(
            actions=[
              "sts:AssumeRole"
            ],
            resources=["{{ kendraアカウント作成したIAM RoleのARN }}"],
          )]
        )
      }
    )

    # lambda関数の定義
    lambda_function = _lambda.Function(
      self, 'kendra-sts-lambda-function',
      code=_lambda.Code.from_ecr_image(
        lambda_ecr_repository,
        tag='latest'
      ),
      handler=_lambda.Handler.FROM_IMAGE,
      runtime=_lambda.Runtime.FROM_IMAGE,
      role=lambda_role,
	  ## kendaraアカウントで作成したロールARNを環境変数に設定。※envファイル推奨
      environment={
        "ROLE_ARN": "{{ kendraアカウント作成したIAM RoleのARN }}",
      },
      memory_size=1024,
      timeout=Duration.minutes(5),
    )

import json
import boto3
import os
import hashlib
import logging
from datetime import timedelta
from dotenv import load_dotenv
from momento import (
  CacheClient,
  Configurations,
  CredentialProvider,
  PreviewVectorIndexClient,
)
from langchain.chat_models import ChatOpenAI
from langchain.memory import MomentoChatMessageHistory, ConversationBufferMemory
from langchain.retrievers import AmazonKendraRetriever
from langchain.chains import RetrievalQA, ConversationalRetrievalChain

# 今回は説明を割愛しますが、OpenAIやmomentoなどのAPIキーは環境変数を利用しています。
load_dotenv()

logger = logging.getLogger(__name__)
def handler(event, context):
  # ※実際には、フロントからbodyで受け取ります。
  index_id =　"{{ kendraのindexID }}"
  human_input = "{{ ユーザーの入力値 }}"
  session_id = "{{ 会話履歴のためのセッションID }}"

  # llmモデル（OpenAI）
  llm = ChatOpenAI(
    model_name=os.environ["OPENAI_API_MODEL"],
    temperature=float(os.environ["OPENAI_API_TEMPERATURE"]),
  )

  # kendraからドキュメント参照
  ## 今回のメイン!!
  retriever = initialize_kendraindex(index_id)

  # 会話履歴（momento cacheを使用 ※今回は説明割愛）
  history = MomentoChatMessageHistory.from_client_params(
    session_id,
    os.environ["MOMENTO_CACHE"],
    timedelta(hours=int(os.environ["MOMENTO_TTL"])),
  )
  memory = ConversationBufferMemory(
    memory_key="chat_history",
    chat_memory=history,
    input_key="question",
    output_key="answer",
    return_messages=True,
  )

  # 回答生成
  qa = ConversationalRetrievalChain.from_llm(
    llm=llm,
    retriever=retriever,
    memory=memory,
    return_source_documents=True,
  )
  res = qa({"question": human_input})

  return {
    "statusCode": 200,
	"body": json.dumps(res["answer"])
  }


# 他アカウントKendraに接続するためのsts
# clientプロパティを与え、LangChainのAmazonKendraRetorieverモデル使用
def initialize_kendraindex(index_id):
  # kendra-account-roleの一時セッションを付与
  sts_connection = boto3.client('sts')
  kendra_account = sts_connection.assume_role(
    RoleArn = os.environ["ROLE_ARN"],　# 環境変数より
    RoleSessionName ="cross_acct_lambda",
  )
  ACCESS_KEY = kendra_account['Credentials']['AccessKeyId']
  SECRET_KEY = kendra_account['Credentials']['SecretAccessKey']
  SESSION_TOKEN = kendra_account['Credentials']['SessionToken']

  client = boto3.client('kendra',
    aws_access_key_id=ACCESS_KEY,
    aws_secret_access_key=SECRET_KEY,
    aws_session_token=SESSION_TOKEN,
  )
  # kendraはデフォルトが英語なので、言語設定必須
  attribute_filter = {"EqualsTo": {"Key": "_language_code","Value": {"StringValue": "ja"}}}
  
  # langchainのAmazonKendraRetorieverモデルにclientオプションを設定
  retriever = AmazonKendraRetriever(
    client=client,
    index_id=index_id,
    attribute_filter=attribute_filter,
  )
  return retriever