[AI Assistance] Add types for Function calling

Provide the ability to create function calling agents

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

Author: Lightning00Blade

front_end/core/host/AidaClient.test.ts

@@ -154,10 +154,27 @@ describeWithEnvironment('AidaClient', () => {
     const provider = new Host.AidaClient.AidaClient();
     const results = await getAllResults(provider);
     assert.deepStrictEqual(results, [
-      {explanation: 'hello ', metadata: {rpcGlobalId: 123}, completed: false},
-      {explanation: 'hello brave ', metadata: {rpcGlobalId: 123}, completed: false},
-      {explanation: 'hello brave new world!', metadata: {rpcGlobalId: 123}, completed: false},
-      {explanation: 'hello brave new world!', metadata: {rpcGlobalId: 123}, completed: true},
+      {
+        explanation: 'hello ',
+        metadata: {rpcGlobalId: 123},
+        completed: false,
+      },
+      {
+        explanation: 'hello brave ',
+        metadata: {rpcGlobalId: 123},
+        completed: false,
+      },
+      {
+        explanation: 'hello brave new world!',
+        metadata: {rpcGlobalId: 123},
+        completed: false,
+      },
+      {
+        explanation: 'hello brave new world!',
+        metadata: {rpcGlobalId: 123},
+        functionCall: undefined,
+        completed: true,
+      },
     ]);
   });
 
@@ -175,8 +192,17 @@ describeWithEnvironment('AidaClient', () => {
     const provider = new Host.AidaClient.AidaClient();
     const results = await getAllResults(provider);
     assert.deepStrictEqual(results, [
-      {explanation: 'hello world', metadata: {rpcGlobalId: 123}, completed: false},
-      {explanation: 'hello world', metadata: {rpcGlobalId: 123}, completed: true},
+      {
+        explanation: 'hello world',
+        metadata: {rpcGlobalId: 123},
+        completed: false,
+      },
+      {
+        explanation: 'hello world',
+        metadata: {rpcGlobalId: 123},
+        functionCall: undefined,
+        completed: true,
+      },
     ]);
   });
 
@@ -255,6 +281,7 @@ describeWithEnvironment('AidaClient', () => {
             'If it were so, it was a grievous fault,\n' +
             'And grievously hath Caesar answer’d it.\n',
         metadata: {rpcGlobalId: 123},
+        functionCall: undefined,
         completed: true,
       },
     ]);
@@ -315,6 +342,7 @@ describeWithEnvironment('AidaClient', () => {
             },
           ],
         },
+        functionCall: undefined,
         completed: true,
       },
     ]);

front_end/core/host/AidaClient.ts

@@ -31,10 +31,53 @@ export interface Content {
   role: Role;
 }
 
-export interface Part {
-  text?: string;
+export type Part = {
+  text: string,
+}|{
+  functionResponse: {
+    name: string,
+    response: Record<string, unknown>,
+  },
+}|{
   // Inline media bytes.
-  inlineData?: MediaBlob;
+  inlineData: MediaBlob,
+};
+
+export const enum ParametersTypes {
+  STRING = 1,
+  NUMBER = 2,
+  INTEGER = 3,
+  BOOLEAN = 4,
+  ARRAY = 5,
+  OBJECT = 6,
+}
+
+interface BaseFunctionParam {
+  description: string;
+  nullable?: boolean;
+}
+
+interface FunctionPrimitiveParams extends BaseFunctionParam {
+  type: ParametersTypes.BOOLEAN|ParametersTypes.INTEGER|ParametersTypes.STRING|ParametersTypes.BOOLEAN;
+}
+interface FunctionObjectParam extends BaseFunctionParam {
+  type: ParametersTypes.OBJECT;
+  // TODO: this can be also be ObjectParams
+  properties: {[Key in string]: FunctionPrimitiveParams};
+}
+// TODO: Add FunctionArrayParam
+
+/**
+ * More about function declaration can be read at
+ * https://ai.google.dev/gemini-api/docs/function-calling
+ */
+export interface FunctionDeclaration {
+  name: string;
+  /**
+   * A description for the LLM to understand what the specific function will do once called.
+   */
+  description: string;
+  parameters: FunctionObjectParam|FunctionPrimitiveParams;
 }
 
 // Raw media bytes.
@@ -83,12 +126,14 @@ export enum UserTier {
 }
 
 export interface AidaRequest {
+  client: string;
   // eslint-disable-next-line @typescript-eslint/naming-convention
   current_message?: Content;
   preamble?: string;
   // eslint-disable-next-line @typescript-eslint/naming-convention
   historical_contexts?: Content[];
-  client: string;
+  // eslint-disable-next-line @typescript-eslint/naming-convention
+  function_declarations?: FunctionDeclaration[];
   options?: {
     temperature?: number,
     // eslint-disable-next-line @typescript-eslint/naming-convention
@@ -145,6 +190,11 @@ export interface AttributionMetadata {
   citations: Citation[];
 }
 
+export interface AidaFunctionCallResponse {
+  name: string;
+  args: Record<string, unknown>;
+}
+
 export interface AidaResponseMetadata {
   rpcGlobalId?: number;
   attributionMetadata?: AttributionMetadata[];
@@ -153,6 +203,7 @@ export interface AidaResponseMetadata {
 export interface AidaResponse {
   explanation: string;
   metadata: AidaResponseMetadata;
+  functionCall?: AidaFunctionCallResponse;
   completed: boolean;
 }
 
@@ -261,6 +312,7 @@ export class AidaClient {
     let chunk;
     const text = [];
     let inCodeChunk = false;
+    let functionCall: AidaFunctionCallResponse|undefined = undefined;
     const metadata: AidaResponseMetadata = {rpcGlobalId: 0};
     while ((chunk = await stream.read())) {
       let textUpdated = false;
@@ -315,6 +367,11 @@ export class AidaClient {
           }
           text.push(result.codeChunk.code);
           textUpdated = true;
+        } else if ('functionCallChunk' in result) {
+          functionCall = {
+            name: result.functionCallChunk.functionCall.name,
+            args: result.functionCallChunk.functionCall.args,
+          };
         } else if ('error' in result) {
           throw new Error(`Server responded: ${JSON.stringify(result)}`);
         } else {
@@ -332,6 +389,7 @@ export class AidaClient {
     yield {
       explanation: text.join('') + (inCodeChunk ? CODE_CHUNK_SEPARATOR : ''),
       metadata,
+      functionCall,
       completed: true,
     };
   }

front_end/panels/freestyler/AiAgent.test.ts

@@ -132,7 +132,7 @@ describeWithEnvironment('AiAgent', () => {
         serverSideLoggingEnabled: false,
       });
       const request = agent.buildRequest({text: 'test input'});
-      assert.strictEqual(request.current_message?.parts[0].text, 'test input');
+      assert.deepStrictEqual(request.current_message?.parts[0], {text: 'test input'});
       assert.strictEqual(request.historical_contexts, undefined);
     });
 
@@ -149,7 +149,7 @@ describeWithEnvironment('AiAgent', () => {
         aidaClient: {} as Host.AidaClient.AidaClient,
       });
       const request = agent.buildRequest({text: 'test input'});
-      assert.strictEqual(request.current_message?.parts[0].text, 'test input');
+      assert.deepStrictEqual(request.current_message?.parts[0], {text: 'test input'});
       assert.strictEqual(request.preamble, 'preamble');
       assert.strictEqual(request.historical_contexts, undefined);
     });
@@ -191,7 +191,7 @@ describeWithEnvironment('AiAgent', () => {
         },
       ];
       const request = agent.buildRequest({text: 'test input'});
-      assert.strictEqual(request.current_message?.parts[0].text, 'test input');
+      assert.deepStrictEqual(request.current_message?.parts[0], {text: 'test input'});
       assert.deepStrictEqual(request.historical_contexts, [
         {
           parts: [{text: 'test'}],
@@ -239,7 +239,7 @@ describeWithEnvironment('AiAgent', () => {
         },
       ];
       const request = agent.buildRequest({text: 'test input'});
-      assert.strictEqual(request.current_message?.parts[0].text, 'test input');
+      assert.deepStrictEqual(request.current_message?.parts[0], {text: 'test input'});
       assert.deepStrictEqual(request.historical_contexts, undefined);
     });
 
@@ -300,7 +300,7 @@ describeWithEnvironment('AiAgent', () => {
         },
       ];
       const request = agent.buildRequest({text: 'test input'});
-      assert.strictEqual(request.current_message?.parts[0].text, 'test input');
+      assert.deepStrictEqual(request.current_message?.parts[0], {text: 'test input'});
       assert.deepStrictEqual(request.historical_contexts, [
         {
           parts: [{text: 'test2'}],

front_end/panels/freestyler/AiAgent.ts

@@ -156,6 +156,10 @@ export abstract class ConversationContext<T> {
   }
 }
 
+interface AgentFunctionDefinition extends Host.AidaClient.FunctionDeclaration {
+  method: (...args: any[]) => Record<string, unknown>;
+}
+
 export abstract class AiAgent<T> {
   static validTemperature(temperature: number|undefined): number|undefined {
     return typeof temperature === 'number' && temperature >= 0 ? temperature : undefined;
@@ -171,6 +175,7 @@ export abstract class AiAgent<T> {
   abstract readonly clientFeature: Host.AidaClient.ClientFeature;
   abstract readonly userTier: string|undefined;
   abstract handleContextDetails(select: ConversationContext<T>|null): AsyncGenerator<ContextResponse, void, void>;
+  functionDefinitions: AgentFunctionDefinition[]|undefined;
   #generatedFromHistory = false;
 
   /**
@@ -231,6 +236,16 @@ export abstract class AiAgent<T> {
     return this.#generatedFromHistory;
   }
 
+  get functionDeclarations(): Host.AidaClient.FunctionDeclaration[]|undefined {
+    return this.functionDefinitions?.map(call => {
+      return {
+        name: call.name,
+        description: call.description,
+        parameters: call.parameters,
+      } satisfies Host.AidaClient.FunctionDeclaration;
+    });
+  }
+
   serialized(): SerializedAgent {
     return {
       id: this.id,
@@ -256,8 +271,17 @@ export abstract class AiAgent<T> {
     for await (rawResponse of this.#aidaClient.fetch(request, options)) {
       response = rawResponse.explanation;
       rpcId = rawResponse.metadata.rpcGlobalId ?? rpcId;
+
+      if (rawResponse.functionCall) {
+        throw new Error('Function calling not supported yet');
+      }
+
       const parsedResponse = this.parseResponse(response);
-      yield {rpcId, parsedResponse, completed: rawResponse.completed};
+      yield {
+        rpcId,
+        parsedResponse,
+        completed: rawResponse.completed,
+      };
     }
 
     debugLog({
@@ -273,16 +297,22 @@ export abstract class AiAgent<T> {
     localStorage.setItem('freestylerStructuredLog', JSON.stringify(this.#structuredLog));
   }
 
-  buildRequest(opts: BuildRequestOptions): Host.AidaClient.AidaRequest {
-    const currentMessage = {parts: [{text: opts.text}], role: Host.AidaClient.Role.USER};
+  buildRequest(part: Host.AidaClient.Part): Host.AidaClient.AidaRequest {
+    const currentMessage: Host.AidaClient.Content = {
+      parts: [part],
+      role: Host.AidaClient.Role.USER,
+    };
     const history = this.#chatHistoryForAida;
+    const declarations = this.functionDeclarations;
     const request: Host.AidaClient.AidaRequest = {
+      client: Host.AidaClient.CLIENT_NAME,
       // eslint-disable-next-line @typescript-eslint/naming-convention
       current_message: currentMessage,
       preamble: this.preamble,
       // eslint-disable-next-line @typescript-eslint/naming-convention
       historical_contexts: history.length ? history : undefined,
-      client: Host.AidaClient.CLIENT_NAME,
+      // eslint-disable-next-line @typescript-eslint/naming-convention
+      ...(declarations ? {function_declarations: declarations} : {}),
       options: {
         temperature: AiAgent.validTemperature(this.options.temperature),
         // eslint-disable-next-line @typescript-eslint/naming-convention
@@ -364,14 +394,18 @@ STOP`;
           flushCurrentStep();
           history.push({
             role: Host.AidaClient.Role.USER,
-            parts: [{text: data.query}],
+            parts: [{
+              text: data.query,
+            }],
           });
           break;
         }
         case ResponseType.ANSWER:
           history.push({
             role: Host.AidaClient.Role.MODEL,
-            parts: [{text: this.formatParsedAnswer({answer: data.text})}],
+            parts: [{
+              text: this.formatParsedAnswer({answer: data.text}),
+            }],
           });
           break;
         case ResponseType.TITLE:

front_end/panels/freestyler/FreestylerAgent.test.ts

@@ -839,9 +839,9 @@ STOP`,
       assert.isUndefined(requests[0].historical_contexts, 'Unexpected historical contexts in the initial request');
       assert.exists(requests[0].current_message);
       assert.lengthOf(requests[0].current_message.parts, 1);
-      assert.strictEqual(
-          requests[0].current_message.parts[0].text,
-          '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+      assert.deepStrictEqual(
+          requests[0].current_message.parts[0],
+          {text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test'},
           'Unexpected input text in the initial request');
       assert.strictEqual(requests[0].current_message.role, Host.AidaClient.Role.USER);
       assert.deepStrictEqual(
@@ -862,8 +862,8 @@ STOP`,
           'Unexpected historical contexts in the follow-up request');
       assert.exists(requests[1].current_message);
       assert.lengthOf(requests[1].current_message.parts, 1);
-      assert.strictEqual(
-          requests[1].current_message.parts[0].text, 'OBSERVATION: test data',
+      assert.deepStrictEqual(
+          requests[1].current_message.parts[0], {text: 'OBSERVATION: test data'},
           'Unexpected input in the follow-up request');
     });
 

test/e2e/freestyler/freestyler_test.ts

@@ -307,7 +307,8 @@ STOP`,
       iframeId: 'iframe',
     });
 
-    assert.deepStrictEqual(
-        result.at(-1)!.request.current_message.parts[0].text, 'OBSERVATION: {"title":"I have a title"}');
+    assert.deepStrictEqual(result.at(-1)!.request.current_message.parts[0], {
+      text: 'OBSERVATION: {"title":"I have a title"}',
+    });
   });
 });