Improve choose rule prompt

apps/web/__tests__/ai-choose-rule.test.ts

@@ -354,6 +354,7 @@ function getEmail({
   content = "content",
 }: { from?: string; subject?: string; content?: string } = {}) {
   return {
+    id: "id",
     from,
     subject,
     content,

apps/web/app/api/google/webhook/process-history-item.test.ts

@@ -164,24 +164,19 @@
   });
 
   it("should skip if message is outbound", async () => {
-    vi.mocked(getThreadMessages).mockResolvedValueOnce([
-      {
-        id: "123",
-        threadId: "thread-123",
-        labelIds: [GmailLabel.SENT],
-        internalDate: "1704067200000", // 2024-01-01T00:00:00Z
-        snippet: "Hello World",
-        historyId: "12345",
-        inline: [],
-        headers: {
-          from: "user@example.com",
-          to: "recipient@example.com",
-          subject: "Test Email",
-          date: "2024-01-01T00:00:00Z",
-        },
-        textPlain: "Hello World",
+    vi.mocked(getMessage).mockResolvedValueOnce({
+      id: "123",
+      threadId: "thread-123",
+      labelIds: [GmailLabel.SENT],
+      payload: {
+        headers: [
+          { name: "From", value: "user@example.com" },
+          { name: "To", value: "recipient@example.com" },
+          { name: "Subject", value: "Test Email" },
+          { name: "Date", value: "2024-01-01T00:00:00Z" },
+        ],
       },
-    ]);
+    });
 
     await processHistoryItem(createHistoryItem(), createOptions());
 
@@ -208,7 +203,7 @@
 
     await processHistoryItem(createHistoryItem(), options);
 
     expect(runColdEmailBlocker).toHaveBeenCalledWith({
       email: expect.objectContaining({
         from: "sender@example.com",
         subject: "Test Email",

apps/web/app/api/google/webhook/process-history-item.ts

@@ -137,7 +137,7 @@ export async function processHistoryItem(
           from: message.headers.from,
           subject: message.headers.subject,
           content,
-          messageId,
+          id: messageId,
           threadId,
           date: internalDateToDate(message.internalDate),
         },

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

@@ -632,6 +632,7 @@ export const generateRulesPromptAction = withActionInstrumentation(
     const snippetsResult = await aiFindSnippets({
       user,
       sentEmails: lastSentMessages.map((message) => ({
+        id: message.id,
         from: message.headers.from,
         replyTo: message.headers["reply-to"],
         cc: message.headers.cc,
@@ -721,6 +722,7 @@ export const reportAiMistakeAction = withActionInstrumentation(
       actualRule,
       expectedRule,
       email: {
+        id: "",
         ...email,
         content,
       },

apps/web/utils/actions/cold-email.ts

@@ -125,7 +125,7 @@ async function checkColdEmail(
       content,
       date: body.date ? new Date(body.date) : undefined,
       threadId: body.threadId || undefined,
-      messageId: body.messageId,
+      id: body.messageId || "",
     },
     user,
     gmail,

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

@@ -5,59 +5,73 @@ import type { User } from "@prisma/client";
 import { stringifyEmail } from "@/utils/stringify-email";
 import type { EmailForLLM } from "@/utils/types";
 import { createScopedLogger } from "@/utils/logger";
+import { Braintrust } from "@/utils/braintrust";
 
 const logger = createScopedLogger("ai-choose-rule");
 
+const braintrust = new Braintrust("choose-rule-1");
+
 type GetAiResponseOptions = {
-  email: {
-    from: string;
-    subject: string;
-    content: string;
-    cc?: string;
-    replyTo?: string;
-  };
+  email: EmailForLLM;
   user: Pick<User, "email" | "about"> & UserAIFields;
   rules: { instructions: string }[];
 };
 
 async function getAiResponse(options: GetAiResponseOptions) {
   const { email, user, rules } = options;
 
-  const rulesWithUnknownRule = [
-    ...rules,
-    {
-      instructions:
-        "None of the other rules match or not enough information to make a decision.",
-    },
-  ];
+  const specialRuleNumber = rules.length + 1;
+
+  const emailSection = stringifyEmail(email, 500);
 
   const system = `You are an AI assistant that helps people manage their emails.
-IMPORTANT: You must strictly follow the exclusions mentioned in each rule.
-- If a rule says to exclude certain types of emails, DO NOT select that rule for those excluded emails.
-- When multiple rules match, choose the more specific one that best matches the email's content.
-- Rules about requiring replies should be prioritized when the email clearly needs a response.
-- If you're unsure, select the last rule (not enough information).
-- It's better to select "not enough information" than to make an incorrect choice.
 
-REMINDER: Pay careful attention to any exclusions mentioned in the rules. If an email matches an exclusion, that rule MUST NOT be selected.
+<instructions>
+IMPORTANT: Follow these instructions carefully when selecting a rule:
 
-These are the rules you can select from:
-${rulesWithUnknownRule
-  .map((rule, i) => `${i + 1}. ${rule.instructions}`)
-  .join("\n")}
+<priority>
+1. Match the email to a SPECIFIC user-defined rule that addresses the email's exact content or purpose.
+2. If the email doesn't match any specific rule but the user has a catch-all rule (like "emails that don't match other criteria"), use that catch-all rule.
+3. Only use rule #${specialRuleNumber} (system fallback) if no user-defined rule can reasonably apply.
+</priority>
 
-${user.about ? `Additional information about the user:\n\n${user.about}` : ""}
+<guidelines>
+- If a rule says to exclude certain types of emails, DO NOT select that rule for those excluded emails.
+- When multiple rules match, choose the more specific one that best matches the email's content.
+- Rules about requiring replies should be prioritized when the email clearly needs a response.
+- Rule #${specialRuleNumber} should ONLY be selected when there is absolutely no user-defined rule that could apply.
+</guidelines>
+</instructions>
+
+<user_rules>
+${rules.map((rule, i) => `${i + 1}. ${rule.instructions}`).join("\n")}
+</user_rules>
+
+<system_fallback>
+${specialRuleNumber}. None of the other rules match or not enough information to make a decision.
+</system_fallback>
+
+${
+  user.about
+    ? `<user_info>
+<about>${user.about}</about>
+<email>${user.email}</email>
+</user_info>`
+    : `<user_info>
+<email>${user.email}</email>
+</user_info>`
+}
 
 <outputFormat>
 Respond with a JSON object with the following fields:
 "reason" - the reason you chose that rule. Keep it concise.
 "rule" - the number of the rule you want to apply
-</outputFormat>
+</outputFormat>`;
 
-Select a rule to apply to the email:`;
+  const prompt = `Select a rule to apply to this email that was sent to me:
 
-  const prompt = `<email>
-${stringifyEmail(email, 500)}
+<email>
+${emailSection}
 </email>`;
 
   logger.trace("Input", { system, prompt });
@@ -91,8 +105,22 @@ ${stringifyEmail(email, 500)}
   });
 
   logger.trace("Response", aiResponse.object);
-  // logger.trace("Usage", aiResponse.usage);
-  // logger.trace("Provider Metadata", aiResponse.experimental_providerMetadata);
+
+  braintrust.insertToDataset({
+    id: email.id,
+    input: {
+      email: emailSection,
+      rules: rules.map((rule, i) => ({
+        ruleNumber: i + 1,
+        instructions: rule.instructions,
+      })),
+      hasAbout: !!user.about,
+      userAbout: user.about,
+      userEmail: user.email,
+      specialRuleNumber,
+    },
+    expected: aiResponse.object.rule,
+  });
 
   return aiResponse.object;
 }

apps/web/utils/ai/reply/generate-nudge.ts

@@ -2,20 +2,15 @@ import { chatCompletion } from "@/utils/llms";
 import type { UserEmailWithAI } from "@/utils/llms/types";
 import { stringifyEmail } from "@/utils/stringify-email";
 import { createScopedLogger } from "@/utils/logger";
+import type { EmailForLLM } from "@/utils/types";
 
 const logger = createScopedLogger("generate-nudge");
 
 export async function aiGenerateNudge({
   messages,
   user,
 }: {
-  messages: {
-    from: string;
-    to: string;
-    subject: string;
-    content: string;
-    date: Date;
-  }[];
+  messages: EmailForLLM[];
   user: UserEmailWithAI;
   onFinish?: (completion: string) => Promise<void>;
 }) {
@@ -32,7 +27,7 @@ ${messages
   .map(
     (msg) => `<email>
 ${stringifyEmail(msg, 3000)}
-<date>${msg.date.toISOString()}</date>
+<date>${msg.date?.toISOString() ?? "unknown"}</date>
 </email>`,
   )
   .join("\n")}

apps/web/utils/ai/reply/generate-reply.ts

@@ -2,21 +2,15 @@ import { chatCompletion } from "@/utils/llms";
 import type { UserEmailWithAI } from "@/utils/llms/types";
 import { stringifyEmail } from "@/utils/stringify-email";
 import { createScopedLogger } from "@/utils/logger";
-
+import type { EmailForLLM } from "@/utils/types";
 const logger = createScopedLogger("generate-reply");
 
 export async function aiGenerateReply({
   messages,
   user,
   instructions,
 }: {
-  messages: {
-    from: string;
-    to: string;
-    subject: string;
-    content: string;
-    date: Date;
-  }[];
+  messages: (EmailForLLM & { to: string })[];
   user: UserEmailWithAI;
   instructions: string | null;
 }) {
@@ -47,7 +41,7 @@ ${messages
   .map(
     (msg) => `<email>
 ${stringifyEmail(msg, 3000)}
-<date>${msg.date.toISOString()}</date>
+<date>${msg.date?.toISOString() ?? "unknown"}</date>
 </email>`,
   )
   .join("\n")}

apps/web/utils/cold-email/is-cold-email.ts

@@ -10,6 +10,7 @@ import { DEFAULT_COLD_EMAIL_PROMPT } from "@/utils/cold-email/prompt";
 import { stringifyEmail } from "@/utils/stringify-email";
 import { createScopedLogger } from "@/utils/logger";
 import { hasPreviousEmailsFromSenderOrDomain } from "@/utils/gmail/message";
+import type { EmailForLLM } from "@/utils/types";
 
 const logger = createScopedLogger("ai-cold-email");
 
@@ -20,14 +21,7 @@ export async function isColdEmail({
   user,
   gmail,
 }: {
-  email: {
-    from: string;
-    subject: string;
-    content: string;
-    date?: Date;
-    threadId?: string;
-    messageId: string | null;
-  };
+  email: EmailForLLM & { threadId?: string };
   user: Pick<User, "id" | "email" | "coldEmailPrompt"> & UserAIFields;
   gmail: gmail_v1.Gmail;
 }): Promise<{
@@ -39,7 +33,7 @@ export async function isColdEmail({
     userId: user.id,
     email: user.email,
     threadId: email.threadId,
-    messageId: email.messageId,
+    messageId: email.id,
   };
 
   logger.info("Checking is cold email", loggerOptions);
@@ -59,11 +53,11 @@ export async function isColdEmail({
   }
 
   const hasPreviousEmail =
-    email.date && email.messageId
+    email.date && email.id
       ? await hasPreviousEmailsFromSenderOrDomain(gmail, {
           from: email.from,
           date: email.date,
-          messageId: email.messageId,
+          messageId: email.id,
         })
       : false;
 
@@ -105,7 +99,7 @@ async function isKnownColdEmailSender({
 }
 
 async function aiIsColdEmail(
-  email: { from: string; subject: string; content: string },
+  email: EmailForLLM,
   user: Pick<User, "email" | "coldEmailPrompt"> & UserAIFields,
 ) {
   const system = `You are an assistant that decides if an email is a cold email or not.
@@ -153,14 +147,7 @@ ${stringifyEmail(email, 500)}
 }
 
 export async function runColdEmailBlocker(options: {
-  email: {
-    from: string;
-    subject: string;
-    content: string;
-    messageId: string;
-    threadId: string;
-    date: Date;
-  };
+  email: EmailForLLM & { threadId: string };
   gmail: gmail_v1.Gmail;
   user: Pick<User, "id" | "email" | "coldEmailPrompt" | "coldEmailBlocker"> &
     UserAIFields;
@@ -173,7 +160,7 @@ export async function runColdEmailBlocker(options: {
 
 export async function blockColdEmail(options: {
   gmail: gmail_v1.Gmail;
-  email: { from: string; messageId: string; threadId: string };
+  email: { from: string; id: string; threadId: string };
   user: Pick<User, "id" | "email" | "coldEmailBlocker">;
   aiReason: string | null;
 }) {
@@ -187,7 +174,7 @@ export async function blockColdEmail(options: {
       fromEmail: email.from,
       userId: user.id,
       reason: aiReason,
-      messageId: email.messageId,
+      messageId: email.id,
       threadId: email.threadId,
     },
   });
@@ -221,7 +208,7 @@ export async function blockColdEmail(options: {
 
     await labelMessage({
       gmail,
-      messageId: email.messageId,
+      messageId: email.id,
       addLabelIds: addLabelIds.length ? addLabelIds : undefined,
       removeLabelIds: removeLabelIds.length ? removeLabelIds : undefined,
     });

apps/web/utils/get-email-from-message.ts

@@ -7,6 +7,7 @@ export function getEmailForLLM(
   contentOptions?: EmailToContentOptions,
 ): EmailForLLM {
   return {
+    id: message.id,
     from: message.headers.from,
     replyTo: message.headers["reply-to"],
     cc: message.headers.cc,

apps/web/utils/types.ts

@@ -102,6 +102,7 @@ export interface ParsedMessageHeaders {
 }
 
 export type EmailForLLM = {
+  id: string;
   from: string;
   replyTo?: string;
   cc?: string;