From d89c84b80ed0b737deb42938e41dae72454c79ff Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Thu, 27 Feb 2025 10:14:32 -0800 Subject: [PATCH 1/7] Fix bedrock agent, need to investigate integration tests --- .../bedrock_agent_with_code_interpreter.py | 3 ++- .../agents/bedrock/bedrock_agent_base.py | 7 ++++++- .../utils/telemetry/agent_diagnostics/decorators.py | 5 ++++- .../bedrock_agent/test_bedrock_agent_integration.py | 10 ++++++---- 4 files changed, 18 insertions(+), 7 deletions(-) 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..eafc3cf08847 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 @@ -48,7 +48,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/semantic_kernel/agents/bedrock/bedrock_agent_base.py b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py index 6fe432af9f0f..9250173abc6e 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py @@ -146,6 +146,11 @@ 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}.") @@ -198,7 +203,7 @@ async def update_agent(self, **kwargs) -> None: ) self.agent_model = BedrockAgentModel(**response["agent"]) - await self._wait_for_agent_status(BedrockAgentStatus.PREPARED) + await self.prepare_agent() except ClientError as e: logger.error(f"Failed to update agent {self.agent_model.agent_id}.") raise e 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..78197ea6bc7b 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 @@ -48,7 +48,7 @@ async def setup_and_teardown(self, request): @pytest.mark.asyncio async def test_update(self): """Test updating the agent.""" - await self.bedrock_agent.update_agent(agentName="updated_agent") + await self.bedrock_agent.update_agent(agentName=f"semantic-kernel-integration-test-agent-{uuid.uuid4()}") assert self.bedrock_agent.agent_model.agent_name == "updated_agent" @pytest.mark.asyncio @@ -79,12 +79,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) From 2f579c85cac466b0d7ba3e64026d0a640e1f9a3a Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Thu, 27 Feb 2025 16:40:42 -0800 Subject: [PATCH 2/7] Cleanning up Bedrock Agent --- .../concepts/agents/bedrock_agent/README.md | 1 - .../bedrock_agent_update_agent.py | 55 ---- .../agents/bedrock/bedrock_agent.py | 240 +++++------------- .../agents/bedrock/bedrock_agent_base.py | 126 ++------- .../ai/bedrock/services/bedrock_base.py | 2 +- .../services/bedrock_chat_completion.py | 3 +- .../services/bedrock_text_completion.py | 3 +- .../services/bedrock_text_embedding.py | 3 +- .../bedrock/services/model_provider/utils.py | 7 - python/semantic_kernel/utils/async_utils.py | 11 + .../test_bedrock_agent_integration.py | 6 - 11 files changed, 105 insertions(+), 352 deletions(-) delete mode 100644 python/samples/concepts/agents/bedrock_agent/bedrock_agent_update_agent.py create mode 100644 python/semantic_kernel/utils/async_utils.py diff --git a/python/samples/concepts/agents/bedrock_agent/README.md b/python/samples/concepts/agents/bedrock_agent/README.md index 2e759a18f919..586b210210a4 100644 --- a/python/samples/concepts/agents/bedrock_agent/README.md +++ b/python/samples/concepts/agents/bedrock_agent/README.md @@ -17,7 +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_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/semantic_kernel/agents/bedrock/bedrock_agent.py b/python/semantic_kernel/agents/bedrock/bedrock_agent.py index cc7cd4248f78..dd98a512f920 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent.py @@ -2,14 +2,18 @@ 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 +from semantic_kernel.agents.bedrock.models.bedrock_agent_status import BedrockAgentStatus +from semantic_kernel.utils.async_utils import run_in_executor + if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -36,16 +40,15 @@ 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.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.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): @@ -58,101 +61,62 @@ class BedrockAgent(BedrockAgentBase, Agent): def __init__( self, - name: str, + agent_model: BedrockAgentModel | dict[str, Any], *, - id: str | None = None, - agent_resource_role_arn: str | None = None, - foundation_model: str | None = None, - kernel: Kernel | None = None, function_choice_behavior: FunctionChoiceBehavior | None = None, + kernel: Kernel | None = None, arguments: KernelArguments | None = None, - instructions: str | None = None, - prompt_template_config: PromptTemplateConfig | None = None, - env_file_path: str | None = None, - env_file_encoding: str | 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 (str): The name of the agent. - id (str, optional): The unique identifier of the agent. - agent_resource_role_arn (str, optional): The ARN of the agent resource role. - Overrides the one in the env file. - foundation_model (str, optional): The foundation model. Overrides the one in the env file. - kernel (Kernel, optional): The kernel to use. + 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. + kernel (Kernel, optional): The kernel to use. arguments (KernelArguments, optional): The kernel arguments. Invoke method arguments take precedence over the arguments provided here. - instructions (str, optional): The instructions for the agent. - prompt_template_config (PromptTemplateConfig, optional): The prompt template configuration. - Cannot be set if instructions is set. - env_file_path (str, optional): The path to the environment file. - env_file_encoding (str, optional): The encoding of the environment file. + 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, } - 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 arguments: args["arguments"] = arguments - if prompt_template: - args["prompt_template"] = prompt_template + 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, 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": @@ -162,89 +126,67 @@ 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. function_choice_behavior (FunctionChoiceBehavior, optional): The function choice behavior for accessing the kernel functions and filters. 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 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, function_choice_behavior=function_choice_behavior, - env_file_path=env_file_path, - env_file_encoding=env_file_encoding, + kernel=kernel, + 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 @@ -262,48 +204,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 9250173abc6e..888a24ed4b5a 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 @@ -17,6 +16,7 @@ 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__) @@ -41,13 +41,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 +54,16 @@ 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. """ args = { - "agent_resource_role_arn": agent_resource_role_arn, - "agent_model": agent_model, + "agent_model": agent_model + if isinstance(agent_model, BedrockAgentModel) + else BedrockAgentModel(**agent_model), "bedrock_runtime_client": bedrock_runtime_client or boto3.client("bedrock-agent-runtime"), "bedrock_client": bedrock_client or boto3.client("bedrock-agent"), **kwargs, @@ -89,56 +86,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, @@ -156,65 +116,13 @@ async def prepare_agent(self) -> None: 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.prepare_agent() - 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, @@ -234,7 +142,7 @@ async def _get_agent(self) -> BedrockAgentModel: 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, @@ -272,7 +180,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, @@ -296,7 +204,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, @@ -325,7 +233,7 @@ async def _create_kernel_function_action_group(self, kernel: Kernel, **kwargs) - return None try: - response = await self._run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.create_agent_action_group, @@ -356,7 +264,7 @@ async def associate_agent_knowledge_base(self, knowledge_base_id: str, **kwargs) ) try: - return await self._run_in_executor( + return await run_in_executor( None, partial( self.bedrock_client.associate_agent_knowledge_base, @@ -380,7 +288,7 @@ async def disassociate_agent_knowledge_base(self, knowledge_base_id: str, **kwar ) try: - await self._run_in_executor( + await run_in_executor( None, partial( self.bedrock_client.disassociate_agent_knowledge_base, @@ -402,7 +310,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, @@ -431,7 +339,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/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..24daf5a115bb 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 @@ -7,6 +7,8 @@ import boto3 +from semantic_kernel.utils.async_utils import run_in_executor + if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -25,7 +27,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 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..4d4937ebddba 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 @@ -9,6 +9,8 @@ import boto3 from pydantic import ValidationError +from semantic_kernel.utils.async_utils import run_in_executor + if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -22,7 +24,6 @@ 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 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..d79b6a345c1a 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 @@ -10,6 +10,8 @@ from numpy import array, ndarray from pydantic import ValidationError +from semantic_kernel.utils.async_utils import run_in_executor + if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -24,7 +26,6 @@ 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 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..349e032c481f --- /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(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)) 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 78197ea6bc7b..4fbadfa361c4 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 @@ -45,12 +45,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=f"semantic-kernel-integration-test-agent-{uuid.uuid4()}") - assert self.bedrock_agent.agent_model.agent_name == "updated_agent" - @pytest.mark.asyncio async def test_invoke(self): """Test invoke of the agent.""" From 90fb7b1eae2ecec3996d07386d56a5be0d2a9641 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Thu, 27 Feb 2025 18:02:23 -0800 Subject: [PATCH 3/7] Need to work on testing next --- .../concepts/agents/bedrock_agent/README.md | 1 - .../bedrock_agent_simple_chat.py | 2 +- .../bedrock_agent_simple_chat_streaming.py | 2 +- .../bedrock_agent_use_existing.py | 51 ------------------- .../bedrock_agent_with_code_interpreter.py | 7 +-- ...k_agent_with_code_interpreter_streaming.py | 7 +-- .../bedrock_agent_with_kernel_function.py | 6 +-- ...ck_agent_with_kernel_function_streaming.py | 6 +-- .../bedrock_mixed_chat_agents.py | 2 +- .../bedrock_mixed_chat_agents_streaming.py | 2 +- .../agents/bedrock/bedrock_agent.py | 10 ++-- .../agents/bedrock/bedrock_agent_base.py | 39 ++++++++++---- python/semantic_kernel/utils/async_utils.py | 2 +- 13 files changed, 48 insertions(+), 89 deletions(-) delete mode 100644 python/samples/concepts/agents/bedrock_agent/bedrock_agent_use_existing.py diff --git a/python/samples/concepts/agents/bedrock_agent/README.md b/python/samples/concepts/agents/bedrock_agent/README.md index 586b210210a4..ea731c4e29c8 100644 --- a/python/samples/concepts/agents/bedrock_agent/README.md +++ b/python/samples/concepts/agents/bedrock_agent/README.md @@ -17,7 +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_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_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 eafc3cf08847..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() 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..755303f429da 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,12 @@ 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, ) + 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_streaming.py b/python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_streaming.py index 5a82ace15dd3..4f08f8795d84 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,12 @@ 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, ) + 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 dd98a512f920..8d2ffc49f6d2 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent.py @@ -19,7 +19,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, @@ -51,7 +50,7 @@ @experimental -class BedrockAgent(BedrockAgentBase, Agent): +class BedrockAgent(BedrockAgentBase): """Bedrock Agent. Manages the interaction with Amazon Bedrock Agent Service. @@ -87,6 +86,7 @@ def __init__( """ args: dict[str, Any] = { "agent_model": agent_model, + **kwargs, } if function_choice_behavior: @@ -133,7 +133,7 @@ async def create_and_prepare_agent( bedrock_client (Any, optional): The Bedrock Client. kernel (Kernel, optional): The kernel 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. prompt_template_config (PromptTemplateConfig, optional): The prompt template configuration. env_file_path (str, optional): The path to the environment file. @@ -155,7 +155,7 @@ async def create_and_prepare_agent( import boto3 from botocore.exceptions import ClientError - bedrock_runtime_client = (bedrock_runtime_client or boto3.client("bedrock-agent-runtime"),) + bedrock_runtime_client = bedrock_runtime_client or boto3.client("bedrock-agent-runtime") bedrock_client = bedrock_client or boto3.client("bedrock-agent") try: @@ -174,7 +174,7 @@ async def create_and_prepare_agent( raise e bedrock_agent = cls( - response, + response["agent"], function_choice_behavior=function_choice_behavior, kernel=kernel, arguments=arguments, diff --git a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py index 888a24ed4b5a..c9ed47f5a246 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py @@ -9,13 +9,12 @@ 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 @@ -23,7 +22,7 @@ @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. @@ -60,10 +59,12 @@ def __init__( bedrock_runtime_client: The Bedrock Runtime Client. kwargs: Additional keyword arguments. """ + agent_model = agent_model if isinstance(agent_model, BedrockAgentModel) else BedrockAgentModel(**agent_model) + args = { - "agent_model": agent_model - if isinstance(agent_model, BedrockAgentModel) - else BedrockAgentModel(**agent_model), + "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, @@ -169,7 +170,9 @@ async def _wait_for_agent_status( 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. Current status: {agent.agent_status}" + ) # endregion Agent Management @@ -193,6 +196,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}.") @@ -217,17 +222,19 @@ 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 @@ -247,6 +254,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}.") @@ -264,7 +273,7 @@ async def associate_agent_knowledge_base(self, knowledge_base_id: str, **kwargs) ) try: - return await run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.associate_agent_knowledge_base, @@ -274,6 +283,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}." @@ -288,7 +301,7 @@ async def disassociate_agent_knowledge_base(self, knowledge_base_id: str, **kwar ) try: - await run_in_executor( + response = await run_in_executor( None, partial( self.bedrock_client.disassociate_agent_knowledge_base, @@ -298,6 +311,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}." diff --git a/python/semantic_kernel/utils/async_utils.py b/python/semantic_kernel/utils/async_utils.py index 349e032c481f..4d2d2d50249a 100644 --- a/python/semantic_kernel/utils/async_utils.py +++ b/python/semantic_kernel/utils/async_utils.py @@ -6,6 +6,6 @@ from typing import Any -async def run_in_executor(self, executor: Any, func: Callable, *args, **kwargs) -> 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)) From cce3111066e8175793f0871bfc20711e07b5e7db Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Fri, 28 Feb 2025 14:45:53 -0800 Subject: [PATCH 4/7] Testing done --- .../agents/bedrock/bedrock_agent.py | 9 +- .../agents/bedrock/bedrock_agent_base.py | 22 +- .../agents/bedrock/bedrock_agent_settings.py | 3 +- .../test_bedrock_agent_integration.py | 23 +- .../unit/agents/bedrock_agent/conftest.py | 12 + .../bedrock_agent/test_bedrock_agent.py | 486 +++++++---- .../bedrock_agent/test_bedrock_agent_base.py | 800 ------------------ .../test_bedrock_agent_settings.py | 18 +- 8 files changed, 391 insertions(+), 982 deletions(-) delete mode 100644 python/tests/unit/agents/bedrock_agent/test_bedrock_agent_base.py diff --git a/python/semantic_kernel/agents/bedrock/bedrock_agent.py b/python/semantic_kernel/agents/bedrock/bedrock_agent.py index 8d2ffc49f6d2..405c17f7be7a 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent.py @@ -11,9 +11,6 @@ from pydantic import ValidationError -from semantic_kernel.agents.bedrock.models.bedrock_agent_status import BedrockAgentStatus -from semantic_kernel.utils.async_utils import run_in_executor - if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -27,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 @@ -40,6 +38,7 @@ from semantic_kernel.exceptions.agent_exceptions import AgentInitializationException, AgentInvokeException from semantic_kernel.functions.kernel_arguments import KernelArguments from semantic_kernel.kernel import Kernel +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, @@ -171,7 +170,7 @@ async def create_and_prepare_agent( ) except ClientError as e: logger.error(f"Failed to create agent {name}.") - raise e + raise AgentInitializationException("Failed to create the Amazon Bedrock Agent.") from e bedrock_agent = cls( response["agent"], @@ -186,7 +185,7 @@ async def create_and_prepare_agent( # 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() + await bedrock_agent.prepare_agent_and_wait_until_prepared() return bedrock_agent diff --git a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py index c9ed47f5a246..79bd27340ec5 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py @@ -93,7 +93,7 @@ def __repr__(self): # region Agent Management - async def _prepare_agent_and_wait_until_prepared(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.") @@ -151,7 +151,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 @@ -164,14 +165,15 @@ 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. Current status: {agent.agent_status}" + f"Agent did not reach status {status} within the specified time." + f" Current status: {self.agent_model.agent_status}" ) # endregion Agent Management @@ -196,7 +198,7 @@ async def create_code_interpreter_action_group(self, **kwargs) -> BedrockActionG ), ) - await self._prepare_agent_and_wait_until_prepared() + await self.prepare_agent_and_wait_until_prepared() return BedrockActionGroupModel(**response["agentActionGroup"]) except ClientError as e: @@ -222,7 +224,7 @@ async def create_user_input_action_group(self, **kwargs) -> BedrockActionGroupMo ), ) - await self._prepare_agent_and_wait_until_prepared() + await self.prepare_agent_and_wait_until_prepared() return BedrockActionGroupModel(**response["agentActionGroup"]) except ClientError as e: @@ -254,7 +256,7 @@ async def create_kernel_function_action_group(self, **kwargs) -> BedrockActionGr ), ) - await self._prepare_agent_and_wait_until_prepared() + await self.prepare_agent_and_wait_until_prepared() return BedrockActionGroupModel(**response["agentActionGroup"]) except ClientError as e: @@ -284,7 +286,7 @@ async def associate_agent_knowledge_base(self, knowledge_base_id: str, **kwargs) ), ) - await self._prepare_agent_and_wait_until_prepared() + await self.prepare_agent_and_wait_until_prepared() return response except ClientError as e: @@ -312,7 +314,7 @@ async def disassociate_agent_knowledge_base(self, knowledge_base_id: str, **kwar ), ) - await self._prepare_agent_and_wait_until_prepared() + await self.prepare_agent_and_wait_until_prepared() return response except ClientError as e: 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/tests/integration/agents/bedrock_agent/test_bedrock_agent_integration.py b/python/tests/integration/agents/bedrock_agent/test_bedrock_agent_integration.py index 4fbadfa361c4..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 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 2dc7aed24e27..20ef46a3de61 100644 --- a/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py +++ b/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py @@ -5,173 +5,414 @@ 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 - -# 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"] +# region Agent Initialization Tests -# 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 BedrockAgent initialization +async def test_bedrock_agent_initialization(bedrock_agent_model_with_id): + agent = BedrockAgent(bedrock_agent_model_with_id) - -# 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(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_existed( +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, +): + 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", + ) + + +# 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, ): - agent = BedrockAgent(name="test_agent", env_file_path="fake_path") + with ( + patch.object(client, "create_agent") as mock_create_agent, + patch.object(BedrockAgent, "prepare_agent_and_wait_until_prepared", new_callable=AsyncMock), + ): + from botocore.exceptions import ClientError + + mock_create_agent.side_effect = ClientError({}, "create_agent") + + 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", + ) - assert agent.agent_model.agent_id is None + +@patch.object(boto3, "client", return_value=Mock()) +async def test_bedrock_agent_prepare_agent_and_wait_until_prepared( + client, + bedrock_agent_unit_test_env, + bedrock_agent_model_with_id, + bedrock_agent_model_with_id_preparing_dict, + bedrock_agent_model_with_id_prepared_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") as mock_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_prepared_dict, + ] - await agent.create_agent() + await agent.prepare_agent_and_wait_until_prepared() - with pytest.raises( - ValueError, match="Agent already exists. Please delete the agent before creating a new one." - ): - await agent.create_agent() + 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_create_with_options( +async def test_bedrock_agent_prepare_agent_and_wait_until_prepared_fail( 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, ): - 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"), + ): + 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, + ] - assert agent.agent_model.agent_id is None + with pytest.raises(TimeoutError): + await agent.prepare_agent_and_wait_until_prepared() + + +# 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(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( - BedrockAgent, "create_code_interpreter_action_group", new_callable=AsyncMock - ) as mock_create_code_interpreter_action_group, + patch.object(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_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 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_action_group") as mock_create_action_group, 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 + 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() + - await agent.create_agent( - enable_code_interpreter=True, - enable_user_input=True, - enable_kernel_function=True, +# 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, "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, "prepare_agent_and_wait_until_prepared" + ) as mock_prepare_agent_and_wait_until_prepared, + ): + await agent.associate_agent_knowledge_base("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_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_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 @@ -184,16 +425,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") @@ -217,16 +451,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): @@ -251,16 +478,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"): @@ -285,16 +505,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 = "" @@ -321,17 +534,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( @@ -371,17 +577,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( @@ -409,3 +608,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): From ad6644196ea5d6be385f03daa8c21dc7aec8147f Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Fri, 28 Feb 2025 15:12:40 -0800 Subject: [PATCH 5/7] Various updates --- .../bedrock_agent_with_kernel_function.py | 1 + ...drock_agent_with_kernel_function_simple.py | 59 +++++++++++++++++++ ...ck_agent_with_kernel_function_streaming.py | 1 + .../agents/bedrock/bedrock_agent.py | 3 +- .../agents/bedrock/bedrock_agent_base.py | 2 +- .../services/bedrock_chat_completion.py | 3 +- .../services/bedrock_text_completion.py | 3 +- .../services/bedrock_text_embedding.py | 3 +- .../bedrock_agent/test_bedrock_agent.py | 6 +- 9 files changed, 70 insertions(+), 11 deletions(-) create mode 100644 python/samples/concepts/agents/bedrock_agent/bedrock_agent_with_kernel_function_simple.py 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 755303f429da..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 @@ -43,6 +43,7 @@ async def main(): INSTRUCTION, kernel=kernel, ) + # 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 4f08f8795d84..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 @@ -43,6 +43,7 @@ async def main(): INSTRUCTION, kernel=kernel, ) + # 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/semantic_kernel/agents/bedrock/bedrock_agent.py b/python/semantic_kernel/agents/bedrock/bedrock_agent.py index acb3a80b7a67..33b57363e193 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent.py @@ -11,8 +11,6 @@ from pydantic import ValidationError -from semantic_kernel.functions.kernel_plugin import KernelPlugin - if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -39,6 +37,7 @@ 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_plugin import KernelPlugin from semantic_kernel.kernel import Kernel from semantic_kernel.utils.async_utils import run_in_executor from semantic_kernel.utils.feature_stage_decorator import experimental diff --git a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py index 79bd27340ec5..94750415a2e7 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py @@ -137,7 +137,7 @@ 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.") 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 24daf5a115bb..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 @@ -7,8 +7,6 @@ import boto3 -from semantic_kernel.utils.async_utils import run_in_executor - if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -46,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 4d4937ebddba..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 @@ -9,8 +9,6 @@ import boto3 from pydantic import ValidationError -from semantic_kernel.utils.async_utils import run_in_executor - if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -28,6 +26,7 @@ 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 d79b6a345c1a..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 @@ -10,8 +10,6 @@ from numpy import array, ndarray from pydantic import ValidationError -from semantic_kernel.utils.async_utils import run_in_executor - if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -29,6 +27,7 @@ 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/tests/unit/agents/bedrock_agent/test_bedrock_agent.py b/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py index 92b00c5e50bb..ddf49aca36ad 100644 --- a/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py +++ b/python/tests/unit/agents/bedrock_agent/test_bedrock_agent.py @@ -20,7 +20,8 @@ # Test case to verify BedrockAgent initialization -async def test_bedrock_agent_initialization(bedrock_agent_model_with_id): +@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) assert agent.name == bedrock_agent_model_with_id.agent_name @@ -30,7 +31,8 @@ async def test_bedrock_agent_initialization(bedrock_agent_model_with_id): # Test case to verify error handling during BedrockAgent initialization with non-auto function choice -async def test_bedrock_agent_initialization_error_with_non_auto_function_choice(bedrock_agent_model_with_id): +@patch.object(boto3, "client", return_value=Mock()) +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( bedrock_agent_model_with_id, From 114f31737419ad953e94e05535538f02f1f9654e Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Fri, 28 Feb 2025 15:18:37 -0800 Subject: [PATCH 6/7] Fix link --- python/samples/concepts/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/samples/concepts/README.md b/python/samples/concepts/README.md index 5cc8e3ec0a42..3ab7044d1d57 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) From 9083fa22451e32bcdd625922867e3e6a39b7902a Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 3 Mar 2025 09:29:52 -0800 Subject: [PATCH 7/7] comments --- python/semantic_kernel/agents/bedrock/bedrock_agent_base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py index 94750415a2e7..708c0d2b01de 100644 --- a/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py +++ b/python/semantic_kernel/agents/bedrock/bedrock_agent_base.py @@ -59,7 +59,9 @@ def __init__( bedrock_runtime_client: The Bedrock Runtime Client. kwargs: Additional keyword arguments. """ - agent_model = agent_model if isinstance(agent_model, BedrockAgentModel) else BedrockAgentModel(**agent_model) + agent_model = ( + agent_model if isinstance(agent_model, BedrockAgentModel) else BedrockAgentModel.model_validate(agent_model) + ) args = { "agent_model": agent_model,