Merge pull request #51 from atomly/feature/document-embedding-task

feat(document): embedding task table

Author: rmolinamir

apps/drizzle-studio/package.json

@@ -0,0 +1,14 @@
+{
+	"name": "@acme/drizzle-studio",
+	"version": "0.0.0",
+	"private": true,
+	"type": "module",
+	"scripts": {
+		"dev": "NODE_ENV=development pnpm with-env pnpm --filter=db studio",
+		"with-env": "dotenv -e ../../.env --"
+	},
+	"devDependencies": {
+		"@acme/db": "workspace:*",
+		"dotenv-cli": "^8.0.0"
+	}
+}

apps/drizzle-studio/turbo.json

@@ -0,0 +1,9 @@
+{
+	"$schema": "https://turborepo.org/schema.json",
+	"extends": ["//"],
+	"tasks": {
+		"dev": {
+			"persistent": true
+		}
+	}
+}

apps/localtunnel/src/middleware/use-nextjs-proxy.ts

@@ -22,42 +22,18 @@ import { signPayload } from "../lib/sign-payload.ts";
 export function useNextjsProxy(app: express.Application) {
 	app.use(async (req, res) => {
 		try {
-			const { table } = req.body;
 			const payload = JSON.stringify(req.body);
 			const signature = signPayload(payload, env.SUPABASE_SECRET);
-			let webhookEndpoint = `${env.NEXT_JS_URL}/api/webhooks/journal-entries`;
-			switch (table) {
-				case "journal_entry": {
-					webhookEndpoint = `${env.NEXT_JS_URL}/api/webhooks/journal-entries`;
-					break;
-				}
-				case "page": {
-					webhookEndpoint = `${env.NEXT_JS_URL}/api/webhooks/page`;
-					break;
-				}
-				default:
-					break;
-			}
-			// Forward the request to Next.js server
-			const response = await fetch(webhookEndpoint, {
+			const webhookEndpoint = new URL(req.url, env.NEXT_JS_URL);
+			// Proxy the request to Next.js server
+			void fetch(webhookEndpoint, {
 				body: payload,
 				headers: {
 					"Content-Type": "application/json",
 					"x-supabase-signature": signature,
 				},
 				method: req.method,
 			});
-
-			// Forward the response back to the client
-			const responseBody = await response.text();
-			res.status(response.status);
-
-			// Forward response headers
-			response.headers.forEach((value, key) => {
-				res.setHeader(key, value);
-			});
-
-			return res.send(responseBody);
 		} catch (error) {
 			console.error("Error proxying webhook:", error);
 			res.status(500).json({ error: "Failed to proxy webhook" });

apps/web/src/app/(app)/pages/_components/page-editor.tsx

@@ -25,6 +25,9 @@ export function PageEditor({
 
 	const { mutate, isPending } = useMutation({
 		...trpc.blocks.saveTransactions.mutationOptions({}),
+		// ! TODO: When the mutation fails we need to revert the changes to the editor just like Notion does.
+		// ! To do this we can use `onError` and `editor.undo()`, without calling the transactions. We might have to get creative.
+		// ! Maybe we can refetch the blocks after an error instead of `undo`?
 		onSuccess: () => {
 			if (pendingChangesRef.current.length > 0) {
 				debouncedMutate();

apps/web/src/app/api/webhooks/journal-entries/route.ts renamed to apps/web/src/app/api/webhooks/journal_entry/route.ts

@@ -1,7 +1,8 @@
-import { and, eq } from "@acme/db";
+import { and, eq, sql } from "@acme/db";
 import {
 	BlockEdge,
 	BlockNode,
+	DocumentEmbeddingTask,
 	zInsertBlockEdge,
 	zInsertBlockNode,
 } from "@acme/db/schema";
@@ -95,6 +96,21 @@ export const blocksRouter = {
 						});
 					}
 				}
+
+				// Anytime we save a transaction we need to update the document embedding task status to debounced.
+				await tx
+					.insert(DocumentEmbeddingTask)
+					.values({
+						document_id: input.document_id,
+						user_id: ctx.session.user.id,
+					})
+					.onConflictDoUpdate({
+						set: {
+							status: "debounced",
+						},
+						target: [DocumentEmbeddingTask.document_id],
+						targetWhere: sql`${DocumentEmbeddingTask.status} != 'completed'`,
+					});
 			});
 		}),
 };

packages/api/src/router/blocks.ts

@@ -22,6 +22,7 @@
 		"check": "biome check",
 		"clean": "git clean -xdf .cache .turbo dist node_modules",
 		"dev": "pnpm build --watch",
+		"prepush": "pnpm build",
 		"push": "pnpm with-env drizzle-kit push",
 		"studio": "pnpm with-env drizzle-kit studio",
 		"typecheck": "tsc --noEmit --emitDeclarationOnly false",

packages/db/package.json

@@ -0,0 +1,69 @@
+import { sql } from "drizzle-orm";
+import { pgEnum, pgTable, text, uniqueIndex } from "drizzle-orm/pg-core";
+import { createInsertSchema, createSelectSchema } from "drizzle-zod";
+import { z } from "zod/v4";
+import { user } from "../auth/user.schema.js";
+import { Document } from "./document.schema.js";
+
+export const DocumentEmbeddingTaskStatus = pgEnum(
+	"document_embedding_task_status",
+	["debounced", "ready", "running", "completed", "failed"],
+);
+
+export const DocumentEmbeddingTask = pgTable(
+	"document_embedding_task",
+	(t) => ({
+		id: t.uuid().notNull().primaryKey().defaultRandom(),
+		user_id: text()
+			.notNull()
+			.references(() => user.id, { onDelete: "cascade" }),
+		document_id: t
+			.uuid()
+			.notNull()
+			.references(() => Document.id, { onDelete: "cascade" }),
+		status: DocumentEmbeddingTaskStatus().notNull().default("debounced"),
+		created_at: t
+			.timestamp({ mode: "string", withTimezone: true })
+			.defaultNow()
+			.notNull(),
+		updated_at: t
+			.timestamp({ mode: "string", withTimezone: true })
+			.defaultNow()
+			.notNull()
+			.$onUpdateFn(() => sql`now()`),
+	}),
+	(t) => [
+		// Unique constraint: only one non-completed task per page
+		// This allows multiple completed tasks but prevents multiple active tasks
+		uniqueIndex("unique_non_completed_task_per_page")
+			.on(t.document_id)
+			.where(sql`${t.status} != 'completed'`),
+	],
+);
+
+export type DocumentEmbeddingTask = typeof DocumentEmbeddingTask.$inferSelect;
+
+export const zInsertDocumentEmbeddingTask = createInsertSchema(
+	DocumentEmbeddingTask,
+	{
+		document_id: z.uuid(),
+	},
+).omit({
+	created_at: true,
+	id: true,
+	updated_at: true,
+});
+
+export const zUpdateDocumentEmbeddingTask = createInsertSchema(
+	DocumentEmbeddingTask,
+	{
+		document_id: z.uuid(),
+	},
+).omit({
+	created_at: true,
+	id: true,
+	updated_at: true,
+	user_id: true,
+});
+
+export const zDocumentEmbeddingTask = createSelectSchema(DocumentEmbeddingTask);

packages/db/src/core/document-embedding-task.schema.ts

@@ -4,6 +4,8 @@ import { createSelectSchema } from "drizzle-zod";
 import { user } from "../auth/user.schema.js";
 import { Page } from "./page.schema.js";
 
+const EMBEDDING_DIMENSIONS = 1536;
+
 export const PageEmbedding = pgTable(
 	"page_embedding",
 	(t) => ({
@@ -23,7 +25,7 @@ export const PageEmbedding = pgTable(
 		// Hash of the entire page's text content for change detection
 		page_text_hash: t.text().notNull(),
 		// The actual embedding vector for this chunk
-		embedding: vector({ dimensions: 1536 }).notNull(),
+		embedding: vector({ dimensions: EMBEDDING_DIMENSIONS }).notNull(),
 		created_at: t
 			.timestamp({ mode: "string", withTimezone: true })
 			.defaultNow()

packages/db/src/core/page-embedding.schema.ts

@@ -25,11 +25,6 @@ export const Page = pgTable("page", (t) => ({
 		.defaultNow()
 		.notNull()
 		.$onUpdateFn(() => sql`now()`),
-	// Separate timestamp for embedding triggers (debounced longer)
-	embed_updated_at: t
-		.timestamp({ mode: "string", withTimezone: true })
-		.defaultNow()
-		.notNull(),
 }));
 
 export type Page = typeof Page.$inferSelect;