feat(js/ai/chat): use predefined agent's models on multi-agent systems (#3777)

Author: marioToribi0

js/ai/src/generate/action.ts

@@ -221,6 +221,11 @@ function applyTransferPreamble(
     return rawRequest;
   }
 
+  // if the transfer preamble has a model, use it for the next request
+  if (transferPreamble?.model) {
+    rawRequest.model = transferPreamble.model;
+  }
+
   return stripUndefinedProps({
     ...rawRequest,
     messages: [
@@ -402,8 +407,10 @@ async function generate(
 
   let nextRequest = {
     ...rawRequest,
+
     messages: [...rawRequest.messages, generatedMessage.toJSON(), toolMessage!],
   };
+
   nextRequest = applyTransferPreamble(nextRequest, transferPreamble);
 
   // then recursively call for another loop

js/ai/src/generate/resolve-tool-requests.ts

@@ -101,7 +101,11 @@ export async function resolveToolRequest(
 
   // if it's a prompt action, go ahead and render the preamble
   if (isPromptAction(tool)) {
-    const preamble = await tool(part.toolRequest.input);
+    const metadata = tool.__action.metadata as Record<string, any>;
+    const preamble = {
+      ...(await tool(part.toolRequest.input)),
+      model: metadata.prompt?.model,
+    };
     const response = {
       toolResponse: {
         name: part.toolRequest.name,

js/doc-snippets/src/multi-agent/multi.ts

@@ -31,6 +31,7 @@ const reservationListTool = reservationTool;
 // Define a prompt that represents a specialist agent
 const reservationAgent = ai.definePrompt({
   name: 'reservationAgent',
+  model: 'googleai/gemini-2.5-flash',
   description: 'Reservation Agent can help manage guest reservations',
   tools: [reservationTool, reservationCancelationTool, reservationListTool],
   system: 'Help guests make and manage reservations',
@@ -43,6 +44,7 @@ const complaintAgent = ai.prompt('complaintAgent');
 // The triage agent is the agent that users interact with initially
 const triageAgent = ai.definePrompt({
   name: 'triageAgent',
+  model: 'googleai/gemini-2.5-flash-lite',
   description: 'Triage Agent',
   tools: [reservationAgent, menuInfoAgent, complaintAgent],
   system: `You are an AI customer service agent for Pavel's Cafe.

js/pnpm-lock.yaml

@@ -0,0 +1,30 @@
+{
+  "name": "multiagents-demo",
+  "version": "1.0.0",
+  "description": "",
+  "main": "lib/index.js",
+  "scripts": {
+    "test": "echo \"Error: no test specified\" && exit 1",
+    "start": "node lib/index.js",
+    "build": "tsc",
+    "build:watch": "tsc --watch",
+    "genkit:dev": "genkit start -- npx tsx --watch src/index.ts"
+  },
+  "keywords": [],
+  "author": "",
+  "license": "ISC",
+  "dependencies": {
+    "genkit": "workspace:*",
+    "@genkit-ai/firebase": "workspace:*",
+    "@genkit-ai/google-cloud": "workspace:*",
+    "@genkit-ai/google-genai": "workspace:*",
+    "express": "^4.20.0",
+    "node-fetch": "3.3.2",
+    "wav": "^1.0.2"
+  },
+  "devDependencies": {
+    "@types/wav": "^1.0.4",
+    "dotenv": "^16.4.5",
+    "typescript": "^5.6.2"
+  }
+}

js/testapps/multiagents-demo/package.json

@@ -0,0 +1,40 @@
+/**
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { ai } from '../config/genkit';
+
+import {
+  fetchCatalogTool,
+  fetchMostPopularItemsTool,
+  getProductDetailsTool,
+  searchCatalogTool,
+} from '../tools';
+
+// Simple agent for browsing and searching products
+export const agent = ai.definePrompt({
+  name: 'catalogAgent',
+  description: 'Catalog Agent can help customers browse and search products',
+  model: 'googleai/gemini-2.5-pro',
+  tools: [
+    fetchCatalogTool,
+    fetchMostPopularItemsTool,
+    searchCatalogTool,
+    getProductDetailsTool,
+  ],
+  system: `You are a helpful catalog agent for TechStore Computer Shop.
+      Help customers browse products, search for items, and get product details.
+      Be friendly and provide accurate information about products, prices, and availability.`,
+});

js/testapps/multiagents-demo/src/agents/catalogAgent.ts

@@ -0,0 +1,20 @@
+/**
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export { agent as catalogAgent } from './catalogAgent';
+export { agent as paymentAgent } from './paymentAgent';
+export { agent as representativeAgent } from './representativeAgent';
+export { agent as triageAgent } from './triageAgent';

js/testapps/multiagents-demo/src/agents/index.ts

@@ -0,0 +1,29 @@
+/**
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { ai } from '../config/genkit';
+import { processPaymentTool } from '../tools';
+
+// Simple agent for processing payments
+export const agent = ai.definePrompt({
+  name: 'paymentAgent',
+  model: 'googleai/gemini-2.5-flash-lite',
+  description: 'Payment Agent can help process payments',
+  tools: [processPaymentTool],
+  system: `You are a payment agent for TechStore Computer Shop.
+      Help customers process payments.
+      Always confirm payment details before processing and provide clear transaction information.`,
+});

js/testapps/multiagents-demo/src/agents/paymentAgent.ts

@@ -0,0 +1,28 @@
+/**
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { ai } from '../config/genkit';
+import { getStoreInfoTool } from '../tools';
+
+// Simple agent for providing store information
+export const agent = ai.definePrompt({
+  name: 'representativeAgent',
+  model: 'googleai/gemini-2.0-flash',
+  description: 'Representative Agent can provide store info',
+  tools: [getStoreInfoTool],
+  system: `You are a customer service representative for TechStore Computer Shop.
+      Help customers provide store information.`,
+});

js/testapps/multiagents-demo/src/agents/representativeAgent.ts

@@ -0,0 +1,34 @@
+/**
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { ai } from '../config/genkit';
+import { agent as catalogAgent } from './catalogAgent';
+import { agent as paymentAgent } from './paymentAgent';
+import { agent as representativeAgent } from './representativeAgent';
+
+// The triage agent routes customers to the appropriate specialist
+export const agent = ai.definePrompt({
+  name: 'triageAgent',
+  description: 'Triage Agent',
+  model: 'googleai/gemini-2.5-flash',
+  tools: [catalogAgent, paymentAgent, representativeAgent],
+  system: `You are an AI customer service agent for TechStore Computer Shop.
+    Greet the user and ask them how you can help. Route them to the appropriate specialist agent:
+    - Use catalogAgent for browsing products, searching items, or getting product information
+    - Use paymentAgent for payment processing
+    - Use representativeAgent for providing store information
+    If you cannot help the customer with the available tools, politely explain so.`,
+});