AI find snippets when generating prompt

apps/web/__tests__/ai-find-snippets.test.ts

@@ -0,0 +1,91 @@
+import { describe, expect, test, vi } from "vitest";
+import { aiFindSnippets } from "@/utils/ai/snippets/find-snippets";
+import type { EmailForLLM } from "@/utils/ai/choose-rule/stringify-email";
+
+// pnpm test ai-find-snippets
+
+vi.mock("server-only", () => ({}));
+
+describe("aiFindSnippets", () => {
+  test("should find snippets in similar emails", async () => {
+    const emails = [
+      getEmail({
+        content:
+          "You can schedule a meeting with me here: https://cal.com/john-smith",
+      }),
+      getEmail({
+        content:
+          "Let's find a time to discuss. You can book a slot at https://cal.com/john-smith",
+      }),
+      getEmail({
+        content:
+          "Thanks for reaching out. Feel free to schedule a meeting at https://cal.com/john-smith",
+      }),
+    ];
+
+    const result = await aiFindSnippets({
+      sentEmails: emails,
+      user: getUser(),
+    });
+
+    expect(result.snippets).toHaveLength(1);
+    expect(result.snippets[0]).toMatchObject({
+      text: expect.stringContaining("cal.com/john-smith"),
+      count: 3,
+    });
+
+    console.log("Returned snippet:");
+    console.log(result.snippets[0]);
+  });
+
+  test("should return empty array for unique emails", async () => {
+    const emails = [
+      getEmail({
+        content:
+          "Hi Sarah, Thanks for the update on Project Alpha. I've reviewed the latest metrics and everything looks on track. Could you share the Q2 projections when you have a moment? Best, Alex",
+      }),
+      getEmail({
+        content:
+          "Just wanted to follow up on the marketing campaign results. The conversion rates are looking promising, but we should discuss optimizing the landing page. Let me know when you're free to chat. Thanks, Alex",
+      }),
+      getEmail({
+        content:
+          "Thanks for looping me in on the client feedback. I'll review the suggestions and share my thoughts during tomorrow's standup. Looking forward to moving this forward. Best regards, Alex",
+      }),
+    ];
+
+    const result = await aiFindSnippets({
+      sentEmails: emails,
+      user: getUser(),
+    });
+
+    expect(result.snippets).toHaveLength(0);
+  });
+});
+
+// helpers
+function getEmail({
+  from = "user@test.com",
+  subject = "Test Subject",
+  content = "Test content",
+  replyTo,
+  cc,
+}: Partial<EmailForLLM> = {}): EmailForLLM {
+  return {
+    from,
+    subject,
+    content,
+    ...(replyTo && { replyTo }),
+    ...(cc && { cc }),
+  };
+}
+
+function getUser() {
+  return {
+    aiModel: null,
+    aiProvider: null,
+    email: "user@test.com",
+    aiApiKey: null,
+    about: null,
+  };
+}

apps/web/app/api/user/bulk-archive/route.ts

@@ -3,7 +3,12 @@ import { NextResponse } from "next/server";
 import type { gmail_v1 } from "@googleapis/gmail";
 import { auth } from "@/app/api/auth/[...nextauth]/auth";
 import { getGmailClient } from "@/utils/gmail/client";
-import { INBOX_LABEL_ID, getOrCreateInboxZeroLabel } from "@/utils/gmail/label";
+import {
+  INBOX_LABEL_ID,
+  getOrCreateInboxZeroLabel,
+  labelVisibility,
+  messageVisibility,
+} from "@/utils/gmail/label";
 import { sleep } from "@/utils/sleep";
 import { withError } from "@/utils/middleware";
 
@@ -24,8 +29,8 @@ async function bulkArchive(body: BulkArchiveBody, gmail: gmail_v1.Gmail) {
   const archivedLabel = await getOrCreateInboxZeroLabel({
     gmail,
     key: "archived",
-    messageListVisibility: "hide",
-    labelListVisibility: "labelHide",
+    messageListVisibility: messageVisibility.hide,
+    labelListVisibility: labelVisibility.labelHide,
   });
 
   if (!archivedLabel.id)

apps/web/utils/actions/ai-rule.ts

@@ -36,9 +36,10 @@ import { aiPromptToRules } from "@/utils/ai/rule/prompt-to-rules";
 import { aiDiffRules } from "@/utils/ai/rule/diff-rules";
 import { aiFindExistingRules } from "@/utils/ai/rule/find-existing-rules";
 import { aiGenerateRulesPrompt } from "@/utils/ai/rule/generate-rules-prompt";
-import { getLabels } from "@/utils/gmail/label";
+import { getLabelById, getLabels, labelVisibility } from "@/utils/gmail/label";
 import { withActionInstrumentation } from "@/utils/actions/middleware";
 import { createScopedLogger } from "@/utils/logger";
+import { aiFindSnippets } from "@/utils/ai/snippets/find-snippets";
 
 const logger = createScopedLogger("ai-rule");
 
@@ -766,7 +767,13 @@ export const generateRulesPromptAction = withActionInstrumentation(
 
     const user = await prisma.user.findUnique({
       where: { id: session.user.id },
-      select: { aiProvider: true, aiModel: true, aiApiKey: true, email: true },
+      select: {
+        aiProvider: true,
+        aiModel: true,
+        aiApiKey: true,
+        email: true,
+        about: true,
+      },
     });
 
     if (!user) return { error: "User not found" };
@@ -775,10 +782,25 @@ export const generateRulesPromptAction = withActionInstrumentation(
     const gmail = getGmailClient(session);
     const lastSent = await getMessages(gmail, {
       query: "in:sent",
-      maxResults: 20,
+      maxResults: 50,
     });
     const gmailLabels = await getLabels(gmail);
     const userLabels = gmailLabels?.filter((label) => label.type === "user");
+
+    const labelsWithCounts: { label: string; threadsTotal: number }[] = [];
+
+    for (const label of userLabels || []) {
+      if (!label.id) continue;
+      if (label.labelListVisibility === labelVisibility.labelHide) continue;
+      const labelById = await getLabelById({ gmail, id: label.id });
+      if (!labelById?.name) continue;
+      if (!labelById.threadsTotal) continue; // Skip labels with 0 threads
+      labelsWithCounts.push({
+        label: labelById.name,
+        threadsTotal: labelById.threadsTotal || 0,
+      });
+    }
+
     const lastSentMessages = (
       await Promise.all(
         lastSent.messages?.map(async (message) => {
@@ -799,11 +821,26 @@ export const generateRulesPromptAction = withActionInstrumentation(
       );
     });
 
+    const snippetsResult = await aiFindSnippets({
+      user,
+      sentEmails: lastSentMessages.map((message) => ({
+        from: message.headers.from,
+        replyTo: message.headers["reply-to"],
+        cc: message.headers.cc,
+        subject: message.headers.subject,
+        content: emailToContent({
+          textHtml: message.textHtml || null,
+          textPlain: message.textPlain || null,
+          snippet: message.snippet,
+        }),
+      })),
+    });
+
     const result = await aiGenerateRulesPrompt({
-      user: { ...user, email: user.email },
+      user,
       lastSentEmails,
-      userLabels:
-        userLabels?.map((label) => label.name).filter(isDefined) || [],
+      snippets: snippetsResult.snippets.map((snippet) => snippet.text),
+      userLabels: labelsWithCounts.map((label) => label.label),
     });
 
     if (isActionError(result)) return { error: result.error };

apps/web/utils/ai/choose-rule/execute.ts

@@ -2,7 +2,12 @@ import type { gmail_v1 } from "@googleapis/gmail";
 import { type EmailForAction, runActionFunction } from "@/utils/ai/actions";
 import prisma from "@/utils/prisma";
 import type { Prisma } from "@prisma/client";
-import { getOrCreateInboxZeroLabel, labelThread } from "@/utils/gmail/label";
+import {
+  getOrCreateInboxZeroLabel,
+  labelThread,
+  labelVisibility,
+  messageVisibility,
+} from "@/utils/gmail/label";
 import { ExecutedRuleStatus } from "@prisma/client";
 import { createScopedLogger } from "@/utils/logger";
 
@@ -30,8 +35,8 @@ export async function executeAct({
     const label = await getOrCreateInboxZeroLabel({
       gmail,
       key: "acted",
-      messageListVisibility: "hide",
-      labelListVisibility: "labelHide",
+      messageListVisibility: messageVisibility.hide,
+      labelListVisibility: labelVisibility.labelHide,
     });
 
     if (!label.id) return;

apps/web/utils/ai/rule/generate-rules-prompt.ts

@@ -1,52 +1,93 @@
 import { z } from "zod";
 import { chatCompletionTools } from "@/utils/llms";
 import type { UserAIFields } from "@/utils/llms/types";
+import type { User } from "@prisma/client";
+import { createScopedLogger } from "@/utils/logger";
+
+const logger = createScopedLogger("ai-generate-rules-prompt");
 
 const parameters = z.object({
   rules: z
     .array(z.string())
     .describe("List of generated rules for email management"),
 });
 
+const parametersSnippets = z.object({
+  rules: z
+    .array(
+      z.object({
+        rule: z.string().describe("The rule to apply to the email"),
+        snippet: z
+          .string()
+          .optional()
+          .describe(
+            "Optional: Include ONLY if this is a snippet-based rule. The exact snippet text this rule is based on.",
+          ),
+      }),
+    )
+    .describe("List of generated rules for email management"),
+});
+
 export async function aiGenerateRulesPrompt({
   user,
   lastSentEmails,
+  snippets,
   userLabels,
 }: {
-  user: UserAIFields & { email: string };
+  user: UserAIFields & Pick<User, "email" | "about">;
   lastSentEmails: string[];
   userLabels: string[];
+  snippets: string[];
 }): Promise<string[]> {
-  const emailSummary = lastSentEmails
-    .map((email, index) => `Email ${index + 1}:\n${email}\n`)
-    .join("\n");
-
   const labelsList = userLabels
-    ? userLabels.map((label) => `- ${label}`).join("\n")
+    ? userLabels
+        .map((label) => `<label><name>${label}</name></label>`)
+        .join("\n")
     : "No labels found";
 
-  const system =
-    "You are an AI assistant that helps people manage their emails by generating rules based on their email behavior and existing labels.";
-
-  const prompt = `
-Analyze the user's email behavior and suggest general rules for managing their inbox effectively. Here's the context:
+  const hasSnippets = snippets.length > 0;
 
-User Email: ${user.email}
+  // When using snippets, we show fewer emails to the AI to avoid overwhelming it
+  const lastSentEmailsCount = hasSnippets ? 20 : 50;
 
-Last 20 Sent Emails:
-${emailSummary}
+  const system =
+    "You are an AI assistant that helps people manage their emails by generating rules based on their email behavior and existing labels.";
 
-User's Labels:
+  const prompt = `Analyze the user's email behavior and suggest general rules for managing their inbox effectively. Here's the context:
+
+<user_email>
+${user.email}
+</user_email>
+${user.about ? `\n<about_user>\n${user.about}\n</about_user>\n` : ""}
+<last_sent_emails>
+${lastSentEmails
+  .slice(0, lastSentEmailsCount)
+  .map((email) => `<email>\n${email}\n</email>`)
+  .join("\n")}
+</last_sent_emails>
+${
+  hasSnippets
+    ? `<user_snippets>\n${snippets
+        .map((snippet) => `<snippet>\n${snippet}\n</snippet>`)
+        .join("\n")}\n</user_snippets>`
+    : ""
+}
+<user_labels>
 ${labelsList}
+</user_labels>
 
-Generate a list of email management rules that would be broadly applicable for this user based on their email behavior and existing labels. The rules should be general enough to apply to various situations, not just specific recent emails. Include actions such as labeling, archiving, forwarding, replying, and drafting responses. Here are some examples of the format and complexity of rules you can create:
+<instructions>
+Generate a list of email management rules that would be broadly applicable for this user based on their email behavior and existing labels. The rules should be general enough to apply to various situations, not just specific recent emails. Include actions such as labeling, archiving, forwarding, replying, and drafting responses.
+</instructions>
 
+<example_rules>
 * Label newsletters as "Newsletter" and archive them
 * If someone asks to schedule a meeting, send them your calendar link
 * For cold emails or unsolicited pitches, draft a polite decline response
 * Label emails related to financial matters as "Finance" and mark as important
 * Forward emails about technical issues to the support team
 * For emails from key clients or partners, label as "VIP" and keep in inbox
+</example_rules>
 
 Focus on creating rules that will help the user organize their inbox more efficiently, save time, and automate responses where appropriate. Consider the following aspects:
 
@@ -56,9 +97,15 @@ Focus on creating rules that will help the user organize their inbox more effici
 4. Forwarding specific types of emails to relevant team members
 5. Prioritizing important or urgent emails
 6. Dealing with newsletters, marketing emails, and potential spam
+${
+  hasSnippets
+    ? "7. Add a rule for each snippet. IMPORTANT: Include the full text of the snippet in your output. The output can be multiple paragraphs long when using snippets."
+    : ""
+}
+
+Your response should only include the list of general rules. Aim for 3-10 broadly applicable rules that would be useful for this user's email management.`;
 
-Your response should only include the list of general rules. Aim for 3-15 broadly applicable rules that would be useful for this user's email management.
-`;
+  logger.trace({ system, prompt });
 
   const aiResponse = await chatCompletionTools({
     userAi: user,
@@ -67,16 +114,28 @@ Your response should only include the list of general rules. Aim for 3-15 broadl
     tools: {
       generate_rules: {
         description: "Generate a list of email management rules",
-        parameters,
+        parameters: hasSnippets ? parametersSnippets : parameters,
       },
     },
-    userEmail: user.email,
+    userEmail: user.email || "",
     label: "Generate rules prompt",
   });
 
-  const parsedRules = aiResponse.toolCalls[0].args as z.infer<
-    typeof parameters
-  >;
+  const args = aiResponse.toolCalls[0].args;
+
+  logger.trace(args);
+
+  return parseRulesResponse(args, hasSnippets);
+}
+
+function parseRulesResponse(args: unknown, hasSnippets: boolean): string[] {
+  if (hasSnippets) {
+    const parsedRules = args as z.infer<typeof parametersSnippets>;
+    return parsedRules.rules.map(({ rule, snippet }) =>
+      snippet ? `${rule}\n<snippet>\n${snippet}\n</snippet>` : rule,
+    );
+  }
 
+  const parsedRules = args as z.infer<typeof parameters>;
   return parsedRules.rules;
 }