528 lines
16 KiB
JavaScript
528 lines
16 KiB
JavaScript
/**
|
|
* 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/.
|
|
*/
|
|
|
|
import { XPCOMUtils } from "resource://gre/modules/XPCOMUtils.sys.mjs";
|
|
import {
|
|
renderPrompt,
|
|
checkMajorVersion,
|
|
} from "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs";
|
|
|
|
const lazy = XPCOMUtils.declareLazy({
|
|
RemoteSettings: "resource://services-settings/remote-settings.sys.mjs",
|
|
setTimeout: "resource://gre/modules/Timer.sys.mjs",
|
|
ChatStore:
|
|
"moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs",
|
|
openAIEngine:
|
|
"moz-src:///browser/components/aiwindow/models/openAIEngine.sys.mjs",
|
|
});
|
|
|
|
ChromeUtils.defineLazyGetter(lazy, "console", function () {
|
|
return console.createInstance({
|
|
prefix: "LLMTelemetry",
|
|
maxLogLevelPref: "browser.smartwindow.telemetryLogLevel",
|
|
});
|
|
});
|
|
|
|
export const TELEMETRY_MAJOR_VERSIONS = Object.freeze({
|
|
wasSuccessful: 1,
|
|
conversationCategory: 1,
|
|
});
|
|
|
|
export const TRIGGER_MAJOR_VERSIONS = Object.freeze({
|
|
uniform_sample: 1,
|
|
uniform_sample_at_turn: 1,
|
|
min_turns: 1,
|
|
});
|
|
|
|
// Registry of named check strategies
|
|
export const TRIGGER_CHECK_STRATEGIES = {
|
|
uniform_sample: (conversation, _params) =>
|
|
conversation.currentTurnIndex?.() == 0,
|
|
uniform_sample_at_turn: (conversation, { turn }) =>
|
|
conversation._telemetryUniformSample === true &&
|
|
conversation.currentTurnIndex?.() == turn,
|
|
min_turns: (conversation, { minTurns }) =>
|
|
(conversation.currentTurnIndex?.() ?? 0) >= minTurns,
|
|
};
|
|
|
|
// Probability that a turn-1 conversation is flagged for uniform sampling.
|
|
export const UNIFORM_SAMPLING_NAME = "uniform_sample";
|
|
const UNKNOWN = "unknown";
|
|
|
|
const RS_TELEMETRY_PROMPTS_COLLECTION = "ai-window-telemetry-prompts";
|
|
const RS_TELEMETRY_TRIGGERS_COLLECTION = "ai-window-telemetry-triggers";
|
|
|
|
// ============================================================
|
|
// Trigger
|
|
// ============================================================
|
|
|
|
/**
|
|
* Represents a single telemetry trigger.
|
|
*
|
|
* A trigger is evaluated against the conversation after each agent response.
|
|
* If it fires and passes the sampling check, the trigger is included in the
|
|
* set passed to runTelemetry.
|
|
*/
|
|
export class Trigger {
|
|
/**
|
|
* @param {string} name - Unique identifier, must match the "triggers" field in RS records.
|
|
* @param {function(object): boolean} checkFn - Returns true if the trigger should fire.
|
|
* @param {number} samplingProbability - Fraction [0, 1] of fires that proceed to evaluation.
|
|
* @param {string} description - Description of trigger
|
|
*/
|
|
constructor(name, checkFn, samplingProbability = 1.0, description = "") {
|
|
this.name = name;
|
|
this.check = checkFn;
|
|
this.samplingProbability = samplingProbability;
|
|
this.description = description;
|
|
}
|
|
}
|
|
|
|
// ============================================================
|
|
// telemetryPromptEngine
|
|
// ============================================================
|
|
|
|
/**
|
|
* Remote Settings records in RS_TELEMETRY_COLLECTION are expected to have:
|
|
* - telemetry_name {string} Unique identifier for the evaluation prompt
|
|
* - triggers {string[]} Trigger names this prompt applies to
|
|
* - prompts {string} System prompt content
|
|
* - model {string} Model ID to use for inference
|
|
*/
|
|
export class TelemetryPromptEngine {
|
|
/** @type {object} */
|
|
#promptRecord;
|
|
|
|
/** @type {object} */
|
|
#engineInstance;
|
|
|
|
/**
|
|
* Builds a telemetryPromptEngine for the given RS record.
|
|
*
|
|
* @param {object} promptRecord
|
|
* @returns {Promise<TelemetryPromptEngine>}
|
|
*/
|
|
static async build(promptRecord) {
|
|
const engine = new TelemetryPromptEngine();
|
|
engine.#promptRecord = promptRecord;
|
|
|
|
const extraHeadersPref = Services.prefs.getStringPref(
|
|
"browser.smartwindow.extraHeaders",
|
|
"{}"
|
|
);
|
|
let extraHeaders = {};
|
|
try {
|
|
extraHeaders = JSON.parse(extraHeadersPref);
|
|
} catch (e) {
|
|
lazy.console.error(
|
|
"Failed to parse extra headers for telemetry engine:",
|
|
e
|
|
);
|
|
}
|
|
|
|
engine.#engineInstance = await lazy.openAIEngine._createEngine({
|
|
apiKey: "",
|
|
backend: "openai",
|
|
baseURL: lazy.openAIEngine.endpoint,
|
|
featureId: "llm-telemetry",
|
|
flowId: null,
|
|
modelId: promptRecord.model ?? null,
|
|
modelRevision: "main",
|
|
taskName: "text-generation",
|
|
serviceType: promptRecord.service_type ?? "ai",
|
|
purpose: promptRecord.purpose ?? "chat",
|
|
extraHeaders,
|
|
});
|
|
|
|
return engine;
|
|
}
|
|
|
|
verifyResult(result) {
|
|
const schema = this.#promptRecord.output_schema;
|
|
const finalMapping = {};
|
|
|
|
let metrics;
|
|
try {
|
|
metrics = JSON.parse(result.finalOutput);
|
|
} catch {
|
|
return Object.fromEntries(Object.keys(schema).map(key => [key, UNKNOWN]));
|
|
}
|
|
|
|
for (const key of Object.keys(schema)) {
|
|
if (!schema[key]?.includes(metrics[key])) {
|
|
finalMapping[key] = UNKNOWN;
|
|
} else {
|
|
finalMapping[key] = metrics[key];
|
|
}
|
|
}
|
|
|
|
return finalMapping;
|
|
}
|
|
|
|
/**
|
|
* Runs the evaluation prompt against the conversation.
|
|
*
|
|
* @param {object} conversation
|
|
* @returns {Promise<object>} Raw LLM response
|
|
*/
|
|
async run(conversation) {
|
|
const fxAccountToken = await lazy.openAIEngine.getFxAccountToken();
|
|
// The LLM-as-judge model needs the raw URLs, not the chat URL tokens.
|
|
const messages = conversation.getMessagesInChatCompletionsFormat({
|
|
applyUrlTokens: false,
|
|
});
|
|
|
|
const prompt = renderPrompt(this.#promptRecord.prompt, {
|
|
chatConversation: JSON.stringify(messages),
|
|
fields: JSON.stringify(Object.keys(this.#promptRecord.output_schema)),
|
|
});
|
|
|
|
const result = await this.#engineInstance.run({
|
|
fxAccountToken,
|
|
args: [{ role: "system", content: prompt }],
|
|
});
|
|
|
|
lazy.console.debug("Returning: ", result);
|
|
return this.verifyResult(result);
|
|
}
|
|
}
|
|
|
|
// ============================================================
|
|
// TelemetryEngine
|
|
// ============================================================
|
|
/**
|
|
* Orchestrates conversation telemetry evaluation. Loads trigger definitions
|
|
* from Remote Settings, checks which triggers fire for a given conversation,
|
|
* and runs LLM-based prompt evaluations for each matched trigger.
|
|
*
|
|
*/
|
|
export class TelemetryEngine {
|
|
_triggers = null;
|
|
|
|
/**
|
|
*
|
|
* @param {object} conversation
|
|
* @returns {Trigger[]}
|
|
*/
|
|
|
|
async _fetchRecords(collection, fallback = []) {
|
|
try {
|
|
const client = lazy.RemoteSettings(collection);
|
|
const records = await client.get();
|
|
return records;
|
|
} catch (e) {
|
|
lazy.console.error("Telemetry: failed to fetch records:", e);
|
|
return fallback;
|
|
}
|
|
}
|
|
|
|
async getTriggerDefinitions() {
|
|
if (this._triggers) {
|
|
return;
|
|
}
|
|
|
|
const triggerRecords = await this._fetchRecords(
|
|
RS_TELEMETRY_TRIGGERS_COLLECTION
|
|
);
|
|
|
|
this._triggers = triggerRecords
|
|
.filter(def => TRIGGER_CHECK_STRATEGIES[def.check])
|
|
.filter(def =>
|
|
checkMajorVersion(def.version, TRIGGER_MAJOR_VERSIONS[def.check])
|
|
)
|
|
.map(
|
|
def =>
|
|
new Trigger(
|
|
def.name,
|
|
conversation =>
|
|
TRIGGER_CHECK_STRATEGIES[def.check](
|
|
conversation,
|
|
def.params ?? {}
|
|
),
|
|
def.sampling_probability ?? 1.0,
|
|
def.description ?? ""
|
|
)
|
|
);
|
|
}
|
|
|
|
_getRandom() {
|
|
return Math.random();
|
|
}
|
|
|
|
async getTriggers(conversation) {
|
|
await this.getTriggerDefinitions();
|
|
|
|
if (!conversation._checkedTelemetryTriggers) {
|
|
conversation._checkedTelemetryTriggers = new Set();
|
|
}
|
|
|
|
const fired = [];
|
|
for (const trigger of this._triggers) {
|
|
if (conversation._checkedTelemetryTriggers.has(trigger.name)) {
|
|
lazy.console.debug(
|
|
"checking:",
|
|
trigger.name,
|
|
"-> skip (already checked)"
|
|
);
|
|
continue;
|
|
}
|
|
lazy.console.debug(
|
|
"checking:",
|
|
trigger.name,
|
|
"-> evaluating",
|
|
conversation._telemetryUniformSample,
|
|
conversation.currentTurnIndex?.()
|
|
);
|
|
if (trigger.check(conversation)) {
|
|
conversation._checkedTelemetryTriggers.add(trigger.name);
|
|
if (this._getRandom() < trigger.samplingProbability) {
|
|
lazy.console.debug("checking:", trigger.name, "-> fired");
|
|
fired.push(trigger);
|
|
if (trigger.name == UNIFORM_SAMPLING_NAME) {
|
|
conversation._telemetryUniformSample = true;
|
|
conversation._telemetryUniformProbability =
|
|
trigger.samplingProbability;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return fired;
|
|
}
|
|
|
|
/**
|
|
* Runs LLM-based evaluations for all prompt records that match any of the
|
|
* fired triggers. Prompts that apply to multiple triggers are deduplicated
|
|
* and run only once.
|
|
*
|
|
* @param {Trigger[]} triggers
|
|
* @param {object} conversation
|
|
* @returns {Promise<Array<{feature: string, result: object}>>}
|
|
*/
|
|
async runTelemetry(triggers, conversation) {
|
|
if (!triggers.length) {
|
|
return [];
|
|
}
|
|
|
|
const allRecords = await this._fetchRecords(
|
|
RS_TELEMETRY_PROMPTS_COLLECTION
|
|
);
|
|
const triggerByName = new Map(triggers.map(t => [t.name, t]));
|
|
|
|
const seen = new Set();
|
|
const promptsToRun = [];
|
|
const samplingProbabilities = new Map();
|
|
for (const record of allRecords) {
|
|
const recordTriggers = record.triggers ?? [];
|
|
if (
|
|
!seen.has(record.telemetry_name) && // sometimes multiple triggers will use the same telemetry
|
|
checkMajorVersion(
|
|
record.version,
|
|
TELEMETRY_MAJOR_VERSIONS[record.telemetry_name]
|
|
)
|
|
) {
|
|
const matchingTriggerName = recordTriggers.find(t =>
|
|
triggerByName.has(t)
|
|
);
|
|
if (matchingTriggerName) {
|
|
seen.add(record.telemetry_name);
|
|
promptsToRun.push(record);
|
|
samplingProbabilities.set(
|
|
record.telemetry_name,
|
|
triggerByName.get(matchingTriggerName).samplingProbability
|
|
);
|
|
}
|
|
}
|
|
}
|
|
const results = await this._runPrompts(promptsToRun, conversation);
|
|
return results.map(r => ({
|
|
...r,
|
|
samplingProbability: samplingProbabilities.get(r.telemetry_name),
|
|
}));
|
|
}
|
|
|
|
/**
|
|
* Runs LLM-based evaluations for all telemetry prompts passed by name.
|
|
* Prompt records that have run_terminal flag set to False will not be run
|
|
*
|
|
* @param {string[]} promptNames
|
|
* @param {object} conversation
|
|
* @returns {Promise<Array<{feature: string, result: object}>>}
|
|
*/
|
|
async runTelemetryByName(promptNames, conversation) {
|
|
if (!promptNames.length) {
|
|
return [];
|
|
}
|
|
|
|
const allRecords = await this._fetchRecords(
|
|
RS_TELEMETRY_PROMPTS_COLLECTION
|
|
);
|
|
|
|
const nameSet = new Set(promptNames);
|
|
const promptsToRun = allRecords
|
|
.filter(record => nameSet.has(record.telemetry_name))
|
|
.filter(record =>
|
|
checkMajorVersion(
|
|
record.version,
|
|
TELEMETRY_MAJOR_VERSIONS[record.telemetry_name]
|
|
)
|
|
)
|
|
.filter(record => record.run_terminal);
|
|
|
|
return this._runPrompts(promptsToRun, conversation);
|
|
}
|
|
|
|
async _runPrompts(promptsToRun, conversation) {
|
|
const results = [];
|
|
for (const record of promptsToRun) {
|
|
let lastError;
|
|
for (let attempt = 0; attempt < 3; attempt++) {
|
|
if (attempt > 0) {
|
|
await new Promise(resolve =>
|
|
lazy.setTimeout(resolve, 1000 * 2 ** attempt + Math.random() * 500)
|
|
);
|
|
}
|
|
try {
|
|
const engine = await TelemetryPromptEngine.build(record);
|
|
const result = await engine.run(conversation);
|
|
results.push({
|
|
telemetry_name: record.telemetry_name,
|
|
result,
|
|
telemetry_version: record.version,
|
|
});
|
|
lastError = null;
|
|
break;
|
|
} catch (error) {
|
|
if (lazy.openAIEngine.is429Error(error)) {
|
|
lastError = error;
|
|
} else {
|
|
lazy.console.error(
|
|
`Telemetry: evaluation failed for ${record.telemetry_name}:`,
|
|
error
|
|
); // other errors fail
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
if (lastError) {
|
|
lazy.console.error(
|
|
`Encountered error running batch telemetry job: ${lastError}`
|
|
);
|
|
}
|
|
}
|
|
|
|
lazy.console.debug("Returning!", results);
|
|
return results;
|
|
}
|
|
}
|
|
|
|
export function normalizeMetadata(metadata = {}) {
|
|
const {
|
|
chat_version = "",
|
|
record_type = "",
|
|
uniform_sampling_probability = 0,
|
|
trigger_sampling_probability = 0,
|
|
triggers = [],
|
|
} = metadata;
|
|
|
|
return {
|
|
chat_version,
|
|
record_type,
|
|
uniform_sampling_probability: Math.round(
|
|
uniform_sampling_probability * 1000
|
|
),
|
|
trigger_sampling_probability: Math.round(
|
|
trigger_sampling_probability * 1000
|
|
),
|
|
triggers: Array.isArray(triggers)
|
|
? JSON.stringify(triggers)
|
|
: (triggers ?? "[]"),
|
|
};
|
|
}
|
|
|
|
export function submitTelemetryResult(
|
|
telemetryResults,
|
|
conversation,
|
|
modelId,
|
|
metadata
|
|
) {
|
|
for (const resultObject of telemetryResults ?? []) {
|
|
const normalized_metadata = normalizeMetadata({
|
|
...metadata,
|
|
trigger_sampling_probability: resultObject?.samplingProbability ?? 0.0,
|
|
});
|
|
|
|
const result = resultObject?.result ?? {};
|
|
for (const [attributeName, attributeValue] of Object.entries(result)) {
|
|
Glean.smartWindow.llmajBasedTelemetry.record({
|
|
...normalized_metadata,
|
|
chat_id: conversation.id,
|
|
model: modelId,
|
|
turn_number: conversation.currentTurnIndex(),
|
|
telemetry_name: resultObject?.telemetry_name ?? "",
|
|
attribute_name: attributeName,
|
|
attribute_value: String(attributeValue ?? UNKNOWN),
|
|
uniform_sampled: (metadata.triggers ?? []).includes(
|
|
UNIFORM_SAMPLING_NAME
|
|
),
|
|
trigger_sampled: (metadata.triggers ?? []).some(
|
|
t => t !== UNIFORM_SAMPLING_NAME
|
|
),
|
|
telemetry_version: resultObject?.telemetry_version ?? "",
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Marks conversations as unprocessed (since new turns have been added) and
|
|
* runs LLM-as-judge-based telemetry.
|
|
*
|
|
* @param {ChatConversation} conversation
|
|
*/
|
|
export async function runLLMaJTelemetry(conversation) {
|
|
const turnIndex = conversation.currentTurnIndex();
|
|
await lazy.ChatStore.markLLMTelemetryUnprocessed(conversation.id).catch(e =>
|
|
console.error("Failed to mark telemetry unprocessed:", e)
|
|
);
|
|
|
|
const telemetryEngine = new TelemetryEngine();
|
|
const triggers = await telemetryEngine.getTriggers(conversation);
|
|
const telemetryStart = ChromeUtils.now();
|
|
telemetryEngine
|
|
.runTelemetry(triggers, conversation)
|
|
.then(results => {
|
|
if (!results.length) {
|
|
return;
|
|
}
|
|
submitTelemetryResult(results, conversation, conversation.engine?.model, {
|
|
record_type: "midChat",
|
|
uniform_sampling_probability: conversation._telemetryUniformProbability,
|
|
triggers: triggers.map(t => t.name),
|
|
chat_version: conversation.systemPromptVersion,
|
|
});
|
|
const prompts = Object.fromEntries(
|
|
results.map(r => [r.telemetry_name, turnIndex])
|
|
);
|
|
const probabilities = Object.fromEntries(
|
|
results.map(r => [r.telemetry_name, r.samplingProbability])
|
|
);
|
|
lazy.ChatStore.updateLLMTelemetryRecord(
|
|
conversation.id,
|
|
prompts,
|
|
probabilities,
|
|
conversation._telemetryUniformProbability
|
|
).catch(e => console.error("Failed to update telemetry record:", e));
|
|
})
|
|
.catch(e => console.error("Telemetry run failed:", e))
|
|
.finally(() => {
|
|
ChromeUtils.addProfilerMarker(
|
|
"SmartWindow",
|
|
{ startTime: telemetryStart },
|
|
"LLMaJTelemetry"
|
|
);
|
|
});
|
|
}
|