commit 53df932c3d25e57cb8bb2eed1118b891c2aa78e3
parent 6f0b6420c08fd4c35c7dc42a25421495523e9985
Author: serge-sans-paille <sguelton@mozilla.com>
Date: Wed, 10 Dec 2025 20:10:36 +0000
Bug 2003303 - Implement Title Generation r=tzhang,ai-models-reviewers,ngrato
Differential Revision: https://phabricator.services.mozilla.com/D275313
Diffstat:
6 files changed, 529 insertions(+), 0 deletions(-)
diff --git a/browser/components/aiwindow/models/TitleGeneration.sys.mjs b/browser/components/aiwindow/models/TitleGeneration.sys.mjs
@@ -0,0 +1,75 @@
+/**
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ */
+
+import {
+ openAIEngine,
+ renderPrompt,
+} from "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs";
+import { titleGenerationPrompt } from "moz-src:///browser/components/aiwindow/models/prompts/TitleGenerationPrompts.sys.mjs";
+
+/**
+ * Generats a default title from the first four words of a message.
+ *
+ * @param {string} message - The user's message
+ * @returns {string} The default title
+ */
+function generateDefaultTitle(message) {
+ if (!message || typeof message !== "string") {
+ return "New Chat";
+ }
+
+ const words = message
+ .trim()
+ .split(/\s+/)
+ .filter(word => !!word.length);
+
+ if (words.length === 0) {
+ return "New Chat";
+ }
+
+ const titleWords = words.slice(0, 4);
+ return titleWords.join(" ") + "...";
+}
+
+/**
+ * Generate a chat title based on the user's message and current tab information.
+ *
+ * @param {string} message - The user's message
+ * @param {object} current_tab - Object containing current tab information
+ * @returns {Promise<string>} The generated chat title
+ */
+export async function generateChatTitle(message, current_tab) {
+ try {
+ // Build the OpenAI engines
+ const engine = await openAIEngine.build();
+
+ const tabInfo = current_tab || { url: "", title: "", description: "" };
+
+ // Render the prompt with actual values
+ const systemPrompt = await renderPrompt(titleGenerationPrompt, {
+ current_tab: JSON.stringify(tabInfo),
+ });
+
+ // Prepare messages for the LLM
+ const messages = [
+ { role: "system", content: systemPrompt },
+ { role: "user", content: message },
+ ];
+
+ // Call the LLM
+ const response = await engine.run({ messages });
+
+ // Extract the generated title from the response
+ const title =
+ response?.choices?.[0]?.message?.content?.trim() ||
+ generateDefaultTitle(message);
+
+ return title;
+ } catch (error) {
+ console.error("Failed to generate chat title:", error);
+ return generateDefaultTitle(message);
+ }
+}
diff --git a/browser/components/aiwindow/models/moz.build b/browser/components/aiwindow/models/moz.build
@@ -19,6 +19,7 @@ MOZ_SRC_FILES += [
"InsightsSchemas.sys.mjs",
"IntentClassifier.sys.mjs",
"SearchBrowsingHistory.sys.mjs",
+ "TitleGeneration.sys.mjs",
"Tools.sys.mjs",
"Utils.sys.mjs",
]
diff --git a/browser/components/aiwindow/models/prompts/TitleGenerationPrompts.sys.mjs b/browser/components/aiwindow/models/prompts/TitleGenerationPrompts.sys.mjs
@@ -0,0 +1,16 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+export const titleGenerationPrompt = `Generate a concise chat title using only the current user message and the current context.
+
+Rules:
+- Fewer than 6 words; reflect the main topic/intent
+- Do not end with punctuation
+- Do not write questions
+- No quotes, brackets, or emojis
+
+Inputs:
+The user is currently viewing this tab page: {current_tab}
+
+Output: Only the title.`;
diff --git a/browser/components/aiwindow/models/prompts/moz.build b/browser/components/aiwindow/models/prompts/moz.build
@@ -8,4 +8,5 @@ with Files("**"):
MOZ_SRC_FILES += [
"AssistantPrompts.sys.mjs",
"insightsPrompts.sys.mjs",
+ "TitleGenerationPrompts.sys.mjs",
]
diff --git a/browser/components/aiwindow/models/tests/xpcshell/test_TitleGeneration.js b/browser/components/aiwindow/models/tests/xpcshell/test_TitleGeneration.js
@@ -0,0 +1,434 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+const { generateChatTitle } = ChromeUtils.importESModule(
+ "moz-src:///browser/components/aiwindow/models/TitleGeneration.sys.mjs"
+);
+
+const { openAIEngine } = ChromeUtils.importESModule(
+ "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs"
+);
+
+const { sinon } = ChromeUtils.importESModule(
+ "resource://testing-common/Sinon.sys.mjs"
+);
+
+/**
+ * Constants for preference keys and test values
+ */
+const PREF_API_KEY = "browser.aiwindow.apiKey";
+const PREF_ENDPOINT = "browser.aiwindow.endpoint";
+const PREF_MODEL = "browser.aiwindow.model";
+
+const API_KEY = "test-api-key";
+const ENDPOINT = "https://api.test-endpoint.com/v1";
+const MODEL = "test-model";
+
+/**
+ * Cleans up preferences after testing
+ */
+registerCleanupFunction(() => {
+ for (let pref of [PREF_API_KEY, PREF_ENDPOINT, PREF_MODEL]) {
+ if (Services.prefs.prefHasUserValue(pref)) {
+ Services.prefs.clearUserPref(pref);
+ }
+ }
+});
+
+/**
+ * Test that generateChatTitle successfully generates a title
+ */
+add_task(async function test_generateChatTitle_success() {
+ Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
+ Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
+ Services.prefs.setStringPref(PREF_MODEL, MODEL);
+
+ const sb = sinon.createSandbox();
+ try {
+ // Mock the engine response
+ const mockResponse = {
+ choices: [
+ {
+ message: {
+ content: "Weather Forecast Query",
+ },
+ },
+ ],
+ };
+
+ const fakeEngineInstance = {
+ run: sb.stub().resolves(mockResponse),
+ };
+
+ sb.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ const message = "What's the weather like today?";
+ const currentTab = {
+ url: "https://weather.example.com",
+ title: "Weather Forecast",
+ description: "Get current weather conditions",
+ };
+
+ const title = await generateChatTitle(message, currentTab);
+
+ Assert.equal(
+ title,
+ "Weather Forecast Query",
+ "Should return the generated title from the LLM"
+ );
+
+ Assert.ok(
+ fakeEngineInstance.run.calledOnce,
+ "Engine run should be called once"
+ );
+
+ // Verify the messages structure passed to the engine
+ const callArgs = fakeEngineInstance.run.firstCall.args[0];
+ Assert.ok(callArgs.messages, "Should pass messages to the engine");
+ Assert.equal(
+ callArgs.messages.length,
+ 2,
+ "Should have system and user messages"
+ );
+ Assert.equal(
+ callArgs.messages[0].role,
+ "system",
+ "First message should be system"
+ );
+ Assert.equal(
+ callArgs.messages[1].role,
+ "user",
+ "Second message should be user"
+ );
+ Assert.equal(
+ callArgs.messages[1].content,
+ message,
+ "User message should contain the input message"
+ );
+
+ // Verify the system prompt contains the tab information
+ const systemContent = callArgs.messages[0].content;
+ Assert.ok(
+ systemContent.includes(currentTab.url),
+ "System prompt should include tab URL"
+ );
+ Assert.ok(
+ systemContent.includes(currentTab.title),
+ "System prompt should include tab title"
+ );
+ Assert.ok(
+ systemContent.includes(currentTab.description),
+ "System prompt should include tab description"
+ );
+ } finally {
+ sb.restore();
+ }
+});
+
+/**
+ * Test that generateChatTitle handles missing tab information
+ */
+add_task(async function test_generateChatTitle_no_tab_info() {
+ Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
+ Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
+ Services.prefs.setStringPref(PREF_MODEL, MODEL);
+
+ const sb = sinon.createSandbox();
+ try {
+ const mockResponse = {
+ choices: [
+ {
+ message: {
+ content: "General Question",
+ },
+ },
+ ],
+ };
+
+ const fakeEngineInstance = {
+ run: sb.stub().resolves(mockResponse),
+ };
+
+ sb.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ const message = "Tell me about AI";
+ const currentTab = null;
+
+ const title = await generateChatTitle(message, currentTab);
+
+ Assert.equal(
+ title,
+ "General Question",
+ "Should return the generated title even without tab info"
+ );
+
+ // Verify the system prompt handles null tab
+ const callArgs = fakeEngineInstance.run.firstCall.args[0];
+ Assert.ok(callArgs.messages, "Should pass messages even with null tab");
+ } finally {
+ sb.restore();
+ }
+});
+
+/**
+ * Test that generateChatTitle handles empty tab fields
+ */
+add_task(async function test_generateChatTitle_empty_tab_fields() {
+ Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
+ Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
+ Services.prefs.setStringPref(PREF_MODEL, MODEL);
+
+ const sb = sinon.createSandbox();
+ try {
+ const mockResponse = {
+ choices: [
+ {
+ message: {
+ content: "Untitled Chat",
+ },
+ },
+ ],
+ };
+
+ const fakeEngineInstance = {
+ run: sb.stub().resolves(mockResponse),
+ };
+
+ sb.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ const message = "Hello";
+ const currentTab = {
+ url: "",
+ title: "",
+ description: "",
+ };
+
+ const title = await generateChatTitle(message, currentTab);
+
+ Assert.equal(title, "Untitled Chat", "Should handle empty tab fields");
+
+ // Verify the system prompt includes the empty tab object
+ const callArgs = fakeEngineInstance.run.firstCall.args[0];
+ Assert.ok(
+ callArgs.messages,
+ "Should pass messages even with empty tab fields"
+ );
+ } finally {
+ sb.restore();
+ }
+});
+
+/**
+ * Test that generateChatTitle handles engine errors gracefully
+ */
+add_task(async function test_generateChatTitle_engine_error() {
+ Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
+ Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
+ Services.prefs.setStringPref(PREF_MODEL, MODEL);
+
+ const sb = sinon.createSandbox();
+ try {
+ const fakeEngineInstance = {
+ run: sb.stub().rejects(new Error("Engine failed")),
+ };
+
+ sb.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ const message = "Test message for error handling";
+ const currentTab = {
+ url: "https://example.com",
+ title: "Example",
+ description: "Test",
+ };
+
+ const title = await generateChatTitle(message, currentTab);
+
+ Assert.equal(
+ title,
+ "Test message for error...",
+ "Should return first four words when engine fails"
+ );
+ } finally {
+ sb.restore();
+ }
+});
+
+/**
+ * Test that generateChatTitle handles malformed engine responses
+ */
+add_task(async function test_generateChatTitle_malformed_response() {
+ Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
+ Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
+ Services.prefs.setStringPref(PREF_MODEL, MODEL);
+
+ const sb = sinon.createSandbox();
+ try {
+ // Test with missing choices
+ const mockResponse1 = {};
+ let fakeEngineInstance = {
+ run: sb.stub().resolves(mockResponse1),
+ };
+ sb.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ let title = await generateChatTitle("test message one two", null);
+ Assert.equal(
+ title,
+ "test message one two...",
+ "Should return first four words for missing choices"
+ );
+
+ // Test with empty choices array
+ sb.restore();
+ const sb2 = sinon.createSandbox();
+ const mockResponse2 = { choices: [] };
+ fakeEngineInstance = {
+ run: sb2.stub().resolves(mockResponse2),
+ };
+ sb2.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ title = await generateChatTitle("another test message here", null);
+ Assert.equal(
+ title,
+ "another test message here...",
+ "Should return first four words for empty choices"
+ );
+
+ // Test with null content
+ sb2.restore();
+ const sb3 = sinon.createSandbox();
+ const mockResponse3 = {
+ choices: [{ message: { content: null } }],
+ };
+ fakeEngineInstance = {
+ run: sb3.stub().resolves(mockResponse3),
+ };
+ sb3.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ title = await generateChatTitle("short test here", null);
+ Assert.equal(
+ title,
+ "short test here...",
+ "Should return first four words for null content"
+ );
+
+ sb3.restore();
+ } finally {
+ sb.restore();
+ }
+});
+
+/**
+ * Test that generateChatTitle trims whitespace from response
+ */
+add_task(async function test_generateChatTitle_trim_whitespace() {
+ Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
+ Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
+ Services.prefs.setStringPref(PREF_MODEL, MODEL);
+
+ const sb = sinon.createSandbox();
+ try {
+ const mockResponse = {
+ choices: [
+ {
+ message: {
+ content: " Title With Spaces \n\n",
+ },
+ },
+ ],
+ };
+
+ const fakeEngineInstance = {
+ run: sb.stub().resolves(mockResponse),
+ };
+
+ sb.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ const title = await generateChatTitle("test", null);
+
+ Assert.equal(
+ title,
+ "Title With Spaces",
+ "Should trim whitespace from generated title"
+ );
+ } finally {
+ sb.restore();
+ }
+});
+
+/**
+ * Test default title generation with fewer than four words
+ */
+add_task(async function test_generateChatTitle_short_message() {
+ Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
+ Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
+ Services.prefs.setStringPref(PREF_MODEL, MODEL);
+
+ const sb = sinon.createSandbox();
+ try {
+ const fakeEngineInstance = {
+ run: sb.stub().rejects(new Error("Engine failed")),
+ };
+
+ sb.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ // Test with three words
+ let title = await generateChatTitle("Hello there friend", null);
+ Assert.equal(
+ title,
+ "Hello there friend...",
+ "Should return three words with ellipsis"
+ );
+
+ // Test with one word
+ title = await generateChatTitle("Hello", null);
+ Assert.equal(title, "Hello...", "Should return one word with ellipsis");
+
+ // Test with empty message
+ title = await generateChatTitle("", null);
+ Assert.equal(
+ title,
+ "New Chat",
+ "Should return 'New Chat' for empty message"
+ );
+
+ // Test with whitespace only
+ title = await generateChatTitle(" ", null);
+ Assert.equal(
+ title,
+ "New Chat",
+ "Should return 'New Chat' for whitespace-only message"
+ );
+ } finally {
+ sb.restore();
+ }
+});
+
+/**
+ * Test default title generation with more than four words
+ */
+add_task(async function test_generateChatTitle_long_message() {
+ Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
+ Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
+ Services.prefs.setStringPref(PREF_MODEL, MODEL);
+
+ const sb = sinon.createSandbox();
+ try {
+ const fakeEngineInstance = {
+ run: sb.stub().rejects(new Error("Engine failed")),
+ };
+
+ sb.stub(openAIEngine, "_createEngine").resolves(fakeEngineInstance);
+
+ const message = "This is a very long message with many words";
+ const title = await generateChatTitle(message, null);
+
+ Assert.equal(
+ title,
+ "This is a very...",
+ "Should return only first four words with ellipsis"
+ );
+ } finally {
+ sb.restore();
+ }
+});
diff --git a/browser/components/aiwindow/models/tests/xpcshell/xpcshell.toml b/browser/components/aiwindow/models/tests/xpcshell/xpcshell.toml
@@ -18,6 +18,8 @@ support-files = []
["test_SearchBrowsingHistory.js"]
+["test_TitleGeneration.js"]
+
["test_Tools_GetOpenTabs.js"]
["test_Tools_SearchBrowsingHistory.js"]
