AI: add support for facts to AidaClient

Change-Id: Ic8d73f09931b823c9faabc3d383ce9e01a8e8fca
Bug: 408172181
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6426044
Commit-Queue: Alex Rudenko <[email protected]>
Reviewed-by: Alex Rudenko <[email protected]>
Auto-Submit: Jack Franklin <[email protected]>

Author: jackfranklin

front_end/core/host/AidaClient.ts

@@ -145,6 +145,25 @@ export enum UserTier {
   PUBLIC = 3,
 }
 
+// Googlers: see the Aida `retrieval` proto; this type is based on that.
+export interface RequestFactMetadata {
+  /**
+   * A description of where the fact comes from.
+   */
+  source: string;
+  /**
+   * Optional: a score to give this fact. Used because if there are more facts than space in the context window, higher scoring facts will be prioritised.
+   */
+  score?: number;
+}
+export interface RequestFact {
+  /**
+   * Content of the fact.
+   */
+  text: string;
+  metadata: RequestFactMetadata;
+}
+
 export type RpcGlobalId = string|number;
 
 export interface AidaRequest {
@@ -156,6 +175,7 @@ export interface AidaRequest {
   historical_contexts?: Content[];
   // eslint-disable-next-line @typescript-eslint/naming-convention
   function_declarations?: FunctionDeclaration[];
+  facts?: RequestFact[];
   options?: {
     temperature?: number,
     // eslint-disable-next-line @typescript-eslint/naming-convention

front_end/models/ai_assistance/agents/AiAgent.test.ts

@@ -185,6 +185,47 @@ describeWithEnvironment('AiAgent', () => {
       assert.isUndefined(request.historical_contexts);
     });
 
+    it('builds a request with a fact', async () => {
+      const agent = new AiAgentMock({
+        aidaClient: mockAidaClient([[{
+          explanation: 'answer',
+        }]]),
+        serverSideLoggingEnabled: true,
+      });
+      const fact: Host.AidaClient.RequestFact = {text: 'This is a fact', metadata: {source: 'devtools'}};
+      agent.addFact(fact);
+      await Array.fromAsync(agent.run('question', {selected: null}));
+      const request = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
+      assert.deepEqual(request.facts, [fact]);
+    });
+
+    it('can manage multiple facts and remove them', async () => {
+      const agent = new AiAgentMock({
+        aidaClient: mockAidaClient([[{
+          explanation: 'answer',
+        }]]),
+        serverSideLoggingEnabled: true,
+      });
+      const f1: Host.AidaClient.RequestFact = {text: 'f1', metadata: {source: 'devtools'}};
+      const f2 = {text: 'f2', metadata: {source: 'devtools'}};
+      agent.addFact(f1);
+      agent.addFact(f2);
+
+      await Array.fromAsync(agent.run('question', {selected: null}));
+      const request1 = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
+      assert.deepEqual(request1.facts, [f1, f2]);
+
+      agent.removeFact(f1);
+      await Array.fromAsync(agent.run('question', {selected: null}));
+      const request2 = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
+      assert.deepEqual(request2.facts, [f2]);
+
+      agent.clearFacts();
+      await Array.fromAsync(agent.run('question', {selected: null}));
+      const request3 = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
+      assert.isUndefined(request3.facts);
+    });
+
     it('builds a request with chat history', async () => {
       const agent = new AiAgentMock({
         aidaClient: mockAidaClient([[{

front_end/models/ai_assistance/agents/AiAgent.ts

@@ -257,6 +257,8 @@ export abstract class AiAgent<T> {
   #id: string = crypto.randomUUID();
   #history: Host.AidaClient.Content[] = [];
 
+  #facts: Set<Host.AidaClient.RequestFact> = new Set<Host.AidaClient.RequestFact>();
+
   constructor(opts: AgentOptions) {
     this.#aidaClient = opts.aidaClient;
     this.#serverSideLoggingEnabled = opts.serverSideLoggingEnabled ?? false;
@@ -268,6 +270,28 @@ export abstract class AiAgent<T> {
     return query;
   }
 
+  currentFacts(): ReadonlySet<Host.AidaClient.RequestFact> {
+    return this.#facts;
+  }
+
+  /**
+   * Add a fact which will be sent for any subsequent requests.
+   * Returns the new list of all facts.
+   * Facts are never automatically removed.
+   */
+  addFact(fact: Host.AidaClient.RequestFact): ReadonlySet<Host.AidaClient.RequestFact> {
+    this.#facts.add(fact);
+    return this.#facts;
+  }
+
+  removeFact(fact: Host.AidaClient.RequestFact): boolean {
+    return this.#facts.delete(fact);
+  }
+
+  clearFacts(): void {
+    this.#facts.clear();
+  }
+
   buildRequest(
       part: Host.AidaClient.Part|Host.AidaClient.Part[],
       role: Host.AidaClient.Role.USER|Host.AidaClient.Role.ROLE_UNSPECIFIED): Host.AidaClient.AidaRequest {
@@ -291,13 +315,14 @@ export abstract class AiAgent<T> {
     const enableAidaFunctionCalling = declarations.length && !this.functionCallEmulationEnabled;
     const userTier = Host.AidaClient.convertToUserTierEnum(this.userTier);
     const premable = userTier === Host.AidaClient.UserTier.TESTERS ? this.preamble : undefined;
+    const facts = Array.from(this.#facts);
     const request: Host.AidaClient.AidaRequest = {
       client: Host.AidaClient.CLIENT_NAME,
-
       current_message: currentMessage,
       preamble: premable,
 
       historical_contexts: history.length ? history : undefined,
+      facts: facts.length ? facts : undefined,
 
       ...(enableAidaFunctionCalling ? {function_declarations: declarations} : {}),
       options: {

front_end/models/ai_assistance/agents/FileAgent.test.ts

@@ -94,6 +94,7 @@ describeWithMockConnection('FileAgent', () => {
                 parts: [{text: 'answer'}],
               },
             ],
+            facts: undefined,
             metadata: {
               disable_user_content_logging: false,
               string_session_id: 'sessionId',

front_end/models/ai_assistance/agents/PerformanceAgent.test.ts

@@ -109,6 +109,7 @@ describeWithEnvironment('PerformanceAgent', () => {
                 parts: [{text: 'answer'}],
               },
             ],
+            facts: undefined,
             metadata: {
               disable_user_content_logging: false,
               string_session_id: 'sessionId',

front_end/models/ai_assistance/agents/StylingAgent.test.ts

@@ -545,6 +545,7 @@ c`;
                 parts: [{text: 'ANSWER: answer'}],
               },
             ],
+            facts: undefined,
             metadata: {
               disable_user_content_logging: false,
               string_session_id: 'sessionId',