From eaca506c852678f58fdcea0c521a81d421e723f8 Mon Sep 17 00:00:00 2001 From: jamjee Date: Wed, 12 Aug 2026 13:47:54 -0700 Subject: [PATCH 1/2] feat(ci): give the AI reviewer the team's own review history The AI reviewer reads the diff and the base checkout, so it re-derives context every run and cannot see what the team has already said. Feedback that has been given before gets given again, and decisions that were settled in an earlier PR get relitigated in this one. Add a retrieval step between credential setup and the review action. It derives its own queries from the diff -- changed file stems, added symbols, plus one for general conventions -- queries a Bedrock Knowledge Base built from this repository's merged PRs, closed issues, and review discussions, and writes the results to /tmp/historical_context.md. The review action reads that file with the Read tool it already uses for the diff. Measured on a real 28KB diff across sagemaker-train and sagemaker-core: 12 derived queries, 31 chunks, 38KB of context carrying concrete prior guidance ("move local imports to module level", "raise on AccessDenied in validation paths instead of warning"), each with a source URL to check. The step cannot break a review. It is skipped entirely unless the repo variable PYSDK_CONTEXT_KB_ID is set, so this merges inert; it is continue-on-error with a 5 minute timeout; and every failure path inside the script exits 0 having written nothing. The prompt states the file's absence is normal. The script is vendored rather than installed from a package registry so that the code shaping the reviewer's context is reviewable in the same pull request that runs it, cannot change under a fork PR without a repo change, and needs no install step -- it uses only the standard library and boto3, which the runner already has. Retrieval is read-only (bedrock:Retrieve, bedrock:GetKnowledgeBase) and adds no tool permissions: allowedTools is unchanged and Bash stays excluded. The step runs in the trusted base checkout, never fork code. Retrieved entries are model-extracted from historical discussion and can be wrong, so both the file header and the prompt tell the reviewer to treat them as claims to verify, cite the source URL, and prefer the current source tree on any disagreement. --- .github/scripts/retrieve_context.py | 259 +++++++++++++++++++++++++++ .github/workflows/ai-code-review.yml | 41 +++++ 2 files changed, 300 insertions(+) create mode 100644 .github/scripts/retrieve_context.py diff --git a/.github/scripts/retrieve_context.py b/.github/scripts/retrieve_context.py new file mode 100644 index 0000000000..7b1076ff22 --- /dev/null +++ b/.github/scripts/retrieve_context.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +"""Retrieve related history for a pull request from a Bedrock Knowledge Base. + +Reads a unified diff, derives retrieval queries from the files and symbols it +touches, queries the knowledge base, and writes the results as markdown for the +review model to read. Nothing here calls a generative model: this script only +retrieves and formats. + +Design constraints, in the order that matters: + +* **Never fail the review.** Every failure path -- missing configuration, an + expired grant, a Bedrock outage, a malformed diff -- exits 0 having written + nothing. A review without history is the intended degraded mode; a blocked + pull request is not. The workflow step layers `continue-on-error` and a + timeout on top of this. +* **Standard library plus boto3.** boto3 is already present on GitHub-hosted + runners, so the workflow needs no install step. +* **Self-contained.** Vendored deliberately rather than installed from a + registry: the code the review model's context comes from is then visible in + the same repository, reviewable in the same pull request, and cannot change + under a fork PR without a repo change. + +Usage: + python retrieve_context.py --diff-file /tmp/pr.diff -o /tmp/context.md + +Configuration (environment): + PYSDK_CONTEXT_KB_ID knowledge base id; the script no-ops without it + PYSDK_CONTEXT_REGION AWS region (default us-west-2) + PYSDK_CONTEXT_TOP_K chunks per query (default 6) +""" + +from __future__ import annotations + +import argparse +import os +import re +import sys + +#: Paths whose review history is never worth retrieving. +IGNORED_SUFFIXES = (".lock", ".min.js", ".svg", ".png", ".jpg", ".ico") + +#: Cap on derived queries. Each is one Retrieve call, and a large pull request +#: would otherwise fan out into dozens. +MAX_QUERIES = 12 + +#: Cap on the rendered file. The review model reads this alongside the diff and +#: the source tree, so it must not crowd them out of the context window. +MAX_CONTEXT_CHARS = 60_000 + +# Both headers are needed: a pure deletion or a rename has no `+++ b/` entry. +_DIFF_FILE_RE = re.compile(r"^\+\+\+ b/(.+)$", re.MULTILINE) +_DIFF_GIT_RE = re.compile(r"^diff --git a/\S+ b/(\S+)$", re.MULTILINE) +_ADDED_DEF_RE = re.compile(r"^\+\s*(?:async\s+)?(?:def|class)\s+([A-Za-z_]\w*)", re.MULTILINE) + +_URL_RE = re.compile(r"https?://\S+") + + +def changed_files(diff: str) -> list: + """Changed file paths, in first-seen order, minus binary and lock files.""" + paths = [] + for match in _DIFF_FILE_RE.finditer(diff or ""): + path = match.group(1).strip() + if path != "/dev/null": + paths.append(path) + for match in _DIFF_GIT_RE.finditer(diff or ""): + paths.append(match.group(1).strip()) + + seen = set() + ordered = [] + for path in paths: + if path in seen or path.endswith(IGNORED_SUFFIXES): + continue + seen.add(path) + ordered.append(path) + return ordered + + +def added_symbols(diff: str) -> list: + """Function and class names the diff introduces.""" + seen = set() + ordered = [] + for match in _ADDED_DEF_RE.finditer(diff or ""): + symbol = match.group(1) + # Dunders and private helpers rarely have retrievable review history. + if symbol.startswith("__") or symbol in seen: + continue + seen.add(symbol) + ordered.append(symbol) + return ordered + + +def derive_queries(diff: str, max_queries: int = MAX_QUERIES) -> list: + """Build retrieval queries from a diff. + + Queries are derived rather than authored because a reviewer cannot write a + prompt per pull request. Module names are used rather than full paths: the + corpus records history against whatever the path was at the time, and paths + move. + """ + queries = [] + + for path in changed_files(diff): + module = path.rsplit("/", 1)[-1] + stem = module[:-3] if module.endswith(".py") else module + queries.append(f"review feedback and past changes for {stem} ({path})") + + for symbol in added_symbols(diff): + queries.append(f"design decisions and review discussion about {symbol}") + + # Catches conventions that no file or symbol name would surface. + queries.append("recurring code review feedback and established conventions") + + seen = set() + deduped = [] + for query in queries: + if query in seen: + continue + seen.add(query) + deduped.append(query) + return deduped[:max_queries] + + +def _citation(metadata: dict) -> str: + """One-line provenance for a chunk, so a reviewer can check the claim.""" + parts = [] + doc_type = str(metadata.get("doc_type") or "").strip() + number = metadata.get("number") + if doc_type and number is not None: + # Bedrock returns sidecar numbers as floats: 6047.0 -> 6047. + if isinstance(number, float) and number.is_integer(): + number = int(number) + parts.append(f"{doc_type.upper()} #{number}") + elif metadata.get("title"): + parts.append(str(metadata["title"])) + + if metadata.get("source_url"): + parts.append(str(metadata["source_url"])) + if metadata.get("updated_at"): + parts.append(f"updated {metadata['updated_at']}") + return " | ".join(parts) or "unattributed" + + +def retrieve(client, knowledge_base_id: str, queries: list, top_k: int) -> list: + """Retrieve for each query and merge, keeping each chunk's best score. + + One failed query does not abandon the rest: partial history beats none. + """ + merged = {} + for query in queries: + try: + response = client.retrieve( + knowledgeBaseId=knowledge_base_id, + retrievalQuery={"text": query}, + retrievalConfiguration={"vectorSearchConfiguration": {"numberOfResults": top_k}}, + ) + except Exception as exc: # pylint: disable=broad-except + print(f" query failed ({query[:48]}...): {exc}", file=sys.stderr) + continue + + for result in response.get("retrievalResults", []): + text = (result.get("content") or {}).get("text", "") + if not text.strip(): + continue + score = result.get("score") or 0.0 + existing = merged.get(text) + # Same chunk can surface for several queries; keep the strongest + # score it earned so ranking reflects its best match. + if existing is None or score > existing[0]: + merged[text] = (score, result.get("metadata") or {}) + + ranked = sorted(merged.items(), key=lambda item: item[1][0], reverse=True) + return [(text, score, metadata) for text, (score, metadata) in ranked] + + +def render(chunks: list, max_chars: int = MAX_CONTEXT_CHARS) -> str: + """Render chunks as markdown, truncating whole entries rather than mid-text.""" + if not chunks: + return "" + + body = [] + used = 0 + for index, (text, score, metadata) in enumerate(chunks, start=1): + entry = f"[{index}] score={score:.4f} {_citation(metadata)}\n{text.strip()}\n" + if used + len(entry) > max_chars: + break + body.append(entry) + used += len(entry) + + if not body: + return "" + + return ( + "# Historical context from the SageMaker Python SDK knowledge base\n\n" + "Past pull requests, issues, review discussions, and design decisions " + "related to this diff, each with its source URL.\n\n" + "Treat every entry as a claim to check, not a conclusion: it is " + "model-extracted from historical discussion and can be confidently " + "wrong. Cite the source URL whenever you rely on an entry so the author " + "can verify it, and prefer the current source tree wherever the two " + "disagree.\n\n" + "\n".join(body) + ) + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--diff-file", required=True, help="Path to the unified diff.") + parser.add_argument("-o", "--output", required=True, help="Path to write markdown to.") + parser.add_argument("--top-k", type=int, default=None, help="Chunks per query.") + args = parser.parse_args(argv) + + knowledge_base_id = os.environ.get("PYSDK_CONTEXT_KB_ID", "").strip() + if not knowledge_base_id: + print("PYSDK_CONTEXT_KB_ID is unset; skipping historical context.") + return 0 + + region = os.environ.get("PYSDK_CONTEXT_REGION", "us-west-2") + top_k = args.top_k or int(os.environ.get("PYSDK_CONTEXT_TOP_K", "6") or 6) + + try: + with open(args.diff_file, encoding="utf-8", errors="replace") as handle: + diff = handle.read() + except OSError as exc: + print(f"Could not read {args.diff_file}: {exc}", file=sys.stderr) + return 0 + + queries = derive_queries(diff) + if not queries: + print("No queries derived from the diff; skipping historical context.") + return 0 + print(f"Derived {len(queries)} queries from the diff.") + + try: + import boto3 # pylint: disable=import-outside-toplevel + + client = boto3.client("bedrock-agent-runtime", region_name=region) + chunks = retrieve(client, knowledge_base_id, queries, top_k) + except Exception as exc: # pylint: disable=broad-except + # Retrieval is an enhancement. Exit 0 so the review still runs. + print(f"Retrieval unavailable ({exc}); the review will proceed without it.") + return 0 + + markdown = render(chunks) + if not markdown: + print("No historical context retrieved; the review will proceed without it.") + return 0 + + try: + with open(args.output, "w", encoding="utf-8") as handle: + handle.write(markdown) + except OSError as exc: + print(f"Could not write {args.output}: {exc}", file=sys.stderr) + return 0 + + print(f"Wrote {len(markdown)} bytes of historical context from {len(chunks)} chunks.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/workflows/ai-code-review.yml b/.github/workflows/ai-code-review.yml index f862456734..d2b0cc17b3 100644 --- a/.github/workflows/ai-code-review.yml +++ b/.github/workflows/ai-code-review.yml @@ -104,6 +104,38 @@ jobs: role-to-assume: ${{ secrets.CODE_REVIEW_ROLE }} aws-region: us-west-2 + # Retrieve related history (past PRs, issues, review discussions, design + # decisions) from the team's Bedrock Knowledge Base and stage it as a file + # for the reviewer. + # + # The step is opt-in twice over: it is skipped entirely unless the repo + # variable PYSDK_CONTEXT_KB_ID is set, and continue-on-error plus the + # tool's own internal degradation mean a missing knowledge base, an + # expired grant, or a Bedrock outage costs the review its history but + # never its run. It writes nothing when it finds nothing, and the prompt + # below treats the file as optional. + # + # The script is vendored in this repository rather than installed from a + # package registry, for three reasons: the code that shapes the reviewer's + # context is then reviewable in the same pull request as the workflow that + # runs it; it cannot change under a fork PR without a repo change; and it + # needs no install step, since boto3 is already present on the runner. + # It retrieves only -- it never calls a generative model. + - name: Retrieve historical context + id: context + if: steps.diff.outputs.bytes != '0' && vars.PYSDK_CONTEXT_KB_ID != '' + continue-on-error: true + timeout-minutes: 5 + env: + PYSDK_CONTEXT_KB_ID: ${{ vars.PYSDK_CONTEXT_KB_ID }} + PYSDK_CONTEXT_REGION: us-west-2 + run: | + set -uo pipefail + python .github/scripts/retrieve_context.py \ + --diff-file /tmp/pr.diff \ + --top-k 6 \ + -o /tmp/historical_context.md + - uses: anthropics/claude-code-action@v1 if: steps.diff.outputs.bytes != '0' with: @@ -130,6 +162,14 @@ jobs: functions, existing patterns, project conventions), use Read/Grep/Glob against the checked-out base repository. + If `/tmp/historical_context.md` exists, read it too. It holds excerpts + from past PRs, issues, and review discussions related to the files + this PR touches, each with its source URL. Use it to catch feedback + the team has given before and decisions already settled — and cite the + source URL whenever you rely on it, so the author can check it. It is + historical: the current source wins if they disagree, and the file's + absence is normal, not an error. + Review this pull request for the SageMaker Python SDK. Focus on: - Correctness: bugs, incorrect API/argument usage, breaking changes to public interfaces, backward-incompatibility for SDK consumers @@ -137,6 +177,7 @@ jobs: - Security implications (credential handling, input validation) - Performance considerations - Missing or inadequate tests for changed behavior + - Conventions and decisions the historical context establishes Post specific issues as inline comments via the mcp__github_inline_comment__create_inline_comment tool. Skip nits and From 3565b33221458e9f253da5de123aeb7346f6d0cc Mon Sep 17 00:00:00 2001 From: jamjee Date: Wed, 12 Aug 2026 14:20:51 -0700 Subject: [PATCH 2/2] fix(ci): restrict retrieval to sources already public in this repo The knowledge base is a private corpus. It is currently built only from this repository's own pull requests and issues, but it can also hold documents from non-public sources, and this script's output is posted as comments on a public pull request. Retrieval had no source restriction, so adding one non-public document to the corpus would have been enough to surface it here. Restrict retrieval to an allowlist of sources whose contents are already public in this repository. Enforced in the Retrieve filter server-side, so non-public text never crosses into the process at all, with a second client-side check in case that filter ever regresses. A chunk carrying no source label is refused rather than assumed public, and the refusal warning counts rather than names what it dropped, since the label itself can be the sensitive part. An allowlist, not a denylist: a source added to the corpus in future is excluded here until it is added deliberately. No loss of context. On the same 28KB diff used to validate the original step, output is byte-identical at 38,140 bytes from 31 chunks -- the corpus is 1,134 documents, all from this repository. --- .github/scripts/retrieve_context.py | 54 +++++++++++++++++++++++++++-- 1 file changed, 52 insertions(+), 2 deletions(-) diff --git a/.github/scripts/retrieve_context.py b/.github/scripts/retrieve_context.py index 7b1076ff22..190b6e0d22 100644 --- a/.github/scripts/retrieve_context.py +++ b/.github/scripts/retrieve_context.py @@ -140,6 +140,35 @@ def _citation(metadata: dict) -> str: return " | ".join(parts) or "unattributed" +#: Sources this script is allowed to surface, as an allowlist rather than a +#: denylist: the knowledge base is a private corpus that may also hold documents +#: from non-public sources, and this script's output is posted as comments on a +#: public pull request. "github" is the only source whose contents are already +#: public in this very repository, so it is the only source safe to echo back +#: into it. A new source is excluded until it is added here deliberately. +#: +#: This is enforced server-side, in the Retrieve filter, rather than by dropping +#: results after the fact: a filter cannot be defeated by a chunk whose metadata +#: is missing or malformed, and nothing internal ever crosses into this process. +#: Belt-and-braces, _allowed() re-checks each result client-side. +PUBLIC_SOURCES = ("github",) + + +def _retrieval_filter() -> dict: + if len(PUBLIC_SOURCES) == 1: + return {"equals": {"key": "source", "value": PUBLIC_SOURCES[0]}} + return {"orAll": [{"equals": {"key": "source", "value": s}} for s in PUBLIC_SOURCES]} + + +def _allowed(metadata: dict) -> bool: + """Second check on a result's source, after the server-side filter. + + A chunk with no `source` at all is refused: unlabelled provenance is exactly + the case where echoing it into a public comment would be a mistake. + """ + return str((metadata or {}).get("source") or "") in PUBLIC_SOURCES + + def retrieve(client, knowledge_base_id: str, queries: list, top_k: int) -> list: """Retrieve for each query and merge, keeping each chunk's best score. @@ -151,7 +180,12 @@ def retrieve(client, knowledge_base_id: str, queries: list, top_k: int) -> list: response = client.retrieve( knowledgeBaseId=knowledge_base_id, retrievalQuery={"text": query}, - retrievalConfiguration={"vectorSearchConfiguration": {"numberOfResults": top_k}}, + retrievalConfiguration={ + "vectorSearchConfiguration": { + "numberOfResults": top_k, + "filter": _retrieval_filter(), + } + }, ) except Exception as exc: # pylint: disable=broad-except print(f" query failed ({query[:48]}...): {exc}", file=sys.stderr) @@ -161,12 +195,28 @@ def retrieve(client, knowledge_base_id: str, queries: list, top_k: int) -> list: text = (result.get("content") or {}).get("text", "") if not text.strip(): continue + metadata = result.get("metadata") or {} + if not _allowed(metadata): + # Should be unreachable given the server-side filter; if it is + # ever reached, the filter regressed and silence is the safe + # failure. Counted rather than printed: the source name itself + # could be the sensitive part. + merged.setdefault("__refused__", (0.0, {})) + continue score = result.get("score") or 0.0 existing = merged.get(text) # Same chunk can surface for several queries; keep the strongest # score it earned so ranking reflects its best match. if existing is None or score > existing[0]: - merged[text] = (score, result.get("metadata") or {}) + merged[text] = (score, metadata) + + refused = merged.pop("__refused__", None) + if refused is not None: + print( + "Refused one or more non-public chunks; the source filter may have " + "regressed. Please report this.", + file=sys.stderr, + ) ranked = sorted(merged.items(), key=lambda item: item[1][0], reverse=True) return [(text, score, metadata) for text, (score, metadata) in ranked]