proper caching and KV

supermemoryai · Apr 13, 2024 · 93e4fd7 · 93e4fd7
1 parent fac02d3
commit 93e4fd7
Show file tree

Hide file tree

Showing 7 changed files with 44 additions and 10 deletions.
diff --git a/apps/cf-ai-backend/src/routes/add.ts b/apps/cf-ai-backend/src/routes/add.ts
@@ -20,7 +20,17 @@ export async function POST(request: Request, store: CloudflareVectorizeStore, _:
 
 	const ourID = `${body.url}-${body.user}`;
 
-	const uuid = Math.random().toString(36).substring(2, 15) + Math.random().toString(36).substring(2, 15);
+	// WHY? Because this helps us to prevent duplicate entries for the same URL and user
+	function seededRandom(seed: string) {
+		let x = [...seed].reduce((acc, cur) => acc + cur.charCodeAt(0), 0);
+		return () => {
+			x = (x * 9301 + 49297) % 233280;
+			return x / 233280;
+		};
+	}
+
+	const random = seededRandom(ourID);
+	const uuid = random().toString(36).substring(2, 15) + random().toString(36).substring(2, 15);
 
 	await env.KV.put(uuid, ourID);
 

diff --git a/apps/cf-ai-backend/src/routes/chat.ts b/apps/cf-ai-backend/src/routes/chat.ts
@@ -67,8 +67,21 @@ export async function POST(request: Request, _: CloudflareVectorizeStore, embedd
 	console.log('highscoreIds', highScoreIds);
 
 	if (sourcesOnly === 'true') {
-		const idsAsStrings = highScoreIds.map((id) => env?.KV.get(id.toString()) ?? id.toString());
-		return new Response(JSON.stringify({ ids: idsAsStrings }), { status: 200 });
+		// Try await env.KV.get(id) for each id in a Promise.all
+		const idsAsStrings = highScoreIds.map(String);
+
+		const storedContent = await Promise.all(
+			idsAsStrings.map(async (id) => {
+				const stored = await env!.KV.get(id);
+				if (stored) {
+					return stored;
+				}
+				return id;
+			}),
+		);
+
+		console.log(storedContent);
+		return new Response(JSON.stringify({ ids: storedContent }), { status: 200 });
 	}
 
 	const vec = await env!.VECTORIZE_INDEX.getByIds(highScoreIds);
@@ -103,7 +116,7 @@ export async function POST(request: Request, _: CloudflareVectorizeStore, embedd
 		history: [...defaultHistory, ...(body.chatHistory ?? [])],
 	});
 
-	const prompt = `Context:\n${preparedContext}\n\nQuestion: ${query}\nAnswer:`;
+	const prompt = `Context:\n${preparedContext ?? ''}\n\nQuestion: ${query}\nAnswer:`;
 
 	const output = await chat.sendMessageStream(prompt);
 

diff --git a/apps/cf-ai-backend/src/routes/query.ts b/apps/cf-ai-backend/src/routes/query.ts
@@ -42,8 +42,20 @@ export async function GET(request: Request, _: CloudflareVectorizeStore, embeddi
 	const highScoreIds = resp.matches.filter(({ score }) => score > 0.3).map(({ id }) => id);
 
 	if (sourcesOnly === 'true') {
-		const idsAsStrings = highScoreIds.map((id) => env?.KV.get(id.toString()) ?? id.toString());
-		return new Response(JSON.stringify({ ids: idsAsStrings }), { status: 200 });
+		const idsAsStrings = highScoreIds.map(String);
+
+		const storedContent = await Promise.all(
+			idsAsStrings.map(async (id) => {
+				const stored = await env!.KV.get(id);
+				if (stored) {
+					return stored;
+				}
+				return id;
+			}),
+		);
+
+		console.log(storedContent);
+		return new Response(JSON.stringify({ ids: storedContent }), { status: 200 });
 	}
 
 	const vec = await env!.VECTORIZE_INDEX.getByIds(highScoreIds);

diff --git a/apps/cf-ai-backend/wrangler.toml b/apps/cf-ai-backend/wrangler.toml
@@ -4,7 +4,7 @@ compatibility_date = "2024-02-23"
 
 [[vectorize]]
 binding = "VECTORIZE_INDEX"
-index_name = "any-vector"
+index_name = "supermem-vector"
 
 [ai]
 binding = "AI"

diff --git a/apps/extension/src/SideBar.tsx b/apps/extension/src/SideBar.tsx
@@ -215,7 +215,6 @@ function SideBar() {
                 <DialogTrigger asChild>
                   <button
                     onClick={() => {
-                      return;
                       sendUrlToAPI();
                       setIsSendingData(true);
                       setTimeout(() => {

diff --git a/apps/web/src/app/api/store/route.ts b/apps/web/src/app/api/store/route.ts
@@ -80,15 +80,14 @@ export async function POST(req: NextRequest) {
     storeToSpace = "none";
   }
 
-  // Count the number of stored content of the user
   const count = await db
     .select({
       count: sql<number>`count(*)`.mapWith(Number),
     })
     .from(storedContent)
     .where(eq(storedContent.user, session.user.id));
 
-  console.log(count[0].count);
+  console.log("count", count[0].count);
 
   const storedContentId = await db.insert(storedContent).values({
     content: data.pageContent,

diff --git a/apps/web/src/lib/utils.ts b/apps/web/src/lib/utils.ts
@@ -19,6 +19,7 @@ export function cleanUrl(url: string) {
 }
 
 export function getIdsFromSource(sourceIds: string[]) {
+  console.log(sourceIds);
   return sourceIds.map((id) => {
     const parts = id.split("-");
     if (parts.length > 1) {