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
*/
/**
* Self-contained flow backing the search_the_web tool. One Exa retrieval, an
* on-demand page-read loop (bounded), grounded answer generation on a pinned
* model, and a non-LLM schema validation. Returns a structured result the main
* assistant uses to decide whether to answer in chat or fall back to a Google
* handoff. All the logic lives here as functions — there is no separate agent
* object because there is no state to carry between calls.
*/
/**
* @import { ChatConversation } from "moz-src:///browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs"
*/
import {
renderPrompt,
MODEL_FEATURES,
parseAndExtractJSON,
} from "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs";
import { sanitizeUntrustedContent } from "moz-src:///browser/components/aiwindow/models/ChatUtils.sys.mjs";
import { openAIEngine } from "moz-src:///browser/components/aiwindow/models/openAIEngine.sys.mjs";
import { ExaSearchProvider } from "moz-src:///browser/components/aiwindow/models/search/SearchProviders.sys.mjs";
import {
GetPageContent,
GET_PAGE_CONTENT,
} from "moz-src:///browser/components/aiwindow/models/Tools.sys.mjs";
const lazy = {};
ChromeUtils.defineLazyGetter(lazy, "console", () =>
console.createInstance({
prefix: "WebSearchFlow",
maxLogLevelPref: "browser.smartwindow.conversation.logLevel",
})
);
ChromeUtils.defineESModuleGetters(lazy, {
buildConversation:
"moz-src:///browser/components/aiwindow/models/PromptLoader.sys.mjs",
loadPrompt:
"moz-src:///browser/components/aiwindow/models/PromptLoader.sys.mjs",
setTimeout: "resource://gre/modules/Timer.sys.mjs",
clearTimeout: "resource://gre/modules/Timer.sys.mjs",
JsonSchema: "resource://gre/modules/JsonSchema.sys.mjs",
});
// Total result pages that may be read in a turn (the V0 "up to 3" cap).
const MAX_PAGES = 3;
// Max generation rounds that may issue a page read. Bounds the loop
// independently of MAX_PAGES so a round requesting several URLs still counts
// against the page cap, not the round cap.
const MAX_READ_ROUNDS = 3;
// Max time to wait for a single page-read batch before falling back. A slow or
// hanging page must not stall the whole search, so on timeout we tell the model
// to answer with what it already has. Pref-backed so tests can shrink it.
const READ_TIMEOUT_PREF = "browser.smartwindow.search.readTimeoutMs";
const READ_TIMEOUT_DEFAULT_MS = 15000;
/**
* JSON schema describing the structured answer the model emits. Used both as
* the model's response format and as the contract validateSearchAnswer checks.
*
* @type {object}
*/
export const SEARCH_ANSWER_SCHEMA = {
name: "SearchAnswer",
schema: {
type: "object",
additionalProperties: false,
required: ["answer", "could_answer", "confidence"],
properties: {
answer: { type: "string" },
could_answer: { type: "boolean" },
confidence: { type: "number", minimum: 0, maximum: 1 },
},
},
};
/**
* Page-read tool offered to the model.
*
* NOTE: intentionally a result-id variant of get_page_content. Results are
* shown to the model with short ids (result_1, result_2, …) and the model
* passes those ids back rather than full URLs, which models reproduce
* unreliably. Ids are mapped back to URLs before fetching. Do NOT replace this
* with the canonical token-based config — the parameter semantics differ.
*
* @type {object}
*/
const GET_PAGE_CONTENT_TOOL = {
type: "function",
function: {
name: GET_PAGE_CONTENT,
description:
"Read the full text of one or more web pages from the provided search " +
"results. Reference results by their id (for example result_1).",
parameters: {
type: "object",
properties: {
result_ids: {
type: "array",
items: {
type: "string",
description:
"A result id shown in the search results, for example result_1.",
},
minItems: 1,
description: "List of result ids to read.",
},
},
required: ["result_ids"],
},
},
};
/**
* One web search result passed to the model.
*
* @typedef {object} SearchResult
* @property {string} title - Result title.
* @property {string} url - Result URL.
* @property {string} snippet - Short text excerpt for the result.
*/
/**
* The structured result returned to the main assistant.
*
* @typedef {object} SearchWorkflowResult
* @property {string} answer - Grounded answer text (empty on failure).
* @property {boolean} could_answer - Model self-assessment of sufficiency.
* @property {number} confidence - Calibrated confidence in [0, 1].
* @property {string[]} searched_urls - URLs Exa returned (code-tracked).
* @property {string[]} read_urls - URLs the flow actually fetched (code-tracked).
* @property {boolean} requiresSearchHandoff - Boolean indicating whether to route to search handoff
* @property {string} [error] - Present when the flow could not run.
*/
/**
* Stable id shown to the model for the result at `index` (result_1, …). Shared
* by the rendered results and the id->URL map so the two stay aligned.
*
* @param {number} index - Zero-based result index.
* @returns {string}
*/
function resultIdFor(index) {
return `result_${index + 1}`;
}
/**
* Returns true only for well-formed http(s) URLs. Used to drop search results
* with malformed or non-web (e.g. javascript:, data:) URLs before they are
* shown to the model or fetched.
*
* @param {unknown} url
* @returns {boolean}
*/
function isValidHttpUrl(url) {
if (typeof url !== "string") {
return false;
}
const parsed = URL.parse(url);
return parsed?.protocol === "https:" || parsed?.protocol === "http:";
}
/**
* Renders the user message handed to the model: the query, optional caller
* context, the current date for freshness judgements, and the formatted
* search results (each tagged with a result id).
*
* @param {object} params
* @param {string} params.query
* @param {string} [params.context]
* @param {string} params.nowISO - Current date as 'YYYY-MM-DD'.
* @param {SearchResult[]} params.results
* @returns {string}
*/
function buildUserMessage({ query, context, nowISO, results }) {
const lines = [`Current date: ${nowISO}`, `Query: ${query}`];
if (context && context.trim()) {
lines.push(`Context: ${context}`);
}
lines.push("", "Search results:");
results.forEach((result, index) => {
const id = resultIdFor(index);
const title = sanitizeUntrustedContent(result.title || "");
const snippet = sanitizeUntrustedContent(result.snippet || "");
const heading = title ? `${title} — ${result.url}` : result.url;
lines.push(`[${id}] ${heading}`);
if (snippet) {
lines.push(` ${snippet}`);
}
});
return lines.join("\n");
}
/**
* Validates the parsed model output against SEARCH_ANSWER_SCHEMA. This is the
* V0 verification step: a non-LLM schema check, not a separate verifier model.
* On any schema mismatch (missing/extra/wrong-typed fields, out-of-range
* confidence) or a semantically unusable answer (claiming an answer with no
* text), it returns the conservative fallback (could_answer false, confidence
* 0) so a malformed result routes to the Google handoff.
*
* @param {object|null} parsed - Raw parsed model output.
* @returns {{answer: string, could_answer: boolean, confidence: number}}
*/
export function validateSearchAnswer(parsed) {
// Conservative fallback: an unusable answer routes to the Google handoff.
const fallback = { answer: "", could_answer: false, confidence: 0 };
// Structural validation against the same schema we request from the model.
const { valid } = lazy.JsonSchema.validate(
parsed,
SEARCH_ANSWER_SCHEMA.schema
);
if (!valid) {
return fallback;
}
// Schema-valid but semantically unusable: claiming an answer with no text.
if (parsed.could_answer && !parsed.answer.trim()) {
return fallback;
}
return {
answer: parsed.answer,
could_answer: parsed.could_answer,
confidence: parsed.confidence,
};
}
/**
* Generates the grounded answer on the pinned answer-generation model. Runs an
* on-demand page-read loop (page reads use tool-calling), then forces the
* structured-output schema on a final, tool-free turn.
*
* A forced json_schema response and tool-calling cannot coexist in one turn —
* with the schema forced the model must emit the final JSON and can't issue a
* get_page_content call — so reads happen with tools and the schema is applied
* only on the final answer turn.
*
* @param {object} params
* @param {string} params.query
* @param {string} [params.context]
* @param {string} params.nowISO
* @param {SearchResult[]} params.results
* @param {string} params.fxAccountToken
* @param {(urls: string[]) => Promise<string[]>} params.readPage - Executes a
* page read and returns the extracted text blocks.
* @param {AbortSignal} [params.signal]
* @param {string|null} [params.flowId]
* @returns {Promise<object|null>} The parsed model output, or null when the
* model produced no parseable answer.
*/
async function generateAnswer({
query,
context,
nowISO,
results,
fxAccountToken,
readPage,
signal,
flowId = null,
}) {
const [conversation, { prompt: rawPrompt }] = await Promise.all([
lazy.buildConversation(MODEL_FEATURES.SEARCH_ANSWER_GENERATION, {
flowId,
}),
lazy.loadPrompt(MODEL_FEATURES.SEARCH_ANSWER_GENERATION),
]);
conversation.setSystemMessage(renderPrompt(rawPrompt, {}));
conversation.addUserMessage(
buildUserMessage({ query, context, nowISO, results })
);
// Map the result ids shown to the model (result_1, …) back to URLs; the
// model references ids, not full URLs (see GET_PAGE_CONTENT_TOOL).
const idToUrl = new Map(
results.map((result, index) => [resultIdFor(index), result.url])
);
let reads = 0;
while (true) {
const allowReads = reads < MAX_READ_ROUNDS;
// Object content so the streaming accumulator can append to `body`.
const assistantMessage = conversation.addAssistantMessage({ body: "" });
const { pendingToolCalls } = await conversation.receiveResponse(
conversation.runWithGenerator({
streamOptions: { enabled: true },
fxAccountToken,
chatId: conversation.id,
tool_choice: allowReads ? "auto" : "none",
tools: allowReads ? [GET_PAGE_CONTENT_TOOL] : [],
signal,
}),
assistantMessage
);
const pageReadCalls =
allowReads && pendingToolCalls
? pendingToolCalls.filter(
call => call.function?.name === GET_PAGE_CONTENT
)
: [];
if (!pageReadCalls.length) {
break;
}
// Nest tool_calls under `body` to match the shape the wire serializer
// reads, so they are preserved on the next request.
assistantMessage.content = {
body: {
tool_calls: pageReadCalls.map(call => ({
id: call.id,
type: "function",
function: {
name: call.function.name,
arguments: call.function.arguments || "{}",
},
})),
},
};
for (const call of pageReadCalls) {
let ids = [];
try {
ids = JSON.parse(call.function.arguments || "{}").result_ids || [];
} catch {
ids = [];
}
const urls = (Array.isArray(ids) ? ids : [])
.map(id => idToUrl.get(id))
.filter(Boolean);
const texts = await readPage(urls);
conversation.addToolMessage({
tool_call_id: call.id,
content: Array.isArray(texts) ? texts.join("\n\n") : String(texts),
name: call.function.name,
});
}
reads++;
}
// Force the structured schema on a final, tool-free turn to get the answer.
conversation.addUserMessage(
"Output only the final JSON object described in your instructions."
);
// Don't pass `signal` to the non-streaming run(): it forwards options to the
// engine process, which cannot structured-clone an AbortSignal. Check for a
// prior abort here instead.
if (signal?.aborted) {
throw new DOMException("Aborted", "AbortError");
}
const response = await conversation.run({
tool_choice: "none",
tools: [],
fxAccountToken,
});
return parseAndExtractJSON(response, null);
}
/**
* Builds a failure result that routes the main assistant to fallback.
*
* @param {string[]} searchedUrls
* @param {string[]} readUrls
* @param {string} message
* @returns {SearchWorkflowResult}
*/
function failure(searchedUrls, readUrls, message) {
return {
answer: "",
could_answer: false,
confidence: 0,
searched_urls: searchedUrls,
read_urls: readUrls,
error: message,
requiresSearchHandoff: false,
};
}
function shouldCallSearchHandoff(conversation) {
return conversation._searchTheWebTurn === conversation.currentTurnIndex();
}
/**
* Tool entrypoint for search_the_web. Runs one Exa retrieval, drives the
* on-demand page reads (bounded by MAX_PAGES), generates and validates the
* grounded answer, and returns the structured result. Errors are returned as a
* result with could_answer false rather than thrown, so the main assistant can
* fall back to the Google handoff.
*
* @param {object} toolParams
* @param {string} toolParams.query - Search query (may be rewritten by the assistant).
* @param {string} [toolParams.context] - Optional caller-supplied context.
* @param {ChatConversation} conversation - Originating conversation; owns the
* anonymous-fetch ledger and security state the page reads depend on.
* @param {AbortSignal} [signal] - Cancels the in-flight answer generation.
* @returns {Promise<SearchWorkflowResult>}
*/
export async function runSearchTheWeb(toolParams, conversation, signal) {
if (shouldCallSearchHandoff(conversation)) {
return {
requiresSearchHandoff: true,
};
}
const query = toolParams?.query;
if (typeof query !== "string" || !query.trim()) {
return failure([], [], "a non-empty query is required");
}
conversation._searchTheWebTurn = conversation.currentTurnIndex();
const context =
typeof toolParams?.context === "string" ? toolParams.context : "";
let results;
try {
const provider = new ExaSearchProvider();
const response = await provider.search(query.trim(), {
maxResults: ExaSearchProvider.MAX_RESULTS,
});
results = Array.isArray(response?.results) ? response.results : [];
} catch (e) {
console.error("[SearchWorkflow] retrieval failed:", e);
return failure([], [], e.message);
}
results = results.filter(item => isValidHttpUrl(item?.url));
const searchedUrls = results.map(item => item.url);
if (!searchedUrls.length) {
return failure([], [], "no search results");
}
// Record all results as seen and add them to the anonymous-fetch ledger,
// then mark the conversation as having seen private + untrusted content. The
// per-turn read limit (MAX_PAGES) caps how many are actually read.
conversation.addSeenUrls(searchedUrls);
conversation.addSerpUrlsForAnonymousFetch(searchedUrls);
conversation.securityProperties.setPrivateData();
conversation.securityProperties.setUntrustedInput();
const fetchableSet = new Set(searchedUrls);
const readUrls = [];
const readSet = new Set();
const readPage = async requestedUrls => {
const remaining = MAX_PAGES - readUrls.length;
if (remaining <= 0) {
return ["Page read limit reached. Answer using what you have gathered."];
}
const fresh = requestedUrls
.filter(url => fetchableSet.has(url) && !readSet.has(url))
.slice(0, remaining);
if (!fresh.length) {
return [
"No further readable pages are available. Answer using what you have gathered.",
];
}
fresh.forEach(url => {
readSet.add(url);
readUrls.push(url);
});
const readTimeoutMs = Services.prefs.getIntPref(
READ_TIMEOUT_PREF,
READ_TIMEOUT_DEFAULT_MS
);
// Reads pages individually (but async) with timeouts as specified so that a
// single read failure doesn't kill the whole batch
const readOne = async url => {
const controller = new AbortController();
let timeoutId;
const timeout = new Promise(resolve => {
timeoutId = lazy.setTimeout(() => {
lazy.console.warn(
`[SearchWorkflow] page read timed out reading ${url} after ${readTimeoutMs}ms; abandoning stuck fetch`
);
controller.abort();
resolve([
`Timed out reading ${url}. Answer with what you have gathered.`,
]);
}, readTimeoutMs);
});
try {
const fetchPromise = GetPageContent.getPageContent(
{ url_list: [url], signal: controller.signal },
conversation
);
// If the timeout wins (stuck load), the fetch stays pending; swallow its
// late settle so it can't surface as an unhandled rejection.
fetchPromise.catch(() => {});
return await Promise.race([fetchPromise, timeout]);
} finally {
lazy.clearTimeout(timeoutId);
}
};
const perUrl = await Promise.all(fresh.map(readOne));
return perUrl.flat();
};
let parsed;
try {
parsed = await generateAnswer({
query: query.trim(),
context,
nowISO: new Date().toISOString().slice(0, 10),
results,
fxAccountToken: await openAIEngine.getFxAccountToken(),
readPage,
signal,
});
} catch (e) {
console.error("[SearchWorkflow] answer generation failed:", e);
return failure(searchedUrls, readUrls, e.message);
}
const validated = validateSearchAnswer(parsed);
lazy.console.log("[Tool] searchTheWeb", {
query,
searched: searchedUrls.length,
read: readUrls.length,
couldAnswer: validated.could_answer,
});
return {
...validated,
searched_urls: searchedUrls,
read_urls: readUrls,
requiresSearchHandoff: false,
};
}