tor-browser

test_smart_assist_engine.js (7549B)

---

/* Any copyright is dedicated to the Public Domain.
* http://creativecommons.org/publicdomain/zero/1.0/ */

const { SmartAssistEngine } = ChromeUtils.importESModule(
 "moz-src:///browser/components/genai/SmartAssistEngine.sys.mjs"
);

const { sinon } = ChromeUtils.importESModule(
 "resource://testing-common/Sinon.sys.mjs"
);

// Prefs
const PREF_API_KEY = "browser.ml.smartAssist.apiKey";
const PREF_ENDPOINT = "browser.ml.smartAssist.endpoint";
const PREF_MODEL = "browser.ml.smartAssist.model";

// Clean prefs after all tests
registerCleanupFunction(() => {
 for (let pref of [PREF_API_KEY, PREF_ENDPOINT, PREF_MODEL]) {
   if (Services.prefs.prefHasUserValue(pref)) {
     Services.prefs.clearUserPref(pref);
   }
 }
});

add_task(async function test_createOpenAIEngine_uses_prefs_and_static_fields() {
 Services.prefs.setStringPref(PREF_API_KEY, "test-key-123");
 Services.prefs.setStringPref(PREF_ENDPOINT, "https://example.test/v1");
 Services.prefs.setStringPref(PREF_MODEL, "gpt-fake");

 const sb = sinon.createSandbox();
 try {
   const fakeEngine = {
     runWithGenerator() {
       throw new Error("not used");
     },
   };
   const stub = sb
     .stub(SmartAssistEngine, "_createEngine")
     .resolves(fakeEngine);

   const engine = await SmartAssistEngine.createOpenAIEngine();

   Assert.strictEqual(
     engine,
     fakeEngine,
     "Should return engine from _createEngine"
   );
   Assert.ok(stub.calledOnce, "_createEngine should be called once");

   const opts = stub.firstCall.args[0];
   Assert.equal(opts.apiKey, "test-key-123", "apiKey should come from pref");
   Assert.equal(
     opts.baseURL,
     "https://example.test/v1",
     "baseURL should come from pref"
   );
   Assert.equal(opts.modelId, "gpt-fake", "modelId should come from pref");
 } finally {
   sb.restore();
 }
});

add_task(async function test_fetchWithHistory_streams_and_forwards_args() {
 const sb = sinon.createSandbox();
 try {
   let capturedArgs = null;
   let capturedStreamOption = null;

   // Fake async generator that yields three text chunks and one empty (ignored)
   const fakeEngine = {
     runWithGenerator({ streamOptions, args }) {
       capturedArgs = args;
       capturedStreamOption = streamOptions;
       async function* gen() {
         yield { text: "Hello" };
         yield { text: " from" };
         yield { text: " fake engine!" };
         yield {}; // ignored by SmartAssistEngine
       }
       return gen();
     },
   };

   sb.stub(SmartAssistEngine, "_createEngine").resolves(fakeEngine);
   sb.stub(SmartAssistEngine, "_getFxAccountToken").resolves("mock_token");

   const messages = [
     { role: "system", content: "You are helpful" },
     { role: "user", content: "Hi there" },
   ];

   // Collect streamed output
   let acc = "";
   for await (const t of SmartAssistEngine.fetchWithHistory(messages)) {
     acc += t;
   }

   Assert.equal(
     acc,
     "Hello from fake engine!",
     "Should concatenate streamed chunks"
   );
   Assert.deepEqual(
     capturedArgs,
     messages,
     "Should forward messages as args to runWithGenerator()"
   );
   Assert.deepEqual(
     capturedStreamOption.enabled,
     true,
     "Should enable streaming in runWithGenerator()"
   );
 } finally {
   sb.restore();
 }
});

add_task(
 async function test_fetchWithHistory_propagates_engine_creation_rejection() {
   const sb = sinon.createSandbox();
   try {
     const err = new Error("creation failed (generic)");
     const stub = sb.stub(SmartAssistEngine, "_createEngine").rejects(err);
     sb.stub(SmartAssistEngine, "_getFxAccountToken").resolves("mock_token");
     const messages = [{ role: "user", content: "Hi" }];

     // Must CONSUME the async generator to trigger the rejection
     const consume = async () => {
       for await (const _message of SmartAssistEngine.fetchWithHistory(
         messages
       )) {
         void _message;
       }
     };

     await Assert.rejects(
       consume(),
       e => e === err,
       "Should propagate the same error thrown by _createEngine"
     );
     Assert.ok(stub.calledOnce, "_createEngine should be called once");
   } finally {
     sb.restore();
   }
 }
);

add_task(async function test_fetchWithHistory_propagates_stream_error() {
 const sb = sinon.createSandbox();
 try {
   const fakeEngine = {
     runWithGenerator() {
       async function* gen() {
         yield { text: "partial" };
         throw new Error("engine stream boom");
       }
       return gen();
     },
   };
   sb.stub(SmartAssistEngine, "_createEngine").resolves(fakeEngine);
   sb.stub(SmartAssistEngine, "_getFxAccountToken").resolves("mock_token");

   const consume = async () => {
     let acc = "";
     for await (const t of SmartAssistEngine.fetchWithHistory([
       { role: "user", content: "x" },
     ])) {
       acc += t;
     }
     return acc;
   };

   await Assert.rejects(
     consume(),
     e => /engine stream boom/.test(e.message),
     "Should propagate errors thrown during streaming"
   );
 } finally {
   sb.restore();
 }
});

add_task(async function test_getPromptIntent_basic() {
 const sb = sinon.createSandbox();
 const cases = [
   { prompt: "please search for news on firefox", expected: "search" },
   { prompt: "Can you FIND me the docs for PageAssist?", expected: "search" }, // case-insensitive
   { prompt: "look up the best pizza in SF", expected: "search" },
   { prompt: "hello there, how are you?", expected: "chat" },
   { prompt: "tell me a joke", expected: "chat" },
 ];

 const fakeEngine = {
   run({ args: [[query]] }) {
     const searchKeywords = [
       "search",
       "find",
       "look",
       "query",
       "locate",
       "explore",
     ];
     const formattedPrompt = query.toLowerCase();
     const isSearch = searchKeywords.some(keyword =>
       formattedPrompt.includes(keyword)
     );

     // Simulate model confidence scores
     if (isSearch) {
       return [
         { label: "search", score: 0.95 },
         { label: "chat", score: 0.05 },
       ];
     }
     return [
       { label: "chat", score: 0.95 },
       { label: "search", score: 0.05 },
     ];
   },
 };

 sb.stub(SmartAssistEngine, "_createEngine").resolves(fakeEngine);

 for (const { prompt, expected } of cases) {
   const intent = await SmartAssistEngine.getPromptIntent(prompt);
   Assert.equal(
     intent,
     expected,
     `getPromptIntent("${prompt}") should return "${expected}"`
   );
 }

 sb.restore();
});

add_task(async function test_preprocessQuery_removes_question_marks() {
 const sb = sinon.createSandbox();

 // Use a minimal fake SmartAssistEngine to test only preprocessing
 const engine = Object.create(SmartAssistEngine);

 const cases = [
   { input: "hello?", expected: "hello" },
   { input: "?prompt", expected: "prompt" },
   { input: "multiple???", expected: "multiple" },
   { input: "mid?dle", expected: "middle" },
   { input: "question? ", expected: "question" },
   { input: " no?  spaces? ", expected: "no  spaces" },
   { input: "???", expected: "" },
   { input: "clean input", expected: "clean input" },
 ];

 for (const { input, expected } of cases) {
   const result = engine._preprocessQuery(input);
   Assert.equal(
     result,
     expected,
     `Expected "${input}" to preprocess to "${expected}", got "${result}"`
   );
 }

 sb.restore();
});