commit 0e7435fc4d4f2f33039882456d67d1074b6854e0 parent 72a7574c6582fe7922ace75c45f1965883983e5f Author: Omar Gonzalez <s9tpepper@apache.org> Date: Thu, 1 Jan 2026 18:54:44 +0000 Bug 2001508 - Add ChatConversation.generatePrompt to integrate LLM context generating functions and update Chat.fetchWithHistory to use ChatConversation for conversation state and prompt submission r=tzhang,ai-frontend-reviewers,ai-models-reviewers,Mardak Differential Revision: https://phabricator.services.mozilla.com/D277129 Diffstat:
16 files changed, 438 insertions(+), 270 deletions(-)
diff --git a/browser/app/profile/firefox.js b/browser/app/profile/firefox.js @@ -2257,7 +2257,7 @@ pref("browser.aiwindow.insights", false); pref("browser.aiwindow.insightsLogLevel", "Warn"); pref("browser.aiwindow.firstrun.autoAdvanceMS", 3000); pref("browser.aiwindow.firstrun.modelChoice", ""); -pref("browser.aiwindow.model", ""); +pref("browser.aiwindow.model", "qwen3-235b-a22b-instruct-2507-maas"); // Block insecure active content on https pages pref("security.mixed_content.block_active_content", true); diff --git a/browser/base/content/test/static/browser_all_files_referenced.js b/browser/base/content/test/static/browser_all_files_referenced.js @@ -335,14 +335,6 @@ var allowlist = [ { file: "moz-src:///browser/components/aiwindow/models/IntentClassifier.sys.mjs", }, - // Bug 2002840 - add function to return real time info injection message & tests (backed out due to unused file) - { - file: "moz-src:///browser/components/aiwindow/models/ChatUtils.sys.mjs", - }, - // Bug 2003623 - Add assistant system prompt - { - file: "moz-src:///browser/components/aiwindow/models/prompts/AssistantPrompts.sys.mjs", - }, // Bug 2004888 - [FirstRun] Create Firstrun.html opening firstrun welcome screen { file: "chrome://browser/content/aiwindow/firstrun.html", @@ -351,10 +343,6 @@ var allowlist = [ { file: "moz-src:///browser/components/aiwindow/models/InsightsHistoryScheduler.sys.mjs", }, - // Bug 2000987 - get user messages from chat source - { - file: "moz-src:///browser/components/aiwindow/models/InsightsChatSource.sys.mjs", - }, // Bug 2003303 - Implement Title Generation (backed out due to unused file) { file: "moz-src:///browser/components/aiwindow/models/TitleGeneration.sys.mjs", diff --git a/browser/components/aiwindow/models/Chat.sys.mjs b/browser/components/aiwindow/models/Chat.sys.mjs @@ -4,6 +4,8 @@ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ +import { ToolRoleOpts } from "moz-src:///browser/components/aiwindow/ui/modules/ChatMessage.sys.mjs"; + /* eslint-disable-next-line mozilla/reject-import-system-module-from-non-system */ import { getFxAccountsSingleton } from "resource://gre/modules/FxAccounts.sys.mjs"; import { openAIEngine } from "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs"; @@ -49,21 +51,21 @@ export const Chat = { * we execute them locally, append results to the conversation, and continue * streaming the model’s follow-up answer. Repeats until no more tool calls. * - * @param {Array<{role:string, content?:string, tool_call_id?:string, tool_calls?:any}>} messages + * @param {ChatConversation} conversation * @yields {string} Assistant text chunks */ - async *fetchWithHistory(messages) { - const engineInstance = await openAIEngine.build(); + async *fetchWithHistory(conversation) { // Note FXA token fetching disabled for now - this is still in progress // We can flip this switch on when more realiable const fxAccountToken = await this._getFxAccountToken(); - // We'll mutate a local copy of the thread as we loop - // We also filter out empty assistant messages because - // these kinds of messages can produce unexpected model responses - let convo = Array.isArray(messages) - ? messages.filter(msg => !(msg.role == "assistant" && !msg.content)) - : []; + // @todo Bug 2007046 + // Update this with correct model id + const modelId = "qwen3-235b-a22b-instruct-2507-maas"; + + const toolRoleOpts = new ToolRoleOpts(modelId); + const currentTurn = conversation.currentTurnIndex(); + const engineInstance = await openAIEngine.build(); // Helper to run the model once (streaming) on current convo const streamModelResponse = () => @@ -72,7 +74,7 @@ export const Chat = { fxAccountToken, tool_choice: "auto", tools: toolsConfig, - args: convo, + args: conversation.getMessagesInOpenAiFormat(), }); // Keep calling until the model finishes without requesting tools @@ -98,25 +100,23 @@ export const Chat = { } // 3) Build the assistant tool_calls message exactly as expected by the API - // Bug 2006159 - Implement parallel tool calling - // TODO: Temporarily only include the first tool call due to quality issue + // + // @todo Bug 2006159 - Implement parallel tool calling + // Temporarily only include the first tool call due to quality issue // with subsequent tool call responses, will include all later once above // ticket is resolved. - const assistantToolMsg = { - role: "assistant", - tool_calls: pendingToolCalls.slice(0, 1).map(toolCall => ({ - id: toolCall.id, - type: "function", - function: { - name: toolCall.function.name, - arguments: toolCall.function.arguments, - }, - })), - }; + const tool_calls = pendingToolCalls.slice(0, 1).map(toolCall => ({ + id: toolCall.id, + type: "function", + function: { + name: toolCall.function.name, + arguments: toolCall.function.arguments, + }, + })); + conversation.addAssistantMessage("function", { tool_calls }); // 4) Execute each tool locally and create a tool message with the result // TODO: Temporarily only execute the first tool call, will run all later - const toolResultMessages = []; for (const toolCall of pendingToolCalls) { const { id, function: functionSpec } = toolCall; const name = functionSpec?.name || ""; @@ -127,11 +127,11 @@ export const Chat = { ? JSON.parse(functionSpec.arguments) : {}; } catch { - toolResultMessages.push({ - role: "tool", + const content = { tool_call_id: id, - content: JSON.stringify({ error: "Invalid JSON arguments" }), - }); + body: { error: "Invalid JSON arguments" }, + }; + conversation.addToolCallMessage(content, currentTurn, toolRoleOpts); continue; } @@ -146,31 +146,17 @@ export const Chat = { result = await toolFunc(toolParams); // Create special tool call log message to show in the UI log panel - const assistantToolCallLogMsg = { - role: "assistant", - content: `Tool Call: ${name} with parameters: ${JSON.stringify( - toolParams - )}`, - type: "tool_call_log", - result, - }; - convo.push(assistantToolCallLogMsg); - yield assistantToolCallLogMsg; + const content = { tool_call_id: id, body: result }; + conversation.addToolCallMessage(content, currentTurn, toolRoleOpts); } catch (e) { result = { error: `Tool execution failed: ${String(e)}` }; + const content = { tool_call_id: id, body: result }; + conversation.addToolCallMessage(content, currentTurn, toolRoleOpts); } - toolResultMessages.push({ - role: "tool", - tool_call_id: id, - content: typeof result === "string" ? result : JSON.stringify(result), - }); - // Bug 2006159 - Implement parallel tool calling, remove after implemented break; } - - convo = [...convo, assistantToolMsg, ...toolResultMessages]; } }, }; diff --git a/browser/components/aiwindow/models/tests/xpcshell/test_Chat.js b/browser/components/aiwindow/models/tests/xpcshell/test_Chat.js @@ -1,6 +1,14 @@ /* Any copyright is dedicated to the Public Domain. * http://creativecommons.org/publicdomain/zero/1.0/ */ +do_get_profile(); + +const { ChatConversation } = ChromeUtils.importESModule( + "moz-src:///browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs" +); +const { SYSTEM_PROMPT_TYPE, MESSAGE_ROLE } = ChromeUtils.importESModule( + "moz-src:///browser/components/aiwindow/ui/modules/ChatConstants.sys.mjs" +); const { Chat } = ChromeUtils.importESModule( "moz-src:///browser/components/aiwindow/models/Chat.sys.mjs" ); @@ -111,14 +119,22 @@ add_task(async function test_Chat_fetchWithHistory_streams_and_forwards_args() { sb.stub(openAIEngine, "build").resolves(fakeEngine); // sb.stub(Chat, "_getFxAccountToken").resolves("mock_token"); - const messages = [ - { role: "system", content: "You are helpful" }, - { role: "user", content: "Hi there" }, - ]; + const conversation = new ChatConversation({ + title: "chat title", + description: "chat desc", + pageUrl: new URL("https://www.firefox.com"), + pageMeta: {}, + }); + conversation.addSystemMessage( + SYSTEM_PROMPT_TYPE.TEXT, + "You are helpful", + 0 + ); + conversation.addUserMessage("Hi there", "https://www.firefox.com", 0); // Collect streamed output let acc = ""; - for await (const chunk of Chat.fetchWithHistory(messages)) { + for await (const chunk of Chat.fetchWithHistory(conversation)) { if (typeof chunk === "string") { acc += chunk; } @@ -130,8 +146,8 @@ add_task(async function test_Chat_fetchWithHistory_streams_and_forwards_args() { "Should concatenate streamed chunks" ); Assert.deepEqual( - capturedArgs, - messages, + [capturedArgs[0].body, capturedArgs[1].body], + [conversation.messages[0].body, conversation.messages[1].body], "Should forward messages as args to runWithGenerator()" ); Assert.deepEqual( @@ -181,26 +197,41 @@ add_task(async function test_Chat_fetchWithHistory_handles_tool_calls() { sb.stub(openAIEngine, "build").resolves(fakeEngine); // sb.stub(Chat, "_getFxAccountToken").resolves("mock_token"); - const messages = [{ role: "user", content: "Use the test tool" }]; + const conversation = new ChatConversation({ + title: "chat title", + description: "chat desc", + pageUrl: new URL("https://www.firefox.com"), + pageMeta: {}, + }); + conversation.addUserMessage( + "Use the test tool", + "https://www.firefox.com", + 0 + ); let textOutput = ""; - let toolCallLogs = []; - for await (const chunk of Chat.fetchWithHistory(messages)) { + for await (const chunk of Chat.fetchWithHistory(conversation)) { if (typeof chunk === "string") { textOutput += chunk; - } else if (chunk?.type === "tool_call_log") { - toolCallLogs.push(chunk); } } + const toolCalls = conversation.messages.filter( + message => + message.role === MESSAGE_ROLE.ASSISTANT && + message?.content?.type === "function" + ); + Assert.equal( textOutput, "I'll help you with that. Tool executed successfully!", "Should yield text from both model calls" ); - Assert.equal(toolCallLogs.length, 1, "Should have one tool call log"); + Assert.equal(toolCalls.length, 1, "Should have one tool call"); Assert.ok( - toolCallLogs[0].content.includes("test_tool"), + toolCalls[0].content.body.tool_calls[0].function.name.includes( + "test_tool" + ), "Tool call log should mention tool name" ); Assert.ok(Chat.toolMap.test_tool.calledOnce, "Tool should be called once"); @@ -228,10 +259,16 @@ add_task( sb.stub(openAIEngine, "build").rejects(err); // sb.stub(Chat, "_getFxAccountToken").resolves("mock_token"); - const messages = [{ role: "user", content: "Hi" }]; + const conversation = new ChatConversation({ + title: "chat title", + description: "chat desc", + pageUrl: new URL("https://www.firefox.com"), + pageMeta: {}, + }); + conversation.addUserMessage("Hi", "https://www.firefox.com", 0); const consume = async () => { - for await (const _chunk of Chat.fetchWithHistory(messages)) { + for await (const _chunk of Chat.fetchWithHistory(conversation)) { void _chunk; } }; @@ -284,10 +321,20 @@ add_task( sb.stub(openAIEngine, "build").resolves(fakeEngine); // sb.stub(Chat, "_getFxAccountToken").resolves("mock_token"); - const messages = [{ role: "user", content: "Test bad JSON" }]; + const conversation = new ChatConversation({ + title: "chat title", + description: "chat desc", + pageUrl: new URL("https://www.firefox.com"), + pageMeta: {}, + }); + conversation.addUserMessage( + "Test bad JSON", + "https://www.firefox.com", + 0 + ); let textOutput = ""; - for await (const chunk of Chat.fetchWithHistory(messages)) { + for await (const chunk of Chat.fetchWithHistory(conversation)) { if (typeof chunk === "string") { textOutput += chunk; } diff --git a/browser/components/aiwindow/ui/actors/AIChatContentParent.sys.mjs b/browser/components/aiwindow/ui/actors/AIChatContentParent.sys.mjs @@ -6,7 +6,7 @@ * JSWindowActor to pass data between AIChatContent singleton and content pages. */ export class AIChatContentParent extends JSWindowActorParent { - async dispatchMessageToChatContent(response) { - return this.sendQuery("AIChatContent:DispatchMessage", response); + dispatchMessageToChatContent(response) { + this.sendAsyncMessage("AIChatContent:DispatchMessage", response); } } diff --git a/browser/components/aiwindow/ui/components/ai-chat-content/ai-chat-content.mjs b/browser/components/aiwindow/ui/components/ai-chat-content/ai-chat-content.mjs @@ -51,10 +51,9 @@ export class AIChatContent extends MozLitElement { handleUserPromptEvent(event) { const { content } = event.detail; - this.conversationState = [ - ...this.conversationState, - { role: "user", content }, - ]; + + this.conversationState.push({ role: "user", content }); + this.requestUpdate(); } @@ -65,17 +64,10 @@ export class AIChatContent extends MozLitElement { */ handleAIResponseEvent(event) { - const { content, latestAssistantMessageIndex } = event.detail; - if (!this.conversationState[latestAssistantMessageIndex]) { - this.conversationState[latestAssistantMessageIndex] = { - role: "assistant", - content: "", - }; - } - this.conversationState[latestAssistantMessageIndex] = { - ...this.conversationState[latestAssistantMessageIndex], - content, - }; + const { ordinal } = event.detail; + + this.conversationState[ordinal] = event.detail; + this.requestUpdate(); } @@ -86,9 +78,12 @@ export class AIChatContent extends MozLitElement { href="chrome://browser/content/aiwindow/components/ai-chat-content.css" /> <div class="chat-content-wrapper"> - ${this.conversationState.map( - msg => html`<ai-chat-message .message=${msg}></ai-chat-message>` - )} + ${this.conversationState.map(msg => { + return html`<ai-chat-message + .message=${msg.content.body} + .role=${msg.role} + ></ai-chat-message>`; + })} </div> `; } diff --git a/browser/components/aiwindow/ui/components/ai-chat-message/ai-chat-message.mjs b/browser/components/aiwindow/ui/components/ai-chat-message/ai-chat-message.mjs @@ -14,7 +14,8 @@ export class AIChatMessage extends MozLitElement { */ static properties = { - message: { type: Object }, + role: { type: String }, + message: { type: String }, }; constructor() { @@ -33,9 +34,9 @@ export class AIChatMessage extends MozLitElement { /> <article> - <div class=${"message-" + this.message.role}> + <div class=${"message-" + this.role}> <!-- TODO: Add markdown parsing here --> - ${this.message.content} + ${this.message} </div> </article> `; diff --git a/browser/components/aiwindow/ui/components/ai-window/ai-window.mjs b/browser/components/aiwindow/ui/components/ai-window/ai-window.mjs @@ -8,16 +8,24 @@ import { MozLitElement } from "chrome://global/content/lit-utils.mjs"; const lazy = {}; ChromeUtils.defineESModuleGetters(lazy, { Chat: "moz-src:///browser/components/aiwindow/models/Chat.sys.mjs", + AIWindow: + "moz-src:///browser/components/aiwindow/ui/modules/AIWindow.sys.mjs", + ChatConversation: + "moz-src:///browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs", + MESSAGE_ROLE: + "moz-src:///browser/components/aiwindow/ui/modules//ChatEnums.sys.mjs", + AssistantRoleOpts: + "moz-src:///browser/components/aiwindow/ui/modules/ChatMessage.sys.mjs", + getRoleLabel: + "moz-src:///browser/components/aiwindow/ui/modules/ChatUtils.sys.mjs", }); -/** - * State Management Strategy: - * - * - On initialization, this component will call `.renderState()` from ChatStore to hydrate the UI - * - Currently using temporary local state (`this.conversationState`) that only tracks the current conversation turn - * - When ChatStore is integrated, rely on it as the source of truth for full conversation history - * - When calling Chat.sys.mjs to fetch responses, supplement the request with complete history from ChatStore - */ +ChromeUtils.defineLazyGetter(lazy, "log", function () { + return console.createInstance({ + prefix: "ChatStore", + maxLogLevelPref: "browser.aiwindow.chatStore.loglevel", + }); +}); /** * A custom element for managing AI Window @@ -25,14 +33,17 @@ ChromeUtils.defineESModuleGetters(lazy, { export class AIWindow extends MozLitElement { static properties = { userPrompt: { type: String }, - conversationState: { type: Array }, }; + #browser; + #conversation; + constructor() { super(); - this._browser = null; + this.userPrompt = ""; - this.conversationState = []; + this.#browser = null; + this.#conversation = new lazy.ChatConversation({}); } connectedCallback() { @@ -54,21 +65,21 @@ export class AIWindow extends MozLitElement { const container = this.renderRoot.querySelector("#browser-container"); container.appendChild(browser); - this._browser = browser; + this.#browser = browser; } /** - * Adds a new message to the conversation history. + * Persists the current conversation state to the database. * - * @param {object} chatEntry - A message object to add to the conversation - * @param {("system"|"user"|"assistant")} chatEntry.role - The role of the message sender - * @param {string} chatEntry.content - The text content of the message + * @private */ - - // TODO - can remove this method after ChatStore is integrated - #updateConversationState = chatEntry => { - this.conversationState = [...this.conversationState, chatEntry]; - }; + async #updateConversation() { + await lazy.AIWindow.chatStore + .updateConversation(this.#conversation) + .catch(updateError => { + lazy.log.error(`Error updating conversation: ${updateError.message}`); + }); + } /** * Fetches an AI response based on the current user prompt. @@ -85,39 +96,39 @@ export class AIWindow extends MozLitElement { } // Handle User Prompt - await this.#dispatchMessageToChatContent({ - role: "user", - content: this.userPrompt, + this.#dispatchMessageToChatContent({ + role: lazy.MESSAGE_ROLE.USER, + content: { + body: this.userPrompt, + }, }); - // TODO - can remove this call after ChatStore is integrated - this.#updateConversationState({ role: "user", content: formattedPrompt }); - this.userPrompt = ""; - - // Create an empty assistant placeholder. - // TODO - can remove this call after ChatStore is integrated - this.#updateConversationState({ role: "assistant", content: "" }); - const latestAssistantMessageIndex = this.conversationState.length - 1; - - let acc = ""; + const nextTurnIndex = this.#conversation.currentTurnIndex() + 1; try { - // TODO - replace with ChatStore integration IE pass chatstore.getConversationState(this.userPrompt) - const stream = lazy.Chat.fetchWithHistory(this.conversationState); + const stream = lazy.Chat.fetchWithHistory( + await this.#conversation.generatePrompt(this.userPrompt) + ); + this.#updateConversation(); + + this.userPrompt = ""; + + // @todo + // fill out these assistant message flags + const assistantRoleOpts = new lazy.AssistantRoleOpts(); + this.#conversation.addAssistantMessage( + "text", + "", + nextTurnIndex, + assistantRoleOpts + ); + for await (const chunk of stream) { - acc += chunk; - - // TODO - can remove this after ChatStore is integrated - this.conversationState[latestAssistantMessageIndex] = { - ...this.conversationState[latestAssistantMessageIndex], - content: acc, - }; - - // TODO - can pass chatstore.getLastturnIndex() instead of latestAssistantMessageIndex after ChatStore is integrated - await this.#dispatchMessageToChatContent({ - role: "assistant", - content: acc, - latestAssistantMessageIndex, - }); + const currentMessage = this.#conversation.messages.at(-1); + currentMessage.content.body += chunk; + + this.#updateConversation(); + this.#dispatchMessageToChatContent(currentMessage); + this.requestUpdate?.(); } } catch (e) { @@ -133,23 +144,23 @@ export class AIWindow extends MozLitElement { * @private */ - async #getAIChatContentActor() { - if (!this._browser) { - console.warn("AI browser not set, cannot get AIChatContent actor"); + #getAIChatContentActor() { + if (!this.#browser) { + lazy.log.warn("AI browser not set, cannot get AIChatContent actor"); return null; } - const windowGlobal = this._browser.browsingContext?.currentWindowGlobal; + const windowGlobal = this.#browser.browsingContext?.currentWindowGlobal; if (!windowGlobal) { - console.warn("No window global found for AI browser"); + lazy.log.warn("No window global found for AI browser"); return null; } try { return windowGlobal.getActor("AIChatContent"); } catch (error) { - console.error("Failed to get AIChatContent actor:", error); + lazy.log.error("Failed to get AIChatContent actor:", error); return null; } } @@ -157,13 +168,19 @@ export class AIWindow extends MozLitElement { /** * Dispatches a message to the AIChatContent actor. * - * @param {object} message - message to dispatch to chat content actor + * @param {ChatMessage} message - message to dispatch to chat content actor * @returns */ - async #dispatchMessageToChatContent(message) { - const actor = await this.#getAIChatContentActor(); - return await actor.dispatchMessageToChatContent(message); + #dispatchMessageToChatContent(message) { + const actor = this.#getAIChatContentActor(); + + if (typeof message.role !== "string") { + const roleLabel = lazy.getRoleLabel(message.role).toLowerCase(); + message.role = roleLabel; + } + + return actor.dispatchMessageToChatContent(message); } /** diff --git a/browser/components/aiwindow/ui/modules/AIWindow.sys.mjs b/browser/components/aiwindow/ui/modules/AIWindow.sys.mjs @@ -7,6 +7,7 @@ import { XPCOMUtils } from "resource://gre/modules/XPCOMUtils.sys.mjs"; const AIWINDOW_URL = "chrome://browser/content/aiwindow/aiWindow.html"; const AIWINDOW_URI = Services.io.newURI(AIWINDOW_URL); + const lazy = {}; ChromeUtils.defineESModuleGetters(lazy, { ChatStore: diff --git a/browser/components/aiwindow/ui/modules/ChatConstants.sys.mjs b/browser/components/aiwindow/ui/modules/ChatConstants.sys.mjs @@ -23,45 +23,9 @@ export const DB_FILE_NAME = "chat-store.sqlite"; */ export const PREF_BRANCH = "browser.aiWindow.chatHistory"; -/** - * @typedef ConversationStatus - * @property {number} ACTIVE - An active conversation - * @property {number} ARCHIVE - An archived conversation - * @property {number} DELETED - A deleted conversation - */ - -/** - * @type {ConversationStatus} - */ -export const CONVERSATION_STATUS = Object.freeze({ - ACTIVE: 0, - ARCHIVED: 1, - DELETED: 2, -}); - -/** - * @typedef {0 | 1 | 2 | 3} MessageRole - */ - -/** - * @enum {MessageRole} - */ -export const MESSAGE_ROLE = Object.freeze({ - USER: 0, - ASSISTANT: 1, - SYSTEM: 2, - TOOL: 3, -}); - -/** - * @typedef {0 | 1 | 2} InsightsFlagSource - */ - -/** - * @type {InsightsFlagSource} - */ -export const INSIGHTS_FLAG_SOURCE = Object.freeze({ - GLOBAL: 0, - CONVERSATION: 1, - MESSAGE_ONCE: 2, -}); +export { + CONVERSATION_STATUS, + MESSAGE_ROLE, + INSIGHTS_FLAG_SOURCE, + SYSTEM_PROMPT_TYPE, +} from "./ChatEnums.sys.mjs"; diff --git a/browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs b/browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs @@ -3,8 +3,19 @@ * 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/. */ -import { makeGuid } from "./ChatUtils.sys.mjs"; -import { CONVERSATION_STATUS, MESSAGE_ROLE } from "./ChatConstants.sys.mjs"; +import { assistantPrompt } from "moz-src:///browser/components/aiwindow/models/prompts/AssistantPrompts.sys.mjs"; + +import { + constructRelevantInsightsContextMessage, + constructRealTimeInfoInjectionMessage, +} from "moz-src:///browser/components/aiwindow/models/ChatUtils.sys.mjs"; + +import { makeGuid, getRoleLabel } from "./ChatUtils.sys.mjs"; +import { + CONVERSATION_STATUS, + MESSAGE_ROLE, + SYSTEM_PROMPT_TYPE, +} from "./ChatConstants.sys.mjs"; import { AssistantRoleOpts, ChatMessage, @@ -143,15 +154,9 @@ export class ChatConversation { * * @param {string} contentBody - The user message content * @param {string?} [pageUrl=""] - The current page url when message was submitted - * @param {number?} [turnIndex=0] - The conversation turn/cycle * @param {UserRoleOpts} [userOpts=new UserRoleOpts()] - User message options */ - addUserMessage( - contentBody, - pageUrl = "", - turnIndex = 0, - userOpts = new UserRoleOpts() - ) { + addUserMessage(contentBody, pageUrl = "", userOpts = new UserRoleOpts()) { const content = { type: "text", body: contentBody, @@ -159,7 +164,11 @@ export class ChatConversation { let url = URL.parse(pageUrl); - this.addMessage(MESSAGE_ROLE.USER, content, url, turnIndex, userOpts); + let currentTurn = this.currentTurnIndex(); + const newTurnIndex = + this.#messages.length === 1 ? currentTurn : currentTurn + 1; + + this.addMessage(MESSAGE_ROLE.USER, content, url, newTurnIndex, userOpts); } /** @@ -167,17 +176,15 @@ export class ChatConversation { * * @param {string} type - The assistant message type: text|function * @param {string} contentBody - The assistant message content - * @param {number} turnIndex - The current conversation turn/cycle * @param {AssistantRoleOpts} [assistantOpts=new AssistantRoleOpts()] - ChatMessage options specific to assistant messages */ addAssistantMessage( type, contentBody, - turnIndex, assistantOpts = new AssistantRoleOpts() ) { const content = { - type: "text", + type, body: contentBody, }; @@ -185,7 +192,7 @@ export class ChatConversation { MESSAGE_ROLE.ASSISTANT, content, "", - turnIndex, + this.currentTurnIndex(), assistantOpts ); } @@ -194,11 +201,16 @@ export class ChatConversation { * Add a tool call message to the conversation * * @param {object} content - The tool call object to be saved as JSON - * @param {number} turnIndex - The current conversation turn/cycle * @param {ToolRoleOpts} [toolOpts=new ToolRoleOpts()] - Message opts for a tool role message */ - addToolCallMessage(content, turnIndex, toolOpts = new ToolRoleOpts()) { - this.addMessage(MESSAGE_ROLE.TOOL, content, "", turnIndex, toolOpts); + addToolCallMessage(content, toolOpts = new ToolRoleOpts()) { + this.addMessage( + MESSAGE_ROLE.TOOL, + content, + "", + this.currentTurnIndex(), + toolOpts + ); } /** @@ -206,12 +218,44 @@ export class ChatConversation { * * @param {string} type - The assistant message type: text|injected_insights|injected_real_time_info * @param {string} contentBody - The system message object to be saved as JSON - * @param {number} turnIndex - The current conversation turn/cycle */ - addSystemMessage(type, contentBody, turnIndex) { + addSystemMessage(type, contentBody) { const content = { type, body: contentBody }; - this.addMessage(MESSAGE_ROLE.SYSTEM, content, "", turnIndex); + this.addMessage(MESSAGE_ROLE.SYSTEM, content, "", this.currentTurnIndex()); + } + + /** + * Takes a new prompt and generates LLM context messages before + * adding new user prompt to messages. + * + * @param {string} prompt - new user prompt + * @param {URL} pageUrl - The URL of the page when prompt was submitted + */ + async generatePrompt(prompt, pageUrl) { + if (!this.#messages.length) { + this.addSystemMessage(SYSTEM_PROMPT_TYPE.TEXT, assistantPrompt); + } + + const nextConversationTurn = this.currentTurnIndex() + 1; + + const realTime = await constructRealTimeInfoInjectionMessage(); + if (realTime.content) { + this.addSystemMessage(SYSTEM_PROMPT_TYPE.REAL_TIME, realTime.content); + } + + const insightsContext = await constructRelevantInsightsContextMessage(); + if (insightsContext?.content) { + this.addSystemMessage( + SYSTEM_PROMPT_TYPE.INSIGHTS, + insightsContext.content, + nextConversationTurn + ); + } + + this.addUserMessage(prompt, pageUrl, nextConversationTurn); + + return this; } /** @@ -257,6 +301,26 @@ export class ChatConversation { return sites.length ? sites.pop() : null; } + /** + * Converts the persisted message data to OpenAI API format + * + * @returns {Array<{ role: string, content: string }>} + */ + getMessagesInOpenAiFormat() { + return this.#messages + .filter(message => { + return !( + message.role === MESSAGE_ROLE.ASSISTANT && !message?.content?.body + ); + }) + .map(message => { + return { + role: getRoleLabel(message.role).toLowerCase(), + content: message.content?.body ?? message.content, + }; + }); + } + #updateActiveBranchTipMessageId() { this.activeBranchTipMessageId = this.messages .filter(m => m.isActiveBranch) diff --git a/browser/components/aiwindow/ui/modules/ChatEnums.sys.mjs b/browser/components/aiwindow/ui/modules/ChatEnums.sys.mjs @@ -0,0 +1,60 @@ +/* + 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/. */ + +/** + * @typedef ConversationStatus + * @property {number} ACTIVE - An active conversation + * @property {number} ARCHIVE - An archived conversation + * @property {number} DELETED - A deleted conversation + */ + +/** + * @type {ConversationStatus} + */ +export const CONVERSATION_STATUS = Object.freeze({ + ACTIVE: 0, + ARCHIVED: 1, + DELETED: 2, +}); + +/** + * @typedef {0 | 1 | 2 | 3} MessageRole + */ + +/** + * @enum {MessageRole} + */ +export const MESSAGE_ROLE = Object.freeze({ + USER: 0, + ASSISTANT: 1, + SYSTEM: 2, + TOOL: 3, +}); + +/** + * @typedef {0 | 1 | 2} InsightsFlagSource + */ + +/** + * @type {InsightsFlagSource} + */ +export const INSIGHTS_FLAG_SOURCE = Object.freeze({ + GLOBAL: 0, + CONVERSATION: 1, + MESSAGE_ONCE: 2, +}); + +/** + * @typedef { "text" | "injected_insights" | "injected_real_time_info" } SystemPromptType + */ + +/** + * @type {SystemPromptType} + */ +export const SYSTEM_PROMPT_TYPE = Object.freeze({ + TEXT: "text", + INSIGHTS: "injected_insights", + REAL_TIME: "injected_real_time_info", +}); diff --git a/browser/components/aiwindow/ui/modules/moz.build b/browser/components/aiwindow/ui/modules/moz.build @@ -1,6 +0,0 @@ -# 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/. - -with Files("**"): - BUG_COMPONENT = ("Core", "Machine Learning: Frontend") diff --git a/browser/components/aiwindow/ui/moz.build b/browser/components/aiwindow/ui/moz.build @@ -16,6 +16,7 @@ MOZ_SRC_FILES += [ "modules/AIWindowMenu.sys.mjs", "modules/ChatConstants.sys.mjs", "modules/ChatConversation.sys.mjs", + "modules/ChatEnums.sys.mjs", "modules/ChatMessage.sys.mjs", "modules/ChatMigrations.sys.mjs", "modules/ChatSql.sys.mjs", diff --git a/browser/components/aiwindow/ui/test/browser/browser_actor_user_prompt.js b/browser/components/aiwindow/ui/test/browser/browser_actor_user_prompt.js @@ -25,7 +25,7 @@ add_task(async function test_user_prompt_dispatch() { // The method should return true for successful dispatch Assert.equal( result, - true, + undefined, // actor is async instead of query now? "dispatchUserPrompt should complete successfully" ); }); @@ -49,7 +49,7 @@ add_task(async function test_streaming_ai_response() { const result = await actor.dispatchMessageToChatContent(streamingResponse); Assert.equal( result, - true, + undefined, // actor is async instead of query now? "Streaming AI response should be dispatched successfully" ); }); diff --git a/browser/components/aiwindow/ui/test/xpcshell/test_ChatConversation.js b/browser/components/aiwindow/ui/test/xpcshell/test_ChatConversation.js @@ -147,13 +147,13 @@ add_task(function test_ChatConversation_addUserMessage() { const conversation = new ChatConversation({}); const content = "user to assistant msg"; - conversation.addUserMessage(content, "https://www.mozilla.com", 0); + conversation.addUserMessage(content, "https://www.mozilla.com"); const message = conversation.messages[0]; Assert.withSoftAssertions(function (soft) { soft.equal(message.role, MESSAGE_ROLE.USER); - soft.equal(message.turnIndex, 0); + soft.equal(message.turnIndex, 1); soft.deepEqual(message.pageUrl, new URL("https://www.mozilla.com")); soft.deepEqual(message.content, { type: "text", @@ -166,7 +166,7 @@ add_task(function test_revisionRootMessageId_ChatConversation_addUserMessage() { const conversation = new ChatConversation({}); const content = "user to assistant msg"; - conversation.addUserMessage(content, "https://www.firefox.com", 0); + conversation.addUserMessage(content, "https://www.firefox.com"); const message = conversation.messages[0]; @@ -180,7 +180,6 @@ add_task(function test_opts_ChatConversation_addUserMessage() { conversation.addUserMessage( content, "https://www.firefox.com", - 0, new UserRoleOpts("321") ); @@ -193,7 +192,7 @@ add_task(function test_ChatConversation_addAssistantMessage() { const conversation = new ChatConversation({}); const content = "response from assistant"; - conversation.addAssistantMessage("text", content, 0); + conversation.addAssistantMessage("text", content); const message = conversation.messages[0]; @@ -244,7 +243,7 @@ add_task(function test_opts_ChatConversation_addAssistantMessage() { ["insight"], ["search"] ); - conversation.addAssistantMessage("text", content, 0, assistantOpts); + conversation.addAssistantMessage("text", content, assistantOpts); const message = conversation.messages[0]; @@ -300,7 +299,7 @@ add_task(function test_ChatConversation_addToolCallMessage() { const content = { random: "tool call specific keys", }; - conversation.addToolCallMessage(content, 0); + conversation.addToolCallMessage(content); const message = conversation.messages[0]; @@ -321,7 +320,7 @@ add_task(function test_opts_ChatConversation_addToolCallMessage() { const content = { random: "tool call specific keys", }; - conversation.addToolCallMessage(content, 0, new ToolRoleOpts("the-model-id")); + conversation.addToolCallMessage(content, new ToolRoleOpts("the-model-id")); const message = conversation.messages[0]; @@ -346,7 +345,7 @@ add_task(function test_ChatConversation_addSystemMessage() { const content = { random: "system call specific keys", }; - conversation.addSystemMessage("text", content, 0); + conversation.addSystemMessage("text", content); const message = conversation.messages[0]; @@ -365,12 +364,12 @@ add_task(function test_ChatConversation_getSitesList() { const conversation = new ChatConversation({}); const content = "user to assistant msg"; - conversation.addUserMessage(content, "https://www.mozilla.com", 0); - conversation.addUserMessage(content, "https://www.mozilla.com", 1); - conversation.addUserMessage(content, "https://www.firefox.com", 2); - conversation.addUserMessage(content, "https://www.cnn.com", 3); - conversation.addUserMessage(content, "https://www.espn.com", 4); - conversation.addUserMessage(content, "https://www.espn.com", 5); + conversation.addUserMessage(content, "https://www.mozilla.com"); + conversation.addUserMessage(content, "https://www.mozilla.com"); + conversation.addUserMessage(content, "https://www.firefox.com"); + conversation.addUserMessage(content, "https://www.cnn.com"); + conversation.addUserMessage(content, "https://www.espn.com"); + conversation.addUserMessage(content, "https://www.espn.com"); const sites = conversation.getSitesList(); @@ -386,12 +385,12 @@ add_task(function test_ChatConversation_getMostRecentPageVisited() { const conversation = new ChatConversation({}); const content = "user to assistant msg"; - conversation.addUserMessage(content, "https://www.mozilla.com", 0); - conversation.addUserMessage(content, "https://www.mozilla.com", 1); - conversation.addUserMessage(content, "https://www.firefox.com", 2); - conversation.addUserMessage(content, "https://www.cnn.com", 3); - conversation.addUserMessage(content, "https://www.espn.com", 4); - conversation.addUserMessage(content, "https://www.espn.com", 5); + conversation.addUserMessage(content, "https://www.mozilla.com"); + conversation.addUserMessage(content, "https://www.mozilla.com"); + conversation.addUserMessage(content, "https://www.firefox.com"); + conversation.addUserMessage(content, "https://www.cnn.com"); + conversation.addUserMessage(content, "https://www.espn.com"); + conversation.addUserMessage(content, "https://www.espn.com"); const mostRecentPageVisited = conversation.getMostRecentPageVisited(); @@ -402,9 +401,9 @@ add_task(function test_noBrowsing_ChatConversation_getMostRecentPageVisited() { const conversation = new ChatConversation({}); const content = "user to assistant msg"; - conversation.addUserMessage(content, "about:aiwindow", 0); - conversation.addUserMessage(content, "", 1); - conversation.addUserMessage(content, null, 2); + conversation.addUserMessage(content, "about:aiwindow"); + conversation.addUserMessage(content, ""); + conversation.addUserMessage(content, null); const mostRecentPageVisited = conversation.getMostRecentPageVisited(); @@ -416,12 +415,12 @@ add_task(function test_ChatConversation_renderState() { const content = "user to assistant msg"; - conversation.addUserMessage(content, "about:aiwindow", 0); - conversation.addToolCallMessage("some content", 0); - conversation.addAssistantMessage("text", "a response", 0); - conversation.addUserMessage(content, "about:aiwindow", 1); - conversation.addSystemMessage("text", "some system message", 1); - conversation.addAssistantMessage("text", "a response", 1); + conversation.addUserMessage(content, "about:aiwindow"); + conversation.addToolCallMessage("some content"); + conversation.addAssistantMessage("text", "a response"); + conversation.addUserMessage(content, "about:aiwindow"); + conversation.addSystemMessage("text", "some system message"); + conversation.addAssistantMessage("text", "a response"); const renderState = conversation.renderState(); @@ -438,16 +437,67 @@ add_task(function test_ChatConversation_currentTurnIndex() { const content = "user to assistant msg"; - conversation.addUserMessage(content, "about:aiwindow", 0); - conversation.addAssistantMessage("text", "a response", 0); - conversation.addUserMessage(content, "about:aiwindow", 2); - conversation.addAssistantMessage("text", "a response", 2); - conversation.addUserMessage(content, "about:aiwindow", 1); - conversation.addAssistantMessage("text", "a response", 1); - conversation.addUserMessage(content, "about:aiwindow", 4); - conversation.addAssistantMessage("text", "a response", 4); - conversation.addUserMessage(content, "about:aiwindow", 3); - conversation.addAssistantMessage("text", "a response", 3); + conversation.addSystemMessage("text", "the system prompt"); + conversation.addUserMessage(content, "about:aiwindow"); + conversation.addAssistantMessage("text", "a response"); + conversation.addUserMessage(content, "about:aiwindow"); + conversation.addAssistantMessage("text", "a response"); + conversation.addUserMessage(content, "about:aiwindow"); + conversation.addAssistantMessage("text", "a response"); + conversation.addUserMessage(content, "about:aiwindow"); + conversation.addAssistantMessage("text", "a response"); + conversation.addUserMessage(content, "about:aiwindow"); + conversation.addAssistantMessage("text", "a response"); Assert.deepEqual(conversation.currentTurnIndex(), 4); }); + +add_task(function test_ChatConversation_helpersTurnIndexing() { + const conversation = new ChatConversation({}); + + conversation.addSystemMessage("text", "the system prompt"); + conversation.addUserMessage("a user's prompt", "https://www.somesite.com"); + conversation.addToolCallMessage({ some: "tool call details" }); + conversation.addAssistantMessage("text", "the llm response"); + conversation.addUserMessage( + "a user's second prompt", + "https://www.somesite.com" + ); + conversation.addToolCallMessage({ some: "more tool call details" }); + conversation.addAssistantMessage("text", "the second llm response"); + + Assert.withSoftAssertions(function (soft) { + soft.equal(conversation.messages.length, 7); + + soft.equal(conversation.messages[0].turnIndex, 0); + soft.equal(conversation.messages[1].turnIndex, 0); + soft.equal(conversation.messages[2].turnIndex, 0); + soft.equal(conversation.messages[3].turnIndex, 0); + soft.equal(conversation.messages[4].turnIndex, 1); + soft.equal(conversation.messages[5].turnIndex, 1); + soft.equal(conversation.messages[6].turnIndex, 1); + }); +}); + +add_task(function test_ChatConversation_getMessagesInOpenAiFormat() { + const conversation = new ChatConversation({}); + conversation.addSystemMessage("text", "the system prompt"); + conversation.addUserMessage("a user's prompt", "https://www.somesite.com"); + conversation.addToolCallMessage({ some: "tool call details" }); + conversation.addAssistantMessage("text", "the llm response"); + conversation.addUserMessage("a user's second prompt", "some question"); + conversation.addToolCallMessage({ some: "more tool call details" }); + conversation.addAssistantMessage("text", "the second llm response"); + + const openAiFormat = conversation.getMessagesInOpenAiFormat(); + + Assert.deepEqual(openAiFormat, [ + { role: "system", content: "the system prompt" }, + { role: "user", content: "a user's prompt" }, + { role: "tool", content: { some: "tool call details" } }, + { role: "assistant", content: "the llm response" }, + { role: "user", content: "a user's second prompt" }, + { role: "tool", content: { some: "more tool call details" } }, + { role: "assistant", content: "the second llm response" }, + ]); +});