From f8f9306ce7364c9857511008361027035f685a7b Mon Sep 17 00:00:00 2001
From: Eliezer Steinbock <3090527+elie222@users.noreply.github.com>
Date: Fri, 22 Nov 2024 15:00:45 -0500
Subject: [PATCH 1/8] Handle ai args better for duplicate action types

---
 apps/web/app/api/google/webhook/group-rule.ts | 23 ++++----
 .../web/app/api/google/webhook/static-rule.ts | 23 ++++----
 .../utils/ai/choose-rule/ai-choose-args.ts    | 56 +++++++++----------
 apps/web/utils/ai/choose-rule/choose.ts       | 36 ++++++------
 4 files changed, 64 insertions(+), 74 deletions(-)

diff --git a/apps/web/app/api/google/webhook/group-rule.ts b/apps/web/app/api/google/webhook/group-rule.ts
index df99732017..ec44b0a7d2 100644
--- a/apps/web/app/api/google/webhook/group-rule.ts
+++ b/apps/web/app/api/google/webhook/group-rule.ts
@@ -3,7 +3,6 @@ import type { ParsedMessage } from "@/utils/types";
 import type { User } from "@prisma/client";
 import { emailToContent } from "@/utils/mail";
 import {
-  getActionItemsFromAiArgsResponse,
   getActionsWithParameters,
   getArgsAiResponse,
 } from "@/utils/ai/choose-rule/ai-choose-args";
@@ -61,19 +60,17 @@ export async function handleGroupRule({
   };
 
   // generate args
-  const aiArgsResponse =
-    getActionsWithParameters(match.rule.actions).length > 0
-      ? await getArgsAiResponse({
-          email,
-          selectedRule: match.rule,
-          user,
-        })
-      : undefined;
+  const shouldAiGenerateArgs =
+    getActionsWithParameters(match.rule.actions).length > 0;
+  const aiArgsResponse = shouldAiGenerateArgs
+    ? await getArgsAiResponse({
+        email,
+        selectedRule: match.rule,
+        user,
+      })
+    : match.rule.actions;
 
-  const actionItems = getActionItemsFromAiArgsResponse(
-    aiArgsResponse,
-    match.rule.actions,
-  );
+  const actionItems = aiArgsResponse || match.rule.actions;
 
   // handle action
   // TODO isThread check to skip
diff --git a/apps/web/app/api/google/webhook/static-rule.ts b/apps/web/app/api/google/webhook/static-rule.ts
index f0772c811e..113c36426c 100644
--- a/apps/web/app/api/google/webhook/static-rule.ts
+++ b/apps/web/app/api/google/webhook/static-rule.ts
@@ -2,7 +2,6 @@ import type { gmail_v1 } from "@googleapis/gmail";
 import type { ParsedMessage, RuleWithActions } from "@/utils/types";
 import { RuleType, type User } from "@prisma/client";
 import {
-  getActionItemsFromAiArgsResponse,
   getActionsWithParameters,
   getArgsAiResponse,
 } from "@/utils/ai/choose-rule/ai-choose-args";
@@ -53,19 +52,17 @@ export async function handleStaticRule({
   };
 
   // generate args
-  const aiArgsResponse =
-    getActionsWithParameters(staticRule.actions).length > 0
-      ? await getArgsAiResponse({
-          email,
-          selectedRule: staticRule,
-          user,
-        })
-      : undefined;
+  const shouldAiGenerateArgs =
+    getActionsWithParameters(staticRule.actions).length > 0;
+  const aiArgsResponse = shouldAiGenerateArgs
+    ? await getArgsAiResponse({
+        email,
+        selectedRule: staticRule,
+        user,
+      })
+    : staticRule.actions;
 
-  const actionItems = getActionItemsFromAiArgsResponse(
-    aiArgsResponse,
-    staticRule.actions,
-  );
+  const actionItems = aiArgsResponse || staticRule.actions;
 
   // handle action
   // TODO isThread check to skip
diff --git a/apps/web/utils/ai/choose-rule/ai-choose-args.ts b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
index 814ec571a2..282372c9ce 100644
--- a/apps/web/utils/ai/choose-rule/ai-choose-args.ts
+++ b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
@@ -1,7 +1,7 @@
 import { z } from "zod";
 import type { UserAIFields } from "@/utils/llms/types";
 import type { ActionItem } from "@/utils/ai/actions";
-import type { Action, ActionType, User } from "@prisma/client";
+import type { Action, User } from "@prisma/client";
 import { chatCompletionTools } from "@/utils/llms";
 import {
   type EmailForLLM,
@@ -9,11 +9,6 @@ import {
 } from "@/utils/ai/choose-rule/stringify-email";
 import { type RuleWithActions, isDefined } from "@/utils/types";
 
-type AIGeneratedArgs = Record<
-  ActionType,
-  Record<keyof Omit<ActionItem, "type">, string>
->;
-
 // Returns parameters for a zod.object for the rule that must be AI generated
 function getToolParametersForRule(actions: Action[]) {
   const actionsWithParameters = getActionsWithParameters(actions);
@@ -21,16 +16,22 @@ function getToolParametersForRule(actions: Action[]) {
   // handle duplicate keys. e.g. "draft_email" and "draft_email" becomes: "draft_email" and "draft_email_2"
   // this is quite an edge case but need to handle regardless for when it happens
   const typeCount: Record<string, number> = {};
-  const parameters: Record<string, z.ZodObject<any>> = {};
+  const parameters: Record<
+    string,
+    { action: Action; parameters: z.ZodObject<Record<string, z.ZodString>> }
+  > = {};
 
   for (const action of actionsWithParameters) {
     // count how many times we have already had this type
     typeCount[action.type] = (typeCount[action.type] || 0) + 1;
-    parameters[
+    const key =
       typeCount[action.type] === 1
         ? action.type
-        : `${action.type}_${typeCount[action.type]}`
-    ] = action.parameters;
+        : `${action.type}_${typeCount[action.type]}`;
+    parameters[key] = {
+      action: action.action,
+      parameters: action.parameters,
+    };
   }
 
   return parameters;
@@ -48,6 +49,7 @@ export function getActionsWithParameters(actions: Action[]) {
       return {
         type: action.type,
         parameters,
+        action,
       };
     })
     .filter(isDefined);
@@ -87,7 +89,7 @@ export async function getArgsAiResponse({
   email: EmailForLLM;
   user: Pick<User, "email" | "about"> & UserAIFields;
   selectedRule: RuleWithActions;
-}) {
+}): Promise<ActionItem[] | undefined> {
   console.log(
     `Generating args for rule ${selectedRule.name} (${selectedRule.id})`,
   );
@@ -130,7 +132,14 @@ ${stringifyEmail(email, 3000)}
     tools: {
       apply_rule: {
         description: "Apply the rule with the given arguments",
-        parameters: z.object(parameters),
+        parameters: z.object(
+          Object.fromEntries(
+            Object.entries(parameters).map(([key, { parameters }]) => [
+              key,
+              parameters,
+            ]),
+          ),
+        ),
       },
     },
     label: "Args for rule",
@@ -141,25 +150,12 @@ ${stringifyEmail(email, 3000)}
 
   if (!toolCall?.toolName) return;
 
-  return toolCall.args;
-}
-
-export function getActionItemsFromAiArgsResponse(
-  response: AIGeneratedArgs | undefined,
-  ruleActions: Action[],
-) {
-  return ruleActions.map((ra) => {
-    // use prefilled values where we have them
-    const a = response?.[ra.type] || ({} as any);
-
+  const actionItems = Object.entries(parameters).map(([key, { action }]) => {
     return {
-      type: ra.type,
-      label: ra.labelPrompt ? a.label : ra.label,
-      subject: ra.subjectPrompt ? a.subject : ra.subject,
-      content: ra.contentPrompt ? a.content : ra.content,
-      to: ra.toPrompt ? a.to : ra.to,
-      cc: ra.ccPrompt ? a.cc : ra.cc,
-      bcc: ra.bccPrompt ? a.bcc : ra.bcc,
+      ...action,
+      ...toolCall.args[key],
     };
   });
+
+  return actionItems;
 }
diff --git a/apps/web/utils/ai/choose-rule/choose.ts b/apps/web/utils/ai/choose-rule/choose.ts
index 908ef38403..4a40eb4a38 100644
--- a/apps/web/utils/ai/choose-rule/choose.ts
+++ b/apps/web/utils/ai/choose-rule/choose.ts
@@ -1,7 +1,6 @@
 import type { ActionItem } from "@/utils/ai/actions";
 import {
   getArgsAiResponse,
-  getActionItemsFromAiArgsResponse,
   getActionsWithParameters,
 } from "@/utils/ai/choose-rule/ai-choose-args";
 import { getAiResponse } from "@/utils/ai/choose-rule/ai-choose-rule";
@@ -46,22 +45,23 @@ export async function chooseRule(options: ChooseRuleOptions): Promise<
   const shouldAiGenerateArgs =
     getActionsWithParameters(selectedRule.actions).length > 0;
 
-  const aiArgsResponse = shouldAiGenerateArgs
-    ? await getArgsAiResponse({
-        ...options,
-        email,
-        selectedRule,
-      })
-    : undefined;
+  if (shouldAiGenerateArgs) {
+    const aiArgsResponse = await getArgsAiResponse({
+      ...options,
+      email,
+      selectedRule,
+    });
 
-  const actionItems = getActionItemsFromAiArgsResponse(
-    aiArgsResponse,
-    selectedRule.actions,
-  );
-
-  return {
-    rule: selectedRule,
-    actionItems,
-    reason: aiResponse?.reason,
-  };
+    return {
+      rule: selectedRule,
+      actionItems: aiArgsResponse || [],
+      reason: aiResponse?.reason,
+    };
+  } else {
+    return {
+      rule: selectedRule,
+      actionItems: selectedRule.actions,
+      reason: aiResponse?.reason,
+    };
+  }
 }

From 8c87b8a89df8bbad3b70c6cc48fe6b3c013f0191 Mon Sep 17 00:00:00 2001
From: Eliezer Steinbock <3090527+elie222@users.noreply.github.com>
Date: Fri, 22 Nov 2024 15:24:20 -0500
Subject: [PATCH 2/8] Don't add extra fields to db

---
 apps/web/utils/ai/choose-rule/ai-choose-args.ts | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/apps/web/utils/ai/choose-rule/ai-choose-args.ts b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
index 282372c9ce..3e1a682f49 100644
--- a/apps/web/utils/ai/choose-rule/ai-choose-args.ts
+++ b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
@@ -151,10 +151,16 @@ ${stringifyEmail(email, 3000)}
   if (!toolCall?.toolName) return;
 
   const actionItems = Object.entries(parameters).map(([key, { action }]) => {
-    return {
-      ...action,
-      ...toolCall.args[key],
+    const actionItem: ActionItem = {
+      type: action.type,
+      label: toolCall.args[key].label || action.label,
+      subject: toolCall.args[key].subject || action.subject,
+      content: toolCall.args[key].content || action.content,
+      to: toolCall.args[key].to || action.to,
+      cc: toolCall.args[key].cc || action.cc,
+      bcc: toolCall.args[key].bcc || action.bcc,
     };
+    return actionItem;
   });
 
   return actionItems;

From 4910a1b63d454e3bf05309a1b829daeba2359932 Mon Sep 17 00:00:00 2001
From: Eliezer Steinbock <3090527+elie222@users.noreply.github.com>
Date: Fri, 22 Nov 2024 15:24:26 -0500
Subject: [PATCH 3/8] Prettier logger

---
 .../app/api/google/webhook/process-history.ts | 52 +++++++++++--------
 1 file changed, 29 insertions(+), 23 deletions(-)

diff --git a/apps/web/app/api/google/webhook/process-history.ts b/apps/web/app/api/google/webhook/process-history.ts
index 45aac8993f..63c418ca9b 100644
--- a/apps/web/app/api/google/webhook/process-history.ts
+++ b/apps/web/app/api/google/webhook/process-history.ts
@@ -25,6 +25,14 @@ import { blockUnsubscribedEmails } from "@/app/api/google/webhook/block-unsubscr
 import { categorizeSender } from "@/utils/actions/categorize";
 import { unwatchEmails } from "@/app/api/google/watch/controller";
 
+const scope = "Process History";
+const log = (message: string) => {
+  console.log(`[${scope}]: ${message}`);
+};
+const logError = (message: string, error?: unknown, extra?: unknown) => {
+  console.error(`[${scope}]: ${message}`, error, extra);
+};
+
 export async function processHistoryForUser(
   decodedData: {
     emailAddress: string;
@@ -70,7 +78,7 @@ export async function processHistoryForUser(
   });
 
   if (!account) {
-    console.error(`Google webhook: Account not found. email: ${email}`);
+    logError(`Account not found. email: ${email}`);
     return NextResponse.json({ ok: true });
   }
 
@@ -79,7 +87,7 @@ export async function processHistoryForUser(
     : undefined;
 
   if (!premium) {
-    console.log(`Google webhook: Account not premium. email: ${email}`);
+    log(`Account not premium. email: ${email}`);
     await unwatchEmails(account);
     return NextResponse.json({ ok: true });
   }
@@ -94,9 +102,7 @@ export async function processHistoryForUser(
   );
 
   if (!userHasAiAccess && !userHasColdEmailAccess) {
-    console.debug(
-      `Google webhook: does not have hasAiOrColdEmailAccess. email: ${email}`,
-    );
+    console.debug(`does not have hasAiOrColdEmailAccess. email: ${email}`);
     await unwatchEmails(account);
     return NextResponse.json({ ok: true });
   }
@@ -107,13 +113,13 @@ export async function processHistoryForUser(
     account.user.coldEmailBlocker !== ColdEmailSetting.DISABLED;
   if (!hasAutomationRules && !shouldBlockColdEmails) {
     console.debug(
-      `Google webhook: has no rules set and cold email blocker disabled. email: ${email}`,
+      `has no rules set and cold email blocker disabled. email: ${email}`,
     );
     return NextResponse.json({ ok: true });
   }
 
   if (!account.access_token || !account.refresh_token) {
-    console.error(
+    logError(
       `Missing access or refresh token. User needs to re-authenticate. email: ${email}`,
     );
     return NextResponse.json({ ok: true });
@@ -121,7 +127,7 @@ export async function processHistoryForUser(
 
   if (!account.user.email) {
     // shouldn't ever happen
-    console.error("Missing user email.", email);
+    logError(`Missing user email: ${email}`);
     return NextResponse.json({ ok: true });
   }
 
@@ -145,8 +151,8 @@ export async function processHistoryForUser(
         historyId - 500, // avoid going too far back
       ).toString();
 
-    console.log(
-      `Webhook: Listing history... Start: ${startHistoryId} lastSyncedHistoryId: ${account.user.lastSyncedHistoryId} gmailHistoryId: ${startHistoryId} email: ${email}`,
+    log(
+      `Listing history... Start: ${startHistoryId} lastSyncedHistoryId: ${account.user.lastSyncedHistoryId} gmailHistoryId: ${startHistoryId} email: ${email}`,
     );
 
     const history = await gmail.users.history.list({
@@ -160,8 +166,8 @@ export async function processHistoryForUser(
     });
 
     if (history.data.history) {
-      console.log(
-        `Webhook: Processing... email: ${email} startHistoryId: ${startHistoryId} historyId: ${history.data.historyId}`,
+      log(
+        `Processing... email: ${email} startHistoryId: ${startHistoryId} historyId: ${history.data.historyId}`,
       );
 
       await processHistory({
@@ -186,8 +192,8 @@ export async function processHistoryForUser(
         },
       });
     } else {
-      console.log(
-        `Webhook: No history. startHistoryId: ${startHistoryId}. ${JSON.stringify(decodedData)}`,
+      log(
+        `No history. startHistoryId: ${startHistoryId}. ${JSON.stringify(decodedData)}`,
       );
 
       // important to save this or we can get into a loop with never receiving history
@@ -197,12 +203,12 @@ export async function processHistoryForUser(
       });
     }
 
-    console.log(`Webhook: Completed. ${JSON.stringify(decodedData)}`);
+    log(`Completed. ${JSON.stringify(decodedData)}`);
 
     return NextResponse.json({ ok: true });
   } catch (error) {
     captureException(error, { extra: { decodedData } }, email);
-    console.error("Error processing webhook", error, decodedData);
+    logError("Error processing webhook", error, decodedData);
     return NextResponse.json({ error: true });
     // be careful about calling an error here with the wrong settings, as otherwise PubSub will call the webhook over and over
     // return NextResponse.error();
@@ -260,7 +266,7 @@ async function processHistory(options: ProcessHistoryOptions) {
           { extra: { email, messageId: m.message?.id } },
           email,
         );
-        console.error(`Error processing history item. email: ${email}`, error);
+        logError(`Error processing history item. email: ${email}`, error);
       }
     }
   }
@@ -292,7 +298,7 @@ async function processHistoryItem(
   if (!messageId) return;
   if (!threadId) return;
 
-  console.log(
+  log(
     `Getting message... email: ${user.email} messageId: ${messageId} threadId: ${threadId}`,
   );
 
@@ -309,7 +315,7 @@ async function processHistoryItem(
 
     // if the rule has already been executed, skip
     if (hasExistingRule) {
-      console.log("Skipping. Rule already exists.");
+      log("Skipping. Rule already exists.");
       return;
     }
 
@@ -326,7 +332,7 @@ async function processHistoryItem(
     });
 
     if (blocked) {
-      console.log(
+      log(
         `Skipping. Blocked unsubscribed email. email: ${user.email} messageId: ${messageId} threadId: ${threadId}`,
       );
       return;
@@ -341,7 +347,7 @@ async function processHistoryItem(
     );
 
     if (shouldRunBlocker) {
-      console.log("Running cold email blocker...");
+      log("Running cold email blocker...");
 
       const hasPreviousEmail = await hasPreviousEmailsFromSenderOrDomain(
         gmail,
@@ -386,7 +392,7 @@ async function processHistoryItem(
     }
 
     if (hasAutomationRules && hasAiAutomationAccess) {
-      console.log("Running rules...");
+      log("Running rules...");
 
       await runRulesOnMessage({
         gmail,
@@ -399,7 +405,7 @@ async function processHistoryItem(
   } catch (error: any) {
     // gmail bug or snoozed email: https://stackoverflow.com/questions/65290987/gmail-api-getmessage-method-returns-404-for-message-gotten-from-listhistory-meth
     if (error.message === "Requested entity was not found.") {
-      console.log(
+      log(
         `Message not found. email: ${user.email} messageId: ${messageId} threadId: ${threadId}`,
       );
       return;

From 77f5a7e92566f87bf801af79691aa82a406e3546 Mon Sep 17 00:00:00 2001
From: Eliezer Steinbock <3090527+elie222@users.noreply.github.com>
Date: Fri, 22 Nov 2024 16:04:38 -0500
Subject: [PATCH 4/8] Adjust copy

---
 apps/web/app/(app)/automation/Rules.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/apps/web/app/(app)/automation/Rules.tsx b/apps/web/app/(app)/automation/Rules.tsx
index 6050abea37..242b86353f 100644
--- a/apps/web/app/(app)/automation/Rules.tsx
+++ b/apps/web/app/(app)/automation/Rules.tsx
@@ -245,7 +245,7 @@ export function Rules() {
           <Button asChild variant="outline">
             <Link href="/automation?tab=prompt">
               <PenIcon className="mr-2 hidden h-4 w-4 md:block" />
-              Add Rule via Prompt File
+              Add Rule via Prompt Tab
             </Link>
           </Button>
           <Button asChild variant="outline">

From f4105c46b43baac05f4959afcb444d43b72dab76 Mon Sep 17 00:00:00 2001
From: Eliezer Steinbock <3090527+elie222@users.noreply.github.com>
Date: Fri, 22 Nov 2024 16:24:43 -0500
Subject: [PATCH 5/8] Adjust logger

---
 .../app/api/google/webhook/process-history.ts | 41 ++++++++-----------
 .../utils/ai/choose-rule/ai-choose-args.ts    | 29 +++++++++----
 apps/web/utils/logger.ts                      | 38 +++++++++++++++++
 3 files changed, 76 insertions(+), 32 deletions(-)
 create mode 100644 apps/web/utils/logger.ts

diff --git a/apps/web/app/api/google/webhook/process-history.ts b/apps/web/app/api/google/webhook/process-history.ts
index 63c418ca9b..5d4317b9af 100644
--- a/apps/web/app/api/google/webhook/process-history.ts
+++ b/apps/web/app/api/google/webhook/process-history.ts
@@ -24,14 +24,9 @@ import { runRulesOnMessage } from "@/utils/ai/choose-rule/run-rules";
 import { blockUnsubscribedEmails } from "@/app/api/google/webhook/block-unsubscribed-emails";
 import { categorizeSender } from "@/utils/actions/categorize";
 import { unwatchEmails } from "@/app/api/google/watch/controller";
+import { createScopeLogger } from "@/utils/logger";
 
-const scope = "Process History";
-const log = (message: string) => {
-  console.log(`[${scope}]: ${message}`);
-};
-const logError = (message: string, error?: unknown, extra?: unknown) => {
-  console.error(`[${scope}]: ${message}`, error, extra);
-};
+const logger = createScopeLogger("Process History");
 
 export async function processHistoryForUser(
   decodedData: {
@@ -78,7 +73,7 @@ export async function processHistoryForUser(
   });
 
   if (!account) {
-    logError(`Account not found. email: ${email}`);
+    logger.error(`Account not found. email: ${email}`);
     return NextResponse.json({ ok: true });
   }
 
@@ -87,7 +82,7 @@ export async function processHistoryForUser(
     : undefined;
 
   if (!premium) {
-    log(`Account not premium. email: ${email}`);
+    logger.log(`Account not premium. email: ${email}`);
     await unwatchEmails(account);
     return NextResponse.json({ ok: true });
   }
@@ -119,7 +114,7 @@ export async function processHistoryForUser(
   }
 
   if (!account.access_token || !account.refresh_token) {
-    logError(
+    logger.error(
       `Missing access or refresh token. User needs to re-authenticate. email: ${email}`,
     );
     return NextResponse.json({ ok: true });
@@ -127,7 +122,7 @@ export async function processHistoryForUser(
 
   if (!account.user.email) {
     // shouldn't ever happen
-    logError(`Missing user email: ${email}`);
+    logger.error(`Missing user email: ${email}`);
     return NextResponse.json({ ok: true });
   }
 
@@ -151,7 +146,7 @@ export async function processHistoryForUser(
         historyId - 500, // avoid going too far back
       ).toString();
 
-    log(
+    logger.log(
       `Listing history... Start: ${startHistoryId} lastSyncedHistoryId: ${account.user.lastSyncedHistoryId} gmailHistoryId: ${startHistoryId} email: ${email}`,
     );
 
@@ -166,7 +161,7 @@ export async function processHistoryForUser(
     });
 
     if (history.data.history) {
-      log(
+      logger.log(
         `Processing... email: ${email} startHistoryId: ${startHistoryId} historyId: ${history.data.historyId}`,
       );
 
@@ -192,7 +187,7 @@ export async function processHistoryForUser(
         },
       });
     } else {
-      log(
+      logger.log(
         `No history. startHistoryId: ${startHistoryId}. ${JSON.stringify(decodedData)}`,
       );
 
@@ -203,12 +198,12 @@ export async function processHistoryForUser(
       });
     }
 
-    log(`Completed. ${JSON.stringify(decodedData)}`);
+    logger.log(`Completed. ${JSON.stringify(decodedData)}`);
 
     return NextResponse.json({ ok: true });
   } catch (error) {
     captureException(error, { extra: { decodedData } }, email);
-    logError("Error processing webhook", error, decodedData);
+    logger.error("Error processing webhook", error, decodedData);
     return NextResponse.json({ error: true });
     // be careful about calling an error here with the wrong settings, as otherwise PubSub will call the webhook over and over
     // return NextResponse.error();
@@ -266,7 +261,7 @@ async function processHistory(options: ProcessHistoryOptions) {
           { extra: { email, messageId: m.message?.id } },
           email,
         );
-        logError(`Error processing history item. email: ${email}`, error);
+        logger.error(`Error processing history item. email: ${email}`, error);
       }
     }
   }
@@ -298,7 +293,7 @@ async function processHistoryItem(
   if (!messageId) return;
   if (!threadId) return;
 
-  log(
+  logger.log(
     `Getting message... email: ${user.email} messageId: ${messageId} threadId: ${threadId}`,
   );
 
@@ -315,7 +310,7 @@ async function processHistoryItem(
 
     // if the rule has already been executed, skip
     if (hasExistingRule) {
-      log("Skipping. Rule already exists.");
+      logger.log("Skipping. Rule already exists.");
       return;
     }
 
@@ -332,7 +327,7 @@ async function processHistoryItem(
     });
 
     if (blocked) {
-      log(
+      logger.log(
         `Skipping. Blocked unsubscribed email. email: ${user.email} messageId: ${messageId} threadId: ${threadId}`,
       );
       return;
@@ -347,7 +342,7 @@ async function processHistoryItem(
     );
 
     if (shouldRunBlocker) {
-      log("Running cold email blocker...");
+      logger.log("Running cold email blocker...");
 
       const hasPreviousEmail = await hasPreviousEmailsFromSenderOrDomain(
         gmail,
@@ -392,7 +387,7 @@ async function processHistoryItem(
     }
 
     if (hasAutomationRules && hasAiAutomationAccess) {
-      log("Running rules...");
+      logger.log("Running rules...");
 
       await runRulesOnMessage({
         gmail,
@@ -405,7 +400,7 @@ async function processHistoryItem(
   } catch (error: any) {
     // gmail bug or snoozed email: https://stackoverflow.com/questions/65290987/gmail-api-getmessage-method-returns-404-for-message-gotten-from-listhistory-meth
     if (error.message === "Requested entity was not found.") {
-      log(
+      logger.log(
         `Message not found. email: ${user.email} messageId: ${messageId} threadId: ${threadId}`,
       );
       return;
diff --git a/apps/web/utils/ai/choose-rule/ai-choose-args.ts b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
index 3e1a682f49..0d8518b605 100644
--- a/apps/web/utils/ai/choose-rule/ai-choose-args.ts
+++ b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
@@ -8,6 +8,9 @@ import {
   stringifyEmail,
 } from "@/utils/ai/choose-rule/stringify-email";
 import { type RuleWithActions, isDefined } from "@/utils/types";
+import { createScopeLogger } from "@/utils/logger";
+
+const logger = createScopeLogger("AI Choose Args");
 
 // Returns parameters for a zod.object for the rule that must be AI generated
 function getToolParametersForRule(actions: Action[]) {
@@ -90,14 +93,14 @@ export async function getArgsAiResponse({
   user: Pick<User, "email" | "about"> & UserAIFields;
   selectedRule: RuleWithActions;
 }): Promise<ActionItem[] | undefined> {
-  console.log(
+  logger.log(
     `Generating args for rule ${selectedRule.name} (${selectedRule.id})`,
   );
 
   const parameters = getToolParametersForRule(selectedRule.actions);
 
   if (!Object.keys(parameters).length) {
-    console.log(
+    logger.log(
       `Skipping. No parameters for rule ${selectedRule.name} (${selectedRule.id})`,
     );
     return;
@@ -123,7 +126,10 @@ ${stringifyEmail(email, 3000)}
 </email>
 `;
 
-  console.log("Calling chat completion tools");
+  logger.log("Calling chat completion tools");
+
+  logger.trace(`System: ${system}`);
+  logger.trace(`Prompt: ${prompt}`);
 
   const aiResponse = await chatCompletionTools({
     userAi: user,
@@ -150,18 +156,23 @@ ${stringifyEmail(email, 3000)}
 
   if (!toolCall?.toolName) return;
 
+  const toolCallArgs = toolCall.args;
+  logger.trace(`Tool call args: ${JSON.stringify(toolCallArgs, null, 2)}`);
+
   const actionItems = Object.entries(parameters).map(([key, { action }]) => {
     const actionItem: ActionItem = {
       type: action.type,
-      label: toolCall.args[key].label || action.label,
-      subject: toolCall.args[key].subject || action.subject,
-      content: toolCall.args[key].content || action.content,
-      to: toolCall.args[key].to || action.to,
-      cc: toolCall.args[key].cc || action.cc,
-      bcc: toolCall.args[key].bcc || action.bcc,
+      label: toolCallArgs[key].label || action.label,
+      subject: toolCallArgs[key].subject || action.subject,
+      content: toolCallArgs[key].content || action.content,
+      to: toolCallArgs[key].to || action.to,
+      cc: toolCallArgs[key].cc || action.cc,
+      bcc: toolCallArgs[key].bcc || action.bcc,
     };
     return actionItem;
   });
 
+  logger.trace(`Action items: ${JSON.stringify(actionItems, null, 2)}`);
+
   return actionItems;
 }
diff --git a/apps/web/utils/logger.ts b/apps/web/utils/logger.ts
new file mode 100644
index 0000000000..c0bf2667b7
--- /dev/null
+++ b/apps/web/utils/logger.ts
@@ -0,0 +1,38 @@
+type LogLevel = "log" | "error" | "warn" | "trace";
+type LogMessage = string | Record<string, unknown>;
+
+const colors = {
+  log: "\x1b[0m", // white
+  error: "\x1b[31m", // red
+  warn: "\x1b[33m", // yellow
+  trace: "\x1b[36m", // cyan
+  reset: "\x1b[0m",
+} as const;
+
+export function createScopeLogger(scope: string) {
+  const formatMessage = (level: LogLevel, message: LogMessage) => {
+    const prefix = `[${scope}]: `;
+
+    if (process.env.NODE_ENV === "development") {
+      return `${colors[level]}${prefix} ${
+        typeof message === "string" ? message : JSON.stringify(message, null, 2)
+      }${colors.reset}`;
+    }
+
+    return `${prefix} ${
+      typeof message === "string" ? message : JSON.stringify(message)
+    }`;
+  };
+
+  return {
+    log: (message: LogMessage) => console.log(formatMessage("log", message)),
+    error: (message: LogMessage, error?: unknown, extra?: unknown) =>
+      console.error(formatMessage("error", message), error, extra),
+    warn: (message: LogMessage) => console.warn(formatMessage("warn", message)),
+    trace: (message: LogMessage) => {
+      if (process.env.NODE_ENV === "development") {
+        console.log(formatMessage("trace", message));
+      }
+    },
+  };
+}

From 26e99bc2edae8abffc9af4af22dc2e81e2831278 Mon Sep 17 00:00:00 2001
From: Eliezer Steinbock <3090527+elie222@users.noreply.github.com>
Date: Fri, 22 Nov 2024 17:33:29 -0500
Subject: [PATCH 6/8] Fix ai generated fields

---
 apps/web/app/api/google/webhook/group-rule.ts |  23 +--
 .../web/app/api/google/webhook/static-rule.ts |  23 +--
 .../utils/ai/choose-rule/ai-choose-args.ts    | 139 ++++++++++++------
 .../ai/choose-rule/choose-and-execute.ts      |  15 +-
 apps/web/utils/ai/choose-rule/choose.ts       |  36 ++---
 apps/web/utils/logger.ts                      |  17 ++-
 6 files changed, 145 insertions(+), 108 deletions(-)

diff --git a/apps/web/app/api/google/webhook/group-rule.ts b/apps/web/app/api/google/webhook/group-rule.ts
index ec44b0a7d2..d9c6aca756 100644
--- a/apps/web/app/api/google/webhook/group-rule.ts
+++ b/apps/web/app/api/google/webhook/group-rule.ts
@@ -2,10 +2,7 @@ import type { gmail_v1 } from "@googleapis/gmail";
 import type { ParsedMessage } from "@/utils/types";
 import type { User } from "@prisma/client";
 import { emailToContent } from "@/utils/mail";
-import {
-  getActionsWithParameters,
-  getArgsAiResponse,
-} from "@/utils/ai/choose-rule/ai-choose-args";
+import { getActionItemsWithAiArgs } from "@/utils/ai/choose-rule/ai-choose-args";
 import {
   findMatchingGroup,
   getGroups,
@@ -59,18 +56,12 @@ export async function handleGroupRule({
     }),
   };
 
-  // generate args
-  const shouldAiGenerateArgs =
-    getActionsWithParameters(match.rule.actions).length > 0;
-  const aiArgsResponse = shouldAiGenerateArgs
-    ? await getArgsAiResponse({
-        email,
-        selectedRule: match.rule,
-        user,
-      })
-    : match.rule.actions;
-
-  const actionItems = aiArgsResponse || match.rule.actions;
+  // get action items with args
+  const actionItems = await getActionItemsWithAiArgs({
+    email,
+    user,
+    selectedRule: match.rule,
+  });
 
   // handle action
   // TODO isThread check to skip
diff --git a/apps/web/app/api/google/webhook/static-rule.ts b/apps/web/app/api/google/webhook/static-rule.ts
index 113c36426c..15cfe203a9 100644
--- a/apps/web/app/api/google/webhook/static-rule.ts
+++ b/apps/web/app/api/google/webhook/static-rule.ts
@@ -1,10 +1,7 @@
 import type { gmail_v1 } from "@googleapis/gmail";
 import type { ParsedMessage, RuleWithActions } from "@/utils/types";
 import { RuleType, type User } from "@prisma/client";
-import {
-  getActionsWithParameters,
-  getArgsAiResponse,
-} from "@/utils/ai/choose-rule/ai-choose-args";
+import { getActionItemsWithAiArgs } from "@/utils/ai/choose-rule/ai-choose-args";
 import { emailToContent } from "@/utils/mail";
 import { saveExecutedRule } from "@/utils/ai/choose-rule/choose-and-execute";
 import { executeAct } from "@/utils/ai/choose-rule/execute";
@@ -51,18 +48,12 @@ export async function handleStaticRule({
     }),
   };
 
-  // generate args
-  const shouldAiGenerateArgs =
-    getActionsWithParameters(staticRule.actions).length > 0;
-  const aiArgsResponse = shouldAiGenerateArgs
-    ? await getArgsAiResponse({
-        email,
-        selectedRule: staticRule,
-        user,
-      })
-    : staticRule.actions;
-
-  const actionItems = aiArgsResponse || staticRule.actions;
+  // get action items with args
+  const actionItems = await getActionItemsWithAiArgs({
+    email,
+    user,
+    selectedRule: staticRule,
+  });
 
   // handle action
   // TODO isThread check to skip
diff --git a/apps/web/utils/ai/choose-rule/ai-choose-args.ts b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
index 0d8518b605..e504d56e3e 100644
--- a/apps/web/utils/ai/choose-rule/ai-choose-args.ts
+++ b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
@@ -12,19 +12,30 @@ import { createScopeLogger } from "@/utils/logger";
 
 const logger = createScopeLogger("AI Choose Args");
 
-// Returns parameters for a zod.object for the rule that must be AI generated
-function getToolParametersForRule(actions: Action[]) {
-  const actionsWithParameters = getActionsWithParameters(actions);
+export type ActionRequiringAi = {
+  actionId: string;
+  type: string;
+  zodParameters: z.ZodObject<Record<string, z.ZodString>>;
+};
+
+export type ActionWithAiArgs = ActionRequiringAi & {
+  args: Record<string, string>;
+};
 
+// Returns parameters for a zod.object for the rule that must be AI generated
+function getToolParametersForRule(actionsRequiringAi: ActionRequiringAi[]) {
   // handle duplicate keys. e.g. "draft_email" and "draft_email" becomes: "draft_email" and "draft_email_2"
   // this is quite an edge case but need to handle regardless for when it happens
   const typeCount: Record<string, number> = {};
   const parameters: Record<
     string,
-    { action: Action; parameters: z.ZodObject<Record<string, z.ZodString>> }
+    {
+      actionId: string;
+      zodParameters: z.ZodObject<Record<string, z.ZodString>>;
+    }
   > = {};
 
-  for (const action of actionsWithParameters) {
+  for (const action of actionsRequiringAi) {
     // count how many times we have already had this type
     typeCount[action.type] = (typeCount[action.type] || 0) + 1;
     const key =
@@ -32,27 +43,28 @@ function getToolParametersForRule(actions: Action[]) {
         ? action.type
         : `${action.type}_${typeCount[action.type]}`;
     parameters[key] = {
-      action: action.action,
-      parameters: action.parameters,
+      actionId: action.actionId,
+      zodParameters: action.zodParameters,
     };
   }
 
   return parameters;
 }
 
-export function getActionsWithParameters(actions: Action[]) {
+// Returns the actions with the args that require ai-generation
+function getActionsWithZodParameters(actions: Action[]) {
   return actions
     .map((action) => {
       const fields = getParameterFieldsForAction(action);
 
       if (!Object.keys(fields).length) return;
 
-      const parameters = z.object(fields);
+      const zodParameters = z.object(fields);
 
       return {
+        actionId: action.id,
         type: action.type,
-        parameters,
-        action,
+        zodParameters,
       };
     })
     .filter(isDefined);
@@ -69,37 +81,73 @@ function getParameterFieldsForAction({
   const fields: Record<string, z.ZodString> = {};
 
   if (typeof labelPrompt === "string")
-    fields.label = z.string().describe("The email label");
+    fields.label = z.string().describe(labelPrompt || "The email label");
   if (typeof subjectPrompt === "string")
-    fields.subject = z.string().describe("The email subject");
+    fields.subject = z.string().describe(subjectPrompt || "The email subject");
   if (typeof contentPrompt === "string")
-    fields.content = z.string().describe("The email content");
+    fields.content = z.string().describe(contentPrompt || "The email content");
   if (typeof toPrompt === "string")
-    fields.to = z.string().describe("The email recipient(s)");
+    fields.to = z.string().describe(toPrompt || "The email recipient(s)");
   if (typeof ccPrompt === "string")
-    fields.cc = z.string().describe("The cc recipient(s)");
+    fields.cc = z.string().describe(ccPrompt || "The cc recipient(s)");
   if (typeof bccPrompt === "string")
-    fields.bcc = z.string().describe("The bcc recipient(s)");
+    fields.bcc = z.string().describe(bccPrompt || "The bcc recipient(s)");
 
   return fields;
 }
 
-export async function getArgsAiResponse({
+// Returns the action items with the ai-generated args where needed
+export async function getActionItemsWithAiArgs({
+  email,
+  user,
+  selectedRule,
+}: {
+  email: EmailForLLM;
+  user: Pick<User, "email" | "about"> & UserAIFields;
+  selectedRule: RuleWithActions;
+}): Promise<ActionItem[]> {
+  // Get the actions that require ai-generated args
+  const actionsRequiringAi = getActionsWithZodParameters(selectedRule.actions);
+
+  // Get the ai-generated args
+  let aiGeneratedActionItems: ActionWithAiArgs[] | undefined;
+  if (actionsRequiringAi.length) {
+    aiGeneratedActionItems = await getArgsAiResponse({
+      email,
+      user,
+      selectedRule,
+      actionsRequiringAi,
+    });
+  }
+
+  // Combine static args with ai-generated args
+  return selectedRule.actions.map((action) => {
+    const aiGeneratedActionItem = aiGeneratedActionItems?.find(
+      (item) => item.actionId === action.id,
+    );
+    return {
+      ...action,
+      ...aiGeneratedActionItem?.args,
+    };
+  });
+}
+
+async function getArgsAiResponse({
   email,
   user,
   selectedRule,
+  actionsRequiringAi,
 }: {
   email: EmailForLLM;
   user: Pick<User, "email" | "about"> & UserAIFields;
   selectedRule: RuleWithActions;
-}): Promise<ActionItem[] | undefined> {
+  actionsRequiringAi: ActionRequiringAi[];
+}): Promise<ActionWithAiArgs[] | undefined> {
   logger.log(
     `Generating args for rule ${selectedRule.name} (${selectedRule.id})`,
   );
 
-  const parameters = getToolParametersForRule(selectedRule.actions);
-
-  if (!Object.keys(parameters).length) {
+  if (!actionsRequiringAi.length) {
     logger.log(
       `Skipping. No parameters for rule ${selectedRule.name} (${selectedRule.id})`,
     );
@@ -126,10 +174,20 @@ ${stringifyEmail(email, 3000)}
 </email>
 `;
 
-  logger.log("Calling chat completion tools");
+  const parameters = getToolParametersForRule(actionsRequiringAi);
+  const zodParameters = z.object(
+    Object.fromEntries(
+      Object.entries(parameters).map(([key, { zodParameters }]) => [
+        key,
+        zodParameters,
+      ]),
+    ),
+  );
 
+  logger.log("Calling chat completion tools");
   logger.trace(`System: ${system}`);
   logger.trace(`Prompt: ${prompt}`);
+  logger.trace("Zod parameters:", JSON.stringify(zodParameters, null, 2));
 
   const aiResponse = await chatCompletionTools({
     userAi: user,
@@ -138,14 +196,7 @@ ${stringifyEmail(email, 3000)}
     tools: {
       apply_rule: {
         description: "Apply the rule with the given arguments",
-        parameters: z.object(
-          Object.fromEntries(
-            Object.entries(parameters).map(([key, { parameters }]) => [
-              key,
-              parameters,
-            ]),
-          ),
-        ),
+        parameters: zodParameters,
       },
     },
     label: "Args for rule",
@@ -159,20 +210,20 @@ ${stringifyEmail(email, 3000)}
   const toolCallArgs = toolCall.args;
   logger.trace(`Tool call args: ${JSON.stringify(toolCallArgs, null, 2)}`);
 
-  const actionItems = Object.entries(parameters).map(([key, { action }]) => {
-    const actionItem: ActionItem = {
-      type: action.type,
-      label: toolCallArgs[key].label || action.label,
-      subject: toolCallArgs[key].subject || action.subject,
-      content: toolCallArgs[key].content || action.content,
-      to: toolCallArgs[key].to || action.to,
-      cc: toolCallArgs[key].cc || action.cc,
-      bcc: toolCallArgs[key].bcc || action.bcc,
-    };
-    return actionItem;
+  const results = actionsRequiringAi.map((actionRequiringAi) => {
+    const actionParameter = Object.entries(parameters).find(
+      ([_, v]) => v.actionId === actionRequiringAi.actionId,
+    );
+
+    if (!actionParameter) return { ...actionRequiringAi, args: {} };
+
+    const [key] = actionParameter; // e.g. key = "draft_email"
+    const args: Record<string, string> = toolCallArgs[key]; // e.g. { label: "Draft", subject: "X", content: "Y" }
+
+    return { ...actionRequiringAi, args };
   });
 
-  logger.trace(`Action items: ${JSON.stringify(actionItems, null, 2)}`);
+  logger.trace(`Results: ${JSON.stringify(results, null, 2)}`);
 
-  return actionItems;
+  return results;
 }
diff --git a/apps/web/utils/ai/choose-rule/choose-and-execute.ts b/apps/web/utils/ai/choose-rule/choose-and-execute.ts
index a61718c8e3..6bf090c614 100644
--- a/apps/web/utils/ai/choose-rule/choose-and-execute.ts
+++ b/apps/web/utils/ai/choose-rule/choose-and-execute.ts
@@ -87,7 +87,20 @@ export async function saveExecutedRule(
   plannedAct: Awaited<ReturnType<typeof chooseRule>>,
 ) {
   const data: Prisma.ExecutedRuleCreateInput = {
-    actionItems: { createMany: { data: plannedAct.actionItems || [] } },
+    actionItems: {
+      createMany: {
+        data:
+          plannedAct.actionItems?.map((a) => ({
+            type: a.type,
+            label: a.label,
+            subject: a.subject,
+            content: a.content,
+            to: a.to,
+            cc: a.cc,
+            bcc: a.bcc,
+          })) || [],
+      },
+    },
     messageId,
     threadId,
     automated: !!plannedAct.rule?.automate,
diff --git a/apps/web/utils/ai/choose-rule/choose.ts b/apps/web/utils/ai/choose-rule/choose.ts
index 4a40eb4a38..1744fc865e 100644
--- a/apps/web/utils/ai/choose-rule/choose.ts
+++ b/apps/web/utils/ai/choose-rule/choose.ts
@@ -1,8 +1,5 @@
 import type { ActionItem } from "@/utils/ai/actions";
-import {
-  getArgsAiResponse,
-  getActionsWithParameters,
-} from "@/utils/ai/choose-rule/ai-choose-args";
+import { getActionItemsWithAiArgs } from "@/utils/ai/choose-rule/ai-choose-args";
 import { getAiResponse } from "@/utils/ai/choose-rule/ai-choose-rule";
 import type { UserAIFields } from "@/utils/llms/types";
 import type { RuleWithActions } from "@/utils/types";
@@ -42,26 +39,15 @@ export async function chooseRule(options: ChooseRuleOptions): Promise<
 
   if (!selectedRule) return { reason: aiResponse?.reason };
 
-  const shouldAiGenerateArgs =
-    getActionsWithParameters(selectedRule.actions).length > 0;
-
-  if (shouldAiGenerateArgs) {
-    const aiArgsResponse = await getArgsAiResponse({
-      ...options,
-      email,
-      selectedRule,
-    });
+  const actionItems = await getActionItemsWithAiArgs({
+    email,
+    user,
+    selectedRule,
+  });
 
-    return {
-      rule: selectedRule,
-      actionItems: aiArgsResponse || [],
-      reason: aiResponse?.reason,
-    };
-  } else {
-    return {
-      rule: selectedRule,
-      actionItems: selectedRule.actions,
-      reason: aiResponse?.reason,
-    };
-  }
+  return {
+    rule: selectedRule,
+    actionItems,
+    reason: aiResponse?.reason,
+  };
 }
diff --git a/apps/web/utils/logger.ts b/apps/web/utils/logger.ts
index c0bf2667b7..8d03f56f38 100644
--- a/apps/web/utils/logger.ts
+++ b/apps/web/utils/logger.ts
@@ -25,13 +25,18 @@ export function createScopeLogger(scope: string) {
   };
 
   return {
-    log: (message: LogMessage) => console.log(formatMessage("log", message)),
-    error: (message: LogMessage, error?: unknown, extra?: unknown) =>
-      console.error(formatMessage("error", message), error, extra),
-    warn: (message: LogMessage) => console.warn(formatMessage("warn", message)),
-    trace: (message: LogMessage) => {
+    log: (message: LogMessage, ...args: unknown[]) =>
+      console.log(formatMessage("log", message), ...args),
+
+    error: (message: LogMessage, ...args: unknown[]) =>
+      console.error(formatMessage("error", message), ...args),
+
+    warn: (message: LogMessage, ...args: unknown[]) =>
+      console.warn(formatMessage("warn", message), ...args),
+
+    trace: (message: LogMessage, ...args: unknown[]) => {
       if (process.env.NODE_ENV === "development") {
-        console.log(formatMessage("trace", message));
+        console.log(formatMessage("trace", message), ...args);
       }
     },
   };

From bec7cbfc50434b79c3b42194af5e6d0ab9d317da Mon Sep 17 00:00:00 2001
From: Eliezer Steinbock <3090527+elie222@users.noreply.github.com>
Date: Fri, 22 Nov 2024 17:39:17 -0500
Subject: [PATCH 7/8] Adjust file order

---
 .../utils/ai/choose-rule/ai-choose-args.ts    | 155 +++++++++---------
 1 file changed, 79 insertions(+), 76 deletions(-)

diff --git a/apps/web/utils/ai/choose-rule/ai-choose-args.ts b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
index e504d56e3e..db9ea1764a 100644
--- a/apps/web/utils/ai/choose-rule/ai-choose-args.ts
+++ b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
@@ -9,6 +9,7 @@ import {
 } from "@/utils/ai/choose-rule/stringify-email";
 import { type RuleWithActions, isDefined } from "@/utils/types";
 import { createScopeLogger } from "@/utils/logger";
+import { zodToJsonSchema } from "zod-to-json-schema";
 
 const logger = createScopeLogger("AI Choose Args");
 
@@ -22,80 +23,6 @@ export type ActionWithAiArgs = ActionRequiringAi & {
   args: Record<string, string>;
 };
 
-// Returns parameters for a zod.object for the rule that must be AI generated
-function getToolParametersForRule(actionsRequiringAi: ActionRequiringAi[]) {
-  // handle duplicate keys. e.g. "draft_email" and "draft_email" becomes: "draft_email" and "draft_email_2"
-  // this is quite an edge case but need to handle regardless for when it happens
-  const typeCount: Record<string, number> = {};
-  const parameters: Record<
-    string,
-    {
-      actionId: string;
-      zodParameters: z.ZodObject<Record<string, z.ZodString>>;
-    }
-  > = {};
-
-  for (const action of actionsRequiringAi) {
-    // count how many times we have already had this type
-    typeCount[action.type] = (typeCount[action.type] || 0) + 1;
-    const key =
-      typeCount[action.type] === 1
-        ? action.type
-        : `${action.type}_${typeCount[action.type]}`;
-    parameters[key] = {
-      actionId: action.actionId,
-      zodParameters: action.zodParameters,
-    };
-  }
-
-  return parameters;
-}
-
-// Returns the actions with the args that require ai-generation
-function getActionsWithZodParameters(actions: Action[]) {
-  return actions
-    .map((action) => {
-      const fields = getParameterFieldsForAction(action);
-
-      if (!Object.keys(fields).length) return;
-
-      const zodParameters = z.object(fields);
-
-      return {
-        actionId: action.id,
-        type: action.type,
-        zodParameters,
-      };
-    })
-    .filter(isDefined);
-}
-
-function getParameterFieldsForAction({
-  labelPrompt,
-  subjectPrompt,
-  contentPrompt,
-  toPrompt,
-  ccPrompt,
-  bccPrompt,
-}: Action) {
-  const fields: Record<string, z.ZodString> = {};
-
-  if (typeof labelPrompt === "string")
-    fields.label = z.string().describe(labelPrompt || "The email label");
-  if (typeof subjectPrompt === "string")
-    fields.subject = z.string().describe(subjectPrompt || "The email subject");
-  if (typeof contentPrompt === "string")
-    fields.content = z.string().describe(contentPrompt || "The email content");
-  if (typeof toPrompt === "string")
-    fields.to = z.string().describe(toPrompt || "The email recipient(s)");
-  if (typeof ccPrompt === "string")
-    fields.cc = z.string().describe(ccPrompt || "The cc recipient(s)");
-  if (typeof bccPrompt === "string")
-    fields.bcc = z.string().describe(bccPrompt || "The bcc recipient(s)");
-
-  return fields;
-}
-
 // Returns the action items with the ai-generated args where needed
 export async function getActionItemsWithAiArgs({
   email,
@@ -187,7 +114,7 @@ ${stringifyEmail(email, 3000)}
   logger.log("Calling chat completion tools");
   logger.trace(`System: ${system}`);
   logger.trace(`Prompt: ${prompt}`);
-  logger.trace("Zod parameters:", JSON.stringify(zodParameters, null, 2));
+  logger.trace("Zod parameters:", zodToJsonSchema(zodParameters));
 
   const aiResponse = await chatCompletionTools({
     userAi: user,
@@ -208,6 +135,7 @@ ${stringifyEmail(email, 3000)}
   if (!toolCall?.toolName) return;
 
   const toolCallArgs = toolCall.args;
+
   logger.trace(`Tool call args: ${JSON.stringify(toolCallArgs, null, 2)}`);
 
   const results = actionsRequiringAi.map((actionRequiringAi) => {
@@ -223,7 +151,82 @@ ${stringifyEmail(email, 3000)}
     return { ...actionRequiringAi, args };
   });
 
-  logger.trace(`Results: ${JSON.stringify(results, null, 2)}`);
+  const resultsForLogging = results.map(({ zodParameters, ...rest }) => rest);
+  logger.trace(`Results: ${JSON.stringify(resultsForLogging, null, 2)}`);
 
   return results;
 }
+
+// Returns parameters for a zod.object for the rule that must be AI generated
+function getToolParametersForRule(actionsRequiringAi: ActionRequiringAi[]) {
+  // handle duplicate keys. e.g. "draft_email" and "draft_email" becomes: "draft_email" and "draft_email_2"
+  // this is quite an edge case but need to handle regardless for when it happens
+  const typeCount: Record<string, number> = {};
+  const parameters: Record<
+    string,
+    {
+      actionId: string;
+      zodParameters: z.ZodObject<Record<string, z.ZodString>>;
+    }
+  > = {};
+
+  for (const action of actionsRequiringAi) {
+    // count how many times we have already had this type
+    typeCount[action.type] = (typeCount[action.type] || 0) + 1;
+    const key =
+      typeCount[action.type] === 1
+        ? action.type
+        : `${action.type}_${typeCount[action.type]}`;
+    parameters[key] = {
+      actionId: action.actionId,
+      zodParameters: action.zodParameters,
+    };
+  }
+
+  return parameters;
+}
+
+// Returns the actions with the args that require ai-generation
+function getActionsWithZodParameters(actions: Action[]) {
+  return actions
+    .map((action) => {
+      const fields = getParameterFieldsForAction(action);
+
+      if (!Object.keys(fields).length) return;
+
+      const zodParameters = z.object(fields);
+
+      return {
+        actionId: action.id,
+        type: action.type,
+        zodParameters,
+      };
+    })
+    .filter(isDefined);
+}
+
+function getParameterFieldsForAction({
+  labelPrompt,
+  subjectPrompt,
+  contentPrompt,
+  toPrompt,
+  ccPrompt,
+  bccPrompt,
+}: Action) {
+  const fields: Record<string, z.ZodString> = {};
+
+  if (typeof labelPrompt === "string")
+    fields.label = z.string().describe(labelPrompt || "The email label");
+  if (typeof subjectPrompt === "string")
+    fields.subject = z.string().describe(subjectPrompt || "The email subject");
+  if (typeof contentPrompt === "string")
+    fields.content = z.string().describe(contentPrompt || "The email content");
+  if (typeof toPrompt === "string")
+    fields.to = z.string().describe(toPrompt || "The email recipient(s)");
+  if (typeof ccPrompt === "string")
+    fields.cc = z.string().describe(ccPrompt || "The cc recipient(s)");
+  if (typeof bccPrompt === "string")
+    fields.bcc = z.string().describe(bccPrompt || "The bcc recipient(s)");
+
+  return fields;
+}

From a889a7d15ea426241a7bc26480e6a5fba5ad9cbb Mon Sep 17 00:00:00 2001
From: Eliezer Steinbock <3090527+elie222@users.noreply.github.com>
Date: Fri, 22 Nov 2024 17:44:51 -0500
Subject: [PATCH 8/8] adjust logger

---
 .../app/api/google/webhook/process-history.ts |  4 ++--
 apps/web/utils/ai/actions.ts                  |  6 +++++-
 .../utils/ai/choose-rule/ai-choose-args.ts    | 19 ++++++++++++-------
 apps/web/utils/logger.ts                      |  2 +-
 4 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/apps/web/app/api/google/webhook/process-history.ts b/apps/web/app/api/google/webhook/process-history.ts
index 5d4317b9af..94ab84f6db 100644
--- a/apps/web/app/api/google/webhook/process-history.ts
+++ b/apps/web/app/api/google/webhook/process-history.ts
@@ -24,9 +24,9 @@ import { runRulesOnMessage } from "@/utils/ai/choose-rule/run-rules";
 import { blockUnsubscribedEmails } from "@/app/api/google/webhook/block-unsubscribed-emails";
 import { categorizeSender } from "@/utils/actions/categorize";
 import { unwatchEmails } from "@/app/api/google/watch/controller";
-import { createScopeLogger } from "@/utils/logger";
+import { createScopedLogger } from "@/utils/logger";
 
-const logger = createScopeLogger("Process History");
+const logger = createScopedLogger("Process History");
 
 export async function processHistoryForUser(
   decodedData: {
diff --git a/apps/web/utils/ai/actions.ts b/apps/web/utils/ai/actions.ts
index 698d760813..8bf7c5ced9 100644
--- a/apps/web/utils/ai/actions.ts
+++ b/apps/web/utils/ai/actions.ts
@@ -9,6 +9,9 @@ import {
 } from "@/utils/gmail/label";
 import { markSpam } from "@/utils/gmail/spam";
 import type { Attachment } from "@/utils/types/mail";
+import { createScopedLogger } from "@/utils/logger";
+
+const logger = createScopedLogger("ai-actions");
 
 export type EmailForAction = {
   threadId: string;
@@ -369,7 +372,8 @@ export const runActionFunction = async (
   action: ActionItem,
   userEmail: string,
 ) => {
-  console.log("Running action", { actionType: action.type, userEmail });
+  logger.log("Running action", { actionType: action.type, userEmail });
+  logger.trace("Running action:", action);
 
   const { type, ...args } = action;
   switch (type) {
diff --git a/apps/web/utils/ai/choose-rule/ai-choose-args.ts b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
index db9ea1764a..174062fc47 100644
--- a/apps/web/utils/ai/choose-rule/ai-choose-args.ts
+++ b/apps/web/utils/ai/choose-rule/ai-choose-args.ts
@@ -8,10 +8,9 @@ import {
   stringifyEmail,
 } from "@/utils/ai/choose-rule/stringify-email";
 import { type RuleWithActions, isDefined } from "@/utils/types";
-import { createScopeLogger } from "@/utils/logger";
-import { zodToJsonSchema } from "zod-to-json-schema";
+import { createScopedLogger } from "@/utils/logger";
 
-const logger = createScopeLogger("AI Choose Args");
+const logger = createScopedLogger("AI Choose Args");
 
 export type ActionRequiringAi = {
   actionId: string;
@@ -48,7 +47,7 @@ export async function getActionItemsWithAiArgs({
   }
 
   // Combine static args with ai-generated args
-  return selectedRule.actions.map((action) => {
+  const results = selectedRule.actions.map((action) => {
     const aiGeneratedActionItem = aiGeneratedActionItems?.find(
       (item) => item.actionId === action.id,
     );
@@ -57,6 +56,12 @@ export async function getActionItemsWithAiArgs({
       ...aiGeneratedActionItem?.args,
     };
   });
+
+  logger.trace(
+    `getActionItemsWithAiArgs. Results: ${JSON.stringify(results, null, 2)}`,
+  );
+
+  return results;
 }
 
 async function getArgsAiResponse({
@@ -114,7 +119,7 @@ ${stringifyEmail(email, 3000)}
   logger.log("Calling chat completion tools");
   logger.trace(`System: ${system}`);
   logger.trace(`Prompt: ${prompt}`);
-  logger.trace("Zod parameters:", zodToJsonSchema(zodParameters));
+  // logger.trace("Zod parameters:", zodToJsonSchema(zodParameters));
 
   const aiResponse = await chatCompletionTools({
     userAi: user,
@@ -151,8 +156,8 @@ ${stringifyEmail(email, 3000)}
     return { ...actionRequiringAi, args };
   });
 
-  const resultsForLogging = results.map(({ zodParameters, ...rest }) => rest);
-  logger.trace(`Results: ${JSON.stringify(resultsForLogging, null, 2)}`);
+  // const resultsForLogging = results.map(({ zodParameters, ...rest }) => rest);
+  // logger.trace(`Results: ${JSON.stringify(resultsForLogging, null, 2)}`);
 
   return results;
 }
diff --git a/apps/web/utils/logger.ts b/apps/web/utils/logger.ts
index 8d03f56f38..09ecf97a57 100644
--- a/apps/web/utils/logger.ts
+++ b/apps/web/utils/logger.ts
@@ -9,7 +9,7 @@ const colors = {
   reset: "\x1b[0m",
 } as const;
 
-export function createScopeLogger(scope: string) {
+export function createScopedLogger(scope: string) {
   const formatMessage = (level: LogLevel, message: LogMessage) => {
     const prefix = `[${scope}]: `;