ask_bedrock_llm_with_knowledge_base_updated.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:
    # Structure response
    kb_responses = [
        {
            "text": result['content']['text'],
            "url": result['location']['confluenceLocation']['url']
        } for result in kb_response['retrievalResults']
    ]
    if enable_rerank:
        # Rerank the knowledge base results
        kb_responses = rerank_text(
            flat_conversation,
            kb_responses,
            bedrock_client
        )
    return kb_responses