Conversation
Add .github/scripts/kb_retrieve.py: builds retrieval queries from the PR title, changed file paths, and def/class names on added diff lines, runs bedrock-agent-runtime retrieve (top-5 per query) against the public review knowledge base, dedupes by source location, renders capped markdown, and fails safe (writes the out file and exits 0 on any error so retrieval can never break the review). Includes a unittest suite with a fake client covering query extraction, dedupe, cap, and the failure path.
Assume the scoped PUBLIC_KB_ROLE, run kb_retrieve.py to fetch team precedent into /tmp/kb-context.md, then re-assume the least-privilege CODE_REVIEW_ROLE for inference so retrieval and model access never share a credential. Extend the reviewer prompt to read the KB context, cite its Source URLs, and treat it as reference material subject to the same prompt-injection rule as the diff. PR title is passed via env, never interpolated into run scripts; the pull_request_target safety model is unchanged.
Add ai-code-review-internal.yml: same collaborator gate and pull_request_target safety model as the public workflow. Uploads the PR diff and metadata to a private S3 bucket via the INTERNAL_REVIEW_ROLE, starts and polls the pysdk-internal-ai-review CodeBuild project, and upserts a single marker comment carrying only a console link. The report body and retrieved KB content never appear in logs or comments.
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Add knowledge-base-grounded AI code review (public + internal)
What this changes
This PR wires the SageMaker Python SDK's AI code review into two Bedrock
Knowledge Bases so review comments are grounded in team precedent and
conventions instead of the model's priors alone.
1. Public review gets KB context (
.github/workflows/ai-code-review.yml).github/scripts/kb_retrieve.py(Python 3, boto3 only) runs afterthe PR diff is fetched. It builds retrieval queries from the PR title, the
changed file paths, and the
def/class/async defnames added in the diff,calls
bedrock-agent-runtime:Retrieve(top-5 per query) against the publicKB, dedupes hits by source URL/location, renders them as markdown, and caps
the output (~12k chars) to
/tmp/kb-context.md./tmp/kb-context.mdafter the diff, groundscomments in it, and cites the matching
Source:URL. The file is treated asreference material, not instructions — the same prompt-injection rule that
governs the diff applies to it.
file and exits 0, so retrieval can never fail the review.
2. New internal review (
.github/workflows/ai-code-review-internal.yml)pull_request_targetsafety model(base-SHA context only, no fork code executed, no Bash).
account, starts the
pysdk-internal-ai-reviewCodeBuild project, polls forcompletion, and upserts one marker comment (
<!-- pysdk-internal-ai-review -->)containing only a console link to the report.
or PR comments (this repo is public). All model output stays inside AWS.
Security model
pull_request_targetis required to hold the OIDC roles and post comments.Both workflows keep the identical collaborator gate: fork/external PRs require
maintainer approval via the
manual-approvalenvironment before any role orsecret is exposed.
read-only via the GitHub API.
env:and neverinterpolated into
run:scripts; the metadata JSON is built withjqso thosevalues are encoded, not spliced.
bedrock:Retrieveon the public KBonly and is assumed just long enough to fetch context; the workflow then
re-assumes the least-privilege
CODE_REVIEW_ROLE(InvokeModel only) for thereview, so retrieval and inference never share one credential.
PutObjectunderinput/(no Get, noList) and
StartBuild/BatchGetBuildson the one project, so a compromisedworkflow cannot read any internal review back out.
Repo settings a maintainer must add
Set these once, after the KB/CodeBuild infrastructure is deployed. Values come
from the SSM parameters written by the CDK stacks.
PUBLIC_KB_ROLEarn:aws:iam::303192504279:role/PySDKPublicReviewKBRole-us-west-2INTERNAL_REVIEW_ROLEarn:aws:iam::303192504279:role/PySDKInternalReviewTriggerRole-us-west-2PUBLIC_KB_IDQWOL4KFTGHExisting secrets already used by the review workflows (
COLLAB_CHECK_TOKEN,CODE_REVIEW_ROLE, and the defaultGITHUB_TOKEN) are unchanged.Testing
.github/scripts/test_kb_retrieve.py(unittest, fake bedrock client) coversquery extraction, dedupe (by URL and by location), the output cap, and the
fail-safe path. Run:
Both workflow YAML files validated with PyYAML.