From e6736b820131aaa5a069c74becec49701aa1bdf1 Mon Sep 17 00:00:00 2001 From: Tony Chen Date: Mon, 29 Sep 2025 19:22:01 +1300 Subject: [PATCH] feat: adding OpenAI support --- src/index.ts | 647 +++++++++++++++++++++++------------- wrangler.toml.example | 19 ++ wrangler.toml.example.clear | 4 + 3 files changed, 431 insertions(+), 239 deletions(-) diff --git a/src/index.ts b/src/index.ts index 72a0ee0..930a899 100644 --- a/src/index.ts +++ b/src/index.ts @@ -6,103 +6,172 @@ created on: 2024 Oct 07 Last updated: 2024 Oct 07 */ import { WorkerEntrypoint } from "cloudflare:workers"; -import {RPCEmailMessage} from './rpcEmail'; +import { RPCEmailMessage } from "./rpcEmail"; -import indexHtml from './index.html'; +import indexHtml from "./index.html"; interface Env { - // If you set another name in wrangler.toml as the value for 'binding', - // replace "DB" with the variable name you defined. - DB: D1Database; - FrontEndAdminID: string; - FrontEndAdminPassword: string; - barkTokens: string; - barkUrl: string; - GoogleAPIKey: string; - UseBark: string; + // If you set another name in wrangler.toml as the value for 'binding', + // replace "DB" with the variable name you defined. + DB: D1Database; + FrontEndAdminID: string; + FrontEndAdminPassword: string; + barkTokens: string; + barkUrl: string; + GoogleAPIKey: string; + UseBark: string; + GeminiModel?: string; + OpenAIAPIKey?: string; + OpenAIBaseUrl?: string; + OpenAIModel?: string; +} + +interface ModelResponse { + ok: boolean; + status: number; + statusText: string; + payload: unknown; +} + +const DEFAULT_GEMINI_MODEL = "gemini-2.0-flash"; +const DEFAULT_OPENAI_BASE_URL = "https://api.openai.com"; +const DEFAULT_OPENAI_MODEL = "gpt-4o-mini"; + +// Normalize model output into JSON text // 将模型输出规范化为可解析的 JSON 文本 +function extractJsonFromText(rawText: string): Record | null { + let candidate = rawText.trim(); + const jsonMatch = candidate.match(/```json\s*([\s\S]*?)\s*```/); + if (jsonMatch && jsonMatch[1]) { + candidate = jsonMatch[1].trim(); + console.log(`Extracted JSON Text: "${candidate}"`); + } else { + console.log(`Assuming entire text is JSON: "${candidate}"`); + } + + try { + return JSON.parse(candidate); + } catch (parseError) { + console.error("JSON parsing error:", parseError); + console.log(`Problematic JSON Text: "${candidate}"`); + return null; + } +} + +// Read structured content from Gemini response // 读取 Gemini 响应中的结构化文本 +function extractGeminiText(payload: unknown): string | null { + const candidate = (payload as any)?.candidates?.[0]?.content?.parts?.[0]; + if (!candidate || typeof candidate.text !== "string") { + console.error("Gemini response is missing expected data structure"); + return null; + } + return candidate.text; +} + +// Read structured content from OpenAI response // 读取 OpenAI 响应中的结构化文本 +function extractOpenAIText(payload: unknown): string | null { + const message = (payload as any)?.choices?.[0]?.message; + if (!message) { + console.error("OpenAI response is missing choices"); + return null; + } + + const content = message.content; + if (typeof content === "string") { + return content; + } + + if (Array.isArray(content)) { + const textPart = content.find((part: any) => part?.type === "text" && typeof part.text === "string"); + if (textPart) { + return textPart.text; + } + } + + if (typeof message?.content?.[0]?.text === "string") { + return message.content[0].text; + } + + console.error("OpenAI response is missing text content"); + return null; +} + +// Remove trailing slashes to avoid double separators // 移除结尾斜杠防止重复拼接 +function normaliseBaseUrl(baseUrl: string): string { + return baseUrl.replace(/\/+$/, ""); } export default class extends WorkerEntrypoint { - async fetch(request: Request): Promise { - const env: Env = this.env; - const ctx: ExecutionContext = this.ctx; - // 将依赖 env 的常量移到函数内部 - const FrontEndAdminID = env.FrontEndAdminID; - const FrontEndAdminPassword = env.FrontEndAdminPassword; + async fetch(request: Request): Promise { + const env: Env = this.env; + const FrontEndAdminID = env.FrontEndAdminID; + const FrontEndAdminPassword = env.FrontEndAdminPassword; - // 提取 Authorization 头 - const authHeader = request.headers.get('Authorization'); + // Basic-auth gate for the admin console // 使用 Basic Auth 保护管理界面 + const authHeader = request.headers.get("Authorization"); - // 如果没有 Authorization 头，提示进行身份验证 - if (!authHeader) { - return new Response('Unauthorized', { - status: 401, - headers: { - 'WWW-Authenticate': 'Basic realm="User Visible Realm"', - }, - }); - } + if (!authHeader) { + return new Response("Unauthorized", { + status: 401, + headers: { + "WWW-Authenticate": "Basic realm=\"User Visible Realm\"", + }, + }); + } - // 检查 Authorization 头是否使用 Basic 认证 - if (!authHeader.startsWith('Basic ')) { - return new Response('Unauthorized', { - status: 401, - headers: { - 'WWW-Authenticate': 'Basic realm="User Visible Realm"', - }, - }); - } + if (!authHeader.startsWith("Basic ")) { + return new Response("Unauthorized", { + status: 401, + headers: { + "WWW-Authenticate": "Basic realm=\"User Visible Realm\"", + }, + }); + } - // 解码 base64 编码的凭据 - const base64Credentials = authHeader.substring('Basic '.length); - const decodedCredentials = atob(base64Credentials); + const base64Credentials = authHeader.substring("Basic ".length); + const decodedCredentials = atob(base64Credentials); + const [username, password] = decodedCredentials.split(":"); - // 将凭据分割为用户名和密码 - const [username, password] = decodedCredentials.split(':'); + if (username !== FrontEndAdminID || password !== FrontEndAdminPassword) { + return new Response("Unauthorized", { + status: 401, + headers: { + "WWW-Authenticate": "Basic realm=\"User Visible Realm\"", + }, + }); + } - // 验证凭据 - if ( - username !== FrontEndAdminID || - password !== FrontEndAdminPassword - ) { - return new Response('Unauthorized', { - status: 401, - headers: { - 'WWW-Authenticate': 'Basic realm="User Visible Realm"', - }, - }); - } + try { + const { results } = await env.DB.prepare( + "SELECT from_org, to_addr, topic, code, created_at FROM code_mails ORDER BY created_at DESC" + ).all(); - try { - const { results } = await env.DB.prepare( - 'SELECT from_org, to_addr, topic, code, created_at FROM code_mails ORDER BY created_at DESC' - ).all(); + let dataHtml = ""; + for (const row of results) { + const codeLinkParts = row.code.split(","); + let codeLinkContent; - let dataHtml = ''; - for (const row of results) { - const codeLinkParts = row.code.split(','); - let codeLinkContent; + if (codeLinkParts.length > 1) { + const [code, link] = codeLinkParts; + codeLinkContent = `${code}
${row.topic}`; + } else if (row.code.startsWith("http")) { + codeLinkContent = `${row.topic}`; + } else { + codeLinkContent = row.code; + } - if (codeLinkParts.length > 1) { - const [code, link] = codeLinkParts; - codeLinkContent = `${code}
${row.topic}`; - } else if (row.code.startsWith('http')) { - codeLinkContent = `${row.topic}`; - } else { - codeLinkContent = row.code; - } - - dataHtml += ` + dataHtml += ` ${row.from_org} ${row.to_addr} ${row.topic} ${codeLinkContent} ${row.created_at} `; - } + } - let responseHtml = indexHtml - .replace('{{TABLE_HEADERS}}', ` + const responseHtml = indexHtml + .replace( + "{{TABLE_HEADERS}}", + ` From To @@ -110,48 +179,150 @@ export default class extends WorkerEntrypoint { Code/Link Receive Time (GMT) - `) - .replace('{{DATA}}', dataHtml); + ` + ) + .replace("{{DATA}}", dataHtml); - return new Response(responseHtml, { - headers: { - 'Content-Type': 'text/html', - }, - }); - } catch (error) { - console.error('Error querying database:', error); - return new Response('Internal Server Error', { status: 500 }); - } - } + return new Response(responseHtml, { + headers: { + "Content-Type": "text/html", + }, + }); + } catch (error) { + console.error("Error querying database:", error); + return new Response("Internal Server Error", { status: 500 }); + } + } - // 主要功能 - async email(message: ForwardableEmailMessage): Promise { + // Call Gemini for primary extraction // 调用 Gemini 执行主解析 + private async callGemini(prompt: string, apiKey: string, model: string): Promise { + const endpoint = `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent?key=${apiKey}`; + const init: RequestInit = { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + contents: [ + { + parts: [{ text: prompt }], + }, + ], + }), + }; - const env: Env = this.env; - const useBark = env.UseBark.toLowerCase() === 'true'; // true or false - const GoogleAPIKey = env.GoogleAPIKey; // "xxxxxxxxxxxxxxxxxxxxxxxx" + try { + const response = await fetch(endpoint, init); + let payload: unknown = null; + try { + payload = await response.json(); + } catch (parseError) { + console.error("Gemini JSON parse error:", parseError); + } + return { + ok: response.ok, + status: response.status, + statusText: response.statusText, + payload, + }; + } catch (error) { + console.error("Gemini request threw:", error); + return { + ok: false, + status: 0, + statusText: "FETCH_ERROR", + payload: null, + }; + } + } - const rawEmail = message instanceof RPCEmailMessage ? (message).rawEmail : await new Response(message.raw).text(); - const message_id = message.headers.get("Message-ID"); + // Call OpenAI as the fallback extractor // 调用 OpenAI 作为兜底解析 + private async callOpenAI( + prompt: string, + baseUrl: string, + apiKey: string, + model: string + ): Promise { + const endpoint = `${normaliseBaseUrl(baseUrl)}/v1/chat/completions`; + const init: RequestInit = { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${apiKey}`, + }, + body: JSON.stringify({ + model, + temperature: 0, + response_format: { type: "json_object" }, + messages: [ + { + role: "system", + content: "You return only valid JSON that matches the requested schema.", + }, + { + role: "user", + content: prompt, + }, + ], + }), + }; - // 将电子邮件保存到数据库 - const {success} = await env.DB.prepare( - `INSERT INTO raw_mails (from_addr, to_addr, raw, message_id) VALUES (?, ?, ?, ?)` - ).bind( - message.from, message.to, rawEmail, message_id // 将电子邮件详细信息绑定到 SQL 语句 - ).run(); + try { + const response = await fetch(endpoint, init); + let payload: unknown = null; + try { + payload = await response.json(); + } catch (parseError) { + console.error("OpenAI JSON parse error:", parseError); + } + return { + ok: response.ok, + status: response.status, + statusText: response.statusText, + payload, + }; + } catch (error) { + console.error("OpenAI request threw:", error); + return { + ok: false, + status: 0, + statusText: "FETCH_ERROR", + payload: null, + }; + } + } - // 检查电子邮件是否成功保存 - if (!success) { - message.setReject(`Failed to save message from ${message.from} to ${message.to}`); // 如果保存失败，则拒绝消息 - console.log(`Failed to save message from ${message.from} to ${message.to}`); // 记录保存失败 - } + // Primary email handler // 主要邮件处理入口 + async email(message: ForwardableEmailMessage): Promise { + const env: Env = this.env; + const useBark = env.UseBark.toLowerCase() === "true"; + const googleApiKey = env.GoogleAPIKey; + const geminiModel = env.GeminiModel || DEFAULT_GEMINI_MODEL; + const openAIKey = env.OpenAIAPIKey; + const openAIBaseUrl = env.OpenAIBaseUrl || DEFAULT_OPENAI_BASE_URL; + const openAIModel = env.OpenAIModel || DEFAULT_OPENAI_MODEL; - // 调用AI，让AI抓取验证码，让AI返回`title`和`code` - // title: 邮件是哪个公司/组织发来的验证码, 比如'Netflix' - // code: 验证码/链接/密码，比如'123456'or'https://example.com/verify?code=123456',如都有则返回'code, link' - // topic: 邮件主题，比如'line register verification' - const aiPrompt = ` + // Pull raw email content // 获取原始邮件内容 + const rawEmail = + message instanceof RPCEmailMessage + ? (message as RPCEmailMessage).rawEmail + : await new Response(message.raw).text(); + const messageId = message.headers.get("Message-ID"); + + // Persist raw mail payload for auditing // 将原始邮件持久化以便审计 + const { success } = await env.DB.prepare( + "INSERT INTO raw_mails (from_addr, to_addr, raw, message_id) VALUES (?, ?, ?, ?)" + ) + .bind(message.from, message.to, rawEmail, messageId) + .run(); + + if (!success) { + message.setReject(`Failed to save message from ${message.from} to ${message.to}`); + console.log(`Failed to save message from ${message.from} to ${message.to}`); + } + + // Prompt instructs model how to format extraction // 提示词说明提取格式和字段要求 + const aiPrompt = ` Email content: ${rawEmail}. Please read the email and extract the following information: @@ -177,146 +348,144 @@ export default class extends WorkerEntrypoint { } `; - try { - // 添加重试机制 - const maxRetries = 3; - let retryCount = 0; - let extractedData = null; + try { + const maxRetries = 3; + let retryCount = 0; + let extractedData: Record | null = null; + let geminiStatusError = false; - while (retryCount < maxRetries && !extractedData) { - // 调用 Google AI API 来获取 title, code, topic - const aiResponse = await fetch(`https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key=${GoogleAPIKey}`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json' - }, - body: JSON.stringify({ - "contents": [ - { - "parts": [ - {"text": aiPrompt} - ] - } - ] - }) - }); + // First attempt Gemini with simple retries // 先使用 Gemini 并进行简单重试 + while (retryCount < maxRetries && !extractedData && !geminiStatusError) { + const geminiResult = await this.callGemini(aiPrompt, googleApiKey, geminiModel); + console.log(`Gemini response attempt ${retryCount + 1}:`, geminiResult.payload); - const aiData = await aiResponse.json(); - console.log(`AI response attempt ${retryCount + 1}:`, aiData); - // 检测ai返回格式是否正确 - if ( - aiData && - aiData.candidates && - aiData.candidates[0] && - aiData.candidates[0].content && - aiData.candidates[0].content.parts && - aiData.candidates[0].content.parts[0] - ) { - let extractedText = aiData.candidates[0].content.parts[0].text; - console.log(`Extracted Text before parsing: "${extractedText}"`); + if (!geminiResult.ok) { + geminiStatusError = true; + console.error( + `Gemini request failed with status ${geminiResult.status} ${geminiResult.statusText}` + ); + break; + } - // Use regex to extract JSON content from code blocks - const jsonMatch = extractedText.match(/```json\s*([\s\S]*?)\s*```/); - if (jsonMatch && jsonMatch[1]) { - extractedText = jsonMatch[1].trim(); - console.log(`Extracted JSON Text: "${extractedText}"`); - } else { - // If no code block, assume the entire text is JSON - extractedText = extractedText.trim(); - console.log(`Assuming entire text is JSON: "${extractedText}"`); - } + const extractedText = extractGeminiText(geminiResult.payload); + if (extractedText) { + console.log(`Extracted Text before parsing: "${extractedText}"`); + const parsed = extractJsonFromText(extractedText); + if (parsed) { + extractedData = parsed; + console.log("Parsed Extracted Data:", extractedData); + break; + } + } - // Parse - try { - extractedData = JSON.parse(extractedText); - console.log(`Parsed Extracted Data:`, extractedData); - } catch (parseError) { - console.error("JSON parsing error:", parseError); - console.log(`Problematic JSON Text: "${extractedText}"`); - } + retryCount += 1; + if (retryCount < maxRetries) { + console.log("Retrying AI request..."); + } else { + console.error("Max retries reached. Unable to get valid AI response."); + } + } - } else { - console.error("AI response is missing expected data structure"); - } + if (!extractedData && geminiStatusError && openAIKey) { + // Gemini failed with non-200, switch to OpenAI fallback // 当 Gemini 返回非 200 时切换到 OpenAI 兜底 + const openAIResult = await this.callOpenAI(aiPrompt, openAIBaseUrl, openAIKey, openAIModel); + console.log("OpenAI fallback response:", openAIResult.payload); + if (!openAIResult.ok) { + console.error( + `OpenAI fallback failed with status ${openAIResult.status} ${openAIResult.statusText}` + ); + } else { + const openAIText = extractOpenAIText(openAIResult.payload); + if (openAIText) { + console.log(`OpenAI text before parsing: "${openAIText}"`); + const parsed = extractJsonFromText(openAIText); + if (parsed) { + extractedData = parsed; + console.log("Parsed Extracted Data from OpenAI:", extractedData); + } + } + } + } else if (!extractedData && geminiStatusError) { + console.error("Gemini request failed and OpenAI fallback is not configured."); + } - if (!extractedData) { - retryCount++; - if (retryCount < maxRetries) { - console.log("Retrying AI request..."); - } else { - console.error("Max retries reached. Unable to get valid AI response."); - } - } - } + if (extractedData) { + // Only persist when a code exists // 仅在存在验证码时写入数据库 + if ((extractedData as any).codeExist === 1) { + const title = (extractedData as any).title || "Unknown Organization"; + const code = (extractedData as any).code || "No Code Found"; + const topic = (extractedData as any).topic || "No Topic Found"; - // extract formatted data - if (extractedData) { - if (extractedData.codeExist === 1) { - const title = extractedData.title || "Unknown Organization"; - const code = extractedData.code || "No Code Found"; - const topic = extractedData.topic || "No Topic Found"; + // Store parsed metadata for UI display // 保存解析结果供前端展示 + const { success: codeMailSuccess } = await env.DB.prepare( + "INSERT INTO code_mails (from_addr, from_org, to_addr, code, topic, message_id) VALUES (?, ?, ?, ?, ?, ?)" + ) + .bind(message.from, title, message.to, code, topic, messageId) + .run(); - // save extracted data to the database - const { success: codeMailSuccess } = await env.DB.prepare( - `INSERT INTO code_mails (from_addr, from_org, to_addr, code, topic, message_id) VALUES (?, ?, ?, ?, ?, ?)` - ).bind( - message.from, title, message.to, code, topic, message_id - ).run(); + if (!codeMailSuccess) { + message.setReject( + `Failed to save extracted code for message from ${message.from} to ${message.to}` + ); + console.log( + `Failed to save extracted code for message from ${message.from} to ${message.to}` + ); + } - if (!codeMailSuccess) { - message.setReject(`Failed to save extracted code for message from ${message.from} to ${message.to}`); - console.log(`Failed to save extracted code for message from ${message.from} to ${message.to}`); - } + if (useBark) { + // Fan-out Bark notifications for each token // 为每个 token 发送 Bark 推送 + const barkUrl = env.barkUrl; + const barkTokens = env.barkTokens + .replace(/^\[|\]$/g, "") + .split(",") + .map((token) => token.trim()); - // Send title and code to Bark using GET request for each token - if (useBark) { - const barkUrl = env.barkUrl; // "https://api.day.app" - // [token1, token2] - const barkTokens = env.barkTokens - .replace(/^\[|\]$/g, '') - .split(',') - .map(token => token.trim()); + const barkUrlEncodedTitle = encodeURIComponent(title); + const barkUrlEncodedCode = encodeURIComponent(code); - const barkUrlEncodedTitle = encodeURIComponent(title); - const barkUrlEncodedCode = encodeURIComponent(code); + for (const token of barkTokens) { + const barkRequestUrl = `${barkUrl}/${token}/${barkUrlEncodedTitle}/${barkUrlEncodedCode}`; - for (const token of barkTokens) { - const barkRequestUrl = `${barkUrl}/${token}/${barkUrlEncodedTitle}/${barkUrlEncodedCode}`; + const barkResponse = await fetch(barkRequestUrl, { + method: "GET", + }); - const barkResponse = await fetch(barkRequestUrl, { - method: "GET" - }); + if (barkResponse.ok) { + console.log( + `Successfully sent notification to Bark for token ${token} for message from ${message.from} to ${message.to}` + ); + const responseData = await barkResponse.json(); + console.log("Bark response:", responseData); + } else { + console.error( + `Failed to send notification to Bark for token ${token}: ${barkResponse.status} ${barkResponse.statusText}` + ); + } + } + } + } else { + console.log("No code found in this email, skipping Bark notification."); + } + } else { + console.error("Failed to extract data from AI response after retries."); + } + } catch (e) { + console.error("Error calling AI or saving to database:", e); + } + } - if (barkResponse.ok) { - console.log(`Successfully sent notification to Bark for token ${token} for message from ${message.from} to ${message.to}`); - const responseData = await barkResponse.json(); - console.log("Bark response:", responseData); - } else { - console.error(`Failed to send notification to Bark for token ${token}: ${barkResponse.status} ${barkResponse.statusText}`); - } - } - } - - - } else { - console.log("No code found in this email, skipping Bark notification."); - } - } else { - console.error("Failed to extract data from AI response after retries."); - } - } catch (e) { - console.error("Error calling AI or saving to database:", e); - } - } - - // 暴露rpc接口，处理来自其他worker的邮件请求 - async rpcEmail(requestBody: string): Promise { - console.log(`Received RPC email , request body: ${requestBody}`); - const bodyObject =JSON.parse(requestBody); - const headersObject = bodyObject.headers; - const headers = new Headers(headersObject); - const rpcEmailMessage: RPCEmailMessage = new RPCEmailMessage(bodyObject.from, bodyObject.to, bodyObject.rawEmail, headers); - await this.email(rpcEmailMessage); - } - -} \ No newline at end of file + // Expose RPC helper for other workers // 暴露 RPC 接口供其他 Worker 调用 + async rpcEmail(requestBody: string): Promise { + console.log(`Received RPC email , request body: ${requestBody}`); + const bodyObject = JSON.parse(requestBody); + const headersObject = bodyObject.headers; + const headers = new Headers(headersObject); + const rpcEmailMessage: RPCEmailMessage = new RPCEmailMessage( + bodyObject.from, + bodyObject.to, + bodyObject.rawEmail, + headers + ); + await this.email(rpcEmailMessage); + } +} diff --git a/wrangler.toml.example b/wrangler.toml.example index 5cfdaae..798ff6a 100644 --- a/wrangler.toml.example +++ b/wrangler.toml.example @@ -58,6 +58,25 @@ GoogleAPIKey = "xxxxxxxxxxx" # your google AI API, you can have one from https://aistudio.google.com # 你的谷歌AI API，你可以从 https://aistudio.google.com 获取 +GeminiModel = "gemini-2.0-flash" +# GeminiModel: String +# optional override for the Gemini model, defaults to "gemini-2.0-flash" +# 可选配置，覆盖默认的 Gemini 模型名称，默认值为 "gemini-2.0-flash" + +OpenAIAPIKey = "" +# OpenAIAPIKey: String +# optional key for OpenAI fallback when Gemini fails with a non-200 response +# 可选配置，当 Gemini 返回非 200 时使用 OpenAI 作为兜底 + +OpenAIBaseUrl = "https://api.openai.com" +# OpenAIBaseUrl: String +# optional, change only if you proxy the OpenAI API +# 可选配置，仅在自建代理时更改 + +OpenAIModel = "gpt-4o-mini" +# OpenAIModel: String +# optional override for the OpenAI model, defaults to "gpt-4o-mini" +# 可选配置，覆盖默认的 OpenAI 模型 # Bind a D1 database. D1 is Cloudflare’s native serverless SQL database. # Docs: https://developers.cloudflare.com/workers/wrangler/configuration/#d1-databases diff --git a/wrangler.toml.example.clear b/wrangler.toml.example.clear index 1c8e44a..90f4470 100644 --- a/wrangler.toml.example.clear +++ b/wrangler.toml.example.clear @@ -13,6 +13,10 @@ barkTokens = "[token1, token2]" FrontEndAdminID = "admin" FrontEndAdminPassword = "password" GoogleAPIKey = "xxxxxxxxxxx" +GeminiModel = "gemini-2.0-flash" +OpenAIAPIKey = "" +OpenAIBaseUrl = "https://api.openai.com" +OpenAIModel = "gpt-4o-mini" [[d1_databases]] binding = "DB"