XLinfini/sousa-gecko
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Bug 2063111 - updates to align with API contract r=mshillabeer,jlewis

Browse Source 
Differential Revision: https://phabricator.services.mozilla.com/D322800
This commit is contained in:
Omar Gonzalez
2026-09-12 06:51:50 +00:00
committed by ogonzalez@mozilla.com
parent a1ede75a77
commit 9a17117b7b
11 changed files with 503 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files
browser/components/aiwindow/models/SmartFormFillModel.sys.mjs
+61 -85
View File
@@ -13,6 +13,7 @@ import {
buildConversation,
loadPrompt,
} from "moz-src:///browser/components/aiwindow/models/PromptLoader.sys.mjs";
import { UrlTokenizer } from "moz-src:///browser/components/aiwindow/ui/modules/UrlTokenizer.sys.mjs";
/**
* Reports the model and prompt version a request is about to be sent with.
@@ -141,11 +142,17 @@ import {
* @property {string} tabContent The tab content
*/
/**
* @typedef {object} MemoryDataForValueGen
* @property {string} id Memory ID
* @property {string} memory_summary Memory summary
*/
/**
* @typedef {object} Context
* @property {string} [pageText] Text of the current page
* @property {Array<TabCandidate>} [relevantTabs] Tabs for context
* @property {Array<string>} [memories] List of memories
* @property {Array<MemoryDataForValueGen>} [memories] List of memories
*/
/**
@@ -161,14 +168,9 @@ import {
* @typedef {object} FieldValue
* @property {string} id The stable field ID
* @property {"fill_from_token" | "select_option" | "generate" | "skip"} action The action the LLM decided for the value
* @property {string} [token] Candidate token, present only when action is
* "fill_from_token"
* @property {"high" | "medium" | "low"} confidence The LLM's value confidence
* @property {string} [optionId] Select option stable ID, present only when
* action is
* "select_option"
* @property {string} [value] Generated value, present only when action is
* "generate"
* @property {string} value Candidate token, option ID, generated value, or an
* empty string when the action is "skip"
*/
/**
@@ -264,6 +266,8 @@ function queueValuesBatchRequest(request, options) {
return promise;
}
const TITLE_CHAR_LIMIT = 100;
const FIELD_CLASSIFICATION_RESPONSE_SCHEMA = {
type: "object",
properties: {
@@ -326,75 +330,21 @@ const FORM_VALUES_RESPONSE_SCHEMA = {
fields: {
type: "array",
items: {
oneOf: [
{
type: "object",
properties: {
id: { type: "string" },
action: {
type: "string",
enum: ["fill_from_token"],
},
token: { type: "string" },
confidence: {
type: "string",
enum: ["high", "medium", "low"],
},
},
required: ["id", "action", "token", "confidence"],
additionalProperties: false,
type: "object",
properties: {
id: { type: "string" },
action: {
type: "string",
enum: ["fill_from_token", "select_option", "generate", "skip"],
},
{
type: "object",
properties: {
id: { type: "string" },
action: {
type: "string",
enum: ["generate"],
},
value: { type: "string" },
confidence: {
type: "string",
enum: ["high", "medium", "low"],
},
},
required: ["id", "action", "value", "confidence"],
additionalProperties: false,
confidence: {
type: "string",
enum: ["high", "medium", "low"],
},
{
type: "object",
properties: {
id: { type: "string" },
action: {
type: "string",
enum: ["select_option"],
},
optionId: { type: "string" },
confidence: {
type: "string",
enum: ["high", "medium", "low"],
},
},
required: ["id", "action", "optionId", "confidence"],
additionalProperties: false,
},
{
type: "object",
properties: {
id: { type: "string" },
action: {
type: "string",
enum: ["skip"],
},
confidence: {
type: "string",
enum: ["high", "medium", "low"],
},
},
required: ["id", "action", "confidence"],
additionalProperties: false,
},
],
value: { type: "string" },
},
required: ["id", "action", "confidence", "value"],
additionalProperties: false,
},
},
},
@@ -409,10 +359,14 @@ const FORM_VALUES_RESPONSE_SCHEMA = {
* @param {object} [param1={}]
* @param {AbortSignal} [param1.signal]
* @param {ModelInfoCallback} [param1.onDispatch]
* @param {UrlTokenizer} param1.urlTokenizer
*
* @returns {Promise<GenerateFormValuesBatchResponse>}
*/
async function generateFormValuesBatch(request, { signal, onDispatch } = {}) {
async function generateFormValuesBatch(
request,
{ signal, onDispatch, urlTokenizer } = {}
) {
signal?.throwIfAborted();
const conversation = await buildConversation(MODEL_FEATURES.SMART_FORM_FILL);
@@ -432,12 +386,21 @@ async function generateFormValuesBatch(request, { signal, onDispatch } = {}) {
]);
signal?.throwIfAborted();
const url = urlTokenizer.encodeToken(request.page.url);
const relevantTabs = request.context.relevantTabs.map(tab => {
return {
...tab,
title: tab.title.substring(0, TITLE_CHAR_LIMIT),
url: urlTokenizer.encodeToken(tab.url),
};
});
const userPrompt = renderPrompt(userPromptTemplate, {
title: request.page.title,
url: request.page.url,
url,
title: request.page.title.substring(0, TITLE_CHAR_LIMIT),
pageText: request.context.pageText ?? "",
memories: JSON.stringify(request.context.memories ?? []),
pageContext: JSON.stringify(request.context.relevantTabs ?? []),
pageContext: JSON.stringify(relevantTabs ?? []),
candidateTokens: JSON.stringify(request.candidates),
fields: JSON.stringify(request.fields),
});
@@ -516,9 +479,11 @@ export const SmartFormFillModel = {
]);
signal?.throwIfAborted();
const urlTokenizer = new UrlTokenizer();
const url = urlTokenizer.encodeToken(request.page.url);
const userPrompt = renderPrompt(userPromptTemplate, {
title: request.page.title,
url: request.page.url,
url,
title: request.page.title.substring(0, TITLE_CHAR_LIMIT),
fields: JSON.stringify(request.fields),
});
@@ -577,11 +542,21 @@ export const SmartFormFillModel = {
]);
signal?.throwIfAborted();
const urlTokenizer = new UrlTokenizer();
const url = urlTokenizer.encodeToken(request.page.url);
const tabs = request.tabs.map(tab => {
return {
...tab,
title: tab.title.substring(0, TITLE_CHAR_LIMIT),
url: urlTokenizer.encodeToken(tab.url),
};
});
const userPrompt = renderPrompt(userPromptTemplate, {
title: request.page.title,
url: request.page.url,
url,
title: request.page.title.substring(0, TITLE_CHAR_LIMIT),
fields: JSON.stringify(request.fields),
tabs: JSON.stringify(request.tabs),
tabs: JSON.stringify(tabs),
max_selected_tabs: request.maxSelectedTabs,
});
@@ -626,6 +601,7 @@ export const SmartFormFillModel = {
async generateFormValues(request, { signal, onDispatch } = {}) {
signal?.throwIfAborted();
const urlTokenizer = new UrlTokenizer();
const requests = [];
for (
let index = 0;
@@ -641,7 +617,7 @@ export const SmartFormFillModel = {
index + MAX_FIELDS_PER_GENERATION_REQUEST
),
},
{ signal, onDispatch }
{ signal, onDispatch, urlTokenizer }
)
);
}
browser/components/aiwindow/models/tests/browser/browser_smart_form_fill_model.js
+181 -94
View File
@@ -99,10 +99,47 @@ function makeRequest(id) {
},
],
candidates: [],
context: {},
context: {
relevantTabs: [],
},
};
}
function makeFields(count) {
return Array.from({ length: count }, (_, index) => ({
id: `field-${index}`,
label: `Field ${index}`,
inputType: "text",
options: [],
type: "unknown",
classificationConfidence: "low",
}));
}
function generateValuesForFields(fields) {
return SmartFormFillModel.generateFormValues({
task: "generate",
page: {
title: "Batch test",
url: "https://example.com/batch",
},
fields,
candidates: [],
context: {
pageText: "",
relevantTabs: [],
memories: [],
},
});
}
function getRequestedFields(request) {
const marker = "Fields to fill:\n";
const content = request.args.at(-1).content;
return JSON.parse(content.slice(content.lastIndexOf(marker) + marker.length));
}
describe("SmartFormFillModel", () => {
let mockEngineMan;
@@ -120,7 +157,8 @@ describe("SmartFormFillModel", () => {
});
describe("classifyFields", () => {
it("uses the field detection prompts and returns classifications", async () => {
it("builds a field classification request", async () => {
const title = "T".repeat(101);
const fields = [
{
id: "field-1",
@@ -134,7 +172,7 @@ describe("SmartFormFillModel", () => {
task: "classify",
enumVersion: "sff-fieldtypes-1",
page: {
title: "Example form",
title,
url: "https://example.com/form",
},
fields,
@@ -157,8 +195,8 @@ describe("SmartFormFillModel", () => {
Assert.deepEqual(request.args[1], {
role: "user",
content:
"Page title:\nExample form\n\n" +
"Page url:\nhttps://example.com/form\n\n" +
`Page title:\n${title.substring(0, 100)}\n\n` +
"Page url:\nEXAMPLE_COM_FORM_1\n\n" +
`Fields to classify:\n${JSON.stringify(fields)}`,
});
Assert.deepEqual(request.tools, []);
@@ -171,32 +209,15 @@ describe("SmartFormFillModel", () => {
);
Assert.deepEqual(responseFormat.json_schema.schema.required, ["fields"]);
respond(
JSON.stringify({
fields: [
{
id: "field-1",
type: "email",
confidence: "high",
},
],
})
);
Assert.deepEqual(await requestPromise, {
fields: [
{
id: "field-1",
type: "email",
confidence: "high",
},
],
});
respond(JSON.stringify({ fields: [] }));
await requestPromise;
});
});
describe("findRelevantTabs", () => {
it("uses the tab selection prompts and returns relevant tabs", async () => {
it("builds a relevant tabs request", async () => {
const pageTitle = "P".repeat(101);
const tabTitle = "T".repeat(101);
const fields = [
{
id: "field-1",
@@ -208,14 +229,21 @@ describe("SmartFormFillModel", () => {
const tabs = [
{
id: "tab-1",
title: "Professional profile",
title: tabTitle,
url: "https://example.com/profile",
},
];
const shapedTabs = [
{
...tabs[0],
title: tabTitle.substring(0, 100),
url: "EXAMPLE_COM_PROFILE_1",
},
];
const requestPromise = SmartFormFillModel.findRelevantTabs({
task: "select_tabs",
page: {
title: "Job application",
title: pageTitle,
url: "https://example.com/jobs/apply",
},
maxSelectedTabs: 3,
@@ -237,12 +265,12 @@ describe("SmartFormFillModel", () => {
role: "user",
content:
"FORM THE USER IS FILLING\n" +
"Title: Job application\n" +
"URL: https://example.com/jobs/apply\n" +
`Title: ${pageTitle.substring(0, 100)}\n` +
"URL: EXAMPLE_COM_JOBS_APPLY_1\n" +
"Fields it is asking for:\n" +
`${JSON.stringify(fields)}\n\n` +
"THE USER'S OTHER OPEN TABS (title and URL only)\n" +
`${JSON.stringify(tabs)}\n\n` +
`${JSON.stringify(shapedTabs)}\n\n` +
"maxSelectedTabs: 3\n\n" +
"Which of these tabs could help fill this form? " +
"Include every tab that plausibly could.",
@@ -259,27 +287,8 @@ describe("SmartFormFillModel", () => {
"selectedTabs",
]);
respond(
JSON.stringify({
selectedTabs: [
{
id: "tab-1",
relevance: "high",
reason: "Your professional profile",
},
],
})
);
Assert.deepEqual(await requestPromise, {
selectedTabs: [
{
id: "tab-1",
relevance: "high",
reason: "Your professional profile",
},
],
});
respond(JSON.stringify({ selectedTabs: [] }));
await requestPromise;
});
});
@@ -292,8 +301,10 @@ describe("SmartFormFillModel", () => {
_clearRemoteClientForTesting();
});
it("uses the value generation prompts and returns fill actions", async () => {
const candidates = [{ token: "$EMAIL_1", type: "email" }];
it("builds a value generation request with the expected schema", async () => {
const pageTitle = "P".repeat(101);
const tabTitle = "T".repeat(101);
const candidates = [{ token: "§EMAIL_1§", type: "email" }];
const fields = [
{
id: "field-email",
@@ -314,16 +325,23 @@ describe("SmartFormFillModel", () => {
];
const relevantTabs = [
{
title: "Example role",
title: tabTitle,
url: "https://example.com/jobs/role",
tabContent: "The role focuses on browser engineering.",
},
];
const shapedRelevantTabs = [
{
...relevantTabs[0],
title: tabTitle.substring(0, 100),
url: "EXAMPLE_COM_JOBS_ROLE_1",
},
];
const requestPromise = SmartFormFillModel.generateFormValues({
task: "generate",
page: {
title: "Job application",
url: "https://example.com/jobs/apply",
title: pageTitle,
url: "https://example.com/jobs/generate",
},
fields,
candidates,
@@ -346,11 +364,11 @@ describe("SmartFormFillModel", () => {
Assert.deepEqual(request.args[1], {
role: "user",
content:
"Current page title:\nJob application\n\n" +
"Current page url:\nhttps://example.com/jobs/apply\n\n" +
`Current page title:\n${pageTitle.substring(0, 100)}\n\n` +
"Current page url:\nEXAMPLE_COM_JOBS_GENERATE_1\n\n" +
"Current page text:\nApply for the example role.\n\n" +
"Relevant memories about the user:\n[]\n\n" +
`Relevant open tabs:\n${JSON.stringify(relevantTabs)}\n\n` +
`Relevant open tabs:\n${JSON.stringify(shapedRelevantTabs)}\n\n` +
`Available candidate tokens:\n${JSON.stringify(candidates)}\n\n` +
`Fields to fill:\n${JSON.stringify(fields)}`,
});
@@ -364,48 +382,117 @@ describe("SmartFormFillModel", () => {
"tabs_used",
"fields",
]);
const fieldSchema =
responseFormat.json_schema.schema.properties.fields.items;
Assert.deepEqual(fieldSchema.required, [
"id",
"action",
"confidence",
"value",
]);
Assert.deepEqual(fieldSchema.properties.action.enum, [
"fill_from_token",
"select_option",
"generate",
"skip",
]);
respond(
JSON.stringify({
memories_used: [],
fields: [
{
id: "field-email",
action: "fill_from_token",
token: "$EMAIL_1",
confidence: "high",
},
{
id: "field-reason",
tabs_used: [],
fields: [],
})
);
await requestPromise;
});
it("splits fields into batches and combines their results", async () => {
const fields = makeFields(21);
const requestPromise = generateValuesForFields(fields);
const batchSizes = [];
for (let index = 0; index < 2; index++) {
const { request, respond } = await mockEngineMan.captureRequest({
purpose: PURPOSE,
});
const requestedFields = getRequestedFields(request);
batchSizes.push(requestedFields.length);
respond(
JSON.stringify({
memories_used: [],
tabs_used: ["tab-1"],
fields: requestedFields.map(({ id }) => ({
id,
action: "generate",
value: "I am interested in browser engineering.",
confidence: "high",
},
],
value: `Value for ${id}`,
})),
})
);
}
const result = await requestPromise;
Assert.deepEqual(
batchSizes,
[20, 1],
"Fields should be split at the batch limit"
);
Assert.deepEqual(
result.fields.map(({ id }) => id),
fields.map(({ id }) => id),
"Fields from every batch should be combined in request order"
);
Assert.deepEqual(
result.tabs_used,
["tab-1"],
"Tab IDs reported by multiple batches should be deduplicated"
);
Assert.deepEqual(
result.batches,
{ total: 2, failed: 0 },
"Batch metadata should report both successful requests"
);
});
it("keeps successful results when another batch fails", async () => {
const fields = makeFields(21);
const requestPromise = generateValuesForFields(fields);
const { request, respond } = await mockEngineMan.captureRequest({
purpose: PURPOSE,
});
const successfulFields = getRequestedFields(request);
respond(
JSON.stringify({
memories_used: [],
tabs_used: [],
fields: successfulFields.map(({ id }) => ({
id,
action: "generate",
confidence: "high",
value: `Value for ${id}`,
})),
})
);
Assert.deepEqual(await requestPromise, {
memories_used: [],
// Absent from the answer above, so the merged result reports it empty.
tabs_used: [],
fields: [
{
id: "field-email",
action: "fill_from_token",
token: "$EMAIL_1",
confidence: "high",
},
{
id: "field-reason",
action: "generate",
value: "I am interested in browser engineering.",
confidence: "high",
},
],
// One batch, because the form fits in a single request.
batches: { total: 1, failed: 0 },
});
await mockEngineMan.captureRequest({ purpose: PURPOSE });
mockEngineMan.rejectAllRequests();
const result = await requestPromise;
Assert.deepEqual(
result.fields.map(({ id }) => id),
successfulFields.map(({ id }) => id),
"Successful batch results should be preserved"
);
Assert.deepEqual(
result.batches,
{ total: 2, failed: 1 },
"Batch metadata should report the failed request"
);
});
it("limits concurrent value generation requests globally", async () => {
browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs
+19 -100
View File
@@ -30,6 +30,7 @@ import {
import { EventEmitter } from "resource://gre/modules/EventEmitter.sys.mjs";
import { Conversation } from "moz-src:///browser/components/aiwindow/models/Conversation.sys.mjs";
import { consumeStreamChunk } from "moz-src:///browser/components/aiwindow/models/TokenStreamParser.sys.mjs";
import { UrlTokenizer } from "moz-src:///browser/components/aiwindow/ui/modules/UrlTokenizer.sys.mjs";
/** @typedef {import("moz-src:///browser/components/aiwindow/models/SearchBrowsingHistory.sys.mjs").HistoryRow} HistoryRow */
@@ -141,49 +142,14 @@ export class ChatConversation extends Conversation {
promptEmbeddedMemories = [];
/**
* A mapping of a URL to its unique URL token. URL tokens are used as shortened
* versions of URLs to help the model deal with very long URLs. Very long URLs are
* problematic since they are hard for a model to repeat back without making mistakes
* or hallucinating details about the URL. There is also additional cost for every
* additional token in the context. Long URLs can also contain prompt injections since
* they can be of an arbitrary size. URL Tokens help solve all of these issues.
* URL Tokenizer instance for sending URLs to LLM.
* The tokens created by UrlTokenizer are cached to the conversation while
* the conversation is loaded in memory. These tokens are in-memory only and not
* serialized to storage.
*
* URL tokens are only generated while a message is "in flight" to and from the language
* model. When tool calls are handled, messages rendered, and messages stored they are
* all done with the URL tokens expanded into full URLs.
*
* There are no guarantees that a URL in this list isn't just hallucinated by the model.
* Any URL the language model invents can be present in this list. The only guarantee
* is that a token maps to some kind of arbitrary URL.
*
* Example mapping:
* https://github.com/mozilla/ -> GITHUB_COM_MOZILLA_1
*
* @type {Map<string, string>}
* @type {UrlTokenizer}
*/
urlToToken = new Map();
/**
* The reverse mapping for a token back to its original URL.
*
* e.g. GITHUB_COM_MOZILLA_1 -> https://github.com/mozilla/
*
* @type {Map<string, string>}
*/
tokenToUrl = new Map();
/**
* A mapping of the base URL token to how many counts there are for it. It's
* used to generate the final number on URL tokens.
*
* e.g.
*
* https://github.com/mozilla/ -> GITHUB_COM_MOZILLA_1
* https://github.com/mozilla#not-part-of-token -> GITHUB_COM_MOZILLA_2
*
* @type {Map<string, number>}
*/
#baseTokenCounts = new Map();
urlTokenizer = null;
/**
* Conversation-level pool of history results keyed by URL, accumulated across
@@ -262,6 +228,7 @@ export class ChatConversation extends Conversation {
securityProperties,
});
this.urlTokenizer = new UrlTokenizer();
this.title = title;
this.description = description;
this.pageUrl = pageUrl;
@@ -352,63 +319,7 @@ export class ChatConversation extends Conversation {
* @returns {string} The short token for the URL (e.g. "GITHUB_COM_1")
*/
convertUrlToToken(url) {
const seenToken = this.urlToToken.get(url);
if (seenToken) {
return seenToken;
}
let baseToken = "";
// Attempt to convert the URL into a base token.
const parsedUrl = URL.parse(url);
if (parsedUrl) {
if (parsedUrl.protocol !== "http:" && parsedUrl.protocol !== "https:") {
// Go ahead and handle URL tokens for more complicated URLs that
// aren't probably supported in the chat interface, but would be useful
// to disambiguate from the HTTP(s) varieties.
baseToken +=
// e.g. "ftp:" -> "FTP"
parsedUrl.protocol.toUpperCase().replace(":", "");
}
// Convert the hostname into a token.
const hostToken = parsedUrl.hostname
.replace(/^www\./, "")
.toUpperCase()
.replace(/[.\-]/g, "_")
.substring(0, 100);
if (hostToken) {
baseToken = baseToken ? `${baseToken}_${hostToken}` : hostToken;
}
// Add on the parts of the URL to the token.
for (let part of parsedUrl.pathname.split("/")) {
if (!part) {
continue;
}
const partToken = part.toUpperCase().replace(/[^A-Z0-9]/g, "_");
const nextToken = `${baseToken}_${partToken}`;
if (nextToken.length > 100) {
break;
}
baseToken = nextToken;
}
} else {
baseToken = "INVALID_URL";
}
let count = this.#baseTokenCounts.get(baseToken) ?? 0;
count += 1;
this.#baseTokenCounts.set(baseToken, count);
const tokenFinal = `${baseToken}_${count}`;
this.urlToToken.set(url, tokenFinal);
this.tokenToUrl.set(tokenFinal, url);
return tokenFinal;
return this.urlTokenizer.encodeToken(url);
}
/**
@@ -420,7 +331,7 @@ export class ChatConversation extends Conversation {
const { plainText, tokens } = consumeStreamChunk(
chunk,
parserState,
this.tokenToUrl
this.urlTokenizer.tokenToUrl
);
if (plainText && currentMessage?.content) {
@@ -477,7 +388,7 @@ export class ChatConversation extends Conversation {
if (result.currentMessage?.content?.body) {
// Expand URL tokens and remove any hallucinated ones.
if (this.urlToToken.size) {
if (this.urlTokenizer.urlToToken.size) {
result.currentMessage.content.body = stripUnresolvedUrlTokens(
result.currentMessage.content.body
);
@@ -1135,6 +1046,14 @@ export class ChatConversation extends Conversation {
return this.messages.filter(m => CHAT_ROLES.includes(m.role)).length;
}
get tokenToUrl() {
return this.urlTokenizer.tokenToUrl;
}
get urlToToken() {
return this.urlTokenizer.urlToToken;
}
/**
* Returns the contextMentions count from the most recent user message in
* the conversation, or 0 if none.
browser/components/aiwindow/ui/modules/SmartFormFillController.sys.mjs
+8 -6
View File
@@ -526,7 +526,7 @@ export class SmartFormFillController {
* @param {PageInfo} pageInfo
* @param {Array<FieldData>} fields
*
* @returns {Promise<Array<string>>}
* @returns {Promise<Array<{id: string, memory_summary: string}>>}
*/
// eslint-disable-next-line no-unused-private-class-members -- will be enabled in v0+
async #getMemories(pageInfo, fields) {
@@ -552,9 +552,11 @@ export class SmartFormFillController {
const relevantMemories =
await lazy.MemoriesManager.getRelevantMemories(contextMessage);
return relevantMemories.map(
relevant_memory => relevant_memory.memory_summary
);
return relevantMemories.map(relevant_memory => {
const { id, memory_summary } = relevant_memory;
return { id, memory_summary };
});
}
/**
@@ -584,7 +586,7 @@ export class SmartFormFillController {
let value;
switch (result.action) {
case "fill_from_token":
value = result.token ? valuesByToken.get(result.token) : undefined;
value = result.value ? valuesByToken.get(result.value) : undefined;
break;
case "generate":
@@ -640,7 +642,7 @@ export class SmartFormFillController {
const count = (typeCounts.get(type) ?? 0) + 1;
typeCounts.set(type, count);
const token = `$${type.toUpperCase().replaceAll("-", "_")}_${count}`;
const token = `§${type.toUpperCase().replaceAll("-", "_")}_${count}§`;
candidates.push({ token, type });
valuesByToken.set(token, value);
browser/components/aiwindow/ui/modules/UrlTokenizer.sys.mjs
+132
View File
@@ -0,0 +1,132 @@
/**
* 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/.
*/
/**
* Tokenizes URLs to send to an LLM
*/
export class UrlTokenizer {
/**
* A mapping of a URL to its unique URL token. URL tokens are used as shortened
* versions of URLs to help the model deal with very long URLs. Very long URLs are
* problematic since they are hard for a model to repeat back without making mistakes
* or hallucinating details about the URL. There is also additional cost for every
* additional token in the context. Long URLs can also contain prompt injections since
* they can be of an arbitrary size. URL Tokens help solve all of these issues.
*
* URL tokens are only generated while a message is "in flight" to and from the language
* model. When tool calls are handled, messages rendered, and messages stored they are
* all done with the URL tokens expanded into full URLs.
*
* There are no guarantees that a URL in this list isn't just hallucinated by the model.
* Any URL the language model invents can be present in this list. The only guarantee
* is that a token maps to some kind of arbitrary URL.
*
* Example mapping:
* https://github.com/mozilla/ -> GITHUB_COM_MOZILLA_1
*
* @type {Map<string, string>}
*/
urlToToken;
/**
* The reverse mapping for a token back to its original URL.
*
* e.g. GITHUB_COM_MOZILLA_1 -> https://github.com/mozilla/
*
* @type {Map<string, string>}
*/
tokenToUrl;
/**
* A mapping of the base URL token to how many counts there are for it. It's
* used to generate the final number on URL tokens.
*
* e.g.
*
* https://github.com/mozilla/ -> GITHUB_COM_MOZILLA_1
* https://github.com/mozilla#not-part-of-token -> GITHUB_COM_MOZILLA_2
*
* @type {Map<string, number>}
*/
#baseTokenCounts = new Map();
constructor() {
this.urlToToken = new Map();
this.tokenToUrl = new Map();
}
/**
* Converts a URL into a token. It first computes a base token from
* the hostname and path parts, then appends a monotonically increasing number on the
* end to make it unique.
*
* URL token analysis:
* https://docs.google.com/document/d/1kwf2PH1APyUR4wrvv6lJIhA12bkQoNVV5KFwAPtWubw/edit?tab=t.0#heading=h.yhx5pggnwgne
*
* @param {string} url - The full URL to register
*
* @returns {string} The short token for the URL (e.g. "GITHUB_COM_1")
*/
encodeToken(url) {
const seenToken = this.urlToToken.get(url);
if (seenToken) {
return seenToken;
}
let baseToken = "";
// Attempt to convert the URL into a base token.
const parsedUrl = URL.parse(url);
if (parsedUrl) {
if (parsedUrl.protocol !== "http:" && parsedUrl.protocol !== "https:") {
// Go ahead and handle URL tokens for more complicated URLs that
// aren't probably supported in the chat interface, but would be useful
// to disambiguate from the HTTP(s) varieties.
baseToken +=
// e.g. "ftp:" -> "FTP"
parsedUrl.protocol.toUpperCase().replace(":", "");
}
// Convert the hostname into a token.
const hostToken = parsedUrl.hostname
.replace(/^www\./, "")
.toUpperCase()
.replace(/[.\-]/g, "_")
.substring(0, 100);
if (hostToken) {
baseToken = baseToken ? `${baseToken}_${hostToken}` : hostToken;
}
// Add on the parts of the URL to the token.
for (let part of parsedUrl.pathname.split("/")) {
if (!part) {
continue;
}
const partToken = part.toUpperCase().replace(/[^A-Z0-9]/g, "_");
const nextToken = `${baseToken}_${partToken}`;
if (nextToken.length > 100) {
break;
}
baseToken = nextToken;
}
} else {
baseToken = "INVALID_URL";
}
let count = this.#baseTokenCounts.get(baseToken) ?? 0;
count += 1;
this.#baseTokenCounts.set(baseToken, count);
const tokenFinal = `${baseToken}_${count}`;
this.urlToToken.set(url, tokenFinal);
this.tokenToUrl.set(tokenFinal, url);
return tokenFinal;
}
}
browser/components/aiwindow/ui/moz.build
+1
View File
@@ -65,6 +65,7 @@ MOZ_SRC_FILES += [
"modules/ToolActionLog.sys.mjs",
"modules/ToolUI.sys.mjs",
"modules/ToolUITelemetry.sys.mjs",
"modules/UrlTokenizer.sys.mjs",
]
TESTING_JS_MODULES += [
browser/components/aiwindow/ui/test/browser/browser_smartformfill_stored_values.js
+2 -2
View File
@@ -55,8 +55,8 @@ async function saveStoredValues(address) {
*/
function respondWithTokenFills(mockEngineManager) {
return respondWithGeneratedFields(mockEngineManager, [
{ action: "fill_from_token", token: "$EMAIL_1", confidence: "high" },
{ action: "fill_from_token", token: "$NAME_1", confidence: "high" },
{ action: "fill_from_token", value: "§EMAIL_1§", confidence: "high" },
{ action: "fill_from_token", value: "§NAME_1§", confidence: "high" },
]);
}
browser/components/aiwindow/ui/test/browser/head_smartformfill_autocomplete.js
+8 -2
View File
@@ -8,6 +8,9 @@
const { MockEngineManager } = ChromeUtils.importESModule(
"resource://testing-common/AIWindowTestUtils.sys.mjs"
);
const { UrlTokenizer } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/ui/modules/UrlTokenizer.sys.mjs"
);
const FORM_URL =
"https://example.com/browser/browser/components/aiwindow/ui/test/browser/test_smartformfill_autocomplete.html";
@@ -183,9 +186,12 @@ function respondToMetadataRequest(
return;
case RELEVANT_TABS_SCHEMA: {
const source = selectedSourceUrl
? requestData.tabs.find(tab => tab.url === selectedSourceUrl)
const selectedSourceToken = selectedSourceUrl
? new UrlTokenizer().encodeToken(selectedSourceUrl, false)
: null;
const source = requestData.tabs.find(
tab => tab.url === selectedSourceToken
);
respond(
JSON.stringify({
selectedTabs: source
browser/components/aiwindow/ui/test/browser/head_smartformfill_form_review.js
+2 -1
View File
@@ -93,7 +93,8 @@ function respondToFormReviewMetadataRequest(schemaName, requestData, respond) {
return;
case RELEVANT_TABS_SCHEMA: {
const source = requestData.tabs.find(tab => tab.url === SOURCE_URL);
const sourceToken = new UrlTokenizer().encodeToken(SOURCE_URL);
const source = requestData.tabs.find(tab => tab.url === sourceToken);
respond(
JSON.stringify({
selectedTabs: source
browser/components/aiwindow/ui/test/xpcshell/test_UrlTokenizer.js
+87
View File
@@ -0,0 +1,87 @@
/* Any copyright is dedicated to the Public Domain.
* http://creativecommons.org/publicdomain/zero/1.0/ */
const { UrlTokenizer } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/ui/modules/UrlTokenizer.sys.mjs"
);
add_task(async function test_UrlTokenizer_encodeToken() {
const cases = [
{
message: "Works for a URL with a path.",
url: "http://www.github.com/foo/bar/baz",
expected: "GITHUB_COM_FOO_BAR_BAZ_1",
},
{
message:
"Returns a new number for a URL that is different but creates the same token.",
url: "http://www.github.com/foo/bar/baz?ignored",
expected: "GITHUB_COM_FOO_BAR_BAZ_2",
},
{
message: "Returns the exact same token given another URL",
url: "http://www.github.com/foo/bar/baz",
expected: "GITHUB_COM_FOO_BAR_BAZ_1",
},
{
message:
"Returns a different token given the same URL with a different protocol",
url: "https://www.github.com/foo/bar/baz",
expected: "GITHUB_COM_FOO_BAR_BAZ_3",
},
{
message: "Can handle about URLs.",
url: "about:config",
expected: "ABOUT_CONFIG_1",
},
{
message: "Uses non-http protocols",
url: "ftp://github.com/foo/bar/baz",
expected: "FTP_GITHUB_COM_FOO_BAR_BAZ_1",
},
{
message: "Uses invalid protocols",
url: "asdf://github.com/foo/bar/baz",
expected: "ASDF_GITHUB_COM_FOO_BAR_BAZ_1",
},
{
message: "Ignores the port.",
url: "http://github.com:1234/ignore/port",
expected: "GITHUB_COM_IGNORE_PORT_1",
},
{
message: "Ignores the params.",
url: "http://www.github.com/ignore/params?token=xxx",
expected: "GITHUB_COM_IGNORE_PARAMS_1",
},
{
message: "Ignores the hash.",
url: "http://www.github.com/ignore/hash/part?token=xxx#hash",
expected: "GITHUB_COM_IGNORE_HASH_PART_1",
},
{
message: "Truncates text in the host from 110 to 100.",
url: `http://www.${"a".repeat(110)}.com/foo`,
expected: "A".repeat(100) + "_1",
},
{
message: "Skips text in the path that is too long",
url: `http://github.com/skip/long/path/` + "A".repeat(100),
expected: "GITHUB_COM_SKIP_LONG_PATH_1",
},
];
const urlTokenizer = new UrlTokenizer({});
for (const { message, url, expected } of cases) {
const token = urlTokenizer.encodeToken(url);
Assert.equal(token, expected, message);
const decodedUrl = urlTokenizer.tokenToUrl.get(token);
Assert.equal(
decodedUrl,
url,
`Expected the decoded token (${token}) to equal '${url}' but got: ${decodedUrl}`
);
}
});
browser/components/aiwindow/ui/test/xpcshell/xpcshell.toml
+2
View File
@@ -41,6 +41,8 @@ run-if = [
["test_ToolUIUpdate.js"]
["test_UrlTokenizer.js"]
["test_chat-utils.js"]
["test_ui_modules_AIWindow_placesEvents.js"]