✨ Added multi tenancy to vector dbs

panoratech · Sep 16, 2024 · a666b60 · a666b60
1 parent c9f0652
commit a666b60
Show file tree

Hide file tree

Showing 20 changed files with 386 additions and 169 deletions.
diff --git a/.env.example b/.env.example
@@ -178,19 +178,22 @@ NEXT_PUBLIC_DISTRIBUTION=selfhost # selfhost or managed
 ## pinecone
 PINECONE_API_KEY=
 PINECONE_INDEX_NAME=
-
 ## qdrant
 QDRANT_BASE_URL=
 QDRANT_API_KEY=
+QDRANT_COLLECTION_NAME=
 ## chroma
 CHROMADB_URL=
+CHROMADB_COLLECTION_NAME=
 ## weaviate
 WEAVIATE_URL=
 WEAVIATE_API_KEY=
+WEAVIATE_CLASS_NAME=
 # turbopuffer
 TURBOPUFFER_API_KEY=
 # milvus
 MILVUS_ADDRESS=
+MILVUS_COLLECTION_NAME=
 
 # EMBEDDINGS 
 JINA_API_KEY=

diff --git a/apps/webapp/src/components/Configuration/RAGSettings/RAGItemDisplay.tsx b/apps/webapp/src/components/Configuration/RAGSettings/RAGItemDisplay.tsx
@@ -24,6 +24,8 @@ const formSchema = z.object({
   url: z.string().optional(),
   indexName: z.string().optional(),
   embeddingApiKey: z.string().optional(),
+  collectionName: z.string().optional(),
+  className: z.string().optional(),
 });
 
 interface ItemDisplayProps {
@@ -183,11 +185,11 @@ export function RAGItemDisplay({ item, type }: ItemDisplayProps) {
             case 'pinecone':
                 return ['apiKey', 'indexName'];
             case 'qdrant':
-                return ['apiKey', 'baseUrl'];
+                return ['apiKey', 'baseUrl', 'collectionName'];
             case 'chromadb':
-                return ['url'];
+                return ['url', 'collectionName'];
             case 'weaviate':
-                return ['apiKey', 'url'];
+                return ['apiKey', 'url', 'className'];
             case 'openai_ada_small_1536':
             case 'openai_ada_large_3072':
             case 'openai_ada_002':
@@ -210,13 +212,16 @@ export function RAGItemDisplay({ item, type }: ItemDisplayProps) {
             case 'qdrant':
                 form.setValue("apiKey", data[0]);
                 form.setValue("baseUrl", data[1]);
+                form.setValue("collectionName", data[1]);
                 break;
             case 'chromadb':
                 form.setValue("url", data[0]);
+                form.setValue("collectionName", data[1]);
                 break;
             case 'weaviate':
                 form.setValue("apiKey", data[0]);
                 form.setValue("url", data[1]);
+                form.setValue("className", data[1]);
                 break;
             case 'openai_ada_small_1536':
             case 'openai_ada_large_3072':

diff --git a/apps/webapp/src/components/Configuration/RAGSettings/RAGItemList.tsx b/apps/webapp/src/components/Configuration/RAGSettings/RAGItemList.tsx
@@ -3,7 +3,7 @@ import { cn } from "@/lib/utils"
 import { Button } from "@/components/ui/button"
 import { ScrollArea } from "@/components/ui/scroll-area"
 import { vectorDatabases, embeddingModels } from "./utils"
-import { useRagItem } from "./useRAGItem"
+import { useRagItem } from "./useRagItem"
 
 interface RAGItemListProps {
   items: (typeof vectorDatabases[number] | typeof embeddingModels[number])[];

diff --git a/apps/webapp/src/components/Configuration/RAGSettings/RAGLayout.tsx b/apps/webapp/src/components/Configuration/RAGSettings/RAGLayout.tsx
@@ -8,7 +8,7 @@ import * as React from "react"
 import { RAGItemDisplay } from "./RAGItemDisplay"
 import { RAGItemList } from "./RAGItemList"
 import { embeddingModels, vectorDatabases } from "./utils"
-import { useRagItem } from "./useRAGItem"
+import { useRagItem } from "./useRagItem"
 
 interface Props {
     items: (typeof vectorDatabases[number] | typeof embeddingModels[number])[];

diff --git a/apps/webapp/src/components/Configuration/RAGSettings/RAGSettingsPage.tsx b/apps/webapp/src/components/Configuration/RAGSettings/RAGSettingsPage.tsx
@@ -1,5 +1,5 @@
 import * as React from "react"
-import { RAGLayout } from "./RagLayout";
+import { RAGLayout } from "./RAGLayout";
 import { embeddingModels, TabType, vectorDatabases } from "./utils";
 
 export default function RAGSettingsPage() {

diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml
@@ -205,11 +205,15 @@ services:
       PINECONE_INDEX_NAME: ${PINECONE_INDEX_NAME}
       QDRANT_BASE_URL: ${QDRANT_BASE_URL}
       QDRANT_API_KEY: ${QDRANT_API_KEY}
+      QDRANT_COLLECTION_NAME: ${QDRANT_COLLECTION_NAME}
       CHROMADB_URL: ${CHROMADB_URL}
+      CHROMADB_COLLECTION_NAME: ${CHROMADB_COLLECTION_NAME}
       WEAVIATE_URL: ${WEAVIATE_URL}
       WEAVIATE_API_KEY: ${WEAVIATE_API_KEY}
+      WEAVIATE_CLASS_NAME: ${WEAVIATE_CLASS_NAME}
       TURBOPUFFER_API_KEY: ${TURBOPUFFER_API_KEY}
       MILVUS_ADDRESS: ${MILVUS_ADDRESS}
+      MILVUS_COLLECTION_NAME: ${MILVUS_COLLECTION_NAME}
 
     restart: unless-stopped
     ports:

diff --git a/docker-compose.source.yml b/docker-compose.source.yml
@@ -205,11 +205,15 @@ services:
       PINECONE_INDEX_NAME: ${PINECONE_INDEX_NAME}
       QDRANT_BASE_URL: ${QDRANT_BASE_URL}
       QDRANT_API_KEY: ${QDRANT_API_KEY}
+      QDRANT_COLLECTION_NAME: ${QDRANT_COLLECTION_NAME}
       CHROMADB_URL: ${CHROMADB_URL}
+      CHROMADB_COLLECTION_NAME: ${CHROMADB_COLLECTION_NAME}
       WEAVIATE_URL: ${WEAVIATE_URL}
       WEAVIATE_API_KEY: ${WEAVIATE_API_KEY}
+      WEAVIATE_CLASS_NAME: ${WEAVIATE_CLASS_NAME}
       TURBOPUFFER_API_KEY: ${TURBOPUFFER_API_KEY}
       MILVUS_ADDRESS: ${MILVUS_ADDRESS}
+      MILVUS_COLLECTION_NAME: ${MILVUS_COLLECTION_NAME}
 
     restart: unless-stopped
     ports:

diff --git a/docker-compose.yml b/docker-compose.yml
@@ -199,11 +199,15 @@ services:
       PINECONE_INDEX_NAME: ${PINECONE_INDEX_NAME}
       QDRANT_BASE_URL: ${QDRANT_BASE_URL}
       QDRANT_API_KEY: ${QDRANT_API_KEY}
+      QDRANT_COLLECTION_NAME: ${QDRANT_COLLECTION_NAME}
       CHROMADB_URL: ${CHROMADB_URL}
+      CHROMADB_COLLECTION_NAME: ${CHROMADB_COLLECTION_NAME}
       WEAVIATE_URL: ${WEAVIATE_URL}
       WEAVIATE_API_KEY: ${WEAVIATE_API_KEY}
+      WEAVIATE_CLASS_NAME: ${WEAVIATE_CLASS_NAME}
       TURBOPUFFER_API_KEY: ${TURBOPUFFER_API_KEY}
       MILVUS_ADDRESS: ${MILVUS_ADDRESS}
+      MILVUS_COLLECTION_NAME: ${MILVUS_COLLECTION_NAME}
 
     restart: unless-stopped
     ports:

diff --git a/packages/api/src/@core/@core-services/environment/environment.service.ts b/packages/api/src/@core/@core-services/environment/environment.service.ts
@@ -65,13 +65,19 @@ export class EnvironmentService {
     };
   }
 
-  getChromaCreds(): string {
-    return this.configService.get<string>('CHROMADB_URL');
+  getChromaCreds() {
+    return {
+      url: this.configService.get<string>('CHROMADB_URL'),
+      collectionName: this.configService.get<string>(
+        'CHROMADB_COLLECTION_NAME',
+      ),
+    };
   }
 
   getMilvusCreds() {
     return {
       address: this.configService.get<string>('MILVUS_ADDRESS'),
+      collectionName: this.configService.get<string>('MILVUS_COLLECTION_NAME'),
     };
   }
   getPineconeCreds() {
@@ -85,6 +91,7 @@ export class EnvironmentService {
     return {
       url: this.configService.get<string>('WEAVIATE_URL'),
       apiKey: this.configService.get<string>('WEAVIATE_API_KEY'),
+      className: this.configService.get<string>('WEAVIATE_CLASS_NAME'),
     };
   }
 
@@ -96,6 +103,7 @@ export class EnvironmentService {
     return {
       baseUrl: this.configService.get<string>('QDRANT_BASE_URL'),
       apiKey: this.configService.get<string>('QDRANT_API_KEY'),
+      collectionName: this.configService.get<string>('QDRANT_COLLECTION_NAME'),
     };
   }
 

diff --git a/packages/api/src/@core/rag/document.processor.ts b/packages/api/src/@core/rag/document.processor.ts
@@ -19,9 +19,13 @@ export class ProcessDocumentProcessor {
 
   @Process('batchDocs')
   async processDocuments(
-    job: Job<{ filesInfo: FileInfo[]; projectId: string }>,
+    job: Job<{
+      filesInfo: FileInfo[];
+      projectId: string;
+      linkedUserId: string;
+    }>,
   ) {
-    const { filesInfo, projectId } = job.data;
+    const { filesInfo, projectId, linkedUserId } = job.data;
     const results = [];
 
     for (const fileInfo of filesInfo) {
@@ -40,7 +44,7 @@ export class ProcessDocumentProcessor {
         // console.log(`chunks for ${fileInfo.id} are ` + JSON.stringify(chunks));
         const embeddings = await this.embeddingService.generateEmbeddings(
           chunks,
-          projectId
+          projectId,
         );
         // Split embeddings into smaller batches
         const batchSize = 100; // Adjust this value as needed
@@ -52,6 +56,7 @@ export class ProcessDocumentProcessor {
             batchChunks,
             batchEmbeddings,
             projectId,
+            linkedUserId,
           );
         }
         results.push(`Successfully processed document ${fileInfo.id}`);

diff --git a/packages/api/src/@core/rag/rag.controller.ts b/packages/api/src/@core/rag/rag.controller.ts
@@ -1,16 +1,8 @@
-import { Controller, Post, Body, UseGuards } from '@nestjs/common';
-import { RagService } from './rag.service';
 import { ApiKeyAuthGuard } from '@@core/auth/guards/api-key.guard';
-import {
-  ApiBody,
-  ApiOperation,
-  ApiParam,
-  ApiQuery,
-  ApiTags,
-  ApiHeader,
-  //ApiKeyAuth,
-} from '@nestjs/swagger';
 import { ConnectionUtils } from '@@core/connections/@utils';
+import { Body, Controller, Headers, Post, UseGuards } from '@nestjs/common';
+import { ApiHeader } from '@nestjs/swagger';
+import { RagService } from './rag.service';
 
 @Controller('rag')
 export class RagController {
@@ -21,9 +13,26 @@ export class RagController {
   ) {}
 
   @Post('query')
+  @ApiHeader({
+    name: 'x-connection-token',
+    required: true,
+    description: 'The connection token',
+    example: 'b008e199-eda9-4629-bd41-a01b6195864a',
+  })
   @UseGuards(ApiKeyAuthGuard)
-  async queryEmbeddings(@Body() body: { query: string; topK?: number }) {
-    return this.documentEmbeddingService.queryEmbeddings(body.query, body.topK);
+  async queryEmbeddings(
+    @Body() body: { query: string; topK?: number },
+    @Headers('x-connection-token') connection_token: string,
+  ) {
+    const { linkedUserId, remoteSource, connectionId, projectId } =
+      await this.connectionUtils.getConnectionMetadataFromConnectionToken(
+        connection_token,
+      );
+    return this.documentEmbeddingService.queryEmbeddings(
+      body.query,
+      body.topK,
+      linkedUserId,
+    );
   }
 
   /*

diff --git a/packages/api/src/@core/rag/rag.service.ts b/packages/api/src/@core/rag/rag.service.ts
@@ -14,11 +14,12 @@ export class RagService {
     private s3Service: S3Service,
   ) {}
 
-  async queryEmbeddings(query: string, topK = 5) {
+  async queryEmbeddings(query: string, topK = 5, linkedUserId: string) {
     const queryEmbedding = await this.embeddingService.embedQuery(query);
     const results = await this.vectorDatabaseService.queryEmbeddings(
       queryEmbedding,
       topK,
+      linkedUserId,
     );
     return results.map((match: any) => ({
       chunk: match.metadata.text,

diff --git a/packages/api/src/@core/rag/vecdb/chromadb/chromadb.service.ts b/packages/api/src/@core/rag/vecdb/chromadb/chromadb.service.ts
@@ -1,52 +1,55 @@
 import { EnvironmentService } from '@@core/@core-services/environment/environment.service';
 import { ProcessedChunk } from '@@core/rag/types';
 import { Injectable } from '@nestjs/common';
-import { ChromaClient } from 'chromadb';
+import { ChromaClient, Collection } from 'chromadb';
 
 @Injectable()
 export class ChromaDBService {
   private client: ChromaClient;
+  private collection: Collection;
 
-  constructor(private envService: EnvironmentService) {
-    //this.initialize();
+  constructor(private envService: EnvironmentService) {}
+
+  async onModuleInit() {
+    return;
   }
 
-  async initialize() {
+  async initialize(credentials: string[]) {
     this.client = new ChromaClient({
-      path: this.envService.getChromaCreds(),
+      path: credentials[0],
+    });
+    this.collection = await this.client.getOrCreateCollection({
+      name: credentials[1],
     });
   }
 
   async storeEmbeddings(
     fileId: string,
     chunks: ProcessedChunk[],
     embeddings: number[][],
+    linkedUserId: string,
   ) {
-    const collection = await this.client.createCollection({ name: fileId });
-    await collection.add({
+    await this.collection.add({
       ids: chunks.map((_, i) => `${fileId}_${i}`),
       embeddings: embeddings,
       metadatas: chunks.map((chunk) => ({
         text: chunk.text,
         ...chunk.metadata,
+        user_id: `ns_${linkedUserId}`,
       })),
     });
   }
 
-  async queryEmbeddings(queryEmbedding: number[], topK: number) {
-    const collections = await this.client.listCollections();
-    const results = await Promise.all(
-      collections.map(async (collection) => {
-        const collectionInstance = await this.client.getCollection({
-          name: collection.name,
-        });
-        const result = await collectionInstance.query({
-          queryEmbeddings: [queryEmbedding],
-          nResults: topK,
-        });
-        return result.metadatas[0];
-      }),
-    );
-    return results.flat().slice(0, topK);
+  async queryEmbeddings(
+    queryEmbedding: number[],
+    topK: number,
+    linkedUserId: string,
+  ) {
+    const result = await this.collection.query({
+      queryEmbeddings: [queryEmbedding],
+      nResults: topK,
+      where: { user_id: `ns_${linkedUserId}` },
+    });
+    return result.metadatas[0];
   }
 }