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.

Is it your intention to directly query the knowledge base or instead have a Bedrock agent interact with the customer?

If the latter, then you likely need to reconsider your approach.

Based on the code you've shared, you're not calling the Bedrock agent/assistant.

Rather, you're using the Bedrock API's RetrieveAndGenerate method to directly query the knowledge base.

Ideally, what you should do is create a Bedrock Agent that is supported by your existing knowledge based. You then specify the context in which this agent works. That is, you give it a prompt that tells the Agent what its purpose is; e.g.:

You are a product information assistant. Assist the customer with questions they have about our products. Always seek confirmation from the customer regarding which product they're enquiring about.

Then in your code use the InvokeAgent method to call the agent, which in turn will use the provided knowledge bases to get the required data and based on the user's input will return the correct information.

Have a read through the Create Bedrock Agents guide.