Pengetahuan AWS Bedrock - Skrip Ujian Asas

Ini ialah skrip ujian yang mudah tetapi berguna untuk membantu anda menguji dan mengesahkan persediaan KB asas AWS anda dengan cepat. Cuma kemas kini rantau AWS anda jika berbeza, dan palamkan ID KB Berock anda.

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"\nTesting 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("\nRaw Response:")
            print(json.dumps(response, indent=2, default=str))
            # Process and print retrieved results
            print("\nProcessed Results:")
            if 'retrievalResults' in response:
                for i, result in enumerate(response['retrievalResults'], 1):
                    print(f"\nResult {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"\nTesting 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 = "\n".join([
                    f"Reference {i+1}:\n{result.get('content', {}).get('text', '')}\n"
                    for i, result in enumerate(kb_response['retrievalResults'])
                ])
            # Prepare Bedrock prompt
            enhanced_prompt = (
                f"Using the following references:\n\n{context}\n\n"
                f"Please answer this question: {query_text}\n"
                "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"\n\nHuman: {enhanced_prompt}\n\nAssistant:",
                    "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("\nBedrock 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("\nTest 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("\nTest 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()

Atas ialah kandungan terperinci Pengetahuan AWS Bedrock - Skrip Ujian Asas. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!