Supabase vector store 2

apps/prompt-hack/pages/api/game.ts

@@ -19,6 +19,7 @@ const rb = new RebuffSdk({
     apikey: getEnvironmentVariable("OPENAI_API_KEY"),
     model: "gpt-3.5-turbo",
   },
+  vectorStore: "pinecone",
   pinecone: {
     environment: getEnvironmentVariable("PINECONE_ENVIRONMENT"),
     apikey: getEnvironmentVariable("PINECONE_API_KEY"),

javascript-sdk/src/api.ts

@@ -6,7 +6,7 @@ import {
 } from "./interface";
 import fetch from "node-fetch";
 import crypto from "crypto";
-import { ApiConfig } from "./config";
+import { ApiConfig } from "./lib/config";
 
 function encodeString(message: string): string {
   return Buffer.from(message, "utf-8").toString("hex");

javascript-sdk/src/config.ts → javascript-sdk/src/lib/config.ts

@@ -4,11 +4,16 @@ export interface ApiConfig {
 }
 
 export interface SdkConfig {
-  pinecone: {
+  vectorStore: string;
+  pinecone?: {
     apikey: string;
     environment: string;
     index: string;
   };
+  supabase?: {
+    serviceKey: string;
+    url: string;
+  };
   openai: {
     apikey: string;
     model: string;

javascript-sdk/src/lib/vectordb.ts

@@ -1,31 +1,34 @@
 import { OpenAIEmbeddings } from "langchain/embeddings/openai";
-import { PineconeClient } from "@pinecone-database/pinecone";
 import { PineconeStore } from "langchain/vectorstores/pinecone";
+import { SupabaseVectorStore } from "langchain/vectorstores/supabase";
+import { VectorStore } from "langchain/vectorstores/base";
+import { PineconeClient } from "@pinecone-database/pinecone";
+import { createClient } from "@supabase/supabase-js";
+import { SdkConfig } from "./lib/config";
 
 
-export default async function initPinecone(
+async function initPinecone(
   environment: string,
   apiKey: string,
   index: string,
-  openaiApiKey: string,
+  openaiEmbeddings: OpenAIEmbeddings,
 ): Promise<PineconeStore> {
   if (!environment) {
     throw new Error("Pinecone environment definition missing");
   }
   if (!apiKey) {
     throw new Error("Pinecone apikey definition missing");
   }
+  if (!index) {
+    throw new Error("Pinecone index definition missing");
+  }
   try {
     const pinecone = new PineconeClient();
 
     await pinecone.init({
       environment,
       apiKey,
     });
-    const openaiEmbeddings = new OpenAIEmbeddings({
-      openAIApiKey: openaiApiKey,
-      modelName: "text-embedding-ada-002" 
-    });
     const pineconeIndex = pinecone.Index(index);
     const vectorStore = await PineconeStore.fromExistingIndex(
       openaiEmbeddings,
@@ -38,3 +41,54 @@ export default async function initPinecone(
     throw new Error("Failed to initialize Pinecone Client");
   }
 }
+
+async function initSupabase(
+  serviceKey: string,
+  url: string,
+  openaiEmbeddings: OpenAIEmbeddings,
+): Promise<SupabaseVectorStore> {
+  if (!serviceKey) {
+    throw new Error("Supabase service key definition missing");
+  }
+  if (!url) {
+    throw new Error("Supabase URL definition missing");
+  }
+  try {
+    const client = createClient(url, serviceKey);
+    const vectorStore = new SupabaseVectorStore(openaiEmbeddings, {
+      client,
+      tableName: "documents",
+    });
+
+    return vectorStore;
+  } catch (error) {
+    console.log("error", error);
+    throw new Error("Failed to initialize Supabase client");
+  }
+}
+
+export default async function initVectorStore(
+  config: SdkConfig
+): Promise<VectorStore> {
+  const openaiEmbeddings = new OpenAIEmbeddings({
+    openAIApiKey: config.openai.apikey,
+    modelName: "text-embedding-ada-002" 
+  });
+  switch (config.vectorStore.toLowerCase()) {
+    case "pinecone":
+      return await initPinecone(
+        config.pinecone.environment,
+        config.pinecone.apikey,
+        config.pinecone.index,
+        openaiEmbeddings
+      );
+    case "supabase":
+      return await initSupabase(
+        config.supabase.serviceKey,
+        config.supabase.url,
+        openaiEmbeddings
+      );
+    default:
+      throw new Error("Unsupported vector store: " + config.vectorStore);
+  }
+}

javascript-sdk/src/sdk.ts

@@ -5,8 +5,8 @@ import {
   RebuffError,
 } from "./interface";
 import crypto from "crypto";
-import { SdkConfig } from "./config";
-import initPinecone from "./lib/vectordb";
+import { SdkConfig } from "./lib/config";
+import initVectorStore from "./lib/vectordb";
 import {
   callOpenAiToDetectPI,
   detectPiUsingVectorDatabase,
@@ -115,6 +115,14 @@ export default class RebuffSdk implements Rebuff {
       modelScore > maxModelScore ||
       vectorScore.topScore > maxVectorScore;
 
+    if (injectionDetected) {
+      await this.logLeakage(userInput, {
+        heuristicScore: heuristicScore.toString(),
+        modelScore: modelScore.toString(),
+        vectorScore: vectorScore.topScore.toString(),
+      });
+    }
+
     return {
       heuristicScore,
       modelScore,

package.json

@@ -41,7 +41,7 @@
     "@supabase/auth-helpers-nextjs": "^0.6.0",
     "@supabase/auth-helpers-react": "^0.3.1",
     "@supabase/auth-ui-react": "^0.2.8",
-    "@supabase/supabase-js": "^2.8.0",
+    "@supabase/supabase-js": "^2.33.2",
     "@tabler/icons-react": "^2.17.0",
     "@types/react": "18.0.28",
     "@types/react-dom": "18.0.11",

server/lib/rebuff.ts

@@ -6,9 +6,14 @@ export const rebuff = new RebuffSdk({
     apikey: getEnvironmentVariable("OPENAI_API_KEY"),
     model: "gpt-3.5-turbo",
   },
+  vectorStore: getEnvironmentVariable("VECTOR_STORE"),
   pinecone: {
-    environment: getEnvironmentVariable("PINECONE_ENVIRONMENT"),
-    apikey: getEnvironmentVariable("PINECONE_API_KEY"),
-    index: getEnvironmentVariable("PINECONE_INDEX_NAME"),
+    environment: process.env["PINECONE_ENVIRONMENT"],
+    apikey: process.env["PINECONE_API_KEY"],
+    index: process.env["PINECONE_INDEX_NAME"],
+  },
+  supabase: {
+    serviceKey: process.env["SUPABASE_SERVICE_KEY"],
+    url: process.env["NEXT_PUBLIC_SUPABASE_URL"],
   },
 });

server/sql_setup/functions/match_documents.sql

@@ -0,0 +1,28 @@
+-- Only necessary if using Supabase as the vector store
+
+create function match_documents (
+  query_embedding vector(1536),
+  match_count int DEFAULT null,
+  filter jsonb DEFAULT '{}'
+) returns table (
+  id bigint,
+  content text,
+  metadata jsonb,
+  similarity float
+)
+language plpgsql
+as $$
+#variable_conflict use_column
+begin
+  return query
+  select
+    id,
+    content,
+    metadata,
+    1 - (documents.embedding <=> query_embedding) as similarity
+  from documents
+  where metadata @> filter
+  order by documents.embedding <=> query_embedding
+  limit match_count;
+end;
+$$;

server/sql_setup/tables/documents.sql

@@ -0,0 +1,11 @@
+-- Only necessary if using Supabase as the vector store
+
+-- Enable the pgvector extension to work with embedding vectors
+create extension vector;
+
+create table documents (
+  id bigserial primary key,
+  content text,
+  metadata jsonb,
+  embedding vector(1536)
+);