[AI Assistance] Simplify agent.run

Moves the history handling logic outside the run itself.
As we are storing all the yields.
Refactor Error yields creation to insure we collect metrics correctly.



Bug: 360751542
Change-Id: I4aa3afd4430725d8cc14c876b1196f5b4a708244
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6235030
Reviewed-by: Alex Rudenko <[email protected]>
Commit-Queue: Nikolay Vitkov <[email protected]>

Author: Lightning00Blade

front_end/panels/ai_assistance/agents/AiAgent.ts

@@ -333,9 +333,66 @@ export abstract class AiAgent<T> {
   get isHistoryEntry(): boolean {
     return this.#generatedFromHistory;
   }
-
   async * run(initialQuery: string, options: {
     signal?: AbortSignal, selected: ConversationContext<T>|null,
+  }): AsyncGenerator<ResponseData, void, void> {
+    for await (const response of this.#run(initialQuery, options)) {
+      this.#addHistory(response);
+      yield response;
+    }
+  }
+
+  async * runFromHistory(): AsyncGenerator<ResponseData, void, void> {
+    if (this.isEmpty) {
+      return;
+    }
+
+    this.#generatedFromHistory = true;
+    for (const entry of this.#history) {
+      yield entry;
+    }
+  }
+
+  parseResponse(response: Host.AidaClient.AidaResponse): ParsedResponse {
+    if (response.functionCalls && response.completed) {
+      throw new Error('Function calling not supported yet');
+    }
+    return {
+      answer: response.explanation,
+    };
+  }
+
+  /**
+   * Declare a function that the AI model can call.
+   * @param name - The name of the function
+   * @param declaration - the function declaration. Currently functions must:
+   * 1. Return an object of serializable key/value pairs. You cannot return
+   *    anything other than a plain JavaScript object that can be serialized.
+   * 2. Take one parameter which is an object that can have
+   *    multiple keys and values. For example, rather than a function being called
+   *    with two args, `foo` and `bar`, you should instead have the function be
+   *    called with one object with `foo` and `bar` keys.
+   */
+  protected declareFunction<Args extends Record<string, unknown>, ReturnType = unknown>(
+      name: string, declaration: FunctionDeclaration<Args, ReturnType>): void {
+    if (this.#functionDeclarations.has(name)) {
+      throw new Error(`Duplicate function declaration ${name}`);
+    }
+    this.#functionDeclarations.set(name, declaration as FunctionDeclaration<Record<string, unknown>, ReturnType>);
+  }
+
+  protected formatParsedAnswer({answer}: ParsedAnswer): string {
+    return answer;
+  }
+
+  protected handleAction(action: string, options?: {signal?: AbortSignal}):
+      AsyncGenerator<SideEffectResponse, ActionResponse, void>;
+  protected handleAction(): never {
+    throw new Error('Unexpected action found');
+  }
+
+  async * #run(initialQuery: string, options: {
+    signal?: AbortSignal, selected: ConversationContext<T>|null,
   }): AsyncGenerator<ResponseData, void, void> {
     if (this.#generatedFromHistory) {
       throw new Error('History entries are read-only.');
@@ -360,24 +417,17 @@ export abstract class AiAgent<T> {
     // Request is built here to capture history up to this point.
     let request = this.buildRequest(query, Host.AidaClient.Role.USER);
 
-    const response = {
+    yield {
       type: ResponseType.USER_QUERY,
       query: initialQuery,
-    } as const;
-    this.#addHistory(response);
-    yield response;
+    };
 
-    for await (const response of this.handleContextDetails(options.selected)) {
-      this.#addHistory(response);
-      yield response;
-    }
+    yield* this.handleContextDetails(options.selected);
 
     for (let i = 0; i < MAX_STEPS; i++) {
-      const queryResponse = {
+      yield {
         type: ResponseType.QUERYING,
-      } as const;
-      this.#addHistory(queryResponse);
-      yield queryResponse;
+      };
 
       let rpcId: Host.AidaClient.RpcGlobalId|undefined;
       let parsedResponse: ParsedResponse|undefined = undefined;
@@ -400,44 +450,33 @@ export abstract class AiAgent<T> {
       } catch (err) {
         debugLog('Error calling the AIDA API', err);
 
-        this.#removeLastRunParts();
+        let error = ErrorType.UNKNOWN;
         if (err instanceof Host.AidaClient.AidaAbortError) {
-          const response = this.#createAbortResponse();
-          this.#addHistory(response);
-          yield response;
-          break;
+          error = ErrorType.ABORT;
+        } else if (err instanceof Host.AidaClient.AidaBlockError) {
+          error = ErrorType.BLOCK;
         }
-
-        const error = (err instanceof Host.AidaClient.AidaBlockError) ? ErrorType.BLOCK : ErrorType.UNKNOWN;
-        const response = {
-          type: ResponseType.ERROR,
-          error,
-        } as const;
-        this.#addHistory(response);
-        Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
-        yield response;
+        yield this.#createErrorResponse(error);
 
         break;
       }
 
       this.#partsHistory.push(request.current_message);
 
       if (parsedResponse && 'answer' in parsedResponse && Boolean(parsedResponse.answer)) {
-        const response = {
-          type: ResponseType.ANSWER,
-          text: parsedResponse.answer,
-          suggestions: parsedResponse.suggestions,
-          rpcId,
-        } as const;
         this.#partsHistory.push({
           parts: [{
             text: this.formatParsedAnswer(parsedResponse),
           }],
           role: Host.AidaClient.Role.MODEL,
         });
         Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceAnswerReceived);
-        this.#addHistory(response);
-        yield response;
+        yield {
+          type: ResponseType.ANSWER,
+          text: parsedResponse.answer,
+          suggestions: parsedResponse.suggestions,
+          rpcId,
+        };
         break;
       } else if (parsedResponse && !('answer' in parsedResponse)) {
         const {
@@ -447,23 +486,19 @@ export abstract class AiAgent<T> {
         } = parsedResponse;
 
         if (title) {
-          const response = {
+          yield {
             type: ResponseType.TITLE,
             title,
             rpcId,
-          } as const;
-          this.#addHistory(response);
-          yield response;
+          };
         }
 
         if (thought) {
-          const response = {
+          yield {
             type: ResponseType.THOUGHT,
             thought,
             rpcId,
-          } as const;
-          this.#addHistory(response);
-          yield response;
+          };
         }
 
         this.#partsHistory.push({
@@ -476,14 +511,9 @@ export abstract class AiAgent<T> {
         if (action) {
           const result = yield* this.handleAction(action, {signal: options.signal});
           if (options?.signal?.aborted) {
-            this.#removeLastRunParts();
-            const response = this.#createAbortResponse();
-            this.#addHistory(response);
-
-            yield response;
+            yield this.#createErrorResponse(ErrorType.ABORT);
             break;
           }
-          this.#addHistory(result);
           query = {text: `${OBSERVATION_PREFIX} ${result.output}`};
           // Capture history state for the next iteration query.
           request = this.buildRequest(query, Host.AidaClient.Role.USER);
@@ -494,12 +524,11 @@ export abstract class AiAgent<T> {
           const result = yield* this.#callFunction(functionCall.name, functionCall.args);
 
           if (result.result) {
-            const response = {
+            yield {
               type: ResponseType.ACTION,
               output: JSON.stringify(result.result),
               canceled: false,
-            } as const;
-            yield response;
+            };
           }
 
           query = {
@@ -510,23 +539,11 @@ export abstract class AiAgent<T> {
           };
           request = this.buildRequest(query, Host.AidaClient.Role.ROLE_UNSPECIFIED);
         } catch {
-          this.#removeLastRunParts();
-          const response = {
-            type: ResponseType.ERROR,
-            error: ErrorType.UNKNOWN,
-          } as const;
-          Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
-          yield response;
+          yield this.#createErrorResponse(ErrorType.UNKNOWN);
           break;
         }
       } else {
-        this.#removeLastRunParts();
-        const response = {
-          type: ResponseType.ERROR,
-          error: i - 1 === MAX_STEPS ? ErrorType.MAX_STEPS : ErrorType.UNKNOWN,
-        } as const;
-        Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
-        yield response;
+        yield this.#createErrorResponse(i - 1 === MAX_STEPS ? ErrorType.MAX_STEPS : ErrorType.UNKNOWN);
         break;
       }
     }
@@ -536,55 +553,6 @@ export abstract class AiAgent<T> {
     }
   }
 
-  async * runFromHistory(): AsyncGenerator<ResponseData, void, void> {
-    if (this.isEmpty) {
-      return;
-    }
-
-    this.#generatedFromHistory = true;
-    for (const entry of this.#history) {
-      yield entry;
-    }
-  }
-
-  parseResponse(response: Host.AidaClient.AidaResponse): ParsedResponse {
-    if (response.functionCalls && response.completed) {
-      throw new Error('Function calling not supported yet');
-    }
-    return {
-      answer: response.explanation,
-    };
-  }
-
-  /**
-   * Declare a function that the AI model can call.
-   * @param name - The name of the function
-   * @param declaration - the function declaration. Currently functions must:
-   * 1. Return an object of serializable key/value pairs. You cannot return
-   *    anything other than a plain JavaScript object that can be serialized.
-   * 2. Take one parameter which is an object that can have
-   *    multiple keys and values. For example, rather than a function being called
-   *    with two args, `foo` and `bar`, you should instead have the function be
-   *    called with one object with `foo` and `bar` keys.
-   */
-  protected declareFunction<Args extends Record<string, unknown>, ReturnType = unknown>(
-      name: string, declaration: FunctionDeclaration<Args, ReturnType>): void {
-    if (this.#functionDeclarations.has(name)) {
-      throw new Error(`Duplicate function declaration ${name}`);
-    }
-    this.#functionDeclarations.set(name, declaration as FunctionDeclaration<Record<string, unknown>, ReturnType>);
-  }
-
-  protected formatParsedAnswer({answer}: ParsedAnswer): string {
-    return answer;
-  }
-
-  protected handleAction(action: string, options?: {signal?: AbortSignal}):
-      AsyncGenerator<SideEffectResponse, ActionResponse, void>;
-  protected handleAction(): never {
-    throw new Error('Unexpected action found');
-  }
-
   async * #callFunction(name: string, args: Record<string, unknown>, options?: {
     signal?: AbortSignal,
     approved?: boolean,
@@ -606,39 +574,32 @@ export abstract class AiAgent<T> {
     if (call.displayInfoFromArgs) {
       const {title, thought, code, suggestions} = call.displayInfoFromArgs(args);
       if (title) {
-        const response = {
+        yield {
           type: ResponseType.TITLE,
           title,
-        } as const;
-        this.#addHistory(response);
-        yield response;
+        };
       }
 
       if (thought) {
-        const response = {
+        yield {
           type: ResponseType.THOUGHT,
           thought,
-        } as const;
-        this.#addHistory(response);
-        yield response;
+        };
       }
 
       if (code) {
-        const response = {
+        yield {
           type: ResponseType.ACTION,
           code,
           canceled: false,
-        } as const;
-        this.#addHistory(response);
-        yield response;
+        };
       }
 
       if (suggestions) {
-        const response = {
+        yield {
           type: ResponseType.SUGGESTIONS,
           suggestions,
-        } as const;
-        yield response;
+        };
       }
     }
 
@@ -671,12 +632,11 @@ export abstract class AiAgent<T> {
 
       const approvedRun = await sideEffectConfirmationPromiseWithResolvers.promise;
       if (!approvedRun) {
-        const response = {
+        yield {
           type: ResponseType.ACTION,
           code: '',
           canceled: true,
-        } as const;
-        yield response;
+        };
         return {
           result: 'Error: User denied code execution with side effects.',
         };
@@ -775,10 +735,15 @@ STOP`;
     }));
   }
 
-  #createAbortResponse(): ResponseData {
+  #createErrorResponse(error: ErrorType): ResponseData {
+    this.#removeLastRunParts();
+    if (error !== ErrorType.ABORT) {
+      Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
+    }
+
     return {
       type: ResponseType.ERROR,
-      error: ErrorType.ABORT,
+      error,
     };
   }
 }