feat(embedding): new embedding process for documents

apps/drizzle-studio/package.json

@@ -1,14 +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"
-   }
+  "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/web/next.config.js

@@ -10,7 +10,13 @@ const config = {
   /** We already do linting and typechecking as separate tasks in CI */
   eslint: { ignoreDuringBuilds: true },
   /** These packages won't be bundled in the server build */
-  serverExternalPackages: ["@mastra/*"],
+  /** @see https://nextjs.org/docs/app/api-reference/config/next-config-js/serverExternalPackages */
+  serverExternalPackages: [
+    "@mastra/*",
+    "@blocknote/server-util",
+    "@blocknote/react",
+    "@blocknote/core",
+  ],
   /** Enables hot reloading for local packages without a build step */
   transpilePackages: ["@acme/api", "@acme/auth", "@acme/db"],
   typescript: { ignoreBuildErrors: true },

apps/web/package.json

@@ -25,13 +25,15 @@
     "@blocknote/core": "^0.35.0",
     "@blocknote/mantine": "^0.35.0",
     "@blocknote/react": "^0.35.0",
+    "@blocknote/server-util": "^0.35.0",
     "@daveyplate/better-auth-ui": "^2.1.0",
     "@hookform/resolvers": "^5.0.1",
     "@mantine/core": "^8.1.3",
     "@mastra/core": "^0.12.1",
     "@mastra/libsql": "^0.12.0",
     "@mastra/loggers": "^0.10.5",
     "@mastra/memory": "^0.12.0",
+    "@mastra/rag": "^1.1.0",
     "@ngrok/ngrok": "^1.5.1",
     "@radix-ui/react-avatar": "^1.1.10",
     "@radix-ui/react-dialog": "^1.1.14",
@@ -61,6 +63,7 @@
     "react-hook-form": "^7.57.0",
     "react-virtuoso": "^4.13.0",
     "remark-gfm": "^4.0.1",
+    "remove-markdown": "^0.6.2",
     "sonner": "^2.0.5",
     "superjson": "2.2.2",
     "tailwind-merge": "^3.3.0",

apps/web/src/ai/mastra/agents/journl-agent.ts

@@ -1,5 +1,5 @@
 import { Agent } from "@mastra/core/agent";
-import { model } from "~/ai/providers/openai/llm";
+import { model } from "~/ai/providers/openai/text";
 import { semanticJournalSearch } from "../tools/semantic-journal-search";
 import { semanticPageSearch } from "../tools/semantic-page-search";
 import { temporalJournalSearch } from "../tools/temporal-journal-search";

apps/web/src/ai/mastra/tools/semantic-journal-search.ts

@@ -15,8 +15,11 @@ export const semanticJournalSearch = createTool({
     });
 
     return results.map((result) => ({
-      ...result,
-      link: `/journal/${result.date}`,
+      content: result.embedding.chunk_markdown_text,
+      date: result.journal_entry.date,
+      id: result.journal_entry.id,
+      link: `/journal/${result.journal_entry.date}`,
+      similarity: result.similarity,
     }));
   },
   id: "semantic-journal-search",

apps/web/src/ai/mastra/tools/semantic-page-search.ts

@@ -15,8 +15,11 @@ export const semanticPageSearch = createTool({
     });
 
     return result.map((result) => ({
-      ...result,
-      link: `/pages/${result.page_id}`,
+      content: result.embedding.chunk_markdown_text,
+      link: `/pages/${result.page.id}`,
+      page_id: result.page.id,
+      page_title: result.page.title,
+      similarity: result.similarity,
     }));
   },
   id: "semantic-page-search",
@@ -31,6 +34,7 @@ export const semanticPageSearch = createTool({
   outputSchema: z.array(
     z.object({
       content: z.string(),
+      link: z.string(),
       page_id: z.string(),
       page_title: z.string(),
       similarity: z.number(),

apps/web/src/ai/providers/openai/embedding.ts

@@ -1,4 +1,4 @@
 import { openai } from "@ai-sdk/openai";
 
-// TODO: Move this to a shared package called `@acme/ai`.
+// ! TODO: Move this to a shared package called `@acme/ai`.
 export const model = openai.embedding("text-embedding-3-small");

apps/web/src/ai/providers/openai/llm.ts → apps/web/src/ai/providers/openai/text.ts

@@ -1,4 +1,4 @@
 import { openai } from "@ai-sdk/openai";
 
-// TODO: Move this to a shared package called `@acme/ai`.
+// ! TODO: Move this to a shared package called `@acme/ai`.
 export const model = openai("gpt-4o-mini");

apps/web/src/app/(app)/@header/_components/header-search-modal.tsx

@@ -4,6 +4,7 @@ import { X } from "lucide-react";
 import Link from "next/link";
 import { useRouter } from "next/navigation";
 import { useEffect, useState } from "react";
+import { useDebounce } from "use-debounce";
 import { Badge } from "~/components/ui/badge";
 import { Button } from "~/components/ui/button";
 import {
@@ -30,22 +31,27 @@ import { useTRPC } from "~/trpc/react";
 const MIN_QUERY_LENGTH = 2;
 const DEFAULT_THRESHOLD = 0.25;
 const DEFAULT_LIMIT = 10;
+const DEFAULT_DEBOUNCE_TIME = 500;
+
 type HeaderSearchButtonProps = React.ComponentProps<typeof Dialog> & {
   children: React.ReactNode;
   limit?: number;
   threshold?: number;
+  debounceTime?: number;
 };
 
 export function HeaderSearchButton({
   children,
   limit = DEFAULT_LIMIT,
   threshold = DEFAULT_THRESHOLD,
+  debounceTime = DEFAULT_DEBOUNCE_TIME,
   ...rest
 }: HeaderSearchButtonProps) {
   const isMobile = useIsMobile();
   const [isOpen, setIsOpen] = useState(false);
   const [query, setQuery] = useState("");
   const router = useRouter();
+  const [debouncedQuery] = useDebounce(query, debounceTime);
 
   // Keyboard shortcut handlers
   useEffect(() => {
@@ -65,13 +71,17 @@ export function HeaderSearchButton({
   }, [isOpen]);
 
   const trpc = useTRPC();
-  const { data: notes, isLoading } = useQuery({
+  const {
+    data: notes,
+    isLoading,
+    isError,
+  } = useQuery({
     ...trpc.notes.getSimilarNotes.queryOptions({
       limit,
-      query,
+      query: debouncedQuery,
       threshold,
     }),
-    enabled: query.length > MIN_QUERY_LENGTH,
+    enabled: debouncedQuery.length > MIN_QUERY_LENGTH,
   });
 
   return (
@@ -97,7 +107,6 @@ export function HeaderSearchButton({
                 autoCorrect="off"
                 spellCheck={false}
                 autoFocus
-                value={query}
                 onValueChange={(value) => setQuery(value)}
               />
             </div>
@@ -113,7 +122,9 @@ export function HeaderSearchButton({
                   </div>
                 ) : !notes ? (
                   <span className="text-muted-foreground text-sm">
-                    Write something to search
+                    {isError
+                      ? "Something went wrong. Please try again later."
+                      : "Write something to search."}
                   </span>
                 ) : (
                   <span className="text-muted-foreground text-sm">

apps/web/src/app/(app)/journal/_components/journal-entry.tsx

@@ -1,9 +1,9 @@
 "use client";
 
+import type { PlaceholderJournalEntry } from "@acme/api";
 import type { JournalEntry } from "@acme/db/schema";
 import { useMutation } from "@tanstack/react-query";
 import Link from "next/link";
-import type { PlaceholderJournalEntry } from "node_modules/@acme/api/src/router/journal";
 import type React from "react";
 import {
   type ComponentProps,

apps/web/src/app/_components/app-providers.tsx

@@ -6,7 +6,7 @@ import { ThreadRuntime } from "~/components/ai/thread-runtime";
 export function AppProviders({ children }: { children: React.ReactNode }) {
   return (
     <MantineProvider>
-      <ThreadRuntime api="/api/chat" initialMessages={[]}>
+      <ThreadRuntime api="/api/ai/journl-agent" initialMessages={[]}>
         {children}
       </ThreadRuntime>
     </MantineProvider>

apps/web/src/app/api/supabase/embed-document/route.ts

@@ -0,0 +1,157 @@
+import {
+  zDocumentEmbeddingTask,
+  type zInsertDocumentEmbedding,
+} from "@acme/db/schema";
+import { ServerBlockNoteEditor } from "@blocknote/server-util";
+import { type ChunkParams, MDocument } from "@mastra/rag";
+import { embedMany } from "ai";
+import { NextResponse } from "next/server";
+import removeMarkdown from "remove-markdown";
+import type { z } from "zod/v4";
+import { model } from "~/ai/providers/openai/embedding";
+import { schema } from "~/components/editor/block-schema";
+import { embedder } from "~/trpc/server";
+import { handler } from "../_lib/webhook-handler";
+
+const CHUNK_PARAMS: ChunkParams = {
+  extract: {
+    keywords: true,
+    summary: true,
+    title: true,
+  },
+  strategy: "semantic-markdown",
+};
+const REMOVE_MARKDOWN_PARAMS: Parameters<typeof removeMarkdown>[1] = {
+  /* GitHub-Flavored Markdown */
+  gfm: true,
+  /* Char to insert instead of stripped list leaders */
+  listUnicodeChar: "",
+  /* Strip list leaders */
+  stripListLeaders: true,
+  /* Replace images with alt-text, if present */
+  useImgAltText: true,
+};
+
+/**
+ * This webhook will embed a document when the task is marked as ready by the Supabase Cronjob.
+ *
+ * @privateRemarks
+ *
+ * There are three jobs that are responsible for the lifecycle of the document embedding task:
+ *
+ * 1. `document-embedding-scheduler`: This job will mark debounced tasks as ready after 2 minutes without updates.
+ * ```sql
+ * -- Debounced → Ready (after 2 minutes)
+ * UPDATE document_embedding_task
+ * SET
+ *   status = 'ready',
+ *   updated_at = NOW()
+ * WHERE
+ *   status = 'debounced'
+ *   AND updated_at < NOW() - INTERVAL '2 minutes';
+ * ```
+ *
+ * 2. `document-embedding-retrier`: This job will retry failed tasks that have failed less than 3 times every 5 minutes.
+ * ```sql
+ * -- Failed → Debounced (with retry increment, max 3 retries)
+ * UPDATE document_embedding_task
+ * SET
+ *   status = 'debounced',
+ *   retries = retries + 1,
+ *   updated_at = NOW()
+ * WHERE
+ *   status = 'failed'
+ *   AND retries < 3
+ *   AND updated_at < NOW() - INTERVAL '5 minutes';
+ * ```
+ *
+ * 3. `document-embedding-sentinel`: This job picks up tasks that been ready for more than 15 minutes without updates.
+ * ```sql
+ * -- Stuck Ready → Failed (after 15 minutes)
+ * UPDATE document_embedding_task
+ * SET
+ *   status = 'failed',
+ *   updated_at = NOW()
+ * WHERE
+ *   status = 'ready'
+ *   AND updated_at < NOW() - INTERVAL '15 minutes';
+ * ```
+ */
+// ! TODO: Track embeddings token usage.
+export const POST = handler(zDocumentEmbeddingTask, async (payload) => {
+  if (payload.type === "DELETE" || payload.record.status !== "ready") {
+    return NextResponse.json({ success: true });
+  }
+
+  try {
+    const document = await embedder.document.getById({
+      id: payload.record.document_id,
+      user_id: payload.record.user_id,
+    });
+
+    if (!document?.tree) {
+      throw new Error("Document not found");
+    }
+
+    const editor = ServerBlockNoteEditor.create({
+      schema,
+    });
+
+    const markdown = await editor.blocksToMarkdownLossy(document.tree);
+
+    const mDocument =
+      await MDocument.fromMarkdown(markdown).chunk(CHUNK_PARAMS);
+
+    // ! TODO: Here's where we get the usage tokens, we must process these in a different transaction
+    // ! To guarantee that the usage is tracked regardless of the success of the embedding.
+    const { embeddings, usage: _usage } = await embedMany({
+      maxRetries: 5,
+      model,
+      values: mDocument.map((chunk) => chunk.text),
+    });
+
+    const insertions: z.infer<typeof zInsertDocumentEmbedding>[] = [];
+
+    for (const [index, embedding] of embeddings.entries()) {
+      const chunk = mDocument.at(index);
+
+      if (!chunk) {
+        continue;
+      }
+
+      console.debug("DocumentEmbedding 👀", {
+        chunk,
+      });
+
+      insertions.push({
+        chunk_id: index,
+        chunk_markdown_text: chunk.text,
+        chunk_raw_text: removeMarkdown(chunk.text, REMOVE_MARKDOWN_PARAMS),
+        document_id: document.id,
+        metadata: {
+          documentTitle: chunk.metadata.documentTitle,
+          excerptKeywords: chunk.metadata.excerptKeywords,
+          sectionSummary: chunk.metadata.sectionSummary,
+        },
+        user_id: document.user_id,
+        vector: embedding,
+      });
+    }
+
+    await embedder.documentEmbedding.embedDocument({
+      document_id: payload.record.document_id,
+      embeddings: insertions,
+      task_id: payload.record.id,
+      user_id: payload.record.user_id,
+    });
+  } catch (error) {
+    console.error("Error embedding document 👀", error);
+
+    await embedder.documentEmbeddingTask.updateStatus({
+      id: payload.record.id,
+      status: "failed",
+    });
+  }
+
+  return NextResponse.json({ success: true });
+});

apps/web/src/app/api/trpc/[trpc]/route.ts

@@ -1,4 +1,4 @@
-import { appRouter, createTRPCContext } from "@acme/api";
+import { apiRouter, createTRPCContext } from "@acme/api";
 import { fetchRequestHandler } from "@trpc/server/adapters/fetch";
 import type { NextRequest } from "next/server";
 
@@ -35,7 +35,7 @@ const handler = async (req: NextRequest) => {
       console.error(`>>> tRPC Error on '${path}'`, error);
     },
     req,
-    router: appRouter,
+    router: apiRouter,
   });
 
   setCorsHeaders(response);