diff --git a/.doc_gen/metadata/bedrock-agent_metadata.yaml b/.doc_gen/metadata/bedrock-agent_metadata.yaml index be0ff03d284..b48dc120f03 100644 --- a/.doc_gen/metadata/bedrock-agent_metadata.yaml +++ b/.doc_gen/metadata/bedrock-agent_metadata.yaml @@ -521,3 +521,68 @@ bedrock-agent_GettingStartedWithBedrockPrompts: services: bedrock-agent: {CreatePrompt, CreatePromptVersion, DeletePrompt} bedrock-runtime: {Converse} + +bedrock-agent_CreateKnowledgeBase: + languages: + Python: + versions: + - sdk_version: 3 + github: python/example_code/bedrock-agent + excerpts: + - description: Create an Amazon Bedrock knowledge base. + snippet_tags: + - python.example_code.bedrock-agent.create_knowledge_base + services: + bedrock-agent: {CreateKnowledgeBase} + +bedrock-agent_GetKnowledgeBase: + languages: + Python: + versions: + - sdk_version: 3 + github: python/example_code/bedrock-agent + excerpts: + - description: Get an Amazon Bedrock knowledge base. + snippet_tags: + - python.example_code.bedrock-agent.get_knowledge_base + services: + bedrock-agent: {GetKnowledgeBase} + +bedrock-agent_UpdateKnowledgeBase: + languages: + Python: + versions: + - sdk_version: 3 + github: python/example_code/bedrock-agent + excerpts: + - description: Update an Amazon Bedrock knowledge base. + snippet_tags: + - python.example_code.bedrock-agent.update_knowledge_base + services: + bedrock-agent: {UpdateKnowledgeBase} + +bedrock-agent_DeleteKnowledgeBase: + languages: + Python: + versions: + - sdk_version: 3 + github: python/example_code/bedrock-agent + excerpts: + - description: Delete an Amazon Bedrock knowledge base. + snippet_tags: + - python.example_code.bedrock-agent.delete_knowledge_base + services: + bedrock-agent: {DeleteKnowledgeBase} + +bedrock-agent_ListKnowledgeBases: + languages: + Python: + versions: + - sdk_version: 3 + github: python/example_code/bedrock-agent + excerpts: + - description: List Amazon Bedrock knowledge Bases. + snippet_tags: + - python.example_code.bedrock-agent.list_knowledge_bases + services: + bedrock-agent: {ListKnowledgeBases} diff --git a/python/example_code/bedrock-agent/README.md b/python/example_code/bedrock-agent/README.md index afc62852977..262f5d581ba 100644 --- a/python/example_code/bedrock-agent/README.md +++ b/python/example_code/bedrock-agent/README.md @@ -44,6 +44,7 @@ Code excerpts that show you how to call individual service functions. - [CreateFlow](flows/flow.py#L18) - [CreateFlowAlias](flows/flow_alias.py#L15) - [CreateFlowVersion](flows/flow_version.py#L18) +- [CreateKnowledgeBase](knowledge_bases/knowledge_base.py#L30) - [CreatePrompt](prompts/prompt.py#L18) - [CreatePromptVersion](prompts/prompt.py#L84) - [DeleteAgent](bedrock_agent_wrapper.py#L118) @@ -51,10 +52,12 @@ Code excerpts that show you how to call individual service functions. - [DeleteFlow](flows/flow.py#L155) - [DeleteFlowAlias](flows/flow_alias.py#L98) - [DeleteFlowVersion](flows/flow_version.py#L91) +- [DeleteKnowledgeBase](knowledge_bases/knowledge_base.py#L169) - [DeletePrompt](prompts/prompt.py#L159) - [GetAgent](bedrock_agent_wrapper.py#L161) - [GetFlow](flows/flow.py#L192) - [GetFlowVersion](flows/flow_version.py#L54) +- [GetKnowledgeBase](knowledge_bases/knowledge_base.py#L90) - [GetPrompt](prompts/prompt.py#L124) - [ListAgentActionGroups](bedrock_agent_wrapper.py#L208) - [ListAgentKnowledgeBases](bedrock_agent_wrapper.py#L237) @@ -62,11 +65,13 @@ Code excerpts that show you how to call individual service functions. - [ListFlowAliases](flows/flow_alias.py#L132) - [ListFlowVersions](flows/flow_version.py#L128) - [ListFlows](flows/flow.py#L229) +- [ListKnowledgeBases](knowledge_bases/knowledge_base.py#L199) - [ListPrompts](prompts/prompt.py#L191) - [PrepareAgent](bedrock_agent_wrapper.py#L266) - [PrepareFlow](flows/flow.py#L58) - [UpdateFlow](flows/flow.py#L112) - [UpdateFlowAlias](flows/flow_alias.py#L55) +- [UpdateKnowledgeBase](knowledge_bases/knowledge_base.py#L120) ### Scenarios @@ -102,6 +107,11 @@ This example shows you how to do the following: - Delete all created resources. + +The flow includes a prompt node that generates a playlist for a chosen genre +and number of songs. The example creates the nodes and permissions +for the flow. + Start the example by running the following at a command prompt: diff --git a/python/example_code/bedrock-agent/knowledge_bases/knowledge_base.py b/python/example_code/bedrock-agent/knowledge_bases/knowledge_base.py new file mode 100644 index 00000000000..f2256c92803 --- /dev/null +++ b/python/example_code/bedrock-agent/knowledge_bases/knowledge_base.py @@ -0,0 +1,456 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Purpose + +Shows how to use the AWS SDK for Python (Boto3) with Amazon Bedrock to work with +knowledge bases in your AWS account. + +This example demonstrates how to: +- Create a Bedrock Agent client +- Get details of a knowledge base +- Update a knowledge base +- List knowledge bases in your account +- Delete a knowledge base +""" + +import argparse +import logging +from pprint import pprint +import boto3 +import uuid +import time +from botocore.exceptions import ClientError +from knowledge_bases.roles import create_knowledge_base_role, delete_knowledge_base_role + +logger = logging.getLogger(__name__) + + +# snippet-start:[python.example_code.bedrock-agent.create_knowledge_base] +def create_knowledge_base(bedrock_agent_client, name, role_arn, description=None): + """ + Creates a new knowledge base. + + Args: + bedrock_agent_client: The Boto3 Bedrock Agent client. + name (str): The name of the knowledge base. + role_arn (str): The ARN of the IAM role that the knowledge base assumes to access resources. + description (str, optional): A description of the knowledge base. + + Returns: + dict: The details of the created knowledge base. + """ + try: + kwargs = { + "name": name, + "roleArn": role_arn, + "knowledgeBaseConfiguration": { + "type": "VECTOR", + "vectorKnowledgeBaseConfiguration": { + "embeddingModelArn": "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1" + } + }, + "storageConfiguration": { + "type": "OPENSEARCH_SERVERLESS", + # Note: You will need to create an OpenSearch Serverless collection first and replace this ARN + # with your actual collection ARN from the OpenSearch console. If you use the console instead, + # you can use the quick-create flow to have Knowledge Bases create the collection for you. + "opensearchServerlessConfiguration": { + "collectionArn": "arn:aws:aoss:us-east-1::123456789012:collection/abcdefgh12345678defgh", + "fieldMapping": { + "metadataField": "metadata", + "textField": "text", + "vectorField": "vector" + }, + "vectorIndexName": "test-uuid" + }, + }, + "clientToken": "test-client-token-" + str(uuid.uuid4()) + } + + if description: + kwargs["description"] = description + + response = bedrock_agent_client.create_knowledge_base(**kwargs) + + logger.info("Created knowledge base with ID: %s", response["knowledgeBase"]["knowledgeBaseId"]) + return response["knowledgeBase"] + + except ClientError as err: + logger.error( + "Couldn't create knowledge base. Here's why: %s: %s", + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + raise +# snippet-end:[python.example_code.bedrock-agent.create_knowledge_base] + + +# snippet-start:[python.example_code.bedrock-agent.get_knowledge_base] +def get_knowledge_base(bedrock_agent_client, knowledge_base_id): + """ + Gets details about a specific knowledge base. + + Args: + bedrock_agent_client: The Boto3 Bedrock Agent client. + knowledge_base_id (str): The ID of the knowledge base. + + Returns: + dict: The details of the knowledge base. + """ + try: + response = bedrock_agent_client.get_knowledge_base( + knowledgeBaseId=knowledge_base_id + ) + + logger.info("Retrieved knowledge base: %s", knowledge_base_id) + return response["knowledgeBase"] + except ClientError as err: + logger.error( + "Couldn't get knowledge base %s. Here's why: %s: %s", + knowledge_base_id, + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + raise +# snippet-end:[python.example_code.bedrock-agent.get_knowledge_base] + + +# snippet-start:[python.example_code.bedrock-agent.update_knowledge_base] +def update_knowledge_base(bedrock_agent_client, knowledge_base_id, name=None, description=None, role_arn=None): + """ + Updates an existing knowledge base. + + Args: + bedrock_agent_client: The Boto3 Bedrock Agent client. + knowledge_base_id (str): The ID of the knowledge base to update. + name (str, optional): The new name for the knowledge base. + description (str, optional): The new description for the knowledge base. + role_arn (str, optional): The new IAM role ARN for the knowledge base. + + Returns: + dict: The details of the updated knowledge base. + """ + try: + kwargs = { + "knowledgeBaseId": knowledge_base_id, + "knowledgeBaseConfiguration": { + "type": "VECTOR", + "vectorKnowledgeBaseConfiguration": { + "embeddingModelArn": "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1" + } + } + } + + if name: + kwargs["name"] = name + if description: + kwargs["description"] = description + if role_arn: + kwargs["roleArn"] = role_arn + + response = bedrock_agent_client.update_knowledge_base(**kwargs) + + logger.info("Updated knowledge base: %s", knowledge_base_id) + return response["knowledgeBase"] + + except ClientError as err: + logger.error( + "Couldn't update knowledge base %s. Here's why: %s: %s", + knowledge_base_id, + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + raise +# snippet-end:[python.example_code.bedrock-agent.update_knowledge_base] + + +# snippet-start:[python.example_code.bedrock-agent.delete_knowledge_base] +def delete_knowledge_base(bedrock_agent_client, knowledge_base_id): + """ + Deletes a knowledge base. + + Args: + bedrock_agent_client: The Boto3 Bedrock Agent client. + knowledge_base_id (str): The ID of the knowledge base to delete. + + Returns: + bool: True if the deletion was successful. + """ + try: + bedrock_agent_client.delete_knowledge_base( + knowledgeBaseId=knowledge_base_id + ) + + logger.info("Deleted knowledge base: %s", knowledge_base_id) + return True + except ClientError as err: + logger.error( + "Couldn't delete knowledge base %s. Here's why: %s: %s", + knowledge_base_id, + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + raise +# snippet-end:[python.example_code.bedrock-agent.delete_knowledge_base] + + +# snippet-start:[python.example_code.bedrock-agent.list_knowledge_bases] +def list_knowledge_bases(bedrock_agent_client, max_results=None): + """ + Lists the knowledge bases in your AWS account. + + Args: + bedrock_agent_client: The Boto3 Bedrock Agent client. + max_results (int, optional): The maximum number of knowledge bases to return. + + Returns: + list: A list of knowledge base details. + """ + try: + kwargs = {} + if max_results is not None: + kwargs["maxResults"] = max_results + + # Initialize an empty list to store all knowledge bases + all_knowledge_bases = [] + + # Use paginator to handle pagination automatically + paginator = bedrock_agent_client.get_paginator('list_knowledge_bases') + page_iterator = paginator.paginate(**kwargs) + + # Iterate through each page of results + for page in page_iterator: + all_knowledge_bases.extend(page.get('knowledgeBaseSummaries', [])) + + logger.info("Found %s knowledge bases.", len(all_knowledge_bases)) + return all_knowledge_bases + except ClientError as err: + logger.error( + "Couldn't list knowledge bases. Here's why: %s: %s", + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + raise +# snippet-end:[python.example_code.bedrock-agent.list_knowledge_bases] + + +def run_knowledge_base_scenario(): + """ + 1. Create an IAM role for the knowledge base + 2. Create a knowledge base + 3. Get details of the knowledge base + 4. Update the knowledge base + 5. Delete the knowledge base and IAM role + """ + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + print("-" * 88) + print("Welcome to the Amazon Bedrock Knowledge Bases scenario.") + print("-" * 88) + + # Create clients + bedrock_agent_client = boto3.client(service_name="bedrock-agent") + iam_client = boto3.client(service_name="iam") + + # Generate unique names for resources + kb_name = "example-knowledge-base-" + str(uuid.uuid4().hex[:8]) + role_name = "example-kb-role" + str(uuid.uuid4().hex[:8]) + + knowledge_base_id = None + + try: + # Step 1: Create IAM role + print("\nCreating IAM role: " + role_name + " ...") + role = create_knowledge_base_role(iam_client, role_name) + role_arn = role['Arn'] + print("Created role with ARN: " + role_arn) + + # Wait for role to propagate + print("Waiting for role to propagate...") + time.sleep(10) + + # Step 2: Create knowledge base + print("Creating knowledge base: " + kb_name + " ...") + kb = create_knowledge_base( + bedrock_agent_client, + kb_name, + role_arn, + "Example knowledge base for demonstration" + ) + knowledge_base_id = kb["knowledgeBaseId"] + print("Created knowledge base with ID: " + knowledge_base_id) + + # Step 3: Get knowledge base details + print("\nGetting details for knowledge base: " + knowledge_base_id + " ...") + kb_details = get_knowledge_base(bedrock_agent_client, knowledge_base_id) + print("Knowledge base details:") + pprint(kb_details) + + # Step 4: Update knowledge base + new_name = kb_name + "-updated" + print("\nUpdating knowledge base name to: " + new_name + " ...") + updated_kb = update_knowledge_base( + bedrock_agent_client, + knowledge_base_id, + new_name, + "Updated description for the knowledge base", + role_arn + ) + print("Updated knowledge base details:") + pprint(updated_kb) + + # Step 5: List knowledge bases + print("\nListing all knowledge bases:") + all_kbs = list_knowledge_bases(bedrock_agent_client) + print("Found " + len(all_kbs) + " knowledge bases.") + + print("\nCleaning up resources...") + if knowledge_base_id: + print("Deleting knowledge base " + knowledge_base_id + " ...") + delete_knowledge_base(bedrock_agent_client, knowledge_base_id) + print("Knowledge base " + knowledge_base_id + " deleted successfully.") + + print("Deleting IAM role " + role_name + " ...") + delete_knowledge_base_role(iam_client, role_name) + print("Role " + role_name + " deleted successfully.") + + print("\nScenario completed successfully!") + + except ClientError as error: + print("Operation failed: " + str(error)) + # Clean up resources on error + if knowledge_base_id: + try: + print("Attempting to delete knowledge base " + knowledge_base_id + " ...") + delete_knowledge_base(bedrock_agent_client, knowledge_base_id) + except Exception as e: + print("Failed to delete knowledge base: " + e) + + try: + print("Attempting to delete IAM role: " + role_name + " ...") + delete_knowledge_base_role(iam_client, role_name) + except Exception as e: + print("Failed to delete IAM role: " + str(e)) + + print("-" * 88) + + +def main(): + """ + Shows how to use the Bedrock Agent API to work with knowledge bases. + """ + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + parser = argparse.ArgumentParser( + description="Work with knowledge bases in your AWS account." + ) + parser.add_argument( + "--action", + choices=["list", "create", "get", "update", "delete", "scenario"], + default="list", + help="The action to perform on knowledge bases.", + ) + parser.add_argument( + "--knowledge-base-id", + help="The ID of the knowledge base (required for get, update, and delete actions).", + ) + parser.add_argument( + "--name", + help="The name of the knowledge base (required for create, optional for update).", + ) + parser.add_argument( + "--role-arn", + help="The ARN of the IAM role for the knowledge base (required for create, optional for update).", + ) + parser.add_argument( + "--description", + help="A description of the knowledge base (optional for create and update).", + ) + parser.add_argument( + "--max-results", + type=int, + help="The maximum number of knowledge bases to return when listing.", + ) + args = parser.parse_args() + + print("-" * 88) + print("Welcome to the Amazon Bedrock Knowledge Bases example.") + print("-" * 88) + + bedrock_agent_client = boto3.client(service_name="bedrock-agent") + + try: + if args.action == "scenario": + run_knowledge_base_scenario() + elif args.action == "list": + print("Listing knowledge bases in your AWS account...") + knowledge_bases = list_knowledge_bases(bedrock_agent_client, args.max_results) + + if knowledge_bases: + print("Found " + len(knowledge_bases) + " knowledge bases") + for kb in knowledge_bases: + print("\n" + "-" * 40) + pprint(kb) + else: + print("No knowledge bases found in your account.") + + elif args.action == "create": + if not args.name or not args.role_arn: + print("Error: --name and --role-arn are required for create action.") + return + + print("Creating knowledge base " + args.name + " ...") + kb = create_knowledge_base( + bedrock_agent_client, + args.name, + args.role_arn, + args.description + ) + print("Knowledge base created successfully:") + pprint(kb) + + elif args.action == "get": + if not args.knowledge_base_id: + print("Error: --knowledge-base-id is required for get action.") + return + + print("Getting details for knowledge base " + args.knowledge_base_id + " ...") + kb = get_knowledge_base(bedrock_agent_client, args.knowledge_base_id) + print("Knowledge base details:") + pprint(kb) + + elif args.action == "update": + if not args.knowledge_base_id: + print("Error: --knowledge-base-id is required for update action.") + return + + print("Updating knowledge base " + args.knowledge_base_id + " ...") + kb = update_knowledge_base( + bedrock_agent_client, + args.knowledge_base_id, + args.name, + args.description, + args.role_arn + ) + print("Knowledge base updated successfully:") + pprint(kb) + + elif args.action == "delete": + if not args.knowledge_base_id: + print("Error: --knowledge-base-id is required for delete action.") + return + + print("Deleting knowledge base " + args.knowledge_base_id + " ...") + if delete_knowledge_base(bedrock_agent_client, args.knowledge_base_id): + print("Knowledge base " + args.knowledge_base_id + " deleted successfully.") + + except ClientError as error: + print("Operation failed: " + error) + + print("-" * 88) + + +if __name__ == "__main__": + main() diff --git a/python/example_code/bedrock-agent/knowledge_bases/roles.py b/python/example_code/bedrock-agent/knowledge_bases/roles.py new file mode 100644 index 00000000000..5916931425a --- /dev/null +++ b/python/example_code/bedrock-agent/knowledge_bases/roles.py @@ -0,0 +1,164 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +IAM Role Management for Amazon Bedrock Knowledge Bases + +This module provides functionality to create, update, and delete IAM roles specifically +configured for Amazon Bedrock knowledge bases. It handles the complete lifecycle of IAM roles +including trust relationships, inline policies, and permissions management. +""" +import logging +import json +from botocore.exceptions import ClientError + +logging.basicConfig( + level=logging.INFO +) +logger = logging.getLogger(__name__) + +# snippet-start:[python.example_code.bedrock-agent.knowledge_base_iam_role] + +def create_knowledge_base_role(client, role_name): + """ + Creates an IAM role for Amazon Bedrock with permissions to manage knowledge bases. + + Args: + client: The Boto3 IAM client. + role_name (str): Name for the new IAM role. + Returns: + dict: The created role information including ARN. + """ + + # Trust relationship policy - allows Amazon Bedrock service to assume this role. + trust_policy = { + "Version": "2012-10-17", + "Statement": [{ + "Effect": "Allow", + "Principal": { + "Service": "bedrock.amazonaws.com" + }, + "Action": "sts:AssumeRole" + }] + } + + # Basic inline policy for knowledge base operations + knowledge_base_policy = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "bedrock:InvokeModel", + "bedrock:Retrieve", + "bedrock:RetrieveAndGenerate", + "aoss:APIAccessAll", + "aoss:CreateIndex", + "aoss:DeleteIndex", + "aoss:UpdateIndex", + "aoss:BatchGetCollection", + "aoss:CreateCollection", + "aoss:DeleteCollection" + ], + # Using * as placeholder - Later you can update with specific ARNs. + "Resource": "*" + } + ] + } + + try: + # Create the IAM role with trust policy + logging.info("Creating role: %s", role_name) + role = client.create_role( + RoleName=role_name, + AssumeRolePolicyDocument=json.dumps(trust_policy), + Description="Role for Amazon Bedrock knowledge base operations" + ) + + # Attach inline policy to the role + print("Attaching inline policy") + client.put_role_policy( + RoleName=role_name, + PolicyName=f"{role_name}-policy", + PolicyDocument=json.dumps(knowledge_base_policy) + ) + + logging.info("Created Role ARN: %s", role['Role']['Arn']) + return role['Role'] + + except ClientError as e: + logging.warning("Error creating role: %s", str(e)) + raise + except Exception as e: + logging.warning("Unexpected error: %s", str(e)) + raise + + +def update_role_policy(client, role_name, resource_arns): + """ + Updates an IAM role's inline policy with specific resource ARNs. + + Args: + client: The Boto3 IAM client. + role_name (str): Name of the existing role. + resource_arns (list): List of resource ARNs to allow access to. + """ + + updated_policy = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "bedrock:InvokeModel", + "bedrock:Retrieve", + "bedrock:RetrieveAndGenerate", + "aoss:APIAccessAll", + "aoss:CreateIndex", + "aoss:DeleteIndex", + "aoss:UpdateIndex", + "aoss:BatchGetCollection", + "aoss:CreateCollection", + "aoss:DeleteCollection" + ], + "Resource": resource_arns + } + ] + } + + try: + client.put_role_policy( + RoleName=role_name, + PolicyName=f"{role_name}-policy", + PolicyDocument=json.dumps(updated_policy) + ) + logging.info("Updated policy for role: %s", role_name) + + except ClientError as e: + logging.warning("Error updating role policy: %s", str(e)) + raise + + +def delete_knowledge_base_role(client, role_name): + """ + Deletes an IAM role. + + Args: + client: The Boto3 IAM client. + role_name (str): Name of the role to delete. + """ + try: + # Detach and delete inline policies + policies = client.list_role_policies(RoleName=role_name)['PolicyNames'] + for policy_name in policies: + client.delete_role_policy(RoleName=role_name, PolicyName=policy_name) + + # Delete the role + client.delete_role(RoleName=role_name) + logging.info("Deleted role: %s", role_name) + + except ClientError as e: + logging.info("Error Deleting role: %s", str(e)) + raise + +# snippet-end:[python.example_code.bedrock-agent.knowledge_base_iam_role] diff --git a/python/example_code/bedrock-agent/knowledge_bases/scenario_get_started_with_knowledge_bases.py b/python/example_code/bedrock-agent/knowledge_bases/scenario_get_started_with_knowledge_bases.py new file mode 100644 index 00000000000..08f1976b9d0 --- /dev/null +++ b/python/example_code/bedrock-agent/knowledge_bases/scenario_get_started_with_knowledge_bases.py @@ -0,0 +1,30 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Purpose + +Shows how to use the AWS SDK for Python (Boto3) with Amazon Bedrock to run a complete +knowledge base scenario. + +This example demonstrates how to: +- Create an IAM role for the knowledge base +- Create a knowledge base +- Get details of a knowledge base +- Update a knowledge base +- List knowledge bases in your account +- Delete the knowledge base and IAM role +""" + +import logging +from knowledge_base import run_knowledge_base_scenario + +def main(): + """ + Runs the Amazon Bedrock Knowledge Bases scenario. + """ + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + run_knowledge_base_scenario() + +if __name__ == "__main__": + main() diff --git a/python/example_code/bedrock-agent/test/conftest.py b/python/example_code/bedrock-agent/test/conftest.py index 5c47f19f061..bdc0fa16aed 100644 --- a/python/example_code/bedrock-agent/test/conftest.py +++ b/python/example_code/bedrock-agent/test/conftest.py @@ -90,4 +90,40 @@ class FakePromptData: "genre": "pop", "number": "1" } - OUTPUT_TEXT ="Here's a playlist with one song" \ No newline at end of file + OUTPUT_TEXT ="Here's a playlist with one song" + + +class FakeKnowledgeBaseData: + """Test data for knowledge base tests.""" + # Define fake knowledge base IDs for testing + KB_ID_1 = "FAKE_KB_ID_1" + KB_ID_2 = "FAKE_KB_ID_2" + + # Define fake knowledge base names for testing + KB_NAME_1 = "FakeKnowledgeBase1" + KB_NAME_2 = "FakeKnowledgeBase2" + + # Define fake ARNs for the knowledge bases + KB_ARN_1 = f"arn:aws:bedrock:us-east-1:123456789012:knowledge-base/{KB_ID_1}" + KB_ARN_2 = f"arn:aws:bedrock:us-east-1:123456789012:knowledge-base/{KB_ID_2}" + + # Other common attributes for knowledge bases + DESCRIPTION = "A fake knowledge base for testing." + CREATED_AT = "2025-03-29T21:34:43.048609+00:00" + UPDATED_AT = "2025-03-30T21:34:43.048609+00:00" + STATUS = "ACTIVE" + + # Data source configuration + DATA_SOURCE_ID = "FAKE_DATA_SOURCE_ID" + DATA_SOURCE_NAME = "FakeDataSource" + DATA_SOURCE_TYPE = "S3" + S3_CONFIGURATION = { + "bucketName": "fake-kb-bucket", + "inclusionPrefixes": ["documents/"] + } + + # Vector store configuration + EMBEDDING_MODEL_ARN = "arn:aws:bedrock:us-east-1:123456789012:embedding-model/amazon.titan-embed-text-v1" + VECTOR_STORE_CONFIGURATION = { + "embeddingModelArn": EMBEDDING_MODEL_ARN + } diff --git a/python/example_code/bedrock-agent/test/test_knowledge_base.py b/python/example_code/bedrock-agent/test/test_knowledge_base.py new file mode 100644 index 00000000000..ff67e5e715d --- /dev/null +++ b/python/example_code/bedrock-agent/test/test_knowledge_base.py @@ -0,0 +1,498 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Unit tests for knowledge_base.py. + +This file contains comprehensive tests for the Amazon Bedrock Knowledge Base API operations: +1. CreateKnowledgeBase - Tests creating a knowledge base with valid parameters and error handling +2. GetKnowledgeBase - Tests retrieving details of a specific knowledge base and error handling +3. UpdateKnowledgeBase - Tests updating a knowledge base with new name, description, and role ARN +4. DeleteKnowledgeBase - Tests deleting a knowledge base and error handling +5. ListKnowledgeBases - Tests listing knowledge bases in the account and error handling + +Each test function: +- Creates a Bedrock Agent client and stubber +- Sets up test data using the FakeKnowledgeBaseData class from conftest.py +- Defines expected parameters and mock responses +- Configures the stubber with these parameters and responses +- Tests both success and error cases using pytest's parametrize feature +- Verifies the results match the expected values + +These tests follow the same pattern as the existing tests in the codebase, +ensuring consistency with the project's testing approach. +""" + +import boto3 +from botocore.exceptions import ClientError +import pytest +import uuid +from unittest.mock import patch, MagicMock + +from knowledge_bases.knowledge_base import ( + create_knowledge_base, + get_knowledge_base, + update_knowledge_base, + delete_knowledge_base, + list_knowledge_bases +) + +from conftest import FakeKnowledgeBaseData + +@pytest.mark.parametrize("error_code", [None, "TestException"]) +def test_create_knowledge_base(make_stubber, error_code): + """ + Test the create_knowledge_base function. + + This test verifies that: + 1. The function correctly calls the CreateKnowledgeBase API + 2. The function properly processes and returns the knowledge base details + 3. Error handling works as expected when the API call fails + + Parameters: + make_stubber: Fixture that creates a stubber for the Bedrock Agent client + error_code: Simulated error code (None for success case, string for error case) + """ + # Create a Bedrock Agent client and stubber + bedrock_agent_client = boto3.client("bedrock-agent") + bedrock_agent_stubber = make_stubber(bedrock_agent_client) + + # Create test data + kb_data = FakeKnowledgeBaseData() + name = kb_data.KB_NAME_1 + role_arn = "arn:aws:iam::123456789012:role/KnowledgeBaseRole" + description = kb_data.DESCRIPTION + + # Mock uuid.uuid4 to return a consistent value for testing + with patch('uuid.uuid4', return_value=MagicMock(return_value="test-uuid")): + + # Define the expected parameters for the API call + expected_params = { + "name": name, + "roleArn": role_arn, + "knowledgeBaseConfiguration": { + "type": "VECTOR", + "vectorKnowledgeBaseConfiguration": { + "embeddingModelArn": "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1" + }, + }, + "storageConfiguration": { + "type": "OPENSEARCH_SERVERLESS", + "opensearchServerlessConfiguration": { + "collectionArn": "arn:aws:aoss:us-east-1::123456789012:collection/abcdefgh12345678defgh", + "fieldMapping": { + "metadataField": "metadata", + "textField": "text", + "vectorField": "vector" + }, + "vectorIndexName": "test-uuid" + } + }, + # Add a longer client token that meets minimum length requirement + "clientToken": "test-client-token-" + str(uuid.uuid4()) + } + + if description: + expected_params["description"] = description + + # Define the mock response that the API would return + response = { + "knowledgeBase": { + "knowledgeBaseId": kb_data.KB_ID_1, + "name": name, + "description": description, + "roleArn": role_arn, + "knowledgeBaseArn": kb_data.KB_ARN_1, + "status": kb_data.STATUS, + "createdAt": kb_data.CREATED_AT, + "updatedAt": kb_data.UPDATED_AT, + "knowledgeBaseConfiguration": { + "type": "VECTOR", + "vectorKnowledgeBaseConfiguration": { + "embeddingModelArn": "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1" + } + } + } + } + + # Configure the stubber with the expected parameters and response + bedrock_agent_stubber.stub_create_knowledge_base( + expected_params, response, error_code=error_code + ) + + if error_code is None: + # Test the success case + result = create_knowledge_base(bedrock_agent_client, name, role_arn, description) + + # Verify the results contain the expected knowledge base details + assert result["knowledgeBaseId"] == kb_data.KB_ID_1 + assert result["name"] == name + assert result["description"] == description + assert result["roleArn"] == role_arn + assert result["knowledgeBaseArn"] == kb_data.KB_ARN_1 + else: + # Test the error case + with pytest.raises(ClientError) as exc_info: + create_knowledge_base(bedrock_agent_client, name, role_arn, description) + + # Verify the error code matches what we expected + assert exc_info.value.response["Error"]["Code"] == error_code + +@pytest.mark.parametrize("error_code", [None, "TestException"]) +def test_get_knowledge_base(make_stubber, error_code): + """ + Test the get_knowledge_base function. + + This test verifies that: + 1. The function correctly calls the GetKnowledgeBase API + 2. The function properly processes and returns the knowledge base details + 3. Error handling works as expected when the API call fails + + Parameters: + make_stubber: Fixture that creates a stubber for the Bedrock Agent client + error_code: Simulated error code (None for success case, string for error case) + """ + # Create a Bedrock Agent client and stubber + bedrock_agent_client = boto3.client("bedrock-agent") + bedrock_agent_stubber = make_stubber(bedrock_agent_client) + + # Create test data + kb_data = FakeKnowledgeBaseData() + knowledge_base_id = kb_data.KB_ID_1 + + # Define the expected parameters for the API call + expected_params = { + "knowledgeBaseId": knowledge_base_id + } + + # Define the mock response that the API would return + response = { + "knowledgeBase": { + "knowledgeBaseId": knowledge_base_id, + "name": kb_data.KB_NAME_1, + "description": kb_data.DESCRIPTION, + "roleArn": "arn:aws:iam::123456789012:role/KnowledgeBaseRole", + "knowledgeBaseArn": kb_data.KB_ARN_1, + "status": kb_data.STATUS, + "createdAt": kb_data.CREATED_AT, + "updatedAt": kb_data.UPDATED_AT, + "knowledgeBaseConfiguration": { + "type": "VECTOR", + "vectorKnowledgeBaseConfiguration": { + "embeddingModelArn": kb_data.EMBEDDING_MODEL_ARN + } + } + } + } + + # Configure the stubber with the expected parameters and response + bedrock_agent_stubber.stub_get_knowledge_base( + expected_params, response, error_code=error_code + ) + + if error_code is None: + # Test the success case + result = get_knowledge_base(bedrock_agent_client, knowledge_base_id) + + # Verify the results contain the expected knowledge base details + assert result["knowledgeBaseId"] == knowledge_base_id + assert result["name"] == kb_data.KB_NAME_1 + assert result["description"] == kb_data.DESCRIPTION + assert result["knowledgeBaseArn"] == kb_data.KB_ARN_1 + assert result["status"] == kb_data.STATUS + else: + # Test the error case + with pytest.raises(ClientError) as exc_info: + get_knowledge_base(bedrock_agent_client, knowledge_base_id) + + # Verify the error code matches what we expected + assert exc_info.value.response["Error"]["Code"] == error_code + + +@pytest.mark.parametrize("error_code", [None, "TestException"]) +def test_update_knowledge_base(make_stubber, error_code): + """ + Test the update_knowledge_base function. + + This test verifies that: + 1. The function correctly calls the UpdateKnowledgeBase API + 2. The function properly processes and returns the updated knowledge base details + 3. Error handling works as expected when the API call fails + + Parameters: + make_stubber: Fixture that creates a stubber for the Bedrock Agent client + error_code: Simulated error code (None for success case, string for error case) + """ + # Create a Bedrock Agent client and stubber + bedrock_agent_client = boto3.client("bedrock-agent") + bedrock_agent_stubber = make_stubber(bedrock_agent_client) + + # Create test data + kb_data = FakeKnowledgeBaseData() + knowledge_base_id = kb_data.KB_ID_1 + name = "Updated KB Name" + description = "Updated description" + role_arn = "arn:aws:iam::123456789012:role/UpdatedKnowledgeBaseRole" + + # Mock uuid.uuid4 to return a consistent value for testing + with patch('uuid.uuid4', return_value=MagicMock(return_value="test-uuid")): + # Define the expected parameters for the API call + expected_params = { + "knowledgeBaseId": knowledge_base_id, + "name": name, + "description": description, + "roleArn": role_arn, + "knowledgeBaseConfiguration": { + "type": "VECTOR", + "vectorKnowledgeBaseConfiguration": { + "embeddingModelArn": "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1" + } + } + } + + # Define the mock response that the API would return + response = { + "knowledgeBase": { + "knowledgeBaseId": knowledge_base_id, + "name": name, + "description": description, + "roleArn": role_arn, + "knowledgeBaseArn": kb_data.KB_ARN_1, + "knowledgeBaseConfiguration": { + "type": "VECTOR", + "vectorKnowledgeBaseConfiguration": { + "embeddingModelArn": "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1" + } + }, + "status": kb_data.STATUS, + "createdAt": kb_data.CREATED_AT, + "updatedAt": kb_data.UPDATED_AT, + + } + } + + # Configure the stubber with the expected parameters and response + bedrock_agent_stubber.stub_update_knowledge_base( + expected_params, response, error_code=error_code + ) + + if error_code is None: + # Test the success case + result = update_knowledge_base( + bedrock_agent_client, + knowledge_base_id, + name=name, + description=description, + role_arn=role_arn + ) + + # Verify the results contain the expected knowledge base details + assert result["knowledgeBaseId"] == knowledge_base_id + assert result["name"] == name + assert result["description"] == description + assert result["roleArn"] == role_arn + assert result["knowledgeBaseArn"] == kb_data.KB_ARN_1 + else: + # Test the error case + with pytest.raises(ClientError) as exc_info: + update_knowledge_base( + bedrock_agent_client, + knowledge_base_id, + name=name, + description=description, + role_arn=role_arn + ) + + # Verify the error code matches what we expected + assert exc_info.value.response["Error"]["Code"] == error_code# make_stubber: Fixture that creates a stubber for the Bedrock Agent client +# error_code: Simulated error code (None for success case, string for error case) +# """ +# # Create a Bedrock Agent client and stubber +# bedrock_agent_client = boto3.client("bedrock-agent") +# bedrock_agent_stubber = make_stubber(bedrock_agent_client) + +# # Create test data +# kb_data = FakeKnowledgeBaseData() +# knowledge_base_id = kb_data.KB_ID_1 +# name = "Updated KB Name" +# description = "Updated description" +# role_arn = "arn:aws:iam::123456789012:role/UpdatedKnowledgeBaseRole" + +# # Mock uuid.uuid4 to return a consistent value for testing +# with patch('uuid.uuid4', return_value=MagicMock(return_value="test-uuid")): +# # Define the expected parameters for the API call +# expected_params = { +# "knowledgeBaseId": knowledge_base_id, +# "name": name, +# "description": description, +# "roleArn": role_arn, +# "clientToken": str(uuid.uuid4()) +# } + +# # Define the mock response that the API would return +# response = { +# "knowledgeBase": { +# "knowledgeBaseId": knowledge_base_id, +# "name": name, +# "description": description, +# "roleArn": role_arn, +# "knowledgeBaseArn": kb_data.KB_ARN_1, +# "status": kb_data.STATUS, +# "createdAt": kb_data.CREATED_AT, +# "updatedAt": kb_data.UPDATED_AT +# } +# } + +# # Configure the stubber with the expected parameters and response +# bedrock_agent_stubber.stub_update_knowledge_base( +# expected_params, response, error_code=error_code +# ) + +# if error_code is None: +# # Test the success case +# result = update_knowledge_base( +# bedrock_agent_client, +# knowledge_base_id, +# name=name, +# description=description, +# role_arn=role_arn +# ) + +# # Verify the results contain the expected knowledge base details +# assert result["knowledgeBaseId"] == knowledge_base_id +# assert result["name"] == name +# assert result["description"] == description +# assert result["roleArn"] == role_arn +# assert result["knowledgeBaseArn"] == kb_data.KB_ARN_1 +# else: +# # Test the error case +# with pytest.raises(ClientError) as exc_info: +# update_knowledge_base( +# bedrock_agent_client, +# knowledge_base_id, +# name=name, +# description=description, +# role_arn=role_arn +# ) + +# # Verify the error code matches what we expected +# assert exc_info.value.response["Error"]["Code"] == error_code + + +@pytest.mark.parametrize("error_code", [None, "TestException"]) +def test_delete_knowledge_base(make_stubber, error_code): + """ + Test the delete_knowledge_base function. + + This test verifies that: + 1. The function correctly calls the DeleteKnowledgeBase API + 2. The function properly processes the response + 3. Error handling works as expected when the API call fails + + Parameters: + make_stubber: Fixture that creates a stubber for the Bedrock Agent client + error_code: Simulated error code (None for success case, string for error case) + """ + # Create a Bedrock Agent client and stubber + bedrock_agent_client = boto3.client("bedrock-agent") + bedrock_agent_stubber = make_stubber(bedrock_agent_client) + + # Create test data + kb_data = FakeKnowledgeBaseData() + knowledge_base_id = kb_data.KB_ID_1 + + # Define the expected parameters for the API call + expected_params = { + "knowledgeBaseId": knowledge_base_id + } + + # Define the mock response that the API would return (empty for delete operation) + response = { + "knowledgeBaseId": knowledge_base_id, + "status": "Deleting" + } + + # Configure the stubber with the expected parameters and response + bedrock_agent_stubber.stub_delete_knowledge_base( + expected_params, response, error_code=error_code + ) + + if error_code is None: + # Test the success case + result = delete_knowledge_base(bedrock_agent_client, knowledge_base_id) + + # Verify the function returns True on success + assert result is True + else: + # Test the error case + with pytest.raises(ClientError) as exc_info: + delete_knowledge_base(bedrock_agent_client, knowledge_base_id) + + # Verify the error code matches what we expected + assert exc_info.value.response["Error"]["Code"] == error_code + + +@pytest.mark.parametrize("error_code", [None, "TestException"]) +def test_list_knowledge_bases(make_stubber, error_code): + """ + Test the list_knowledge_bases function. + + This test verifies that: + 1. The function correctly calls the ListKnowledgeBases API + 2. The function properly processes and returns the knowledge base summaries + 3. Error handling works as expected when the API call fails + + Parameters: + make_stubber: Fixture that creates a stubber for the Bedrock Agent client + error_code: Simulated error code (None for success case, string for error case) + """ + # Create a Bedrock Agent client and stubber + bedrock_agent_client = boto3.client("bedrock-agent") + bedrock_agent_stubber = make_stubber(bedrock_agent_client) + + # Create test data for knowledge bases + kb_data = FakeKnowledgeBaseData() + + # Define the expected parameters for the API call (empty for basic list) + expected_params = {} + + # Define the mock response that the API would return + response = { + "knowledgeBaseSummaries": [ + { + "knowledgeBaseId": kb_data.KB_ID_1, + "name": kb_data.KB_NAME_1, + "description": kb_data.DESCRIPTION, + "status": kb_data.STATUS, + "updatedAt": kb_data.UPDATED_AT, + }, + { + "knowledgeBaseId": kb_data.KB_ID_2, + "name": kb_data.KB_NAME_2, + "description": kb_data.DESCRIPTION, + "status": kb_data.STATUS, + "updatedAt": kb_data.UPDATED_AT, + } + ] + } + + # Configure the stubber with the expected parameters and response + bedrock_agent_stubber.stub_list_knowledge_bases( + expected_params, response, error_code=error_code + ) + + if error_code is None: + # Test the success case + result = list_knowledge_bases(bedrock_agent_client) + + # Verify the results contain the expected knowledge bases + assert len(result) == 2 + assert result[0]["knowledgeBaseId"] == kb_data.KB_ID_1 + assert result[1]["knowledgeBaseId"] == kb_data.KB_ID_2 + else: + # Test the error case + with pytest.raises(ClientError) as exc_info: + list_knowledge_bases(bedrock_agent_client) + + # Verify the error code matches what we expected + assert exc_info.value.response["Error"]["Code"] == error_code diff --git a/python/test_tools/bedrock_agent_stubber.py b/python/test_tools/bedrock_agent_stubber.py index 11f37987851..cf9e732b429 100644 --- a/python/test_tools/bedrock_agent_stubber.py +++ b/python/test_tools/bedrock_agent_stubber.py @@ -175,3 +175,63 @@ def stub_list_prompts(self, expected_params, response, error_code=None): self._stub_bifurcator( "list_prompts", expected_params, response, error_code=error_code ) + + def stub_list_knowledge_bases(self, expected_params, response, error_code=None): + """ + Adds a response for the ListKnowledgeBases operation. + + :param expected_params: The parameters that are expected in the request. + :param response: The response to return when the expected parameters are seen. + :param error_code: The error code to return, if any. + """ + self._stub_bifurcator( + "list_knowledge_bases", expected_params, response, error_code=error_code + ) + + def stub_create_knowledge_base(self, expected_params, response, error_code=None): + """ + Adds a response for the CreateKnowledgeBase operation. + + :param expected_params: The parameters that are expected in the request. + :param response: The response to return when the expected parameters are seen. + :param error_code: The error code to return, if any. + """ + self._stub_bifurcator( + "create_knowledge_base", expected_params, response, error_code=error_code + ) + + def stub_get_knowledge_base(self, expected_params, response, error_code=None): + """ + Adds a response for the GetKnowledgeBase operation. + + :param expected_params: The parameters that are expected in the request. + :param response: The response to return when the expected parameters are seen. + :param error_code: The error code to return, if any. + """ + self._stub_bifurcator( + "get_knowledge_base", expected_params, response, error_code=error_code + ) + + def stub_update_knowledge_base(self, expected_params, response, error_code=None): + """ + Adds a response for the UpdateKnowledgeBase operation. + + :param expected_params: The parameters that are expected in the request. + :param response: The response to return when the expected parameters are seen. + :param error_code: The error code to return, if any. + """ + self._stub_bifurcator( + "update_knowledge_base", expected_params, response, error_code=error_code + ) + + def stub_delete_knowledge_base(self, expected_params, response, error_code=None): + """ + Adds a response for the DeleteKnowledgeBase operation. + + :param expected_params: The parameters that are expected in the request. + :param response: The response to return when the expected parameters are seen. + :param error_code: The error code to return, if any. + """ + self._stub_bifurcator( + "delete_knowledge_base", expected_params, response, error_code=error_code + )