[アップデート] Amazon Bedrock Knowledge bases がストリーミングレスポンスをサポートしました #AWSreInvent
こんにちは! AWS 事業本部コンサルティング部のたかくに(@takakuni_)です。
re:Invent 2024 の季節ですね。
Amazon Bedrock Knowledge bases がストリーミングレスポンスをサポートしました。やったね
何が嬉しいのか
今まで Amazon Bedrock Knowledge bases ではストリーミングレスポンスをサポートしておらず、以下の 2 つの API がサポートしていました。
- Retrieve(ベクトル検索結果のみ取得する API)
- RetrieveAndGenerate(ベクトル検索 + 回答を生成する API)
今回新たに RetrieveAndGenerateStream API が登場しました。
Today, we are announcing the support of RetrieveAndGenerateStream API in Bedrock Knowledge Bases. This new streaming API allows Bedrock Knowledge Base customers to receive the response as it is being generated by the Large Language Model (LLM), rather than waiting for the complete response.
やってみる
早速ですがやってみましょう。と言いつつも、 API リファレンスに RetrieveAndGenerateStream がないですね。
そのような場合は API Changes を追いましょう。今後、何かの役に立つかもしれません(?)
Usage 書いてありますね。
client.retrieve_and_generate_stream(
input={
'text': 'string'
},
retrieveAndGenerateConfiguration={
'externalSourcesConfiguration': {
'generationConfiguration': {
'additionalModelRequestFields': {
'string': {...}|[...]|123|123.4|'string'|True|None
},
'guardrailConfiguration': {
'guardrailId': 'string',
'guardrailVersion': 'string'
},
'inferenceConfig': {
'textInferenceConfig': {
'maxTokens': 123,
'stopSequences': [
'string',
],
'temperature': ...,
'topP': ...
}
},
'promptTemplate': {
'textPromptTemplate': 'string'
}
},
'modelArn': 'string',
'sources': [
{
'byteContent': {
'contentType': 'string',
'data': b'bytes',
'identifier': 'string'
},
's3Location': {
'uri': 'string'
},
'sourceType': 'S3'|'BYTE_CONTENT'
},
]
},
'knowledgeBaseConfiguration': {
'generationConfiguration': {
'additionalModelRequestFields': {
'string': {...}|[...]|123|123.4|'string'|True|None
},
'guardrailConfiguration': {
'guardrailId': 'string',
'guardrailVersion': 'string'
},
'inferenceConfig': {
'textInferenceConfig': {
'maxTokens': 123,
'stopSequences': [
'string',
],
'temperature': ...,
'topP': ...
}
},
'promptTemplate': {
'textPromptTemplate': 'string'
}
},
'knowledgeBaseId': 'string',
'modelArn': 'string',
'orchestrationConfiguration': {
'additionalModelRequestFields': {
'string': {...}|[...]|123|123.4|'string'|True|None
},
'inferenceConfig': {
'textInferenceConfig': {
'maxTokens': 123,
'stopSequences': [
'string',
],
'temperature': ...,
'topP': ...
}
},
'promptTemplate': {
'textPromptTemplate': 'string'
},
'queryTransformationConfiguration': {
'type': 'QUERY_DECOMPOSITION'
}
},
'retrievalConfiguration': {
'vectorSearchConfiguration': {
'filter': {
'andAll': [
{'... recursive ...'},
],
'equals': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'greaterThan': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'greaterThanOrEquals': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'in': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'lessThan': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'lessThanOrEquals': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'listContains': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'notEquals': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'notIn': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'orAll': [
{'... recursive ...'},
],
'startsWith': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
},
'stringContains': {
'key': 'string',
'value': {...}|[...]|123|123.4|'string'|True|None
}
},
'implicitFilterConfiguration': {
'metadataAttributes': [
{
'description': 'string',
'key': 'string',
'type': 'STRING'|'NUMBER'|'BOOLEAN'|'STRING_LIST'
},
],
'modelArn': 'string'
},
'numberOfResults': 123,
'overrideSearchType': 'HYBRID'|'SEMANTIC',
'rerankingConfiguration': {
'bedrockRerankingConfiguration': {
'metadataConfiguration': {
'selectionMode': 'SELECTIVE'|'ALL',
'selectiveModeConfiguration': {
'fieldsToExclude': [
{
'fieldName': 'string'
},
],
'fieldsToInclude': [
{
'fieldName': 'string'
},
]
}
},
'modelConfiguration': {
'additionalModelRequestFields': {
'string': {...}|[...]|123|123.4|'string'|True|None
},
'modelArn': 'string'
},
'numberOfRerankedResults': 123
},
'type': 'BEDROCK_RERANKING_MODEL'
}
}
}
},
'type': 'KNOWLEDGE_BASE'|'EXTERNAL_SOURCES'
},
sessionConfiguration={
'kmsKeyArn': 'string'
},
sessionId='string'
)
作成したコードは以下になります。
import json
import boto3
KNOWLEDGEBASE_ID = "ES2HX0CTWR" # ナレッジベース ID を記載
response_model_arn = "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-5-sonnet-20241022-v2:0"
bedrock_agent = boto3.client("bedrock-agent-runtime", region_name="us-west-2")
prompt = (
"おばあさんはどこに行きましたか?"
)
response = bedrock_agent.retrieve_and_generate_stream(
input={"text": prompt},
retrieveAndGenerateConfiguration={
"type": "KNOWLEDGE_BASE",
"knowledgeBaseConfiguration": {
"knowledgeBaseId": KNOWLEDGEBASE_ID,
"modelArn": response_model_arn,
},
}
)
stream = response.get("stream")
if stream:
for event in stream:
if "output" in event:
print(event['output']['text'], end="", flush=True)
ストリーミングできていますね。
ちなみにですが、最後にソースチャンクが出力されていました。引用元も表示したい場合は別途処理が必要そうです。
(app-py3.12) takakuni.shinnosuke@HL01556 app % python index.py
{'output': {'text': '\nおば'}}
{'output': {'text': 'あさんは'}}
{'output': {'text': '川'}}
{'output': {'text': 'へ'}}
{'output': {'text': '洗濯に'}}
{'output': {'text': '行きました。'}}
{'output': {'text': '\n'}}
{'citation': {'citation': {'generatedResponsePart': {'textResponsePart': {'span': {'end': 17, 'start': 0}, 'text': '\nおばあさんは川へ洗濯に行きました。'}}, 'retrievedReferences': [{'content': {'text': 'むかし、むかし、あるところに、おじいさんとおばあさんがありました。まいにち、おじいさんは山へしば刈りに、おばあさんは川へ洗濯に行きました。 ある日、おばあさんが、川のそばで、せっせと洗濯をしていますと、川上から、大きな桃が一つ、「ドンブラコッコ、スッコッコ。ドンブラコッコ、スッコッコ。」と流れて来ました。 「おやおや、これはみごとな桃だこと。おじいさんへのおみやげに、どれどれ、うちへ持って帰りましょう。」 おばあさんは、そう言いながら、腰をかがめて桃を取ろうとしましたが、遠くって手がとどきません。 おばあさんはそこで、「あっちの水は、かあらいぞ。こっちの水は、ああまいぞ。'}, 'location': {'s3Location': {'uri': 's3://rerank-dev-us-west-2-kb-datasource-123456789012/桃太郎第1章.txt'}, 'type': 'S3'}, 'metadata': {'x-amz-bedrock-kb-source-uri': 's3://rerank-dev-us-west-2-kb-datasource-123456789012/桃太郎第1章.txt', 'x-amz-bedrock-kb-data-source-id': '7IX0SSZAKY'}}]}}}
まとめ
以上、「Amazon Bedrock Knowledge bases がストリーミングレスポンスをサポートしました。」でした。
re:Invent 2024 でいろんなアップデートが来ているので、どんどんお届けできればと思います。
このブログがどなたかの参考になれば幸いです。
AWS 事業本部コンサルティング部のたかくに(@takakuni_)でした!