From 74904769ddb7c42f3dc0ddcc23482620baf41f2c Mon Sep 17 00:00:00 2001 From: 0xBuooy Date: Mon, 18 May 2026 11:13:14 +0800 Subject: [PATCH] Configure AI model selection Add a centralized AI model config that defines defaults for the email agent, auto-drafting, prompt-injection scanning, and draft verification. The helper supports Worker env var overrides while keeping the existing models as fallbacks. Update the agent and AI helper call sites to resolve models through the config instead of hard-coded strings, and expose the corresponding model variables in wrangler.jsonc for deploy-time tuning. --- workers/agent/index.ts | 17 ++++++++++++----- workers/lib/ai-config.ts | 30 ++++++++++++++++++++++++++++++ workers/lib/ai.ts | 20 ++++++++++++++++---- wrangler.jsonc | 6 +++++- 4 files changed, 63 insertions(+), 10 deletions(-) create mode 100644 workers/lib/ai-config.ts 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": [ {