feat: expose defineFormat on Genkit instance, allowing to define custom formats

js/ai/src/formats/index.ts

@@ -18,18 +18,20 @@ import { JSONSchema } from '@genkit-ai/core';
 import { Registry } from '@genkit-ai/core/registry';
 import { OutputOptions } from '../generate.js';
 import { MessageData, TextPart } from '../model.js';
-import { arrayFormatter } from './array';
-import { enumFormatter } from './enum';
-import { jsonFormatter } from './json';
-import { jsonlFormatter } from './jsonl';
-import { textFormatter } from './text';
-import { Formatter } from './types';
+import { arrayFormatter } from './array.js';
+import { enumFormatter } from './enum.js';
+import { jsonFormatter } from './json.js';
+import { jsonlFormatter } from './jsonl.js';
+import { textFormatter } from './text.js';
+import { type Formatter } from './types.js';
+
+export { type Formatter };
 
 export function defineFormat(
   registry: Registry,
   options: { name: string } & Formatter['config'],
   handler: Formatter['handler']
-) {
+): { config: Formatter['config']; handler: Formatter['handler'] } {
   const { name, ...config } = options;
   const formatter = { config, handler };
   registry.registerValue('format', name, formatter);

js/genkit/package.json

@@ -80,6 +80,12 @@
       "import": "./lib/schema.mjs",
       "default": "./lib/schema.js"
     },
+    "./formats": {
+      "types": "./lib/formats.d.ts",
+      "require": "./lib/formats.js",
+      "import": "./lib/formats.mjs",
+      "default": "./lib/formats.js"
+    },
     "./retriever": {
       "types": "./lib/retriever.d.ts",
       "require": "./lib/retriever.js",
@@ -170,6 +176,9 @@
       "schema": [
         "lib/schema"
       ],
+      "formats": [
+        "lib/formats"
+      ],
       "retriever": [
         "lib/retriever"
       ],

js/genkit/src/formats.ts

@@ -0,0 +1,17 @@
+/**
+ * Copyright 2024 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 { type FormatArgument, type Formatter } from '@genkit-ai/ai/formats';

js/genkit/src/genkit.ts

@@ -68,7 +68,11 @@ import {
   EvaluatorAction,
   EvaluatorFn,
 } from '@genkit-ai/ai/evaluator';
-import { configureFormats } from '@genkit-ai/ai/formats';
+import {
+  configureFormats,
+  defineFormat,
+  Formatter,
+} from '@genkit-ai/ai/formats';
 import {
   defineModel,
   DefineModelOptions,
@@ -247,6 +251,48 @@ export class Genkit {
     return defineSchema(this.registry, name, schema);
   }
 
+  /**
+   * Defines and registers a custom model output formatter.
+   *
+   * Here's an example of a custom JSON output formatter:
+   *
+   * ```ts
+   * import { extractJson } from 'genkit/extract';
+   *
+   * ai.defineFormat(
+   *   { name: 'customJson' },
+   *   (schema) => {
+   *     let instructions: string | undefined;
+   *     if (schema) {
+   *       instructions = `Output should be in JSON format and conform to the following schema:
+   * \`\`\`
+   * ${JSON.stringify(schema)}
+   * \`\`\`
+   * `;
+   *     }
+   *     return {
+   *       parseChunk: (chunk) => extractJson(chunk.accumulatedText),
+   *       parseMessage: (message) => extractJson(message.text),
+   *       instructions,
+   *     };
+   *   }
+   * );
+   *
+   * const { output } = await ai.generate({
+   *   prompt: 'Invent a menu item for a pirate themed restaurant.',
+   *   output: { format: 'customJson', schema: MenuItemSchema },
+   * });
+   * ```
+   */
+  defineFormat(
+    options: {
+      name: string;
+    } & Formatter['config'],
+    handler: Formatter['handler']
+  ): { config: Formatter['config']; handler: Formatter['handler'] } {
+    return defineFormat(this.registry, options, handler);
+  }
+
   /**
    * Defines and registers a schema from a JSON schema.
    *

js/genkit/tests/formats_test.ts

@@ -0,0 +1,77 @@
+/**
+ * Copyright 2024 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 assert from 'node:assert';
+import { beforeEach, describe, it } from 'node:test';
+import { Genkit, genkit } from '../src/genkit';
+import { defineEchoModel } from './helpers';
+
+describe('formats', () => {
+  let ai: Genkit;
+
+  beforeEach(() => {
+    ai = genkit({});
+    defineEchoModel(ai);
+  });
+
+  it('lets you define and use a custom output format', async () => {
+    ai.defineFormat(
+      {
+        name: 'banana',
+        format: 'banana',
+      },
+      (schema) => {
+        let instructions: string | undefined;
+
+        if (schema) {
+          instructions = `Output should be in banana format`;
+        }
+
+        return {
+          parseChunk: (chunk) => {
+            return `banana: ${chunk.content[0].text}`;
+          },
+
+          parseMessage: (message) => {
+            return `banana: ${message.content[0].text}`;
+          },
+
+          instructions,
+        };
+      }
+    );
+
+    const { output } = await ai.generate({
+      model: 'echoModel',
+      prompt: 'hi',
+      output: { format: 'banana' },
+    });
+
+    assert.strictEqual(output, 'banana: Echo: hi');
+
+    const { response, stream } = await ai.generateStream({
+      model: 'echoModel',
+      prompt: 'hi',
+      output: { format: 'banana' },
+    });
+    const chunks: string[] = [];
+    for await (const chunk of stream) {
+      chunks.push(`${chunk.output}`);
+    }
+    assert.deepStrictEqual(chunks, ['banana: 3', 'banana: 2', 'banana: 1']);
+    assert.strictEqual((await response).output, 'banana: Echo: hi');
+  });
+});