Mock genkit streaming responses in Storybook

README.md

Genkit Streaming in Storybook

This is my current best attempt to allow iterating on the user experience of a UI that loads streaming responses from an LLM and may have some interesting intermediate states while the response trickles in.

Component

• https://genkit.dev/docs/client/

import { streamFlow } from 'genkit/beta/client';

// [...]

const flow = streamFlow<FlowOutputType>({
  url: '/api/my-flow',
  input: flowInput,
});

Stories

Using msw

Plugin: https://storybook.js.org/addons/msw-storybook-addon

import {http} from 'msw';

import {
  createGenkitStreamingResponse,
  resetDefaultTokenDelay,
  setDefaultTokenDelay,
} from './mock-genkit.ts';

export const Example: Story = {
  args: {
    // [General arguments...]

    // Make the delay between tokens configurable to explore the exact UX.
    tokenDelay: 100,
  },
  beforeEach: (context) => {
    // It's tough to pass args via the msw storybook plugin, so this just uses
    // global state that gets reset afterwards.
    setDefaultTokenDelay(context.args.tokenDelay ?? 0);
    return () => {
      resetDefaultTokenDelay();
    };
  },
  parameters: {
    msw: {
      handlers: {
        chat: http.post('/api/my-flow', async ({request}) => {
          const jsonBody = await request.json();
          const {
            data: flowInput,
          } = jsonBody as {data: FlowInputType};

          const flowOutput: FlowOutputType = getMockResponse(flowInput);
          return createGenkitStreamingResponse(flowOutput);
        }),
      },
    },
  },
};

Using fetchMock

Uses manual setup because the storybook plugin for fetch-mock doesn't support Storybook v9 yet as of today.

import fetchMock from 'fetch-mock';

import {createGenkitStreamingResponse} from './mock-genkit.ts';

export const Example: Story = {
  args: {
    // [General arguments...]

    // Make the delay between tokens configurable to explore the exact UX.
    tokenDelay: 100,
  },
  beforeEach: (context) => {
    const tokenDelay = context.args.tokenDelay ?? 0;
    // fetchMock.config.allowRelativeUrls = true;
    fetchMock.mockGlobal();
    fetchMock.post(`${location.origin}/api/my-flow`, async ({options}) => {
      if (typeof options.body !== 'string') return 500;
      const {
        data: flowInput
      } = JSON.parse(options.body) as {data: FlowInputType};
      const flowOutput: FlowOutputType = getMockResponse(flowInput);
      return createGenkitStreamingResponse(flowOutput, tokenDelay);
    });
    return () => {
      fetchMock.hardReset();
    };
  },
};

mock-genkit.ts

const INIITIAL_DEFAULT_TOKEN_DELAY = 0;

let defaultTokenDelay = INIITIAL_DEFAULT_TOKEN_DELAY;

export function setDefaultTokenDelay(delay: number) {
  defaultTokenDelay = delay;
}

export function resetDefaultTokenDelay() {
  defaultTokenDelay = INIITIAL_DEFAULT_TOKEN_DELAY;
}

function toIncompleteRegExp(str: string): RegExp {
  return new RegExp(
    `\\b${str
      .slice(0, -2)
      .split('')
      .reduceRight((acc, cur) => {
        // Match the current character potentially followed by the suffix.
        return `(?:${cur}${acc}?)`;
      }, str.at(-2))}$`,
  );
}

const JSON_KEYWORDS = ['true', 'false', 'null'].map((token) => ({
  pattern: toIncompleteRegExp(token),
  token,
}));

// There's likely properly tested libraries for this but the ones I found
// immediately seemed much too large for a quick storybook hack.
function repairJSONObject(partialJSON: string): string {
  // Find missing trailing quotes and braces in the given JSON string.
  // Return once the result is valid JSON.
  let repaired = partialJSON;
  const maxAttempts = 10;
  let attempts = 0;

  while (attempts < maxAttempts) {
    try {
      JSON.parse(repaired);
      return repaired; // Valid JSON
    } catch (e) {
      if (!(e instanceof SyntaxError)) {
        throw e;
      }
      if (e.message.includes('Unterminated string')) {
        repaired += '"';
      } else if (e.message.includes('after array element')) {
        repaired += ']';
      } else if (e.message.includes("'}'")) {
        repaired += '}';
      } else if (e.message === 'Unexpected end of JSON input') {
        const keyword = JSON_KEYWORDS.find(kw => kw.pattern.test(repaired));
        if (keyword) {
          // 'tru', 'fa', 'nul', etc.
          repaired = repaired.replace(keyword.pattern, keyword.token);
        } else if (repaired.endsWith(',')) {
          // '[1,2,'
          repaired = repaired.slice(0, -1) + ']';
        } else {
        // '[1,2' or '['
          repaired += ']';
        }
      } else if (e.message.includes('Expected double-quoted property name')) {
        // '{"x":2,'
        if (repaired.endsWith(',')) {
          repaired = repaired.slice(0, -1);
        }
        repaired += '}';
      } else if (e.message.includes("':'")) {
        repaired += ':null}';
      } else {
        console.error(e);
        break; // Unknown error, stop trying
      }
      attempts++;
    }
  }
  return '{}'; // Return a minimally valid version.
}

export function createGenkitStreamingResponse<T>(result: T, tokenDelay: number = defaultTokenDelay) {
  const resultJSON = JSON.stringify(result);
  const contentTokens = resultJSON.split(' ');
  const encoder = new TextEncoder();
  const stream = new ReadableStream({
    start(controller) {
      let contentLength = tokenDelay > 0 ? 0 : contentTokens.length;
      function sendNextChunk() {
        ++contentLength;
        if (contentLength >= contentTokens.length) {
          controller.enqueue(encoder.encode(`data: {"result":${resultJSON}}\n\n`));
          controller.close();
        } else {
          const partialJSON = contentTokens.slice(0, contentLength).join(' ');
          const repairedJSON = repairJSONObject(partialJSON);
          controller.enqueue(encoder.encode(`data: {"message":${repairedJSON}}\n\n`));
          setTimeout(sendNextChunk, tokenDelay);
        }
      }
      setTimeout(sendNextChunk, tokenDelay);
    },
  });
  return new Response(stream, {
    status: 200,
    headers: {
      'Content-Type': 'text/event-stream',
    },
  });
}