diff --git a/README.md b/README.md index cf93621..b05e5a5 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ message meets a specified criterion. - **AI classification rule** – adds the "AI classification" term with `matches` and `doesn't match` operators. - **Configurable endpoint** – set the classification service URL on the options page. +- **Prompt templates** – choose between several model formats or provide your own custom template. - **Filter editor integration** – patches Thunderbird's filter editor to accept text criteria for AI classification. - **Result caching** – avoids duplicate requests for already-evaluated messages. diff --git a/_locales/en-US/messages.json b/_locales/en-US/messages.json index 02b6d40..2aeebef 100644 --- a/_locales/en-US/messages.json +++ b/_locales/en-US/messages.json @@ -4,5 +4,14 @@ "doesntMatch": { "message": "doesn't match" }, "options.title": { "message": "AI Filter Options" }, "options.endpoint": { "message": "Endpoint" }, + "options.template": { "message": "Prompt template" }, + "options.customTemplate": { "message": "Custom template" }, + "options.systemInstructions": { "message": "System instructions" }, + "options.reset": { "message": "Reset to default" }, + "options.placeholders": { "message": "Placeholders: {{system}}, {{email}}, {{operator}}, {{query}}" }, + "template.openai": { "message": "OpenAI / ChatML" }, + "template.qwen": { "message": "Qwen" }, + "template.mistral": { "message": "Mistral" }, + "template.custom": { "message": "Custom" }, "options.save": { "message": "Save" } } diff --git a/background.js b/background.js index 7508a17..3526d63 100644 --- a/background.js +++ b/background.js @@ -14,7 +14,7 @@ console.log("[ai-filter] background.js loaded – ready to classify"); (async () => { try { - const store = await browser.storage.local.get(["endpoint"]); + const store = await browser.storage.local.get(["endpoint", "templateName", "customTemplate", "customSystemPrompt"]); await browser.aiFilter.initConfig(store); console.log("[ai-filter] configuration loaded", store); try { diff --git a/modules/ExpressionSearchFilter.jsm b/modules/ExpressionSearchFilter.jsm index 7fc305d..ca4c0fd 100644 --- a/modules/ExpressionSearchFilter.jsm +++ b/modules/ExpressionSearchFilter.jsm @@ -8,6 +8,19 @@ var { FileUtils } = ChromeUtils.importESModule("resource://gre/modules/Fil var EXPORTED_SYMBOLS = ["AIFilter", "ClassificationTerm"]; +const SYSTEM_PREFIX = `You are an email-classification assistant. +Read the email below and the classification criterion provided by the user. +`; + +const DEFAULT_CUSTOM_SYSTEM_PROMPT = "Determine whether the email satisfies the user's criterion."; + +const SYSTEM_SUFFIX = ` +Return ONLY a JSON object on a single line of the form: +{"match": true} - if the email satisfies the criterion +{"match": false} - otherwise + +Do not add any other keys, text, or formatting.`; + class CustomerTermBase { constructor(nameId, operators) { this.extension = ExtensionParent.GlobalManager.getExtension("ai-filter@example"); @@ -100,31 +113,58 @@ function getPlainText(msgHdr) { } let gEndpoint = "http://127.0.0.1:5000/v1/classify"; +let gTemplateName = "openai"; +let gCustomTemplate = ""; +let gCustomSystemPrompt = DEFAULT_CUSTOM_SYSTEM_PROMPT; +let gTemplateText = ""; + +function loadTemplate(name) { + try { + let url = `resource://aifilter/prompt_templates/${name}.txt`; + let xhr = new XMLHttpRequest(); + xhr.open("GET", url, false); + xhr.send(); + if (xhr.status === 0 || xhr.status === 200) { + return xhr.responseText; + } + } catch (e) { + console.error(`[ai-filter][ExpressionSearchFilter] Failed to load template '${name}':`, e); + } + return ""; +} + function setConfig(config = {}) { if (config.endpoint) { gEndpoint = config.endpoint; } + if (config.templateName) { + gTemplateName = config.templateName; + } + if (typeof config.customTemplate === "string") { + gCustomTemplate = config.customTemplate; + } + if (typeof config.customSystemPrompt === "string") { + gCustomSystemPrompt = config.customSystemPrompt; + } + gTemplateText = gTemplateName === "custom" ? gCustomTemplate : loadTemplate(gTemplateName); console.log(`[ai-filter][ExpressionSearchFilter] Endpoint set to ${gEndpoint}`); + console.log(`[ai-filter][ExpressionSearchFilter] Template set to ${gTemplateName}`); } -function buildPrompt(body, criterion) { +function buildSystemPrompt() { + return SYSTEM_PREFIX + (gCustomSystemPrompt || DEFAULT_CUSTOM_SYSTEM_PROMPT) + SYSTEM_SUFFIX; +} + +function buildPrompt(body, criterion, opName) { console.log(`[ai-filter][ExpressionSearchFilter] Building prompt with criterion: "${criterion}"`); - return `<|im_start|>system -You are an email-classification assistant. -Read the email below and the classification criterion provided by the user. - -Return ONLY a JSON object on a single line of the form: -{"match": true} - if the email satisfies the criterion -{"match": false} - otherwise - -Do not add any other keys, text, or formatting.<|im_end|> -<|im_start|>user -**Email Contents** -\`\`\` -${body} -\`\`\` -Classification Criteria: ${criterion}<|im_end|> -<|im_start|>assistant`; + const data = { + system: buildSystemPrompt(), + email: body, + operator: opName, + query: criterion, + }; + let template = gTemplateText || loadTemplate(gTemplateName); + return template.replace(/{{\s*(\w+)\s*}}/g, (m, key) => data[key] || ""); } class ClassificationTerm extends CustomerTermBase { @@ -148,7 +188,7 @@ class ClassificationTerm extends CustomerTermBase { let body = getPlainText(msgHdr); let payload = JSON.stringify({ - prompt: buildPrompt(body, value), + prompt: buildPrompt(body, value, opName), max_tokens: 4096, temperature: 1.31, top_p: 1, diff --git a/options/options.html b/options/options.html index cddd30a..d567186 100644 --- a/options/options.html +++ b/options/options.html @@ -6,6 +6,17 @@
+
+ +
+ +
diff --git a/options/options.js b/options/options.js index 8c03ef8..4e3e3a0 100644 --- a/options/options.js +++ b/options/options.js @@ -1,14 +1,54 @@ document.addEventListener('DOMContentLoaded', async () => { - let { endpoint = 'http://127.0.0.1:5000/v1/classify' } = await browser.storage.local.get(['endpoint']); - document.getElementById('endpoint').value = endpoint; -}); + const defaults = await browser.storage.local.get([ + 'endpoint', + 'templateName', + 'customTemplate', + 'customSystemPrompt' + ]); + document.getElementById('endpoint').value = defaults.endpoint || 'http://127.0.0.1:5000/v1/classify'; -document.getElementById('save').addEventListener('click', async () => { - const endpoint = document.getElementById('endpoint').value; - await browser.storage.local.set({ endpoint }); - try { - await browser.aiFilter.initConfig({ endpoint }); - } catch (e) { - console.error('[ai-filter][options] failed to apply config', e); + const templates = { + openai: browser.i18n.getMessage('template.openai'), + qwen: browser.i18n.getMessage('template.qwen'), + mistral: browser.i18n.getMessage('template.mistral'), + custom: browser.i18n.getMessage('template.custom') + }; + const templateSelect = document.getElementById('template'); + for (const [value, label] of Object.entries(templates)) { + const opt = document.createElement('option'); + opt.value = value; + opt.textContent = label; + templateSelect.appendChild(opt); } + templateSelect.value = defaults.templateName || 'openai'; + + const customBox = document.getElementById('custom-template-container'); + const customTemplate = document.getElementById('custom-template'); + customTemplate.value = defaults.customTemplate || ''; + + function updateVisibility() { + customBox.style.display = templateSelect.value === 'custom' ? 'block' : 'none'; + } + templateSelect.addEventListener('change', updateVisibility); + updateVisibility(); + + const DEFAULT_SYSTEM = 'Determine whether the email satisfies the user\'s criterion.'; + const systemBox = document.getElementById('system-instructions'); + systemBox.value = defaults.customSystemPrompt || DEFAULT_SYSTEM; + document.getElementById('reset-system').addEventListener('click', () => { + systemBox.value = DEFAULT_SYSTEM; + }); + + document.getElementById('save').addEventListener('click', async () => { + const endpoint = document.getElementById('endpoint').value; + const templateName = templateSelect.value; + const customTemplateText = customTemplate.value; + const customSystemPrompt = systemBox.value; + await browser.storage.local.set({ endpoint, templateName, customTemplate: customTemplateText, customSystemPrompt }); + try { + await browser.aiFilter.initConfig({ endpoint, templateName, customTemplate: customTemplateText, customSystemPrompt }); + } catch (e) { + console.error('[ai-filter][options] failed to apply config', e); + } + }); }); diff --git a/prompt_templates/mistral.txt b/prompt_templates/mistral.txt new file mode 100644 index 0000000..4eb1f4b --- /dev/null +++ b/prompt_templates/mistral.txt @@ -0,0 +1,7 @@ +[INST] {{system}} + +Email: +{{email}} + +Criterion ({{operator}}): {{query}} +[/INST] diff --git a/prompt_templates/openai.txt b/prompt_templates/openai.txt new file mode 100644 index 0000000..2088d8e --- /dev/null +++ b/prompt_templates/openai.txt @@ -0,0 +1,9 @@ +<|im_start|>system +{{system}}<|im_end|> +<|im_start|>user +**Email Contents** +``` +{{email}} +``` +Classification Criterion ({{operator}}): {{query}}<|im_end|> +<|im_start|>assistant diff --git a/prompt_templates/qwen.txt b/prompt_templates/qwen.txt new file mode 100644 index 0000000..ca18018 --- /dev/null +++ b/prompt_templates/qwen.txt @@ -0,0 +1,10 @@ +<|im_start|>system +{{system}} +<|im_end|> +<|im_start|>user +Email: +{{email}} + +Criterion ({{operator}}): {{query}} +<|im_end|> +<|im_start|>assistant