[Ai Assistance] Use function calling for StylingAgent

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

Author: Lightning00Blade

front_end/core/host/AidaClient.ts

@@ -10,24 +10,27 @@ import type {AidaClientResult, SyncInformation} from './InspectorFrontendHostAPI
 import {bindOutputStream} from './ResourceLoader.js';
 
 export enum Role {
-  // Unspecified role.
+  /** Provide this role when giving a function call response  */
   ROLE_UNSPECIFIED = 0,
-  // The user.
+  /** Tags the content came from the user */
   USER = 1,
-  // The model.
+  /** Tags the content came from the LLM */
   MODEL = 2,
 }
 
 export const enum Rating {
+  // Resets the vote to null in the logs
   SENTIMENT_UNSPECIFIED = 'SENTIMENT_UNSPECIFIED',
   POSITIVE = 'POSITIVE',
   NEGATIVE = 'NEGATIVE',
 }
 
-// A `Content` represents a single turn message.
+/**
+ * A `Content` represents a single turn message.
+ */
 export interface Content {
   parts: Part[];
-  // The producer of the content.
+  /** The producer of the content. */
   role: Role;
 }
 
@@ -44,7 +47,7 @@ export type Part = {
     response: Record<string, unknown>,
   },
 }|{
-  // Inline media bytes.
+  /** Inline media bytes. */
   inlineData: MediaBlob,
 };
 
@@ -71,23 +74,23 @@ interface FunctionArrayParam extends BaseFunctionParam {
   items: FunctionPrimitiveParams;
 }
 
-export interface FunctionObjectParam extends BaseFunctionParam {
+export interface FunctionObjectParam<T extends string|number|symbol = string> extends BaseFunctionParam {
   type: ParametersTypes.OBJECT;
   // TODO: this can be also be ObjectParams
-  properties: {[Key in string]: FunctionPrimitiveParams|FunctionArrayParam};
+  properties: Record<T, FunctionPrimitiveParams|FunctionArrayParam>;
 }
 
 /**
  * More about function declaration can be read at
  * https://ai.google.dev/gemini-api/docs/function-calling
  */
-export interface FunctionDeclaration {
+export interface FunctionDeclaration<T extends string|number|symbol = string> {
   name: string;
   /**
    * A description for the LLM to understand what the specific function will do once called.
    */
   description: string;
-  parameters: FunctionObjectParam|FunctionPrimitiveParams|FunctionArrayParam;
+  parameters: FunctionObjectParam<T>;
 }
 
 // Raw media bytes.

front_end/panels/ai_assistance/AiAssistancePanel.ts

@@ -40,7 +40,7 @@ import {
 } from './agents/NetworkAgent.js';
 import {PatchAgent, ProjectContext} from './agents/PatchAgent.js';
 import {CallTreeContext, PerformanceAgent} from './agents/PerformanceAgent.js';
-import {NodeContext, StylingAgent} from './agents/StylingAgent.js';
+import {NodeContext, StylingAgent, StylingAgentWithFunctionCalling} from './agents/StylingAgent.js';
 import aiAssistancePanelStyles from './aiAssistancePanel.css.js';
 import {
   AiHistoryStorage,
@@ -334,6 +334,13 @@ export class AiAssistancePanel extends UI.Panel.Panel {
           ...options,
           changeManager: this.#changeManager,
         });
+        if (isAiAssistanceStylingWithFunctionCallingEnabled()) {
+          agent = new StylingAgentWithFunctionCalling({
+            ...options,
+            changeManager: this.#changeManager,
+          });
+        }
+
         break;
       }
       case AgentType.NETWORK: {
@@ -939,9 +946,13 @@ export class AiAssistancePanel extends UI.Panel.Panel {
             commitStep();
             break;
           }
+          case ResponseType.SUGGESTIONS: {
+            systemMessage.suggestions = data.suggestions;
+            break;
+          }
           case ResponseType.SIDE_EFFECT: {
             step.isLoading = false;
-            step.code = data.code;
+            step.code ??= data.code;
             step.sideEffect = {
               onAnswer: data.confirm,
             };
@@ -950,8 +961,8 @@ export class AiAssistancePanel extends UI.Panel.Panel {
           }
           case ResponseType.ACTION: {
             step.isLoading = false;
-            step.code = data.code;
-            step.output = data.output;
+            step.code ??= data.code;
+            step.output ??= data.output;
             step.canceled = data.canceled;
             if (isAiAssistanceChangeSummariesEnabled() && this.#currentAgent && !this.#currentAgent.isHistoryEntry) {
               this.#viewProps.changeSummary = this.#changeManager.formatChanges(this.#currentAgent.id);
@@ -960,7 +971,7 @@ export class AiAssistancePanel extends UI.Panel.Panel {
             break;
           }
           case ResponseType.ANSWER: {
-            systemMessage.suggestions = data.suggestions;
+            systemMessage.suggestions ??= data.suggestions;
             systemMessage.answer = data.text;
             systemMessage.rpcId = data.rpcId;
             // When there is an answer without any thinking steps, we don't want to show the thinking step.
@@ -1075,7 +1086,21 @@ function isAiAssistanceServerSideLoggingEnabled(): boolean {
   return localStorage.getItem('aiAssistance_enableServerSideLogging') !== 'false';
 }
 
+function setAiAssistanceStylingWithFunctionCalling(enabled: boolean): void {
+  if (enabled) {
+    localStorage.setItem('aiAssistance_stylingFunctionCalling', 'true');
+  } else {
+    localStorage.setItem('aiAssistance_stylingFunctionCalling', 'false');
+  }
+}
+
+function isAiAssistanceStylingWithFunctionCallingEnabled(): boolean {
+  return localStorage.getItem('aiAssistance_stylingFunctionCalling') === 'true';
+}
+
 // @ts-ignore
 globalThis.setAiAssistanceServerSideLoggingEnabled = setAiAssistanceServerSideLoggingEnabled;
 // @ts-ignore
 globalThis.setAiAssistanceChangeSummariesEnabled = setAiAssistanceChangeSummariesEnabled;
+// @ts-ignore
+globalThis.setAiAssistanceStylingWithFunctionCalling = setAiAssistanceStylingWithFunctionCalling;