Merge pull request #11 from coolbeevip/issue-10

coolbeevip · web-flow · commit 05ab1bc464d7 · 2025-02-28T15:41:12.000+08:00
Add custom rule support to summary prompt
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml
@@ -0,0 +1,49 @@
+name: CI
+
+on:
+  push:
+    paths:
+      - "src/**"
+      - "*.py"
+      - "*.toml"
+      - "requirements.txt"
+      - ".github/workflows/**"
+  pull_request:
+    paths:
+      - "src/**"
+      - "*.py"
+      - "*.toml"
+      - "requirements.txt"
+      - ".github/workflows/**"
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check Out Repo
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+          cache: 'pip'
+      - name: Install Dependencies
+        run: make init
+      - name: Run Linter
+        run: make lint
+
+  test:
+    needs: lint
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check Out Repo
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+          cache: 'pip'
+      - name: Install Dependencies
+        run: make init
+      - name: Run Tests
+        run: make test
diff --git a/Makefile b/Makefile
@@ -9,8 +9,8 @@ fmt:
 	@black ./src ./tests ./gitlab_bot.py
 	@isort --profile black ./src ./tests ./gitlab_bot.py
 
-coverage: lint
-	@pytest --cov=an_copilot tests
+test: lint
+	@pytest --cov=src tests
 
 i18n:
 	xgettext -d base -o src/locales/gitlab-bot.pot *.py
diff --git a/README.md b/README.md
@@ -180,10 +180,11 @@ These parameters are used to connect Azure OpenAI models.
 
 This parameter specifies the AI model backend to use. The default value is `openai`, but you can switch to `google` for Google models or `azure-openai` for Azure OpenAI models.
 
-**`BOT_GITLAB_MERGE_REQUEST_SUMMARY_ENABLED` / `BOT_GITLAB_MERGE_REQUEST_SUMMARY_LANGUAGE`**
+**`BOT_GITLAB_MERGE_REQUEST_SUMMARY_ENABLED` / `BOT_GITLAB_MERGE_REQUEST_SUMMARY_LANGUAGE` / `BOT_GITLAB_MERGE_REQUEST_SUMMARY_PROMPT_RULE` **
 
 The `BOT_GITLAB_MERGE_REQUEST_SUMMARY_ENABLED` parameter controls whether the system generates a summary for GitLab merge requests. By default, it is set to `true`.
 The `BOT_GITLAB_MERGE_REQUEST_SUMMARY_LANGUAGE` parameter specifies the language in which the generated merge request summaries will be written. By default, it is set to `English`. However, you can change it to any preferred language code if needed.
+The `BOT_GITLAB_MERGE_REQUEST_SUMMARY_PROMPT_RULE` The parameters specify the prompt rules for generating a merge request summary. You can change it to any preferred prompt rule as needed, for example: `1. Please summarize and condense highly\n2. Please output in English`.
 
 **`BOT_GITLAB_MERGE_REQUEST_EMAIL_USERNAME_NOT_MATCH_ENABLED`**
 
diff --git a/README_ZH.md b/README_ZH.md
@@ -179,10 +179,11 @@ BOT_GIT_COMMIT_SUBJECT_EXAMPLES_MARKDOWN="* feat: 添加认证模块\n* fix: 解
 
 该参数指定要使用的 AI 模型后端。默认值为 `openai`，但您可以切换为 `google`（用于 Google 模型）或 `azure-openai`（用于 Azure OpenAI 模型）。
 
-**`BOT_GITLAB_MERGE_REQUEST_SUMMARY_ENABLED` / `BOT_GITLAB_MERGE_REQUEST_SUMMARY_LANGUAGE`**
+**`BOT_GITLAB_MERGE_REQUEST_SUMMARY_ENABLED` / `BOT_GITLAB_MERGE_REQUEST_SUMMARY_LANGUAGE` / `BOT_GITLAB_MERGE_REQUEST_SUMMARY_PROMPT_RULE`**
 
 `BOT_GITLAB_MERGE_REQUEST_SUMMARY_ENABLED` 参数控制系统是否生成 GitLab 合并请求的摘要。默认设置为 `true`。
 `BOT_GITLAB_MERGE_REQUEST_SUMMARY_LANGUAGE` 参数指定生成的合并请求摘要的语言。默认设置为 `English`。但您可以根据需要更改为任何首选的语言代码。
+`BOT_GITLAB_MERGE_REQUEST_SUMMARY_PROMPT_RULE` 参数指定生成合并请求摘要的提示规则。您可以根据需要更改为任何首选的提示规则，例如：`1. 请高度概括并总结\n2. 请使用中文输出`。
 
 **`BOT_GITLAB_MERGE_REQUEST_EMAIL_USERNAME_NOT_MATCH_ENABLED`**
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -36,6 +36,11 @@ lint = [
     "isort>=5,<6",
     "black>=23,<24"
 ]
+test = [
+    "pytest==7.4.3",
+    "coverage==7.2.7",
+    "pytest-cov==4.1.0"
+]
 package = [
     "flit==3.9.0"
 ]
diff --git a/src/config.py b/src/config.py
@@ -135,6 +135,10 @@
     "BOT_GITLAB_MERGE_REQUEST_SUMMARY_LANGUAGE", "English"
 )
 
+bot_gitlab_merge_request_summary_prompt_rule = os.getenv(
+    "BOT_GITLAB_MERGE_REQUEST_SUMMARY_PROMPT_RULE", None
+)
+
 # gitlab merge request email username not match check
 bot_gitlab_merge_request_email_username_not_match_enabled = (
     os.getenv(
diff --git a/src/llm.py b/src/llm.py
@@ -24,14 +24,15 @@
     azure_openai_api_key,
     azure_openai_api_version,
     azure_openai_endpoint,
+    bot_gitlab_merge_request_summary_prompt_rule,
     bot_language,
     google_api_key,
     model_name,
     openai_api_base,
     openai_api_key,
 )
 from src.i18n import _
-from src.prompts.prompts import get_prompt_text
+from src.prompts.prompts import get_summary_prompt_text
 
 AI = None
 
@@ -77,8 +78,10 @@ def ai_diffs_summary(git_diff) -> str:
     logging.info(f"diffs_string: {diff_string}")
     messages = [
         SystemMessage(
-            content=get_prompt_text(
-                name="DIFFS_SUMMARY_PROMPTS_TEMPLATE", lang=bot_language
+            content=get_summary_prompt_text(
+                name="DIFFS_SUMMARY_PROMPTS_TEMPLATE",
+                lang=bot_language,
+                custom_rule=bot_gitlab_merge_request_summary_prompt_rule,
             )
         ),
         HumanMessage(content=diff_string),
diff --git a/src/prompts/DIFFS_SUMMARY_PROMPTS_TEMPLATE_en.txt b/src/prompts/DIFFS_SUMMARY_PROMPTS_TEMPLATE_en.txt
@@ -1,6 +1,10 @@
 You are an intelligent robot proficient in multiple programming languages and good at reading git-dff reports.
 
+## Rules
+
+<RULES>
 1. Understand and summarize the author’s intention based on the report content.
 2. Then analyze whether there are apparent flaws based on the report content and give suggestions.
+</RULES>
 
 Next, I will send you the complete content of the git-diff report. Please respond according to the rules above after receiving it.
diff --git a/src/prompts/DIFFS_SUMMARY_PROMPTS_TEMPLATE_zh.txt b/src/prompts/DIFFS_SUMMARY_PROMPTS_TEMPLATE_zh.txt
@@ -1,6 +1,10 @@
 你是一个精通多种编程语言的智能机器人，擅长阅读 git-dff 报告。
 
+## 规则
+
+<RULES>
 1. 先根据报告内容理解作者意图并总结
 2. 再根据报告内容分析是否存在明显的缺陷并给出建议
+</RULES>
 
 接下来我会发给你 git-diff 报告完整内容，收到后请按照上边的规则回复。
diff --git a/src/prompts/prompts.py b/src/prompts/prompts.py
@@ -1,8 +1,20 @@
 import os
+import re
 
 
-def get_prompt_text(name: str, lang: str) -> str:
+def get_summary_prompt_text(name: str, lang: str, custom_rule: str = None) -> str:
     current_directory = os.path.dirname(os.path.realpath(__file__))
     path = os.path.join(current_directory, f"{name}_{lang}.txt")
     with open(path, "r", encoding="UTF-8") as file:
-        return file.read()
+        summary_prompt = file.read()
+        if custom_rule:
+            # replace summary_prompt_rule with the actual rule
+            summary_prompt = re.sub(
+                r"<RULES>.*?</RULES>", f"{custom_rule}", summary_prompt, flags=re.DOTALL
+            )
+        else:
+            # Remove the tags
+            summary_prompt = summary_prompt.replace("<RULES>\n", "").replace(
+                "</RULES>\n", ""
+            )
+        return summary_prompt
diff --git a/tests/test_llm.py b/tests/test_llm.py
@@ -1,4 +1,5 @@
 import json
+import os
 import unittest
 from unittest import TestCase
 
@@ -7,11 +8,14 @@
 
 class TestLLM(TestCase):
     def test_ai_diffs_summary(self):
-        with open("git_diff.json", "r") as f:
+        test_dir = os.path.dirname(os.path.abspath(__file__))
+        file_path = os.path.join(test_dir, "git_diff.json")
+
+        with open(file_path, "r") as f:
             git_diff_json = json.loads(f.read())
             summary_report = ai_diffs_summary(git_diff_json)
-            self.assertIn("AI 摘要(实验)", summary_report)
-            self.assertIn("AI 模型", summary_report)
+            self.assertIn("AI Summary(experiment)", summary_report)
+            self.assertIn("Powered by", summary_report)
             print(summary_report)
 
 
diff --git a/tests/test_prompts.py b/tests/test_prompts.py
@@ -1,18 +1,45 @@
 import unittest
 from unittest import TestCase
 
-from src.prompts.prompts import get_prompt_text
+from src.prompts.prompts import get_summary_prompt_text
 
 
 class TestPrompts(TestCase):
-    def test_load_prompt(self):
+    def test_load_default_prompt(self):
         self.assertIn(
             "你是一个精通多种编程语言的智能机器人",
-            get_prompt_text("DIFFS_SUMMARY_PROMPTS_TEMPLATE", "zh"),
+            get_summary_prompt_text("DIFFS_SUMMARY_PROMPTS_TEMPLATE", "zh"),
         )
         self.assertIn(
             "You are an intelligent robot proficient in multiple programming languages and good at reading git-dff reports",
-            get_prompt_text("DIFFS_SUMMARY_PROMPTS_TEMPLATE", "en"),
+            get_summary_prompt_text("DIFFS_SUMMARY_PROMPTS_TEMPLATE", "en"),
+        )
+        self.assertNotIn(
+            "<RULES>", get_summary_prompt_text("DIFFS_SUMMARY_PROMPTS_TEMPLATE", "en")
+        )
+        self.assertNotIn(
+            "</RULES>", get_summary_prompt_text("DIFFS_SUMMARY_PROMPTS_TEMPLATE", "zh")
+        )
+
+    def test_load_custom_prompt(self):
+        custom_rules = "1. RULE1\n2. RULE2"
+        self.assertIn(
+            "1. RULE1\n2. RULE2",
+            get_summary_prompt_text(
+                "DIFFS_SUMMARY_PROMPTS_TEMPLATE", "zh", custom_rules
+            ),
+        )
+        self.assertIn(
+            "1. RULE1\n2. RULE2",
+            get_summary_prompt_text(
+                "DIFFS_SUMMARY_PROMPTS_TEMPLATE", "en", custom_rules
+            ),
+        )
+        self.assertNotIn(
+            "<RULES>", get_summary_prompt_text("DIFFS_SUMMARY_PROMPTS_TEMPLATE", "en")
+        )
+        self.assertNotIn(
+            "</RULES>", get_summary_prompt_text("DIFFS_SUMMARY_PROMPTS_TEMPLATE", "zh")
         )
 
 

Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,11 @@ lint = [`
`36`	`36`	`"isort>=5,<6",`
`37`	`37`	`"black>=23,<24"`
`38`	`38`	`]`
	`39`	`+test = [`
	`40`	`+ "pytest==7.4.3",`
	`41`	`+ "coverage==7.2.7",`
	`42`	`+ "pytest-cov==4.1.0"`
	`43`	`+]`
`39`	`44`	`package = [`
`40`	`45`	`"flit==3.9.0"`
`41`	`46`	`]`