Consolidate action subtype metadata setting in defineAction

js/ai/src/embedder.ts

@@ -16,7 +16,6 @@
 
 import { Action, defineAction } from '@genkit-ai/core';
 import { lookupAction } from '@genkit-ai/core/registry';
-import { setCustomMetadataAttributes } from '@genkit-ai/core/tracing';
 import * as z from 'zod';
 import { Document, DocumentData, DocumentDataSchema } from './document.js';
 
@@ -83,13 +82,11 @@ export function defineEmbedder<
         info: options.info,
       },
     },
-    (i) => {
-      setCustomMetadataAttributes({ subtype: 'embedder' });
-      return runner(
+    (i) =>
+      runner(
         i.input.map((dd) => new Document(dd)),
         i.options
-      );
-    }
+      )
   );
   const ewm = withMetadata(
     embedder as Action<typeof EmbedRequestSchema, typeof EmbedResponseSchema>,

js/ai/src/evaluator.ts

@@ -17,11 +17,7 @@
 import { Action, defineAction } from '@genkit-ai/core';
 import { logger } from '@genkit-ai/core/logging';
 import { lookupAction } from '@genkit-ai/core/registry';
-import {
-  SPAN_TYPE_ATTR,
-  runInNewSpan,
-  setCustomMetadataAttributes,
-} from '@genkit-ai/core/tracing';
+import { SPAN_TYPE_ATTR, runInNewSpan } from '@genkit-ai/core/tracing';
 import * as z from 'zod';
 
 export const ATTR_PREFIX = 'genkit';
@@ -143,7 +139,6 @@ export function defineEvaluator<
       metadata: metadata,
     },
     async (i) => {
-      setCustomMetadataAttributes({ subtype: 'evaluator' });
       let evalResponses: EvalResponses = [];
       for (let index = 0; index < i.dataset.length; index++) {
         const datapoint = i.dataset[index];

js/ai/src/model.ts

@@ -21,7 +21,6 @@ import {
   StreamingCallback,
 } from '@genkit-ai/core';
 import { toJsonSchema } from '@genkit-ai/core/schema';
-import { setCustomMetadataAttributes } from '@genkit-ai/core/tracing';
 import { performance } from 'node:perf_hooks';
 import { z } from 'zod';
 import { conformOutput, validateSupport } from './model/middleware.js';
@@ -311,7 +310,6 @@ export function defineModel<
       },
     },
     (input) => {
-      setCustomMetadataAttributes({ subtype: 'model' });
       const startTimeMs = performance.now();
       return runner(input, getStreamingCallback())
         .then((response) => {

js/ai/src/prompt.ts

@@ -14,9 +14,8 @@
  * limitations under the License.
  */
 
-import { Action, action, JSONSchema7 } from '@genkit-ai/core';
-import { lookupAction, registerAction } from '@genkit-ai/core/registry';
-import { setCustomMetadataAttributes } from '@genkit-ai/core/tracing';
+import { Action, defineAction, JSONSchema7 } from '@genkit-ai/core';
+import { lookupAction } from '@genkit-ai/core/registry';
 import z from 'zod';
 import { GenerateOptions } from './generate';
 import { GenerateRequest, GenerateRequestSchema, ModelArgument } from './model';
@@ -52,20 +51,17 @@ export function definePrompt<I extends z.ZodTypeAny>(
   },
   fn: PromptFn<I>
 ): PromptAction<I> {
-  const a = action(
+  const a = defineAction(
     {
+      actionType: 'prompt',
       name,
       description,
       inputSchema,
       inputJsonSchema,
       metadata: { ...(metadata || { prompt: {} }), type: 'prompt' },
     },
-    (i: I): Promise<GenerateRequest> => {
-      setCustomMetadataAttributes({ subtype: 'prompt' });
-      return fn(i);
-    }
+    fn
   );
-  registerAction('prompt', a);
   return a as PromptAction<I>;
 }
 

js/ai/src/retriever.ts

@@ -16,7 +16,6 @@
 
 import { Action, defineAction } from '@genkit-ai/core';
 import { lookupAction } from '@genkit-ai/core/registry';
-import { setCustomMetadataAttributes } from '@genkit-ai/core/tracing';
 import * as z from 'zod';
 import { Document, DocumentData, DocumentDataSchema } from './document.js';
 import { EmbedderInfo } from './embedder.js';
@@ -128,10 +127,7 @@ export function defineRetriever<
         info: options.info,
       },
     },
-    (i) => {
-      setCustomMetadataAttributes({ subtype: 'retriever' });
-      return runner(new Document(i.query), i.options);
-    }
+    (i) => runner(new Document(i.query), i.options)
   );
   const rwm = retrieverWithMetadata(
     retriever as Action<
@@ -169,13 +165,11 @@ export function defineIndexer<IndexerOptions extends z.ZodTypeAny>(
         embedderInfo: options.embedderInfo,
       },
     },
-    (i) => {
-      setCustomMetadataAttributes({ subtype: 'indexer' });
-      return runner(
+    (i) =>
+      runner(
         i.documents.map((dd) => new Document(dd)),
         i.options
-      );
-    }
+      )
   );
   const iwm = indexerWithMetadata(
     indexer as Action<typeof IndexerRequestSchema, z.ZodVoid>,

js/ai/src/tool.ts

@@ -128,10 +128,7 @@ export function defineTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
       outputJsonSchema,
       metadata: { ...(metadata || {}), type: 'tool' },
     },
-    (i) => {
-      setCustomMetadataAttributes({ subtype: 'tool' });
-      return fn(i);
-    }
+    (i) => fn(i)
   );
   return a as ToolAction<I, O>;
 }

js/core/src/action.ts

@@ -21,7 +21,11 @@ import * as z from 'zod';
 import { ActionType, lookupPlugin, registerAction } from './registry.js';
 import { parseSchema } from './schema.js';
 import * as telemetry from './telemetry.js';
-import { SPAN_TYPE_ATTR, runInNewSpan } from './tracing.js';
+import {
+  SPAN_TYPE_ATTR,
+  runInNewSpan,
+  setCustomMetadataAttributes,
+} from './tracing.js';
 
 export { Status, StatusCodes, StatusSchema } from './statusTypes.js';
 export { JSONSchema7 };
@@ -176,7 +180,10 @@ export function defineAction<
   },
   fn: (input: z.infer<I>) => Promise<z.infer<O>>
 ): Action<I, O> {
-  const act = action(config, fn);
+  const act = action(config, (i: I): Promise<z.infer<O>> => {
+    setCustomMetadataAttributes({ subtype: config.actionType });
+    return fn(i);
+  });
   act.__action.actionType = config.actionType;
   registerAction(config.actionType, act);
   return act;