Use langchain for pinecone

javascript-sdk/src/lib/detect.ts

@@ -1,51 +1,32 @@
 import stringSimilarity from "string-similarity";
 import { normalizeString } from "./prompts";
-import { PineconeClient } from "@pinecone-database/pinecone";
 import { OpenAIApi } from "openai";
 import { RebuffError } from "../interface";
+import { VectorStore } from "langchain/vectorstores/base";
 
 export async function detectPiUsingVectorDatabase(
   input: string,
   similarityThreshold: number,
-  pinecone: PineconeClient,
-  pineconeIndex: string,
-  openai: OpenAIApi
+  vectorStore: VectorStore
 ): Promise<{ topScore: number; countOverMaxVectorScore: number }> {
   try {
-    // Create embedding from input
-    const emb = await openai.createEmbedding({
-      model: "text-embedding-ada-002",
-      input: input,
-    });
-
-    // Get Pinecone Index
-    const index = pinecone.Index(pineconeIndex);
-
-    // Query similar embeddings
-    const queryResponse = await index.query({
-      queryRequest: {
-        vector: emb.data.data[0].embedding,
-        topK: 20,
-        includeValues: true,
-      },
-    });
+    const topK = 20;
+    const results = await vectorStore.similaritySearchWithScore(input, topK);
 
     let topScore = 0;
     let countOverMaxVectorScore = 0;
 
-    if (queryResponse.matches != undefined) {
-      for (const match of queryResponse.matches) {
-        if (match.score == undefined) {
-          continue;
-        }
+    for (const [_, score] of results) {
+      if (score == undefined) {
+        continue;
+      }
 
-        if (match.score > topScore) {
-          topScore = match.score;
-        }
+      if (score > topScore) {
+        topScore = score;
+      }
 
-        if (match.score >= similarityThreshold && match.score > topScore) {
-          countOverMaxVectorScore++;
-        }
+      if (score >= similarityThreshold && score > topScore) {
+        countOverMaxVectorScore++;
       }
     }
 

javascript-sdk/src/lib/vectordb.ts

@@ -1,9 +1,14 @@
+import { OpenAIEmbeddings } from "langchain/embeddings/openai";
 import { PineconeClient } from "@pinecone-database/pinecone";
+import { PineconeStore } from "langchain/vectorstores/pinecone";
+
 
 export default async function initPinecone(
   environment: string,
-  apiKey: string
-): Promise<PineconeClient> {
+  apiKey: string,
+  index: string,
+  openaiApiKey: string,
+): Promise<PineconeStore> {
   if (!environment) {
     throw new Error("Pinecone environment definition missing");
   }
@@ -17,8 +22,17 @@ export default async function initPinecone(
       environment,
       apiKey,
     });
+    const openaiEmbeddings = new OpenAIEmbeddings({
+      openAIApiKey: openaiApiKey,
+      modelName: "text-embedding-ada-002" 
+    });
+    const pineconeIndex = pinecone.Index(index);
+    const vectorStore = await PineconeStore.fromExistingIndex(
+      openaiEmbeddings,
+      { pineconeIndex }
+    );
 
-    return pinecone;
+    return vectorStore;
   } catch (error) {
     console.log("error", error);
     throw new Error("Failed to initialize Pinecone Client");

javascript-sdk/src/sdk.ts

@@ -7,8 +7,6 @@ import {
 import crypto from "crypto";
 import { SdkConfig } from "./config";
 import initPinecone from "./lib/vectordb";
-import { v4 as uuidv4 } from "uuid";
-import { PineconeClient } from "@pinecone-database/pinecone";
 import {
   callOpenAiToDetectPI,
   detectPiUsingVectorDatabase,
@@ -17,38 +15,29 @@ import {
 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";
 
 function generateCanaryWord(length = 8): string {
   // Generate a secure random hexadecimal canary word
   return crypto.randomBytes(length / 2).toString("hex");
 }
 
 export default class RebuffSdk implements Rebuff {
-  private pinecone: {
-    conn?: PineconeClient;
-    index: string;
-  };
+  private vectorStore: VectorStore | undefined;
+  private sdkConfig: SdkConfig;
 
   private openai: {
     conn: OpenAIApi;
     model: string;
   };
 
   constructor(config: SdkConfig) {
-    this.pinecone = { index: "" };
+    this.sdkConfig = config;
     this.openai = {
       conn: getOpenAIInstance(config.openai.apikey),
       model: config.openai.model || "gpt-3.5-turbo",
     };
-    (async () => {
-      this.pinecone = {
-        conn: await initPinecone(
-          config.pinecone.environment,
-          config.pinecone.apikey
-        ),
-        index: config.pinecone.index,
-      };
-    })();
   }
 
   async detectInjection({
@@ -118,9 +107,7 @@ export default class RebuffSdk implements Rebuff {
       ? await detectPiUsingVectorDatabase(
           userInput,
           maxVectorScore,
-          await this.getPineconeConnection(),
-          this.pinecone.index,
-          this.openai.conn
+          await this.getVectorStore()
         )
       : { topScore: 0, countOverMaxVectorScore: 0 };
     const injectionDetected =
@@ -169,48 +156,26 @@ export default class RebuffSdk implements Rebuff {
     return false;
   }
 
-  // Calling functions immediately after constructor can cause issues if pinecone connection needs time
-  async getPineconeConnection(): Promise<PineconeClient> {
-    if (this.pinecone.conn) {
-      return this.pinecone.conn;
-    }
-    // Wait 1 second for pinecone connection to be initialized by constructor
-    await new Promise((resolve) => setTimeout(resolve, 1000));
-    // If pinecone connection is still not initialized, throw an error
-    if (this.pinecone.conn) {
-      return this.pinecone.conn;
+  async getVectorStore(): Promise<VectorStore> {
+    if (this.vectorStore) {
+      return this.vectorStore;
     }
-    throw new RebuffError("Pinecone connection not initialized yet");
+    this.vectorStore = await initPinecone(
+      this.sdkConfig.pinecone.environment,
+      this.sdkConfig.pinecone.apikey,
+      this.sdkConfig.pinecone.index,
+      this.sdkConfig.openai.apikey
+    );
+    return this.vectorStore
   }
 
   async logLeakage(
     input: string,
     metaData: Record<string, string>
   ): Promise<void> {
-    const emb = await this.openai.conn.createEmbedding({
-      model: "text-embedding-ada-002",
-      input: input,
-    });
-
-    // Get Pinecone Index
-    const index = (await this.getPineconeConnection()).Index(
-      this.pinecone.index
-    );
-
-    // Insert embedding into index
-    await index.upsert({
-      upsertRequest: {
-        vectors: [
-          {
-            id: uuidv4(),
-            values: emb.data.data[0].embedding,
-            metadata: {
-              input: input,
-              ...metaData,
-            },
-          },
-        ],
-      },
-    });
+    await (await this.getVectorStore()).addDocuments([new Document({
+      metadata: metaData,
+      pageContent: input,
+    })]);
   }
 }