diff --git a/workers/agent/index.ts b/workers/agent/index.ts index ab6f59b9..77fb6e87 100644 --- a/workers/agent/index.ts +++ b/workers/agent/index.ts @@ -13,6 +13,7 @@ import { createWorkersAI } from "workers-ai-provider"; import { z } from "zod"; import type { EmailFull, EmailMetadata } from "../lib/schemas"; import { verifyDraft, isPromptInjection } from "../lib/ai"; +import { getAIModel } from "../lib/ai-config"; import { getMailboxStub, stripHtmlToText, @@ -281,7 +282,7 @@ export class EmailAgent extends AIChatAgent { const systemPrompt = await getSystemPrompt(env, mailboxId); const result = streamText({ - model: workersai("@cf/moonshotai/kimi-k2.5"), + model: workersai(getAIModel("emailAgent", env)), system: systemPrompt, messages: await convertToModelMessages(this.messages), tools, @@ -347,7 +348,9 @@ export class EmailAgent extends AIChatAgent { try { const email = (await stub.getEmail(emailData.emailId)) as EmailFull | null; if (email?.body) { - const isInjection = await isPromptInjection(env.AI, email.body); + const isInjection = await isPromptInjection(env.AI, email.body, { + model: getAIModel("promptInjectionScanner", env), + }); if (isInjection) { console.warn("Skipping auto-draft due to detected prompt injection:", emailData.emailId); @@ -395,7 +398,9 @@ export class EmailAgent extends AIChatAgent { // could plant an injection in an earlier email in the thread // that gets included in the agent's prompt. if (threadContext) { - const threadInjection = await isPromptInjection(env.AI, threadContext); + const threadInjection = await isPromptInjection(env.AI, threadContext, { + model: getAIModel("promptInjectionScanner", env), + }); if (threadInjection) { console.warn("Skipping auto-draft due to prompt injection in thread context:", emailData.threadId); const newMessages = [ @@ -463,7 +468,7 @@ Based on the email content and thread context above, draft a reply using draft_r try { const result = await generateText({ - model: workersai("@cf/moonshotai/kimi-k2.5"), + model: workersai(getAIModel("autoDraft", env)), system: systemPrompt, messages: await convertToModelMessages(messages), tools, @@ -478,7 +483,9 @@ Based on the email content and thread context above, draft a reply using draft_r if (!draftToolCalled && result.text.trim()) { // Model generated a draft inline as text -- verify with AI - const sanitizedText = await verifyDraft(env.AI, result.text.trim()); + const sanitizedText = await verifyDraft(env.AI, result.text.trim(), { + model: getAIModel("draftVerifier", env), + }); if (!sanitizedText) { // Inline text was entirely agent commentary, skip } else { diff --git a/workers/lib/ai-config.ts b/workers/lib/ai-config.ts new file mode 100644 index 00000000..104e8511 --- /dev/null +++ b/workers/lib/ai-config.ts @@ -0,0 +1,30 @@ +// Copyright (c) 2026 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +export const AI_MODELS = { + emailAgent: "@cf/moonshotai/kimi-k2.5", + autoDraft: "@cf/moonshotai/kimi-k2.5", + promptInjectionScanner: "@cf/meta/llama-3.1-8b-instruct-fast", + draftVerifier: "@cf/meta/llama-4-scout-17b-16e-instruct", +} as const; + +export type AIModelKey = keyof typeof AI_MODELS; + +const ENV_MODEL_KEYS = { + emailAgent: "EMAIL_AGENT_MODEL", + autoDraft: "AUTO_DRAFT_MODEL", + promptInjectionScanner: "PROMPT_INJECTION_MODEL", + draftVerifier: "DRAFT_VERIFIER_MODEL", +} as const satisfies Record; + +export function getAIModel(key: AIModelKey, env?: object): string { + const envKey = ENV_MODEL_KEYS[key]; + const configuredModel = env + ? (env as Record)[envKey] + : undefined; + + return typeof configuredModel === "string" && configuredModel.trim() + ? configuredModel.trim() + : AI_MODELS[key]; +} diff --git a/workers/lib/ai.ts b/workers/lib/ai.ts index 78ac4ddd..00e7ae4b 100644 --- a/workers/lib/ai.ts +++ b/workers/lib/ai.ts @@ -10,6 +10,7 @@ */ import { escapeHtml, stripHtmlToText, textToHtml } from "./email-helpers"; +import { getAIModel } from "./ai-config"; // ── Prompt Injection Scanner ─────────────────────────────────────── @@ -21,16 +22,21 @@ Return ONLY "NO" if it is a normal email (even if angry, confused, or containing Respond with exactly one word: YES or NO.`; -export async function isPromptInjection(ai: Ai, bodyHtml: string | null | undefined): Promise { +export async function isPromptInjection( + ai: Ai, + bodyHtml: string | null | undefined, + options: { model?: string } = {}, +): Promise { if (!bodyHtml) return false; const plainText = stripHtmlToText(bodyHtml).trim(); if (plainText.length < 10) return false; + const model = options.model ?? getAIModel("promptInjectionScanner"); try { const response = (await ai.run( // @ts-expect-error — model string not in generated union - "@cf/meta/llama-3.1-8b-instruct-fast", + model, { messages: [ { role: "system", content: INJECTION_PROMPT }, @@ -119,7 +125,11 @@ function splitQuotedBlock(html: string): { reply: string; quoted: string } { * Verify and clean a draft email body using AI. * Falls back to returning the original body if the AI call fails. */ -export async function verifyDraft(ai: Ai, body: string): Promise { +export async function verifyDraft( + ai: Ai, + body: string, + options: { model?: string } = {}, +): Promise { if (!body || !body.trim()) return body; // Separate the quoted reply block so the AI only reviews the user's text @@ -133,10 +143,12 @@ export async function verifyDraft(ai: Ai, body: string): Promise { // Skip very short replies — nothing to verify if (replyText.trim().length < 20) return body; + const model = options.model ?? getAIModel("draftVerifier"); try { const response = (await ai.run( - "@cf/meta/llama-4-scout-17b-16e-instruct", + // @ts-expect-error — model string may come from Worker configuration + model, { messages: [ { role: "system", content: VERIFIER_PROMPT }, diff --git a/wrangler.jsonc b/wrangler.jsonc index 53a65cd0..ce568448 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -14,7 +14,11 @@ // TEAM_DOMAIN may be the base Access URL or the full /cdn-cgi/access/certs URL. // The worker now fails closed outside local development if Access is not configured. "DOMAINS": "example.com", - "EMAIL_ADDRESSES": [] + "EMAIL_ADDRESSES": [], + "EMAIL_AGENT_MODEL": "@cf/moonshotai/kimi-k2.5", + "AUTO_DRAFT_MODEL": "@cf/moonshotai/kimi-k2.5", + "PROMPT_INJECTION_MODEL": "@cf/meta/llama-3.1-8b-instruct-fast", + "DRAFT_VERIFIER_MODEL": "@cf/meta/llama-4-scout-17b-16e-instruct" }, "send_email": [ {