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-runtime/README.md b/python/example_code/bedrock-agent-runtime/README.md index f0657454c6b..9bc0ffdd17d 100644 --- a/python/example_code/bedrock-agent-runtime/README.md +++ b/python/example_code/bedrock-agent-runtime/README.md @@ -48,6 +48,13 @@ Code excerpts that show you how to call individual service functions. - [InvokeAgent](bedrock_agent_runtime_wrapper.py#L33) - [InvokeFlow](bedrock_agent_runtime_wrapper.py#L71) +### Scenarios + +Code examples that show you how to accomplish a specific task by calling multiple +functions within the same service. + +- [Create and invoke a flow](../bedrock-agent/flows/playlist_flow.py) + @@ -80,6 +87,29 @@ python flows/flow-conversation.py +#### Create and invoke a flow + +This example shows you how to do the following: + +- Create an execution role for the flow. +- Create the flow. +- Deploy the fully configured flow. +- Invoke the flow with user-provided prompts. +- Delete all created resources. + + + + +Start the example by running the following at a command prompt: + +``` +python ../bedrock-agent/flows/playlist_flow.py +``` + + + + + ### Tests ⚠ Running tests might result in charges to your AWS account. diff --git a/python/example_code/bedrock-agent/README.md b/python/example_code/bedrock-agent/README.md index 9c9068056c2..c72d2a0326c 100644 --- a/python/example_code/bedrock-agent/README.md +++ b/python/example_code/bedrock-agent/README.md @@ -2,15 +2,13 @@ ## Overview -Shows how to use the AWS SDK for Python (Boto3) to work with Amazon Bedrock agents, flows, and managed prompts. +Shows how to use the AWS SDK for Python (Boto3) to work with Amazon Bedrock Agents. _Amazon Bedrock Agents offer you the ability to build and configure autonomous agents in your application._ - - ## ⚠ Important * Running this code might result in charges to your AWS account. For more details, see [AWS Pricing](https://aws.amazon.com/pricing/) and [Free Tier](https://aws.amazon.com/free/). @@ -43,46 +41,46 @@ Code excerpts that show you how to call individual service functions. - [CreateAgent](bedrock_agent_wrapper.py#L32) - [CreateAgentActionGroup](bedrock_agent_wrapper.py#L61) - [CreateAgentAlias](bedrock_agent_wrapper.py#L96) +- [CreateFlow](flows/flow.py#L18) +- [CreateFlowAlias](flows/flow_alias.py#L15) +- [CreateFlowVersion](flows/flow_version.py#L18) +- [CreateKnowledgeBase](knowledge_bases/knowledge_base.py#L29) +- [CreatePrompt](prompts/prompt.py#L18) +- [CreatePromptVersion](prompts/prompt.py#L84) - [DeleteAgent](bedrock_agent_wrapper.py#L118) - [DeleteAgentAlias](bedrock_agent_wrapper.py#L139) +- [DeleteFlow](flows/flow.py#L155) +- [DeleteFlowAlias](flows/flow_alias.py#L98) +- [DeleteFlowVersion](flows/flow_version.py#L91) +- [DeleteKnowledgeBase](knowledge_bases/knowledge_base.py#L165) +- [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#L86) +- [GetPrompt](prompts/prompt.py#L124) - [ListAgentActionGroups](bedrock_agent_wrapper.py#L208) - [ListAgentKnowledgeBases](bedrock_agent_wrapper.py#L237) - [ListAgents](bedrock_agent_wrapper.py#L185) +- [ListFlowAliases](flows/flow_alias.py#L132) +- [ListFlowVersions](flows/flow_version.py#L128) +- [ListFlows](flows/flow.py#L229) +- [ListKnowledgeBases](knowledge_bases/knowledge_base.py#L195) +- [ListPrompts](prompts/prompt.py#L191) - [PrepareAgent](bedrock_agent_wrapper.py#L266) -- [InvokeFlow](flows/run_flow.py#L23) -- [CreateFlow](flows/flow.py#L18) -- [PrepareFlow](flows/flow.py#L58) +- [PrepareFlow](flows/flow.py#L58) - [UpdateFlow](flows/flow.py#L112) -- [DeleteFlow](flows/flow.py#L156) -- [GetFlow](flows/flow.py#L192) -- [ListFlows](flows/flow.py#L229) -- [CreateFlowVersion](flows/flow_version.py#L18) -- [GetFlowVersion](flows/flow_version.py#L54) -- [DeleteFlowVersion](flows/flow_version.py#L91) -- [ListFlowVersions](flows/flow_version.py#L128) -- [CreateFlowAlias](flows/flow_alias.py#L15) -- [UpdateFlowAlias](flows/flow_alias.py#L55) -- [DeleteFlowAlias](flows/flow_alias.py#L98) -- [ListFlowAliases](flows/flow_alias.py#L132) -- [FlowConversation](flows/flow-conversation.py) -- [CreatePrompt](prompts/prompt.py#L32) -- [CreatePromptVersion](prompts/prompt.py#L61) -- [GetPrompt](prompts/prompt.py#L71) -- [DeletePrompt](prompts/prompt.py#L154) -- [ListPrompts](prompts/prompt.py#L187) - - - +- [UpdateFlowAlias](flows/flow_alias.py#L55) +- [UpdateKnowledgeBase](knowledge_bases/knowledge_base.py#L116) ### Scenarios Code examples that show you how to accomplish a specific task by calling multiple functions within the same service. -- [Create and invoke an agent](scenario_get_started_with_agents.py) - [Create and invoke a flow](flows/playlist_flow.py) - [Create and invoke a managed prompt](prompts/scenario_get_started_with_prompts.py) +- [Create and invoke an agent](scenario_get_started_with_agents.py) @@ -98,63 +96,59 @@ functions within the same service. -#### Create and invoke an agent +#### Create and invoke a flow This example shows you how to do the following: -- Create an execution role for the agent. -- Create the agent and deploy a DRAFT version. -- Create a Lambda function that implements the agent's capabilities. -- Create an action group that connects the agent to the Lambda function. -- Deploy the fully configured agent. -- Invoke the agent with user-provided prompts. +- Create an execution role for the flow. +- Create the flow. +- Deploy the fully configured flow. +- Invoke the flow with user-provided prompts. - Delete all created resources. - - - -Start the example by running the following at a command prompt: - -``` -python scenario_get_started_with_agents.py -``` - - - - - + -#### Create and invoke a flow - -Shows how to create a simple flow that generates music playlists. 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: ``` python flows/playlist_flow.py ``` + + + + When prompted, enter the genre of music and the number of songs you want in the playlist. Optionally, the script can delete the resources that it creates. If you want to use the flow later, such as in the Amazon Bedrock console, enter `n` when the script prompts you to delete resources. Note that you will then need to manually delete the resources. + +#### Create and invoke a managed prompt -#### List flows +This example shows you how to do the following: -Shows how to List Amazon Bedrock flows, versions of a flow, and aliases of a flow. +- Create a managed prompt. +- Create a version of the prompt. +- Invoke the prompt using the version. +- Clean up resources (optional). + + + Start the example by running the following at a command prompt: ``` -python flows/list_flows.py +python prompts/scenario_get_started_with_prompts.py ``` -The example first lists the flows in the current AWS Region. It -then prompts for a flow ID, which you can get from the list of flows. Finally, the example lists the flow versions and flow aliases for the flow ID that you entered. -#### Create and invoke a managed prompt + + This example shows you how to do the following: @@ -184,18 +178,32 @@ You can also specify a different AWS region or model ID: python prompts/scenario_get_started_with_prompts.py --region us-west-2 --model-id anthropic.claude-3-sonnet-20240229-v1:0 ``` -#### List prompts + -Shows how to list Amazon Bedrock managed prompts and versions of a prompt. +#### Create and invoke an agent + +This example shows you how to do the following: + +- Create an execution role for the agent. +- Create the agent and deploy a DRAFT version. +- Create a Lambda function that implements the agent's capabilities. +- Create an action group that connects the agent to the Lambda function. +- Deploy the fully configured agent. +- Invoke the agent with user-provided prompts. +- Delete all created resources. + + + Start the example by running the following at a command prompt: ``` -python prompts/list_prompts.py +python scenario_get_started_with_agents.py ``` -The example first lists the prompts in the current AWS Region. + + ### Tests 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..486e8b663e9 --- /dev/null +++ b/python/example_code/bedrock-agent/knowledge_bases/knowledge_base.py @@ -0,0 +1,455 @@ +# 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 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", + "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..70044514b14 --- /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:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1", + "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/example_code/bedrock-runtime/README.md b/python/example_code/bedrock-runtime/README.md index 784ae837f07..46d2f4a6c37 100644 --- a/python/example_code/bedrock-runtime/README.md +++ b/python/example_code/bedrock-runtime/README.md @@ -181,6 +181,28 @@ Start the example by running the following at a command prompt: python ../bedrock-agent/prompts/scenario_get_started_with_prompts.py ``` + + + +#### Create and invoke a managed prompt + +This example shows you how to do the following: + +- Create a managed prompt. +- Create a version of the prompt. +- Invoke the prompt using the version. +- Clean up resources (optional). + + + + +Start the example by running the following at a command prompt: + +``` +python ../bedrock-agent/prompts/scenario_get_started_with_prompts.py +``` + + 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 + )