最近、有人チャットボットにRAG(Retrieval-Augmented Generation)としてAmazon Bedrock Knowledge Baseを組み込み、相手の最新のメッセージに対して複数の返信候補を提案する機能を追加しました。提案精度を向上させるためにいろいろチューニングしたので、それらの手法を共有します。
GPTやClaudeに聞けば一発で教えてくれるだろ！と思うかもしれませんが、知識カットオフでこれらのLLMが把握していない情報もあったりするので、現時点では少しはこの記事が役に立つと思います。
コード全文は「実装例: Bedrock Knowledge Baseサービスとクライアント」に書いたので、そこだけ知りたい方は飛ばしてご覧ください(TypeScriptの例となります)。
あと、もっといい方法があるよ！といった指摘やその他意見があれば、自分の勉強にもなるのでコメントをいただけるとうれしいです。
1. ナレッジベース作成時にFoundation Models as a Parser(FMP)を選択する
S3をデータソースとして選択する際、Notionからエクスポートされた情報(マークダウンや画像)を正確にインデックス化する必要があったので、Foundation Models as a Parser(FMP)を選択しました。
画像内のテキスト情報もインデックス化されて検索対象となるので、スクリーンショットや図表を含むドキュメントを扱う場合に有効です。
2. 検索結果数の最適化
OpenSearch Serverlessをナレッジベースとして使用しているのですが、デフォルトの検索結果数(5)では質問に対して十分な情報を取得できないケースがありました。
vectorSearchConfigurationのnumberOfResultsパラメータをデフォルトの5から10に増やすことで、回答の失敗率が下がりました。
retrievalConfiguration: {
vectorSearchConfiguration: {
numberOfResults: 10, // 回答失敗率を減らすためにベクトル検索の結果数を10に設定
}
}
検索結果数を増やすと処理時間が若干増加してレスポンス時間も伸びますが(自分の場合は1秒くらい)、回答品質の向上と失敗率の低減というメリットのほうが大きかったので10にしました。ユースケースによっては5〜20の間で調整したほうがいいと思います。
3. クエリ分解(Query Decomposition)の活用
複雑な質問や複数の要素を含む質問に対して、Bedrockのクエリ分解機能を活用することで検索精度を向上させることができます。
orchestrationConfiguration: {
queryTransformationConfiguration: {
type: 'QUERY_DECOMPOSITION' as const
}
}
クエリ分解を有効にすると、「AとBの違いを教えてください」といった複合的な質問を「Aとは何か」「Bとは何か」「AとBの違いは何か」のような複数のサブクエリに分解し、それぞれに対して個別に検索を行います。これによって、より関連性の高い情報を取得できるようになります。
注意点
クエリ分解を有効にしたら、レスポンス時間が5秒くらい伸びてしまいました。リアルタイム性が要求される用途では、この機能の有効化については検討したほうがよさそうです。
4. システムプロンプトに会話履歴を動的に含める
generationConfigurationのtextPromptTemplateを活用して、システムプロンプトを動的に構成することで、より自然な対話を実現できます。特に会話履歴を含めることで、文脈を考慮した応答が可能になります。
以下のコードでは、最新のユーザーメッセージを除いた会話履歴を動的にシステムプロンプトに含めています。
private buildPrompt(chats: Array<ChatEntity>): { systemPrompt: string; userPrompt: string } {
const latestVisitorChatIndex = [...chats].reverse().findIndex((chat) => chat.sender === 'visitor')
if (latestVisitorChatIndex === -1) {
throw new BadRequestException('最新メッセージが見つからず、AIが返信候補を生成できませんでした')
}
const latestVisitorChat = chats[chats.length - 1 - latestVisitorChatIndex]
const previousChats = chats.slice(0, chats.length - 1 - latestVisitorChatIndex)
const formattedHistory =
previousChats.length > 0
? previousChats.map((chat) => `${chat.sender}: ${chat.chat ?? ''}`).join('\n')
: 'No previous conversation history available.'
const systemPrompt = `
Human: You are a question answering agent. I will provide you with a set of search results and a user's question.
Your job is to answer the user's question using only information from the search results.
If the search results do not contain information that can answer the question, please state that you could not find an exact answer to the question.
Do not include any explanations, introductions, or additional text before or after the JSON output.
Your task is to respond to potential customers in a polite and professional manner that builds trust.
Each response should be concise (up to 3 lines) and use proper honorific language (e.g., "〜いただけますでしょうか", "〜でございます").
Always include a word of gratitude or empathy in response to the customer's statement.
Here are the search results in numbered order:
$search_results$
Here is the conversation history:
${formattedHistory}
Here is the user's question:
$query$
Generate the response in Japanese language.
Ensure the response follows this exact JSON format:
{"replies":["返信1","返信2","返信3"]}
Assistant:
`.trim()
return {
systemPrompt,
userPrompt: latestVisitorChat.chat ?? ''
}
}
buildPromptでシステムプロンプト(systemPrompt)とユーザープロンプト(userPrompt)を作成してbedrockKnowledgeBaseServiceに渡しています。
const { systemPrompt, userPrompt } = this.buildPrompt(chats)
const llmResponse = await this.bedrockKnowledgeBaseService.queryKnowledgeBase(systemPrompt, userPrompt)
注意点
カスタムプロンプトテンプレートを使用すると、デフォルトの$output_format_instructions$が上書きされて、citations(引用情報)がレスポンスに含まれなくなります。引用情報が必要な場合は、プロンプト内に$output_format_instructions$を含める必要がありますが、その場合はレスポンス形式の完全なカスタマイズができなくなるトレードオフがあります。
現在のBedrockでは、カスタムフォーマット(JSONなど)と引用情報の両方を同時に取得することは難しいため、要件に応じて優先度を決める必要があります。
5. キーワード検索とベクトル検索のハイブリッド化
Bedrock Knowledge Baseでは、ベクトル検索とキーワード検索を組み合わせたハイブリッド検索が可能です。
retrievalConfiguration: {
vectorSearchConfiguration: {
numberOfResults: 10,
overrideSearchType: 'HYBRID' as const
}
}
ベクトル検索は意味的な類似性に優れており、キーワード検索は正確なキーワードマッチングに強みがあります。両者を組み合わせることで、より関連性の高い情報を抽出できる可能性が高まります。
デフォルトではBedrockが検索戦略を決めてくれるとのことだったので、一旦いじらずにデフォルト設定をしています。
By default, Amazon Bedrock decides a search strategy for you.
6. リランキング(未検証)
こちらはコスト増の懸念から試していないのですが、検索結果の取得後にリランキング(再順位付け)を行うことで、より関連性の高い順に結果を並べ替えることができます。リランキングは専用のリランカーモデルを使用するため、Claudeとは別のモデルを組み合わせて使用することになります。
retrievalConfiguration: {
vectorSearchConfiguration: {
numberOfResults: 10,
overrideSearchType: 'HYBRID' as const
},
rerankingConfiguration: {
type: 'BEDROCK_RERANKING_MODEL' as const,
bedrockRerankingConfiguration: {
modelConfiguration: {
modelArn: 'arn:aws:bedrock:region:account:reranker/model-id'
},
numberOfRerankedResults: 5 // リランキング後に返す結果の数
}
}
}
リランキングでは、検索で取得した結果に対して、クエリとの関連性をより精密に評価し直します。ベクトル検索やキーワード検索が大まかな関連性を見つけるのに対して、リランカーモデルはクエリと文書の意味的な関係をより深く理解して評価します。
注意点
リランキングを使用するとコストと処理時間が増加します。ただし、より関連性の高い結果に絞り込むことでLLMへの入力を最適化でき、結果的に回答精度の向上につながる可能性があります。現時点では試していませんが、今後の改善策として検討しています。
実装例: Bedrock Knowledge Baseサービスとクライアント
以下は、上記の最適化を適用したTypeScriptの実装例です。
Bedrock Knowledge Baseサービス
@Injectable()
export class BedrockKnowledgeBaseService {
constructor(private readonly bedrockKnowledgeBaseClient: BedrockKnowledgeBaseClient) {}
/**
* ナレッジベースに基づいて LLM(Claude)に問い合わせる
* @param systemPrompt システムプロンプト
* @param userPrompt ユーザーのクエリ
* @returns AIの応答テキスト
*/
async queryKnowledgeBase(systemPrompt: string, userPrompt: string): Promise<string> {
try {
const ragConfig = {
type: 'KNOWLEDGE_BASE' as const,
knowledgeBaseConfiguration: {
knowledgeBaseId: process.env.BEDROCK_KNOWLEDGE_BASE_ID,
modelArn: process.env.BEDROCK_KNOWLEDGE_BASE_MODEL_ARN,
orchestrationConfiguration: {
queryTransformationConfiguration: { type: 'QUERY_DECOMPOSITION' as const }
},
retrievalConfiguration: {
vectorSearchConfiguration: {
numberOfResults: 10,
overrideSearchType: 'HYBRID' as const
}
},
generationConfiguration: {
promptTemplate: {
textPromptTemplate: systemPrompt
},
inferenceConfig: {
textInferenceConfig: {
maxTokens: 500,
temperature: 0.2,
topP: 0.5,
topK: 10
}
}
}
}
}
const input: RetrieveAndGenerateCommandInput = {
input: { text: userPrompt },
retrieveAndGenerateConfiguration: ragConfig
}
const response = await this.bedrockKnowledgeBaseClient.sendQuery(input)
if (!response) {
throw new InternalServerErrorException('Failed to get a valid response from Bedrock.')
}
return response
} catch (error) {
console.error('Error querying Bedrock API:', error)
throw new InternalServerErrorException('Failed to get response from Bedrock.')
}
}
}
Bedrock Knowledge Baseクライアント
@Injectable()
export class BedrockKnowledgeBaseClient {
private readonly client: BedrockAgentRuntimeClient
constructor() {
this.client = new BedrockAgentRuntimeClient({
region: process.env.REGION
})
}
/**
* Bedrock のナレッジベースに問い合わせる
* @param input RetrieveAndGenerateCommandInput
* @returns AIの応答
*/
async sendQuery(input: RetrieveAndGenerateCommandInput): Promise<string> {
try {
const command = new RetrieveAndGenerateCommand(input)
const response = await this.client.send(command)
if (!response.output?.text) {
throw new InternalServerErrorException('Failed to get a valid response from Bedrock.')
}
return response.output.text
} catch (error) {
console.error('Error querying Bedrock API:', error)
throw new InternalServerErrorException('Failed to get response from Bedrock.')
}
}
}
今後の検討事項
検索スコアの調整(閾値設定)で不要な結果を除外することができたりしないのかな?と思ったのですが、Bedrock Knowledge Baseではできなそうでした(間違っていたらご指摘お願いします)。
記事内にも書いたのですが、クエリ分解でレイテンシが大きくなってしまったので(レスポンスまでに10秒強)、その対策をどうしようか現在検討中です。それこそ、Bedrockのみバージニアリージョンで利用して、クロスリージョンによる通信のオーバーヘッドも受け入れつつ、3.5 Haikuなどのモデルを使うか、Llamaなどのオープンソースモデルをカスタムモデルインポートで使おうか、などと考えています。今のところTokyoリージョンで使えるモデルが少なくて困りますね。
あと、ナレッジベースのデータ元がNotionに集約されていて、現在は手動でマークダウンやPNGをエクスポートしているのですが、さすがに手間なのでNotion APIで自動取得するような仕組みをつくりたいと思っています。
現在はオンライン評価のみなので、オフラインで自動評価(LLM as a Judge)もしたいですね。課題がたくさん..
参考資料