import boto3 import json import time from datetime import datetime def test_kb_setup():     """Test function to verify Bedrock Knowledge Base setup and queries"""     # Initialize clients     bedrock_agent = boto3.client('bedrock-agent-runtime', region_name='us-east-1')     bedrock_runtime = boto3.client('bedrock-runtime', region_name='us-east-1')     # Your Knowledge Base ID     kb_id = "**your-knowledge-base-id**"  # Replace with your actual KB ID     def test_kb_query(query_text):         """Test a single knowledge base query"""         print(f" Testing query: '{query_text}'")         print("-" * 50)         try:             # Query the knowledge base             response = bedrock_agent.retrieve(                 knowledgeBaseId=kb_id,                 retrievalQuery={'text': query_text},                 retrievalConfiguration={                     'vectorSearchConfiguration': {                         'numberOfResults': 3                     }                 }             )             # Print raw response for debugging             print(" Raw Response:")             print(json.dumps(response, indent=2, default=str))             # Process and print retrieved results             print(" Processed Results:")             if 'retrievalResults' in response:                 for i, result in enumerate(response['retrievalResults'], 1):                     print(f" Result {i}:")                     print(f"Score: {result.get('score', 'N/A')}")                     print(f"Content: {result.get('content', {}).get('text', 'N/A')}")                     print(f"Location: {result.get('location', 'N/A')}")             else:                 print("No results found in response")             return True         except Exception as e:             print(f"Error during query: {str(e)}")             return False     def test_kb_with_bedrock(query_text):         """Test knowledge base integration with Bedrock"""         print(f" Testing KB + Bedrock integration for: '{query_text}'")         print("-" * 50)         try:             # First get KB results             kb_response = bedrock_agent.retrieve(                 knowledgeBaseId=kb_id,                 retrievalQuery={'text': query_text},                 retrievalConfiguration={                     'vectorSearchConfiguration': {                         'numberOfResults': 3                     }                 }             )             # Format context from KB results             context = ""             if 'retrievalResults' in kb_response:                 context = " ".join([                     f"Reference {i+1}: {result.get('content', {}).get('text', '')} "                     for i, result in enumerate(kb_response['retrievalResults'])                 ])             # Prepare Bedrock prompt             enhanced_prompt = (                 f"Using the following references:  {context}  "                 f"Please answer this question: {query_text} "                 "Base your response on the provided references and clearly cite them when used."             )             # Get Bedrock response             bedrock_response = bedrock_runtime.invoke_model(                 modelId="anthropic.claude-v2",                 body=json.dumps({                     "prompt": f"  Human: {enhanced_prompt}  Assistant:",                     "max_tokens_to_sample": 500,                     "temperature": 0.7,                     "top_p": 1,                 }),                 contentType="application/json",                 accept="application/json",             )             response_body = json.loads(bedrock_response.get('body').read())             final_response = response_body.get('completion', '').strip()             print(" Bedrock Response:")             print(final_response)             return True         except Exception as e:             print(f"Error during KB + Bedrock integration: {str(e)}")             return False     # Run test queries     test_queries = [         "What are our company's remote work policies?",         "Tell me about employee benefits",         "What is the vacation policy?",         "How does the performance review process work?",         "What are the working hours?"     ]     print("Starting Knowledge Base Tests")     print("=" * 50)     # Test 1: Basic KB Queries     print(" Test 1: Basic Knowledge Base Queries")     for query in test_queries:         success = test_kb_query(query)         if not success:             print(f"Failed on query: {query}")     # Test 2: KB + Bedrock Integration     print(" Test 2: Knowledge Base + Bedrock Integration")     for query in test_queries:         success = test_kb_with_bedrock(query)         if not success:             print(f"Failed on integration test: {query}") if __name__ == "__main__":     test_kb_setup()