ask_bedrock_llm_with_knowledge_base.py

# Function to retrieve info from RAG with knowledge base
def ask_bedrock_llm_with_knowledge_base(flat_conversation, knowledge_base_id, bedrock_client) -> str:
    # uses embedding model to retrieve vectors
    response = bedrock_client.retrieve(
        retrievalQuery={
          'text': flat_conversation
        },
        knowledgeBaseId=knowledge_base_id,
        retrievalConfiguration={
            'vectorSearchConfiguration': {
                'numberOfResults': knowledgeBaseContextNumberOfResults,
                #'overrideSearchType': "HYBRID", # optional, default is the KB chooses the search type
            }
        },
    )

    return response