How to Build an AI Agent for Code Review

1. If a function in handlers/ queries the database directly (not through repository layer), flag it
2. If a secret (API key, password) appears in code (not env config), block and alert human
3. If a new endpoint has no rate limiting, flag as security risk
4. If a query loops over results and calls database per row, flag N+1
5. If a change touches auth/* but has no test addition, flag as risky

When analyzing PR:
1. Extract files changed
2. Parse imports to find related modules
3. Fetch last 3 commits to those modules
4. Vector search for similar patterns in codebase
5. Look up module in architecture registry
6. Inject all of this as context into the LLM prompt

# 1. Index your codebase
python scripts/embed_codebase.py --output-db postgres://...

# 2. Query at review time
results = db.query("""
  SELECT file, content, 1 - (embedding <=> query_embedding) as similarity
  FROM code_chunks
  WHERE 1 - (embedding <=> query_embedding) > 0.75
  ORDER BY similarity DESC
  LIMIT 10
""")

# 1. Embed at CI time
pinecone_client.upsert(
  vectors=[
    (chunk_id, embedding, {"file": path, "content": code})
    for chunk_id, embedding, code, path in embeddings
  ]
)

# 2. Query at review time
results = index.query(query_embedding, top_k=10, include_metadata=True)

1. PR opens → webhook fires
2. Agent fetches changed files from GitHub
3. For each changed file:
   a. Vector search for similar patterns in codebase
   b. AST parse to find all imports and function calls
   c. Fetch full content of imported modules
4. Build context: [changed code] + [similar patterns] + [related modules] + [your rules]
5. Send to Claude or GPT-4 with your review prompt
6. Parse response, post comments to PR

from langgraph.graph import StateGraph, START, END
from typing import TypedDict, Annotated
import operator

class CodeReviewState(TypedDict):
    pr_number: int
    files_changed: list[str]
    file_contents: dict[str, str]
    context: str
    review_comments: Annotated[list[str], operator.add]
    issues_found: Annotated[list[dict], operator.add]

def fetch_pr_files(state: CodeReviewState) -> CodeReviewState:
    """Fetch changed files from GitHub API"""
    files = github_client.get_pr_files(state["pr_number"])
    contents = {f["filename"]: f["content"] for f in files}
    return {**state, "files_changed": list(contents.keys()), "file_contents": contents}

def retrieve_context(state: CodeReviewState) -> CodeReviewState:
    """Semantic search + AST to build context"""
    context_parts = []

    for file_path in state["files_changed"]:
        # Vector search for similar code
        similar = vector_db.search(state["file_contents"][file_path], top_k=5)
        context_parts.append(f"# Similar patterns in codebase:\n{similar}")

        # AST analysis for imports
        imports = extract_imports(state["file_contents"][file_path])
        related_code = {imp: state["file_contents"].get(imp, "NOT FOUND") for imp in imports}
        context_parts.append(f"# Related imports:\n{related_code}")

    return {**state, "context": "\n".join(context_parts)}

def analyze_with_llm(state: CodeReviewState) -> CodeReviewState:
    """Call Claude to generate review"""
    prompt = f"""
You are an expert code reviewer. Review this PR against these rules:
{YOUR_REVIEW_RULES}

Changed files:
{state['files_changed']}

Context (similar code, related modules):
{state['context']}

Full code:
{state['file_contents']}

For each issue found, respond with:
- Issue: [brief description]
- Severity: [critical|high|medium|low]
- Location: [file:line]
- Suggestion: [how to fix]

Be specific. Reference your rules and the code exactly.
"""

    response = claude_client.messages.create(
        model="claude-opus-4",
        max_tokens=2000,
        messages=[{"role": "user", "content": prompt}]
    )

    issues = parse_issues(response.content[0].text)
    comments = format_comments(issues)

    return {**state, "issues_found": issues, "review_comments": comments}

def post_comments(state: CodeReviewState) -> CodeReviewState:
    """Post review comments to PR"""
    for comment in state["review_comments"]:
        github_client.create_review_comment(state["pr_number"], comment)
    return state

# Build graph
graph = StateGraph(CodeReviewState)
graph.add_node("fetch_files", fetch_pr_files)
graph.add_node("retrieve_context", retrieve_context)
graph.add_node("analyze", analyze_with_llm)
graph.add_node("post", post_comments)

graph.add_edge(START, "fetch_files")
graph.add_edge("fetch_files", "retrieve_context")
graph.add_edge("retrieve_context", "analyze")
graph.add_edge("analyze", "post")
graph.add_edge("post", END)

agent = graph.compile()

result = agent.invoke({
    "pr_number": 1234,
    "files_changed": [],
    "file_contents": {},
    "context": "",
    "review_comments": [],
    "issues_found": []
})

from fastapi import FastAPI, Request
from pydantic import BaseModel

app = FastAPI()

class GitHubWebhook(BaseModel):
    action: str
    pull_request: dict
    repository: dict

@app.post("/webhooks/github")
async def handle_pr(payload: GitHubWebhook):
    if payload["action"] != "opened":
        return {"status": "ignored"}

    pr_number = payload["pull_request"]["number"]
    repo = payload["repository"]["full_name"]

    # Run agent asynchronously
    result = agent.invoke({
        "pr_number": pr_number,
        "files_changed": [],
        "file_contents": {},
        "context": "",
        "review_comments": [],
        "issues_found": []
    })

    return {"status": "reviewed", "issues": len(result["issues_found"])}

# Install dependencies
pip install fastapi uvicorn langgraph

# Run server
uvicorn app:app --host 0.0.0.0 --port 8000

modules:
  handlers:
    rules:
      - must_use: repository_layer
      - must_have: rate_limiting
      - must_test: all_changes
    risk_level: high

  repository:
    rules:
      - must_not: call_handlers
      - must_use: typed_queries
    risk_level: medium

  utils:
    rules:
      - must_be: pure_functions
    risk_level: low

security:
  threat_model:
    - api_injection
    - auth_bypass
    - secrets_in_code

  review_extra_strict:
    - auth/
    - payments/
    - admin/

import yaml

with open("docs/architecture.yaml") as f:
    architecture = yaml.safe_load(f)

# In analyze_with_llm():
prompt += f"\n\nArchitecture rules:\n{architecture}"

# At index time, embed your architecture docs
architecture_embeddings = {
    "stateless": embed("API handlers must be stateless. Use dependency injection for state."),
    "layer_separation": embed("Handlers should not query the database directly. Use repository layer."),
}

# At review time, for each change:
change_embedding = embed(changed_code)
violations = []

for rule_name, rule_embedding in architecture_embeddings.items():
    similarity = cosine_similarity(change_embedding, rule_embedding)
    if similarity > 0.8:  # High relevance
        violations.append({rule_name, similarity})

def analyze_with_llm(state: CodeReviewState) -> CodeReviewState:
    try:
        # ... existing code ...
    except RateLimitError:
        # Backoff and retry
        import time
        time.sleep(60)
        return analyze_with_llm(state)

    except TokenLimitError:
        # File too large. Summarize instead.
        summary = summarize_large_file(state["file_contents"])
        state["file_contents"] = {k: summary if len(v) > 20000 else v
                                 for k, v in state["file_contents"].items()}
        return analyze_with_llm(state)

    except Exception as e:
        # Post failure comment and alert
        github_client.create_review_comment(
            state["pr_number"],
            f"Code review agent failed: {str(e)}. Check logs."
        )
        logger.error(f"Review failed for PR {state['pr_number']}: {e}")
        return state

from concurrent.futures import ThreadPoolExecutor, TimeoutError
import threading

executor = ThreadPoolExecutor(max_workers=4)

@app.post("/webhooks/github")
async def handle_pr(payload: GitHubWebhook):
    try:
        future = executor.submit(agent.invoke, initial_state)
        result = future.result(timeout=300)  # 5 minute timeout
    except TimeoutError:
        logger.warning(f"Review timeout for PR {payload['pull_request']['number']}")
        # Retry later or post "review timed out" comment

    return {"status": "done"}

issues = parse_issues(response.content[0].text)

# Filter to high-confidence only
issues = [
    issue for issue in issues
    if issue["confidence"] >= 0.85  # Claude's own confidence score
]

# Agent flags: "Database called in handler (layer violation)"
# But check: does the handler have a comment explaining why?

file_content = state["file_contents"][issue["file"]]
context_around_issue = extract_lines(file_content,
                                     issue["line"] - 3,
                                     issue["line"] + 3)

if "TODO:" in context_around_issue or "HACK:" in context_around_issue:
    issue["severity"] = "low"  # Developer already knows about it
    issue["skip"] = True

SEVERITY_WEIGHTS = {
    "handlers/auth": 2.0,      # Double severity in auth code
    "handlers/payments": 2.0,
    "utils/": 0.5,             # Half severity in utils
}

for issue in issues:
    for module, weight in SEVERITY_WEIGHTS.items():
        if issue["file"].startswith(module):
            issue["severity_score"] *= weight

import hmac
import hashlib

@app.post("/webhooks/github")
async def handle_pr(request: Request):
    signature = request.headers.get("X-Hub-Signature-256")
    body = await request.body()

    expected = "sha256=" + hmac.new(
        WEBHOOK_SECRET.encode(),
        body,
        hashlib.sha256
    ).hexdigest()

    if not hmac.compare_digest(signature, expected):
        return {"error": "invalid signature"}, 401

    # ... handle PR ...