[Python] How to make my AWS Bedrock RAG AI Assistant ask which product the user is asking about

I added a RAG AI assistant using AWS Bedrock to our website so customers can ask questions and get answers about the four products we sell. Each product has the same two documents, like: product_A_user_guide.pdf, product_A_specs.pdf, product_B_user_guide.pdf, product_B_specs.pdf, etc.

The products offer similar features but differ in how to perform those features. So, when asking the AI assistant a question without specifying the product it doesn’t always return the correct result. For example, if you ask a question that pertains to all four products, like “how do I run a calibration test”, it might return an answer from product A, B, C or D. How can I make my assistant smarter so that it asks the user “tell me which product you’re asking about: A, B, C or D?” if the user doesn’t specify it in their question?

This is my code:

bedrock = boto3.client("bedrock-agent-runtime", region_name="us-west-1")

def lambda_handler(event, context):

model_id = "amazon.titan-text-premier-v1:0"
question = event["queryStringParameters"]["question"]
session_id = event["queryStringParameters"]["session_id”] if “session_id" in event["queryStringParameters"] else None

kb_id = os.environ["KNOWLEDGE_BASE_ID"]
region = "us-west-1"
model_arn = f"arn:aws:bedrock:{region}::foundation-model/{model_id}"

# Query the knowledge base
response = queryKB(question, kb_id, model_arn, session_id)

# Extract the generated text and session ID from the response
generated_text = response["output"]["text"].strip()
session_id = response.get("sessionId", "")
citations = response["citations"]

headers = {
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Credentials": True
}

# Return the response in the expected format
return {
“statusCode": 200,
"headers": headers,
"body": json.dumps({
"question": question,
"answer": generated_text,
"citations": citations,
"sessionId": session_id
}, ensure_ascii=False)
}

def queryKB(question, kbId, model_arn, sessionId=None):
return bedrock.retrieve_and_generate(
input={
"text": question
},
retrieveAndGenerateConfiguration={
“type": "KNOWLEDGE_BASE",
"knowledgeBaseConfiguration": {
"knowledgeBaseId": kbId,
"modelArn": model_arn
}
},
sessionId=sessionId if sessionId else None
)


Thanks for your help.

Continue reading...