Jeremy Janella
119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
# requires openai/anthropic to run; didnt implement those. the script is just for guardrail review so that isnt nessecary
|
|
|
|
from pydantic import BaseModel, StringConstraints
|
|
from typing import Annotated
|
|
|
|
class LLMQuery(BaseModel, validate_assignment=True):
|
|
prompt: Annotated[str, StringConstraints(pattern=r"^[0-9.\+\-\*/\(\)]+$")]
|
|
response: Annotated[str, StringConstraints(pattern=r"^[0-9\.]+$")] = None
|
|
|
|
SYSTEM_PROMPT = '''You are a calculator. Please compute the result of the following mathematical expression.
|
|
Only respond with the result, no other text.
|
|
|
|
'''
|
|
|
|
import json
|
|
|
|
PRODUCTS = [
|
|
'SnaccPayload',
|
|
'PowerChex',
|
|
'CrunchPacket',
|
|
'ByteChew',
|
|
'CryptoChunks'
|
|
]
|
|
|
|
# https://gist.githubusercontent.com/mspublic/f942cb80a83a7880cff689def8334c7b/raw/30527efde9e1d2b11ede9d9bf205d98c2ba9d550/Instruction-Bypass.yara
|
|
PROMPT_INJECTION_REGEX = r"^.*(Ignore|Disregard|Skip|Forget|Neglect|Overlook|Omit|Bypass)\s*(prior|previous|preceding|above|foregoing|earlier|initial)?\s*(content|text|instructions|instruction|directives|directive|commands|command).*$"
|
|
|
|
# https://gist.githubusercontent.com/mspublic/f942cb80a83a7880cff689def8334c7b/raw/30527efde9e1d2b11ede9d9bf205d98c2ba9d550/jailbreaks.json
|
|
JAILBREAKS = json.load(open('./jailbreaks.json')).get('jailbreak')
|
|
|
|
# check if any item in the passed list is contained in the passed string
|
|
def list_check(list, string):
|
|
return any(item in string for item in list)
|
|
|
|
# detect jailbreak based on string similarity
|
|
from difflib import SequenceMatcher
|
|
def jailbreak_detect(prompt, jailbreak_threshold=0.5):
|
|
return any(SequenceMatcher(None, jailbreak, prompt).ratio() > jailbreak_threshold for jailbreak in JAILBREAKS)
|
|
|
|
class GuardrailPromptException(Exception):
|
|
pass
|
|
|
|
class GuardrailResponseException(Exception):
|
|
pass
|
|
|
|
|
|
from pydantic import BaseModel, field_validator
|
|
import re
|
|
|
|
class LLMQuery(BaseModel, validate_assignment=True):
|
|
prompt: str
|
|
response: str = None
|
|
|
|
@field_validator("prompt")
|
|
@classmethod
|
|
def validate_prompt(cls, prompt: str) -> str:
|
|
prompt = prompt.strip()
|
|
|
|
# Check for PacketSnacc product
|
|
if not list_check(PRODUCTS, prompt):
|
|
raise GuardrailPromptException("No PacketSnacc product mentioned in user prompt.")
|
|
|
|
# Check for prompt injection
|
|
if re.search(PROMPT_INJECTION_REGEX, prompt, re.IGNORECASE):
|
|
raise GuardrailPromptException("Prompt injection attempt detected.")
|
|
|
|
# Check for jailbreaking
|
|
if jailbreak_detect(prompt):
|
|
raise GuardrailPromptException("Jailbreak attempt detected.")
|
|
|
|
return prompt
|
|
|
|
@field_validator("response")
|
|
@classmethod
|
|
def validate_response(cls, response: str) -> str:
|
|
# Check for competitors
|
|
if list_check(COMPETITORS, response):
|
|
raise GuardrailResponseException("Output does not satisfy company policy.")
|
|
|
|
# Check for profanity
|
|
if list_check(PROFANITY, response):
|
|
raise GuardrailResponseException("Profane language detected.")
|
|
|
|
# Check for credit card information
|
|
if re.search(r"^.*[0-9]{13,19}.*$", response):
|
|
raise GuardrailResponseException("Information leakage detected.")
|
|
|
|
# Remove HTML-Tags
|
|
response = re.sub(r"<.*?>", "", response)
|
|
|
|
return response
|
|
|
|
COMPETITORS = [
|
|
"SnackOverflow",
|
|
"NullBite",
|
|
"CyberChow"
|
|
]
|
|
|
|
# https://raw.githubusercontent.com/zacanger/profane-words/refs/heads/master/words.json
|
|
PROFANITY = json.load(open('./words.json'))
|
|
|
|
def query_llm(system_prompt:str, prompt: str) -> str:
|
|
# openai/anthorpic etc; not relevent
|
|
return response
|
|
|
|
def protected_query_llm(prompt: str) -> LLMQuery:
|
|
query = LLMQuery(prompt=prompt)
|
|
query.response = query_llm(SYSTEM_PROMPT, query.prompt)
|
|
return query
|
|
|
|
|
|
while 1:
|
|
try:
|
|
prompt = input("> ")
|
|
query_obj = protected_query_llm(prompt)
|
|
print(query_obj.response)
|
|
except Exception as e:
|
|
print(f'Error: {e}')
|