Merge pull request #3 from wagesj45/codex/enhance-thunderbird-ai-add-on-for-multiple-model-prompt-form

Add prompt template support

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.

_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" }
 }

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 {

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,

options/options.html

@@ -6,6 +6,17 @@
 </head>
 <body>
     <label>Endpoint: <input id="endpoint" type="text"></label><br>
+    <label>Prompt template:
+        <select id="template"></select>
+    </label><br>
+    <div id="custom-template-container" style="display:none">
+        <label>Custom template</label><br>
+        <textarea id="custom-template" rows="6" cols="60"></textarea>
+        <p>Placeholders: {{system}}, {{email}}, {{operator}}, {{query}}</p>
+    </div>
+    <label>System instructions:</label><br>
+    <textarea id="system-instructions" rows="4" cols="60"></textarea>
+    <button id="reset-system">Reset to default</button><br>
     <button id="save">Save</button>
     <script src="options.js"></script>
 </body>

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 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;
-    await browser.storage.local.set({ endpoint });
+        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 });
+            await browser.aiFilter.initConfig({ endpoint, templateName, customTemplate: customTemplateText, customSystemPrompt });
         } catch (e) {
             console.error('[ai-filter][options] failed to apply config', e);
         }
+    });
 });

prompt_templates/mistral.txt

@@ -0,0 +1,7 @@
+[INST] {{system}}
+
+Email:
+{{email}}
+
+Criterion ({{operator}}): {{query}}
+[/INST]

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

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