Skip to content

fix(bedrock_agent): fix querystring field resolution #6777

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
Show file tree
Hide file tree
Changes from 2 commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,19 @@ def query_string_parameters(self) -> dict[str, str]:
parameters = self.get("parameters") or []
return {x["name"]: x["value"] for x in parameters}

@cached_property
def resolved_query_string_parameters(self) -> dict[str, list[str]]:
"""
Override the base implementation to prevent splitting parameter values by commas.

For Bedrock Agent events, parameters are already properly structured and should not
be split by commas as they might contain commas as part of their actual values
(e.g., SQL queries).
"""
# Return each parameter value as a single-item list without splitting by commas
parameters = self.get("parameters") or []
return {x["name"]: [x["value"]] for x in parameters}

@property
def resolved_headers_field(self) -> dict[str, Any]:
return {}
Expand Down
47 changes: 46 additions & 1 deletion tests/functional/event_handler/_pydantic/test_bedrock_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from typing_extensions import Annotated

from aws_lambda_powertools.event_handler import BedrockAgentResolver, BedrockResponse, Response, content_types
from aws_lambda_powertools.event_handler.openapi.params import Body
from aws_lambda_powertools.event_handler.openapi.params import Body, Query
from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
from tests.functional.utils import load_event

Expand Down Expand Up @@ -343,3 +343,48 @@ def handler() -> Optional[Dict]:

# THEN the OpenAPI schema must contain the "x-requireConfirmation" extension at the operation level
assert schema["paths"]["/"]["get"]["x-requireConfirmation"] == "ENABLED"


def test_bedrock_agent_with_comma_parameters():
# GIVEN a Bedrock Agent resolver
app = BedrockAgentResolver()
received_query = None

@app.post("/sql-query", description="Run a SQL query")
def run_sql_query(query: Annotated[str, Query()]):
nonlocal received_query
received_query = query
return {"result": "Query executed"}

# WHEN calling the event handler with a parameter containing commas
event = {
"actionGroup": "TestActionGroup",
"messageVersion": "1.0",
"sessionId": "12345678912345",
"sessionAttributes": {},
"promptSessionAttributes": {},
"inputText": "Run a SQL query",
"agent": {
"alias": "TEST",
"name": "test",
"version": "1",
"id": "test123",
},
"httpMethod": "POST",
"apiPath": "/sql-query",
"parameters": [
{
"name": "query",
"type": "string",
"value": "SELECT a.source_name, b.thing FROM table",
},
],
}

result = app(event, {})

# THEN the parameter with commas should be correctly passed to the handler
assert received_query == "SELECT a.source_name, b.thing FROM table"
assert result["response"]["httpStatusCode"] == 200
body = json.loads(result["response"]["responseBody"]["application/json"]["body"])
assert body["result"] == "Query executed"
Original file line number Diff line number Diff line change
Expand Up @@ -62,3 +62,35 @@ 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_event_with_comma_parameters():
event = {
"actionGroup": "TestActionGroup",
"messageVersion": "1.0",
"sessionId": "12345678912345",
"sessionAttributes": {},
"promptSessionAttributes": {},
"inputText": "Run a SQL query",
"agent": {
"alias": "TEST",
"name": "test",
"version": "1",
"id": "test123",
},
"httpMethod": "POST",
"apiPath": "/sql-query",
"parameters": [
{
"name": "query",
"type": "string",
"value": "SELECT a.source_name, b.thing FROM table",
},
],
}

parsed_event = BedrockAgentEvent(event)

assert parsed_event.query_string_parameters["query"] == "SELECT a.source_name, b.thing FROM table"
assert parsed_event.resolved_query_string_parameters["query"] == ["SELECT a.source_name, b.thing FROM table"]
assert len(parsed_event.resolved_query_string_parameters["query"]) == 1
Loading