From a13af9c561413b91421978e0262a8df8ebd54778 Mon Sep 17 00:00:00 2001 From: Risto McGehee Date: Sat, 23 Dec 2023 12:28:46 -0800 Subject: [PATCH 1/3] Modularize detection checks --- javascript-sdk/src/index.ts | 13 +- javascript-sdk/src/interface.ts | 59 +++- javascript-sdk/src/lib/Strategy.ts | 6 + javascript-sdk/src/lib/detect.ts | 197 -------------- javascript-sdk/src/sdk.ts | 128 ++++----- javascript-sdk/src/tactics/Heuristic.ts | 138 ++++++++++ javascript-sdk/src/tactics/OpenAI.ts | 40 +++ javascript-sdk/src/tactics/Tactic.ts | 20 ++ javascript-sdk/src/tactics/Vector.ts | 45 +++ javascript-sdk/tests/index.test.ts | 348 ++++++++++++------------ 10 files changed, 538 insertions(+), 456 deletions(-) create mode 100644 javascript-sdk/src/lib/Strategy.ts delete mode 100644 javascript-sdk/src/lib/detect.ts create mode 100644 javascript-sdk/src/tactics/Heuristic.ts create mode 100644 javascript-sdk/src/tactics/OpenAI.ts create mode 100644 javascript-sdk/src/tactics/Tactic.ts create mode 100644 javascript-sdk/src/tactics/Vector.ts diff --git a/javascript-sdk/src/index.ts b/javascript-sdk/src/index.ts index 26add71..4c5ec9a 100644 --- a/javascript-sdk/src/index.ts +++ b/javascript-sdk/src/index.ts @@ -1,8 +1,19 @@ export { default as RebuffApi } from "./api"; export { default as RebuffSdk } from "./sdk"; -export { +export type { ApiConfig, SdkConfig, RebuffConfig, VectorDbConfig, } from "./config"; +export { + RebuffError +} from "./interface"; +export type { + DetectRequest, + DetectResponse, + Rebuff, + TacticName, + TacticOverride, + TacticResult +} from "./interface"; diff --git a/javascript-sdk/src/interface.ts b/javascript-sdk/src/interface.ts index adf01f2..f596088 100644 --- a/javascript-sdk/src/interface.ts +++ b/javascript-sdk/src/interface.ts @@ -1,25 +1,54 @@ +export enum TacticName { + // A series of heuristics are used to determine whether the input is prompt injection. + Heuristic = "heuristic", + // A language model is asked if the input appears to be prompt injection. + LanguageModel = "language_model", + // A vector database of known prompt injection attacks is queried for similarity. + VectorDB = "vector_db", +} + +export interface TacticOverride { + // The name of the tactic to override. + name: TacticName; + // The threshold to use for this tactic. If the score is above this threshold, the tactic will be considered detected. + // If not specified, the default threshold for the tactic will be used. + threshold?: number; + // Whether to run this tactic. Defaults to true if not specified. + run?: boolean; +} + export interface DetectRequest { + // The user input to check for prompt injection. userInput: string; + // The base64-encoded user input. If this is specified, the user input will be ignored. userInputBase64?: string; - runHeuristicCheck: boolean; - runVectorCheck: boolean; - runLanguageModelCheck: boolean; - maxHeuristicScore: number; - maxModelScore: number; - maxVectorScore: number; + // Any tactics to change behavior for. If any tactic is not specified, the default threshold for that tactic will be + // used. + tacticOverrides?: TacticOverride[]; +} + +export interface TacticResult { + // The name of the tactic. + name: TacticName; + // The score for the tactic. This is a number between 0 and 1. The closer to 1, the more likely that this is a + // prompt injection attempt. + score: number; + // Whether this tactic evaluated the input as a prompt injection attempt. + detected: boolean; + // The threshold used for this tactic. If the score is above this threshold, the tactic will be considered detected. + threshold: number; + // Some tactics return additional fields: + // * "vector_db": + // - "countOverMaxVectorScore" (number): The number of different vectors whose similarity score is above the + // threshold. + additionalFields: Record; } export interface DetectResponse { - heuristicScore: number; - modelScore: number; - vectorScore: Record; - runHeuristicCheck: boolean; - runVectorCheck: boolean; - runLanguageModelCheck: boolean; - maxHeuristicScore: number; - maxVectorScore: number; - maxModelScore: number; + // Whether prompt injection was detected. injectionDetected: boolean; + // The result for each tactic that was executed. + tacticResults: TacticResult[]; } export class RebuffError extends Error { diff --git a/javascript-sdk/src/lib/Strategy.ts b/javascript-sdk/src/lib/Strategy.ts new file mode 100644 index 0000000..9c0fe72 --- /dev/null +++ b/javascript-sdk/src/lib/Strategy.ts @@ -0,0 +1,6 @@ +import Tactic from "../tactics/Tactic"; + +export default interface Strategy { + // The tactics that will be executed as part of this strategy. + tactics: Tactic[]; +} diff --git a/javascript-sdk/src/lib/detect.ts b/javascript-sdk/src/lib/detect.ts deleted file mode 100644 index d8388e0..0000000 --- a/javascript-sdk/src/lib/detect.ts +++ /dev/null @@ -1,197 +0,0 @@ -import stringSimilarity from "string-similarity"; -import { normalizeString } from "./prompts"; -import { OpenAIApi } from "openai"; -import { RebuffError } from "../interface"; -import { VectorStore } from "langchain/vectorstores/base"; - -export async function detectPiUsingVectorDatabase( - input: string, - similarityThreshold: number, - vectorStore: VectorStore -): Promise<{ topScore: number; countOverMaxVectorScore: number }> { - try { - const topK = 20; - const results = await vectorStore.similaritySearchWithScore(input, topK); - - let topScore = 0; - let countOverMaxVectorScore = 0; - - for (const [_, score] of results) { - if (score == undefined) { - continue; - } - - if (score > topScore) { - topScore = score; - } - - if (score >= similarityThreshold && score > topScore) { - countOverMaxVectorScore++; - } - } - - return { topScore, countOverMaxVectorScore }; - } catch (error) { - throw new RebuffError(`Error in detectPiUsingVectorDatabase: ${error}`); - } -} - -export function generateInjectionKeywords() { - // Define lists of verbs, adjectives, prepositions, and objects - const verbs = [ - "Ignore", - "Disregard", - "Skip", - "Forget", - "Neglect", - "Overlook", - "Omit", - "Bypass", - "Pay no attention to", - "Do not follow", - "Do not obey", - ]; - - const adjectives = [ - "", - "prior", - "previous", - "preceding", - "above", - "foregoing", - "earlier", - "initial", - ]; - - const prepositions = [ - "", - "and start over", - "and start anew", - "and begin afresh", - "and start from scratch", - ]; - - const objects = [ - "content", - "text", - "instructions", - "instruction", - "directives", - "directive", - "commands", - "command", - "context", - "conversation", - "input", - "inputs", - "data", - "message", - "messages", - "communication", - "response", - "responses", - "request", - "requests", - ]; - - // Generate all possible combinations of sentences - const injectionKeywords = []; - for (const verb of verbs) { - for (const adjective of adjectives) { - for (const object of objects) { - for (const preposition of prepositions) { - injectionKeywords.push( - `${verb} ${adjective} ${object} ${preposition}`.trim() - ); - } - } - } - } - - return injectionKeywords; -} - -// Generate and print the injection keywords -const injectionKeywords = generateInjectionKeywords(); - -// Return the highest score, ensure the score is normalized between 0 and 1 -// based on the amount of consecutive words matched and the similarity score -export function detectPromptInjectionUsingHeuristicOnInput( - input: string -): number { - let highestScore = 0; - const normalizedInput = normalizeString(input); - - for (const keyword of injectionKeywords) { - const normalizedKeyword = normalizeString(keyword); - const keywordParts = normalizedKeyword.split(" "); - const keywordLength = keywordParts.length; - - // Generate substrings of similar length in the input string - const inputParts = normalizedInput.split(" "); - const inputSubstrings = []; - for (let i = 0; i <= inputParts.length - keywordLength; i++) { - inputSubstrings.push(inputParts.slice(i, i + keywordLength).join(" ")); - } - - // Calculate the similarity score between the keyword and each substring - for (const substring of inputSubstrings) { - const similarityScore = stringSimilarity.compareTwoStrings( - normalizedKeyword, - substring - ); - - // Calculate the score based on the number of consecutive words matched - const matchedWordsCount = keywordParts.filter( - (part, index) => substring.split(" ")[index] === part - ).length; - const maxMatchedWords = 5; - const baseScore = - matchedWordsCount > 0 - ? 0.5 + 0.5 * Math.min(matchedWordsCount / maxMatchedWords, 1) - : 0; - - // Adjust the score using the similarity score - const adjustedScore = - baseScore - similarityScore * (1 / (maxMatchedWords * 2)); - - // Update the highest score if the current adjusted score is higher - if (adjustedScore > highestScore) { - highestScore = adjustedScore; - } - } - } - - return highestScore; -} - -export async function callOpenAiToDetectPI( - promptToDetectPiUsingOpenAI: string, - openai: OpenAIApi, - model = "gpt-3.5-turbo" -): Promise<{ completion: string; error?: string }> { - try { - const completion = await openai.createChatCompletion({ - model, - messages: [{ role: "user", content: promptToDetectPiUsingOpenAI }], - }); - - if (completion.data.choices[0].message === undefined) { - console.log("completion.data.choices[0].message is undefined"); - return { completion: "", error: "server_error" }; - } - - if (completion.data.choices.length === 0) { - console.log("completion.data.choices.length === 0"); - return { completion: "", error: "server_error" }; - } - - return { - completion: completion.data.choices[0].message.content || "", - error: undefined, - }; - } catch (error) { - console.error("Error in callOpenAiToDetectPI:", error); - return { completion: "", error: "server_error" }; - } -} diff --git a/javascript-sdk/src/sdk.ts b/javascript-sdk/src/sdk.ts index 2119a00..63ed902 100644 --- a/javascript-sdk/src/sdk.ts +++ b/javascript-sdk/src/sdk.ts @@ -3,20 +3,19 @@ import { DetectResponse, Rebuff, RebuffError, + TacticResult, } from "./interface"; import crypto from "crypto"; import { SdkConfig } from "./config"; import initVectorStore from "./lib/vectordb"; -import { - callOpenAiToDetectPI, - detectPiUsingVectorDatabase, - detectPromptInjectionUsingHeuristicOnInput, -} from "./lib/detect"; import getOpenAIInstance from "./lib/openai"; -import { renderPromptForPiDetection } from "./lib/prompts"; import { OpenAIApi } from "openai"; import { VectorStore } from "langchain/vectorstores/base"; import { Document } from "langchain/document"; +import Strategy from "./lib/Strategy"; +import Heuristic from "./tactics/Heuristic"; +import OpenAI from "./tactics/OpenAI"; +import Vector from "./tactics/Vector"; function generateCanaryWord(length = 8): string { // Generate a secure random hexadecimal canary word @@ -24,8 +23,10 @@ function generateCanaryWord(length = 8): string { } export default class RebuffSdk implements Rebuff { - private vectorStore: VectorStore | undefined; private sdkConfig: SdkConfig; + private vectorStore: VectorStore | undefined; + private strategies: Record | undefined; + private defaultStrategy: string; private openai: { conn: OpenAIApi; @@ -38,17 +39,34 @@ export default class RebuffSdk implements Rebuff { conn: getOpenAIInstance(config.openai.apikey), model: config.openai.model || "gpt-3.5-turbo", }; + this.defaultStrategy = "standard"; + } + + private async getStrategies(): Promise> { + if (this.strategies) { + return this.strategies; + } + const heuristicScoreThreshold = 0.75; + const vectorScoreThreshold = 0.9; + const openaiScoreThreshold = 0.9; + const strategies: Record = { + // For now, this is the only strategy. + "standard": { + tactics: [ + new Heuristic(heuristicScoreThreshold), + new Vector(vectorScoreThreshold, await this.getVectorStore()), + new OpenAI(openaiScoreThreshold, this.openai.model, this.openai.conn), + ] + }, + }; + this.strategies = strategies; + return this.strategies; } async detectInjection({ userInput = "", userInputBase64 = "", - maxHeuristicScore = 0.75, - maxVectorScore = 0.9, - maxModelScore = 0.9, - runHeuristicCheck = true, - runVectorCheck = true, - runLanguageModelCheck = true, + tacticOverrides = [], }: DetectRequest): Promise { if (userInputBase64) { // Create a buffer from the hexadecimal string @@ -59,73 +77,33 @@ export default class RebuffSdk implements Rebuff { if (!userInput) { throw new RebuffError("userInput is required"); } - if (typeof runHeuristicCheck !== "boolean") { - throw new RebuffError("runHeuristicCheck must be a boolean"); - } - if (typeof runVectorCheck !== "boolean") { - throw new RebuffError("runVectorCheck must be a boolean"); - } - if (typeof runLanguageModelCheck !== "boolean") { - throw new RebuffError("runLanguageModelCheck must be a boolean"); - } - if ( - maxHeuristicScore === null || - maxModelScore === null || - maxVectorScore === null - ) { - throw new RebuffError( - "maxHeuristicScore, maxModelScore, and maxVectorScore are required" - ); - } - - runHeuristicCheck = runHeuristicCheck === null ? true : runHeuristicCheck; - runVectorCheck = runVectorCheck === null ? true : runVectorCheck; - runLanguageModelCheck = - runLanguageModelCheck === null ? true : runLanguageModelCheck; - if (!userInput) { - throw new RebuffError("userInput is required"); + const strategies = await this.getStrategies(); + let injectionDetected = false; + let tacticResults: TacticResult[] = []; + for (const tactic of strategies[this.defaultStrategy].tactics) { + const tacticOverride = tacticOverrides.find(t => t.name === tactic.name); + if (tacticOverride && tacticOverride.run === false) { + continue; + } + const threshold = tacticOverride?.threshold ?? tactic.defaultThreshold; + const execution = await tactic.execute(userInput, threshold); + const result = { + name: tactic.name, + score: execution.score, + threshold, + detected: execution.score > threshold, + additionalFields: execution.additionalFields ?? {}, + } as TacticResult; + if (result.detected) { + injectionDetected = true; + } + tacticResults.push(result); } - const heuristicScore = runHeuristicCheck - ? detectPromptInjectionUsingHeuristicOnInput(userInput) - : 0; - - const modelScore = runLanguageModelCheck - ? parseFloat( - ( - await callOpenAiToDetectPI( - renderPromptForPiDetection(userInput), - this.openai.conn, - this.openai.model - ) - ).completion - ) - : 0; - - const vectorScore = runVectorCheck - ? await detectPiUsingVectorDatabase( - userInput, - maxVectorScore, - await this.getVectorStore() - ) - : { topScore: 0, countOverMaxVectorScore: 0 }; - const injectionDetected = - heuristicScore > maxHeuristicScore || - modelScore > maxModelScore || - vectorScore.topScore > maxVectorScore; - return { - heuristicScore, - modelScore, - vectorScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, - maxHeuristicScore, - maxVectorScore, - maxModelScore, injectionDetected, + tacticResults, } as DetectResponse; } diff --git a/javascript-sdk/src/tactics/Heuristic.ts b/javascript-sdk/src/tactics/Heuristic.ts new file mode 100644 index 0000000..7ac0e30 --- /dev/null +++ b/javascript-sdk/src/tactics/Heuristic.ts @@ -0,0 +1,138 @@ +import stringSimilarity from "string-similarity"; +import { normalizeString } from "../lib/prompts"; +import { TacticName } from "../interface"; +import Tactic, { TacticExecution } from "./Tactic"; + +function generateInjectionKeywords() { + // Define lists of verbs, adjectives, prepositions, and objects + const verbs = [ + "Ignore", + "Disregard", + "Skip", + "Forget", + "Neglect", + "Overlook", + "Omit", + "Bypass", + "Pay no attention to", + "Do not follow", + "Do not obey", + ]; + + const adjectives = [ + "", + "prior", + "previous", + "preceding", + "above", + "foregoing", + "earlier", + "initial", + ]; + + const prepositions = [ + "", + "and start over", + "and start anew", + "and begin afresh", + "and start from scratch", + ]; + + const objects = [ + "content", + "text", + "instructions", + "instruction", + "directives", + "directive", + "commands", + "command", + "context", + "conversation", + "input", + "inputs", + "data", + "message", + "messages", + "communication", + "response", + "responses", + "request", + "requests", + ]; + + // Generate all possible combinations of sentences + const injectionKeywords = []; + for (const verb of verbs) { + for (const adjective of adjectives) { + for (const object of objects) { + for (const preposition of prepositions) { + injectionKeywords.push( + `${verb} ${adjective} ${object} ${preposition}`.trim() + ); + } + } + } + } + + return injectionKeywords; +} + +// Generate and print the injection keywords +const injectionKeywords = generateInjectionKeywords(); + +export default class Heuristic implements Tactic { + name = TacticName.Heuristic; + defaultThreshold: number; + + constructor(threshold: number) { + this.defaultThreshold = threshold; + } + + execute(input: string): Promise { + let highestScore = 0; + const normalizedInput = normalizeString(input); + + for (const keyword of injectionKeywords) { + const normalizedKeyword = normalizeString(keyword); + const keywordParts = normalizedKeyword.split(" "); + const keywordLength = keywordParts.length; + + // Generate substrings of similar length in the input string + const inputParts = normalizedInput.split(" "); + const inputSubstrings = []; + for (let i = 0; i <= inputParts.length - keywordLength; i++) { + inputSubstrings.push(inputParts.slice(i, i + keywordLength).join(" ")); + } + + // Calculate the similarity score between the keyword and each substring + for (const substring of inputSubstrings) { + const similarityScore = stringSimilarity.compareTwoStrings( + normalizedKeyword, + substring + ); + + // Calculate the score based on the number of consecutive words matched + const matchedWordsCount = keywordParts.filter( + (part, index) => substring.split(" ")[index] === part + ).length; + const maxMatchedWords = 5; + const baseScore = + matchedWordsCount > 0 + ? 0.5 + 0.5 * Math.min(matchedWordsCount / maxMatchedWords, 1) + : 0; + + // Adjust the score using the similarity score + const adjustedScore = + baseScore - similarityScore * (1 / (maxMatchedWords * 2)); + + // Update the highest score if the current adjusted score is higher + if (adjustedScore > highestScore) { + highestScore = adjustedScore; + } + } + } + + return Promise.resolve({ score: highestScore }); + } +} diff --git a/javascript-sdk/src/tactics/OpenAI.ts b/javascript-sdk/src/tactics/OpenAI.ts new file mode 100644 index 0000000..7e473a3 --- /dev/null +++ b/javascript-sdk/src/tactics/OpenAI.ts @@ -0,0 +1,40 @@ +import { RebuffError, TacticName } from "../interface"; +import { OpenAIApi } from "openai"; +import { renderPromptForPiDetection } from "../lib/prompts"; +import Tactic, { TacticExecution } from "./Tactic"; + +export default class OpenAI implements Tactic { + name = TacticName.LanguageModel; + private openai: OpenAIApi; + private model: string; + public defaultThreshold: number; + + constructor(threshold: number, model: string, openai: OpenAIApi) { + this.defaultThreshold = threshold; + this.model = model; + this.openai = openai; + } + + async execute(input: string): Promise { + try { + const completion = await this.openai.createChatCompletion({ + model: this.model, + messages: [{ role: "user", content: renderPromptForPiDetection(input) }], + }); + + if (completion.data.choices.length === 0) { + throw new Error("completion.data.choices.length === 0"); + } + if (completion.data.choices[0].message === undefined) { + throw new Error("completion.data.choices[0].message is undefined"); + } + + // FIXME: Handle when parseFloat returns NaN. + const score = parseFloat(completion.data.choices[0].message.content || ""); + return { score }; + } catch (error) { + console.error("Error in callOpenAiToDetectPI:", error); + throw new RebuffError("Error in getting score for large language model"); + } + } +} diff --git a/javascript-sdk/src/tactics/Tactic.ts b/javascript-sdk/src/tactics/Tactic.ts new file mode 100644 index 0000000..ea98ba9 --- /dev/null +++ b/javascript-sdk/src/tactics/Tactic.ts @@ -0,0 +1,20 @@ +import { TacticName } from "../interface"; + +export interface TacticExecution { + // A score between 0 and 1, inclusive, representing the likelihood that the input is prompt + // injection. The closer to 1, the more likely that is it prompt injection. + score: number; + // Optional additional fields that can be used to return additional information about the + // execution of this tactic. + additionalFields?: Record; +} + +export default interface Tactic { + // The name of this tactic. + name: TacticName; + // The threshold to use to determine if the tactic has detected prompt injection. At runtime, + // the caller can provide a different threshold to override the default. + defaultThreshold: number; + // Execute the tactic on the given input. + execute(input: string, thresholdOverride?: number): Promise; +} diff --git a/javascript-sdk/src/tactics/Vector.ts b/javascript-sdk/src/tactics/Vector.ts new file mode 100644 index 0000000..15e4e12 --- /dev/null +++ b/javascript-sdk/src/tactics/Vector.ts @@ -0,0 +1,45 @@ +import { VectorStore } from "langchain/vectorstores/base"; +import { RebuffError, TacticName } from "../interface"; +import Tactic, { TacticExecution } from "./Tactic"; + + +export default class Vector implements Tactic { + name = TacticName.VectorDB; + defaultThreshold: number; + + private vectorStore: VectorStore; + + constructor(threshold: number, vectorStore: VectorStore) { + this.defaultThreshold = threshold; + this.vectorStore = vectorStore; + } + + async execute(input: string, thresholdOverride?: number): Promise { + const threshold = thresholdOverride || this.defaultThreshold; + try { + const topK = 20; + const results = await this.vectorStore.similaritySearchWithScore(input, topK); + + let topScore = 0; + let countOverMaxVectorScore = 0; + + for (const [_, score] of results) { + if (score == undefined) { + continue; + } + + if (score > topScore) { + topScore = score; + } + + if (score >= threshold) { + countOverMaxVectorScore++; + } + } + + return { score: topScore, additionalFields: { countOverMaxVectorScore } }; + } catch (error) { + throw new RebuffError(`Error in getting score for vector tactic: ${error}`); + } + } +} diff --git a/javascript-sdk/tests/index.test.ts b/javascript-sdk/tests/index.test.ts index 22efbf3..f72f76e 100644 --- a/javascript-sdk/tests/index.test.ts +++ b/javascript-sdk/tests/index.test.ts @@ -1,7 +1,7 @@ /* eslint-disable @typescript-eslint/no-unused-expressions */ import { describe } from "mocha"; import { expect } from "chai"; -import { DetectRequest, DetectResponse } from "../src/interface"; +import { DetectRequest, DetectResponse, TacticName, TacticResult } from "../src/interface"; import RebuffSDK from "../src/sdk"; import { getEnvironmentVariable } from "./helpers"; @@ -47,48 +47,62 @@ describe("Rebuff API tests", function () { const request: DetectRequest = { userInput: "abc", userInputBase64: "abc", - runHeuristicCheck: true, - runVectorCheck: false, - runLanguageModelCheck: true, - maxHeuristicScore: 0.5, - maxModelScore: 0.8, - maxVectorScore: 0.0, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: 0.7 }, + { name: TacticName.LanguageModel, run: false }, + ], }; + expect(request).to.have.property("userInput"); expect(request).to.have.property("userInputBase64"); - expect(request).to.have.property("runHeuristicCheck"); - expect(request).to.have.property("runVectorCheck"); - expect(request).to.have.property("runLanguageModelCheck"); - expect(request).to.have.property("maxHeuristicScore"); - expect(request).to.have.property("maxModelScore"); - expect(request).to.have.property("maxVectorScore"); + expect(request).to.have.property("tacticOverrides"); + if (request.tacticOverrides) { + expect(request.tacticOverrides).to.be.an("array"); + expect(request.tacticOverrides).to.have.lengthOf(2); + expect(request.tacticOverrides[0]).to.have.property("name"); + expect(request.tacticOverrides[0]).to.have.property("threshold"); + expect(request.tacticOverrides[1]).to.have.property("name"); + expect(request.tacticOverrides[1]).to.have.property("run"); + } }); }); describe("DetectResponse", () => { it("should have the correct properties", () => { const response: DetectResponse = { - heuristicScore: 0.5, - modelScore: 0.8, - vectorScore: { abc: 0.9 }, - runHeuristicCheck: true, - runVectorCheck: false, - runLanguageModelCheck: true, - maxHeuristicScore: 0.5, - maxModelScore: 0.8, - maxVectorScore: 0.0, injectionDetected: false, + tacticResults: [ + { + name: TacticName.Heuristic, + score: 0.5, + detected: false, + threshold: 0.9, + additionalFields: {}, + }, + { + name: TacticName.LanguageModel, + score: 0.8, + detected: false, + threshold: 0.9, + additionalFields: {}, + } + ], }; - expect(response).to.have.property("heuristicScore"); - expect(response).to.have.property("modelScore"); - expect(response).to.have.property("vectorScore"); - expect(response).to.have.property("runHeuristicCheck"); - expect(response).to.have.property("runVectorCheck"); - expect(response).to.have.property("runLanguageModelCheck"); - expect(response).to.have.property("maxHeuristicScore"); - expect(response).to.have.property("maxModelScore"); - expect(response).to.have.property("maxVectorScore"); + expect(response).to.have.property("injectionDetected"); + expect(response).to.have.property("tacticResults"); + expect(response.tacticResults).to.be.an("array"); + expect(response.tacticResults).to.have.lengthOf(2); + expect(response.tacticResults[0]).to.have.property("name"); + expect(response.tacticResults[0]).to.have.property("score"); + expect(response.tacticResults[0]).to.have.property("detected"); + expect(response.tacticResults[0]).to.have.property("threshold"); + expect(response.tacticResults[0]).to.have.property("additionalFields"); + expect(response.tacticResults[1]).to.have.property("name"); + expect(response.tacticResults[1]).to.have.property("score"); + expect(response.tacticResults[1]).to.have.property("detected"); + expect(response.tacticResults[1]).to.have.property("threshold"); + expect(response.tacticResults[1]).to.have.property("additionalFields"); }); }); @@ -148,27 +162,28 @@ describe("Rebuff API tests", function () { const maxHeuristicScore = 0.5; const maxVectorScore = 0.95; const maxModelScore = 0.9; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = true; const detectResponse = await rb.detectInjection({ userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, - } - ); + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, threshold: maxModelScore }, + ], + }); expect(detectResponse.injectionDetected).to.be.true; - // Check if the 'heuristicScore' attribute is present in the result object - expect(detectResponse).to.have.property("heuristicScore"); + // Check that the one of the results has a name of heuristic + const heuristicResult = detectResponse.tacticResults.find( + result => result.name === TacticName.Heuristic + ); + expect(heuristicResult).to.not.be.undefined; // Ensure that the heuristic score is greater than 0.5 - expect(detectResponse.heuristicScore).to.be.greaterThan(0.50); + if (heuristicResult) { + expect(heuristicResult).to.have.property("score"); + expect(heuristicResult.score).to.be.greaterThan(0.5); + } }); }); @@ -178,26 +193,28 @@ describe("Rebuff API tests", function () { const maxHeuristicScore = 0.1; const maxVectorScore = 0.9; const maxModelScore = 0.1; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = true; const detectResponse = await rb.detectInjection( { userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, threshold: maxModelScore }, + ], } ); - // Check if the 'heuristicScore' attribute is present in the result object - expect(detectResponse).to.have.property("heuristicScore"); + // Check that the one of the results has a name of heuristic + const heuristicResult = detectResponse.tacticResults.find( + result => result.name === TacticName.Heuristic + ); + expect(heuristicResult).to.not.be.undefined; // Ensure that the heuristicScore is less than 0.1 - expect(detectResponse.heuristicScore).to.be.lessThan(0.1); + if (heuristicResult) { + expect(heuristicResult).to.have.property("score"); + expect(heuristicResult.score).to.be.lessThan(0.1); + } }); }); }); @@ -208,26 +225,28 @@ describe("Rebuff API tests", function () { const maxHeuristicScore = 0.5; const maxVectorScore = 0.95; const maxModelScore = 0.9; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = true; const detectResponse = await rb.detectInjection({ userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, threshold: maxModelScore }, + ], } ); expect(detectResponse.injectionDetected).to.be.true; - // Check if the 'modelScore' attribute is present in the result object - expect(detectResponse).to.have.property("modelScore"); + // Check that the one of the results has a name of LanguageModel + const modelResult = detectResponse.tacticResults.find( + result => result.name === TacticName.LanguageModel + ); + expect(modelResult).to.not.be.undefined; // Ensure that the model score is greater than 0.9 - expect(detectResponse.modelScore).to.be.greaterThan(0.9); + if (modelResult) { + expect(modelResult).to.have.property("score"); + expect(modelResult.score).to.be.greaterThan(0.9); + } }); }); @@ -237,26 +256,26 @@ describe("Rebuff API tests", function () { const maxHeuristicScore = 0.1; const maxVectorScore = 0.9; const maxModelScore = 0.1; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = true; - const detectResponse = await rb.detectInjection( - { - userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, - } + const detectResponse = await rb.detectInjection({ + userInput, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, threshold: maxModelScore }, + ], + }); + + // Check that the one of the results has a name of LanguageModel + const modelResult = detectResponse.tacticResults.find( + result => result.name === TacticName.LanguageModel ); - - // Check if the 'modelScore' attribute is present in the result object - expect(detectResponse).to.have.property("modelScore"); + expect(modelResult).to.not.be.undefined; // Ensure that the model score is less than 0.1 - expect(detectResponse.heuristicScore).to.be.lessThan(0.1); + if (modelResult) { + expect(modelResult).to.have.property("score"); + expect(modelResult.score).to.be.lessThan(0.1); + } }); }); }); @@ -267,27 +286,28 @@ describe("Rebuff API tests", function () { const maxHeuristicScore = 0.5; const maxVectorScore = 0.95; const maxModelScore = 0.9; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = true; const detectResponse = await rb.detectInjection({ userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, - } - ); + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, threshold: maxModelScore }, + ], + }); expect(detectResponse.injectionDetected).to.be.true; - // Check if the 'vectorScore' attribute is present in the result object - expect(detectResponse).to.have.property("vectorScore"); + // Check that the one of the results has a name of VectorDB + const vectorResult = detectResponse.tacticResults.find( + result => result.name === TacticName.VectorDB + ); + expect(vectorResult).to.not.be.undefined; - // Ensure that the vector score is greater than 0.95 - expect(detectResponse.vectorScore.topScore).to.be.greaterThan(0.95); + // Ensure that the model score is greater than 0.95 + if (vectorResult) { + expect(vectorResult).to.have.property("score"); + expect(vectorResult.score).to.be.greaterThan(0.95); + } }); }); @@ -297,26 +317,28 @@ describe("Rebuff API tests", function () { const maxHeuristicScore = 0.1; const maxVectorScore = 0.9; const maxModelScore = 0.1; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = true; const detectResponse = await rb.detectInjection( { userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, threshold: maxModelScore }, + ], } ); - // Check if the 'vectorScore' attribute is present in the result object - expect(detectResponse).to.have.property("vectorScore"); + // Check that the one of the results has a name of VectorDB + const vectorResult = detectResponse.tacticResults.find( + result => result.name === TacticName.VectorDB + ); + expect(vectorResult).to.not.be.undefined; - // Ensure that the vector score is less than 0.9 - expect(detectResponse.vectorScore.topScore).to.be.lessThan(0.9); + // Ensure that the model score is less than 0.9 + if (vectorResult) { + expect(vectorResult).to.have.property("score"); + expect(vectorResult.score).to.be.lessThan(0.9); + } }); }); }); @@ -326,28 +348,29 @@ describe("Rebuff API tests", function () { it("should detect prompt injection", async () => { const maxHeuristicScore = 0.5; const maxVectorScore = 0.95; - const maxModelScore = 0.9; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = false; const detectResponse = await rb_chroma.detectInjection({ userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, run: false }, + ], } ); expect(detectResponse.injectionDetected).to.be.true; - // Check if the 'vectorScore' attribute is present in the result object - expect(detectResponse).to.have.property("vectorScore"); + // Check that the one of the results has a name of VectorDB + const vectorResult = detectResponse.tacticResults.find( + result => result.name === TacticName.VectorDB + ); + expect(vectorResult).to.not.be.undefined; - // Ensure that the vector score is greater than 0.95 - expect(detectResponse.vectorScore.topScore).to.be.greaterThan(0.95); + // Ensure that the model score is greater than 0.95 + if (vectorResult) { + expect(vectorResult).to.have.property("score"); + expect(vectorResult.score).to.be.greaterThan(0.95); + } }); }); @@ -356,27 +379,28 @@ describe("Rebuff API tests", function () { it("should not detect prompt injection", async () => { const maxHeuristicScore = 0.1; const maxVectorScore = 0.9; - const maxModelScore = 0.1; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = false; const detectResponse = await rb_chroma.detectInjection( { userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, run: false }, + ], } ); - // Check if the 'vectorScore' attribute is present in the result object - expect(detectResponse).to.have.property("vectorScore"); + // Check that the one of the results has a name of VectorDB + const vectorResult = detectResponse.tacticResults.find( + result => result.name === TacticName.VectorDB + ); + expect(vectorResult).to.not.be.undefined; - // Ensure that the vector score is less than 0.9 - expect(detectResponse.vectorScore.topScore).to.be.lessThan(0.9); + // Ensure that the model score is less than 0.9 + if (vectorResult) { + expect(vectorResult).to.have.property("score"); + expect(vectorResult.score).to.be.lessThan(0.9); + } }); }); }); @@ -385,50 +409,38 @@ describe("Rebuff API tests", function () { prompt_injection_inputs.forEach(function (userInput) { it("should detect prompt injection", async () => { - // Test the isInjectionDetected method const maxHeuristicScore = 0.5; const maxVectorScore = 0.95; const maxModelScore = 0.9; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = true; - - const isInjectionDetected = await rb.detectInjection({ + const detectResponse = await rb.detectInjection({ userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, threshold: maxModelScore }, + ], }); - expect(isInjectionDetected.injectionDetected).to.be.true; + expect(detectResponse.injectionDetected).to.be.true; }); }); benign_inputs.forEach(function (userInput) { it("should not detect prompt injection", async () => { - // Test the isInjectionDetected method const maxHeuristicScore = 0.1; const maxVectorScore = 0.9; const maxModelScore = 0.1; - const runHeuristicCheck = true; - const runVectorCheck = true; - const runLanguageModelCheck = true; - const isInjectionDetected = await rb.detectInjection({ + const detectResponse = await rb.detectInjection({ userInput, - maxHeuristicScore, - maxVectorScore, - maxModelScore, - runHeuristicCheck, - runVectorCheck, - runLanguageModelCheck, + tacticOverrides: [ + { name: TacticName.Heuristic, threshold: maxHeuristicScore }, + { name: TacticName.VectorDB, threshold: maxVectorScore }, + { name: TacticName.LanguageModel, threshold: maxModelScore }, + ], }); - expect(isInjectionDetected.injectionDetected).to.be.false; + expect(detectResponse.injectionDetected).to.be.false; }); }); }); }); - From 06440663c2c3e446d32c7b1323442146ec27c3be Mon Sep 17 00:00:00 2001 From: Risto McGehee Date: Wed, 27 Dec 2023 21:18:09 -0800 Subject: [PATCH 2/3] Use npx for tests --- javascript-sdk/tests/setup-and-run-tests.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/javascript-sdk/tests/setup-and-run-tests.sh b/javascript-sdk/tests/setup-and-run-tests.sh index 565a6d5..e9ddf6c 100755 --- a/javascript-sdk/tests/setup-and-run-tests.sh +++ b/javascript-sdk/tests/setup-and-run-tests.sh @@ -10,10 +10,10 @@ docker rm chroma || true docker run -d -p 8000:8000 --name chroma chromadb/chroma echo "Waiting for Chroma server to start..." -ts-node tests/wait-for-chroma.ts +npx ts-node tests/wait-for-chroma.ts echo "Inserting Chroma test vectors..." -ts-node tests/insert-chroma-vectors.ts +npx ts-node tests/insert-chroma-vectors.ts echo "Running tests..." -mocha --require ts-node/register tests/**/*.test.ts +npx mocha --require ts-node/register tests/**/*.test.ts From 07a8d4291f0d77d4884deea4074df13320daf486 Mon Sep 17 00:00:00 2001 From: Risto McGehee Date: Fri, 12 Jan 2024 21:17:56 -0800 Subject: [PATCH 3/3] Fix test error --- javascript-sdk/tests/setup-and-run-tests.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/javascript-sdk/tests/setup-and-run-tests.sh b/javascript-sdk/tests/setup-and-run-tests.sh index e9ddf6c..fa7b966 100755 --- a/javascript-sdk/tests/setup-and-run-tests.sh +++ b/javascript-sdk/tests/setup-and-run-tests.sh @@ -10,10 +10,10 @@ docker rm chroma || true docker run -d -p 8000:8000 --name chroma chromadb/chroma echo "Waiting for Chroma server to start..." -npx ts-node tests/wait-for-chroma.ts +node --loader ts-node/esm tests/wait-for-chroma.ts echo "Inserting Chroma test vectors..." -npx ts-node tests/insert-chroma-vectors.ts +node --loader ts-node/esm tests/insert-chroma-vectors.ts echo "Running tests..." npx mocha --require ts-node/register tests/**/*.test.ts