[AI Assistance] Refactor function calling

Bug: 385082492
Change-Id: I8bb2b553b7825d2a2fa88071dcd7d77fde7ed71d
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6110476
Commit-Queue: Alex Rudenko <[email protected]>
Reviewed-by: Wolfgang Beyer <[email protected]>
Auto-Submit: Alex Rudenko <[email protected]>
Commit-Queue: Wolfgang Beyer <[email protected]>

Author: OrKoN

front_end/core/host/AidaClient.ts

@@ -68,7 +68,7 @@ export interface FunctionPrimitiveParams extends BaseFunctionParam {
 
 interface FunctionArrayParam extends BaseFunctionParam {
   type: ParametersTypes.ARRAY;
-  items: FunctionPrimitiveParams[];
+  items: FunctionPrimitiveParams;
 }
 
 export interface FunctionObjectParam extends BaseFunctionParam {
@@ -87,7 +87,7 @@ export interface FunctionDeclaration {
    * A description for the LLM to understand what the specific function will do once called.
    */
   description: string;
-  parameters: FunctionObjectParam|FunctionPrimitiveParams;
+  parameters: FunctionObjectParam|FunctionPrimitiveParams|FunctionArrayParam;
 }
 
 // Raw media bytes.
@@ -141,7 +141,7 @@ export type RpcGlobalId = string|number;
 export interface AidaRequest {
   client: string;
   // eslint-disable-next-line @typescript-eslint/naming-convention
-  current_message?: Content;
+  current_message: Content;
   preamble?: string;
   // eslint-disable-next-line @typescript-eslint/naming-convention
   historical_contexts?: Content[];

front_end/panels/ai_assistance/agents/AiAgent.test.ts

@@ -551,7 +551,7 @@ describeWithEnvironment('AiAgent', () => {
 
       constructor(opts: AiAssistance.AgentOptions) {
         super(opts);
-        this.functionDeclarations.set('testFn', {
+        this.declareFunction('testFn', {
           description: 'test fn description',
           parameters: {type: Host.AidaClient.ParametersTypes.OBJECT, properties: {}, description: 'arg description'},
           handler: this.#test.bind(this),

front_end/panels/ai_assistance/agents/AiAgent.ts

@@ -155,14 +155,12 @@ export abstract class ConversationContext<T> {
   }
 }
 
-export type FunctionDeclaration<Args extends Record<string, unknown>, ReturnType = Record<string, unknown>> = {
+export type FunctionDeclaration<Args, ReturnType> = {
   description: string,
   parameters: Host.AidaClient.FunctionObjectParam,
   handler: (args: Args) => Promise<ReturnType>,
 };
 
-export type FunctionDeclarations = Map<string, FunctionDeclaration<Record<string, unknown>, Record<string, unknown>>>;
-
 export abstract class AiAgent<T> {
   static validTemperature(temperature: number|undefined): number|undefined {
     return typeof temperature === 'number' && temperature >= 0 ? temperature : undefined;
@@ -178,7 +176,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>;
-  protected functionDeclarations: FunctionDeclarations = new Map();
+  #functionDeclarations = new Map<string, FunctionDeclaration<unknown, unknown>>();
   #generatedFromHistory = false;
 
   /**
@@ -203,7 +201,21 @@ export abstract class AiAgent<T> {
   }
 
   get chatHistoryForTesting(): Array<Host.AidaClient.Content> {
-    return this.#buildChatHistoryForAida();
+    return this.buildChatHistoryForAida();
+  }
+
+  declareFunction<Args, ReturnType>(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<unknown, unknown>);
+  }
+
+  async callFunction(name: string, args: unknown): Promise<Record<string, unknown>> {
+    const call = this.#functionDeclarations.get(name);
+    return (call ? await call.handler(args) : {
+             error: `Function ${name} is not found.`,
+           }) as Record<string, unknown>;
   }
 
   set chatNewHistoryForTesting(history: HistoryEntryStorage) {
@@ -241,7 +253,7 @@ export abstract class AiAgent<T> {
 
   #buildFunctionDeclarationsForAida(): Host.AidaClient.FunctionDeclaration[] {
     const result: Host.AidaClient.FunctionDeclaration[] = [];
-    for (const [name, definition] of this.functionDeclarations.entries()) {
+    for (const [name, definition] of this.#functionDeclarations.entries()) {
       result.push({
         name,
         description: definition.description,
@@ -278,11 +290,6 @@ export abstract class AiAgent<T> {
     for await (aidaResponse of this.#aidaClient.fetch(request, options)) {
       response = aidaResponse.explanation;
       rpcId = aidaResponse.metadata.rpcGlobalId ?? rpcId;
-
-      if (aidaResponse.functionCalls) {
-        throw new Error('Function calling not supported yet');
-      }
-
       const parsedResponse = this.parseResponse(aidaResponse);
       yield {
         rpcId,
@@ -309,7 +316,7 @@ export abstract class AiAgent<T> {
       parts: [part],
       role: Host.AidaClient.Role.USER,
     };
-    const history = this.#buildChatHistoryForAida();
+    const history = this.buildChatHistoryForAida();
     const declarations = this.#buildFunctionDeclarationsForAida();
     const request: Host.AidaClient.AidaRequest = {
       client: Host.AidaClient.CLIENT_NAME,
@@ -380,7 +387,7 @@ STOP`;
     return text;
   }
 
-  #buildChatHistoryForAida(): Host.AidaClient.Content[] {
+  buildChatHistoryForAida(): Host.AidaClient.Content[] {
     const history: Array<Host.AidaClient.Content> = [];
     let currentParsedStep: ParsedStep = {};
     let lastRunStartIdx = 0;