Skip to content

Fixes #657

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
elie222 merged 5 commits into from
Aug 7, 2025
Merged

Fixes #657

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ac009a8
fix build
elie222 Aug 6, 2025
facb791
fix build
elie222 Aug 6, 2025
6052565
fix build
elie222 Aug 6, 2025
0f1f7fd
ai generate object fixes
elie222 Aug 6, 2025
5f54c27
fix
elie222 Aug 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 29 additions & 5 deletions apps/web/utils/ai/categorize-sender/ai-categorize-senders.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
import { z } from "zod";
import { chatCompletionObject } from "@/utils/llms";
import { isDefined } from "@/utils/types";
import type { EmailAccountWithAI } from "@/utils/llms/types";
import type { Category } from "@prisma/client";
import { formatCategoriesForPrompt } from "@/utils/ai/categorize-sender/format-categories";
import { createScopedLogger } from "@/utils/logger";
import { extractEmailAddress } from "@/utils/email";
import { getModel } from "@/utils/llms/model";
import { generateObject } from "ai";
import { saveAiUsage } from "@/utils/usage";

const logger = createScopedLogger("ai-categorize-senders");

Expand Down Expand Up @@ -89,15 +91,37 @@ ${formatCategoriesForPrompt(categories)}

logger.trace("Categorize senders", { system, prompt });

const aiResponse = await chatCompletionObject({
userAi: emailAccount.user,
// const aiResponse = await chatCompletionObject({
// userAi: emailAccount.user,
// system,
// prompt,
// schema: categorizeSendersSchema,
// userEmail: emailAccount.email,
// usageLabel: "Categorize senders bulk",
// });

const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);

const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema: categorizeSendersSchema,
userEmail: emailAccount.email,
usageLabel: "Categorize senders bulk",
providerOptions,
});

if (aiResponse.usage) {
await saveAiUsage({
email: emailAccount.email,
usage: aiResponse.usage,
provider,
model,
label: "Categorize senders bulk",
});
}

logger.trace("Categorize senders response", {
senders: aiResponse.object.senders,
});
Expand Down
34 changes: 29 additions & 5 deletions apps/web/utils/ai/categorize-sender/ai-categorize-single-sender.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
import { generateObject } from "ai";
import { z } from "zod";
import { chatCompletionObject } from "@/utils/llms";
import type { EmailAccountWithAI } from "@/utils/llms/types";
import type { Category } from "@prisma/client";
import { formatCategoriesForPrompt } from "@/utils/ai/categorize-sender/format-categories";
import { createScopedLogger } from "@/utils/logger";
import { getModel } from "@/utils/llms/model";
import { saveAiUsage } from "@/utils/usage";

const logger = createScopedLogger("aiCategorizeSender");

Expand Down Expand Up @@ -57,15 +59,37 @@ ${formatCategoriesForPrompt(categories)}

logger.trace("aiCategorizeSender", { system, prompt });

const aiResponse = await chatCompletionObject({
userAi: emailAccount.user,
// const aiResponse = await chatCompletionObject({
// userAi: emailAccount.user,
// system,
// prompt,
// schema: categorizeSenderSchema,
// userEmail: emailAccount.email,
// usageLabel: "Categorize sender",
// });

const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);

const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema: categorizeSenderSchema,
userEmail: emailAccount.email,
usageLabel: "Categorize sender",
providerOptions,
});

if (aiResponse.usage) {
await saveAiUsage({
email: emailAccount.email,
usage: aiResponse.usage,
provider,
model,
label: "Categorize sender",
});
}

if (!categories.find((c) => c.name === aiResponse.object.category))
return null;

Expand Down
77 changes: 52 additions & 25 deletions apps/web/utils/ai/choose-rule/ai-choose-rule.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import { z } from "zod";
import type { EmailAccountWithAI } from "@/utils/llms/types";
import { chatCompletionObject } from "@/utils/llms";
import { stringifyEmail } from "@/utils/stringify-email";
import type { EmailForLLM } from "@/utils/types";
import { createScopedLogger } from "@/utils/logger";
import type { ModelType } from "@/utils/llms/model";
import { getModel, type ModelType } from "@/utils/llms/model";
import { generateObject } from "ai";
import { saveAiUsage } from "@/utils/usage";
// import { Braintrust } from "@/utils/braintrust";

const logger = createScopedLogger("ai-choose-rule");
Expand Down Expand Up @@ -80,36 +81,62 @@ ${emailSection}

logger.trace("Input", { system, prompt });

const aiResponse = await chatCompletionObject({
userAi: emailAccount.user,
modelType,
messages: [
{
role: "system",
content: system,
// This will cache if the user has a very long prompt. Although usually won't do anything as it's hard for this prompt to reach 1024 tokens
// https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching#cache-limitations
// NOTE: Needs permission from AWS to use this. Otherwise gives error: "You do not have access to explicit prompt caching"
// Currently only available to select customers: https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html
// providerOptions: {
// bedrock: { cachePoint: { type: "ephemeral" } },
// anthropic: { cacheControl: { type: "ephemeral" } },
// },
},
{
role: "user",
content: prompt,
},
],
// const aiResponse = await chatCompletionObject({
// userAi: emailAccount.user,
// modelType,
// messages: [
// {
// role: "system",
// content: system,
// // This will cache if the user has a very long prompt. Although usually won't do anything as it's hard for this prompt to reach 1024 tokens
// // https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching#cache-limitations
// // NOTE: Needs permission from AWS to use this. Otherwise gives error: "You do not have access to explicit prompt caching"
// // Currently only available to select customers: https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html
// // providerOptions: {
// // bedrock: { cachePoint: { type: "ephemeral" } },
// // anthropic: { cacheControl: { type: "ephemeral" } },
// // },
// },
// {
// role: "user",
// content: prompt,
// },
// ],
// schema: z.object({
// reason: z.string(),
// ruleName: z.string().nullish(),
// noMatchFound: z.boolean().nullish(),
// }),
// userEmail: emailAccount.email,
// usageLabel: "Choose rule",
// });

const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);

const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema: z.object({
reason: z.string(),
ruleName: z.string().nullish(),
noMatchFound: z.boolean().nullish(),
}),
userEmail: emailAccount.email,
usageLabel: "Choose rule",
providerOptions,
});
Comment on lines +114 to 128
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Aug 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Pass modelType parameter to getModel call.

The modelType parameter (defaulted to "default" on line 23) should be passed to the getModel call to ensure the correct model type is used.

Apply this fix:

  const { provider, model, llmModel, providerOptions } = getModel(
    emailAccount.user,
+   modelType,
  );
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);
const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema: z.object({
reason: z.string(),
ruleName: z.string().nullish(),
noMatchFound: z.boolean().nullish(),
}),
userEmail: emailAccount.email,
usageLabel: "Choose rule",
providerOptions,
});
const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
modelType,
);
const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema: z.object({
reason: z.string(),
ruleName: z.string().nullish(),
noMatchFound: z.boolean().nullish(),
}),
providerOptions,
});
🤖 Prompt for AI Agents 
In apps/web/utils/ai/choose-rule/ai-choose-rule.ts around lines 114 to 128, the
call to getModel is missing the modelType parameter, which defaults to "default"
on line 23. Update the getModel call to include the modelType argument to ensure
the correct model type is used when retrieving the model, provider, llmModel,
and providerOptions.


if (aiResponse.usage) {
await saveAiUsage({
email: emailAccount.email,
usage: aiResponse.usage,
provider,
model,
label: "Choose rule",
});
}

logger.trace("Response", aiResponse.object);

// braintrust.insertToDataset({
Expand Down
34 changes: 29 additions & 5 deletions apps/web/utils/ai/choose-rule/ai-detect-recurring-pattern.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
import { z } from "zod";
import type { EmailAccountWithAI } from "@/utils/llms/types";
import { chatCompletionObject } from "@/utils/llms";
import type { EmailForLLM } from "@/utils/types";
import { stringifyEmail } from "@/utils/stringify-email";
import { createScopedLogger } from "@/utils/logger";
import { getModel } from "@/utils/llms/model";
import { generateObject } from "ai";
import { saveAiUsage } from "@/utils/usage";

const logger = createScopedLogger("detect-recurring-pattern");

Expand Down Expand Up @@ -98,15 +100,37 @@ ${stringifyEmail(email, 500)}
logger.trace("Input", { system, prompt });

try {
const aiResponse = await chatCompletionObject({
userAi: emailAccount.user,
// const aiResponse = await chatCompletionObject({
// userAi: emailAccount.user,
// system,
// prompt,
// schema,
// userEmail: emailAccount.email,
// usageLabel: "Detect recurring pattern",
// });

const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);

const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema,
userEmail: emailAccount.email,
usageLabel: "Detect recurring pattern",
providerOptions,
});

if (aiResponse.usage) {
await saveAiUsage({
email: emailAccount.email,
usage: aiResponse.usage,
provider,
model,
label: "Detect recurring pattern",
});
}

logger.trace("Response", aiResponse.object);

// braintrust.insertToDataset({
Expand Down
6 changes: 5 additions & 1 deletion apps/web/utils/ai/choose-rule/choose-args.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,11 @@ export async function getActionItemsWithAiArgs({
modelType,
});

return combineActionsWithAiArgs(selectedRule.actions, result, draft);
return combineActionsWithAiArgs(
selectedRule.actions,
result as ActionArgResponse,
draft,
);
}

function combineActionsWithAiArgs(
Expand Down
34 changes: 29 additions & 5 deletions apps/web/utils/ai/clean/ai-clean-select-labels.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import { z } from "zod";
import { chatCompletionObject } from "@/utils/llms";
import type { EmailAccountWithAI } from "@/utils/llms/types";
import { createScopedLogger } from "@/utils/logger";
import { getModel } from "@/utils/llms/model";
import { generateObject } from "ai";
import { saveAiUsage } from "@/utils/usage";

const logger = createScopedLogger("ai/clean/select-labels");

Expand Down Expand Up @@ -31,15 +33,37 @@ ${instructions}

logger.trace("Input", { system, prompt });

const aiResponse = await chatCompletionObject({
userAi: emailAccount.user,
// const aiResponse = await chatCompletionObject({
// userAi: emailAccount.user,
// system,
// prompt,
// schema,
// userEmail: emailAccount.email,
// usageLabel: "Clean - Select Labels",
// });

const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);

const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema,
userEmail: emailAccount.email,
usageLabel: "Clean - Select Labels",
providerOptions,
});
Comment on lines +45 to 55
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Aug 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Implement error handling for AI generation failures.

The refactored code lacks fallback mechanisms for AI failures, which is a requirement per the coding guidelines for LLM-related functions.

Consider adding error handling:

+ try {
    const aiResponse = await generateObject({
      model: llmModel,
      system,
      prompt,
      schema,
      providerOptions,
    });
+ } catch (error) {
+   logger.error("AI generation failed", { error });
+   // Return empty array as fallback for labels
+   return [];
+ }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);
const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema,
userEmail: emailAccount.email,
usageLabel: "Clean - Select Labels",
providerOptions,
});
const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);
try {
const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema,
providerOptions,
});
} catch (error) {
logger.error("AI generation failed", { error });
// Return empty array as fallback for labels
return [];
}
🤖 Prompt for AI Agents 
In apps/web/utils/ai/clean/ai-clean-select-labels.ts around lines 45 to 55, the
code calls generateObject without error handling, which risks unhandled failures
during AI generation. Wrap the generateObject call in a try-catch block to catch
any errors, and implement a fallback mechanism or return a default value in the
catch block to ensure graceful failure handling as per coding guidelines.


if (aiResponse.usage) {
await saveAiUsage({
email: emailAccount.email,
usage: aiResponse.usage,
provider,
model,
label: "Clean - Select Labels",
});
}

logger.trace("Result", { response: aiResponse.object });

return aiResponse.object.labels;
Expand Down
34 changes: 29 additions & 5 deletions apps/web/utils/ai/clean/ai-clean.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
import { z } from "zod";
import { chatCompletionObject } from "@/utils/llms";
import { generateObject } from "ai";
import type { EmailAccountWithAI } from "@/utils/llms/types";
import { createScopedLogger } from "@/utils/logger";
import type { EmailForLLM } from "@/utils/types";
import { stringifyEmailSimple } from "@/utils/stringify-email";
import { formatDateForLLM, formatRelativeTimeForLLM } from "@/utils/date";
import { preprocessBooleanLike } from "@/utils/zod";
import { getModel } from "@/utils/llms/model";
import { saveAiUsage } from "@/utils/usage";
// import { Braintrust } from "@/utils/braintrust";

const logger = createScopedLogger("ai/clean");
Expand Down Expand Up @@ -92,15 +94,37 @@ The current date is ${currentDate}.

logger.trace("Input", { system, prompt });

const aiResponse = await chatCompletionObject({
userAi: emailAccount.user,
// const aiResponse = await chatCompletionObject({
// userAi: emailAccount.user,
// system,
// prompt,
// schema,
// userEmail: emailAccount.email,
// usageLabel: "Clean",
// });

const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);

const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema,
userEmail: emailAccount.email,
usageLabel: "Clean",
providerOptions,
});
Comment on lines +106 to 116
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Aug 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add error handling and fallback for AI generation failures.

The function lacks error handling for AI generation failures, which could cause the entire email processing to fail. Given this function makes critical decisions about archiving emails, a fallback strategy is essential.

Add comprehensive error handling:

+ try {
    const aiResponse = await generateObject({
      model: llmModel,
      system,
      prompt,
      schema,
      providerOptions,
    });
+ } catch (error) {
+   logger.error("AI generation failed for email archiving", { 
+     error, 
+     messageId,
+     emailAccount: emailAccount.email 
+   });
+   // Fallback: don't archive when AI fails (safer default)
+   return { archive: false };
+ }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);
const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema,
userEmail: emailAccount.email,
usageLabel: "Clean",
providerOptions,
});
const { provider, model, llmModel, providerOptions } = getModel(
emailAccount.user,
);
try {
const aiResponse = await generateObject({
model: llmModel,
system,
prompt,
schema,
providerOptions,
});
} catch (error) {
logger.error("AI generation failed for email archiving", {
error,
messageId,
emailAccount: emailAccount.email,
});
// Fallback: don't archive when AI fails (safer default)
return { archive: false };
}
🤖 Prompt for AI Agents 
In apps/web/utils/ai/clean/ai-clean.ts around lines 106 to 116, the call to
generateObject lacks error handling, risking failure of the entire email
processing if AI generation fails. Wrap the generateObject call in a try-catch
block to catch any errors, log or handle the error appropriately, and implement
a fallback strategy such as returning a default response or skipping AI-based
decisions to ensure the function continues safely.


if (aiResponse.usage) {
await saveAiUsage({
email: emailAccount.email,
usage: aiResponse.usage,
provider,
model,
label: "Clean",
});
}

logger.trace("Result", { response: aiResponse.object });

// braintrust.insertToDataset({
Expand Down
Loading
Loading