Skip to content

feat(bedrock_agent): add new Amazon Bedrock Agents Functions Resolver #6564

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 27 commits into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 16 commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
41bc401
feat(bedrock_agent): create bedrock agents functions data class
anafalcao Apr 25, 2025
bed8f3f
create resolver
anafalcao Apr 25, 2025
a3765f0
mypy
anafalcao Apr 25, 2025
44d80f8
add response
anafalcao Apr 25, 2025
abbc100
add name param to tool
anafalcao Apr 28, 2025
e42ceff
add response optional fields
anafalcao Apr 29, 2025
86c7ab7
bedrockfunctionresponse and response state
anafalcao Apr 30, 2025
34948d7
remove body message
anafalcao May 1, 2025
24978cb
add parser
anafalcao May 1, 2025
45f85f6
add test for required fields
anafalcao May 5, 2025
b420a90
Merge branch 'develop' into feat/bedrock_functions
anafalcao May 5, 2025
84bb6b0
add more tests for parser and resolver
anafalcao May 5, 2025
20bbe9f
Merge branch 'feat/bedrock_functions' of https://github.com/aws-power…
anafalcao May 5, 2025
d463304
add validation response state
anafalcao May 5, 2025
b4ab6b9
Merge branch 'develop' into feat/bedrock_functions
leandrodamascena May 6, 2025
39e0d36
Merge branch 'develop' into feat/bedrock_functions
leandrodamascena May 8, 2025
54a7edf
params injection
anafalcao May 8, 2025
fdde207
doc event handler, parser and data class
anafalcao May 15, 2025
c8b1b2f
fix doc typo
anafalcao May 15, 2025
db7d6b9
fix doc typo
anafalcao May 15, 2025
266ebcb
mypy
anafalcao May 15, 2025
4211b72
Merge branch 'develop' into feat/bedrock_functions
leandrodamascena May 26, 2025
20215ed
Small refactor + documentation
leandrodamascena May 26, 2025
ef31cb5
Small refactor + documentation
leandrodamascena May 26, 2025
c200a3a
Small refactor + documentation
leandrodamascena May 26, 2025
ca700c8
Small refactor + documentation
leandrodamascena May 26, 2025
c992463
Aligning Python implementation with TS
leandrodamascena May 27, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions aws_lambda_powertools/event_handler/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -26,7 +30,9 @@
"ALBResolver",
"ApiGatewayResolver",
"BedrockAgentResolver",
"BedrockAgentFunctionResolver",
"BedrockResponse",
"BedrockFunctionResponse",
"CORSConfig",
"LambdaFunctionUrlResolver",
"Response",
Expand Down
188 changes: 188 additions & 0 deletions aws_lambda_powertools/event_handler/bedrock_agent_function.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,188 @@
from __future__ import annotations

from typing import TYPE_CHECKING, Any

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 : str | 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: str | None = None,
) -> None:
if response_state is not None and response_state not in ["FAILURE", "REPROMPT"]:
raise ValueError("responseState must be None, '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

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:
if not description:
raise ValueError("Tool description is required")

function_name = name or func.__name__
if function_name in self._tools:
raise ValueError(f"Tool '{function_name}' already registered")

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")

Check warning on line 176 in aws_lambda_powertools/event_handler/bedrock_agent_function.py

View check run for this annotation

Codecov / codecov/patch

aws_lambda_powertools/event_handler/bedrock_agent_function.py#L176

Added line #L176 was not covered by tests

function_name = self.current_event.function

try:
result = self._tools[function_name]["function"]()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A critical part of this agent's workflow is passing parameters to functions. These parameters contain values ​​that will be useful for the function in question. In this discussion, we talk about this:

Support for injecting parameters into the function signature.

Can you pls add this support?

return BedrockFunctionsResponseBuilder(result).build(self.current_event)
except Exception as e:
return BedrockFunctionsResponseBuilder(
BedrockFunctionResponse(
body=f"Error: {str(e)}",
),
).build(self.current_event)
2 changes: 2 additions & 0 deletions aws_lambda_powertools/utilities/data_classes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -59,6 +60,7 @@
"AppSyncResolverEventsEvent",
"ALBEvent",
"BedrockAgentEvent",
"BedrockAgentFunctionEvent",
"CloudWatchAlarmData",
"CloudWatchAlarmEvent",
"CloudWatchAlarmMetric",
Expand Down
Original file line number Diff line number Diff line change
@@ -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 {}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -20,6 +20,7 @@
"ApiGatewayV2Envelope",
"ApiGatewayWebSocketEnvelope",
"BedrockAgentEnvelope",
"BedrockAgentFunctionEnvelope",
"CloudWatchLogsEnvelope",
"DynamoDBStreamEnvelope",
"EventBridgeEnvelope",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
2 changes: 2 additions & 0 deletions aws_lambda_powertools/utilities/parser/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
)
from .bedrock_agent import (
BedrockAgentEventModel,
BedrockAgentFunctionEventModel,
BedrockAgentModel,
BedrockAgentPropertyModel,
BedrockAgentRequestBodyModel,
Expand Down Expand Up @@ -208,6 +209,7 @@
"BedrockAgentEventModel",
"BedrockAgentRequestBodyModel",
"BedrockAgentRequestMediaModel",
"BedrockAgentFunctionEventModel",
"S3BatchOperationJobModel",
"S3BatchOperationModel",
"S3BatchOperationTaskModel",
Expand Down
18 changes: 18 additions & 0 deletions aws_lambda_powertools/utilities/parser/models/bedrock_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Loading
Loading