diff --git a/README.md b/README.md index c0bba3a..7443d4d 100644 --- a/README.md +++ b/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 diff --git a/saist.rules b/saist.rules new file mode 100644 index 0000000..3e940f5 --- /dev/null +++ b/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." diff --git a/saist/main.py b/saist/main.py index 947393c..47dc488 100644 --- a/saist/main.py +++ b/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: diff --git a/saist/util/rules.py b/saist/util/rules.py new file mode 100644 index 0000000..0b4c607 --- /dev/null +++ b/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! +