MCP Terminal Server

MIT License
Overview InspectNew Schema Related Servers Reviews Score
js
ai
tests
generate
/**
 * 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 { stripUndefinedProps, z } from '@genkit-ai/core';
import { Registry } from '@genkit-ai/core/registry';
import * as assert from 'assert';
import { readFileSync } from 'fs';
import { beforeEach, describe, it } from 'node:test';
import { parse } from 'yaml';
import {
  GenerateAction,
  defineGenerateAction,
} from '../../src/generate/action.js';
import {
  GenerateActionOptionsSchema,
  GenerateResponseChunkData,
  GenerateResponseChunkSchema,
  GenerateResponseSchema,
} from '../../src/model.js';
import { defineTool } from '../../src/tool.js';
import { ProgrammableModel, defineProgrammableModel } from '../helpers.js';

const SpecSuiteSchema = z
  .object({
    tests: z.array(
      z
        .object({
          name: z.string(),
          input: GenerateActionOptionsSchema,
          streamChunks: z
            .array(z.array(GenerateResponseChunkSchema))
            .optional(),
          modelResponses: z.array(GenerateResponseSchema),
          expectResponse: GenerateResponseSchema.optional(),
          stream: z.boolean().optional(),
          expectChunks: z.array(GenerateResponseChunkSchema).optional(),
        })
        .strict()
    ),
  })
  .strict();

describe('spec', () => {
  let registry: Registry;
  let pm: ProgrammableModel;

  beforeEach(() => {
    registry = new Registry();
    defineGenerateAction(registry);
    pm = defineProgrammableModel(registry);
    defineTool(
      registry,
      { name: 'testTool', description: 'description' },
      async () => 'tool called'
    );
  });

  SpecSuiteSchema.parse(
    parse(readFileSync('../../tests/specs/generate.yaml', 'utf-8'))
  ).tests.forEach((test) => {
    it(test.name, async () => {
      if (test.modelResponses || test.streamChunks) {
        let reqCounter = 0;
        pm.handleResponse = async (req, sc) => {
          if (test.streamChunks && sc) {
            test.streamChunks[reqCounter].forEach(sc);
          }
          return test.modelResponses?.[reqCounter++]!;
        };
      }
      const action = (await registry.lookupAction(
        '/util/generate'
      )) as GenerateAction;

      if (test.stream) {
        const { output, stream } = action.stream(test.input);

        const chunks = [] as GenerateResponseChunkData[];
        for await (const chunk of stream) {
          chunks.push(stripUndefinedProps(chunk));
        }

        assert.deepStrictEqual(chunks, test.expectChunks);

        assert.deepStrictEqual(
          stripUndefinedProps(await output),
          test.expectResponse
        );
      } else {
        const response = await action(test.input);

        assert.deepStrictEqual(
          stripUndefinedProps(response),
          test.expectResponse
        );
      }
    });
  });
});