Source code
Revision control
Copy as Markdown
Other Tools
/* 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
import {
DECLINE_KINDS,
TOOLS,
} from "moz-src:///browser/components/widgets/Tools.sys.mjs";
import { Workspace } from "moz-src:///browser/components/widgets/Workspace.sys.mjs";
const lazy = {};
ChromeUtils.defineESModuleGetters(lazy, {
AuthProvider:
"moz-src:///browser/components/aiwindow/models/AuthProvider.sys.mjs",
Conversation:
"moz-src:///browser/components/aiwindow/models/Conversation.sys.mjs",
GeneratedWidgetsFeature:
"moz-src:///browser/components/widgets/GeneratedWidgetsFeature.sys.mjs",
MODEL_FEATURES: "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs",
SERVICE_TYPES: "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs",
openAIEngine:
"moz-src:///browser/components/aiwindow/models/openAIEngine.sys.mjs",
parseVersion: "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs",
});
ChromeUtils.defineLazyGetter(lazy, "logger", () =>
console.createInstance({
logModule: "WidgetGenerator",
maxLogLevel: "Warn",
maxLogLevelPref: "browser.widgets.generated.loglevel",
})
);
const ENDPOINT_PREF = "browser.widgets.generated.endpoint";
const API_KEY_PREF = "browser.widgets.generated.apiKey";
const MODEL_PREF = "browser.widgets.generated.model";
const MAX_TURNS_PREF = "browser.widgets.generated.maxTurns";
const LOCAL_PROMPT_PREF = "browser.widgets.generated.localPrompt";
const COLLECTION = "widget-generation";
const LOCAL_PROMPT_URL = "chrome://browser/content/widgets/prompt.md";
/**
* The record version this build reads. Bump it when the tools or file checks
* change.
*/
export const PROMPT_MAJOR_VERSION = 2;
/**
* The codes of a WidgetGeneratorError.
*/
export const GENERATION_ERRORS = Object.freeze({
/** AI Controls block custom widgets. */
BLOCKED: "blocked",
/** No account is signed in. */
AUTH: "auth",
/** The model can't be reached. */
SERVICE: "service",
/** The model's files don't make a widget. */
GENERATION: "generation",
/** The model declined a request that needs more than a widget can do. */
UNSUPPORTED: DECLINE_KINDS.UNSUPPORTED,
/** The model declined a request that doesn't describe a widget. */
UNCLEAR: DECLINE_KINDS.UNCLEAR,
/** The model declined a request under the acceptable use rules. */
POLICY: DECLINE_KINDS.POLICY,
});
/**
* An error generating a widget.
*/
export class WidgetGeneratorError extends Error {
name = "WidgetGeneratorError";
/**
* @param {string} message
* @param {string} code - one of GENERATION_ERRORS
* @param {Error} [cause] - the error that led to this one
*/
constructor(message, code, cause) {
super(message, { cause });
this.code = code;
}
}
/**
* Turns a prompt into a widget's files.
*/
class _WidgetGenerator {
/**
* @param {object} options
* @param {string} options.prompt - what the user asks for
* @param {AbortSignal} [options.signal] - aborting it ends generation at the
* next step
* @returns {Promise<Map<string, string>>} the files the model wrote, by path
*/
async generate({ prompt, signal }) {
signal?.throwIfAborted();
this.#throwIfBlocked();
let endpoint = Services.prefs.getStringPref(ENDPOINT_PREF, "");
let connection = {};
let token = null;
if (endpoint) {
// A custom connection set from prefs
connection = {
baseURL: endpoint,
apiKey: Services.prefs.getStringPref(API_KEY_PREF, ""),
};
} else {
// The default openAIEngine connection using an account token.
token = await lazy.AuthProvider.getToken();
if (!token) {
throw new WidgetGeneratorError(
"No account is signed in",
GENERATION_ERRORS.AUTH
);
}
}
signal?.throwIfAborted();
let config = await this.#loadConfig();
signal?.throwIfAborted();
let engine = await this.#buildEngine({
...connection,
model: Services.prefs.getStringPref(MODEL_PREF, "") || config.model,
});
signal?.throwIfAborted();
this.#throwIfBlocked();
let conversation = this.#startConversation({
...config,
engine,
prompt,
});
let files = await this.#runConversation(conversation, { token, signal });
if (!files.size) {
throw new WidgetGeneratorError(
"The model wrote no files",
GENERATION_ERRORS.GENERATION
);
}
return files;
}
#throwIfBlocked() {
if (lazy.GeneratedWidgetsFeature.isBlocked) {
throw new WidgetGeneratorError(
"Custom widgets are blocked in AI Controls",
GENERATION_ERRORS.BLOCKED
);
}
}
/**
* Reads the latest generator record for PROMPT_MAJOR_VERSION. Without one,
* the local prompt is used and the server picks the model and parameters.
* LOCAL_PROMPT_PREF uses the local prompt with any record.
*
* @returns {Promise<{systemPrompt: string, model: ?string, parameters: ?object}>}
*/
async #loadConfig() {
let record = await this.#loadRecord();
let systemPrompt =
record && !Services.prefs.getBoolPref(LOCAL_PROMPT_PREF, false)
? record.prompt
: await this.#loadLocalPrompt();
return {
systemPrompt,
model: record?.model,
parameters: record?.parameters,
};
}
/**
* Reads the latest generator record for PROMPT_MAJOR_VERSION.
*
* @returns {Promise<?object>} the record, or null when none can be read
*/
async #loadRecord() {
let records;
try {
records = await lazy.RemoteSettings(COLLECTION).get();
} catch (error) {
lazy.logger.warn("The generator records can't be read", error);
return null;
}
let record = records
.map(entry => ({ entry, version: lazy.parseVersion(entry.version) }))
.filter(({ version }) => version?.major === PROMPT_MAJOR_VERSION)
.sort((a, b) => b.version.minor - a.version.minor)[0]?.entry;
if (!record) {
lazy.logger.warn("No generator record found");
}
return record ?? null;
}
/**
* Reads the prompt packaged with Firefox.
*
* @returns {Promise<string>}
*/
async #loadLocalPrompt() {
let response = await fetch(LOCAL_PROMPT_URL);
return response.text();
}
/**
* Builds the engine for the model.
*
* @param {object} options
* @param {string} options.model
* @param {string} [options.baseURL]
* @param {string} [options.apiKey]
*/
async #buildEngine({ model, baseURL, apiKey }) {
try {
return await lazy.openAIEngine.build({
model,
baseURL,
apiKey,
feature: lazy.MODEL_FEATURES.WIDGET_GENERATION,
serviceType: lazy.SERVICE_TYPES.WIDGET_GENERATION,
});
} catch (error) {
throw new WidgetGeneratorError(
"The engine setup failed",
GENERATION_ERRORS.SERVICE,
error
);
}
}
/**
* Starts a conversation with the system prompt as the system message.
*
* @param {object} options
* @param {openAIEngine} options.engine
* @param {string} options.systemPrompt
* @param {object} options.parameters
* @param {string} options.prompt - what the user asks for
*/
#startConversation({ engine, systemPrompt, parameters, prompt }) {
let conversation = new lazy.Conversation({
feature: lazy.MODEL_FEATURES.WIDGET_GENERATION,
engine,
parameters,
});
conversation.setSystemMessage(systemPrompt);
conversation.addUserMessage(prompt);
return conversation;
}
/**
* Sends the conversation to the model for its next reply.
*
* @param {object} conversation - the conversation holding the request
* @param {?string} token - null when the endpoint pref is set
* @returns {Promise<{finalOutput: string, toolCalls: object[]}>} the reply
*/
async #runTurn(conversation, token) {
try {
return await conversation.run({ token, tools: TOOLS });
} catch (error) {
throw new WidgetGeneratorError(
"The model request failed",
GENERATION_ERRORS.SERVICE,
error
);
}
}
/**
* Runs the conversation until a reply calls no tools or the model
* declines the request.
*
* @param {object} conversation - the conversation holding the request
* @param {object} options
* @param {?string} options.token - null when the endpoint pref is set
* @param {AbortSignal} [options.signal] - checked after each reply
* @returns {Promise<Map<string, string>>} the files the model wrote, by path
*/
async #runConversation(conversation, { token, signal }) {
let workspace = new Workspace();
let lastError;
let maxTurns = Services.prefs.getIntPref(MAX_TURNS_PREF, 30);
for (let turn = 0; turn < maxTurns; turn++) {
let response = await this.#runTurn(conversation, token);
signal?.throwIfAborted();
this.#throwIfBlocked();
let text = response?.finalOutput?.trim();
if (text) {
lazy.logger.debug(text);
}
let calls = response?.toolCalls ?? [];
if (!calls.length) {
if (!workspace.changed && lastError) {
throw new WidgetGeneratorError(
lastError,
GENERATION_ERRORS.GENERATION
);
}
return workspace.files;
}
conversation.addAssistantMessage({}, { tool_calls: calls });
for (let call of calls) {
let name = call.function?.name;
let { reply, error, decline } = workspace.run(call);
if (decline) {
lazy.logger.debug("Declined", decline.kind, decline.reason);
throw new WidgetGeneratorError(
decline.reason || decline.kind,
decline.kind
);
}
lastError = error;
if (error) {
lazy.logger.debug(error, name, call.function?.arguments);
}
conversation.addToolMessage({
tool_call_id: call.id,
content: error ?? reply,
name,
});
}
}
throw new WidgetGeneratorError(
"The model was still calling tools after the last turn",
GENERATION_ERRORS.GENERATION
);
}
}
export const WidgetGenerator = new _WidgetGenerator();