added image support in the backend

This commit is contained in:
Dhravya 2024-06-16 18:17:54 -05:00
parent c5b31e5435
commit 9588768b70
4 changed files with 90 additions and 9 deletions

View file

@ -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],
},
);

View file

@ -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(

View file

@ -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;

View file

@ -1,6 +1,7 @@
{
"compilerOptions": {
"lib": ["ES2020"],
"types": ["@cloudflare/workers-types"]
"types": ["@cloudflare/workers-types"],
"downlevelIteration": true
}
}