feat: AI Chat UI (#25)

Author: ChristopherChudzicki

.eslintrc.js

@@ -46,6 +46,7 @@ module.exports = {
           "**/*.test.tsx",
           "**/src/setupJest.ts",
           "**/jest-setup.ts",
+          "**/jsdom-extended.ts",
           "**/test-utils/**",
           "**/test-utils/**",
           "**/webpack.config.js",

.storybook/main.ts

@@ -1,5 +1,13 @@
 import { StorybookConfig } from "@storybook/react-webpack5"
 import TsconfigPathsPlugin from "tsconfig-paths-webpack-plugin"
+import { exec as execCb } from "child_process"
+import { promisify } from "util"
+
+const exec = promisify(execCb)
+const getGitSha = async (): Promise<string> => {
+  const { stdout } = await exec("git rev-parse HEAD")
+  return stdout.trim()
+}
 
 const config: StorybookConfig = {
   stories: ["../src/**/*.mdx", "../src/**/*.stories.tsx"],
@@ -26,6 +34,11 @@ const config: StorybookConfig = {
   typescript: {
     reactDocgen: "react-docgen-typescript",
   },
+  env: async () => {
+    return {
+      STORYBOOK_GIT_SHA: await getGitSha(),
+    }
+  },
 }
 
 export default config

jest.config.ts

@@ -9,7 +9,7 @@ const config: Config.InitialOptions = {
     "jest-watch-typeahead/testname",
   ],
   setupFilesAfterEnv: ["./jest-setup.ts"],
-  testEnvironment: "jsdom",
+  testEnvironment: "<rootDir>/jsdom-extended.ts",
   transform: {
     "^.+\\.(t|j)sx?$": "@swc/jest",
   },

package.json

@@ -39,6 +39,20 @@
         "default": "./dist/esm/index.js"
       }
     },
+    "./ai": {
+      "import": {
+        "types": "./dist/esm/ai.d.ts",
+        "default": "./dist/esm/ai.js"
+      },
+      "require": {
+        "types": "./dist/cjs/ai.d.ts",
+        "default": "./dist/cjs/ai.js"
+      },
+      "default": {
+        "types": "./dist/esm/ai.d.ts",
+        "default": "./dist/esm/ai.js"
+      }
+    },
     "./type-augmentation": {
       "import": "./dist/type-augmentation/index.d.ts",
       "require": "./dist/type-augmentation/index.d.ts",
@@ -55,27 +69,30 @@
     "@mui/system": "^6.1.6",
     "@remixicon/react": "^4.2.0",
     "@types/jest": "^29.5.14",
+    "ai": "^4.0.13",
     "classnames": "^2.5.1",
     "lodash": "^4.17.21",
-    "material-ui-popup-state": "^5.1.0",
-    "tiny-invariant": "^1.3.1"
+    "react-markdown": "^9.0.1",
+    "tiny-invariant": "^1.3.1",
+    "zod": "^3.23.8"
   },
   "devDependencies": {
     "@chromatic-com/storybook": "^1.9.0",
     "@faker-js/faker": "^9.0.0",
-    "@storybook/addon-actions": "^8.4.2",
-    "@storybook/addon-essentials": "^8.4.2",
-    "@storybook/addon-interactions": "^8.4.2",
-    "@storybook/addon-links": "^8.4.2",
-    "@storybook/addon-onboarding": "^8.4.2",
+    "@jest/environment": "^29.7.0",
+    "@storybook/addon-actions": "^8.4.7",
+    "@storybook/addon-essentials": "^8.4.7",
+    "@storybook/addon-interactions": "^8.4.7",
+    "@storybook/addon-links": "^8.4.7",
+    "@storybook/addon-onboarding": "^8.4.7",
     "@storybook/addon-webpack5-compiler-swc": "^1.0.5",
-    "@storybook/blocks": "^8.4.2",
-    "@storybook/nextjs": "^8.4.2",
-    "@storybook/preview-api": "^8.4.2",
-    "@storybook/react": "^8.4.2",
-    "@storybook/react-webpack5": "^8.4.2",
-    "@storybook/test": "^8.4.2",
-    "@storybook/types": "^8.4.2",
+    "@storybook/blocks": "^8.4.7",
+    "@storybook/nextjs": "^8.4.7",
+    "@storybook/preview-api": "^8.4.7",
+    "@storybook/react": "^8.4.7",
+    "@storybook/react-webpack5": "^8.4.7",
+    "@storybook/test": "^8.4.7",
+    "@storybook/types": "^8.4.7",
     "@swc/jest": "^0.2.37",
     "@testing-library/jest-dom": "^6.6.3",
     "@testing-library/react": "^16.0.1",
@@ -99,7 +116,7 @@
     "jest": "^29.7.0",
     "jest-environment-jsdom": "^29.5.0",
     "jest-extended": "^4.0.2",
-    "jest-fail-on-console": "^3.2.0",
+    "jest-fail-on-console": "^3.3.1",
     "jest-watch-typeahead": "^2.2.2",
     "next": "^15.0.2",
     "prettier": "^3.3.3",

src/ai.ts

@@ -0,0 +1,2 @@
+export { AiChat } from "./components/AiChat/AiChat"
+export type { AiChatProps } from "./components/AiChat/AiChat"

src/components/AiChat/AiChat.mdx

@@ -0,0 +1,25 @@
+import { Meta, Title, Primary, Controls, Stories } from "@storybook/blocks"
+
+import * as AiChat from "./AiChat.stories"
+import { gitLink } from "../../story-utils"
+
+<Meta of={AiChat} />
+
+<Title />
+
+Exported from `smoot-design/ai`, the AiChat component is a chat interface
+for use with AI services. It can be used with text-streaming or JSON APIs.
+
+This demo shows the AiChat component with a simple text-streaming API.
+
+<Primary />
+
+## Inputs
+
+The component accepts the following inputs (props):
+
+<Controls />
+
+See <a href={gitLink("src/components/AiChat/types.ts")}>AiChat/types.ts</a> for all Typescript interface definitions.
+
+<Stories includePrimary={false} />

src/components/AiChat/AiChat.stories.tsx

@@ -0,0 +1,87 @@
+import * as React from "react"
+import type { Meta, StoryObj } from "@storybook/react"
+import { AiChat } from "./AiChat"
+import type { AiChatProps } from "./types"
+import { mockJson, mockStreaming } from "./story-utils"
+import styled from "@emotion/styled"
+
+const TEST_API_STREAMING = "http://localhost:4567/streaming"
+const TEST_API_JSON = "http://localhost:4567/json"
+
+const INITIAL_MESSAGES: AiChatProps["initialMessages"] = [
+  {
+    content: "Hi! What are you interested in learning about?",
+    role: "assistant",
+  },
+]
+
+const STARTERS = [
+  { content: "I'm interested in quantum computing" },
+  { content: "I want to understand global warming. " },
+  { content: "I am curious about AI applications for business" },
+]
+
+const Container = styled.div({
+  width: "100%",
+  height: "350px",
+})
+
+const meta: Meta<typeof AiChat> = {
+  title: "smoot-design/AiChat",
+  component: AiChat,
+  render: (args) => <AiChat {...args} />,
+  decorators: (Story) => {
+    return (
+      <Container>
+        <Story />
+      </Container>
+    )
+  },
+  args: {
+    initialMessages: INITIAL_MESSAGES,
+    requestOpts: { apiUrl: TEST_API_STREAMING },
+    conversationStarters: STARTERS,
+  },
+  argTypes: {
+    conversationStarters: {
+      control: { type: "object", disable: true },
+    },
+    initialMessages: {
+      control: { type: "object", disable: true },
+    },
+    requestOpts: {
+      control: { type: "object", disable: true },
+      table: { readonly: true }, // See above
+    },
+  },
+  beforeEach: () => {
+    const originalFetch = window.fetch
+    window.fetch = (url, opts) => {
+      if (url === TEST_API_STREAMING) {
+        return mockStreaming()
+      } else if (url === TEST_API_JSON) {
+        return mockJson()
+      }
+      return originalFetch(url, opts)
+    }
+  },
+}
+
+export default meta
+
+type Story = StoryObj<typeof AiChat>
+
+export const StreamingResponses: Story = {}
+
+/**
+ * Here `AiChat` is used with a non-streaming JSON API. The JSON is converted
+ * to text via `parseContent`.
+ */
+export const JsonResponses: Story = {
+  args: {
+    requestOpts: { apiUrl: TEST_API_JSON },
+    parseContent: (content: unknown) => {
+      return JSON.parse(content as string).message
+    },
+  },
+}

src/components/AiChat/AiChat.test.tsx

@@ -0,0 +1,162 @@
+// This was giving false positives
+/* eslint-disable testing-library/await-async-utils */
+import { render, screen, waitFor } from "@testing-library/react"
+import user from "@testing-library/user-event"
+import { AiChat } from "./AiChat"
+import { ThemeProvider } from "../ThemeProvider/ThemeProvider"
+import * as React from "react"
+import { AiChatProps } from "./types"
+import { faker } from "@faker-js/faker/locale/en"
+
+const counter = jest.fn() // use jest.fn as counter because it resets on each test
+const mockFetch = jest.mocked(
+  jest.fn(() => {
+    const count = counter.mock.calls.length
+    counter()
+    return Promise.resolve(
+      new Response(`AI Response ${count}`, {
+        headers: {
+          "Content-Type": "application/json",
+        },
+      }),
+    )
+  }) as typeof fetch,
+)
+window.fetch = mockFetch
+jest.mock("react-markdown", () => {
+  return {
+    __esModule: true,
+    default: ({ children }: { children: string }) => <div>{children}</div>,
+  }
+})
+
+const getMessages = (): HTMLElement[] => {
+  return Array.from(document.querySelectorAll(".MitAiChat--message"))
+}
+const getConversationStarters = (): HTMLElement[] => {
+  return Array.from(
+    document.querySelectorAll("button.MitAiChat--conversationStarter"),
+  )
+}
+const whenCount = async <T,>(cb: () => T[], count: number) => {
+  return await waitFor(() => {
+    const result = cb()
+    expect(result).toHaveLength(count)
+    return result
+  })
+}
+
+describe("AiChat", () => {
+  const setup = (props: Partial<AiChatProps> = {}) => {
+    const initialMessages: AiChatProps["initialMessages"] = [
+      { role: "assistant", content: faker.lorem.sentence() },
+    ]
+    const conversationStarters: AiChatProps["conversationStarters"] = [
+      { content: faker.lorem.sentence() },
+      { content: faker.lorem.sentence() },
+    ]
+    render(
+      <AiChat
+        initialMessages={initialMessages}
+        conversationStarters={conversationStarters}
+        requestOpts={{ apiUrl: "http://localhost:4567/test" }}
+        {...props}
+      />,
+      { wrapper: ThemeProvider },
+    )
+
+    return { initialMessages, conversationStarters }
+  }
+
+  test("Clicking conversation starters and sending chats", async () => {
+    const { initialMessages, conversationStarters } = setup()
+
+    const scrollBy = jest.spyOn(HTMLElement.prototype, "scrollBy")
+
+    const initialMessageEls = getMessages()
+    expect(initialMessageEls.length).toBe(1)
+    expect(initialMessageEls[0]).toHaveTextContent(initialMessages[0].content)
+
+    const starterEls = getConversationStarters()
+    expect(starterEls.length).toBe(2)
+    expect(starterEls[0]).toHaveTextContent(conversationStarters[0].content)
+    expect(starterEls[1]).toHaveTextContent(conversationStarters[1].content)
+
+    const chosen = faker.helpers.arrayElement([0, 1])
+
+    await user.click(starterEls[chosen])
+    expect(scrollBy).toHaveBeenCalled()
+    scrollBy.mockReset()
+
+    const messageEls = await whenCount(getMessages, 3)
+
+    expect(messageEls[0]).toHaveTextContent(initialMessages[0].content)
+    expect(messageEls[1]).toHaveTextContent(
+      conversationStarters[chosen].content,
+    )
+    expect(messageEls[2]).toHaveTextContent("AI Response 0")
+
+    await user.click(screen.getByPlaceholderText("Type a message..."))
+    await user.paste("User message")
+    await user.click(screen.getByRole("button", { name: "Send" }))
+    expect(scrollBy).toHaveBeenCalled()
+
+    const afterSending = await whenCount(getMessages, 5)
+    expect(afterSending[3]).toHaveTextContent("User message")
+    expect(afterSending[4]).toHaveTextContent("AI Response 1")
+  })
+
+  test("transformBody is called before sending requests", async () => {
+    const fakeBody = { message: faker.lorem.sentence() }
+    const apiUrl = faker.internet.url()
+    const transformBody = jest.fn(() => fakeBody)
+    const { initialMessages } = setup({
+      requestOpts: { apiUrl, transformBody },
+    })
+
+    await user.click(screen.getByPlaceholderText("Type a message..."))
+    await user.paste("User message")
+    await user.click(screen.getByRole("button", { name: "Send" }))
+
+    expect(transformBody).toHaveBeenCalledWith([
+      expect.objectContaining(initialMessages[0]),
+      expect.objectContaining({ content: "User message", role: "user" }),
+    ])
+    expect(mockFetch).toHaveBeenCalledTimes(1)
+    expect(mockFetch).toHaveBeenCalledWith(
+      apiUrl,
+      expect.objectContaining({
+        body: JSON.stringify(fakeBody),
+      }),
+    )
+  })
+
+  test("parseContent is called on the API-received message content", async () => {
+    const fakeBody = { message: faker.lorem.sentence() }
+    const apiUrl = faker.internet.url()
+    const transformBody = jest.fn(() => fakeBody)
+    const { initialMessages, conversationStarters } = setup({
+      requestOpts: { apiUrl, transformBody },
+      parseContent: jest.fn((content) => `Parsed: ${content}`),
+    })
+
+    await user.click(getConversationStarters()[0])
+
+    await whenCount(getMessages, initialMessages.length + 2)
+
+    await user.click(screen.getByPlaceholderText("Type a message..."))
+    await user.paste("User message")
+    await user.click(screen.getByRole("button", { name: "Send" }))
+
+    await whenCount(getMessages, initialMessages.length + 4)
+
+    const messagesTexts = getMessages().map((el) => el.textContent)
+    expect(messagesTexts).toEqual([
+      initialMessages[0].content,
+      conversationStarters[0].content,
+      "Parsed: AI Response 0",
+      "User message",
+      "Parsed: AI Response 1",
+    ])
+  })
+})