-
Notifications
You must be signed in to change notification settings - Fork 437
feat(data_classes): add support for Bedrock Agents event #3262
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
Merged
Merged
Changes from 7 commits
Commits
Show all changes
12 commits
Select commit
Hold shift + click to select a range
030aed6
feat(data_classes): add support for Bedrock Agents events
rubenfonseca 2b4cf35
fix: add tests
rubenfonseca 2671337
fix: add docs
rubenfonseca 2eacca1
Merge branch 'develop' into rf/bedrock-events
leandrodamascena a7eb9e2
fix: changed to the new payload schema
rubenfonseca 5d5ce0a
fix: example
rubenfonseca db71067
fix: add more tests
rubenfonseca 4e9110d
fix: remove response
rubenfonseca fd524b4
fix: part 2
rubenfonseca ead0e9a
fix: comments
rubenfonseca 89700ec
fix: remove response example
rubenfonseca 82d4b75
Merge branch 'develop' into rf/bedrock-events
leandrodamascena File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
143 changes: 143 additions & 0 deletions
143
aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,143 @@ | ||
from typing import Dict, List | ||
|
||
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 BedrockAgentProperty(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 BedrockAgentRequestMedia(DictWrapper): | ||
@property | ||
def properties(self) -> List[BedrockAgentProperty]: | ||
return [BedrockAgentProperty(x) for x in self["properties"]] | ||
|
||
|
||
class BedrockAgentRequestBody(DictWrapper): | ||
@property | ||
def content(self) -> Dict[str, BedrockAgentRequestMedia]: | ||
return {k: BedrockAgentRequestMedia(v) for k, v in self["content"].items()} | ||
|
||
|
||
class BedrockAgentEvent(DictWrapper): | ||
""" | ||
Bedrock Agent input event | ||
|
||
See https://docs.aws.amazon.com/bedrock/latest/userguide/agents-create.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 api_path(self) -> str: | ||
return self["apiPath"] | ||
|
||
@property | ||
def http_method(self) -> str: | ||
return self["httpMethod"] | ||
|
||
@property | ||
def parameters(self) -> List[BedrockAgentProperty]: | ||
return [BedrockAgentProperty(x) for x in self["parameters"]] | ||
|
||
@property | ||
def request_body(self) -> BedrockAgentRequestBody: | ||
rubenfonseca marked this conversation as resolved.
Show resolved
Hide resolved
|
||
return BedrockAgentRequestBody(self["requestBody"]) | ||
|
||
@property | ||
def agent(self) -> BedrockAgentInfo: | ||
return BedrockAgentInfo(self["agent"]) | ||
|
||
@property | ||
def session_attributes(self) -> Dict[str, str]: | ||
return self["sessionAttributes"] | ||
|
||
@property | ||
def prompt_session_attributes(self) -> Dict[str, str]: | ||
return self["promptSessionAttributes"] | ||
|
||
|
||
class BedrockAgentResponseMedia(DictWrapper): | ||
@property | ||
def body(self) -> str: | ||
return self["body"] | ||
|
||
|
||
class BedrockAgentResponsePayload(DictWrapper): | ||
@property | ||
def action_group(self) -> str: | ||
return self["actionGroup"] | ||
|
||
@property | ||
def api_path(self) -> str: | ||
return self["apiPath"] | ||
|
||
@property | ||
def http_method(self) -> str: | ||
return self["httpMethod"] | ||
|
||
@property | ||
def http_status_code(self) -> int: | ||
return self["httpStatusCode"] | ||
|
||
@property | ||
def response_body(self) -> Dict[str, BedrockAgentResponseMedia]: | ||
return {k: BedrockAgentResponseMedia(v) for k, v in self["responseBody"].items()} | ||
|
||
|
||
class BedrockAgentResponseEvent(DictWrapper): | ||
rubenfonseca marked this conversation as resolved.
Show resolved
Hide resolved
|
||
""" | ||
Bedrock Agent output event | ||
|
||
See: https://docs.aws.amazon.com/bedrock/latest/userguide/agents-create.html | ||
""" | ||
|
||
@property | ||
def message_version(self) -> str: | ||
return self["messageVersion"] | ||
|
||
@property | ||
def response(self) -> BedrockAgentResponsePayload: | ||
return BedrockAgentResponsePayload(self["response"]) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,25 @@ | ||
from aws_lambda_powertools import Logger | ||
from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent, event_source | ||
from aws_lambda_powertools.utilities.typing import LambdaContext | ||
|
||
logger = Logger() | ||
|
||
|
||
@event_source(data_class=BedrockAgentEvent) | ||
def lambda_handler(event: BedrockAgentEvent, context: LambdaContext) -> dict: | ||
input_text = event.input_text | ||
|
||
logger.info(f"Bedrock Agent {event.action_group} invoked with input", input_text=input_text) | ||
|
||
return { | ||
"message_version": "1.0", | ||
"responses": [ | ||
{ | ||
"action_group": event.action_group, | ||
"api_path": event.api_path, | ||
"http_method": event.http_method, | ||
"http_status_code": 200, | ||
"response_body": {"application/json": {"body": "This is the response"}}, | ||
}, | ||
], | ||
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,16 @@ | ||
{ | ||
"actionGroup": "ClaimManagementActionGroup", | ||
"messageVersion": "1.0", | ||
"sessionId": "12345678912345", | ||
"sessionAttributes": {}, | ||
"promptSessionAttributes": {}, | ||
"inputText": "I want to claim my insurance", | ||
"agent": { | ||
"alias": "TSTALIASID", | ||
"name": "test", | ||
"version": "DRAFT", | ||
"id": "8ZXY0W8P1H" | ||
}, | ||
"httpMethod": "GET", | ||
"apiPath": "/claims" | ||
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,35 @@ | ||
{ | ||
"actionGroup": "ClaimManagementActionGroup", | ||
"messageVersion": "1.0", | ||
"sessionId": "12345678912345", | ||
"sessionAttributes": {}, | ||
"promptSessionAttributes": {}, | ||
"inputText": "Send reminders to all pending documents", | ||
"agent": { | ||
"alias": "TSTALIASID", | ||
"name": "test", | ||
"version": "DRAFT", | ||
"id": "8ZXY0W8P1H" | ||
}, | ||
"httpMethod": "POST", | ||
"apiPath": "/send-reminders", | ||
"requestBody": { | ||
"content": { | ||
"application/json": { | ||
"properties": [ | ||
{ | ||
"name": "claimId", | ||
"type": "string", | ||
"value": "20" | ||
}, | ||
{ | ||
"name": "pendingDocuments", | ||
"type": "string", | ||
"value": "social number and vat" | ||
} | ||
] | ||
} | ||
} | ||
}, | ||
"parameters": [] | ||
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,14 @@ | ||
{ | ||
"messageVersion": "1.0", | ||
"response": { | ||
"actionGroup": "ClaimManagementActionGroup", | ||
"apiPath": "/claims/{claimId}/identify-missing-documents", | ||
"httpMethod": "GET", | ||
"httpStatusCode": 200, | ||
"responseBody": { | ||
"application/json": { | ||
"body": "This is the response" | ||
} | ||
} | ||
} | ||
} |
rubenfonseca marked this conversation as resolved.
Show resolved
Hide resolved
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,57 @@ | ||
from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent | ||
from tests.functional.utils import load_event | ||
|
||
|
||
def test_bedrock_agent_event(): | ||
raw_event = load_event("bedrockAgentEvent.json") | ||
parsed_event = BedrockAgentEvent(raw_event) | ||
|
||
assert parsed_event.session_id == "12345678912345" | ||
assert parsed_event.input_text == "I want to claim my insurance" | ||
assert parsed_event.message_version == "1.0" | ||
assert parsed_event.http_method == "GET" | ||
assert parsed_event.api_path == "/claims" | ||
assert parsed_event.session_attributes == {} | ||
assert parsed_event.prompt_session_attributes == {} | ||
assert parsed_event.action_group == "ClaimManagementActionGroup" | ||
rubenfonseca marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
agent = parsed_event.agent | ||
assert agent.alias == "TSTALIASID" | ||
assert agent.name == "test" | ||
assert agent.version == "DRAFT" | ||
assert agent.id == "8ZXY0W8P1H" | ||
|
||
|
||
def test_bedrock_agent_event_with_post(): | ||
raw_event = load_event("bedrockAgentPostEvent.json") | ||
parsed_event = BedrockAgentEvent(raw_event) | ||
|
||
assert parsed_event.session_id == "12345678912345" | ||
assert parsed_event.input_text == "Send reminders to all pending documents" | ||
assert parsed_event.message_version == "1.0" | ||
assert parsed_event.http_method == "POST" | ||
assert parsed_event.api_path == "/send-reminders" | ||
assert parsed_event.session_attributes == {} | ||
assert parsed_event.prompt_session_attributes == {} | ||
assert parsed_event.action_group == "ClaimManagementActionGroup" | ||
|
||
agent = parsed_event.agent | ||
assert agent.alias == "TSTALIASID" | ||
assert agent.name == "test" | ||
assert agent.version == "DRAFT" | ||
assert agent.id == "8ZXY0W8P1H" | ||
|
||
request_body = parsed_event.request_body.content | ||
assert "application/json" in request_body | ||
|
||
json_request = request_body["application/json"] | ||
properties = json_request.properties | ||
assert len(properties) == 2 | ||
|
||
assert properties[0].name == "claimId" | ||
assert properties[0].type == "string" | ||
assert properties[0].value == "20" | ||
|
||
assert properties[1].name == "pendingDocuments" | ||
assert properties[1].type == "string" | ||
assert properties[1].value == "social number and vat" |
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.