feat: Add support for JSON file uploads

n4ze3m · Sep 5, 2024 · 6bc983b · 6bc983b
1 parent 1c3003e
commit 6bc983b
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 3 deletions.
diff --git a/app/ui/src/components/Common/BotForm.tsx b/app/ui/src/components/Common/BotForm.tsx
@@ -117,7 +117,7 @@ export const BotForm = ({
             }}
           >
             <Upload.Dragger
-              accept={`.pdf,.docx,.csv,.txt,.mp3,.mp4,.zip`}
+              accept={`.pdf,.docx,.csv,.txt,.mp3,.mp4,.zip,.json`}
               multiple={true}
               maxCount={botConfig?.fileUploadSizeLimit}
               beforeUpload={(file) => {
@@ -132,6 +132,7 @@ export const BotForm = ({
                   "video/mpeg",
                   "application/zip",
                   "application/x-zip-compressed",
+                  "application/json",
                 ]
                   .map((type) => type.toLowerCase())
                   .join(", ");
@@ -163,8 +164,8 @@ export const BotForm = ({
                   <InboxIcon className="h-10 w-10 text-gray-400" />
                 </p>
                 <p className="ant-upload-text">
-                  Click or drag PDF, Docx, CSV , TXT, MP3, MP4, Zip files to
-                  this
+                  Click or drag PDF, Docx, CSV, TXT, MP3, MP4, Zip, or JSON
+                  files to this area
                 </p>
                 <p className="ant-upload-hint">
                   {`Support is available for a single or bulk upload of up to ${botConfig?.fileUploadSizeLimit}

diff --git a/server/src/queue/controllers/json.controller.ts b/server/src/queue/controllers/json.controller.ts
@@ -0,0 +1,55 @@
+import { QSource } from "../type";
+import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
+import { DialoqbaseVectorStore } from "../../utils/store";
+import { embeddings } from "../../utils/embeddings";
+import { PrismaClient } from "@prisma/client";
+import { getModelInfo } from "../../utils/get-model-info";
+import * as fs from "fs/promises"
+
+export const jsonQueueController = async (
+    source: QSource,
+    prisma: PrismaClient
+) => {
+    console.log("loading json");
+
+    const location = source.location!;
+
+    const json = await fs.readFile(location, "utf-8");
+
+    const textSplitter = new RecursiveCharacterTextSplitter({
+        chunkSize: source.chunkSize,
+        chunkOverlap: source.chunkOverlap,
+    });
+
+    const chunks = await textSplitter.splitDocuments([
+        {
+            pageContent: json,
+            metadata: {
+                source: location,
+            },
+        },
+    ]);
+
+    const embeddingInfo = await getModelInfo({
+        model: source.embedding,
+        prisma,
+        type: "embedding",
+    });
+
+    if (!embeddingInfo) {
+        throw new Error("Embedding not found. Please verify the embedding id");
+    }
+
+    await DialoqbaseVectorStore.fromDocuments(
+        chunks,
+        embeddings(
+            embeddingInfo.model_provider!.toLowerCase(),
+            embeddingInfo.model_id,
+            embeddingInfo?.config
+        ),
+        {
+            botId: source.botId,
+            sourceId: source.id,
+        }
+    );
+};
diff --git a/server/src/queue/controllers/zip.controller.ts b/server/src/queue/controllers/zip.controller.ts
@@ -26,6 +26,8 @@ function getMimeType(filename: string): string {
             return "audio/mpeg";
         case ".zip":
             return "application/zip";
+        case ".json":
+            return "application/json";
         default:
             return "none";
     }

diff --git a/server/src/queue/index.ts b/server/src/queue/index.ts
@@ -16,6 +16,7 @@ import { sitemapQueueController } from "./controllers/sitemap.controller";
 import { SandboxedJob } from "bullmq";
 import { getRagSettings } from "../utils/rag-settings";
 import { zipQueueController } from "./controllers/zip.controller";
+import { jsonQueueController } from "./controllers/json.controller";
 
 const prisma = new PrismaClient();
 
@@ -84,6 +85,9 @@ export default async function queueHandler(job: SandboxedJob) {
           case  "zip":
             await zipQueueController(source, prisma);
             break;
+          case "json":
+            await jsonQueueController(source, prisma);
+            break;
           default:
             break;
         }

diff --git a/server/src/utils/fileType.ts b/server/src/utils/fileType.ts
@@ -20,6 +20,8 @@ export const fileTypeFinder = (mimeType: string) => {
       return "zip";
     case "application/x-zip-compressed":
       return "zip";
+    case "application/json":
+      return "json";
     default:
       return "none";
   }