WIP RAG testing frontend

backend/endpoints/v1/tools/index.js

@@ -1,8 +1,13 @@
+const { DocumentVectors } = require("../../../models/documentVectors");
 const { Organization } = require("../../../models/organization");
 const {
   OrganizationConnection,
 } = require("../../../models/organizationConnection");
+const {
+  OrganizationWorkspace,
+} = require("../../../models/organizationWorkspace");
 const { Queue } = require("../../../models/queue");
+const { SystemSettings } = require("../../../models/systemSettings");
 const {
   userFromSession,
   validSessionForUser,
@@ -14,6 +19,8 @@ const {
 const {
   organizationResetJob,
 } = require("../../../utils/jobs/organizationResetJob");
+const { OpenAi } = require("../../../utils/openAi");
+const { selectConnector } = require("../../../utils/vectordatabases/providers");
 
 process.env.NODE_ENV === "development"
   ? require("dotenv").config({ path: `.env.${process.env.NODE_ENV}` })
@@ -58,13 +65,11 @@ function toolEndpoints(app) {
           organization_id: Number(organization.id),
         });
         if (!originalConnector) {
-          response
-            .status(200)
-            .json({
-              success: false,
-              message:
-                "No vector database is connected to the original organization.",
-            });
+          response.status(200).json({
+            success: false,
+            message:
+              "No vector database is connected to the original organization.",
+          });
           return;
         }
 
@@ -83,13 +88,11 @@ function toolEndpoints(app) {
           organization_id: Number(destinationOrg.id),
         });
         if (!destinationConnector) {
-          response
-            .status(200)
-            .json({
-              success: false,
-              message:
-                "No vector database is connected to the destination organization.",
-            });
+          response.status(200).json({
+            success: false,
+            message:
+              "No vector database is connected to the destination organization.",
+          });
           return;
         }
 
@@ -149,12 +152,10 @@ function toolEndpoints(app) {
           organization_id: Number(organization.id),
         });
         if (!connector) {
-          response
-            .status(200)
-            .json({
-              success: false,
-              message: "No vector database is connected to this organization.",
-            });
+          response.status(200).json({
+            success: false,
+            message: "No vector database is connected to this organization.",
+          });
           return;
         }
 
@@ -181,6 +182,112 @@ function toolEndpoints(app) {
       }
     }
   );
+
+  app.post(
+    "/v1/tools/org/:orgSlug/workspace-similarity-search",
+    [validSessionForUser],
+    async function (request, response) {
+      try {
+        let queryVector;
+        const { orgSlug } = request.params;
+        const {
+          workspaceId,
+          input,
+          inputType = "text",
+          topK = 3,
+        } = reqBody(request);
+        const user = await userFromSession(request);
+        if (!user || user.role !== "admin") {
+          response.sendStatus(403).end();
+          return;
+        }
+
+        const organization = await Organization.getWithOwner(user.id, {
+          slug: orgSlug,
+        });
+        if (!organization) {
+          response.status(200).json({ results: [], error: "No org found." });
+          return;
+        }
+
+        const workspace = await OrganizationWorkspace.get({
+          id: workspaceId,
+          organization_id: organization.id,
+        });
+        if (!workspace) {
+          response
+            .status(200)
+            .json({ results: [], error: "No workspace found." });
+          return;
+        }
+
+        const connector = await OrganizationConnection.get({
+          organization_id: Number(organization.id),
+        });
+        if (!connector) {
+          response.status(200).json({
+            results: [],
+            error: "No vector database is connected to this organization.",
+          });
+          return;
+        }
+
+        if (inputType === "text") {
+          if (input?.length === 0) {
+            response.status(200).json({
+              results: [],
+              error: "No input data to embed.",
+            });
+            return;
+          }
+
+          const openAiKey = (
+            await SystemSettings.get({ label: "open_ai_api_key" })
+          )?.value;
+          if (!openAiKey) {
+            response.status(200).json({
+              results: [],
+              error: "No embedding API key set - cannot embed text data.",
+            });
+            return;
+          }
+
+          const openai = new OpenAi(openAiKey);
+          queryVector = await openai.embedTextChunk(input);
+        } else {
+          queryVector = input;
+        }
+
+        if (!queryVector || queryVector?.length === 0) {
+          response.status(200).json({
+            results: [],
+            error: "Failed to embed or parse input data.",
+          });
+          return;
+        }
+
+        const vectorDb = selectConnector(connector);
+        const searchResults = await vectorDb.similarityResponse(
+          workspace.fname,
+          queryVector,
+          topK
+        );
+        const results = searchResults.vectorIds.map((_, i) => {
+          return {
+            vectorId: searchResults.vectorIds[i],
+            text: searchResults.contextTexts[i],
+            metadata: searchResults.sourceDocuments[i],
+            score: searchResults.scores[i],
+          };
+        });
+
+        response.status(200).json({ results, error: null });
+      } catch (e) {
+        console.log(e.message, e);
+        response.sendStatus(500).end();
+      }
+    }
+  );
 }
 
 module.exports = { toolEndpoints };

backend/utils/search/documentEmbeddings/semantic.js

@@ -38,11 +38,11 @@ async function semanticSearch(document, query) {
 
   // From similarity search we can find all document vector DB items to infer their associated
   // document record.
-  const searchString = searchResults.vectorIds
-    .map((vid) => `'${vid}'`)
-    .join(",");
   const fragments = await DocumentVectors.where(
-    { vectorId: { in: queryString }, document_id: Number(document.id) },
+    {
+      vectorId: { in: searchResults?.vectorIds || [] },
+      document_id: Number(document.id),
+    },
     100
   );
   return { fragments, error: null };

backend/utils/search/workspaceDocuments/semantic.js

@@ -33,11 +33,8 @@ async function semanticSearch(workspace, query) {
 
   // From similarity search we can find all document vector DB items to infer their associated
   // document record.
-  const searchString = searchResults.vectorIds
-    .map((vid) => `'${vid}'`)
-    .join(",");
   const matchingDocumentVectors = await DocumentVectors.where({
-    vectorId: { in: { searchString } },
+    vectorId: { in: searchResults?.vectorIds || [] },
   });
   const docDbIds = new Set();
   matchingDocumentVectors.forEach((record) => docDbIds.add(record.document_id));

backend/utils/vectordatabases/providers/chroma/index.js

@@ -18,6 +18,13 @@ class Chroma {
     return { type, settings };
   }
 
+  distanceToScore(distance = null) {
+    if (distance === null || typeof distance !== "number") return 0.0;
+    if (distance >= 1.0) return 1;
+    if (distance <= 0) return 0;
+    return 1 - distance;
+  }
+
   async connect() {
     const { ChromaClient } = require("chromadb");
     const { type, settings } = this.config;
@@ -221,23 +228,25 @@ class Chroma {
     }
   }
 
-  async similarityResponse(namespace, queryVector) {
+  async similarityResponse(namespace, queryVector, topK = 4) {
     const { client } = await this.connect();
     const collection = await client.getCollection({ name: namespace });
     const result = {
       vectorIds: [],
       contextTexts: [],
       sourceDocuments: [],
+      scores: [],
     };
 
     const response = await collection.query({
       queryEmbeddings: queryVector,
-      nResults: 4,
+      nResults: topK,
     });
     response.ids[0].forEach((_, i) => {
       result.vectorIds.push(response.ids[0][i]);
       result.contextTexts.push(response.documents[0][i]);
       result.sourceDocuments.push(response.metadatas[0][i]);
+      result.scores.push(this.distanceToScore(response.distances[0][i]));
     });
 
     return result;

backend/utils/vectordatabases/providers/pinecone/index.js

@@ -374,18 +374,19 @@ class Pinecone {
     }
   }
 
-  async similarityResponse(namespace, queryVector) {
+  async similarityResponse(namespace, queryVector, topK = 4) {
     const { pineconeIndex } = await this.connect();
     const result = {
       vectorIds: [],
       contextTexts: [],
       sourceDocuments: [],
+      scores: [],
     };
     const response = await pineconeIndex.query({
       queryRequest: {
         namespace,
         vector: queryVector,
-        topK: 4,
+        topK,
         includeMetadata: true,
       },
     });
@@ -394,6 +395,7 @@ class Pinecone {
       result.vectorIds.push(match.id);
       result.contextTexts.push(match.metadata.text);
       result.sourceDocuments.push(match);
+      result.scores.push(match.score);
     });
 
     return result;

backend/utils/vectordatabases/providers/qdrant/index.js

@@ -195,17 +195,18 @@ class QDrant {
     }
   }
 
-  async similarityResponse(namespace, queryVector) {
+  async similarityResponse(namespace, queryVector, topK = 4) {
     const { client } = await this.connect();
     const result = {
       vectorIds: [],
       contextTexts: [],
       sourceDocuments: [],
+      scores: [],
     };
 
     const responses = await client.search(namespace, {
       vector: queryVector,
-      limit: 4,
+      limit: topK,
       with_payload: true,
     });
 
@@ -216,6 +217,7 @@ class QDrant {
         id: response.id,
       });
       result.vectorIds.push(response.id);
+      result.scores.push(response.score);
     });
 
     return result;

backend/utils/vectordatabases/providers/weaviate/index.js

@@ -379,28 +379,31 @@ class Weaviate {
     }
   }
 
-  async similarityResponse(namespace, queryVector) {
+  async similarityResponse(namespace, queryVector, topK = 4) {
     const { client } = await this.connect();
     const className = this.camelCase(namespace);
     const result = {
       vectorIds: [],
       contextTexts: [],
       sourceDocuments: [],
+      scores: [],
     };
     const fieldsForCollection = await this.fieldNamesForCollection(namespace);
-    const queryString = `${fieldsForCollection.join(" ")} _additional { id }`;
+    const queryString = `${fieldsForCollection.join(
+      " "
+    )} _additional { id certainty }`;
     const queryResponse = await client.graphql
       .get()
       .withClassName(className)
       .withFields(queryString)
       .withNearVector({ vector: queryVector })
-      .withLimit(4)
+      .withLimit(topK)
       .do();
 
     const responses = queryResponse?.data?.Get?.[className];
     responses.forEach((response) => {
       const {
-        _additional: { id },
+        _additional: { id, certainty },
         ...rest
       } = response;
       result.contextTexts.push(rest?.text || "");
@@ -409,6 +412,7 @@ class Weaviate {
         id,
       });
       result.vectorIds.push(id);
+      result.scores.push(certainty);
     });
 
     return result;

frontend/src/App.tsx

@@ -28,6 +28,7 @@ const MigrateConnectionView = lazy(
   () => import('./pages/Tools/MigrateConnection')
 );
 const ResetConnectionView = lazy(() => import('./pages/Tools/ResetConnection'));
+const RAGDriftTestingView = lazy(() => import('./pages/Tools/RAGDrift'));
 
 function App() {
   return (
@@ -88,6 +89,10 @@ function App() {
             path="/dashboard/:slug/tools/db-reset"
             element={<PrivateRoute Component={ResetConnectionView} />}
           />
+          <Route
+            path="/dashboard/:slug/tools/rag-drift"
+            element={<PrivateRoute Component={RAGDriftTestingView} />}
+          />
 
           <Route path="/auth/sign-up" element={<SignUp />} />
           <Route path="/auth/sign-in" element={<SignIn />} />

frontend/src/models/organization.ts

@@ -318,3 +318,11 @@ const Organization = {
 };
 
 export default Organization;
+export interface IOrganization {
+  id: number;
+  name: string;
+  slug: string;
+  uuid: string;
+  createdAt: string;
+  lastUpdatedAt: string;
+}

frontend/src/models/tools.ts

@@ -34,6 +34,65 @@ const Tools = {
         return { success: false, message: e.message };
       });
   },
+  ragTests: async (
+    slug: string
+  ): Promise<{ ragTests: []; message: null | string }> => {
+    return fetch(`${API_BASE}/v1/tools/org/${slug}/rag-tests`, {
+      method: 'GET',
+      cache: 'no-cache',
+      headers: baseHeaders(),
+    })
+      .then((res) => res.json())
+      .catch((e) => {
+        console.error(e);
+        return { ragTests: [], message: e.message };
+      });
+  },
+  newRAGTest: async (
+    slug: string,
+    settings: object
+  ): Promise<{ test: object | null; error: null | string }> => {
+    return fetch(`${API_BASE}/v1/tools/org/${slug}/rag-tests/create`, {
+      method: 'POST',
+      cache: 'no-cache',
+      headers: baseHeaders(),
+      body: JSON.stringify({ settings }),
+    })
+      .then((res) => res.json())
+      .catch((e) => {
+        console.error(e);
+        return { test: null, error: e.message };
+      });
+  },
+
+  // Generic Uitls
+  workspaceSimilaritySearch: async (
+    orgSlug: string,
+    input: string,
+    inputType: 'vector' | 'text' = 'text',
+    workspaceId: number,
+    topK: number = 3
+  ): Promise<{ results: []; error: null | string }> => {
+    return fetch(
+      `${API_BASE}/v1/tools/org/${orgSlug}/workspace-similarity-search`,
+      {
+        method: 'POST',
+        cache: 'no-cache',
+        headers: baseHeaders(),
+        body: JSON.stringify({
+          topK,
+          input,
+          inputType,
+          workspaceId,
+        }),
+      }
+    )
+      .then((res) => res.json())
+      .catch((e) => {
+        console.error(e);
+        return { results: [], error: e.message };
+      });
+  },
 };
 
 export default Tools;

frontend/src/models/workspace.ts

@@ -200,3 +200,14 @@ const Workspace = {
 };
 
 export default Workspace;
+export interface IWorkspace {
+  id: number;
+  name: string;
+  slug: string;
+  fname: string;
+  uuid: string;
+  organization_id: number;
+  createdAt: string;
+  lastUpdatedAt: number;
+  documentCount: number;
+}

frontend/src/pages/Tools/ToolsList/index.tsx

@@ -9,21 +9,21 @@ export default function ToolsList({ organization }: { organization: any }) {
       <div className="flex items-start justify-between">
         <div className="mb-6 px-6">
           <h4 className="text-3xl font-semibold text-black dark:text-white">
-            Organization Tools
+            Advanced management tools
           </h4>
           <p className="text-gray-600">
             below are a list of advanced {APP_NAME} only tools and services that
-            you can use to manage your vectorized data.
+            you can use to manage your connected vector database.
           </p>
         </div>
       </div>
 
       <div className="px-6">
         <ToolItem
-          title="Reset vector database"
-          description="Wipe out all existing data in your vector database in one click."
+          title="Retrieval Drift Testing & Alerts"
+          description="Catch context drift in production before it reaches your customers as your vector store changes over time."
           available={true}
-          linkTo={paths.tools.resetTool(organization)}
+          linkTo={paths.tools.ragDrift(organization)}
         />
         <ToolItem
           title="Migrate vector database to another provider"
@@ -37,9 +37,10 @@ export default function ToolsList({ organization }: { organization: any }) {
           available={false}
         />
         <ToolItem
-          title="Workspace Unit Testing & Alerts"
-          description="Ensure that your workspaces responds with the correct context and responses as embeddings change over time using various metrics."
-          available={false}
+          title="Reset vector database"
+          description="Wipe out all existing data in your vector database in one click."
+          available={true}
+          linkTo={paths.tools.resetTool(organization)}
         />
       </div>
     </div>

frontend/src/utils/paths.ts

@@ -43,6 +43,9 @@ const paths = {
     resetTool: function ({ slug }: { slug: string }) {
       return `/dashboard/${slug}/tools/db-reset`;
     },
+    ragDrift: function ({ slug }: { slug: string }) {
+      return `/dashboard/${slug}/tools/rag-drift`;
+    },
   },
   dashboard: function () {
     return '/dashboard';