Consolidate function calling

This PR removes duplication in handling side effects and
function calling by introducing a function calling emulation
mode for StylingAgent.

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

Author: OrKoN

front_end/panels/ai_assistance/agents/AiAgent.ts

@@ -187,13 +187,12 @@ export interface FunctionDeclaration<Args extends Record<string, unknown>, Retur
    */
   parameters: Host.AidaClient.FunctionObjectParam<keyof Args>;
   /**
-   * Provided a way to give information back to
-   * the UI before running the the handler
+   * Provided a way to give information back to the UI.
    */
   displayInfoFromArgs?: (
       args: Args,
       ) => {
-    title?: string, thought?: string, code?: string, suggestions?: [string, ...string[]],
+    title?: string, thought?: string, action?: string, suggestions?: [string, ...string[]],
   };
   /**
    * Function implementation that the LLM will try to execute,
@@ -208,8 +207,25 @@ export interface FunctionDeclaration<Args extends Record<string, unknown>, Retur
   }) => Promise<FunctionCallHandlerResult<ReturnType>>;
 }
 
-const OBSERVATION_PREFIX = 'OBSERVATION:';
+const OBSERVATION_PREFIX = 'OBSERVATION: ';
 
+interface AidaFetchResult {
+  text?: string;
+  functionCall?: Host.AidaClient.AidaFunctionCallResponse;
+  completed: boolean;
+  rpcId?: Host.AidaClient.RpcGlobalId;
+}
+
+/**
+ * AiAgent is a base class for implementing an interaction with AIDA
+ * that involves one or more requests being sent to AIDA optionally
+ * utilizing function calling.
+ *
+ * TODO: missing a test that action code is yielded before the
+ * confirmation dialog.
+ * TODO: missing a test for an error if it took
+ * more than MAX_STEPS iterations.
+ */
 export abstract class AiAgent<T> {
   /** Subclasses need to define these. */
   abstract readonly type: AgentType;
@@ -274,6 +290,7 @@ export abstract class AiAgent<T> {
     function validTemperature(temperature: number|undefined): number|undefined {
       return typeof temperature === 'number' && temperature >= 0 ? temperature : undefined;
     }
+    const enableAidaFunctionCalling = declarations.length && !this.functionCallEmulationEnabled;
     const request: Host.AidaClient.AidaRequest = {
       client: Host.AidaClient.CLIENT_NAME,
 
@@ -282,7 +299,7 @@ export abstract class AiAgent<T> {
 
       historical_contexts: history.length ? history : undefined,
 
-      ...(declarations.length ? {function_declarations: declarations} : {}),
+      ...(enableAidaFunctionCalling ? {function_declarations: declarations} : {}),
       options: {
         temperature: validTemperature(this.options.temperature),
         model_id: this.options.modelId,
@@ -294,8 +311,8 @@ export abstract class AiAgent<T> {
         client_version: Root.Runtime.getChromeVersion(),
       },
 
-      functionality_type: declarations.length ? Host.AidaClient.FunctionalityType.AGENTIC_CHAT :
-                                                Host.AidaClient.FunctionalityType.CHAT,
+      functionality_type: enableAidaFunctionCalling ? Host.AidaClient.FunctionalityType.AGENTIC_CHAT :
+                                                      Host.AidaClient.FunctionalityType.CHAT,
 
       client_feature: this.clientFeature,
     };
@@ -314,13 +331,13 @@ export abstract class AiAgent<T> {
     return this.#origin;
   }
 
-  parseResponse(response: Host.AidaClient.AidaResponse): ParsedResponse {
-    if (response.functionCalls && response.completed) {
-      throw new Error('Function calling not supported yet');
-    }
-    return {
-      answer: response.explanation,
-    };
+  /**
+   * Parses a streaming text response into a
+   * though/action/title/answer/suggestions component. This is only used
+   * by StylingAgent.
+   */
+  parseTextResponse(response: string): ParsedResponse {
+    return {answer: response};
   }
 
   /**
@@ -346,10 +363,14 @@ export abstract class AiAgent<T> {
     return answer;
   }
 
-  protected handleAction(action: string, options?: {signal?: AbortSignal}):
-      AsyncGenerator<SideEffectResponse, ActionResponse, void>;
-  protected handleAction(): never {
-    throw new Error('Unexpected action found');
+  /**
+   * Special mode for StylingAgent that turns custom text output into a
+   * function call.
+   */
+  protected functionCallEmulationEnabled = false;
+  protected emulateFunctionCall(_aidaResponse: Host.AidaClient.AidaResponse): Host.AidaClient.AidaFunctionCallResponse|
+      'no-function-call'|'wait-for-completion' {
+    throw new Error('Unexpected emulateFunctionCall. Only StylingAgent implements function call emulation');
   }
 
   async *
@@ -391,23 +412,30 @@ export abstract class AiAgent<T> {
       };
 
       let rpcId: Host.AidaClient.RpcGlobalId|undefined;
-      let parsedResponse: ParsedResponse|undefined = undefined;
+      let textResponse = '';
       let functionCall: Host.AidaClient.AidaFunctionCallResponse|undefined = undefined;
       try {
         for await (const fetchResult of this.#aidaFetch(request, {signal: options.signal})) {
           rpcId = fetchResult.rpcId;
-          parsedResponse = fetchResult.parsedResponse;
+          textResponse = fetchResult.text ?? '';
           functionCall = fetchResult.functionCall;
 
-          // Only yield partial responses here and do not add partial answers to the history.
-          if (!fetchResult.completed && !fetchResult.functionCall && 'answer' in parsedResponse &&
-              parsedResponse.answer) {
+          if (!functionCall && !fetchResult.completed) {
+            const parsed = this.parseTextResponse(textResponse);
+            const partialAnswer = 'answer' in parsed ? parsed.answer : '';
+            if (!partialAnswer) {
+              continue;
+            }
+            // Only yield partial responses here and do not add partial answers to the history.
             yield {
               type: ResponseType.ANSWER,
-              text: parsedResponse.answer,
+              text: partialAnswer,
               complete: false,
             };
           }
+          if (functionCall) {
+            break;
+          }
         }
       } catch (err) {
         debugLog('Error calling the AIDA API', err);
@@ -425,7 +453,11 @@ export abstract class AiAgent<T> {
 
       this.#history.push(request.current_message);
 
-      if (parsedResponse && 'answer' in parsedResponse && Boolean(parsedResponse.answer)) {
+      if (textResponse) {
+        const parsedResponse = this.parseTextResponse(textResponse);
+        if (!('answer' in parsedResponse)) {
+          throw new Error('Expected a completed response to have an answer');
+        }
         this.#history.push({
           parts: [{
             text: this.formatParsedAnswer(parsedResponse),
@@ -441,66 +473,24 @@ export abstract class AiAgent<T> {
           rpcId,
         };
         break;
-      } else if (parsedResponse && !('answer' in parsedResponse)) {
-        const {
-          title,
-          thought,
-          action,
-        } = parsedResponse;
-
-        if (title) {
-          yield {
-            type: ResponseType.TITLE,
-            title,
-            rpcId,
-          };
-        }
-
-        if (thought) {
-          yield {
-            type: ResponseType.THOUGHT,
-            thought,
-            rpcId,
-          };
-        }
-
-        this.#history.push({
-          parts: [{
-            text: this.#formatParsedStep(parsedResponse),
-          }],
-          role: Host.AidaClient.Role.MODEL,
-        });
+      }
 
-        if (action) {
-          const result = yield* this.handleAction(action, {signal: options.signal});
-          if (options?.signal?.aborted) {
+      if (functionCall) {
+        try {
+          const result = yield* this.#callFunction(functionCall.name, functionCall.args, options);
+          if (options.signal?.aborted) {
             yield this.#createErrorResponse(ErrorType.ABORT);
             break;
           }
-          query = {text: `${OBSERVATION_PREFIX} ${result.output}`};
-          // Capture history state for the next iteration query.
-          request = this.buildRequest(query, Host.AidaClient.Role.USER);
-          yield result;
-        }
-      } else if (functionCall) {
-        try {
-          const result = yield* this.#callFunction(functionCall.name, functionCall.args);
-
-          if (result.result) {
-            yield {
-              type: ResponseType.ACTION,
-              output: JSON.stringify(result.result),
-              canceled: false,
-            };
-          }
-
-          query = {
+          query = this.functionCallEmulationEnabled ? {text: OBSERVATION_PREFIX + result.result} : {
             functionResponse: {
               name: functionCall.name,
               response: result,
             },
           };
-          request = this.buildRequest(query, Host.AidaClient.Role.ROLE_UNSPECIFIED);
+          request = this.buildRequest(
+              query,
+              this.functionCallEmulationEnabled ? Host.AidaClient.Role.USER : Host.AidaClient.Role.ROLE_UNSPECIFIED);
         } catch {
           yield this.#createErrorResponse(ErrorType.UNKNOWN);
           break;
@@ -524,18 +514,31 @@ export abstract class AiAgent<T> {
     if (!call) {
       throw new Error(`Function ${name} is not found.`);
     }
-    this.#history.push({
-      parts: [{
-        functionCall: {
-          name,
-          args,
-        },
-      }],
-      role: Host.AidaClient.Role.MODEL,
-    });
+    if (this.functionCallEmulationEnabled) {
+      if (!call.displayInfoFromArgs) {
+        throw new Error('functionCallEmulationEnabled requires all functions to provide displayInfoFromArgs');
+      }
+      // Emulated function calls are formatted as text.
+      this.#history.push({
+        parts: [{text: this.#formatParsedStep(call.displayInfoFromArgs(args))}],
+        role: Host.AidaClient.Role.MODEL,
+      });
+    } else {
+      this.#history.push({
+        parts: [{
+          functionCall: {
+            name,
+            args,
+          },
+        }],
+        role: Host.AidaClient.Role.MODEL,
+      });
+    }
 
+    let code;
     if (call.displayInfoFromArgs) {
-      const {title, thought, code, suggestions} = call.displayInfoFromArgs(args);
+      const {title, thought, action: callCode} = call.displayInfoFromArgs(args);
+      code = callCode;
       if (title) {
         yield {
           type: ResponseType.TITLE,
@@ -549,7 +552,11 @@ export abstract class AiAgent<T> {
           thought,
         };
       }
+    }
 
+    let result = await call.handler(args, options) as FunctionCallHandlerResult<unknown>;
+
+    if ('requiresApproval' in result) {
       if (code) {
         yield {
           type: ResponseType.ACTION,
@@ -558,17 +565,6 @@ export abstract class AiAgent<T> {
         };
       }
 
-      if (suggestions) {
-        yield {
-          type: ResponseType.SUGGESTIONS,
-          suggestions,
-        };
-      }
-    }
-
-    let result = await call.handler(args, options);
-
-    if ('requiresApproval' in result) {
       const sideEffectConfirmationPromiseWithResolvers = this.confirmSideEffect<boolean>();
 
       void sideEffectConfirmationPromiseWithResolvers.promise.then(result => {
@@ -597,7 +593,8 @@ export abstract class AiAgent<T> {
       if (!approvedRun) {
         yield {
           type: ResponseType.ACTION,
-          code: '',
+          code,
+          output: 'Error: User denied code execution with side effects.',
           canceled: true,
         };
         return {
@@ -611,17 +608,30 @@ export abstract class AiAgent<T> {
       });
     }
 
+    if ('result' in result) {
+      yield {
+        type: ResponseType.ACTION,
+        code,
+        output: typeof result.result === 'string' ? result.result : JSON.stringify(result.result),
+        canceled: false,
+      };
+    }
+
+    if ('error' in result) {
+      yield {
+        type: ResponseType.ACTION,
+        code,
+        output: result.error,
+        canceled: false,
+      };
+    }
+
     return result as {result: unknown};
   }
 
   async *
-      #aidaFetch(request: Host.AidaClient.AidaRequest, options?: {signal?: AbortSignal}): AsyncGenerator<
-          {
-            parsedResponse: ParsedResponse,
-            functionCall?: Host.AidaClient.AidaFunctionCallResponse, completed: boolean,
-            rpcId?: Host.AidaClient.RpcGlobalId,
-          },
-          void, void> {
+      #aidaFetch(request: Host.AidaClient.AidaRequest, options?: {signal?: AbortSignal}):
+          AsyncGenerator<AidaFetchResult, void, void> {
     let aidaResponse: Host.AidaClient.AidaResponse|undefined = undefined;
     let response = '';
     let rpcId: Host.AidaClient.RpcGlobalId|undefined;
@@ -631,19 +641,32 @@ export abstract class AiAgent<T> {
         debugLog('functionCalls.length', aidaResponse.functionCalls.length);
         yield {
           rpcId,
-          parsedResponse: {answer: ''},
           functionCall: aidaResponse.functionCalls[0],
           completed: true,
         };
         break;
       }
 
+      if (this.functionCallEmulationEnabled) {
+        const emulatedFunctionCall = this.emulateFunctionCall(aidaResponse);
+        if (emulatedFunctionCall === 'wait-for-completion') {
+          continue;
+        }
+        if (emulatedFunctionCall !== 'no-function-call') {
+          yield {
+            rpcId,
+            functionCall: emulatedFunctionCall,
+            completed: true,
+          };
+          break;
+        }
+      }
+
       response = aidaResponse.explanation;
       rpcId = aidaResponse.metadata.rpcGlobalId ?? rpcId;
-      const parsedResponse = this.parseResponse(aidaResponse);
       yield {
         rpcId,
-        parsedResponse,
+        text: aidaResponse.explanation,
         completed: aidaResponse.completed,
       };
     }