firefox-main/browser/components/aiwindow/models/tests/xpcshell/test_TitleGeneration.js (file symbol)

Source code

Revision control

Copy as Markdown

Other Tools

Test Info: Warnings

/* 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(
);
/**
* 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 = {
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 = {
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();
}
});