268 enable configuration of the llm evalutation defence in sandbox

backend/.eslintrc.cjs

@@ -22,7 +22,7 @@ module.exports = {
         checksVoidReturn: false,
       },
     ],
-
+    "@typescript-eslint/unbound-method": ["error", { ignoreStatic: true }],
     "func-style": ["error", "declaration"],
     "prefer-template": "error",
   },

backend/jest.config.js

@@ -3,4 +3,5 @@ module.exports = {
   modulePathIgnorePatterns: ["build", "coverage", "node_modules"],
   preset: "ts-jest",
   testEnvironment: "node",
+  silent: true,
 };

backend/src/defence.ts

@@ -3,7 +3,9 @@ import { ChatDefenceReport } from "./models/chat";
 import { DEFENCE_TYPES, DefenceConfig, DefenceInfo } from "./models/defence";
 import { LEVEL_NAMES } from "./models/level";
 import {
-  retrievalQAPrePromptSecure,
+  maliciousPromptEvalPrePrompt,
+  promptInjectionEvalPrePrompt,
+  qAPrePromptSecure,
   systemRoleDefault,
   systemRoleLevel1,
   systemRoleLevel2,
@@ -24,11 +26,20 @@ function getInitialDefences(): DefenceInfo[] {
         value: process.env.EMAIL_WHITELIST ?? "",
       },
     ]),
-    new DefenceInfo(DEFENCE_TYPES.LLM_EVALUATION, []),
+    new DefenceInfo(DEFENCE_TYPES.EVALUATION_LLM_INSTRUCTIONS, [
+      {
+        id: "prompt-injection-evaluator-prompt",
+        value: promptInjectionEvalPrePrompt,
+      },
+      {
+        id: "malicious-prompt-evaluator-prompt",
+        value: maliciousPromptEvalPrePrompt,
+      },
+    ]),
     new DefenceInfo(DEFENCE_TYPES.QA_LLM_INSTRUCTIONS, [
       {
         id: "prePrompt",
-        value: retrievalQAPrePromptSecure,
+        value: qAPrePromptSecure,
       },
     ]),
     new DefenceInfo(DEFENCE_TYPES.RANDOM_SEQUENCE_ENCLOSURE, [
@@ -177,7 +188,7 @@ function getEmailWhitelistVar(defences: DefenceInfo[]) {
   );
 }
 
-function getQALLMprePrompt(defences: DefenceInfo[]) {
+function getQAPrePromptFromConfig(defences: DefenceInfo[]) {
   return getConfigValue(
     defences,
     DEFENCE_TYPES.QA_LLM_INSTRUCTIONS,
@@ -186,10 +197,26 @@ function getQALLMprePrompt(defences: DefenceInfo[]) {
   );
 }
 
+function getPromptInjectionEvalPrePromptFromConfig(defences: DefenceInfo[]) {
+  return getConfigValue(
+    defences,
+    DEFENCE_TYPES.EVALUATION_LLM_INSTRUCTIONS,
+    "prompt-injection-evaluator-prompt",
+    ""
+  );
+}
+
+function getMaliciousPromptEvalPrePromptFromConfig(defences: DefenceInfo[]) {
+  return getConfigValue(
+    defences,
+    DEFENCE_TYPES.EVALUATION_LLM_INSTRUCTIONS,
+    "malicious-prompt-evaluator-prompt",
+    ""
+  );
+}
+
 function isDefenceActive(id: DEFENCE_TYPES, defences: DefenceInfo[]) {
-  return defences.find((defence) => defence.id === id && defence.isActive)
-    ? true
-    : false;
+  return defences.some((defence) => defence.id === id && defence.isActive);
 }
 
 function generateRandomString(string_length: number) {
@@ -230,16 +257,14 @@ function transformRandomSequenceEnclosure(
   const randomString: string = generateRandomString(
     Number(getRandomSequenceEnclosureLength(defences))
   );
-  const introText: string = getRandomSequenceEnclosurePrePrompt(defences);
-  const transformedMessage: string = introText.concat(
+  return getRandomSequenceEnclosurePrePrompt(defences).concat(
     randomString,
     " {{ ",
     message,
     " }} ",
     randomString,
     ". "
   );
-  return transformedMessage;
 }
 
 // function to escape XML characters in user input to prevent hacking with XML tagging on
@@ -274,12 +299,7 @@ function transformXmlTagging(message: string, defences: DefenceInfo[]) {
   const prePrompt = getXMLTaggingPrePrompt(defences);
   const openTag = "<user_input>";
   const closeTag = "</user_input>";
-  const transformedMessage: string = prePrompt.concat(
-    openTag,
-    escapeXml(message),
-    closeTag
-  );
-  return transformedMessage;
+  return prePrompt.concat(openTag, escapeXml(message), closeTag);
 }
 
 //apply defence string transformations to original message
@@ -370,15 +390,29 @@ async function detectTriggeredDefences(
   }
 
   // evaluate the message for prompt injection
-  const evalPrompt = await queryPromptEvaluationModel(message, openAiApiKey);
+  const configPromptInjectionEvalPrePrompt =
+    getPromptInjectionEvalPrePromptFromConfig(defences);
+  const configMaliciousPromptEvalPrePrompt =
+    getMaliciousPromptEvalPrePromptFromConfig(defences);
+
+  const evalPrompt = await queryPromptEvaluationModel(
+    message,
+    configPromptInjectionEvalPrePrompt,
+    configMaliciousPromptEvalPrePrompt,
+    openAiApiKey
+  );
   if (evalPrompt.isMalicious) {
-    if (isDefenceActive(DEFENCE_TYPES.LLM_EVALUATION, defences)) {
-      defenceReport.triggeredDefences.push(DEFENCE_TYPES.LLM_EVALUATION);
+    if (isDefenceActive(DEFENCE_TYPES.EVALUATION_LLM_INSTRUCTIONS, defences)) {
+      defenceReport.triggeredDefences.push(
+        DEFENCE_TYPES.EVALUATION_LLM_INSTRUCTIONS
+      );
       console.debug("LLM evalutation defence active.");
       defenceReport.isBlocked = true;
       defenceReport.blockedReason = `Message blocked by the malicious prompt evaluator.${evalPrompt.reason}`;
     } else {
-      defenceReport.alertedDefences.push(DEFENCE_TYPES.LLM_EVALUATION);
+      defenceReport.alertedDefences.push(
+        DEFENCE_TYPES.EVALUATION_LLM_INSTRUCTIONS
+      );
     }
   }
   return defenceReport;
@@ -391,7 +425,9 @@ export {
   detectTriggeredDefences,
   getEmailWhitelistVar,
   getInitialDefences,
-  getQALLMprePrompt,
+  getQAPrePromptFromConfig,
+  getPromptInjectionEvalPrePromptFromConfig,
+  getMaliciousPromptEvalPrePromptFromConfig,
   getSystemRole,
   isDefenceActive,
   transformMessage,

backend/src/langchain.ts

@@ -14,10 +14,12 @@ import { CHAT_MODELS, ChatAnswer } from "./models/chat";
 import { DocumentsVector } from "./models/document";
 
 import {
-  maliciousPromptTemplate,
-  promptInjectionEvalTemplate,
-  qAcontextTemplate,
-  retrievalQAPrePrompt,
+  maliciousPromptEvalPrePrompt,
+  maliciousPromptEvalMainPrompt,
+  promptInjectionEvalPrePrompt,
+  promptInjectionEvalMainPrompt,
+  qAMainPrompt,
+  qAPrePrompt,
 } from "./promptTemplates";
 import { LEVEL_NAMES } from "./models/level";
 import { PromptEvaluationChainReply, QaChainReply } from "./models/langchain";
@@ -62,21 +64,26 @@ async function getDocuments(filePath: string) {
     chunkSize: 1000,
     chunkOverlap: 0,
   });
-  const splitDocs = await textSplitter.splitDocuments(docs);
-  return splitDocs;
+
+  return await textSplitter.splitDocuments(docs);
 }
 
-// join the configurable preprompt to the context template
-function getQAPromptTemplate(prePrompt: string) {
-  if (!prePrompt) {
+// choose between the provided preprompt and the default preprompt and prepend it to the main prompt and return the PromptTemplate
+function makePromptTemplate(
+  configPrePrompt: string,
+  defaultPrePrompt: string,
+  mainPrompt: string,
+  templateNameForLogging: string
+) {
+  if (!configPrePrompt) {
     // use the default prePrompt
-    prePrompt = retrievalQAPrePrompt;
+    configPrePrompt = defaultPrePrompt;
   }
-  const fullPrompt = prePrompt + qAcontextTemplate;
-  console.debug(`QA prompt: ${fullPrompt}`);
-  const template: PromptTemplate = PromptTemplate.fromTemplate(fullPrompt);
-  return template;
+  const fullPrompt = `${configPrePrompt}\n${mainPrompt}`;
+  console.debug(`${templateNameForLogging}: ${fullPrompt}`);
+  return PromptTemplate.fromTemplate(fullPrompt);
 }
+
 // create and store the document vectors for each level
 async function initDocumentVectors() {
   const docVectors: DocumentsVector[] = [];
@@ -124,24 +131,36 @@ function initQAModel(
     streaming: true,
     openAIApiKey: openAiApiKey,
   });
-  const promptTemplate = getQAPromptTemplate(prePrompt);
+  const promptTemplate = makePromptTemplate(
+    prePrompt,
+    qAPrePrompt,
+    qAMainPrompt,
+    "QA prompt template"
+  );
 
   return RetrievalQAChain.fromLLM(model, documentVectors.asRetriever(), {
     prompt: promptTemplate,
   });
 }
 
 // initialise the prompt evaluation model
-function initPromptEvaluationModel(openAiApiKey: string) {
+function initPromptEvaluationModel(
+  configPromptInjectionEvalPrePrompt: string,
+  configMaliciousPromptEvalPrePrompt: string,
+  openAiApiKey: string
+) {
   if (!openAiApiKey) {
     console.debug(
       "No OpenAI API key set to initialise prompt evaluation model"
     );
     return;
   }
   // create chain to detect prompt injection
-  const promptInjectionPrompt = PromptTemplate.fromTemplate(
-    promptInjectionEvalTemplate
+  const promptInjectionEvalTemplate = makePromptTemplate(
+    configPromptInjectionEvalPrePrompt,
+    promptInjectionEvalPrePrompt,
+    promptInjectionEvalMainPrompt,
+    "Prompt injection eval prompt template"
   );
 
   const promptInjectionChain = new LLMChain({
@@ -150,21 +169,25 @@ function initPromptEvaluationModel(openAiApiKey: string) {
       temperature: 0,
       openAIApiKey: openAiApiKey,
     }),
-    prompt: promptInjectionPrompt,
+    prompt: promptInjectionEvalTemplate,
     outputKey: "promptInjectionEval",
   });
 
   // create chain to detect malicious prompts
-  const maliciousInputPrompt = PromptTemplate.fromTemplate(
-    maliciousPromptTemplate
+  const maliciousPromptEvalTemplate = makePromptTemplate(
+    configMaliciousPromptEvalPrePrompt,
+    maliciousPromptEvalPrePrompt,
+    maliciousPromptEvalMainPrompt,
+    "Malicious input eval prompt template"
   );
+
   const maliciousInputChain = new LLMChain({
     llm: new OpenAI({
       modelName: CHAT_MODELS.GPT_4,
       temperature: 0,
       openAIApiKey: openAiApiKey,
     }),
-    prompt: maliciousInputPrompt,
+    prompt: maliciousPromptEvalTemplate,
     outputKey: "maliciousInputEval",
   });
 
@@ -209,22 +232,31 @@ async function queryDocuments(
 }
 
 // ask LLM whether the prompt is malicious
-async function queryPromptEvaluationModel(input: string, openAIApiKey: string) {
-  const promptEvaluationChain = initPromptEvaluationModel(openAIApiKey);
+async function queryPromptEvaluationModel(
+  input: string,
+  configPromptInjectionEvalPrePrompt: string,
+  configMaliciousPromptEvalPrePrompt: string,
+  openAIApiKey: string
+) {
+  const promptEvaluationChain = initPromptEvaluationModel(
+    configPromptInjectionEvalPrePrompt,
+    configMaliciousPromptEvalPrePrompt,
+    openAIApiKey
+  );
   if (!promptEvaluationChain) {
     console.debug("Prompt evaluation chain not initialised.");
     return { isMalicious: false, reason: "" };
   }
   console.log(`Checking '${input}' for malicious prompts`);
 
   // get start time
-  const startTime = new Date().getTime();
+  const startTime = Date.now();
   console.debug("Calling prompt evaluation model...");
   const response = (await promptEvaluationChain.call({
     prompt: input,
   })) as PromptEvaluationChainReply;
   // log the time taken
-  const endTime = new Date().getTime();
+  const endTime = Date.now();
   console.debug(`Prompt evaluation model call took ${endTime - startTime}ms`);
 
   const promptInjectionEval = formatEvaluationOutput(
@@ -256,7 +288,7 @@ async function queryPromptEvaluationModel(input: string, openAIApiKey: string) {
 function formatEvaluationOutput(response: string) {
   try {
     // split response on first full stop or comma
-    const splitResponse = response.split(/\.|,/);
+    const splitResponse = response.split(/[.,]/);
     const answer = splitResponse[0]?.replace(/\W/g, "").toLowerCase();
     const reason = splitResponse[1]?.trim();
     return {
@@ -276,12 +308,12 @@ function formatEvaluationOutput(response: string) {
 export {
   initQAModel,
   getFilepath,
-  getQAPromptTemplate,
   getDocuments,
   initPromptEvaluationModel,
   queryDocuments,
   queryPromptEvaluationModel,
   formatEvaluationOutput,
   setVectorisedDocuments,
   initDocumentVectors,
+  makePromptTemplate,
 };

backend/src/models/defence.ts

@@ -1,7 +1,7 @@
 enum DEFENCE_TYPES {
   CHARACTER_LIMIT = "CHARACTER_LIMIT",
   EMAIL_WHITELIST = "EMAIL_WHITELIST",
-  LLM_EVALUATION = "LLM_EVALUATION",
+  EVALUATION_LLM_INSTRUCTIONS = "EVALUATION_LLM_INSTRUCTIONS",
   QA_LLM_INSTRUCTIONS = "QA_LLM_INSTRUCTIONS",
   RANDOM_SEQUENCE_ENCLOSURE = "RANDOM_SEQUENCE_ENCLOSURE",
   SYSTEM_ROLE = "SYSTEM_ROLE",

backend/src/openai.ts

@@ -3,7 +3,7 @@ import {
   getSystemRole,
   detectFilterList,
   getFilterList,
-  getQALLMprePrompt,
+  getQAPrePromptFromConfig,
 } from "./defence";
 import { sendEmail, getEmailWhitelist, isEmailInWhitelist } from "./email";
 import { queryDocuments } from "./langchain";
@@ -231,14 +231,14 @@ async function chatGptCallFunction(
         ) as FunctionAskQuestionParams;
         console.debug(`Asking question: ${params.question}`);
         // if asking a question, call the queryDocuments
-        let qaPrompt = "";
+        let configQAPrePrompt = "";
         if (isDefenceActive(DEFENCE_TYPES.QA_LLM_INSTRUCTIONS, defences)) {
-          qaPrompt = getQALLMprePrompt(defences);
+          configQAPrePrompt = getQAPrePromptFromConfig(defences);
         }
         response = (
           await queryDocuments(
             params.question,
-            qaPrompt,
+            configQAPrePrompt,
             currentLevel,
             openAiApiKey
           )