diff --git a/python/samples/concepts/README.md b/python/samples/concepts/README.md index 6a4dd62dc55d..17377b62ec7b 100644 --- a/python/samples/concepts/README.md +++ b/python/samples/concepts/README.md @@ -24,12 +24,12 @@ - [Azure AI Agent Streaming](./agents/azure_ai_agent/azure_ai_agent_streaming.py) #### [Bedrock Agent](../../semantic_kernel/agents/bedrock/bedrock_agent.py) + - [Bedrock Agent Simple Chat Streaming](./agents/bedrock_agent/bedrock_agent_simple_chat_streaming.py) - [Bedrock Agent Simple Chat](./agents/bedrock_agent/bedrock_agent_simple_chat.py) -- [Bedrock Agent Update Agent](./agents/bedrock_agent/bedrock_agent_update_agent.py) -- [Bedrock Agent Use Existing](./agents/bedrock_agent/bedrock_agent_use_existing.py) - [Bedrock Agent With Code Interpreter Streaming](./agents/bedrock_agent/bedrock_agent_with_code_interpreter_streaming.py) - [Bedrock Agent With Code Interpreter](./agents/bedrock_agent/bedrock_agent_with_code_interpreter.py) +- [Bedrock Agent With Kernel Function Simple](./agents/bedrock_agent/bedrock_agent_with_kernel_function_simple.py) - [Bedrock Agent With Kernel Function Streaming](./agents/bedrock_agent/bedrock_agent_with_kernel_function_streaming.py) - [Bedrock Agent With Kernel Function](./agents/bedrock_agent/bedrock_agent_with_kernel_function.py) - [Bedrock Agent Mixed Chat Agents Streaming](./agents/bedrock_agent/bedrock_mixed_chat_agents_streaming.py) diff --git a/python/samples/concepts/agents/bedrock_agent/README.md b/python/samples/concepts/agents/bedrock_agent/README.md index 2e759a18f919..ea731c4e29c8 100644 --- a/python/samples/concepts/agents/bedrock_agent/README.md +++ b/python/samples/concepts/agents/bedrock_agent/README.md @@ -17,8 +17,6 @@ | [bedrock_agent_with_code_interpreter_streaming.py](bedrock_agent_with_code_interpreter_streaming.py) | Example of using the Bedrock agent with a code interpreter and streaming. | | [bedrock_mixed_chat_agents.py](bedrock_mixed_chat_agents.py) | Example of using multiple chat agents in a single script. | | [bedrock_mixed_chat_agents_streaming.py](bedrock_mixed_chat_agents_streaming.py) | Example of using multiple chat agents in a single script with streaming. | -| [bedrock_agent_update_agent.py](bedrock_agent_update_agent.py) | Example of updating an agent. | -| [bedrock_agent_use_existing.py](bedrock_agent_use_existing.py) | Example of using an existing agent. | ## Before running the samples diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_simple_chat.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_simple_chat.py index 7c2f36e33dd2..e50d376b93f0 100644 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_simple_chat.py +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_simple_chat.py @@ -14,7 +14,7 @@ async def main(): - bedrock_agent = await BedrockAgent.create(AGENT_NAME, instructions=INSTRUCTION) + bedrock_agent = await BedrockAgent.create_and_prepare_agent(AGENT_NAME, instructions=INSTRUCTION) session_id = BedrockAgent.create_session_id() try: diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_simple_chat_streaming.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_simple_chat_streaming.py index 5987e825b30f..099b9de75f51 100644 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_simple_chat_streaming.py +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_simple_chat_streaming.py @@ -14,7 +14,7 @@ async def main(): - bedrock_agent = await BedrockAgent.create(AGENT_NAME, instructions=INSTRUCTION) + bedrock_agent = await BedrockAgent.create_and_prepare_agent(AGENT_NAME, instructions=INSTRUCTION) session_id = BedrockAgent.create_session_id() try: diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_update_agent.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_update_agent.py deleted file mode 100644 index 781dff6332dd..000000000000 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_update_agent.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio - -from semantic_kernel.agents.bedrock.bedrock_agent import BedrockAgent - -# This sample shows how to update an existing Bedrock agent. -# This sample uses the following main component(s): -# - a Bedrock agent -# You will learn how to connect to a Bedrock agent and update its properties. - - -# Make sure to replace AGENT_NAME and AGENT_ID with the correct values -AGENT_NAME = "semantic-kernel-bedrock-agent" -INSTRUCTION = "You are a friendly assistant but you don't know anything about AI." -NEW_INSTRUCTION = "You are a friendly assistant and you know a lot about AI." - - -async def main(): - bedrock_agent = await BedrockAgent.create(AGENT_NAME, instructions=INSTRUCTION) - - async def ask_about_ai(): - session_id = BedrockAgent.create_session_id() - async for response in bedrock_agent.invoke( - session_id=session_id, - input_text="What is AI in one sentence?", - ): - print(response) - - try: - print("Before updating the agent:") - await ask_about_ai() - - await bedrock_agent.update_agent(instruction=NEW_INSTRUCTION) - - print("After updating the agent:") - await ask_about_ai() - finally: - # Delete the agent - await bedrock_agent.delete_agent() - - # Sample output (using anthropic.claude-3-haiku-20240307-v1:0): - # Before updating the agent: - # I apologize, but I do not have any information about AI or the ability to define it. - # As I mentioned, I am a friendly assistant without any knowledge about AI. I cannot - # provide a definition for AI in one sentence. If you have a different question I can - # try to assist with, please let me know. - # After updating the agent: - # AI is the field of computer science that aims to create systems and machines that can - # perform tasks that typically require human intelligence, such as learning, problem-solving, - # perception, and decision-making. - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_use_existing.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_use_existing.py deleted file mode 100644 index 969b850519a3..000000000000 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_use_existing.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio - -from semantic_kernel.agents.bedrock.bedrock_agent import BedrockAgent - -# This sample shows how to interact with a Bedrock agent in the simplest way. -# This sample uses the following main component(s): -# - a Bedrock agent that has already been created -# You will learn how to connect to an existing Bedrock agent and talk to it. - - -# Make sure to replace AGENT_NAME and AGENT_ID with the correct values -AGENT_NAME = "semantic-kernel-bedrock-agent" -AGENT_ID = "..." - - -async def main(): - bedrock_agent = await BedrockAgent.retrieve(AGENT_ID, AGENT_NAME) - session_id = BedrockAgent.create_session_id() - - try: - while True: - user_input = input("User:> ") - if user_input == "exit": - print("\n\nExiting chat...") - break - - # Invoke the agent - # The chat history is maintained in the session - async for response in bedrock_agent.invoke( - session_id=session_id, - input_text=user_input, - ): - print(f"Bedrock agent: {response}") - except KeyboardInterrupt: - print("\n\nExiting chat...") - return False - except EOFError: - print("\n\nExiting chat...") - return False - - # Sample output (using anthropic.claude-3-haiku-20240307-v1:0): - # User:> Hi, my name is John. - # Bedrock agent: Hello John. How can I help you? - # User:> What is my name? - # Bedrock agent: Your name is John. - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_code_interpreter.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_code_interpreter.py index 87f42734de04..ad6bf184b9fa 100644 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_code_interpreter.py +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_code_interpreter.py @@ -29,11 +29,8 @@ async def main(): - bedrock_agent = await BedrockAgent.create( - AGENT_NAME, - instructions=INSTRUCTION, - enable_code_interpreter=True, - ) + bedrock_agent = await BedrockAgent.create_and_prepare_agent(AGENT_NAME, instructions=INSTRUCTION) + await bedrock_agent.create_code_interpreter_action_group() session_id = BedrockAgent.create_session_id() @@ -48,7 +45,8 @@ async def main(): ): print(f"Response:\n{response}") assert isinstance(response, ChatMessageContent) # nosec - binary_item = next(item for item in response.items if isinstance(item, BinaryContent)) + if not binary_item: + binary_item = next((item for item in response.items if isinstance(item, BinaryContent)), None) finally: # Delete the agent await bedrock_agent.delete_agent() diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_code_interpreter_streaming.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_code_interpreter_streaming.py index 72d43a1e2372..ca60c477e66e 100644 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_code_interpreter_streaming.py +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_code_interpreter_streaming.py @@ -29,11 +29,8 @@ async def main(): - bedrock_agent = await BedrockAgent.create( - AGENT_NAME, - instructions=INSTRUCTION, - enable_code_interpreter=True, - ) + bedrock_agent = await BedrockAgent.create_and_prepare_agent(AGENT_NAME, instructions=INSTRUCTION) + await bedrock_agent.create_code_interpreter_action_group() session_id = BedrockAgent.create_session_id() diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function.py index 3a257f302a94..928c02054fa7 100644 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function.py +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function.py @@ -38,12 +38,13 @@ async def main(): # Create a kernel kernel = get_kernel() - bedrock_agent = await BedrockAgent.create( + bedrock_agent = await BedrockAgent.create_and_prepare_agent( AGENT_NAME, - instructions=INSTRUCTION, + INSTRUCTION, kernel=kernel, - enable_kernel_function=True, ) + # Note: We still need to create the kernel function action group on the service side. + await bedrock_agent.create_kernel_function_action_group() session_id = BedrockAgent.create_session_id() diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_simple.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_simple.py new file mode 100644 index 000000000000..b214ab5591dc --- /dev/null +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_simple.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from typing import Annotated + +from semantic_kernel.agents.bedrock.bedrock_agent import BedrockAgent +from semantic_kernel.functions.kernel_function_decorator import kernel_function + +# This sample shows how to interact with a Bedrock agent that is capable of using kernel functions. +# Instead of creating a kernel and adding plugins to it, you can directly pass the plugins to the +# agent when creating it. +# This sample uses the following main component(s): +# - a Bedrock agent +# - a kernel function +# - a kernel +# You will learn how to create a new Bedrock agent and ask it a question that requires a kernel function to answer. + +AGENT_NAME = "semantic-kernel-bedrock-agent" +INSTRUCTION = "You are a friendly assistant. You help people find information." + + +class WeatherPlugin: + """Mock weather plugin.""" + + @kernel_function(description="Get real-time weather information.") + def current(self, location: Annotated[str, "The location to get the weather"]) -> str: + """Returns the current weather.""" + return f"The weather in {location} is sunny." + + +async def main(): + bedrock_agent = await BedrockAgent.create_and_prepare_agent( + AGENT_NAME, + INSTRUCTION, + plugins=[WeatherPlugin()], + ) + # Note: We still need to create the kernel function action group on the service side. + await bedrock_agent.create_kernel_function_action_group() + + session_id = BedrockAgent.create_session_id() + + try: + # Invoke the agent + async for response in bedrock_agent.invoke( + session_id=session_id, + input_text="What is the weather in Seattle?", + ): + print(f"Response:\n{response}") + finally: + # Delete the agent + await bedrock_agent.delete_agent() + + # Sample output (using anthropic.claude-3-haiku-20240307-v1:0): + # Response: + # The current weather in Seattle is sunny. + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_streaming.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_streaming.py index 5a82ace15dd3..aa4dce75e0ed 100644 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_streaming.py +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_streaming.py @@ -38,12 +38,13 @@ async def main(): # Create a kernel kernel = get_kernel() - bedrock_agent = await BedrockAgent.create( + bedrock_agent = await BedrockAgent.create_and_prepare_agent( AGENT_NAME, - instructions=INSTRUCTION, + INSTRUCTION, kernel=kernel, - enable_kernel_function=True, ) + # Note: We still need to create the kernel function action group on the service side. + await bedrock_agent.create_kernel_function_action_group() session_id = BedrockAgent.create_session_id() diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_mixed_chat_agents.py b/python/samples/concepts/agents/bedrock_agent/bedrock_mixed_chat_agents.py index 8a600fdf9174..aa29259276d5 100644 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_mixed_chat_agents.py +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_mixed_chat_agents.py @@ -60,7 +60,7 @@ async def main(): instructions=REVIEWER_INSTRUCTIONS, ) - agent_writer = await BedrockAgent.create( + agent_writer = await BedrockAgent.create_and_prepare_agent( COPYWRITER_NAME, instructions=COPYWRITER_INSTRUCTIONS, ) diff --git a/python/samples/concepts/agents/bedrock_agent/bedrock_mixed_chat_agents_streaming.py b/python/samples/concepts/agents/bedrock_agent/bedrock_mixed_chat_agents_streaming.py index 8c695fcc7223..b4ed1668b822 100644 --- a/python/samples/concepts/agents/bedrock_agent/bedrock_mixed_chat_agents_streaming.py +++ b/python/samples/concepts/agents/bedrock_agent/bedrock_mixed_chat_agents_streaming.py @@ -60,7 +60,7 @@ async def main(): instructions=REVIEWER_INSTRUCTIONS, ) - agent_writer = await BedrockAgent.create( + agent_writer = await BedrockAgent.create_and_prepare_agent( COPYWRITER_NAME, instructions=COPYWRITER_INSTRUCTIONS, ) diff --git a/python/semantic_kernel/agents/bedrock/bedrock_agent.py b/python/semantic_kernel/agents/bedrock/bedrock_agent.py index ad4443f87022..33b57363e193 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent.py @@ -2,10 +2,11 @@ import asyncio +import logging import sys import uuid from collections.abc import AsyncIterable -from functools import reduce +from functools import partial, reduce from typing import Any, ClassVar from pydantic import ValidationError @@ -15,7 +16,6 @@ else: from typing_extensions import override # pragma: no cover -from semantic_kernel.agents.agent import Agent from semantic_kernel.agents.bedrock.action_group_utils import ( parse_function_result_contents, parse_return_control_payload, @@ -24,6 +24,7 @@ from semantic_kernel.agents.bedrock.bedrock_agent_settings import BedrockAgentSettings from semantic_kernel.agents.bedrock.models.bedrock_agent_event_type import BedrockAgentEventType from semantic_kernel.agents.bedrock.models.bedrock_agent_model import BedrockAgentModel +from semantic_kernel.agents.bedrock.models.bedrock_agent_status import BedrockAgentStatus from semantic_kernel.agents.channels.agent_channel import AgentChannel from semantic_kernel.agents.channels.bedrock_agent_channel import BedrockAgentChannel from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior @@ -36,20 +37,20 @@ from semantic_kernel.contents.utils.author_role import AuthorRole from semantic_kernel.exceptions.agent_exceptions import AgentInitializationException, AgentInvokeException from semantic_kernel.functions.kernel_arguments import KernelArguments -from semantic_kernel.functions.kernel_function import TEMPLATE_FORMAT_MAP from semantic_kernel.functions.kernel_plugin import KernelPlugin from semantic_kernel.kernel import Kernel -from semantic_kernel.prompt_template.prompt_template_base import PromptTemplateBase -from semantic_kernel.prompt_template.prompt_template_config import PromptTemplateConfig +from semantic_kernel.utils.async_utils import run_in_executor from semantic_kernel.utils.feature_stage_decorator import experimental from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import ( trace_agent_get_response, trace_agent_invocation, ) +logger = logging.getLogger(__name__) + @experimental -class BedrockAgent(BedrockAgentBase, Agent): +class BedrockAgent(BedrockAgentBase): """Bedrock Agent. Manages the interaction with Amazon Bedrock Agent Service. @@ -59,106 +60,68 @@ class BedrockAgent(BedrockAgentBase, Agent): def __init__( self, - name: str, + agent_model: BedrockAgentModel | dict[str, Any], *, - agent_resource_role_arn: str | None = None, - arguments: KernelArguments | None = None, - env_file_encoding: str | None = None, - env_file_path: str | None = None, - id: str | None = None, - instructions: str | None = None, - foundation_model: str | None = None, function_choice_behavior: FunctionChoiceBehavior | None = None, kernel: Kernel | None = None, plugins: list[KernelPlugin | object] | dict[str, KernelPlugin | object] | None = None, - prompt_template_config: PromptTemplateConfig | None = None, + arguments: KernelArguments | None = None, + bedrock_runtime_client: Any | None = None, + bedrock_client: Any | None = None, + **kwargs, ) -> None: """Initialize the Bedrock Agent. Note that this only creates the agent object and does not create the agent in the service. Args: - name: The name of the agent. - agent_resource_role_arn: The ARN of the agent resource role. - Overrides the one in the env file. - arguments: The kernel arguments. - Invoke method arguments take precedence over the arguments provided here. - env_file_path: The path to the environment file. - env_file_encoding: The encoding of the environment file. - foundation_model: The foundation model. Overrides the one in the env file. - function_choice_behavior: The function choice behavior for accessing + agent_model (BedrockAgentModel | dict[str, Any]): The agent model. + function_choice_behavior (FunctionChoiceBehavior, optional): The function choice behavior for accessing the kernel functions and filters. - id: The unique identifier of the agent. - instructions: The instructions for the agent. - kernel: The kernel to use. - plugins: The plugins for the agent. If plugins are included along with a kernel, any plugins - that already exist in the kernel will be overwritten. - prompt_template_config: The prompt template configuration. - Cannot be set if instructions is set. + kernel (Kernel, optional): The kernel to use. + plugins (list[KernelPlugin | object] | dict[str, KernelPlugin | object], optional): The plugins to use. + arguments (KernelArguments, optional): The kernel arguments. + Invoke method arguments take precedence over the arguments provided here. + bedrock_runtime_client: The Bedrock Runtime Client. + bedrock_client: The Bedrock Client. + **kwargs: Additional keyword arguments. """ - try: - bedrock_agent_settings = BedrockAgentSettings.create( - agent_resource_role_arn=agent_resource_role_arn, - foundation_model=foundation_model, - env_file_path=env_file_path, - env_file_encoding=env_file_encoding, - ) - except ValidationError as e: - raise AgentInitializationException("Failed to initialize the Amazon Bedrock Agent settings.") from e - - bedrock_agent_model = BedrockAgentModel( - agentId=id, - agentName=name, - foundationModel=bedrock_agent_settings.foundation_model, - ) - - prompt_template: PromptTemplateBase | None = None - if instructions and prompt_template_config and prompt_template_config.template: - raise AgentInitializationException("Cannot set both instructions and prompt_template_config.template.") - if prompt_template_config: - prompt_template = TEMPLATE_FORMAT_MAP[prompt_template_config.template_format]( - prompt_template_config=prompt_template_config - ) - args: dict[str, Any] = { - "agent_resource_role_arn": bedrock_agent_settings.agent_resource_role_arn, - "name": name, - "agent_model": bedrock_agent_model, + "agent_model": agent_model, + **kwargs, } - if id: - args["id"] = id - if instructions: - args["instructions"] = instructions - if kernel: - args["kernel"] = kernel + if function_choice_behavior: args["function_choice_behavior"] = function_choice_behavior + if kernel: + args["kernel"] = kernel + if plugins: + args["plugins"] = plugins if arguments: args["arguments"] = arguments - if prompt_template: - args["prompt_template"] = prompt_template - if plugins is not None: - args["plugins"] = plugins + if bedrock_runtime_client: + args["bedrock_runtime_client"] = bedrock_runtime_client + if bedrock_client: + args["bedrock_client"] = bedrock_client super().__init__(**args) # region convenience class methods @classmethod - async def create( + async def create_and_prepare_agent( cls, name: str, + instructions: str, *, agent_resource_role_arn: str | None = None, foundation_model: str | None = None, + bedrock_runtime_client: Any | None = None, + bedrock_client: Any | None = None, kernel: Kernel | None = None, + plugins: list[KernelPlugin | object] | dict[str, KernelPlugin | object] | None = None, function_choice_behavior: FunctionChoiceBehavior | None = None, arguments: KernelArguments | None = None, - instructions: str | None = None, - prompt_template_config: PromptTemplateConfig | None = None, - enable_code_interpreter: bool | None = None, - enable_user_input: bool | None = None, - enable_kernel_function: bool | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, ) -> "BedrockAgent": @@ -168,89 +131,69 @@ async def create( Args: name (str): The name of the agent. + instructions (str, optional): The instructions for the agent. agent_resource_role_arn (str, optional): The ARN of the agent resource role. foundation_model (str, optional): The foundation model. + bedrock_runtime_client (Any, optional): The Bedrock Runtime Client. + bedrock_client (Any, optional): The Bedrock Client. kernel (Kernel, optional): The kernel to use. + plugins (list[KernelPlugin | object] | dict[str, KernelPlugin | object], optional): The plugins to use. function_choice_behavior (FunctionChoiceBehavior, optional): The function choice behavior for accessing - the kernel functions and filters. + the kernel functions and filters. Only FunctionChoiceType.AUTO is supported. arguments (KernelArguments, optional): The kernel arguments. - instructions (str, optional): The instructions for the agent. prompt_template_config (PromptTemplateConfig, optional): The prompt template configuration. - enable_code_interpreter (bool, optional): Enable code interpretation. - enable_user_input (bool, optional): Enable user input. - enable_kernel_function (bool, optional): Enable kernel function. env_file_path (str, optional): The path to the environment file. env_file_encoding (str, optional): The encoding of the environment file. Returns: An instance of BedrockAgent with the created agent. """ - bedrock_agent = cls( - name, - agent_resource_role_arn=agent_resource_role_arn, - foundation_model=foundation_model, - kernel=kernel, - function_choice_behavior=function_choice_behavior, - arguments=arguments, - instructions=instructions, - prompt_template_config=prompt_template_config, - env_file_path=env_file_path, - env_file_encoding=env_file_encoding, - ) - - return await bedrock_agent.create_agent( - enable_code_interpreter=enable_code_interpreter, - enable_user_input=enable_user_input, - enable_kernel_function=enable_kernel_function, - ) - - @classmethod - async def retrieve( - cls, - id: str, - name: str, - *, - agent_resource_role_arn: str | None = None, - kernel: Kernel | None = None, - function_choice_behavior: FunctionChoiceBehavior | None = None, - env_file_path: str | None = None, - env_file_encoding: str | None = None, - ) -> "BedrockAgent": - """Retrieve an agent asynchronously. + try: + bedrock_agent_settings = BedrockAgentSettings.create( + agent_resource_role_arn=agent_resource_role_arn, + foundation_model=foundation_model, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + except ValidationError as e: + raise AgentInitializationException("Failed to initialize the Amazon Bedrock Agent settings.") from e - This is a convenience method that creates an instance of BedrockAgent and - then retrieves an existing agent from the service. + import boto3 + from botocore.exceptions import ClientError - Note that: - 1. If the agent has existing action groups that require control returned to the user, - a kernel with the required functions must be provided. - 2. If the agent has not been prepared, client code must prepare the agent by calling `prepare_agent()`. + bedrock_runtime_client = bedrock_runtime_client or boto3.client("bedrock-agent-runtime") + bedrock_client = bedrock_client or boto3.client("bedrock-agent") - If client code want to enable the available action groups, it can call the respective methods: - - `create_code_interpreter_action_group()` - - `create_user_input_action_group()` - - `create_kernel_function_action_group()` + try: + response = await run_in_executor( + None, + partial( + bedrock_client.create_agent, + agentName=name, + foundationModel=bedrock_agent_settings.foundation_model, + agentResourceRoleArn=bedrock_agent_settings.agent_resource_role_arn, + instruction=instructions, + ), + ) + except ClientError as e: + logger.error(f"Failed to create agent {name}.") + raise AgentInitializationException("Failed to create the Amazon Bedrock Agent.") from e - Args: - id (str): The unique identifier of the agent. - name (str): The name of the agent. - agent_resource_role_arn (str, optional): The ARN of the agent resource role. - kernel (Kernel, optional): The kernel to use. - function_choice_behavior (FunctionChoiceBehavior, optional): The function choice behavior for accessing - the kernel functions and filters. - env_file_path (str, optional): The path to the environment file. - env_file_encoding (str, optional): The encoding of the environment file. - """ bedrock_agent = cls( - name, - id=id, - agent_resource_role_arn=agent_resource_role_arn, - kernel=kernel, + response["agent"], function_choice_behavior=function_choice_behavior, - env_file_path=env_file_path, - env_file_encoding=env_file_encoding, + kernel=kernel, + plugins=plugins, + arguments=arguments, + bedrock_runtime_client=bedrock_runtime_client, + bedrock_client=bedrock_client, ) - bedrock_agent.agent_model = await bedrock_agent._get_agent() + + # The agent will first enter the CREATING status. + # When the operation finishes, it will enter the NOT_PREPARED status. + # We need to wait for the agent to reach the NOT_PREPARED status before we can prepare it. + await bedrock_agent._wait_for_agent_status(BedrockAgentStatus.NOT_PREPARED) + await bedrock_agent.prepare_agent_and_wait_until_prepared() return bedrock_agent @@ -268,48 +211,6 @@ def create_session_id(cls) -> str: # endregion - async def create_agent( - self, - *, - enable_code_interpreter: bool | None = None, - enable_user_input: bool | None = None, - enable_kernel_function: bool | None = None, - **kwargs, - ) -> "BedrockAgent": - """Create an agent on the service asynchronously. This will also prepare the agent so that it ready for use. - - Args: - enable_code_interpreter (bool, optional): Enable code interpretation. - enable_user_input (bool, optional): Enable user input. - enable_kernel_function (bool, optional): Enable kernel function. - **kwargs: Additional keyword arguments. - - Returns: - An instance of BedrockAgent with the created agent. - """ - if not self.agent_model.foundation_model: - raise AgentInitializationException("Foundation model is required to create an agent.") - - await self._create_agent( - self.instructions or await self.format_instructions(self.kernel, self.arguments) or "", - **kwargs, - ) - - if enable_code_interpreter: - await self.create_code_interpreter_action_group() - if enable_user_input: - await self.create_user_input_action_group() - if enable_kernel_function: - await self._create_kernel_function_action_group(self.kernel) - - await self.prepare_agent() - - if not self.agent_model.agent_id: - raise AgentInitializationException("Agent ID is not set.") - self.id = self.agent_model.agent_id - - return self - @trace_agent_get_response @override async def get_response( diff --git a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py index 6fe432af9f0f..708c0d2b01de 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py @@ -2,7 +2,6 @@ import asyncio import logging -from collections.abc import Callable from functools import partial from typing import Any, ClassVar @@ -10,20 +9,20 @@ from botocore.exceptions import ClientError from pydantic import Field, field_validator +from semantic_kernel.agents.agent import Agent from semantic_kernel.agents.bedrock.action_group_utils import kernel_function_to_bedrock_function_schema from semantic_kernel.agents.bedrock.models.bedrock_action_group_model import BedrockActionGroupModel from semantic_kernel.agents.bedrock.models.bedrock_agent_model import BedrockAgentModel from semantic_kernel.agents.bedrock.models.bedrock_agent_status import BedrockAgentStatus from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior, FunctionChoiceType -from semantic_kernel.kernel import Kernel -from semantic_kernel.kernel_pydantic import KernelBaseModel +from semantic_kernel.utils.async_utils import run_in_executor from semantic_kernel.utils.feature_stage_decorator import experimental logger = logging.getLogger(__name__) @experimental -class BedrockAgentBase(KernelBaseModel): +class BedrockAgentBase(Agent): """Bedrock Agent Base Class to provide common functionalities for Bedrock Agents.""" # There is a default alias created by Bedrock for the working draft version of the agent. @@ -41,13 +40,10 @@ class BedrockAgentBase(KernelBaseModel): function_choice_behavior: FunctionChoiceBehavior = Field(default=FunctionChoiceBehavior.Auto()) # Agent Model: stores the agent information agent_model: BedrockAgentModel - # Agent ARN: The Amazon Resource Name (ARN) of the agent. - agent_resource_role_arn: str def __init__( self, - agent_resource_role_arn: str, - agent_model: BedrockAgentModel, + agent_model: BedrockAgentModel | dict[str, Any], *, function_choice_behavior: FunctionChoiceBehavior | None = None, bedrock_runtime_client: Any | None = None, @@ -57,16 +53,20 @@ def __init__( """Initialize the Bedrock Agent Base. Args: - agent_resource_role_arn: The Amazon Resource Name (ARN) of the agent. agent_model: The Bedrock Agent Model. function_choice_behavior: The function choice behavior. - bedrock_runtime_client: The Bedrock Runtime Client. bedrock_client: The Bedrock Client. + bedrock_runtime_client: The Bedrock Runtime Client. kwargs: Additional keyword arguments. """ + agent_model = ( + agent_model if isinstance(agent_model, BedrockAgentModel) else BedrockAgentModel.model_validate(agent_model) + ) + args = { - "agent_resource_role_arn": agent_resource_role_arn, "agent_model": agent_model, + "id": agent_model.agent_id, + "name": agent_model.agent_name, "bedrock_runtime_client": bedrock_runtime_client or boto3.client("bedrock-agent-runtime"), "bedrock_client": bedrock_client or boto3.client("bedrock-agent"), **kwargs, @@ -89,56 +89,19 @@ def validate_function_choice_behavior( raise ValueError("Only FunctionChoiceType.AUTO is supported.") return function_choice_behavior - async def _run_in_executor(self, executor: Any, func: Callable, *args, **kwargs) -> Any: - """Run a function in an executor.""" - return await asyncio.get_event_loop().run_in_executor(executor, partial(func, *args, **kwargs)) - def __repr__(self): """Return the string representation of the Bedrock Agent.""" return f"{self.agent_model}" # region Agent Management - async def _create_agent( - self, - instruction: str, - **kwargs, - ) -> BedrockAgentModel: - """Create an agent asynchronously on the Bedrock service.""" - if self.agent_model.agent_id: - # Once the agent is created, the agent_id will be set. - raise ValueError("Agent already exists. Please delete the agent before creating a new one.") - - try: - response = await self._run_in_executor( - None, - partial( - self.bedrock_client.create_agent, - agentName=self.agent_model.agent_name, - foundationModel=self.agent_model.foundation_model, - agentResourceRoleArn=self.agent_resource_role_arn, - instruction=instruction, - **kwargs, - ), - ) - self.agent_model = BedrockAgentModel(**response["agent"]) - - # The agent will first enter the CREATING status. - # When the agent is created, it will enter the NOT_PREPARED status. - await self._wait_for_agent_status(BedrockAgentStatus.NOT_PREPARED) - - return self.agent_model - except ClientError as e: - logger.error(f"Failed to create agent {self.agent_model.agent_name}.") - raise e - - async def prepare_agent(self) -> None: + async def prepare_agent_and_wait_until_prepared(self) -> None: """Prepare the agent for use.""" if not self.agent_model.agent_id: raise ValueError("Agent does not exist. Please create the agent before preparing it.") try: - await self._run_in_executor( + await run_in_executor( None, partial( self.bedrock_client.prepare_agent, @@ -146,70 +109,23 @@ async def prepare_agent(self) -> None: ), ) + # The agent will take some time to enter the PREPARING status after the prepare operation is called. + # We need to wait for the agent to reach the PREPARING status before we can proceed, otherwise we + # will return immediately if the agent is already in PREPARED status. + await self._wait_for_agent_status(BedrockAgentStatus.PREPARING) + # The agent will enter the PREPARED status when the preparation is complete. await self._wait_for_agent_status(BedrockAgentStatus.PREPARED) except ClientError as e: logger.error(f"Failed to prepare agent {self.agent_model.agent_id}.") raise e - async def _create_agent_alias(self, alias_name: str, **kwargs) -> dict[str, Any]: - """Create an agent alias asynchronously. - - Creating an alias is similar to taking a snapshot of the agent's current settings. - Later, users can switch between aliases to use different configurations. - """ - if not self.agent_model.agent_id: - raise ValueError("Agent does not exist. Please create the agent before creating an alias.") - - try: - return await self._run_in_executor( - None, - partial( - self.bedrock_client.create_agent_alias, - agentAliasName=alias_name, - agentId=self.agent_model.agent_id, - **kwargs, - ), - ) - except ClientError as e: - logger.error(f"Failed to create alias {alias_name} for agent {self.agent_model.agent_id}.") - raise e - - async def update_agent(self, **kwargs) -> None: - """Update an agent asynchronously. - - Args: - kwargs: The keyword arguments to update the agent. - """ - if not self.agent_model.agent_id: - raise ValueError("Agent has not been created. Please create the agent before updating it.") - - try: - response = await self._run_in_executor( - None, - partial( - self.bedrock_client.update_agent, - agentId=self.agent_model.agent_id, - agentResourceRoleArn=self.agent_resource_role_arn, - # Use the existing agent name and foundation model if not provided since they are required. - agentName=kwargs.pop("agentName", None) or self.agent_model.agent_name, - foundationModel=kwargs.pop("foundationModel", None) or self.agent_model.foundation_model, - **kwargs, - ), - ) - self.agent_model = BedrockAgentModel(**response["agent"]) - - await self._wait_for_agent_status(BedrockAgentStatus.PREPARED) - except ClientError as e: - logger.error(f"Failed to update agent {self.agent_model.agent_id}.") - raise e - async def delete_agent(self, **kwargs) -> None: """Delete an agent asynchronously.""" if not self.agent_model.agent_id: raise ValueError("Agent does not exist. Please create the agent before deleting it.") try: - await self._run_in_executor( + await run_in_executor( None, partial( self.bedrock_client.delete_agent, @@ -223,13 +139,13 @@ async def delete_agent(self, **kwargs) -> None: logger.error(f"Failed to delete agent {self.agent_model.agent_id}.") raise e - async def _get_agent(self) -> BedrockAgentModel: + async def _get_agent(self) -> None: """Get an agent.""" if not self.agent_model.agent_id: raise ValueError("Agent does not exist. Please create the agent before getting it.") try: - response = await self._run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.get_agent, @@ -237,7 +153,8 @@ async def _get_agent(self) -> BedrockAgentModel: ), ) - return BedrockAgentModel(**response["agent"]) + # Update the agent model + self.agent_model = BedrockAgentModel(**response["agent"]) except ClientError as e: logger.error(f"Failed to get agent {self.agent_model.agent_id}.") raise e @@ -250,13 +167,16 @@ async def _wait_for_agent_status( ) -> None: """Wait for the agent to reach a specific status.""" for _ in range(max_attempts): - agent = await self._get_agent() - if agent.agent_status == status: + await self._get_agent() + if self.agent_model.agent_status == status: return await asyncio.sleep(interval) - raise TimeoutError(f"Agent did not reach status {status} within the specified time.") + raise TimeoutError( + f"Agent did not reach status {status} within the specified time." + f" Current status: {self.agent_model.agent_status}" + ) # endregion Agent Management @@ -267,7 +187,7 @@ async def create_code_interpreter_action_group(self, **kwargs) -> BedrockActionG raise ValueError("Agent does not exist. Please create the agent before creating an action group for it.") try: - response = await self._run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.create_agent_action_group, @@ -280,6 +200,8 @@ async def create_code_interpreter_action_group(self, **kwargs) -> BedrockActionG ), ) + await self.prepare_agent_and_wait_until_prepared() + return BedrockActionGroupModel(**response["agentActionGroup"]) except ClientError as e: logger.error(f"Failed to create code interpreter action group for agent {self.agent_model.agent_id}.") @@ -291,7 +213,7 @@ async def create_user_input_action_group(self, **kwargs) -> BedrockActionGroupMo raise ValueError("Agent does not exist. Please create the agent before creating an action group for it.") try: - response = await self._run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.create_agent_action_group, @@ -304,23 +226,25 @@ async def create_user_input_action_group(self, **kwargs) -> BedrockActionGroupMo ), ) + await self.prepare_agent_and_wait_until_prepared() + return BedrockActionGroupModel(**response["agentActionGroup"]) except ClientError as e: logger.error(f"Failed to create user input action group for agent {self.agent_model.agent_id}.") raise e - async def _create_kernel_function_action_group(self, kernel: Kernel, **kwargs) -> BedrockActionGroupModel | None: + async def create_kernel_function_action_group(self, **kwargs) -> BedrockActionGroupModel | None: """Create a kernel function action group.""" if not self.agent_model.agent_id: raise ValueError("Agent does not exist. Please create the agent before creating an action group for it.") - function_call_choice_config = self.function_choice_behavior.get_config(kernel) + function_call_choice_config = self.function_choice_behavior.get_config(self.kernel) if not function_call_choice_config.available_functions: logger.warning("No available functions. Skipping kernel function action group creation.") return None try: - response = await self._run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.create_agent_action_group, @@ -334,6 +258,8 @@ async def _create_kernel_function_action_group(self, kernel: Kernel, **kwargs) - ), ) + await self.prepare_agent_and_wait_until_prepared() + return BedrockActionGroupModel(**response["agentActionGroup"]) except ClientError as e: logger.error(f"Failed to create kernel function action group for agent {self.agent_model.agent_id}.") @@ -351,7 +277,7 @@ async def associate_agent_knowledge_base(self, knowledge_base_id: str, **kwargs) ) try: - return await self._run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.associate_agent_knowledge_base, @@ -361,6 +287,10 @@ async def associate_agent_knowledge_base(self, knowledge_base_id: str, **kwargs) **kwargs, ), ) + + await self.prepare_agent_and_wait_until_prepared() + + return response except ClientError as e: logger.error( f"Failed to associate agent {self.agent_model.agent_id} with knowledge base {knowledge_base_id}." @@ -375,7 +305,7 @@ async def disassociate_agent_knowledge_base(self, knowledge_base_id: str, **kwar ) try: - await self._run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.disassociate_agent_knowledge_base, @@ -385,6 +315,10 @@ async def disassociate_agent_knowledge_base(self, knowledge_base_id: str, **kwar **kwargs, ), ) + + await self.prepare_agent_and_wait_until_prepared() + + return response except ClientError as e: logger.error( f"Failed to disassociate agent {self.agent_model.agent_id} with knowledge base {knowledge_base_id}." @@ -397,7 +331,7 @@ async def list_associated_agent_knowledge_bases(self, **kwargs) -> dict[str, Any raise ValueError("Agent does not exist. Please create the agent before listing associated knowledge bases.") try: - return await self._run_in_executor( + return await run_in_executor( None, partial( self.bedrock_client.list_agent_knowledge_bases, @@ -426,7 +360,7 @@ async def _invoke_agent( agent_alias = agent_alias or self.WORKING_DRAFT_AGENT_ALIAS try: - return await self._run_in_executor( + return await run_in_executor( None, partial( self.bedrock_runtime_client.invoke_agent, diff --git a/python/semantic_kernel/agents/bedrock/bedrock_agent_settings.py b/python/semantic_kernel/agents/bedrock/bedrock_agent_settings.py index 56635020abc8..a3679478678b 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent_settings.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent_settings.py @@ -23,11 +23,10 @@ class BedrockAgentSettings(KernelBaseSettings): https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html (Env var BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN) - foundation_model: str - The Amazon Bedrock foundation model ID to use. - This is required when creating a new agent. (Env var BEDROCK_AGENT_FOUNDATION_MODEL) """ env_prefix: ClassVar[str] = "BEDROCK_AGENT_" agent_resource_role_arn: str - foundation_model: str | None = None + foundation_model: str diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_base.py b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_base.py index 115ae65409dd..22bb30c60d6a 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_base.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_base.py @@ -4,8 +4,8 @@ from functools import partial from typing import Any, ClassVar -from semantic_kernel.connectors.ai.bedrock.services.model_provider.utils import run_in_executor from semantic_kernel.kernel_pydantic import KernelBaseModel +from semantic_kernel.utils.async_utils import run_in_executor class BedrockBase(KernelBaseModel, ABC): diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_chat_completion.py b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_chat_completion.py index 5c4f3e6cd192..a65034fa1d49 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_chat_completion.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_chat_completion.py @@ -25,7 +25,6 @@ finish_reason_from_bedrock_to_semantic_kernel, format_bedrock_function_name_to_kernel_function_fully_qualified_name, remove_none_recursively, - run_in_executor, update_settings_from_function_choice_configuration, ) from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase @@ -45,6 +44,7 @@ ServiceInvalidRequestError, ServiceInvalidResponseError, ) +from semantic_kernel.utils.async_utils import run_in_executor from semantic_kernel.utils.telemetry.model_diagnostics.decorators import ( trace_chat_completion, trace_streaming_chat_completion, diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_completion.py b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_completion.py index 81092a7c7fa4..e1335f2d01cc 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_completion.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_completion.py @@ -22,11 +22,11 @@ parse_streaming_text_completion_response, parse_text_completion_response, ) -from semantic_kernel.connectors.ai.bedrock.services.model_provider.utils import run_in_executor from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase from semantic_kernel.contents.streaming_text_content import StreamingTextContent from semantic_kernel.contents.text_content import TextContent from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidRequestError +from semantic_kernel.utils.async_utils import run_in_executor from semantic_kernel.utils.telemetry.model_diagnostics.decorators import ( trace_streaming_text_completion, trace_text_completion, diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_embedding.py b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_embedding.py index f963db5c5f0b..7fe759f1c0be 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_embedding.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_embedding.py @@ -24,10 +24,10 @@ get_text_embedding_request_body, parse_text_embedding_response, ) -from semantic_kernel.connectors.ai.bedrock.services.model_provider.utils import run_in_executor from semantic_kernel.connectors.ai.embeddings.embedding_generator_base import EmbeddingGeneratorBase from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidRequestError +from semantic_kernel.utils.async_utils import run_in_executor if TYPE_CHECKING: pass diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/utils.py b/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/utils.py index 7607696559c5..e6425eda1c39 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/utils.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/model_provider/utils.py @@ -1,9 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. -import asyncio import json from collections.abc import Callable, Mapping -from functools import partial from typing import TYPE_CHECKING, Any from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import BedrockChatPromptExecutionSettings @@ -23,11 +21,6 @@ from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings -async def run_in_executor(executor, func, *args, **kwargs): - """Run a function in an executor.""" - return await asyncio.get_event_loop().run_in_executor(executor, partial(func, *args, **kwargs)) - - def remove_none_recursively(data: dict, max_depth: int = 5) -> dict: """Remove None values from a dictionary recursively.""" if max_depth <= 0: diff --git a/python/semantic_kernel/utils/async_utils.py b/python/semantic_kernel/utils/async_utils.py new file mode 100644 index 000000000000..4d2d2d50249a --- /dev/null +++ b/python/semantic_kernel/utils/async_utils.py @@ -0,0 +1,11 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from collections.abc import Callable +from functools import partial +from typing import Any + + +async def run_in_executor(executor: Any, func: Callable, *args, **kwargs) -> Any: + """Run a function in an executor.""" + return await asyncio.get_event_loop().run_in_executor(executor, partial(func, *args, **kwargs)) diff --git a/python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py b/python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py index 4f49b39acd6d..884464cb7efb 100644 --- a/python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py +++ b/python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py @@ -7,6 +7,7 @@ from opentelemetry.trace import get_tracer from semantic_kernel.contents.chat_message_content import ChatMessageContent +from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent from semantic_kernel.utils.feature_stage_decorator import experimental from semantic_kernel.utils.telemetry.agent_diagnostics import gen_ai_attributes @@ -24,7 +25,9 @@ def trace_agent_invocation(invoke_func: Callable) -> Callable: OPERATION_NAME = "invoke_agent" @functools.wraps(invoke_func) - async def wrapper_decorator(*args: Any, **kwargs: Any) -> AsyncIterable: + async def wrapper_decorator( + *args: Any, **kwargs: Any + ) -> AsyncIterable[ChatMessageContent | StreamingChatMessageContent]: agent: "Agent" = args[0] with tracer.start_as_current_span(f"{OPERATION_NAME} {agent.name}") as span: diff --git a/python/tests/integration/agents/bedrock_agent/test_bedrock_agent_integration.py b/python/tests/integration/agents/bedrock_agent/test_bedrock_agent_integration.py index 34acf5da7cac..04d35ccd1bc3 100644 --- a/python/tests/integration/agents/bedrock_agent/test_bedrock_agent_integration.py +++ b/python/tests/integration/agents/bedrock_agent/test_bedrock_agent_integration.py @@ -18,21 +18,18 @@ async def setup_and_teardown(self, request): This is run for each test function, i.e. each test function will have its own instance of the agent. """ - kwargs = {} - if hasattr(request, "param"): - if "enable_code_interpreter" in request.param: - kwargs["enable_code_interpreter"] = request.param.get("enable_code_interpreter") - if "enable_kernel_function" in request.param: - kwargs["enable_kernel_function"] = request.param.get("enable_kernel_function") - if "kernel" in request.param: - kwargs["kernel"] = request.getfixturevalue(request.param.get("kernel")) - try: - self.bedrock_agent = await BedrockAgent.create( - name=f"semantic-kernel-integration-test-agent-{uuid.uuid4()}", - instructions="You are a helpful assistant that help users with their questions.", - **kwargs, + self.bedrock_agent = await BedrockAgent.create_and_prepare_agent( + f"semantic-kernel-integration-test-agent-{uuid.uuid4()}", + "You are a helpful assistant that help users with their questions.", ) + if hasattr(request, "param"): + if "enable_code_interpreter" in request.param: + await self.bedrock_agent.create_code_interpreter_action_group() + if "kernel" in request.param: + self.bedrock_agent.kernel = request.getfixturevalue(request.param.get("kernel")) + if "enable_kernel_function" in request.param: + await self.bedrock_agent.create_kernel_function_action_group() except Exception as e: pytest.fail("Failed to create agent") raise e @@ -45,12 +42,6 @@ async def setup_and_teardown(self, request): pytest.fail(f"Failed to delete agent: {e}") raise e - @pytest.mark.asyncio - async def test_update(self): - """Test updating the agent.""" - await self.bedrock_agent.update_agent(agentName="updated_agent") - assert self.bedrock_agent.agent_model.agent_name == "updated_agent" - @pytest.mark.asyncio async def test_invoke(self): """Test invoke of the agent.""" @@ -79,12 +70,14 @@ async def test_code_interpreter(self): Monkey 6 Dolphin 2 """ + binary_item: BinaryContent | None = None async for message in self.bedrock_agent.invoke(BedrockAgent.create_session_id(), input_text): assert isinstance(message, ChatMessageContent) assert message.role == AuthorRole.ASSISTANT - if not any(item for item in message.items if isinstance(item, BinaryContent)): - # TODO (eavanvalkenburg): redo the assert instead of this. - pytest.xfail(reason="flaky test") + if not binary_item: + binary_item = next((item for item in message.items if isinstance(item, BinaryContent)), None) + + assert binary_item @pytest.mark.asyncio @pytest.mark.parametrize("setup_and_teardown", [{"enable_code_interpreter": True}], indirect=True) diff --git a/python/tests/unit/agents/bedrock_agent/conftest.py b/python/tests/unit/agents/bedrock_agent/conftest.py index 45ff026fcf40..b76ae70b88a5 100644 --- a/python/tests/unit/agents/bedrock_agent/conftest.py +++ b/python/tests/unit/agents/bedrock_agent/conftest.py @@ -78,6 +78,18 @@ def bedrock_agent_model_with_id_prepared_dict(): } +@pytest.fixture +def bedrock_agent_model_with_id_preparing_dict(): + return { + "agent": { + "agentId": "test_agent_id", + "agentName": "test_agent_name", + "foundationModel": "test_foundation_model", + "agentStatus": "PREPARING", + } + } + + @pytest.fixture def bedrock_agent_model_with_id_not_prepared_dict(): return { diff --git a/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py b/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py index be96116aba23..ddf49aca36ad 100644 --- a/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py +++ b/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py @@ -5,205 +5,434 @@ import boto3 import pytest -from semantic_kernel.agents.bedrock.action_group_utils import parse_function_result_contents +from semantic_kernel.agents.bedrock.action_group_utils import ( + kernel_function_to_bedrock_function_schema, + parse_function_result_contents, +) from semantic_kernel.agents.bedrock.bedrock_agent import BedrockAgent from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior from semantic_kernel.contents.function_result_content import FunctionResultContent from semantic_kernel.exceptions.agent_exceptions import AgentInitializationException, AgentInvokeException -from semantic_kernel.prompt_template.prompt_template_config import PromptTemplateConfig +from semantic_kernel.functions.kernel_arguments import KernelArguments +from semantic_kernel.kernel import Kernel +# region Agent Initialization Tests -# Test case to verify the initialization of BedrockAgent -@patch.object(boto3, "client", return_value=Mock()) -async def test_bedrock_agent_initialization(client, bedrock_agent_unit_test_env): - agent = BedrockAgent(name="test_agent", env_file_path="fake_path") - - assert agent.name == "test_agent" - assert agent.agent_model.agent_name == "test_agent" - assert agent.id is not None # This is a UUID generated by the agent class - assert agent.agent_model.agent_id is None # The agent has not been created on the service yet - assert agent.agent_resource_role_arn == bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"] - assert agent.agent_model.foundation_model == bedrock_agent_unit_test_env["BEDROCK_AGENT_FOUNDATION_MODEL"] +# Test case to verify BedrockAgent initialization +@patch.object(boto3, "client", return_value=Mock()) +async def test_bedrock_agent_initialization(client, bedrock_agent_model_with_id): + agent = BedrockAgent(bedrock_agent_model_with_id) -# Test case to verify error handling during BedrockAgent initialization -@pytest.mark.parametrize("exclude_list", [["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"]], indirect=True) -async def test_bedrock_agent_initialization_error(bedrock_agent_unit_test_env): - with pytest.raises(AgentInitializationException, match="Failed to initialize the Amazon Bedrock Agent settings."): - BedrockAgent(name="test_agent", env_file_path="fake_path") - - -# Test case to verify error handling during BedrockAgent initialization with instructions and prompt template -async def test_bedrock_agent_initialization_error_with_instructions_and_prompt_template(bedrock_agent_unit_test_env): - with pytest.raises( - AgentInitializationException, match="Cannot set both instructions and prompt_template_config.template." - ): - BedrockAgent( - name="test_agent", - instructions="test_instructions", - prompt_template_config=PromptTemplateConfig(template="test_template"), - env_file_path="fake_path", - ) + assert agent.name == bedrock_agent_model_with_id.agent_name + assert agent.agent_model.agent_name == bedrock_agent_model_with_id.agent_name + assert agent.agent_model.agent_id == bedrock_agent_model_with_id.agent_id + assert agent.agent_model.foundation_model == bedrock_agent_model_with_id.foundation_model # Test case to verify error handling during BedrockAgent initialization with non-auto function choice @patch.object(boto3, "client", return_value=Mock()) -async def test_bedrock_agent_initialization_error_with_non_auto_function_choice(client, bedrock_agent_unit_test_env): +async def test_bedrock_agent_initialization_error_with_non_auto_function_choice(client, bedrock_agent_model_with_id): with pytest.raises(ValueError, match="Only FunctionChoiceType.AUTO is supported."): BedrockAgent( - name="test_agent", + bedrock_agent_model_with_id, function_choice_behavior=FunctionChoiceBehavior.NoneInvoke(), - env_file_path="fake_path", ) # Test case to verify the creation of BedrockAgent @patch.object(boto3, "client", return_value=Mock()) -async def test_bedrock_agent_create( +@pytest.mark.parametrize( + "kernel, function_choice_behavior, arguments", + [ + (None, None, None), + (Kernel(), None, None), + (Kernel(), FunctionChoiceBehavior.Auto(), None), + (Kernel(), FunctionChoiceBehavior.Auto(), KernelArguments()), + ], +) +async def test_bedrock_agent_create_and_prepare_agent( client, - bedrock_agent_unit_test_env, bedrock_agent_model_with_id_not_prepared_dict, + bedrock_agent_unit_test_env, + kernel, + function_choice_behavior, + arguments, ): - agent = BedrockAgent(name="test_agent", instructions="test_instructions", env_file_path="fake_path") - - assert agent.agent_model.agent_id is None - with ( - patch.object(agent.bedrock_client, "create_agent") as mock_create_agent, - patch.object(agent.bedrock_client, "get_agent") as mock_get_agent, - patch.object(BedrockAgent, "prepare_agent", new_callable=AsyncMock), + patch.object(client, "create_agent") as mock_create_agent, + patch.object(BedrockAgent, "_wait_for_agent_status", new_callable=AsyncMock), + patch.object(BedrockAgent, "prepare_agent_and_wait_until_prepared", new_callable=AsyncMock), ): mock_create_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict - mock_get_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict - await agent.create_agent() + agent = await BedrockAgent.create_and_prepare_agent( + name=bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentName"], + instructions="test_instructions", + bedrock_client=client, + env_file_path="fake_path", + kernel=kernel, + function_choice_behavior=function_choice_behavior, + arguments=arguments, + ) mock_create_agent.assert_called_once_with( - agentName="test_agent", + agentName=bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentName"], foundationModel=bedrock_agent_unit_test_env["BEDROCK_AGENT_FOUNDATION_MODEL"], agentResourceRoleArn=bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], instruction="test_instructions", ) assert agent.agent_model.agent_id == bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentId"] + assert agent.id == bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentId"] + assert agent.agent_model.agent_name == bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentName"] + assert agent.name == bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentName"] + assert ( + agent.agent_model.foundation_model + == bedrock_agent_model_with_id_not_prepared_dict["agent"]["foundationModel"] + ) + assert agent.kernel is not None + assert agent.function_choice_behavior is not None + if arguments: + assert agent.arguments is not None # Test case to verify the creation of BedrockAgent +@pytest.mark.parametrize( + "exclude_list", + [ + ["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], + ["BEDROCK_AGENT_FOUNDATION_MODEL"], + ], + indirect=True, +) @patch.object(boto3, "client", return_value=Mock()) -async def test_bedrock_agent_create_with_plugin_via_constructor( - client, bedrock_agent_unit_test_env, bedrock_agent_model_with_id_not_prepared_dict, custom_plugin_class +async def test_bedrock_agent_create_and_prepare_agent_settings_validation_error( + client, + bedrock_agent_model_with_id_not_prepared_dict, + bedrock_agent_unit_test_env, ): - agent = BedrockAgent( - name="test_agent", instructions="test_instructions", env_file_path="fake_path", plugins=[custom_plugin_class()] - ) + with pytest.raises(AgentInitializationException): + await BedrockAgent.create_and_prepare_agent( + name=bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentName"], + instructions="test_instructions", + env_file_path="fake_path", + ) - assert agent.agent_model.agent_id is None +# Test case to verify the creation of BedrockAgent +@patch.object(boto3, "client", return_value=Mock()) +async def test_bedrock_agent_create_and_prepare_agent_service_exception( + client, + bedrock_agent_model_with_id_not_prepared_dict, + bedrock_agent_unit_test_env, +): with ( - patch.object(agent.bedrock_client, "create_agent") as mock_create_agent, - patch.object(agent.bedrock_client, "get_agent") as mock_get_agent, - patch.object(BedrockAgent, "prepare_agent", new_callable=AsyncMock), + patch.object(client, "create_agent") as mock_create_agent, + patch.object(BedrockAgent, "prepare_agent_and_wait_until_prepared", new_callable=AsyncMock), ): - mock_create_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict - mock_get_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict + from botocore.exceptions import ClientError - await agent.create_agent() + mock_create_agent.side_effect = ClientError({}, "create_agent") - mock_create_agent.assert_called_once_with( - agentName="test_agent", - foundationModel=bedrock_agent_unit_test_env["BEDROCK_AGENT_FOUNDATION_MODEL"], - agentResourceRoleArn=bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - instruction="test_instructions", - ) - assert agent.agent_model.agent_id == bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentId"] - assert agent.kernel.plugins is not None - assert len(agent.kernel.plugins) == 1 + with pytest.raises(AgentInitializationException): + await BedrockAgent.create_and_prepare_agent( + name=bedrock_agent_model_with_id_not_prepared_dict["agent"]["agentName"], + instructions="test_instructions", + bedrock_client=client, + env_file_path="fake_path", + ) @patch.object(boto3, "client", return_value=Mock()) -async def test_bedrock_agent_create_existed( +async def test_bedrock_agent_prepare_agent_and_wait_until_prepared( client, bedrock_agent_unit_test_env, - bedrock_agent_model_with_id_not_prepared_dict, + bedrock_agent_model_with_id, + bedrock_agent_model_with_id_preparing_dict, + bedrock_agent_model_with_id_prepared_dict, ): - agent = BedrockAgent(name="test_agent", env_file_path="fake_path") + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) + + with ( + patch.object(client, "get_agent") as mock_get_agent, + patch.object(client, "prepare_agent") as mock_prepare_agent, + ): + mock_get_agent.side_effect = [ + bedrock_agent_model_with_id_preparing_dict, + bedrock_agent_model_with_id_prepared_dict, + ] + + await agent.prepare_agent_and_wait_until_prepared() - assert agent.agent_model.agent_id is None + mock_prepare_agent.assert_called_once_with(agentId=bedrock_agent_model_with_id.agent_id) + assert mock_get_agent.call_count == 2 + assert agent.agent_model.agent_status == "PREPARED" + + +@patch.object(boto3, "client", return_value=Mock()) +async def test_bedrock_agent_prepare_agent_and_wait_until_prepared_fail( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, + bedrock_agent_model_with_id_preparing_dict, +): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) with ( - patch.object(agent.bedrock_client, "create_agent") as mock_create_agent, - patch.object(agent.bedrock_client, "get_agent") as mock_get_agent, - patch.object(BedrockAgent, "prepare_agent", new_callable=AsyncMock), + patch.object(client, "get_agent") as mock_get_agent, + patch.object(client, "prepare_agent"), ): - mock_create_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict - mock_get_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict + mock_get_agent.side_effect = [ + bedrock_agent_model_with_id_preparing_dict, + bedrock_agent_model_with_id_preparing_dict, + bedrock_agent_model_with_id_preparing_dict, + bedrock_agent_model_with_id_preparing_dict, + bedrock_agent_model_with_id_preparing_dict, + bedrock_agent_model_with_id_preparing_dict, + ] + + with pytest.raises(TimeoutError): + await agent.prepare_agent_and_wait_until_prepared() - await agent.create_agent() - with pytest.raises( - ValueError, match="Agent already exists. Please delete the agent before creating a new one." - ): - await agent.create_agent() +# Test case to verify the creation of a code interpreter action group +@patch.object(boto3, "client", return_value=Mock()) +async def test_create_code_interpreter_action_group( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, + bedrock_action_group_mode_dict, +): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) + + with ( + patch.object(client, "create_agent_action_group") as mock_create_action_group, + patch.object( + BedrockAgent, "prepare_agent_and_wait_until_prepared" + ) as mock_prepare_agent_and_wait_until_prepared, + ): + mock_create_action_group.return_value = bedrock_action_group_mode_dict + action_group_model = await agent.create_code_interpreter_action_group() + + mock_create_action_group.assert_called_once_with( + agentId=agent.agent_model.agent_id, + agentVersion=agent.agent_model.agent_version or "DRAFT", + actionGroupName=f"{agent.agent_model.agent_name}_code_interpreter", + actionGroupState="ENABLED", + parentActionGroupSignature="AMAZON.CodeInterpreter", + ) + assert action_group_model.action_group_id == bedrock_action_group_mode_dict["agentActionGroup"]["actionGroupId"] + mock_prepare_agent_and_wait_until_prepared.assert_called_once() +# Test case to verify the creation of BedrockAgent with plugins @patch.object(boto3, "client", return_value=Mock()) -async def test_bedrock_agent_create_with_options( +async def test_bedrock_agent_create_with_plugin_via_constructor( client, bedrock_agent_unit_test_env, - bedrock_agent_model_with_id_not_prepared_dict, + bedrock_agent_model_with_id, + custom_plugin_class, ): - agent = BedrockAgent(name="test_agent", env_file_path="fake_path") + agent = BedrockAgent( + bedrock_agent_model_with_id, + plugins=[custom_plugin_class()], + bedrock_client=client, + ) + + assert agent.kernel.plugins is not None + assert len(agent.kernel.plugins) == 1 - assert agent.agent_model.agent_id is None + +# Test case to verify the creation of a user input action group +@patch.object(boto3, "client", return_value=Mock()) +async def test_create_user_input_action_group( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, + bedrock_action_group_mode_dict, +): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) with ( - patch.object(agent.bedrock_client, "create_agent") as mock_create_agent, - patch.object(agent.bedrock_client, "get_agent") as mock_get_agent, - patch.object(BedrockAgent, "prepare_agent", new_callable=AsyncMock), + patch.object(agent.bedrock_client, "create_agent_action_group") as mock_create_action_group, patch.object( - BedrockAgent, "create_code_interpreter_action_group", new_callable=AsyncMock - ) as mock_create_code_interpreter_action_group, + BedrockAgent, "prepare_agent_and_wait_until_prepared" + ) as mock_prepare_agent_and_wait_until_prepared, + ): + mock_create_action_group.return_value = bedrock_action_group_mode_dict + action_group_model = await agent.create_user_input_action_group() + + mock_create_action_group.assert_called_once_with( + agentId=agent.agent_model.agent_id, + agentVersion=agent.agent_model.agent_version or "DRAFT", + actionGroupName=f"{agent.agent_model.agent_name}_user_input", + actionGroupState="ENABLED", + parentActionGroupSignature="AMAZON.UserInput", + ) + assert action_group_model.action_group_id == bedrock_action_group_mode_dict["agentActionGroup"]["actionGroupId"] + mock_prepare_agent_and_wait_until_prepared.assert_called_once() + + +# Test case to verify the creation of a kernel function action group +@patch.object(boto3, "client", return_value=Mock()) +async def test_create_kernel_function_action_group( + client, + kernel_with_function, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, + bedrock_action_group_mode_dict, +): + agent = BedrockAgent(bedrock_agent_model_with_id, kernel=kernel_with_function, bedrock_client=client) + + with ( + patch.object(agent.bedrock_client, "create_agent_action_group") as mock_create_action_group, patch.object( - BedrockAgent, "create_user_input_action_group", new_callable=AsyncMock - ) as mock_create_user_input_action_group, + BedrockAgent, "prepare_agent_and_wait_until_prepared" + ) as mock_prepare_agent_and_wait_until_prepared, + ): + mock_create_action_group.return_value = bedrock_action_group_mode_dict + + action_group_model = await agent.create_kernel_function_action_group() + + mock_create_action_group.assert_called_once_with( + agentId=agent.agent_model.agent_id, + agentVersion=agent.agent_model.agent_version or "DRAFT", + actionGroupName=f"{agent.agent_model.agent_name}_kernel_function", + actionGroupState="ENABLED", + actionGroupExecutor={"customControl": "RETURN_CONTROL"}, + functionSchema=kernel_function_to_bedrock_function_schema( + agent.function_choice_behavior.get_config(kernel_with_function) + ), + ) + assert action_group_model.action_group_id == bedrock_action_group_mode_dict["agentActionGroup"]["actionGroupId"] + mock_prepare_agent_and_wait_until_prepared.assert_called_once() + + +# Test case to verify the association of an agent with a knowledge base +@patch.object(boto3, "client", return_value=Mock()) +async def test_associate_agent_knowledge_base( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, +): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) + + with ( + patch.object(agent.bedrock_client, "associate_agent_knowledge_base") as mock_associate_knowledge_base, patch.object( - BedrockAgent, "_create_kernel_function_action_group", new_callable=AsyncMock - ) as mock_create_kernel_function_action_group, + BedrockAgent, "prepare_agent_and_wait_until_prepared" + ) as mock_prepare_agent_and_wait_until_prepared, ): - mock_create_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict - mock_get_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict + await agent.associate_agent_knowledge_base("test_knowledge_base_id") - await agent.create_agent( - enable_code_interpreter=True, - enable_user_input=True, - enable_kernel_function=True, + mock_associate_knowledge_base.assert_called_once_with( + agentId=agent.agent_model.agent_id, + agentVersion=agent.agent_model.agent_version, + knowledgeBaseId="test_knowledge_base_id", ) - - mock_create_code_interpreter_action_group.assert_called_once() - mock_create_user_input_action_group.assert_called_once() - mock_create_kernel_function_action_group.assert_called_once() + mock_prepare_agent_and_wait_until_prepared.assert_called_once() -# Test case to verify the retrieval of BedrockAgent +# Test case to verify the disassociation of an agent with a knowledge base @patch.object(boto3, "client", return_value=Mock()) -async def test_bedrock_agent_retrieve( +async def test_disassociate_agent_knowledge_base( client, bedrock_agent_unit_test_env, bedrock_agent_model_with_id, ): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) + with ( - patch.object(BedrockAgent, "_get_agent") as mock_get_agent, + patch.object(agent.bedrock_client, "disassociate_agent_knowledge_base") as mock_disassociate_knowledge_base, + patch.object( + BedrockAgent, "prepare_agent_and_wait_until_prepared" + ) as mock_prepare_agent_and_wait_until_prepared, ): - mock_get_agent.return_value = bedrock_agent_model_with_id + await agent.disassociate_agent_knowledge_base("test_knowledge_base_id") + mock_disassociate_knowledge_base.assert_called_once_with( + agentId=agent.agent_model.agent_id, + agentVersion=agent.agent_model.agent_version, + knowledgeBaseId="test_knowledge_base_id", + ) + mock_prepare_agent_and_wait_until_prepared.assert_called_once() - agent = await BedrockAgent.retrieve( - bedrock_agent_model_with_id.agent_id, - bedrock_agent_model_with_id.agent_name, - env_file_path="fake_path", + +# Test case to verify listing associated knowledge bases with an agent +@patch.object(boto3, "client", return_value=Mock()) +async def test_list_associated_agent_knowledge_bases( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, +): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) + + with patch.object(agent.bedrock_client, "list_agent_knowledge_bases") as mock_list_knowledge_bases: + await agent.list_associated_agent_knowledge_bases() + + mock_list_knowledge_bases.assert_called_once_with( + agentId=agent.agent_model.agent_id, + agentVersion=agent.agent_model.agent_version, ) - mock_get_agent.assert_called_once() - assert agent.agent_model.agent_id == bedrock_agent_model_with_id.agent_id - assert agent.id == agent.agent_model.agent_id + +# endregion + +# region Agent Deletion Tests + + +@patch.object(boto3, "client", return_value=Mock()) +async def test_delete_agent( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, +): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) + + agent_id = bedrock_agent_model_with_id.agent_id + with patch.object(agent.bedrock_client, "delete_agent") as mock_delete_agent: + await agent.delete_agent() + + mock_delete_agent.assert_called_once_with(agentId=agent_id) + assert agent.agent_model.agent_id is None + + +# Test case to verify error handling when deleting an agent that does not exist +@patch.object(boto3, "client", return_value=Mock()) +async def test_delete_agent_twice_error( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, +): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) + + with patch.object(agent.bedrock_client, "delete_agent"): + await agent.delete_agent() + + with pytest.raises(ValueError): + await agent.delete_agent() + + +# Test case to verify error handling when there is a client error during agent deletion +@patch.object(boto3, "client", return_value=Mock()) +async def test_delete_agent_client_error( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, +): + agent = BedrockAgent(bedrock_agent_model_with_id, bedrock_client=client) + + with patch.object(agent.bedrock_client, "delete_agent") as mock_delete_agent: + from botocore.exceptions import ClientError + + mock_delete_agent.side_effect = ClientError({"Error": {"Code": "500"}}, "delete_agent") + + with pytest.raises(ClientError): + await agent.delete_agent() + + +# endregion + +# region Agent Invoke Tests # Test case to verify the `get_response` method of BedrockAgent @@ -216,16 +445,9 @@ async def test_bedrock_agent_get_response( simple_response, ): with ( - patch.object(BedrockAgent, "_get_agent") as mock_get_agent, patch.object(BedrockAgent, "_invoke_agent", new_callable=AsyncMock) as mock_invoke_agent, ): - mock_get_agent.return_value = bedrock_agent_model_with_id - - agent = await BedrockAgent.retrieve( - bedrock_agent_model_with_id.agent_id, - bedrock_agent_model_with_id.agent_name, - env_file_path="fake_path", - ) + agent = BedrockAgent(bedrock_agent_model_with_id) mock_invoke_agent.return_value = bedrock_agent_non_streaming_simple_response response = await agent.get_response("test_session_id", "test_input_text") @@ -249,16 +471,9 @@ async def test_bedrock_agent_get_response_exception( bedrock_agent_non_streaming_empty_response, ): with ( - patch.object(BedrockAgent, "_get_agent") as mock_get_agent, patch.object(BedrockAgent, "_invoke_agent", new_callable=AsyncMock) as mock_invoke_agent, ): - mock_get_agent.return_value = bedrock_agent_model_with_id - - agent = await BedrockAgent.retrieve( - bedrock_agent_model_with_id.agent_id, - bedrock_agent_model_with_id.agent_name, - env_file_path="fake_path", - ) + agent = BedrockAgent(bedrock_agent_model_with_id) mock_invoke_agent.return_value = bedrock_agent_non_streaming_empty_response with pytest.raises(AgentInvokeException): @@ -283,16 +498,9 @@ async def test_bedrock_agent_invoke( simple_response, ): with ( - patch.object(BedrockAgent, "_get_agent") as mock_get_agent, patch.object(BedrockAgent, "_invoke_agent", new_callable=AsyncMock) as mock_invoke_agent, ): - mock_get_agent.return_value = bedrock_agent_model_with_id - - agent = await BedrockAgent.retrieve( - bedrock_agent_model_with_id.agent_id, - bedrock_agent_model_with_id.agent_name, - env_file_path="fake_path", - ) + agent = BedrockAgent(bedrock_agent_model_with_id) mock_invoke_agent.return_value = bedrock_agent_non_streaming_simple_response async for message in agent.invoke("test_session_id", "test_input_text"): @@ -317,16 +525,9 @@ async def test_bedrock_agent_invoke_stream( simple_response, ): with ( - patch.object(BedrockAgent, "_get_agent") as mock_get_agent, patch.object(BedrockAgent, "_invoke_agent", new_callable=AsyncMock) as mock_invoke_agent, ): - mock_get_agent.return_value = bedrock_agent_model_with_id - - agent = await BedrockAgent.retrieve( - bedrock_agent_model_with_id.agent_id, - bedrock_agent_model_with_id.agent_name, - env_file_path="fake_path", - ) + agent = BedrockAgent(bedrock_agent_model_with_id) mock_invoke_agent.return_value = bedrock_agent_streaming_simple_response full_message = "" @@ -353,17 +554,10 @@ async def test_bedrock_agent_invoke_with_function_call( bedrock_agent_non_streaming_simple_response, ): with ( - patch.object(BedrockAgent, "_get_agent") as mock_get_agent, patch.object(BedrockAgent, "_invoke_agent", new_callable=AsyncMock) as mock_invoke_agent, patch.object(BedrockAgent, "_handle_function_call_contents") as mock_handle_function_call_contents, ): - mock_get_agent.return_value = bedrock_agent_model_with_id - - agent = await BedrockAgent.retrieve( - bedrock_agent_model_with_id.agent_id, - bedrock_agent_model_with_id.agent_name, - env_file_path="fake_path", - ) + agent = BedrockAgent(bedrock_agent_model_with_id) function_result_contents = [ FunctionResultContent( @@ -403,17 +597,10 @@ async def test_bedrock_agent_invoke_stream_with_function_call( bedrock_agent_streaming_simple_response, ): with ( - patch.object(BedrockAgent, "_get_agent") as mock_get_agent, patch.object(BedrockAgent, "_invoke_agent", new_callable=AsyncMock) as mock_invoke_agent, patch.object(BedrockAgent, "_handle_function_call_contents") as mock_handle_function_call_contents, ): - mock_get_agent.return_value = bedrock_agent_model_with_id - - agent = await BedrockAgent.retrieve( - bedrock_agent_model_with_id.agent_id, - bedrock_agent_model_with_id.agent_name, - env_file_path="fake_path", - ) + agent = BedrockAgent(bedrock_agent_model_with_id) function_result_contents = [ FunctionResultContent( @@ -441,3 +628,6 @@ async def test_bedrock_agent_invoke_stream_with_function_call( "returnControlInvocationResults": parse_function_result_contents(function_result_contents), }, ) + + +# endregion diff --git a/python/tests/unit/agents/bedrock_agent/test_bedrock_agent_base.py b/python/tests/unit/agents/bedrock_agent/test_bedrock_agent_base.py deleted file mode 100644 index 6d21752b1daa..000000000000 --- a/python/tests/unit/agents/bedrock_agent/test_bedrock_agent_base.py +++ /dev/null @@ -1,800 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -from unittest.mock import Mock, patch - -import boto3 -import pytest -from botocore.exceptions import ClientError - -from semantic_kernel.agents.bedrock.action_group_utils import kernel_function_to_bedrock_function_schema -from semantic_kernel.agents.bedrock.bedrock_agent_base import BedrockAgentBase -from semantic_kernel.agents.bedrock.models.bedrock_agent_status import BedrockAgentStatus - - -# Test case to verify the creation of an agent -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_agent( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, - bedrock_agent_model_with_id_not_prepared_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with ( - patch.object(bedrock_agent_base.bedrock_client, "create_agent") as mock_create_agent, - patch.object(bedrock_agent_base.bedrock_client, "get_agent") as mock_get_agent, - ): - mock_create_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict - mock_get_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict - await bedrock_agent_base._create_agent("test_instruction") - - assert bedrock_agent_base.agent_model.agent_id == "test_agent_id" - - -# Test case to verify error handling when creating an agent that already exists -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_agent_already_exists( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with pytest.raises(ValueError, match="Agent already exists. Please delete the agent before creating a new one."): - await bedrock_agent_base._create_agent("test_instruction") - - -# Test case to verify error handling when there is a client error during agent creation -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_agent_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "create_agent") as mock_create_agent: - mock_create_agent.side_effect = ClientError({"Error": {"Code": "500"}}, "create_agent") - with pytest.raises(ClientError): - await bedrock_agent_base._create_agent("test_instruction") - - -# Test case to verify the preparation of an agent -@patch.object(boto3, "client", return_value=Mock()) -async def test_prepare_agent( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, - bedrock_agent_model_with_id_prepared_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with ( - patch.object(bedrock_agent_base.bedrock_client, "prepare_agent") as mock_prepare_agent, - patch.object(bedrock_agent_base.bedrock_client, "get_agent") as mock_get_agent, - ): - mock_get_agent.return_value = bedrock_agent_model_with_id_prepared_dict - await bedrock_agent_base.prepare_agent() - - mock_prepare_agent.assert_called_once_with(agentId=bedrock_agent_base.agent_model.agent_id) - bedrock_agent_base.agent_model.agent_status = BedrockAgentStatus.PREPARED - - -# Test case to verify error handling when preparing an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_prepare_agent_no_id(client, bedrock_agent_unit_test_env, bedrock_agent_model): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - bedrock_agent_base.agent_model.agent_id = None - with pytest.raises(ValueError, match="Agent does not exist. Please create the agent before preparing it."): - await bedrock_agent_base.prepare_agent() - - -# Test case to verify error handling when there is a client error during agent preparation -@patch.object(boto3, "client", return_value=Mock()) -async def test_prepare_agent_client_error(client, bedrock_agent_unit_test_env, bedrock_agent_model_with_id): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "prepare_agent") as mock_prepare_agent: - mock_prepare_agent.side_effect = ClientError({"Error": {"Code": "500"}}, "prepare_agent") - with pytest.raises(ClientError): - await bedrock_agent_base.prepare_agent() - - -# Test case to verify the update of an agent -@patch.object(boto3, "client", return_value=Mock()) -async def test_update_agent( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, - bedrock_agent_model_with_id_prepared_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - new_agent_name = "new_agent_name" - bedrock_agent_model_with_id_prepared_dict["agent"]["agentName"] = new_agent_name - - with ( - patch.object(bedrock_agent_base.bedrock_client, "update_agent") as mock_update_agent, - patch.object(bedrock_agent_base.bedrock_client, "get_agent") as mock_get_agent, - ): - mock_update_agent.return_value = bedrock_agent_model_with_id_prepared_dict - mock_get_agent.return_value = bedrock_agent_model_with_id_prepared_dict - await bedrock_agent_base.update_agent(agentName=new_agent_name) - - mock_update_agent.assert_called_once_with( - agentId=bedrock_agent_base.agent_model.agent_id, - agentResourceRoleArn=bedrock_agent_base.agent_resource_role_arn, - agentName=new_agent_name, - foundationModel=bedrock_agent_base.agent_model.foundation_model, - ) - assert bedrock_agent_base.agent_model.agent_status == BedrockAgentStatus.PREPARED - - -# Test case to verify error handling when updating an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_update_agent_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises(ValueError, match="Agent has not been created. Please create the agent before updating it."): - await bedrock_agent_base.update_agent() - - -# Test case to verify error handling when there is a client error during agent update -@patch.object(boto3, "client", return_value=Mock()) -async def test_update_agent_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "update_agent") as mock_update_agent: - mock_update_agent.side_effect = ClientError({"Error": {"Code": "500"}}, "update_agent") - with pytest.raises(ClientError): - await bedrock_agent_base.update_agent() - - -# Test case to verify the deletion of an agent -@patch.object(boto3, "client", return_value=Mock()) -async def test_delete_agent( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - agent_id = bedrock_agent_base.agent_model.agent_id - with patch.object(bedrock_agent_base.bedrock_client, "delete_agent") as mock_delete_agent: - await bedrock_agent_base.delete_agent() - - mock_delete_agent.assert_called_once_with(agentId=agent_id) - assert bedrock_agent_base.agent_model.agent_id is None - - -# Test case to verify error handling when deleting an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_delete_agent_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - bedrock_agent_base.agent_model.agent_id = None - with pytest.raises(ValueError, match="Agent does not exist. Please create the agent before deleting it."): - await bedrock_agent_base.delete_agent() - - -# Test case to verify error handling when there is a client error during agent deletion -@patch.object(boto3, "client", return_value=Mock()) -async def test_delete_agent_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "delete_agent") as mock_delete_agent: - mock_delete_agent.side_effect = ClientError({"Error": {"Code": "500"}}, "delete_agent") - with pytest.raises(ClientError): - await bedrock_agent_base.delete_agent() - - -# Test case to verify the retrieval of an agent -@patch.object(boto3, "client", return_value=Mock()) -async def test_get_agent( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, - bedrock_agent_model_with_id_prepared_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "get_agent") as mock_get_agent: - mock_get_agent.return_value = bedrock_agent_model_with_id_prepared_dict - await bedrock_agent_base._get_agent() - - mock_get_agent.assert_called_once_with(agentId=bedrock_agent_base.agent_model.agent_id) - - -# Test case to verify error handling when retrieving an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_get_agent_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises(ValueError, match="Agent does not exist. Please create the agent before getting it."): - await bedrock_agent_base._get_agent() - - -# Test case to verify error handling when there is a client error during agent retrieval -@patch.object(boto3, "client", return_value=Mock()) -async def test_get_agent_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "get_agent") as mock_get_agent: - mock_get_agent.side_effect = ClientError({"Error": {"Code": "500"}}, "get_agent") - with pytest.raises(ClientError): - await bedrock_agent_base._get_agent() - - -# Test case to verify waiting for an agent to reach a specific status -@patch.object(boto3, "client", return_value=Mock()) -async def test_wait_for_agent_status( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, - bedrock_agent_model_with_id_prepared_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "get_agent") as mock_get_agent: - mock_get_agent.return_value = bedrock_agent_model_with_id_prepared_dict - - await bedrock_agent_base._wait_for_agent_status(BedrockAgentStatus.PREPARED) - - mock_get_agent.assert_called_once() - - -# Test case to verify error handling when waiting for an agent to reach a specific status times out -@patch.object(boto3, "client", return_value=Mock()) -async def test_wait_for_agent_status_timeout( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, - bedrock_agent_model_with_id_not_prepared_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "get_agent") as mock_get_agent: - mock_get_agent.return_value = bedrock_agent_model_with_id_not_prepared_dict - - with pytest.raises( - TimeoutError, - match=f"Agent did not reach status {BedrockAgentStatus.PREPARED} within the specified time.", - ): - await bedrock_agent_base._wait_for_agent_status(BedrockAgentStatus.PREPARED, max_attempts=3) - - assert mock_get_agent.call_count == 3 - - -# Test case to verify the creation of a code interpreter action group -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_code_interpreter_action_group( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, - bedrock_action_group_mode_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "create_agent_action_group") as mock_create_action_group: - mock_create_action_group.return_value = bedrock_action_group_mode_dict - action_group_model = await bedrock_agent_base.create_code_interpreter_action_group() - - mock_create_action_group.assert_called_once_with( - agentId=bedrock_agent_base.agent_model.agent_id, - agentVersion=bedrock_agent_base.agent_model.agent_version or "DRAFT", - actionGroupName=f"{bedrock_agent_base.agent_model.agent_name}_code_interpreter", - actionGroupState="ENABLED", - parentActionGroupSignature="AMAZON.CodeInterpreter", - ) - assert action_group_model.action_group_id == bedrock_action_group_mode_dict["agentActionGroup"]["actionGroupId"] - - -# Test case to verify error handling when creating a code interpreter action group for an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_code_interpreter_action_group_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises( - ValueError, match="Agent does not exist. Please create the agent before creating an action group for it." - ): - await bedrock_agent_base.create_code_interpreter_action_group() - - -# Test case to verify error handling when there is a client error during code interpreter action group creation -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_code_interpreter_action_group_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "create_agent_action_group") as mock_create_action_group: - mock_create_action_group.side_effect = ClientError({"Error": {"Code": "500"}}, "create_agent_action_group") - with pytest.raises(ClientError): - await bedrock_agent_base.create_code_interpreter_action_group() - - -# Test case to verify the creation of a user input action group -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_user_input_action_group( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, - bedrock_action_group_mode_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "create_agent_action_group") as mock_create_action_group: - mock_create_action_group.return_value = bedrock_action_group_mode_dict - action_group_model = await bedrock_agent_base.create_user_input_action_group() - - mock_create_action_group.assert_called_once_with( - agentId=bedrock_agent_base.agent_model.agent_id, - agentVersion=bedrock_agent_base.agent_model.agent_version or "DRAFT", - actionGroupName=f"{bedrock_agent_base.agent_model.agent_name}_user_input", - actionGroupState="ENABLED", - parentActionGroupSignature="AMAZON.UserInput", - ) - assert action_group_model.action_group_id == bedrock_action_group_mode_dict["agentActionGroup"]["actionGroupId"] - - -# Test case to verify error handling when creating a user input action group for an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_user_input_action_group_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises( - ValueError, match="Agent does not exist. Please create the agent before creating an action group for it." - ): - await bedrock_agent_base.create_user_input_action_group() - - -# Test case to verify error handling when there is a client error during user input action group creation -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_user_input_action_group_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "create_agent_action_group") as mock_create_action_group: - mock_create_action_group.side_effect = ClientError({"Error": {"Code": "500"}}, "create_agent_action_group") - with pytest.raises(ClientError): - await bedrock_agent_base.create_user_input_action_group() - - -# Test case to verify the creation of a kernel function action group -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_kernel_function_action_group( - client, - kernel_with_function, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, - bedrock_action_group_mode_dict, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "create_agent_action_group") as mock_create_action_group: - mock_create_action_group.return_value = bedrock_action_group_mode_dict - - action_group_model = await bedrock_agent_base._create_kernel_function_action_group(kernel_with_function) - - mock_create_action_group.assert_called_once_with( - agentId=bedrock_agent_base.agent_model.agent_id, - agentVersion=bedrock_agent_base.agent_model.agent_version or "DRAFT", - actionGroupName=f"{bedrock_agent_base.agent_model.agent_name}_kernel_function", - actionGroupState="ENABLED", - actionGroupExecutor={"customControl": "RETURN_CONTROL"}, - functionSchema=kernel_function_to_bedrock_function_schema( - bedrock_agent_base.function_choice_behavior.get_config(kernel_with_function) - ), - ) - assert action_group_model.action_group_id == bedrock_action_group_mode_dict["agentActionGroup"]["actionGroupId"] - - -# Test case to verify error handling when creating a kernel function action group for an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_kernel_function_action_group_no_id( - client, - kernel_with_function, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises( - ValueError, match="Agent does not exist. Please create the agent before creating an action group for it." - ): - await bedrock_agent_base._create_kernel_function_action_group(kernel_with_function) - - -# Test case to verify error handling when there are no available functions for the kernel function action group -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_kernel_function_action_group_no_functions( - client, - kernel, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "create_agent_action_group") as mock_create_action_group: - action_group_model = await bedrock_agent_base._create_kernel_function_action_group(kernel) - - assert action_group_model is None - mock_create_action_group.assert_not_called() - - -# Test case to verify error handling when there is a client error during kernel function action group creation -@patch.object(boto3, "client", return_value=Mock()) -async def test_create_kernel_function_action_group_client_error( - client, - kernel_with_function, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "create_agent_action_group") as mock_create_action_group: - mock_create_action_group.side_effect = ClientError({"Error": {"Code": "500"}}, "create_agent_action_group") - with pytest.raises(ClientError): - await bedrock_agent_base._create_kernel_function_action_group(kernel_with_function) - - -# Test case to verify the association of an agent with a knowledge base -@patch.object(boto3, "client", return_value=Mock()) -async def test_associate_agent_knowledge_base( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object( - bedrock_agent_base.bedrock_client, "associate_agent_knowledge_base" - ) as mock_associate_knowledge_base: - await bedrock_agent_base.associate_agent_knowledge_base("test_knowledge_base_id") - - mock_associate_knowledge_base.assert_called_once_with( - agentId=bedrock_agent_base.agent_model.agent_id, - agentVersion=bedrock_agent_base.agent_model.agent_version, - knowledgeBaseId="test_knowledge_base_id", - ) - - -# Test case to verify error handling when associating an agent with a knowledge base that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_associate_agent_knowledge_base_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises( - ValueError, match="Agent does not exist. Please create the agent before associating it with a knowledge base." - ): - await bedrock_agent_base.associate_agent_knowledge_base("test_knowledge_base_id") - - -# Test case to verify error handling when there is a client error during knowledge base association -@patch.object(boto3, "client", return_value=Mock()) -async def test_associate_agent_knowledge_base_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object( - bedrock_agent_base.bedrock_client, "associate_agent_knowledge_base" - ) as mock_associate_knowledge_base: - mock_associate_knowledge_base.side_effect = ClientError( - {"Error": {"Code": "500"}}, "associate_agent_knowledge_base" - ) - with pytest.raises(ClientError): - await bedrock_agent_base.associate_agent_knowledge_base("test_knowledge_base_id") - - -# Test case to verify the disassociation of an agent with a knowledge base -@patch.object(boto3, "client", return_value=Mock()) -async def test_disassociate_agent_knowledge_base( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object( - bedrock_agent_base.bedrock_client, "disassociate_agent_knowledge_base" - ) as mock_disassociate_knowledge_base: - await bedrock_agent_base.disassociate_agent_knowledge_base("test_knowledge_base_id") - mock_disassociate_knowledge_base.assert_called_once_with( - agentId=bedrock_agent_base.agent_model.agent_id, - agentVersion=bedrock_agent_base.agent_model.agent_version, - knowledgeBaseId="test_knowledge_base_id", - ) - - -# Test case to verify error handling when disassociating an agent with a knowledge base that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_disassociate_agent_knowledge_base_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises( - ValueError, - match="Agent does not exist. Please create the agent before disassociating it with a knowledge base.", - ): - await bedrock_agent_base.disassociate_agent_knowledge_base("test_knowledge_base_id") - - -# Test case to verify error handling when there is a client error during knowledge base disassociation -@patch.object(boto3, "client", return_value=Mock()) -async def test_disassociate_agent_knowledge_base_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object( - bedrock_agent_base.bedrock_client, "disassociate_agent_knowledge_base" - ) as mock_disassociate_knowledge_base: - mock_disassociate_knowledge_base.side_effect = ClientError( - {"Error": {"Code": "500"}}, "disassociate_agent_knowledge_base" - ) - with pytest.raises(ClientError): - await bedrock_agent_base.disassociate_agent_knowledge_base("test_knowledge_base_id") - - -# Test case to verify listing associated knowledge bases with an agent -@patch.object(boto3, "client", return_value=Mock()) -async def test_list_associated_agent_knowledge_bases( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "list_agent_knowledge_bases") as mock_list_knowledge_bases: - await bedrock_agent_base.list_associated_agent_knowledge_bases() - - mock_list_knowledge_bases.assert_called_once_with( - agentId=bedrock_agent_base.agent_model.agent_id, - agentVersion=bedrock_agent_base.agent_model.agent_version, - ) - - -# Test case to verify error handling when listing associated knowledge bases for an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_list_associated_agent_knowledge_bases_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises( - ValueError, match="Agent does not exist. Please create the agent before listing associated knowledge bases." - ): - await bedrock_agent_base.list_associated_agent_knowledge_bases() - - -# Test case to verify error handling when there is a client error during listing associated knowledge bases -@patch.object(boto3, "client", return_value=Mock()) -async def test_list_associated_agent_knowledge_bases_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_client, "list_agent_knowledge_bases") as mock_list_knowledge_bases: - mock_list_knowledge_bases.side_effect = ClientError({"Error": {"Code": "500"}}, "list_agent_knowledge_bases") - with pytest.raises(ClientError): - await bedrock_agent_base.list_associated_agent_knowledge_bases() - - -# Test case to verify invoking an agent -@patch.object(boto3, "client", return_value=Mock()) -async def test_invoke_agent( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_runtime_client, "invoke_agent") as mock_invoke_agent: - await bedrock_agent_base._invoke_agent("test_session_id", "test_input_text") - - mock_invoke_agent.assert_called_once_with( - agentAliasId=bedrock_agent_base.WORKING_DRAFT_AGENT_ALIAS, - agentId=bedrock_agent_base.agent_model.agent_id, - sessionId="test_session_id", - inputText="test_input_text", - ) - - -# Test case to verify error handling when invoking an agent that does not exist -@patch.object(boto3, "client", return_value=Mock()) -async def test_invoke_agent_no_id( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model, - ) - - with pytest.raises(ValueError, match="Agent does not exist. Please create the agent before invoking it."): - await bedrock_agent_base._invoke_agent("test_session_id", "test_input_text") - - -# Test case to verify error handling when there is a client error during agent invocation -@patch.object(boto3, "client", return_value=Mock()) -async def test_invoke_agent_client_error( - client, - bedrock_agent_unit_test_env, - bedrock_agent_model_with_id, -): - bedrock_agent_base = BedrockAgentBase( - bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], - bedrock_agent_model_with_id, - ) - - with patch.object(bedrock_agent_base.bedrock_runtime_client, "invoke_agent") as mock_invoke_agent: - mock_invoke_agent.side_effect = ClientError({"Error": {"Code": "500"}}, "invoke_agent") - with pytest.raises(ClientError): - await bedrock_agent_base._invoke_agent("test_session_id", "test_input_text") diff --git a/python/tests/unit/agents/bedrock_agent/test_bedrock_agent_settings.py b/python/tests/unit/agents/bedrock_agent/test_bedrock_agent_settings.py index d8719b1e95c2..c56e3fcb878f 100644 --- a/python/tests/unit/agents/bedrock_agent/test_bedrock_agent_settings.py +++ b/python/tests/unit/agents/bedrock_agent/test_bedrock_agent_settings.py @@ -14,16 +14,14 @@ def test_bedrock_agent_settings_from_env_vars(bedrock_agent_unit_test_env): assert settings.foundation_model == bedrock_agent_unit_test_env["BEDROCK_AGENT_FOUNDATION_MODEL"] -@pytest.mark.parametrize("exclude_list", [["BEDROCK_AGENT_FOUNDATION_MODEL"]], indirect=True) -def test_bedrock_agent_settings_from_env_vars_missing_optional(bedrock_agent_unit_test_env): - """Test loading BedrockAgentSettings from environment variables with missing optional fields.""" - settings = BedrockAgentSettings.create(env_file_path="fake_path") - - assert settings.agent_resource_role_arn == bedrock_agent_unit_test_env["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"] - assert settings.foundation_model is None - - -@pytest.mark.parametrize("exclude_list", [["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"]], indirect=True) +@pytest.mark.parametrize( + "exclude_list", + [ + ["BEDROCK_AGENT_AGENT_RESOURCE_ROLE_ARN"], + ["BEDROCK_AGENT_FOUNDATION_MODEL"], + ], + indirect=True, +) def test_bedrock_agent_settings_from_env_vars_missing_required(bedrock_agent_unit_test_env): """Test loading BedrockAgentSettings from environment variables with missing required fields.""" with pytest.raises(ValidationError):