import { google } from "googleapis";
import { getAuthClient } from "./gmail-auth.js";
import { emailSchema, emailCategorySchema, categories } from "./schema.js";
import type { Email, EmailCategory } from "./schema.js";
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import {
pipeline,
ZeroShotClassificationPipeline,
type PipelineType,
} from "@huggingface/transformers";
let classifier: ZeroShotClassificationPipeline | null = null;
async function getClassifier() {
if (!classifier) {
classifier = await pipeline<"zero-shot-classification">(
"zero-shot-classification",
"Xenova/bart-large-mnli"
);
}
return classifier;
}
async function classifyEmail(email: Email): Promise<{
category: EmailCategory;
confidence: number;
}> {
try {
const classifier = await getClassifier();
const emailText = `Subject: ${email.subject}\n Snippet:${email.snippet}\n From: ${email.from}`;
const result = await classifier(emailText, [...categories]);
const category = Array.isArray(result)
? result?.[0]?.labels[0]
: result?.labels[0];
const confidence = Array.isArray(result)
? result?.[0]?.scores[0]
: result?.scores[0];
console.error("Classification from classifier: ", {
category,
confidence,
result,
});
const parsed = emailCategorySchema.safeParse(category);
if (!parsed.success) {
throw new Error(`Invalid category predicted: ${category}`);
}
return { category: parsed.data, confidence: confidence ?? 0 };
} catch (error) {
console.error("Classification failed:", error);
return { category: "unknown", confidence: 0 };
}
}
export async function getUnreadEmails(count: number = 5, server?: McpServer) {
try {
const auth = await getAuthClient();
const gmail = google.gmail({ version: "v1", auth });
const response = await gmail.users.messages.list({
userId: "me",
maxResults: count,
q: "is:unread",
});
let emails = [];
for (const message of response.data.messages || []) {
const emailResponse = await getEmail(message.id!);
const email = {
id: emailResponse.id,
threadId: emailResponse.threadId,
snippet: emailResponse.snippet,
from:
emailResponse.payload?.headers?.find(
(header) => header.name === "From"
)?.value || "",
to:
emailResponse.payload?.headers?.find((header) => header.name === "To")
?.value || "",
subject:
emailResponse.payload?.headers?.find(
(header) => header.name === "Subject"
)?.value || "",
date:
emailResponse.payload?.headers?.find(
(header) => header.name === "Date"
)?.value || "",
category: "unknown" as EmailCategory,
};
const parsedEmail = emailSchema.parse(email);
// Classify email
const classification = await classifyEmail(parsedEmail);
let category = classification.category;
if (classification.confidence < 0.7 && server) {
try {
const response = await server.server.createMessage({
messages: [
{
role: "user",
content: {
type: "text",
text: `Classify this email into exactly one category.
From: ${email.from}
Subject: ${email.subject}
Snippet: ${email.snippet}
Categories:
- urgent: Time-sensitive, requires immediate action
- newsletter: Marketing, promotional, or digest emails
- work: Professional, business-related communication
- personal: Personal correspondence from individuals
Respond with ONLY one word from: urgent, newsletter, work, personal`,
},
},
],
maxTokens: 5,
});
const aiResult =
response.content.type === "text" ? response.content.text : "";
if (aiResult && aiResult.trim().length > 0) {
const parsed = emailCategorySchema.safeParse(
aiResult.toLowerCase().trim()
);
if (parsed.success) {
category = parsed.data;
}
}
} catch (error) {
console.error("AI classification failed", error);
}
}
emails.push({ ...parsedEmail, category });
}
return emails;
} catch (error) {
console.error("Failed to fetch emails:", error);
throw error;
}
}
export async function getEmail(messageId: string) {
const auth = await getAuthClient();
const gmail = google.gmail({ version: "v1", auth });
const response = await gmail.users.messages.get({
userId: "me",
id: messageId,
format: "full",
});
return response.data;
}
export async function saveDraftReplyToGmail(draftEmail: string, email: Email) {
const auth = await getAuthClient();
const gmail = google.gmail({ version: "v1", auth });
const raw = Buffer.from(draftEmail).toString("base64url");
try {
const response = await gmail.users.drafts.create({
userId: "me",
requestBody: {
message: {
raw,
threadId: email.threadId,
},
},
});
return response.data;
} catch (error) {
console.error("Failed to save draft reply:", error);
throw error;
}
}
