[Freestyler] Extract run code

Unify how we do run in AI Assistance to streamline creation
of new Agents

Bug: 369303799
Change-Id: I4aca9f997caf22062e0e812b94a62f955688fdee
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/5904625
Commit-Queue: Nikolay Vitkov <[email protected]>
Reviewed-by: Alex Rudenko <[email protected]>
Reviewed-by: Ergün Erdoğmuş <[email protected]>

Author: Lightning00Blade

front_end/panels/freestyler/AiAgent.test.ts

@@ -11,7 +11,7 @@ import * as Freestyler from './freestyler.js';
 
 const {AiAgent} = Freestyler;
 
-class AiAgentMock extends AiAgent {
+class AiAgentMock extends AiAgent<unknown> {
   override preamble = 'preamble';
 
   clientFeature: Host.AidaClient.ClientFeature = 0;

front_end/panels/freestyler/AiAgent.ts

@@ -25,7 +25,7 @@ export interface AnswerResponse {
   type: ResponseType.ANSWER;
   text: string;
   rpcId?: number;
-  suggestions?: string[];
+  suggestions?: [string, ...string[]];
 }
 
 export interface ErrorResponse {
@@ -99,14 +99,34 @@ type AgentOptions = {
   serverSideLoggingEnabled?: boolean,
 };
 
-export abstract class AiAgent {
+interface ParsedResponseAnswer {
+  answer: string;
+  suggestions?: [string, ...string[]];
+}
+
+interface ParsedResponseStep {
+  thought?: string;
+  title?: string;
+  action?: string;
+}
+
+export type ParsedResponse = ParsedResponseAnswer|ParsedResponseStep;
+
+const MAX_STEP = 10;
+
+export abstract class AiAgent<T> {
+  static validTemperature(temperature: number|undefined): number|undefined {
+    return typeof temperature === 'number' && temperature >= 0 ? temperature : undefined;
+  }
+
   readonly #sessionId: string = crypto.randomUUID();
   #aidaClient: Host.AidaClient.AidaClient;
   #serverSideLoggingEnabled: boolean;
   abstract readonly preamble: string;
   abstract readonly options: AidaRequestOptions;
   abstract readonly clientFeature: Host.AidaClient.ClientFeature;
   abstract readonly userTier: string|undefined;
+  abstract handleContextDetails(select: T|null): AsyncGenerator<ContextResponse, void, void>;
 
   /**
    * Mapping between the unique request id and
@@ -135,7 +155,49 @@ export abstract class AiAgent {
     this.#chatHistory.delete(id);
   }
 
-  addToHistory({
+  addToHistory(options: {
+    id: number,
+    query: string,
+    response: ParsedResponse,
+  }): void {
+    const response = options.response;
+    if ('answer' in response) {
+      this.#storeHistoryEntries({
+        id: options.id,
+        query: options.query,
+        output: response.answer,
+      });
+      return;
+    }
+
+    const {
+      title,
+      thought,
+      action,
+    } = response;
+
+    if (thought) {
+      this.#storeHistoryEntries({
+        id: options.id,
+        query: options.query,
+        output: `THOUGHT: ${thought}
+TITLE: ${title}
+ACTION
+${action}
+STOP`,
+      });
+    } else {
+      this.#storeHistoryEntries({
+        id: options.id,
+        query: options.query,
+        output: `ACTION
+${action}
+STOP`,
+      });
+    }
+  }
+
+  #storeHistoryEntries({
     id,
     query,
     output,
@@ -169,15 +231,15 @@ export abstract class AiAgent {
       options?: {signal?: AbortSignal},
       ): Promise<{
     response: string,
-    rpcId: number|undefined,
+    rpcId?: number,
   }> {
     const request = this.buildRequest({
       input,
     });
 
     let rawResponse: Host.AidaClient.AidaResponse|undefined = undefined;
     let response = '';
-    let rpcId;
+    let rpcId: number|undefined;
     for await (rawResponse of this.#aidaClient.fetch(request, options)) {
       response = rawResponse.explanation;
       rpcId = rawResponse.metadata.rpcGlobalId ?? rpcId;
@@ -223,8 +285,135 @@ export abstract class AiAgent {
     return request;
   }
 
-  static validTemperature(temperature: number|undefined): number|undefined {
-    return typeof temperature === 'number' && temperature >= 0 ? temperature : undefined;
+  handleAction(_action: string, _rpcId?: number): AsyncGenerator<SideEffectResponse, ActionResponse, void> {
+    throw new Error('Unexpected action found');
+  }
+
+  async enhanceQuery(query: string, selected: T|null): Promise<string>;
+  async enhanceQuery(query: string): Promise<string> {
+    return query;
+  }
+
+  parseResponse(response: string): ParsedResponse {
+    return {
+      answer: response,
+    };
+  }
+
+  #runId = 0;
+  async * run(query: string, options: {
+    signal?: AbortSignal, selected: T|null,
+  }): AsyncGenerator<ResponseData, void, void> {
+    yield* this.handleContextDetails(options.selected);
+
+    query = await this.enhanceQuery(query, options.selected);
+    const currentRunId = ++this.#runId;
+
+    for (let i = 0; i < MAX_STEP; i++) {
+      yield {
+        type: ResponseType.QUERYING,
+      };
+
+      let response: string;
+      let rpcId: number|undefined;
+      try {
+        const fetchResult = await this.aidaFetch(
+            query,
+            {signal: options.signal},
+        );
+        response = fetchResult.response;
+        rpcId = fetchResult.rpcId;
+      } catch (err) {
+        debugLog('Error calling the AIDA API', err);
+
+        if (err instanceof Host.AidaClient.AidaAbortError) {
+          this.removeHistoryRun(currentRunId);
+          yield {
+            type: ResponseType.ERROR,
+            error: ErrorType.ABORT,
+            rpcId,
+          };
+          break;
+        }
+
+        yield {
+          type: ResponseType.ERROR,
+          error: ErrorType.UNKNOWN,
+          rpcId,
+        };
+        break;
+      }
+
+      const parsedResponse = this.parseResponse(response);
+
+      this.addToHistory({
+        id: currentRunId,
+        query,
+        response: parsedResponse,
+      });
+      if ('answer' in parsedResponse) {
+        const {
+          answer,
+          suggestions,
+        } = parsedResponse;
+        if (answer) {
+          yield {
+            type: ResponseType.ANSWER,
+            text: answer,
+            rpcId,
+            suggestions,
+          };
+        } else {
+          this.removeHistoryRun(currentRunId);
+          yield {
+            type: ResponseType.ERROR,
+            error: ErrorType.UNKNOWN,
+            rpcId,
+          };
+        }
+
+        break;
+      }
+
+      const {
+        title,
+        thought,
+        action,
+      } = parsedResponse;
+
+      if (title) {
+        yield {
+          type: ResponseType.TITLE,
+          title,
+          rpcId,
+        };
+      }
+
+      if (thought) {
+        yield {
+          type: ResponseType.THOUGHT,
+          thought,
+          rpcId,
+        };
+      }
+
+      if (action) {
+        const result = yield* this.handleAction(action, rpcId);
+        yield result;
+        query = `OBSERVATION: ${result.output}`;
+      }
+
+      if (i === MAX_STEP - 1) {
+        yield {
+          type: ResponseType.ERROR,
+          error: ErrorType.MAX_STEPS,
+        };
+        break;
+      }
+    }
+    if (isDebugMode()) {
+      window.dispatchEvent(new CustomEvent('freestylerdone'));
+    }
   }
 }
 

front_end/panels/freestyler/DrJonesFileAgent.test.ts

@@ -157,7 +157,7 @@ describeWithEnvironment('DrJonesFileAgent', () => {
       });
 
       const uiSourceCode = project.uiSourceCodeForURL(url);
-      const responses = await Array.fromAsync(agent.run('test', {selectedFile: uiSourceCode}));
+      const responses = await Array.fromAsync(agent.run('test', {selected: uiSourceCode}));
 
       assert.deepStrictEqual(responses, [
         {
@@ -173,9 +173,13 @@ File Content:
             },
           ],
         },
+        {
+          type: ResponseType.QUERYING,
+        },
         {
           type: ResponseType.ANSWER,
           text: 'This is the answer',
+          suggestions: undefined,
           rpcId: 123,
         },
       ]);

front_end/panels/freestyler/DrJonesFileAgent.ts

@@ -12,10 +12,7 @@ import {
   type AidaRequestOptions,
   type ContextDetail,
   type ContextResponse,
-  debugLog,
-  ErrorType,
-  isDebugMode,
-  type ResponseData,
+  type ParsedResponse,
   ResponseType,
 } from './AiAgent.js';
 
@@ -70,7 +67,7 @@ const MAX_FILE_SIZE = 50000;
  * One agent instance handles one conversation. Create a new agent
  * instance for a new conversation.
  */
-export class DrJonesFileAgent extends AiAgent {
+export class DrJonesFileAgent extends AiAgent<Workspace.UISourceCode.UISourceCode> {
   readonly preamble = preamble;
   readonly clientFeature = Host.AidaClient.ClientFeature.CHROME_DRJONES_FILE_AGENT;
   get userTier(): string|undefined {
@@ -88,9 +85,9 @@ export class DrJonesFileAgent extends AiAgent {
     };
   }
 
-  *
+  async *
       handleContextDetails(selectedFile: Workspace.UISourceCode.UISourceCode|null):
-          Generator<ContextResponse, void, void> {
+          AsyncGenerator<ContextResponse, void, void> {
     if (!selectedFile) {
       return;
     }
@@ -102,61 +99,16 @@ export class DrJonesFileAgent extends AiAgent {
     };
   }
 
-  async enhanceQuery(query: string, selectedFile: Workspace.UISourceCode.UISourceCode|null): Promise<string> {
+  override async enhanceQuery(query: string, selectedFile: Workspace.UISourceCode.UISourceCode|null): Promise<string> {
     const fileEnchantmentQuery =
         selectedFile ? `# Selected file\n${formatFile(selectedFile)}\n\n# User request\n\n` : '';
     return `${fileEnchantmentQuery}${query}`;
   }
 
-  #runId = 0;
-  async * run(query: string, options: {
-    signal?: AbortSignal, selectedFile: Workspace.UISourceCode.UISourceCode|null,
-  }): AsyncGenerator<ResponseData, void, void> {
-    yield* this.handleContextDetails(options.selectedFile);
-
-    query = await this.enhanceQuery(query, options.selectedFile);
-    const currentRunId = ++this.#runId;
-
-    let response: string;
-    let rpcId: number|undefined;
-    try {
-      const fetchResult = await this.aidaFetch(query, {signal: options.signal});
-      response = fetchResult.response;
-      rpcId = fetchResult.rpcId;
-    } catch (err) {
-      debugLog('Error calling the AIDA API', err);
-      if (err instanceof Host.AidaClient.AidaAbortError) {
-        this.removeHistoryRun(currentRunId);
-        yield {
-          type: ResponseType.ERROR,
-          error: ErrorType.ABORT,
-          rpcId,
-        };
-        return;
-      }
-
-      yield {
-        type: ResponseType.ERROR,
-        error: ErrorType.UNKNOWN,
-        rpcId,
-      };
-      return;
-    }
-
-    this.addToHistory({
-      id: currentRunId,
-      query,
-      output: response,
-    });
-
-    yield {
-      type: ResponseType.ANSWER,
-      text: response,
-      rpcId,
+  override parseResponse(response: string): ParsedResponse {
+    return {
+      answer: response,
     };
-    if (isDebugMode()) {
-      window.dispatchEvent(new CustomEvent('freestylerdone'));
-    }
   }
 }