fix(js/ai): Fixes use of namespaced tools in model calls.

js/ai/src/generate/action.ts

@@ -15,6 +15,7 @@
  */
 
 import {
+  GenkitError,
   getStreamingCallback,
   runWithStreamingCallback,
   z,
@@ -116,6 +117,19 @@ async function generate(
   const tools = await resolveTools(registry, rawRequest.tools);
 
   const resolvedFormat = await resolveFormat(registry, rawRequest.output);
+  // Create a lookup of tool names with namespaces stripped to original names
+  const toolMap = tools.reduce<Record<string, ToolAction>>((acc, tool) => {
+    const name = tool.__action.name;
+    const shortName = name.substring(name.lastIndexOf('/') + 1);
+    if (acc[shortName]) {
+      throw new GenkitError({
+        status: 'INVALID_ARGUMENT',
+        message: `Cannot provide two tools with the same name: '${name}' and '${acc[shortName]}'`,
+      });
+    }
+    acc[shortName] = tool;
+    return acc;
+  }, {});
 
   const request = await actionToGenerateRequest(
     rawRequest,
@@ -184,9 +198,7 @@ async function generate(
         'Tool request expected but not provided in tool request part'
       );
     }
-    const tool = tools?.find(
-      (tool) => tool.__action.name === part.toolRequest?.name
-    );
+    const tool = toolMap[part.toolRequest?.name];
     if (!tool) {
       throw Error(`Tool ${part.toolRequest?.name} not found`);
     }
@@ -238,7 +250,7 @@ async function actionToGenerateRequest(
     messages: options.messages,
     config: options.config,
     docs: options.docs,
-    tools: resolvedTools?.map((tool) => toToolDefinition(tool)) || [],
+    tools: resolvedTools?.map((tool) => toToolDefinition(tool, true)) || [],
     output: {
       ...(resolvedFormat?.config || {}),
       schema: toJsonSchema({

js/ai/src/tool.ts

@@ -134,10 +134,16 @@ export async function lookupToolByName(
  * Converts a tool action to a definition of the tool to be passed to a model.
  */
 export function toToolDefinition(
-  tool: Action<z.ZodTypeAny, z.ZodTypeAny>
+  tool: Action<z.ZodTypeAny, z.ZodTypeAny>,
+  stripNamespace = false
 ): ToolDefinition {
+  let name = tool.__action.name;
+  if (stripNamespace) {
+    name = name.substring(name.lastIndexOf('/') + 1);
+  }
+
   return {
-    name: tool.__action.name,
+    name,
     description: tool.__action.description || '',
     outputSchema: toJsonSchema({
       schema: tool.__action.outputSchema ?? z.void(),

js/ai/tests/generate/generate_test.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { z } from '@genkit-ai/core';
+import { PluginProvider, z } from '@genkit-ai/core';
 import { Registry } from '@genkit-ai/core/registry';
 import assert from 'node:assert';
 import { beforeEach, describe, it } from 'node:test';
@@ -43,6 +43,23 @@ describe('toGenerateRequest', () => {
     }
   );
 
+  const namespacedPlugin: PluginProvider = {
+    name: 'namespaced',
+    initializer: async () => {},
+  };
+  registry.registerPluginProvider('namespaced', namespacedPlugin);
+
+  const namespacedTool = defineTool(
+    registry,
+    {
+      name: 'namespaced/add',
+      description: 'add two numbers together',
+      inputSchema: z.object({ a: z.number(), b: z.number() }),
+      outputSchema: z.number(),
+    },
+    async ({ a, b }) => a + b
+  );
+
   const testCases = [
     {
       should: 'translate a string prompt correctly',
@@ -95,6 +112,37 @@ describe('toGenerateRequest', () => {
         output: {},
       },
     },
+    {
+      should: 'strip namespaces from tools when passing to the model',
+      prompt: {
+        model: 'vertexai/gemini-1.0-pro',
+        tools: ['namespaced/add'],
+        prompt: 'Add 10 and 5.',
+      },
+      expectedOutput: {
+        messages: [{ role: 'user', content: [{ text: 'Add 10 and 5.' }] }],
+        config: undefined,
+        docs: undefined,
+        tools: [
+          {
+            description: 'add two numbers together',
+            inputSchema: {
+              $schema: 'http://json-schema.org/draft-07/schema#',
+              additionalProperties: true,
+              properties: { a: { type: 'number' }, b: { type: 'number' } },
+              required: ['a', 'b'],
+              type: 'object',
+            },
+            name: 'namespaced/add',
+            outputSchema: {
+              $schema: 'http://json-schema.org/draft-07/schema#',
+              type: 'number',
+            },
+          },
+        ],
+        output: {},
+      },
+    },
     {
       should:
         'translate a string prompt correctly with tools referenced by their action',

js/genkit/src/genkit.ts

@@ -463,7 +463,7 @@ export class Genkit {
           if (!response.tools && options.tools) {
             response.tools = (
               await resolveTools(this.registry, options.tools)
-            ).map(toToolDefinition);
+            ).map((t) => toToolDefinition(t));
           }
           if (!response.output && options.output) {
             response.output = {

js/testapps/flow-simple-ai/src/index.ts

@@ -28,6 +28,7 @@ import { initializeApp } from 'firebase-admin/app';
 import { getFirestore } from 'firebase-admin/firestore';
 import { MessageSchema, genkit, run, z } from 'genkit';
 import { logger } from 'genkit/logging';
+import { PluginProvider } from 'genkit/plugin';
 import { Allow, parse } from 'partial-json';
 
 logger.setLogLevel('debug');
@@ -53,6 +54,32 @@ const ai = genkit({
   plugins: [googleAI(), vertexAI()],
 });
 
+const math: PluginProvider = {
+  name: 'math',
+  initializer: async () => {
+    ai.defineTool(
+      {
+        name: 'math/add',
+        description: 'add two numbers',
+        inputSchema: z.object({ a: z.number(), b: z.number() }),
+        outputSchema: z.number(),
+      },
+      async ({ a, b }) => a + b
+    );
+
+    ai.defineTool(
+      {
+        name: 'math/subtract',
+        description: 'subtract two numbers',
+        inputSchema: z.object({ a: z.number(), b: z.number() }),
+        outputSchema: z.number(),
+      },
+      async ({ a, b }) => a - b
+    );
+  },
+};
+ai.registry.registerPluginProvider('math', math);
+
 const app = initializeApp();
 
 export const jokeFlow = ai.defineFlow(
@@ -538,11 +565,18 @@ export const arrayStreamTester = ai.defineStreamingFlow(
   }
 );
 
-// async function main() {
-//   const { stream, output } = arrayStreamTester();
-//   for await (const chunk of stream) {
-//     console.log(chunk);
-//   }
-//   console.log(await output);
-// }
-// main();
+ai.defineFlow(
+  {
+    name: 'math',
+    inputSchema: z.string(),
+    outputSchema: z.string(),
+  },
+  async (query) => {
+    const { text } = await ai.generate({
+      model: gemini15Flash,
+      prompt: query,
+      tools: ['math/add', 'math/subtract'],
+    });
+    return text;
+  }
+);