Create function to init sdk

protectai · Jan 15, 2024 · e7d1ad2 · e7d1ad2
1 parent b54c827
commit e7d1ad2
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 38 deletions.
diff --git a/javascript-sdk/src/sdk.ts b/javascript-sdk/src/sdk.ts
@@ -9,7 +9,6 @@ import crypto from "crypto";
 import { SdkConfig } from "./config";
 import initVectorStore from "./lib/vectordb";
 import getOpenAIInstance from "./lib/openai";
-import { OpenAIApi } from "openai";
 import { VectorStore } from "langchain/vectorstores/base";
 import { Document } from "langchain/document";
 import Strategy from "./lib/Strategy";
@@ -23,44 +22,36 @@ function generateCanaryWord(length = 8): string {
 }
 
 export default class RebuffSdk implements Rebuff {
-  private sdkConfig: SdkConfig;
-  private vectorStore: VectorStore | undefined;
-  private strategies: Record<string, Strategy> | undefined;
+  private vectorStore: VectorStore;
+  private strategies: Record<string, Strategy>;
   private defaultStrategy: string;
 
-  private openai: {
-    conn: OpenAIApi;
-    model: string;
-  };
+  private constructor(strategies: Record<string, Strategy>, vectorStore: VectorStore) {
+    this.vectorStore = vectorStore;
+    this.strategies = strategies;
+    this.defaultStrategy = "standard";
+  }
 
-  constructor(config: SdkConfig) {
-    this.sdkConfig = config;
-    this.openai = {
+  public static async init(config: SdkConfig): Promise<RebuffSdk> {
+    const vectorStore = await initVectorStore(config);
+    const openai = {
       conn: getOpenAIInstance(config.openai.apikey),
       model: config.openai.model || "gpt-3.5-turbo",
     };
-    this.defaultStrategy = "standard";
-  }
-
-  private async getStrategies(): Promise<Record<string, Strategy>> {
-    if (this.strategies) {
-      return this.strategies;
-    }
     const heuristicScoreThreshold = 0.75;
     const vectorScoreThreshold = 0.9;
     const openaiScoreThreshold = 0.9;
-    const strategies: Record<string, Strategy> = {
+    const strategies = {
       // 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),
+          new Vector(vectorScoreThreshold, vectorStore),
+          new OpenAI(openaiScoreThreshold, openai.model, openai.conn),
         ]
       },
     };
-    this.strategies = strategies;
-    return this.strategies;
+    return new RebuffSdk(strategies, vectorStore);
   }
 
   async detectInjection({
@@ -78,10 +69,9 @@ export default class RebuffSdk implements Rebuff {
       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) {
+    for (const tactic of this.strategies[this.defaultStrategy].tactics) {
       const tacticOverride = tacticOverrides.find(t => t.name === tactic.name);
       if (tacticOverride && tacticOverride.run === false) {
         continue;
@@ -134,19 +124,11 @@ export default class RebuffSdk implements Rebuff {
     return false;
   }
 
-  async getVectorStore(): Promise<VectorStore> {
-    if (this.vectorStore) {
-      return this.vectorStore;
-    }
-    this.vectorStore = await initVectorStore(this.sdkConfig);
-    return this.vectorStore
-  }
-
   async logLeakage(
     input: string,
     metaData: Record<string, string>
   ): Promise<void> {
-    await (await this.getVectorStore()).addDocuments([new Document({
+    await this.vectorStore.addDocuments([new Document({
       metadata: metaData,
       pageContent: input,
     })]);

diff --git a/javascript-sdk/tests/index.test.ts b/javascript-sdk/tests/index.test.ts
@@ -1,12 +1,12 @@
 /* eslint-disable @typescript-eslint/no-unused-expressions */
 import { describe } from "mocha";
 import { expect } from "chai";
-import { DetectRequest, DetectResponse, TacticName, TacticResult } from "../src/interface";
+import { DetectRequest, DetectResponse, TacticName } from "../src/interface";
 import RebuffSDK from "../src/sdk";
 import { getEnvironmentVariable } from "./helpers";
 
 // Initialize the Rebuff SDK with a real API token and URL
-const rb = new RebuffSDK({
+const rb = await RebuffSDK.init({
   openai: {
     apikey: getEnvironmentVariable("OPENAI_API_KEY"),
     model: "gpt-3.5-turbo",
@@ -19,7 +19,7 @@ const rb = new RebuffSDK({
     }
   }
 });
-const rb_chroma = new RebuffSDK({
+const rb_chroma = await RebuffSDK.init({
   openai: {
     apikey: getEnvironmentVariable("OPENAI_API_KEY"),
     model: "gpt-3.5-turbo",

diff --git a/javascript-sdk/tsconfig.json b/javascript-sdk/tsconfig.json
@@ -1,6 +1,6 @@
 {
   "compilerOptions": {
-    "target": "es2015",
+    "target": "es2017",
     "lib": ["dom", "dom.iterable", "esnext"],
     "module": "esnext",
     "strict": true,