From fef4ab06a0f665f877ad6a2f905bd175505e1096 Mon Sep 17 00:00:00 2001
From: Anurag Maherchandani <anurag.ait@gmail.com>
Date: Wed, 19 Mar 2025 16:43:13 +0530
Subject: [PATCH] Add custom review rules feature to AI Code Reviewer

---
 README.md   | 14 ++++++++++++++
 action.yml  |  4 ++++
 src/main.ts | 23 ++++++++++++++++++-----
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 566a3ea..6000736 100644
--- a/README.md
+++ b/README.md
@@ -9,6 +9,7 @@ review process.
 - Reviews pull requests using OpenAI's GPT-4 API.
 - Provides intelligent comments and suggestions for improving your code.
 - Filters out files that match specified exclude patterns.
+- Supports custom review rules to tailor the AI's feedback to your project's needs.
 - Easy to set up and integrate into your GitHub workflow.
 
 ## Setup
@@ -44,6 +45,7 @@ jobs:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           OPENAI_API_MODEL: "gpt-4" # Optional: defaults to "gpt-4"
           exclude: "**/*.json, **/*.md" # Optional: exclude patterns separated by commas
+          custom_rules: "Check for proper error handling, Ensure functions have proper docstrings, Verify variable names follow team conventions" # Optional: custom review rules separated by commas
 ```
 
 4. Replace `your-username` with your GitHub username or organization name where the AI Code Reviewer repository is
@@ -58,6 +60,18 @@ jobs:
 The AI Code Reviewer GitHub Action retrieves the pull request diff, filters out excluded files, and sends code chunks to
 the OpenAI API. It then generates review comments based on the AI's response and adds them to the pull request.
 
+### Custom Rules
+
+You can provide custom review rules to tailor the AI's feedback to your project's specific needs. These rules are added to the prompt sent to the OpenAI API.
+
+For example, if your team has specific coding conventions or you want to focus on particular aspects of code quality, you can specify them as custom rules:
+
+```yaml
+custom_rules: "Check for proper error handling, Make sure unit tests are included, Ensure all functions are properly documented"
+```
+
+Each rule should be a clear, concise instruction that guides the AI's review process.
+
 ## Contributing
 
 Contributions are welcome! Please feel free to submit issues or pull requests to improve the AI Code Reviewer GitHub
diff --git a/action.yml b/action.yml
index 976aed7..86e1523 100644
--- a/action.yml
+++ b/action.yml
@@ -15,6 +15,10 @@ inputs:
     description: "Glob patterns to exclude files from the diff analysis"
     required: false
     default: ""
+  custom_rules:
+    description: "Custom review rules to be added to the prompt, separated by commas"
+    required: false
+    default: ""
 runs:
   using: "node16"
   main: "dist/index.js"
diff --git a/src/main.ts b/src/main.ts
index 13677ae..2fdb3eb 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -5,9 +5,13 @@ import { Octokit } from "@octokit/rest";
 import parseDiff, { Chunk, File } from "parse-diff";
 import minimatch from "minimatch";
 
-const GITHUB_TOKEN: string = core.getInput("GITHUB_TOKEN");
-const OPENAI_API_KEY: string = core.getInput("OPENAI_API_KEY");
-const OPENAI_API_MODEL: string = core.getInput("OPENAI_API_MODEL");
+const GITHUB_TOKEN: string = core.getInput("GITHUB_TOKEN", { required: true });
+const OPENAI_API_KEY: string = core.getInput("OPENAI_API_KEY", { required: true });
+const OPENAI_API_MODEL: string = core.getInput("OPENAI_API_MODEL", { required: true });
+const CUSTOM_RULES: string[] = core.getInput("custom_rules")
+  .split(",")
+  .map((rule: string) => rule.trim())
+  .filter((rule: string) => rule !== "");
 
 const octokit = new Octokit({ auth: GITHUB_TOKEN });
 
@@ -79,13 +83,22 @@ async function analyzeCode(
 }
 
 function createPrompt(file: File, chunk: Chunk, prDetails: PRDetails): string {
-  return `Your task is to review pull requests. Instructions:
+  let promptInstructions = `Your task is to review pull requests. Instructions:
 - Provide the response in following JSON format:  {"reviews": [{"lineNumber":  <line_number>, "reviewComment": "<review comment>"}]}
 - Do not give positive comments or compliments.
 - Provide comments and suggestions ONLY if there is something to improve, otherwise "reviews" should be an empty array.
 - Write the comment in GitHub Markdown format.
 - Use the given description only for the overall context and only comment the code.
-- IMPORTANT: NEVER suggest adding comments to the code.
+- IMPORTANT: NEVER suggest adding comments to the code.`;
+
+  if (CUSTOM_RULES.length > 0) {
+    promptInstructions += "\n\nAdditional custom rules:";
+    CUSTOM_RULES.forEach((rule: string) => {
+      promptInstructions += `\n- ${rule}`;
+    });
+  }
+
+  return `${promptInstructions}
 
 Review the following code diff in the file "${
     file.to