Modify prompt using saist.rules

README.md

@@ -155,6 +155,36 @@ This setup will:
 
 ---
 
+## 🧠 Custom Prompt Rules 
+
+You can modify the system prompt used by creating a file named `saist.rules` in your working directory.
+
+This file supports the following YAML keys:
+
+PROMPT_OVERRIDE - Fully replaces the system prompt 
+PROMPT_PRE - Adds text before the existing prompt 
+PROMPT_POST - Adds text after the existing prompt 
+
+If multiple keys are defined, `PRE` and `POST` will wrap around the override text.
+
+---
+
+### 📝 `saist.rules.example`
+
+
+```yaml
+# Example prompt modification
+
+# Use this to completely override the system prompt:
+PROMPT_OVERRIDE: "Identify any potential vulnerabilites in the following code. Focus on input validation"
+
+# Use these to prepend or append context to the prompt
+PROMPT_PRE: "You are reviewing critical code. Please be strict.\n\n"
+PROMPT_POST: "\n\nPlease provide practical advice."
+```
+
+---
+
 ## 🛣️ Roadmap
 
 - Ability to influence the prompts

saist.rules

@@ -0,0 +1,8 @@
+# Example prompt modification
+
+# Use this to completely override the system prompt:
+# PROMPT_OVERRIDE: "Identify any potential vulnerabilites in the following code. Focus on input validation"
+
+# Use these to prepend or append context to the prompt
+PROMPT_PRE: "You are reviewing critical code. Please be strict.\n\n"
+PROMPT_POST: "\n\nPlease provide practical advice."

saist/main.py

@@ -31,6 +31,8 @@
 
 from util.output import print_banner, write_csv
 
+from util.rules import PromptRules
+
 prompts = prompts()
 load_dotenv(".env")
 
@@ -40,11 +42,11 @@ async def analyze_single_file(scm: Scm, adapter: BaseLlmAdapter, filename, patch
     """
     Analyzes a SINGLE file diff with OpenAI, returning a Findings object or None on error.
     """
-    system_prompt = prompts.DETECT
+    system_prompt = PromptRules.apply_rules(prompts.DETECT)
     logger.debug(f"Processing {filename}")
-    prompt = (
+    prompt =(
         f"\n\nFile: {filename}\n{patch_text}\n"
-    )
+        )
     try:
         return (await adapter.prompt_structured(system_prompt, prompt, Findings, [scm.read_file_contents])).findings
     except Exception as e:

saist/util/rules.py

@@ -0,0 +1,49 @@
+import os
+import yaml
+import logging
+
+logger = logging.getLogger(__name__) 
+
+class PromptRules:
+    RulesFile = "saist.rules"
+
+    @staticmethod
+    def apply_rules(prompt: str) -> str:
+        rules = PromptRules.load_rules()
+
+        override = rules.get("PROMPT_OVERRIDE")
+        pre = rules.get("PROMPT_PRE", "")
+        post = rules.get("PROMPT_POST", "")
+
+        if override:
+            final_prompt = f"{pre}{override}{post}"
+        else:
+            final_prompt = f"{pre}{prompt}{post}"
+
+        return final_prompt
+
+    @staticmethod
+    def load_rules():
+        if not os.path.exists(PromptRules.RulesFile):
+            logger.warning("No saist.rules file found.")
+            return {} #return empty rules
+
+        try:
+            with open(PromptRules.RulesFile, 'r') as file:
+                yaml_content = file.read()
+                rules = yaml.safe_load(yaml_content) #using safe_load to prevent exploits (thank you stack overflow)
+
+                keys = [key for key in ["PROMPT_OVERRIDE", "PROMPT_PRE", "PROMPT_POST"] if key in rules]
+            if keys:
+                logger.debug(f"Loaded prompt rules: {', '.join(keys)}")
+            else:
+                logger.debug("No valid keys found.")
+
+            return rules if rules is not None else {}
+        except Exception as ex:
+            logger.error(f"Error reading saist.rules: {ex}")
+            return {}
+
+#I first wrote this code in c# using dictionaries and a similar yaml parsing library for .net and then used a converter for python. 
+#I'm not too well versed with python yet but doing this excersise and seeing how things are converted, learning some syntax and hacking stuff together has been a massive help and has been really enjoyable!
+