Improve AI choose rule prompt and add real world tests

Author: elie222

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

@@ -91,6 +91,222 @@ describe.skipIf(!isAiTest)("aiChooseRule", () => {
       reason: expect.any(String),
     });
   });
+
+  describe("Complex real-world rule scenarios", () => {
+    const recruiters = getRule(
+      "Match emails from recruiters or about job opportunities",
+    );
+    const legal = getRule(
+      "Match emails containing legal documents or contracts",
+    );
+    const requiresResponse = getRule(
+      "Match emails that require a response or action from the recipient",
+    );
+    const productUpdates = getRule(
+      "Match emails about product updates or feature announcements",
+    );
+    const financial = getRule(
+      "Match emails containing financial information or invoices",
+    );
+    const technicalIssues = getRule(
+      "Match emails about technical issues like server downtime or bug reports",
+    );
+    const marketing = getRule(
+      "Match emails containing marketing or promotional content",
+    );
+    const teamUpdates = getRule(
+      "Match emails about team updates or internal communications",
+    );
+    const customerFeedback = getRule(
+      "Match emails about customer feedback or support requests",
+    );
+    const events = getRule(
+      "Match emails containing event invitations or RSVPs",
+    );
+    const projectDeadlines = getRule(
+      "Match emails about project deadlines or milestones",
+    );
+    const urgent = getRule("Match urgent emails requiring immediate attention");
+    const catchAll = getRule("Match emails that don't fit any other category");
+
+    const rules = [
+      recruiters,
+      legal,
+      requiresResponse,
+      productUpdates,
+      financial,
+      technicalIssues,
+      marketing,
+      teamUpdates,
+      customerFeedback,
+      events,
+      projectDeadlines,
+      urgent,
+      catchAll,
+    ];
+
+    test("Should match technical issues", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "Server downtime reported",
+          content:
+            "We're experiencing critical server issues affecting production.",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: technicalIssues,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match financial emails", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "Your invoice for March 2024",
+          content: "Please find attached your invoice for services rendered.",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: financial,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match recruiter emails", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "New job opportunity at Tech Corp",
+          content:
+            "I came across your profile and think you'd be perfect for...",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: recruiters,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match legal documents", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "Please review: Contract for new project",
+          content: "Attached is the contract for your review and signature.",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: legal,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match emails requiring response", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "Team lunch tomorrow?",
+          content: "Would you like to join us for team lunch tomorrow at 12pm?",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: requiresResponse,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match product updates", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "New Feature Release: AI Integration",
+          content: "We're excited to announce our new AI features...",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: productUpdates,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match marketing emails", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "50% off Spring Sale!",
+          content: "Don't miss out on our biggest sale of the season!",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: marketing,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match team updates", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "Weekly Team Update",
+          content: "Here's what the team accomplished this week...",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: teamUpdates,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match customer feedback", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "Customer Feedback: App Performance",
+          content: "I've been experiencing slow loading times...",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: customerFeedback,
+        reason: expect.any(String),
+      });
+    });
+
+    test("Should match event invitations", async () => {
+      const result = await aiChooseRule({
+        rules,
+        email: getEmail({
+          subject: "Invitation: Annual Tech Conference",
+          content: "You're invited to speak at our annual conference...",
+        }),
+        user: getUser(),
+      });
+
+      expect(result).toEqual({
+        rule: events,
+        reason: expect.any(String),
+      });
+    });
+  });
 });
 
 // helpers

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

@@ -34,6 +34,8 @@ async function getAiResponse(options: GetAiResponseOptions) {
   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.
   

apps/web/utils/ai/rule/create-rule-schema.ts

@@ -102,10 +102,12 @@ export const getCreateRuleSchemaWithCategories = (
               "Whether senders in `categoryFilters` should be included or excluded",
             ),
           categoryFilters: z
-            .array(z.enum(availableCategories))
+            .array(z.string())
             .optional()
             .describe(
-              "The categories to match. If multiple categories are specified, the rule will match if ANY of the categories match (OR operation)",
+              `The categories to match. If multiple categories are specified, the rule will match if ANY of the categories match (OR operation). Available categories: ${availableCategories
+                .map((c) => `"${c}"`)
+                .join(", ")}`,
             ),
         })
         .optional()