From 9588768b70be92f9e354a350d14993540405ec61 Mon Sep 17 00:00:00 2001
From: Dhravya <dhravyashah@gmail.com>
Date: Sun, 16 Jun 2024 18:17:54 -0500
Subject: [PATCH] added image support in the backend

---
 apps/cf-ai-backend/src/helper.ts | 11 ++---
 apps/cf-ai-backend/src/index.ts  | 83 +++++++++++++++++++++++++++++++-
 apps/cf-ai-backend/src/types.ts  |  2 +-
 apps/cf-ai-backend/tsconfig.json |  3 +-
 4 files changed, 90 insertions(+), 9 deletions(-)

diff --git a/apps/cf-ai-backend/src/helper.ts b/apps/cf-ai-backend/src/helper.ts
index 5d01c881..1540c9fd 100644
--- a/apps/cf-ai-backend/src/helper.ts
+++ b/apps/cf-ai-backend/src/helper.ts
@@ -103,14 +103,13 @@ export async function batchCreateChunksAndEmbeddings({
   await deleteDocument({ url: body.url, user: body.user, c: context, store });
 
   const random = seededRandom(ourID);
+  const uuid =
+    random().toString(36).substring(2, 15) +
+    random().toString(36).substring(2, 15);
 
   for (let i = 0; i < chunks.length; i++) {
     const chunk = chunks[i];
-    const uuid =
-      random().toString(36).substring(2, 15) +
-      random().toString(36).substring(2, 15) +
-      "-" +
-      i;
+    const chunkId = `${uuid}-${i}`;
 
     const newPageContent = `Title: ${body.title}\nDescription: ${body.description}\nURL: ${body.url}\nContent: ${chunk}`;
 
@@ -129,7 +128,7 @@ export async function batchCreateChunksAndEmbeddings({
         },
       ],
       {
-        ids: [uuid],
+        ids: [chunkId],
       },
     );
 
diff --git a/apps/cf-ai-backend/src/index.ts b/apps/cf-ai-backend/src/index.ts
index 7d3f69e6..94351b70 100644
--- a/apps/cf-ai-backend/src/index.ts
+++ b/apps/cf-ai-backend/src/index.ts
@@ -18,7 +18,12 @@ import { swaggerUI } from "@hono/swagger-ui";
 
 const app = new Hono<{ Bindings: Env }>();
 
-app.get("/doc", swaggerUI({ url: "/doc" }));
+app.get(
+  "/ui",
+  swaggerUI({
+    url: "/doc",
+  }),
+);
 
 // ------- MIDDLEWARES -------
 app.use("*", poweredBy());
@@ -34,6 +39,17 @@ app.use("/api/", async (c, next) => {
 });
 // ------- MIDDLEWARES END -------
 
+const fileSchema = z
+  .instanceof(File)
+  .refine(
+    (file) => file.size <= 10 * 1024 * 1024,
+    "File size should be less than 10MB",
+  ) // Validate file size
+  .refine(
+    (file) => ["image/jpeg", "image/png", "image/gif"].includes(file.type),
+    "Invalid file type",
+  ); // Validate file type
+
 app.get("/", (c) => {
   return c.text("Supermemory backend API is running!");
 });
@@ -57,6 +73,71 @@ app.post("/api/add", zValidator("json", vectorObj), async (c) => {
   return c.json({ status: "ok" });
 });
 
+app.post(
+  "/api/add-with-image",
+  zValidator(
+    "form",
+    z.object({
+      images: z
+        .array(fileSchema)
+        .min(1, "At least one image is required")
+        .optional(),
+      "images[]": z
+        .array(fileSchema)
+        .min(1, "At least one image is required")
+        .optional(),
+      text: z.string().optional(),
+      space: z.string().optional(),
+      url: z.string(),
+      user: z.string(),
+    }),
+    (c) => {
+      console.log(c);
+    },
+  ),
+  async (c) => {
+    const body = c.req.valid("form");
+
+    const { store } = await initQuery(c);
+
+    if (!(body.images || body["images[]"])) {
+      return c.json({ status: "error", message: "No images found" }, 400);
+    }
+
+    const imagePromises = (body.images ?? body["images[]"]).map(
+      async (image) => {
+        const buffer = await image.arrayBuffer();
+        const input = {
+          image: [...new Uint8Array(buffer)],
+          prompt:
+            "What's in this image? caption everything you see in great detail",
+          max_tokens: 1024,
+        };
+        const response = await c.env.AI.run(
+          "@cf/llava-hf/llava-1.5-7b-hf",
+          input,
+        );
+        console.log(response.description);
+        return response.description;
+      },
+    );
+
+    const imageDescriptions = await Promise.all(imagePromises);
+
+    await batchCreateChunksAndEmbeddings({
+      store,
+      body,
+      chunks: [
+        imageDescriptions,
+        ...(body.text ? chunkText(body.text, 1536) : []),
+      ].flat(),
+      context: c,
+    });
+
+    return c.json({ status: "ok" });
+  },
+);
+
 app.get(
   "/api/ask",
   zValidator(
diff --git a/apps/cf-ai-backend/src/types.ts b/apps/cf-ai-backend/src/types.ts
index ca68d9ee..3db62553 100644
--- a/apps/cf-ai-backend/src/types.ts
+++ b/apps/cf-ai-backend/src/types.ts
@@ -2,7 +2,7 @@ import { z } from "zod";
 
 export type Env = {
   VECTORIZE_INDEX: VectorizeIndex;
-  AI: Fetcher;
+  AI: Ai;
   SECURITY_KEY: string;
   OPENAI_API_KEY: string;
   GOOGLE_AI_API_KEY: string;
diff --git a/apps/cf-ai-backend/tsconfig.json b/apps/cf-ai-backend/tsconfig.json
index 2b75d5a0..fcdf6914 100644
--- a/apps/cf-ai-backend/tsconfig.json
+++ b/apps/cf-ai-backend/tsconfig.json
@@ -1,6 +1,7 @@
 {
   "compilerOptions": {
     "lib": ["ES2020"],
-    "types": ["@cloudflare/workers-types"]
+    "types": ["@cloudflare/workers-types"],
+    "downlevelIteration": true
   }
 }