Skip to content

Repository files navigation

Code Review Agent

An automated, RAG-enhanced code review system that analyzes GitHub pull-request diffs for anti-patterns. It fine-tunes a CodeBERT/StarCoder classifier with PEFT/LoRA, exports it to a CPU-optimized quantized ONNX model, and serves reviews through a FastAPI backend, a GitHub PR flow, and a VS Code extension. Repository-specific context is grounded via a PostgreSQL + pgvector retrieval pipeline.

Project status: working MVP. The full pipeline runs end-to-end — diff → anti-pattern classification → repository-grounded context → JSON response or GitHub review comment. The bundled model is trained on a small synthetic dataset so the demo works out of the box; it demonstrates the system, not production accuracy. A real deployment needs a model trained on a real, human-labeled corpus of PR diffs (see Limitations).

Table of Contents

Quick Demo

Two commands. The first builds the model (one-time, ~3 min); the second starts the whole stack.

# Prereqs: Python 3.11, Docker. Create a venv and install first:
python3.11 -m venv .venv && source .venv/bin/activate && pip install -r requirements.txt

# 1) Generate synthetic data, train the demo model, export to ONNX (one-time)
scripts/setup_demo.sh

# 2) Start Postgres+pgvector, migrate, ingest this repo, launch the API
scripts/run_demo.sh

Then, in another terminal, send it some diffs:

scripts/sample_review.sh

Expected output — each diff is classified and grounded with retrieved repo context:

### SQL injection (expect: security)
  label=security  confidence=0.61  time=66ms
  context: 3 repo snippet(s) retrieved
### Nested loop (expect: performance)
  label=performance  confidence=0.40  time=60ms
  context: 3 repo snippet(s) retrieved
### Bare except (expect: error_handling)
  label=error_handling  confidence=0.36  time=53ms
  context: 2 repo snippet(s) retrieved
### Typed helper (expect: clean)
  label=clean  confidence=0.66  time=52ms
  context: 1 repo snippet(s) retrieved

How You Use It

There are three ways to run a review, all working today:

  1. HTTP APIPOST /api/v1/review/diff with a raw diff (no GitHub needed), or POST /api/v1/review/pr with a repo, PR number, and token. Any tool that speaks HTTP can call it. See Running the API.
  2. VS Code extension — press F5 in the extension/ folder to launch an Extension Development Host, then run "Code Review Agent: Review Current Diff" to review your working changes against the API. See VS Code Extension.
  3. GitHub PR flow/api/v1/review/pr fetches a real PR's diff and posts inline review comments back to GitHub.

Architecture

The system runs in two phases:

Offline (training) — PR diffs are parsed into tokenized sequences, a classifier is fine-tuned with LoRA adapters, and the merged model is exported to int8-quantized ONNX for fast CPU inference.

Online (inference) — A diff arrives via the API or a GitHub webhook. The relevant repository context is retrieved from pgvector, the ONNX model classifies each changed hunk, and anti-pattern findings are returned as JSON or posted back as inline GitHub review comments.

GitHub / VS Code ──▶ FastAPI ──▶ pgvector retrieval ──▶ ONNX Runtime ──▶ review comments

See docs:ARCHITECTURE.md for data flows, the database schema, and API contracts, and docs:PRD.md for product goals and SLAs.

Repository Layout

Path Description
ml/ Dataset parsing (data.py), LoRA fine-tuning (train.py), ONNX export (export.py)
backend/api/ FastAPI app (main.py) and routes (routes.py)
backend/inference/ ONNX Runtime inference engine
backend/rag/ LangChain repository ingestion into pgvector
backend/db/ SQLAlchemy models, async session, Alembic migrations
backend/github/ Async GitHub REST client for diffs and review comments
infrastructure/ docker-compose.yml (Postgres + pgvector) and Lambda-ready Dockerfile
extension/ VS Code extension (TypeScript)
docs:*.md PRD, architecture, task roadmap, and changelog

Prerequisites

  • Python 3.11 (the ML/serving stack does not yet support 3.12+)
  • Docker (for the PostgreSQL + pgvector database)
  • Node.js 18+ and npm (for the VS Code extension)

Setup

# 1. Create a virtual environment and install dependencies
python3.11 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

# 2. Start PostgreSQL + pgvector
#    Create infrastructure/.env first (POSTGRES_PASSWORD is required):
#
#      POSTGRES_USER=review_agent
#      POSTGRES_PASSWORD=<choose-a-password>
#      POSTGRES_DB=code_review
#      POSTGRES_PORT=5432
#
docker compose -f infrastructure/docker-compose.yml --env-file infrastructure/.env up -d

# 3. Point the app at the database (required — there is no default)
export DATABASE_URL="postgresql+asyncpg://review_agent:<password>@localhost:5432/code_review"

# 4. Apply the schema
alembic -c backend/db/alembic.ini upgrade head

ML Pipeline

All three stages are driven by JSON config files. Training data is JSONL, one record per line: {"diff": "<unified diff>", "label": <int>}.

# Fine-tune with LoRA (metrics tracked in Weights & Biases; set wandb_enabled=false to skip)
python -m ml.train data/train_config.json

# Merge adapters, export to ONNX, and apply int8 dynamic quantization
python -m ml.export data/export_config.json

Training reports macro and per-class F1/precision/recall and saves the best checkpoint (by f1_macro) alongside its tokenizer. Export produces a self-contained quantized model directory ready for inference.

RAG Ingestion

Index a local checkout of the repository you want to review. This chunks .py and .md files, embeds them with all-MiniLM-L6-v2, and upserts the vectors into pgvector.

export DATABASE_URL="postgresql+asyncpg://review_agent:<password>@localhost:5432/code_review"
python -m backend.rag.ingest data/ingest_config.json

Running the API

export DATABASE_URL="postgresql+asyncpg://review_agent:<password>@localhost:5432/code_review"
export MODEL_DIR=outputs/onnx/quantized   # directory containing the exported ONNX model
# export API_KEY=<secret>                 # optional: require X-API-Key on all routes

python -m uvicorn backend.api.main:app --host 0.0.0.0 --port 8080

Interactive docs are served at http://localhost:8080/docs.

Endpoints

POST /api/v1/review/diff — classify a raw unified diff (no GitHub interaction):

curl -X POST http://localhost:8080/api/v1/review/diff \
  -H "Content-Type: application/json" \
  -d '{"diff": "--- a/app.py\n+++ b/app.py\n@@ -1,2 +1,3 @@\n def f():\n+    eval(x)\n     return 1"}'

POST /api/v1/review/pr — fetch a PR's diff from GitHub, classify each file, and post inline review comments:

curl -X POST http://localhost:8080/api/v1/review/pr \
  -H "Content-Type: application/json" \
  -d '{"repository": "owner/repo", "pull_request_number": 123, "github_token": "ghp_..."}'

If API_KEY is set, include -H "X-API-Key: <secret>" on every request.

Docker

docker build -f infrastructure/Dockerfile -t code-review-agent .
docker run -p 8080:8080 -e DATABASE_URL=... code-review-agent

VS Code Extension

cd extension
npm install
npm run compile

Press F5 in VS Code to launch an Extension Development Host, then run "Code Review Agent: Review Current Diff" from the command palette. It sends the active Git diff to the API (configurable via codeReviewAgent.apiUrl, default http://localhost:8080) and reports findings.

Configuration Reference

Environment variables

Variable Required Description
DATABASE_URL Yes Async SQLAlchemy URL, e.g. postgresql+asyncpg://user:pass@host:5432/code_review
MODEL_DIR API only Directory of the exported ONNX model (default outputs/onnx/quantized)
MAX_LENGTH No Max token length for inference (default 512; the demo model uses 192)
LABEL_NAMES No Comma-separated class names (default clean,performance,security,error_handling,style,logic)
RAG_ENABLED No true to retrieve repo context from pgvector when reviewing (needs DATABASE_URL + ingested data)
EMBEDDING_MODEL No Sentence-transformers model for retrieval (default all-MiniLM-L6-v2)
API_KEY No When set, all API routes require a matching X-API-Key header

Training config (ml/train.py)

Key fields: model_name, tokenizer_name, train_data, eval_data, num_labels, epochs, train_batch_size, learning_rate, lora (r, lora_alpha, lora_dropout, target_modules), and wandb_enabled. The model and tokenizer are fully config-driven — no hardcoded defaults — so CodeBERT and StarCoder can be swapped without code changes.

Limitations

This is an MVP. Known gaps before it is production-ready:

  • The model is trained on synthetic data. The bundled demo model learns from templated examples (scripts/generate_demo_data.py), so it classifies demo-style diffs well but will not generalize to arbitrary real-world code. Production accuracy (the PRD target of >89% per-class F1) requires training on a real, human-labeled corpus of PR diffs.
  • RAG grounds the review, but does not yet feed the classifier. Retrieved repo context is returned alongside the prediction and attached to GitHub comments; the ONNX classifier itself still scores the diff alone. Fusing context into the model input is future work.
  • No automatic GitHub webhook trigger yet. The /review/pr endpoint works when called directly; auto-running on pull_request events (with signature verification) is not built.
  • Not yet deployed. A Lambda-ready Dockerfile exists but there is no deployment pipeline, and there are no automated tests or CI.

Contributing

Engineering standards (typing, formatting, async, error handling) and operational directives are documented in CONTRIBUTING.md. A running technical log of changes lives in docs:CHANGELOG.md.

About

AI-augmented code review agent using fine-tuned CodeBERT/StarCoder and RAG to catch issues with precision

Topics

Resources

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages