diff --git a/aws_lambda_powertools/event_handler/__init__.py b/aws_lambda_powertools/event_handler/__init__.py index 89952580dcc..f374590428d 100644 --- a/aws_lambda_powertools/event_handler/__init__.py +++ b/aws_lambda_powertools/event_handler/__init__.py @@ -12,6 +12,10 @@ ) from aws_lambda_powertools.event_handler.appsync import AppSyncResolver from aws_lambda_powertools.event_handler.bedrock_agent import BedrockAgentResolver, BedrockResponse +from aws_lambda_powertools.event_handler.bedrock_agent_function import ( + BedrockAgentFunctionResolver, + BedrockFunctionResponse, +) from aws_lambda_powertools.event_handler.events_appsync.appsync_events import AppSyncEventsResolver from aws_lambda_powertools.event_handler.lambda_function_url import ( LambdaFunctionUrlResolver, @@ -26,7 +30,9 @@ "ALBResolver", "ApiGatewayResolver", "BedrockAgentResolver", + "BedrockAgentFunctionResolver", "BedrockResponse", + "BedrockFunctionResponse", "CORSConfig", "LambdaFunctionUrlResolver", "Response", diff --git a/aws_lambda_powertools/event_handler/bedrock_agent_function.py b/aws_lambda_powertools/event_handler/bedrock_agent_function.py new file mode 100644 index 00000000000..20c16f48f5d --- /dev/null +++ b/aws_lambda_powertools/event_handler/bedrock_agent_function.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +import inspect +import warnings +from typing import TYPE_CHECKING, Any, Literal + +from aws_lambda_powertools.warnings import PowertoolsUserWarning + +if TYPE_CHECKING: + from collections.abc import Callable + +from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent + + +class BedrockFunctionResponse: + """Response class for Bedrock Agent Functions + + Parameters + ---------- + body : Any, optional + Response body + session_attributes : dict[str, str] | None + Session attributes to include in the response + prompt_session_attributes : dict[str, str] | None + Prompt session attributes to include in the response + response_state : Literal["FAILURE", "REPROMPT"] | None + Response state ("FAILURE" or "REPROMPT") + + Examples + -------- + ```python + @app.tool(description="Function that uses session attributes") + def test_function(): + return BedrockFunctionResponse( + body="Hello", + session_attributes={"userId": "123"}, + prompt_session_attributes={"lastAction": "login"} + ) + ``` + """ + + def __init__( + self, + body: Any = None, + session_attributes: dict[str, str] | None = None, + prompt_session_attributes: dict[str, str] | None = None, + knowledge_bases: list[dict[str, Any]] | None = None, + response_state: Literal["FAILURE", "REPROMPT"] | None = None, + ) -> None: + if response_state is not None and response_state not in ["FAILURE", "REPROMPT"]: + raise ValueError("responseState must be 'FAILURE' or 'REPROMPT'") + + self.body = body + self.session_attributes = session_attributes + self.prompt_session_attributes = prompt_session_attributes + self.knowledge_bases = knowledge_bases + self.response_state = response_state + + +class BedrockFunctionsResponseBuilder: + """ + Bedrock Functions Response Builder. This builds the response dict to be returned by Lambda + when using Bedrock Agent Functions. + """ + + def __init__(self, result: BedrockFunctionResponse | Any) -> None: + self.result = result + + def build(self, event: BedrockAgentFunctionEvent) -> dict[str, Any]: + """Build the full response dict to be returned by the lambda""" + if isinstance(self.result, BedrockFunctionResponse): + body = self.result.body + session_attributes = self.result.session_attributes + prompt_session_attributes = self.result.prompt_session_attributes + knowledge_bases = self.result.knowledge_bases + response_state = self.result.response_state + + else: + body = self.result + session_attributes = None + prompt_session_attributes = None + knowledge_bases = None + response_state = None + + # Per AWS Bedrock documentation, currently only "TEXT" is supported as the responseBody content type + # https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html + response: dict[str, Any] = { + "messageVersion": "1.0", + "response": { + "actionGroup": event.action_group, + "function": event.function, + "functionResponse": {"responseBody": {"TEXT": {"body": str(body if body is not None else "")}}}, + }, + } + + # Add responseState if provided + if response_state: + response["response"]["functionResponse"]["responseState"] = response_state + + # Add session attributes if provided in response or maintain from input + response.update( + { + "sessionAttributes": session_attributes or event.session_attributes or {}, + "promptSessionAttributes": prompt_session_attributes or event.prompt_session_attributes or {}, + }, + ) + + # Add knowledge bases configuration if provided + if knowledge_bases: + response["knowledgeBasesConfiguration"] = knowledge_bases + + return response + + +class BedrockAgentFunctionResolver: + """Bedrock Agent Function resolver that handles function definitions + + Examples + -------- + ```python + from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver + + app = BedrockAgentFunctionResolver() + + @app.tool(description="Gets the current UTC time") + def get_current_time(): + from datetime import datetime + return datetime.utcnow().isoformat() + + def lambda_handler(event, context): + return app.resolve(event, context) + ``` + """ + + def __init__(self) -> None: + self._tools: dict[str, dict[str, Any]] = {} + self.current_event: BedrockAgentFunctionEvent | None = None + self._response_builder_class = BedrockFunctionsResponseBuilder + + def tool( + self, + description: str | None = None, + name: str | None = None, + ) -> Callable: + """Decorator to register a tool function + + Parameters + ---------- + description : str | None + Description of what the tool does + name : str | None + Custom name for the tool. If not provided, uses the function name + """ + + def decorator(func: Callable) -> Callable: + function_name = name or func.__name__ + if function_name in self._tools: + warnings.warn( + f"Tool '{function_name}' already registered. Overwriting with new definition.", + PowertoolsUserWarning, + stacklevel=2, + ) + + self._tools[function_name] = { + "function": func, + "description": description, + } + return func + + return decorator + + def resolve(self, event: dict[str, Any], context: Any) -> dict[str, Any]: + """Resolves the function call from Bedrock Agent event""" + try: + self.current_event = BedrockAgentFunctionEvent(event) + return self._resolve() + except KeyError as e: + raise ValueError(f"Missing required field: {str(e)}") + + def _resolve(self) -> dict[str, Any]: + """Internal resolution logic""" + if self.current_event is None: + raise ValueError("No event to process") + + function_name = self.current_event.function + + try: + parameters = {} + if hasattr(self.current_event, "parameters"): + for param in self.current_event.parameters: + parameters[param.name] = param.value + + func = self._tools[function_name]["function"] + sig = inspect.signature(func) + + valid_params = {} + for name, value in parameters.items(): + if name in sig.parameters: + valid_params[name] = value + + result = func(**valid_params) + return BedrockFunctionsResponseBuilder(result).build(self.current_event) + except Exception as e: + return BedrockFunctionsResponseBuilder( + BedrockFunctionResponse( + body=f"Error: {str(e)}", + ), + ).build(self.current_event) diff --git a/aws_lambda_powertools/utilities/data_classes/__init__.py b/aws_lambda_powertools/utilities/data_classes/__init__.py index 7c1b67e6fa0..da0ef655fea 100644 --- a/aws_lambda_powertools/utilities/data_classes/__init__.py +++ b/aws_lambda_powertools/utilities/data_classes/__init__.py @@ -9,6 +9,7 @@ from .appsync_resolver_events_event import AppSyncResolverEventsEvent from .aws_config_rule_event import AWSConfigRuleEvent from .bedrock_agent_event import BedrockAgentEvent +from .bedrock_agent_function_event import BedrockAgentFunctionEvent from .cloud_watch_alarm_event import ( CloudWatchAlarmConfiguration, CloudWatchAlarmData, @@ -59,6 +60,7 @@ "AppSyncResolverEventsEvent", "ALBEvent", "BedrockAgentEvent", + "BedrockAgentFunctionEvent", "CloudWatchAlarmData", "CloudWatchAlarmEvent", "CloudWatchAlarmMetric", diff --git a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_function_event.py b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_function_event.py new file mode 100644 index 00000000000..ab479c59381 --- /dev/null +++ b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_function_event.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from aws_lambda_powertools.utilities.data_classes.common import DictWrapper + + +class BedrockAgentInfo(DictWrapper): + @property + def name(self) -> str: + return self["name"] + + @property + def id(self) -> str: # noqa: A003 + return self["id"] + + @property + def alias(self) -> str: + return self["alias"] + + @property + def version(self) -> str: + return self["version"] + + +class BedrockAgentFunctionParameter(DictWrapper): + @property + def name(self) -> str: + return self["name"] + + @property + def type(self) -> str: # noqa: A003 + return self["type"] + + @property + def value(self) -> str: + return self["value"] + + +class BedrockAgentFunctionEvent(DictWrapper): + """ + Bedrock Agent Function input event + + Documentation: + https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html + """ + + @property + def message_version(self) -> str: + return self["messageVersion"] + + @property + def input_text(self) -> str: + return self["inputText"] + + @property + def session_id(self) -> str: + return self["sessionId"] + + @property + def action_group(self) -> str: + return self["actionGroup"] + + @property + def function(self) -> str: + return self["function"] + + @property + def parameters(self) -> list[BedrockAgentFunctionParameter]: + parameters = self.get("parameters") or [] + return [BedrockAgentFunctionParameter(x) for x in parameters] + + @property + def agent(self) -> BedrockAgentInfo: + return BedrockAgentInfo(self["agent"]) + + @property + def session_attributes(self) -> dict[str, str]: + return self.get("sessionAttributes", {}) or {} + + @property + def prompt_session_attributes(self) -> dict[str, str]: + return self.get("promptSessionAttributes", {}) or {} diff --git a/aws_lambda_powertools/utilities/parser/envelopes/__init__.py b/aws_lambda_powertools/utilities/parser/envelopes/__init__.py index e1ac8cdbf5e..0bf4b7a5535 100644 --- a/aws_lambda_powertools/utilities/parser/envelopes/__init__.py +++ b/aws_lambda_powertools/utilities/parser/envelopes/__init__.py @@ -2,7 +2,7 @@ from .apigw_websocket import ApiGatewayWebSocketEnvelope from .apigwv2 import ApiGatewayV2Envelope from .base import BaseEnvelope -from .bedrock_agent import BedrockAgentEnvelope +from .bedrock_agent import BedrockAgentEnvelope, BedrockAgentFunctionEnvelope from .cloudwatch import CloudWatchLogsEnvelope from .dynamodb import DynamoDBStreamEnvelope from .event_bridge import EventBridgeEnvelope @@ -20,6 +20,7 @@ "ApiGatewayV2Envelope", "ApiGatewayWebSocketEnvelope", "BedrockAgentEnvelope", + "BedrockAgentFunctionEnvelope", "CloudWatchLogsEnvelope", "DynamoDBStreamEnvelope", "EventBridgeEnvelope", diff --git a/aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py b/aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py index 3d234999116..392c17cc425 100644 --- a/aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py +++ b/aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any from aws_lambda_powertools.utilities.parser.envelopes.base import BaseEnvelope -from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel +from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel, BedrockAgentFunctionEventModel if TYPE_CHECKING: from aws_lambda_powertools.utilities.parser.types import Model @@ -34,3 +34,27 @@ def parse(self, data: dict[str, Any] | Any | None, model: type[Model]) -> Model parsed_envelope: BedrockAgentEventModel = BedrockAgentEventModel.model_validate(data) logger.debug(f"Parsing event payload in `input_text` with {model}") return self._parse(data=parsed_envelope.input_text, model=model) + + +class BedrockAgentFunctionEnvelope(BaseEnvelope): + """Bedrock Agent Function envelope to extract data within input_text key""" + + def parse(self, data: dict[str, Any] | Any | None, model: type[Model]) -> Model | None: + """Parses data found with model provided + + Parameters + ---------- + data : dict + Lambda event to be parsed + model : type[Model] + Data model provided to parse after extracting data using envelope + + Returns + ------- + Model | None + Parsed detail payload with model provided + """ + logger.debug(f"Parsing incoming data with Bedrock Agent Function model {BedrockAgentFunctionEventModel}") + parsed_envelope: BedrockAgentFunctionEventModel = BedrockAgentFunctionEventModel.model_validate(data) + logger.debug(f"Parsing event payload in `input_text` with {model}") + return self._parse(data=parsed_envelope.input_text, model=model) diff --git a/aws_lambda_powertools/utilities/parser/models/__init__.py b/aws_lambda_powertools/utilities/parser/models/__init__.py index 7ea8da2dc22..ad8e3d7a92f 100644 --- a/aws_lambda_powertools/utilities/parser/models/__init__.py +++ b/aws_lambda_powertools/utilities/parser/models/__init__.py @@ -32,6 +32,7 @@ ) from .bedrock_agent import ( BedrockAgentEventModel, + BedrockAgentFunctionEventModel, BedrockAgentModel, BedrockAgentPropertyModel, BedrockAgentRequestBodyModel, @@ -208,6 +209,7 @@ "BedrockAgentEventModel", "BedrockAgentRequestBodyModel", "BedrockAgentRequestMediaModel", + "BedrockAgentFunctionEventModel", "S3BatchOperationJobModel", "S3BatchOperationModel", "S3BatchOperationTaskModel", diff --git a/aws_lambda_powertools/utilities/parser/models/bedrock_agent.py b/aws_lambda_powertools/utilities/parser/models/bedrock_agent.py index 62465162167..1aa5ae07a34 100644 --- a/aws_lambda_powertools/utilities/parser/models/bedrock_agent.py +++ b/aws_lambda_powertools/utilities/parser/models/bedrock_agent.py @@ -36,3 +36,21 @@ class BedrockAgentEventModel(BaseModel): agent: BedrockAgentModel parameters: Optional[List[BedrockAgentPropertyModel]] = None request_body: Optional[BedrockAgentRequestBodyModel] = Field(None, alias="requestBody") + + +class BedrockAgentFunctionEventModel(BaseModel): + """Bedrock Agent Function event model + + Documentation: + https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html + """ + + message_version: str = Field(..., alias="messageVersion") + agent: BedrockAgentModel + input_text: str = Field(..., alias="inputText") + session_id: str = Field(..., alias="sessionId") + action_group: str = Field(..., alias="actionGroup") + function: str + parameters: Optional[List[BedrockAgentPropertyModel]] = None + session_attributes: Dict[str, str] = Field({}, alias="sessionAttributes") + prompt_session_attributes: Dict[str, str] = Field({}, alias="promptSessionAttributes") diff --git a/docs/core/event_handler/bedrock_agents.md b/docs/core/event_handler/bedrock_agents.md index b7626f32f97..324b8ca15dd 100644 --- a/docs/core/event_handler/bedrock_agents.md +++ b/docs/core/event_handler/bedrock_agents.md @@ -1,9 +1,12 @@ --- -title: Agents for Amazon Bedrock +title: Amazon Bedrock Agents description: Core utility --- -Create [Agents for Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/agents.html#agents-how){target="_blank"} using event handlers and auto generation of OpenAPI schemas. +Create [Amazon Bedrock Agents](https://docs.aws.amazon.com/bedrock/latest/userguide/agents.html#agents-how){target="_blank"} using event handlers with two different action groups approaches: + +* OpenAPI schema +* Function details
```mermaid @@ -13,7 +16,8 @@ Create [Agents for Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/us ## Key features -* Minimal boilerplate to build Agents for Amazon Bedrock +* Minimal boilerplate to build Amazon Bedrock Agents +* Support for both OpenAPI-based and Function-based actions * Automatic generation of [OpenAPI schemas](https://www.openapis.org/){target="_blank"} from your business logic code * Built-in data validation for requests and responses * Similar experience to authoring [REST and HTTP APIs](api_gateway.md){target="_blank"} @@ -26,72 +30,53 @@ Create [Agents for Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/us **[OpenAPI schema](https://www.openapis.org/){target="_blank"}** is an industry standard JSON-serialized string that represents the structure and parameters of your API. +**Function details** consist of a list of parameters, defined by their name, [data type](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent_ParameterDetail.html), and whether they are required. The agent uses these configurations to determine what information it needs to elicit from the user. + **Action group** is a collection of two resources where you define the actions that the agent should carry out: an OpenAPI schema to define the APIs that the agent can invoke to carry out its tasks, and a Lambda function to execute those actions. **Large Language Models (LLM)** are very large deep learning models that are pre-trained on vast amounts of data, capable of extracting meanings from a sequence of text and understanding the relationship between words and phrases on it. -**Agent for Amazon Bedrock** is an Amazon Bedrock feature to build and deploy conversational agents that can interact with your customers using Large Language Models (LLM) and AWS Lambda functions. - -## Getting started +**Amazon Bedrock Agent** is an Amazon Bedrock feature to build and deploy conversational agents that can interact with your customers using Large Language Models (LLM) and AWS Lambda functions. !!! tip "All examples shared in this documentation are available within the [project repository](https://github.com/aws-powertools/powertools-lambda-python/tree/develop/examples)" -### Install - -!!! info "This is unnecessary if you're installing Powertools for AWS Lambda (Python) via [Lambda Layer/SAR](../../index.md#lambda-layer){target="_blank"}." +## Choose your Action Group -You need to add `pydantic` as a dependency in your preferred tool _e.g., requirements.txt, pyproject.toml_. At this time, we only support Pydantic V2. +An action group defines actions that the agent can help the user perform. You can define action groups as OpenAPI-based or Function-based. -### Required resources +| Aspect | OpenAPI-based Actions | Function-based Actions | +|--------|---------------------|----------------------| +| Definition Style | `@app.get("/path", description="")`
`@app.post("/path", description="")`| `@app.tool()` | +| Parameter Handling | Path, query, and body parameters | Function parameters | +| Use Case | REST-like APIs, complex request/response structures | Direct function calls, simpler input/output | +| Session Management | Via `BedrockResponse` | Via `BedrockFunctionResponse` | +| Best For | - Complex APIs with multiple endpoints
- When OpenAPI spec is required
- Integration with existing REST APIs | - Simple function-based actions
- Direct LLM-to-function mapping
- When function descriptions are sufficient | -To build Agents for Amazon Bedrock, you will need: - -| Requirement | Description | SAM Supported | CDK Supported | -|-----------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------|:-------------:|:-------------:| -| [Lambda Function](#your-first-agent) | Defines your business logic for the action group | ✅ | ✅ | -| [OpenAPI Schema](#generating-openapi-schemas) | API description, structure, and action group parameters | ❌ | ✅ | -| [Bedrock Service Role](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-permissions.html){target="_blank"} | Allows Amazon Bedrock to invoke foundation models | ✅ | ✅ | -| Agents for Bedrock | The service that will combine all the above to create the conversational agent | ❌ | ✅ | - -=== "Using AWS Serverless Application Model (SAM)" - Using [AWS SAM](https://aws.amazon.com/serverless/sam/){target="_blank"} you can create your Lambda function and the necessary permissions. However, you still have to create your Agent for Amazon Bedrock [using the AWS console](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-create.html){target="_blank"}. - - ```yaml hl_lines="18 26 34 61" - --8<-- "examples/event_handler_bedrock_agents/sam/template.yaml" - ``` - - 1. Amazon Bedrock needs permissions to invoke this Lambda function - 2. Check the [supported foundational models](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-supported.html){target="_blank"} - 3. You need the role ARN when creating the Agent for Amazon Bedrock +## Getting started -=== "Using AWS Cloud Developer Kit (CDK)" - This example uses the [Generative AI CDK constructs](https://awslabs.github.io/generative-ai-cdk-constructs/src/cdk-lib/bedrock/#agents){target="_blank"} to create your Agent with [AWS CDK](https://aws.amazon.com/cdk/){target="_blank"}. - These constructs abstract the underlying permission setup and code bundling of your Lambda function. +### Install - ```python - --8<-- "examples/event_handler_bedrock_agents/cdk/bedrock_agent_stack.py" - ``` +!!! info "This is unnecessary if you're installing Powertools for AWS Lambda (Python) via [Lambda Layer/SAR](../../index.md#lambda-layer){target="_blank"}." - 1. The path to your Lambda function handler - 2. The path to the OpenAPI schema describing your API +If you define the action group setting up an OpenAPI schema, you need to add `pydantic` as a dependency in your preferred tool _e.g., requirements.txt, pyproject.toml_. At this time, we only support Pydantic V2. ### Your first Agent -To create an agent, use the `BedrockAgentResolver` to annotate your actions. +To create an agent, use the `BedrockAgentResolver` or the `BedrockAgentFunctionResolver` to annotate your actions. This is similar to the way [all the other Event Handler](api_gateway.md) resolvers work. -You are required to add a `description` parameter in each endpoint, doing so will improve Bedrock's understanding of your actions. +The resolvers used by Amazon Bedrock Agents are compatible with all Powertools for AWS Lambda [features](../../index.md#features){target="blank"}. +For reference, we use [Logger](../logger.md) and [Tracer](../tracer.md) in this example. -=== "Lambda handler" +**OpenAPI-based actions** - The resolvers used by Agents for Amazon Bedrock are compatible with all Powertools for AWS Lambda [features](../../index.md#features){target="blank"}. - For reference, we use [Logger](../logger.md) and [Tracer](../tracer.md) in this example. +=== "Lambda handler" ```python hl_lines="4 9 12 21" --8<-- "examples/event_handler_bedrock_agents/src/getting_started.py" ``` - 1. `description` is a **required** field that should contain a human readable description of your action + 1. `description` is a **required** field that should contain a human readable description of your action. 2. We take care of **parsing**, **validating**, **routing** and **responding** to the request. === "OpenAPI schema" @@ -114,6 +99,28 @@ You are required to add a `description` parameter in each endpoint, doing so wil --8<-- "examples/event_handler_bedrock_agents/src/getting_started_output.json" ``` +**Function-based actions** +=== "Lambda handler" + + ```python hl_lines="4 9 12 21" + --8<-- "examples/event_handler_bedrock_agents/src/getting_started_functions.py" + ``` + + 1. `name` and `description` are optional here. + 2. We take care of **parsing**, **validating**, **routing** and **responding** to the request. + +=== "Input payload" + + ```json hl_lines="9 12" + --8<-- "examples/event_handler_bedrock_agents/src/input_getting_started_func.json" + ``` + +=== "Output payload" + + ```json hl_lines="9" + --8<-- "examples/event_handler_bedrock_agents/src/getting_started_output_func.json" + ``` + ??? note "What happens under the hood?" Powertools will handle the request from the Agent, parse, validate, and route it to the correct method in your code. The response is then validated and formatted back to the Agent. @@ -130,9 +137,12 @@ You can define the expected format for incoming data and responses by using type Define constraints using standard Python types, [dataclasses](https://docs.python.org/3/library/dataclasses.html) or [Pydantic models](https://docs.pydantic.dev/latest/concepts/models/). Pydantic is a popular library for data validation using Python type annotations. +The examples below uses [Pydantic's EmailStr](https://docs.pydantic.dev/2.0/usage/types/string_types/#emailstr){target="_blank"} to validate the email address passed to the `schedule_meeting` function. +The function then returns a boolean indicating if the meeting was successfully scheduled. + +**OpenAPI-based validation** + === "Lambda handler" - This example uses [Pydantic's EmailStr](https://docs.pydantic.dev/2.0/usage/types/string_types/#emailstr){target="_blank"} to validate the email address passed to the `schedule_meeting` function. - The function then returns a boolean indicating if the meeting was successfully scheduled. ```python hl_lines="1 2 6 16-18" --8<-- "examples/event_handler_bedrock_agents/src/getting_started_with_validation.py" @@ -159,6 +169,27 @@ Pydantic is a popular library for data validation using Python type annotations. ```json hl_lines="10" --8<-- "examples/event_handler_bedrock_agents/src/getting_started_with_validation_output.json" ``` +**Function-based validation** + +Uses direct type hints and focuses on function parameters + +=== "Lambda handler" + + ```python hl_lines="1 3-4 12" + --8<-- "examples/event_handler_bedrock_agents/src/validation_schema_func.py" + ``` + +=== "Input payload" + + ```json + --8<-- "examples/event_handler_bedrock_agents/src/input_validation_schema_func.json" + ``` + +=== "Output payload" + + ```json + --8<-- "examples/event_handler_bedrock_agents/src/output_validation_schema_func.json" + ``` #### When validation fails @@ -166,28 +197,77 @@ If the request validation fails, your event handler will not be called, and an e Similarly, if the response fails validation, your handler will abort the response. ???+ info "What does this mean for my Agent?" - The event handler will always return a response according to the OpenAPI schema. - A validation failure always results in a [422 response](https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/422). - However, how Amazon Bedrock interprets that failure is non-deterministic, since it depends on the characteristics of the LLM being used. + The event handler will always return a response according to the schema (OpenAPI) or type hints (Function-based). + A validation failure in OpenAPI-based actions results in a [422 response](https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/422). + For both approaches, how Amazon Bedrock interprets that failure is non-deterministic, since it depends on the characteristics of the LLM being used. -=== "Input payload" +=== "OpenAPI-based Input payload" ```json hl_lines="11" --8<-- "examples/event_handler_bedrock_agents/src/validation_failure_input.json" ``` -=== "Output payload" +=== "OpenAPI-based Output payload" ```json hl_lines="10" --8<-- "examples/event_handler_bedrock_agents/src/validation_failure_output.json" ``` +=== "Function-based Input payload" + + ```json hl_lines="12" + --8<-- "examples/event_handler_bedrock_agents/src/validation_failure_input_func.json" + ``` + +=== "Function-based Output payload" + + ```json hl_lines="9" + --8<-- "examples/event_handler_bedrock_agents/src/validation_failure_output_func.json" + ``` +
```mermaid --8<-- "docs/core/event_handler/bedrock_agents_validation_sequence_diagram.mermaid" ```
+### Accessing custom request fields + +The event sent by Amazon Bedrock Agents into your Lambda function contains a [number of extra event fields](#request_fields_table), exposed in the `app.current_event` field. + +???+ note "Why is this useful?" + You can for instance identify new conversations (`session_id`) or store and analyze entire conversations (`input_text`). + +=== "Accessing request fields" + + In this example, we [append correlation data](../logger.md#appending-additional-keys) to all generated logs. + This can be used to aggregate logs by `session_id` and observe the entire conversation between a user and the Agent. + + ```python hl_lines="13-16" + --8<-- "examples/event_handler_bedrock_agents/src/accessing_request_fields.py" + ``` + + + +The input event fields available depend on your Agent's configuration (OpenAPI-based or Function-based): + +| Name | Type | Description | OpenAPI | Function | +|------|------|-------------|----------|-----------| +| message_version | str | The version of the message format. Amazon Bedrock only supports version 1.0. | ✅ | ✅ | +| agent | BedrockAgentInfo | Contains information about the name, ID, alias, and version of the agent. | ✅ | ✅ | +| input_text | str | The user input for the conversation turn. | ✅ | ✅ | +| session_id | str | The unique identifier of the agent session. | ✅ | ✅ | +| action_group | str | The name of the action group. | ✅ | ✅ | +| api_path | str | The path to the API operation, as defined in the OpenAPI schema. | ✅ | ❌ | +| http_method | str | The method of the API operation, as defined in the OpenAPI schema. | ✅ | ❌ | +| function | str | The name of the function being called. | ❌ | ✅ | +| parameters | List[Parameter] | Contains parameters with name, type, and value properties. | ✅ | ✅ | +| request_body | BedrockAgentRequestBody | Contains the request body and its properties. | ✅ | ❌ | +| session_attributes | Dict[str, str] | Contains session attributes and their values. | ✅ | ✅ | +| prompt_session_attributes | Dict[str, str] | Contains prompt attributes and their values. | ✅ | ✅ | + +## OpenAPI-based actions + ### Generating OpenAPI schemas Use the `get_openapi_json_schema` function provided by the resolver to produce a JSON-serialized string that represents your OpenAPI schema. @@ -219,7 +299,7 @@ python3 app.py > schema.json ### Crafting effective OpenAPI schemas -Working with Agents for Amazon Bedrock will introduce [non-deterministic behaviour to your system](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-how.html#agents-rt){target="_blank"}. +Working with Amazon Bedrock Agents will introduce [non-deterministic behaviour to your system](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-how.html#agents-rt){target="_blank"}. ???+ note "Why is that?" Amazon Bedrock uses LLMs to understand and respond to user input. @@ -248,44 +328,11 @@ The following video demonstrates the end-to-end process: During the creation process, you should use the schema [previously generated](#generating-openapi-schemas) when prompted for an OpenAPI specification. -## Advanced - -### Accessing custom request fields - -The event sent by Agents for Amazon Bedrock into your Lambda function contains a [number of extra event fields](#request_fields_table), exposed in the `app.current_event` field. +### Advanced -???+ note "Why is this useful?" - You can for instance identify new conversations (`session_id`) or store and analyze entire conversations (`input_text`). +#### Additional metadata -=== "Accessing request fields" - - In this example, we [append correlation data](../logger.md#appending-additional-keys) to all generated logs. - This can be used to aggregate logs by `session_id` and observe the entire conversation between a user and the Agent. - - ```python hl_lines="13-16" - --8<-- "examples/event_handler_bedrock_agents/src/accessing_request_fields.py" - ``` - - -The input event fields are: - -| Name | Type | Description | -|---------------------------|------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| message_version | `str` | The version of the message that identifies the format of the event data going into the Lambda function and the expected format of the response from a Lambda function. Amazon Bedrock only supports version 1.0. | -| agent | `BedrockAgentInfo` | Contains information about the name, ID, alias, and version of the agent that the action group belongs to. | -| input_text | `str` | The user input for the conversation turn. | -| session_id | `str` | The unique identifier of the agent session. | -| action_group | `str` | The name of the action group. | -| api_path | `str` | The path to the API operation, as defined in the OpenAPI schema. | -| http_method | `str` | The method of the API operation, as defined in the OpenAPI schema. | -| parameters | `List[BedrockAgentProperty]` | Contains a list of objects. Each object contains the name, type, and value of a parameter in the API operation, as defined in the OpenAPI schema. | -| request_body | `BedrockAgentRequestBody` | Contains the request body and its properties, as defined in the OpenAPI schema. | -| session_attributes | `Dict[str, str]` | Contains session attributes and their values. | -| prompt_session_attributes | `Dict[str, str]` | Contains prompt attributes and their values. | - -### Additional metadata - -To enrich the view that Agents for Amazon Bedrock has of your Lambda functions, +To enrich the view that Amazon Bedrock Agents has of your Lambda functions, use a combination of [Pydantic Models](https://docs.pydantic.dev/latest/concepts/models/){target="_blank"} and [OpenAPI](https://www.openapis.org/){target="_blank"} type annotations to add constraints to your APIs parameters. ???+ info "When is this useful?" @@ -323,7 +370,7 @@ You can enable user confirmation with Bedrock Agents to have your application as 1. Add an openapi extension -### Fine grained responses +#### OpenAPI-based Responses ???+ info "Note" The default response only includes the essential fields to keep the payload size minimal, as AWS Lambda has a maximum response size of 25 KB. @@ -334,7 +381,7 @@ You can use `BedrockResponse` class to add additional fields as needed, such as --8<-- "examples/event_handler_bedrock_agents/src/working_with_bedrockresponse.py" ``` -## Testing your code +### Testing your code Test your routes by passing an [Agent for Amazon Bedrock proxy event](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html#agents-lambda-input) request: @@ -349,3 +396,33 @@ Test your routes by passing an [Agent for Amazon Bedrock proxy event](https://do ```python hl_lines="14-17" --8<-- "examples/event_handler_bedrock_agents/src/assert_bedrock_agent_response_module.py" ``` + +## Function-based Actions + +The `BedrockAgentFunctionResolver` handles three main aspects: + +* **Function Registration**: Using the `@app.tool()` decorator + +| Field | Required | Description | +|-------|----------|-------------| +| description | No | Human-readable description of what the function does. | +| name | No | Custom name for the function. Defaults to the function name. | + +* **Parameter Processing**: Automatic parsing and validation of input parameters +* **Response Formatting**: Converting function outputs into Bedrock Agent compatible responses + +### Function-based Responses + +???+ info "Note" + The default response only includes the essential fields to keep the payload size minimal, as AWS Lambda has a maximum response size of 25 KB. + +You can use `BedrockFunctionResponse` class to customize your response [with additional fields](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html#agents-lambda-response){target="_blank"}. This class allows you to: + +* Return a response body +* Set session and prompt session attributes +* Set knowledge bases configurations +* Control the response state ("FAILURE" or "REPROMPT") + +```python title="working_with_bedrockresponse.py" title="Customzing your Bedrock Function Response" hl_lines="2 8" +--8<-- "examples/event_handler_bedrock_agents/src/working_bedrock_functions_response.py" +``` diff --git a/docs/core/event_handler/bedrock_agents.mermaid b/docs/core/event_handler/bedrock_agents.mermaid index 19ae2270234..fc4fb2ccdc4 100644 --- a/docs/core/event_handler/bedrock_agents.mermaid +++ b/docs/core/event_handler/bedrock_agents.mermaid @@ -2,12 +2,17 @@ flowchart LR Bedrock[LLM] <-- uses --> Agent You[User input] --> Agent Agent -- consults --> OpenAPI - Agent[Agents for Amazon Bedrock] -- invokes --> Lambda + Agent -- consults --> Functions + Agent[Amazon Bedrock Agents] -- invokes --> Lambda subgraph OpenAPI Schema end + subgraph Functions + ToolDescriptions[Tool Descriptions] + end + subgraph Lambda[Lambda Function] direction TB Parsing[Parameter Parsing] --> Validation @@ -19,10 +24,8 @@ flowchart LR subgraph ActionGroup[Action Group] OpenAPI -. generated from .-> Lambda + Functions -. defined in .-> Lambda end style Code fill:#ffa500,color:black,font-weight:bold,stroke-width:3px - style You stroke:#0F0,stroke-width:2px - - - + style You stroke:#0F0,stroke-width:2px \ No newline at end of file diff --git a/docs/core/event_handler/bedrock_agents_getting_started.mermaid b/docs/core/event_handler/bedrock_agents_getting_started.mermaid index 29f3a26e323..6c6b13de72e 100644 --- a/docs/core/event_handler/bedrock_agents_getting_started.mermaid +++ b/docs/core/event_handler/bedrock_agents_getting_started.mermaid @@ -14,7 +14,10 @@ sequenceDiagram participant Your Code end - Agent->>Lambda: GET /current_time + alt Function-based + Agent->>Lambda: {function: "current_time", parameters: [], ...} + end + activate Lambda Lambda->>Parsing: parses parameters Parsing->>Validation: validates input @@ -26,8 +29,11 @@ sequenceDiagram Routing->>Validation: returns output Validation->>Parsing: validates output Parsing->>Lambda: formats response - Lambda->>Agent: 1709215709 + + alt Function-based + Lambda->>Agent: {response: {functionResponse: {responseBody: {...}}}} + end deactivate Lambda Agent-->>Agent: LLM interaction - Agent->>User: "The current time is 14:08:29 GMT" + Agent->>User: "The current time is 14:08:29 GMT" \ No newline at end of file diff --git a/docs/utilities/data_classes.md b/docs/utilities/data_classes.md index b3ba5a2f474..3fcf940296a 100644 --- a/docs/utilities/data_classes.md +++ b/docs/utilities/data_classes.md @@ -75,7 +75,8 @@ Each event source is linked to its corresponding GitHub file with the full set o | [AppSync Authorizer](#appsync-authorizer) | `AppSyncAuthorizerEvent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/appsync_authorizer_event.py) | | [AppSync Resolver](#appsync-resolver) | `AppSyncResolverEvent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/appsync_resolver_event.py) | | [AWS Config Rule](#aws-config-rule) | `AWSConfigRuleEvent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/aws_config_rule_event.py) | -| [Bedrock Agent](#bedrock-agent) | `BedrockAgent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py) | +| [Bedrock Agent - OpenAPI](#bedrock-agent) | `BedrockAgentEvent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py) | +| [Bedrock Agent - Function](#bedrock-agent) | `BedrockAgentFunctionEvent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/bedrock_agent_function_event.py) | | [CloudFormation Custom Resource](#cloudformation-custom-resource) | `CloudFormationCustomResourceEvent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/cloudformation_custom_resource_event.py) | | [CloudWatch Alarm State Change Action](#cloudwatch-alarm-state-change-action) | `CloudWatchAlarmEvent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/cloud_watch_alarm_event.py) | | [CloudWatch Dashboard Custom Widget](#cloudwatch-dashboard-custom-widget) | `CloudWatchDashboardCustomWidgetEvent` | [Github](https://github.com/aws-powertools/powertools-lambda-python/blob/develop/aws_lambda_powertools/utilities/data_classes/cloud_watch_custom_widget_event.py) | diff --git a/docs/utilities/parser.md b/docs/utilities/parser.md index 4cdf0d452f2..59302f45a34 100644 --- a/docs/utilities/parser.md +++ b/docs/utilities/parser.md @@ -112,7 +112,8 @@ The example above uses `SqsModel`. Other built-in models can be found below. | **APIGatewayWebSocketConnectEventModel** | Lambda Event Source payload for Amazon API Gateway WebSocket API $connect message | | **APIGatewayWebSocketDisconnectEventModel** | Lambda Event Source payload for Amazon API Gateway WebSocket API $disconnect message | | **AppSyncResolverEventModel** | Lambda Event Source payload for AWS AppSync Resolver | -| **BedrockAgentEventModel** | Lambda Event Source payload for Bedrock Agents | +| **BedrockAgentEventModel** | Lambda Event Source payload for Bedrock Agents - OpenAPI-based | +| **BedrockAgentFunctionEventModel** | Lambda Event Source payload for Bedrock Agents - Function-based | | **CloudFormationCustomResourceCreateModel** | Lambda Event Source payload for AWS CloudFormation `CREATE` operation | | **CloudFormationCustomResourceUpdateModel** | Lambda Event Source payload for AWS CloudFormation `UPDATE` operation | | **CloudFormationCustomResourceDeleteModel** | Lambda Event Source payload for AWS CloudFormation `DELETE` operation | diff --git a/examples/event_handler_bedrock_agents/src/accessing_request_fields.py b/examples/event_handler_bedrock_agents/src/accessing_request_fields.py index 529c9343702..feb44eda6cc 100644 --- a/examples/event_handler_bedrock_agents/src/accessing_request_fields.py +++ b/examples/event_handler_bedrock_agents/src/accessing_request_fields.py @@ -8,7 +8,7 @@ app = BedrockAgentResolver() -@app.get("/current_time", description="Gets the current time in seconds") # (1)! +@app.get("/current_time", description="Gets the current time in seconds") def current_time() -> int: logger.append_keys( session_id=app.current_event.session_id, diff --git a/examples/event_handler_bedrock_agents/src/getting_started_functions.py b/examples/event_handler_bedrock_agents/src/getting_started_functions.py new file mode 100644 index 00000000000..4a8dda5c27c --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/getting_started_functions.py @@ -0,0 +1,21 @@ +from time import time + +from aws_lambda_powertools import Logger, Tracer +from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +tracer = Tracer() +logger = Logger() +app = BedrockAgentFunctionResolver() + + +@app.tool(name="currentTime", description="Gets the current time in seconds") # (1)! +@tracer.capture_method +def current_time() -> int: + return int(time()) + + +@logger.inject_lambda_context +@tracer.capture_lambda_handler +def lambda_handler(event: dict, context: LambdaContext): + return app.resolve(event, context) # (2)! diff --git a/examples/event_handler_bedrock_agents/src/getting_started_output_func.json b/examples/event_handler_bedrock_agents/src/getting_started_output_func.json new file mode 100644 index 00000000000..2777dd96add --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/getting_started_output_func.json @@ -0,0 +1,14 @@ +{ + "messageVersion": "1.0", + "response": { + "actionGroup": "CurrentTime", + "function": "CurrentTime", + "functionResponse": { + "responseBody": { + "application/json": { + "body": "1704708165" + } + } + } + } +} \ No newline at end of file diff --git a/examples/event_handler_bedrock_agents/src/input_getting_started_func.json b/examples/event_handler_bedrock_agents/src/input_getting_started_func.json new file mode 100644 index 00000000000..d0b59c614aa --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/input_getting_started_func.json @@ -0,0 +1,16 @@ +{ + "messageVersion": "1.0", + "agent": { + "name": "TimeAgent", + "id": "XLHH72XNF2", + "alias": "TSTALIASID", + "version": "DRAFT" + }, + "inputText": "What is the current time?", + "sessionId": "123456789012345", + "actionGroup": "CurrentTime", + "function": "CurrentTime", + "parameters": [], + "sessionAttributes": {}, + "promptSessionAttributes": {} +} \ No newline at end of file diff --git a/examples/event_handler_bedrock_agents/src/input_validation_schema_func.json b/examples/event_handler_bedrock_agents/src/input_validation_schema_func.json new file mode 100644 index 00000000000..0ecb458bbcc --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/input_validation_schema_func.json @@ -0,0 +1,19 @@ +{ + "messageVersion": "1.0", + "agent": { + "name": "TestAgent", + "id": "XXXXX", + "alias": "LATEST", + "version": "DRAFT" + }, + "inputText": "Schedule a meeting with john@example.com", + "sessionId": "session-id", + "actionGroup": "TestActionGroup", + "function": "schedule_meeting", + "parameters": [ + { + "name": "email", + "value": "john@example.com" + } + ] +} \ No newline at end of file diff --git a/examples/event_handler_bedrock_agents/src/output_validation_schema_func.json b/examples/event_handler_bedrock_agents/src/output_validation_schema_func.json new file mode 100644 index 00000000000..a0beb314635 --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/output_validation_schema_func.json @@ -0,0 +1,17 @@ +{ + "messageVersion": "1.0", + "response": { + "actionGroup": "TestActionGroup", + "function": "schedule_meeting", + "functionResponse": { + "responseBody": { + "TEXT": { + "body": "true" + } + } + } + }, + "sessionAttributes": { + "last_email": "john@example.com" + } +} \ No newline at end of file diff --git a/examples/event_handler_bedrock_agents/src/validation_failure_input_func.json b/examples/event_handler_bedrock_agents/src/validation_failure_input_func.json new file mode 100644 index 00000000000..95759ae2b67 --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/validation_failure_input_func.json @@ -0,0 +1,16 @@ +{ + "messageVersion": "1.0", + "agent": { + "name": "TestAgent", + "id": "XXXXX", + "alias": "LATEST", + "version": "DRAFT" + }, + "parameters": [ + { + "name": "email", + "value": "not-an-email" + } + ], + "function": "schedule_meeting" +} \ No newline at end of file diff --git a/examples/event_handler_bedrock_agents/src/validation_failure_output_func.json b/examples/event_handler_bedrock_agents/src/validation_failure_output_func.json new file mode 100644 index 00000000000..fac4a714698 --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/validation_failure_output_func.json @@ -0,0 +1,14 @@ +{ + "messageVersion": "1.0", + "response": { + "actionGroup": "TestActionGroup", + "function": "schedule_meeting", + "functionResponse": { + "responseBody": { + "TEXT": { + "body": "Error: Invalid value for parameter 'email': value is not a valid email address" + } + } + } + } +} \ No newline at end of file diff --git a/examples/event_handler_bedrock_agents/src/validation_schema_func.py b/examples/event_handler_bedrock_agents/src/validation_schema_func.py new file mode 100644 index 00000000000..36a55d35545 --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/validation_schema_func.py @@ -0,0 +1,22 @@ +from pydantic import EmailStr + +from aws_lambda_powertools import Logger, Tracer +from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver +from aws_lambda_powertools.event_handler.bedrock_agent_function import BedrockFunctionResponse + +tracer = Tracer() +logger = Logger() +app = BedrockAgentFunctionResolver() + + +@app.tool(description="Schedules a meeting with the team") +@tracer.capture_method +def schedule_meeting(email: EmailStr) -> BedrockFunctionResponse: + logger.info("Scheduling a meeting", email=email) + return BedrockFunctionResponse(body=True, session_attributes={"last_email": str(email)}) + + +@logger.inject_lambda_context +@tracer.capture_lambda_handler +def lambda_handler(event: dict, context): + return app.resolve(event, context) diff --git a/examples/event_handler_bedrock_agents/src/working_bedrock_functions_response.py b/examples/event_handler_bedrock_agents/src/working_bedrock_functions_response.py new file mode 100644 index 00000000000..8854469b017 --- /dev/null +++ b/examples/event_handler_bedrock_agents/src/working_bedrock_functions_response.py @@ -0,0 +1,14 @@ +from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver, BedrockFunctionResponse + +app = BedrockAgentFunctionResolver() + + +@app.tool(description="Function that demonstrates response customization") +def custom_response(): + return BedrockFunctionResponse( + body="Hello World", + session_attributes={"user_id": "123"}, + prompt_session_attributes={"last_action": "greeting"}, + response_state="REPROMPT", + knowledge_bases=[{"name": "kb1", "enabled": True}], + ) diff --git a/tests/events/bedrockAgentFunctionEvent.json b/tests/events/bedrockAgentFunctionEvent.json new file mode 100644 index 00000000000..e849c3e6b73 --- /dev/null +++ b/tests/events/bedrockAgentFunctionEvent.json @@ -0,0 +1,32 @@ +{ + "messageVersion": "1.0", + "agent": { + "alias": "PROD", + "name": "hr-assistant-function-def", + "version": "1", + "id": "1234abcd" + }, + "sessionId": "123456789123458", + "sessionAttributes": { + "employeeId": "EMP123" + }, + "promptSessionAttributes": { + "lastInteraction": "2024-02-01T15:30:00Z", + "requestType": "vacation" + }, + "inputText": "I want to request vacation from March 15 to March 20", + "actionGroup": "VacationsActionGroup", + "function": "submitVacationRequest", + "parameters": [ + { + "name": "startDate", + "type": "string", + "value": "2024-03-15" + }, + { + "name": "endDate", + "type": "string", + "value": "2024-03-20" + } + ] +} \ No newline at end of file diff --git a/tests/functional/event_handler/required_dependencies/test_bedrock_agent_functions.py b/tests/functional/event_handler/required_dependencies/test_bedrock_agent_functions.py new file mode 100644 index 00000000000..f369dbb9cd3 --- /dev/null +++ b/tests/functional/event_handler/required_dependencies/test_bedrock_agent_functions.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import pytest + +from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver, BedrockFunctionResponse +from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent +from aws_lambda_powertools.warnings import PowertoolsUserWarning +from tests.functional.utils import load_event + + +def test_bedrock_agent_function_with_string_response(): + # GIVEN a Bedrock Agent Function resolver + app = BedrockAgentFunctionResolver() + + @app.tool() + def test_function(): + assert isinstance(app.current_event, BedrockAgentFunctionEvent) + return "Hello from string" + + # WHEN calling the event handler + raw_event = load_event("bedrockAgentFunctionEvent.json") + raw_event["function"] = "test_function" + result = app.resolve(raw_event, {}) + + # THEN process event correctly with string response + assert result["messageVersion"] == "1.0" + assert result["response"]["actionGroup"] == raw_event["actionGroup"] + assert result["response"]["function"] == "test_function" + assert result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] == "Hello from string" + assert "responseState" not in result["response"]["functionResponse"] + + +def test_bedrock_agent_function_error_handling(): + # GIVEN a Bedrock Agent Function resolver + app = BedrockAgentFunctionResolver() + + @app.tool(description="Function with error handling") + def error_function(): + return BedrockFunctionResponse( + body="Invalid input", + response_state="REPROMPT", + session_attributes={"error": "true"}, + ) + + @app.tool(description="Function that raises error") + def exception_function(): + raise ValueError("Something went wrong") + + # WHEN calling with explicit error response + raw_event = load_event("bedrockAgentFunctionEvent.json") + raw_event["function"] = "error_function" + result = app.resolve(raw_event, {}) + + # THEN include REPROMPT state and session attributes + assert result["response"]["functionResponse"]["responseState"] == "REPROMPT" + assert result["sessionAttributes"] == {"error": "true"} + + +def test_bedrock_agent_function_registration(): + # GIVEN a Bedrock Agent Function resolver + app = BedrockAgentFunctionResolver() + + # WHEN registering with duplicate name + @app.tool(name="custom", description="First registration") + def first_function(): + return "first test" + + # THEN a warning should be issued when registering a duplicate + with pytest.warns(PowertoolsUserWarning, match="Tool 'custom' already registered"): + + @app.tool(name="custom", description="Second registration") + def second_function(): + return "second test" + + # AND the most recent function should be registered + raw_event = load_event("bedrockAgentFunctionEvent.json") + raw_event["function"] = "custom" + result = app.resolve(raw_event, {}) + + # The second function should be used + assert result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] == "second test" + + +def test_bedrock_agent_function_with_optional_fields(): + # GIVEN a Bedrock Agent Function resolver + app = BedrockAgentFunctionResolver() + + @app.tool(description="Function with all optional fields") + def test_function(): + return BedrockFunctionResponse( + body="Hello", + session_attributes={"userId": "123"}, + prompt_session_attributes={"context": "test"}, + knowledge_bases=[ + { + "knowledgeBaseId": "kb1", + "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 5}}, + }, + ], + ) + + # WHEN calling the event handler + raw_event = load_event("bedrockAgentFunctionEvent.json") + raw_event["function"] = "test_function" + result = app.resolve(raw_event, {}) + + # THEN include all optional fields in response + assert result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] == "Hello" + assert result["sessionAttributes"] == {"userId": "123"} + assert result["promptSessionAttributes"] == {"context": "test"} + assert result["knowledgeBasesConfiguration"][0]["knowledgeBaseId"] == "kb1" + + +def test_bedrock_agent_function_invalid_event(): + # GIVEN a Bedrock Agent Function resolver + app = BedrockAgentFunctionResolver() + + # WHEN calling with invalid event + with pytest.raises(ValueError, match="Missing required field"): + app.resolve({}, {}) + + +def test_resolve_raises_value_error_on_missing_required_field(): + """Test that resolve() raises ValueError when a required field is missing from the event""" + # GIVEN a Bedrock Agent Function resolver and an incomplete event + resolver = BedrockAgentFunctionResolver() + incomplete_event = { + "messageVersion": "1.0", + "agent": {"alias": "PROD", "name": "hr-assistant-function-def", "version": "1", "id": "1234abcd"}, + "sessionId": "123456789123458", + } + + # WHEN calling resolve with the incomplete event + # THEN a ValueError is raised with information about the missing field + with pytest.raises(ValueError) as excinfo: + resolver.resolve(incomplete_event, {}) + + assert "Missing required field:" in str(excinfo.value) + + +def test_resolve_with_no_registered_function(): + # GIVEN a Bedrock Agent Function resolver + app = BedrockAgentFunctionResolver() + + # AND a valid event but with a non-existent function + raw_event = { + "messageVersion": "1.0", + "agent": {"name": "TestAgent", "id": "test-id", "alias": "test", "version": "1"}, + "actionGroup": "test_group", + "function": "non_existent_function", + "parameters": [], + } + + # WHEN calling resolve with a non-existent function + result = app.resolve(raw_event, {}) + + # THEN the response should contain an error message + assert "Error: 'non_existent_function'" in result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] + + +def test_bedrock_function_response_state_validation(): + # GIVEN invalid and valid response states + valid_states = ["FAILURE", "REPROMPT"] + invalid_state = "INVALID" + + # WHEN creating responses with valid states + # THEN no error should be raised + for state in valid_states: + try: + BedrockFunctionResponse(body="test", response_state=state) + except ValueError: + pytest.fail(f"Unexpected ValueError for response_state={state}") + + # WHEN creating a response with invalid state + # THEN ValueError should be raised with correct message + with pytest.raises(ValueError) as exc_info: + BedrockFunctionResponse(body="test", response_state=invalid_state) + + assert str(exc_info.value) == "responseState must be 'FAILURE' or 'REPROMPT'" + + +def test_bedrock_agent_function_with_parameters(): + # GIVEN a Bedrock Agent Function resolver + app = BedrockAgentFunctionResolver() + + # Track received parameters + received_params = {} + + @app.tool(description="Function that accepts parameters") + def vacation_request(startDate, endDate): + # Store received parameters for assertion + received_params["startDate"] = startDate + received_params["endDate"] = endDate + return f"Vacation request from {startDate} to {endDate} submitted" + + # WHEN calling the event handler with parameters + raw_event = load_event("bedrockAgentFunctionEvent.json") + raw_event["function"] = "vacation_request" + result = app.resolve(raw_event, {}) + + # THEN parameters should be correctly passed to the function + assert received_params["startDate"] == "2024-03-15" + assert received_params["endDate"] == "2024-03-20" + assert ( + "Vacation request from 2024-03-15 to 2024-03-20 submitted" + in result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] + ) diff --git a/tests/unit/data_classes/required_dependencies/test_bedrock_agent_function_event.py b/tests/unit/data_classes/required_dependencies/test_bedrock_agent_function_event.py new file mode 100644 index 00000000000..e055c894604 --- /dev/null +++ b/tests/unit/data_classes/required_dependencies/test_bedrock_agent_function_event.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent +from tests.functional.utils import load_event + + +def test_bedrock_agent_function_event(): + raw_event = load_event("bedrockAgentFunctionEvent.json") + parsed_event = BedrockAgentFunctionEvent(raw_event) + + # Test basic event properties + assert parsed_event.message_version == raw_event["messageVersion"] + assert parsed_event.session_id == raw_event["sessionId"] + assert parsed_event.input_text == raw_event["inputText"] + assert parsed_event.action_group == raw_event["actionGroup"] + assert parsed_event.function == raw_event["function"] + + # Test agent information + agent = parsed_event.agent + raw_agent = raw_event["agent"] + assert agent.alias == raw_agent["alias"] + assert agent.name == raw_agent["name"] + assert agent.version == raw_agent["version"] + assert agent.id == raw_agent["id"] + + # Test session attributes + assert parsed_event.session_attributes == raw_event["sessionAttributes"] + assert parsed_event.prompt_session_attributes == raw_event["promptSessionAttributes"] + + # Test parameters + parameters = parsed_event.parameters + raw_parameters = raw_event["parameters"] + assert len(parameters) == len(raw_parameters) + + for param, raw_param in zip(parameters, raw_parameters): + assert param.name == raw_param["name"] + assert param.type == raw_param["type"] + assert param.value == raw_param["value"] + + +def test_bedrock_agent_function_event_minimal(): + """Test with minimal required fields""" + minimal_event = { + "messageVersion": "1.0", + "agent": { + "alias": "PROD", + "name": "hr-assistant-function-def", + "version": "1", + "id": "1234abcd-56ef-78gh-90ij-klmn12345678", + }, + "sessionId": "87654321-abcd-efgh-ijkl-mnop12345678", + "inputText": "I want to request vacation", + "actionGroup": "VacationsActionGroup", + "function": "submitVacationRequest", + } + + parsed_event = BedrockAgentFunctionEvent(minimal_event) + + assert parsed_event.session_attributes == {} + assert parsed_event.prompt_session_attributes == {} + assert parsed_event.parameters == [] diff --git a/tests/unit/parser/_pydantic/test_bedrock_agent.py b/tests/unit/parser/_pydantic/test_bedrock_agent.py index 207318952cc..e66a202a53f 100644 --- a/tests/unit/parser/_pydantic/test_bedrock_agent.py +++ b/tests/unit/parser/_pydantic/test_bedrock_agent.py @@ -1,5 +1,5 @@ from aws_lambda_powertools.utilities.parser import envelopes, parse -from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel +from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel, BedrockAgentFunctionEventModel from tests.functional.utils import load_event from tests.unit.parser._pydantic.schemas import MyBedrockAgentBusiness @@ -76,3 +76,48 @@ def test_bedrock_agent_event_with_post(): assert properties[1].name == raw_properties[1]["name"] assert properties[1].type_ == raw_properties[1]["type"] assert properties[1].value == raw_properties[1]["value"] + + +def test_bedrock_agent_function_event(): + raw_event = load_event("bedrockAgentFunctionEvent.json") + model = BedrockAgentFunctionEventModel(**raw_event) + + assert model.message_version == raw_event["messageVersion"] + assert model.session_id == raw_event["sessionId"] + assert model.input_text == raw_event["inputText"] + assert model.action_group == raw_event["actionGroup"] + assert model.function == raw_event["function"] + assert model.session_attributes == {"employeeId": "EMP123"} + assert model.prompt_session_attributes == {"lastInteraction": "2024-02-01T15:30:00Z", "requestType": "vacation"} + + agent = model.agent + raw_agent = raw_event["agent"] + assert agent.alias == raw_agent["alias"] + assert agent.name == raw_agent["name"] + assert agent.version == raw_agent["version"] + assert agent.id_ == raw_agent["id"] + + parameters = model.parameters + assert parameters is not None + assert len(parameters) == 2 + + assert parameters[0].name == "startDate" + assert parameters[0].type_ == "string" + assert parameters[0].value == "2024-03-15" + + assert parameters[1].name == "endDate" + assert parameters[1].type_ == "string" + assert parameters[1].value == "2024-03-20" + + +def test_bedrock_agent_function_event_with_envelope(): + raw_event = load_event("bedrockAgentFunctionEvent.json") + raw_event["inputText"] = '{"username": "Jane", "name": "Doe"}' + parsed_event: MyBedrockAgentBusiness = parse( + event=raw_event, + model=MyBedrockAgentBusiness, + envelope=envelopes.BedrockAgentFunctionEnvelope, + ) + + assert parsed_event.username == "Jane" + assert parsed_event.name == "Doe"