1"""
2Knowledge Base shim Lambda — AgentCore Gateway target.
3
4WHY: the native Gateway KB connector only takes a MANAGED Bedrock KB and a fixed
5retrieval contract. For VECTOR / KENDRA / SQL KBs, or to preserve a managed KB's
6non-default retrieval config (reranker, metadata filter, hybrid override, top-k),
7expose retrieval as this Lambda: it presents one MCP tool, calls
8`bedrock-agent-runtime:Retrieve` against the source KB, and returns MCP-shaped
9passages.
10
11Gateway invokes a Lambda target with the tool arguments flat in `event` and the
12tool name in context.client_context.custom['bedrockAgentCoreToolName']
13("<target>___<tool>"). The handler also tolerates a direct {"query": ...} for local
14testing.
15"""
16# <<< RENDER: delete this whole block after substituting the tokens below.
17#   {{KB_ID}}        - the source knowledge base id

kb_shim.py.tmpl on GitHub

# {{TOP_K}} - numberOfResults from the source KB association (default 5)

# {{SEARCH_TYPE}} - "HYBRID" | "SEMANTIC" | "" (empty means KB default)

# <<< /RENDER

import boto3

# Rendered at migration time — the Gateway lambda code target has no

# environment-variable support, so all config is baked in as literals.

_KB_ID = "{{KB_ID}}"

_TOP_K = int("{{TOP_K}}")

_SEARCH_TYPE = "{{SEARCH_TYPE}}".strip().upper()

_MAX_QUERY_LEN = 1000 # reject oversized queries — abuse / runaway retrieval cost

_runtime = boto3.client("bedrock-agent-runtime")

def _vector_search_config():

cfg = {"numberOfResults": _TOP_K}

if _SEARCH_TYPE in ("HYBRID", "SEMANTIC"):

cfg["overrideSearchType"] = _SEARCH_TYPE

# <<< OPTIONAL: metadata_filter

# Render the source KB association's `filter` here when present.

# cfg["filter"] = {{METADATA_FILTER_JSON}}

# <<< /OPTIONAL: metadata_filter

# <<< OPTIONAL: reranking

# Render the source KB association's `rerankingConfiguration` here when present.

# (Set on retrievalConfiguration, not vectorSearchConfiguration — see the

# bedrock-agent-runtime Retrieve API shape.)

# <<< /OPTIONAL: reranking

return cfg

def lambda_handler(event, context):

# single-tool target: Gateway passes the tool args flat in `event`

args = event if isinstance(event, dict) else {}

query = args.get("query") or args.get("text") or ""

if not isinstance(query, str) or not query:

return {"error": "Missing required argument 'query'."}

if len(query) > _MAX_QUERY_LEN:

return {"error": f"Query exceeds {_MAX_QUERY_LEN} chars."}

vector_cfg = _vector_search_config()

max_results = args.get("max_results")

if isinstance(max_results, int) and max_results > 0:

vector_cfg["numberOfResults"] = max_results

resp = _runtime.retrieve(

knowledgeBaseId=_KB_ID,

retrievalQuery={"text": query},

retrievalConfiguration={"vectorSearchConfiguration": vector_cfg},

)

# Returns KB passage text + source location verbatim. If the KB holds

# sensitive data (PII, financial records), review what it returns and redact

# or filter fields here before returning, and keep this Lambda's CloudWatch

# log group KMS-encrypted (see references/deploy.md) — do not log full passages.

results = [

{

"text": r.get("content", {}).get("text", ""),

"source": r.get("location", {}),

"score": r.get("score"),

}

for r in resp.get("retrievalResults", [])

]

return {"results": results, "count": len(results)}