Skip to content

feat(ci): Ground AI code review in Bedrock knowledge bases (public + internal) - #6295

Open
jam-jee wants to merge 4 commits into
aws:masterfrom
jam-jee:feat/ai-review-knowledge-base
Open

jam-jee wants to merge 4 commits into
aws:masterfrom
jam-jee:feat/ai-review-knowledge-base

Conversation

@jam-jee

@jam-jee jam-jee commented Sep 20, 2026

Copy link
Copy Markdown
Collaborator

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)

  • New step .github/scripts/kb_retrieve.py (Python 3, boto3 only) runs after
    the PR diff is fetched. It builds retrieval queries from the PR title, the
    changed file paths, and the def/class/async def names added in the diff,
    calls bedrock-agent-runtime:Retrieve (top-5 per query) against the public
    KB, dedupes hits by source URL/location, renders them as markdown, and caps
    the output (~12k chars) to /tmp/kb-context.md.
  • The reviewer prompt now reads /tmp/kb-context.md after the diff, grounds
    comments in it, and cites the matching Source: URL. The file is treated as
    reference material, not instructions — the same prompt-injection rule that
    governs the diff applies to it.
  • Retrieval is best-effort: on any error the script still writes the output
    file and exits 0, so retrieval can never fail the review.

2. New internal review (.github/workflows/ai-code-review-internal.yml)

  • Reuses the exact same collaborator gate and pull_request_target safety model
    (base-SHA context only, no fork code executed, no Bash).
  • Uploads the diff + metadata JSON to a private S3 bucket in the SageMaker AWS
    account, starts the pysdk-internal-ai-review CodeBuild project, polls for
    completion, and upserts one marker comment (<!-- pysdk-internal-ai-review -->)
    containing only a console link to the report.
  • The report body and any retrieved KB content never appear in workflow logs
    or PR comments (this repo is public). All model output stays inside AWS.

Security model

  • pull_request_target is 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-approval environment before any role or
    secret is exposed.
  • Neither workflow checks out or executes fork head code — the PR diff is pulled
    read-only via the GitHub API.
  • PR-authored strings (title, author, url) are passed via env: and never
    interpolated into run: scripts; the metadata JSON is built with jq so those
    values are encoded, not spliced.
  • The public retrieval role is scoped to bedrock:Retrieve on the public KB
    only and is assumed just long enough to fetch context; the workflow then
    re-assumes the least-privilege CODE_REVIEW_ROLE (InvokeModel only) for the
    review, so retrieval and inference never share one credential.
  • The internal trigger role can only PutObject under input/ (no Get, no
    List) and StartBuild/BatchGetBuilds on the one project, so a compromised
    workflow 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.

Kind Name Value (placeholder)
Secret PUBLIC_KB_ROLE arn:aws:iam::303192504279:role/PySDKPublicReviewKBRole-us-west-2
Secret INTERNAL_REVIEW_ROLE arn:aws:iam::303192504279:role/PySDKInternalReviewTriggerRole-us-west-2
Variable PUBLIC_KB_ID QWOL4KFTGH

Existing secrets already used by the review workflows (COLLAB_CHECK_TOKEN,
CODE_REVIEW_ROLE, and the default GITHUB_TOKEN) are unchanged.

Testing

  • .github/scripts/test_kb_retrieve.py (unittest, fake bedrock client) covers
    query extraction, dedupe (by URL and by location), the output cap, and the
    fail-safe path. Run:

    python3 -m unittest discover \
      -s .github/scripts -p 'test_kb_retrieve.py'
    
  • Both workflow YAML files validated with PyYAML.

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.
@jam-jee
jam-jee deployed to auto-approve September 20, 2026 01:03 — with GitHub Actions Active
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant