61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
"""
|
|
RAG tool Lambda -- HNCB de-identification demo (steps 4-6).
|
|
|
|
The cloud agent calls this tool via AgentCore Gateway with a TOKEN, never a name.
|
|
This function:
|
|
1. resolves the token against the on-prem vault -> real customer_id
|
|
2. reads the raw customer record (stays here, never returned)
|
|
3. returns a token-keyed, de-identified evidence package
|
|
|
|
If the token can't be resolved, it refuses -- it never guesses or leaks.
|
|
"""
|
|
import os
|
|
import boto3
|
|
|
|
ddb = boto3.resource("dynamodb")
|
|
VAULT = ddb.Table(os.environ["VAULT_TABLE"])
|
|
CUSTOMERS = ddb.Table(os.environ["CUSTOMERS_TABLE"])
|
|
|
|
|
|
def _tool_input(event):
|
|
# AgentCore Gateway passes tool arguments in the event. Accept a few shapes
|
|
# so this is robust across Gateway target styles during the demo.
|
|
for key in ("input", "arguments", "parameters", "body"):
|
|
if isinstance(event.get(key), dict):
|
|
return event[key]
|
|
return event
|
|
|
|
|
|
def lambda_handler(event, _context):
|
|
args = _tool_input(event)
|
|
token = args.get("customer_token")
|
|
period_days = int(args.get("period_days", 90))
|
|
|
|
if not token:
|
|
return {"error": "customer_token is required"}
|
|
|
|
# 1. resolve token -> customer_id (the reversible map lives only here)
|
|
entry = VAULT.get_item(Key={"token": token}).get("Item")
|
|
if not entry:
|
|
return {"error": f"unknown token {token}; refusing to proceed"}
|
|
customer_id = entry["value"] if entry.get("type") == "CUSTOMER" else entry.get("customer_id")
|
|
if not customer_id:
|
|
return {"error": "token did not resolve to a customer record"}
|
|
|
|
# 2. read the raw record (NEVER returned to the caller)
|
|
cust = CUSTOMERS.get_item(Key={"customer_id": customer_id}).get("Item")
|
|
if not cust:
|
|
return {"error": "customer record not found"}
|
|
|
|
# 3. build a de-identified evidence package (summaries + categories only)
|
|
return {
|
|
"customer_token": token,
|
|
"period": f"last {period_days} days",
|
|
"activity_summary": cust.get("activity_summary", "n/a"),
|
|
"amount_band": cust.get("amount_band", "n/a"),
|
|
"product_type": cust.get("product_type", "n/a"),
|
|
"risk_profile": cust.get("risk_profile", "n/a"),
|
|
"sales_constraints": cust.get("sales_constraints", "n/a"),
|
|
"note": "de-identified; no raw PII in this payload",
|
|
}
|