如何在boto3调用带知识库的AWS Bedrock时传入max_token_to_sample参数

你可以通过在retrieveAndGenerateConfiguration的knowledgeBaseConfiguration中添加generationConfiguration字段，传入模型生成相关的参数，包括最大令牌数。需要注意不同模型对应的参数名称有差异：

• Anthropic Claude系列模型使用max_tokens参数
• Amazon Titan系列模型使用maxTokenCount参数
• 其他模型请参考对应模型的官方文档确认参数名

修改后的代码示例（以Claude模型为例）：

aws_session = boto3.Session()
bedrock_agent_client = aws_session.client(service_name="bedrock-agent-runtime", region_name="us-west-2")
response = bedrock_agent_client.retrieve_and_generate(
    input={"text": input_data},
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": config.bedrock.kb_id,
            "modelArn": model_arn,
            "generationConfiguration": {
                "promptTemplateConfiguration": {
                    "textPromptTemplate": ""  # 不需要自定义模板可留空，使用知识库默认模板
                },
                "inferenceConfiguration": {
                    "textInferenceConfiguration": {
                        "max_tokens": 1024  # 设置你需要的最大令牌数
                    }
                }
            }
        },
    },
)
return response

如果使用Amazon Titan模型，只需将textInferenceConfiguration中的参数改为maxTokenCount：

"textInferenceConfiguration": {
    "maxTokenCount": 1024
}

内容的提问来源于stack exchange，提问作者Liang HE