MDX Chat Completion (#212)

This PR adds `MdxChatCompletion`, which produces MDX output as a string.
It does not handle:
* Auto-healing partial MDX (related:
micromark/micromark-extension-mdx-jsx#10)
* Rendering to the UI

I agree with @zkoch that this also might not be the final API we want.

I added docs for this, but as a standalone, this feature isn't super
interesting – it gets fun when we're actually able to render to the UI.
That can be a follow-up.

To verify that I didn't break the previous `UICompletion` component, I
manually tested the nextjs demo at http://localhost:3000/recipe.

.prettierignore

@@ -1,9 +1,10 @@
 .pnp*
 .yarn*
 .vscode*
-*data*
+packages/examples/src/bakeoff/*/data/*
 dist
 .next
 .docusaurus
-*build*
 packages/docs/docs/api
+packages/docs/build
+packages/create-react-app-demo/build/

packages/ai-jsx/package.json

@@ -4,7 +4,7 @@
   "repository": "fixie-ai/ai-jsx",
   "bugs": "https://github.com/fixie-ai/ai-jsx/issues",
   "homepage": "https://ai-jsx.com",
-  "version": "0.5.14",
+  "version": "0.5.15",
   "volta": {
     "extends": "../../package.json"
   },
@@ -245,6 +245,15 @@
         "default": "./dist/cjs/react/completion.cjs"
       }
     },
+    "./react/jit-ui/mdx": {
+      "import": {
+        "types": "./dist/esm/react/jit-ui/mdx.d.ts",
+        "default": "./dist/esm/react/jit-ui/mdx.js"
+      },
+      "require": {
+        "default": "./dist/cjs/react/jit-ui/mdx.cjs"
+      }
+    },
     "./experimental/next": {
       "import": {
         "types": "./dist/esm/experimental/next/index.d.ts",

packages/ai-jsx/src/react/completion.tsx

@@ -7,44 +7,13 @@ import { isJsxBoundary } from './jsx-boundary.js';
 import { AIJSXError, ErrorCode } from '../core/errors.js';
 import z from 'zod';
 
-function reactComponentName(component: React.JSXElementConstructor<any> | string) {
-  return typeof component === 'string' ? component : component.name;
-}
-
 export async function* UICompletion(
   { example, children }: { example: React.ReactNode; children: AI.Node },
   { render, logger }: AI.ComponentContext
 ) {
   yield '';
-  const reactComponents = new Set<React.JSXElementConstructor<any> | string>();
-  function collectComponents(node: React.ReactNode | AI.Node, inReact: boolean) {
-    if (Array.isArray(node)) {
-      node.forEach((node) => collectComponents(node, inReact));
-    }
-
-    if (React.isValidElement(node)) {
-      if (inReact) {
-        reactComponents.add(node.type);
-      }
-
-      const childrenAreReact = (inReact || node.type === AI.React) && !isJsxBoundary(node.type);
-      if ('children' in node.props) {
-        collectComponents(node.props.children, childrenAreReact);
-      }
-    }
-
-    if (AI.isElement(node)) {
-      const childrenAreReact = (inReact || node.tag === AI.React) && !isJsxBoundary(node.tag);
-      if ('children' in node.props) {
-        collectComponents(node.props.children, childrenAreReact);
-      }
-    }
-  }
-
-  collectComponents(example, true);
-
-  const validComponents = Object.fromEntries(Array.from(reactComponents).map((c) => [reactComponentName(c), c]));
 
+  const validComponents = collectComponents(example);
   interface SerializedComponent {
     name: string;
     children: (string | SerializedComponent)[];
@@ -115,3 +84,40 @@ export async function* UICompletion(
     yield component;
   }
 }
+
+/**
+ * Bug: this function doesn't handle symbols, so if `Component` is a fragment, it'll return
+ * `undefined`.
+ */
+function reactComponentName(component: React.JSXElementConstructor<any> | string) {
+  return typeof component === 'string' ? component : component.name;
+}
+
+export function collectComponents(node: React.ReactNode | AI.Node) {
+  const reactComponents = new Set<React.JSXElementConstructor<any> | string>();
+  function collectComponentsRec(node: React.ReactNode | AI.Node, inReact: boolean) {
+    if (Array.isArray(node)) {
+      node.forEach((node) => collectComponentsRec(node, inReact));
+    }
+
+    if (React.isValidElement(node)) {
+      if (inReact) {
+        reactComponents.add(node.type);
+      }
+
+      const childrenAreReact = (inReact || node.type === AI.React) && !isJsxBoundary(node.type);
+      if ('children' in node.props) {
+        collectComponentsRec(node.props.children, childrenAreReact);
+      }
+    }
+
+    if (AI.isElement(node)) {
+      const childrenAreReact = (inReact || node.tag === AI.React) && !isJsxBoundary(node.tag);
+      if ('children' in node.props) {
+        collectComponentsRec(node.props.children, childrenAreReact);
+      }
+    }
+  }
+  collectComponentsRec(node, true);
+  return Object.fromEntries(Array.from(reactComponents).map((c) => [reactComponentName(c), c]));
+}

packages/ai-jsx/src/react/jit-ui/mdx.tsx

@@ -0,0 +1,115 @@
+import * as AI from '../core.js';
+import { ChatCompletion, SystemMessage } from '../../core/completion.js';
+import React from 'react';
+import { collectComponents } from '../completion.js';
+
+/**
+ * Use GPT-4 with this.
+ */
+export function MdxChatCompletion({ children, usageExamples }: { children: AI.Node; usageExamples: React.ReactNode }) {
+  const components = collectComponents(usageExamples);
+  /* prettier-ignore */
+  return <ChatCompletion>
+    <SystemMessage>
+      You are an assistant who can use React components to work with the user. By default, you use markdown. However, if it's useful, you can also mix in the following React components: {Object.keys(components).join(', ')}.
+      All your responses
+      should be in MDX, which is Markdown For the Component Era. Here are instructions for how to use MDX:
+      === Begin instructions
+      {/* Snipped from https://github.com/mdx-js/mdx/blob/main/docs/docs/what-is-mdx.server.mdx. */}
+      MDX allows you to use JSX in your markdown content.
+      You can import components, such as interactive charts or alerts, and embed them
+      within your content.
+      This makes writing long-form content with components a blast.
+
+      More practically MDX can be explained as a format that combines markdown with
+      JSX and looks as follows:
+
+      === Begin example
+      {`
+        Here is some markdown text 
+        <MyComponent id="123" />
+
+        # Here is more markdown text
+
+        <Component
+          open
+          x={1}
+          label={'this is a string, *not* markdown!'}
+          icon={<Icon />}
+        />
+
+        * Markdown list item 1
+        * Markdown list item 2
+        * Markdown list item 3
+      `}
+      === end example
+      === end instructions
+
+      Do not include a starting ```mdx and closing ``` line. Just respond with the MDX itself.
+
+      Do not include extra whitespace that is not needed for the markdown interpretation. For instance, if your component has a prop that's a JSON object, put it all on one line:
+
+      {'<Component prop={[[{"key": "value"}, {"long": "field"}]]} />'}
+
+      This doc tells you the differences between MDX and markdown.
+
+      {/* Adapted from https://github.com/micromark/mdx-state-machine#72-deviations-from-markdown */}
+      === Start doc
+      ### 7.2 Deviations from Markdown
+
+      MDX adds constructs to Markdown but also prohibits certain normal Markdown
+      constructs.
+
+      #### 7.2.2 Indented code
+
+      Indentation to create code blocks is not supported.
+      Instead, use fenced code blocks.
+
+      The reason for this change is so that elements can be indented.
+
+      {/* Commenting out the negative examples because they seem to confuse the LLM. */}
+      {/*
+      Incorrect:
+
+      ```markdown
+          console.log(1)
+      ``` */}
+
+      Correct:
+
+        ```js
+        console.log(1)
+        ```
+
+      #### 7.2.3 Autolinks
+
+      Autolinks are not supported.
+      Instead, use links or references.
+
+      The reason for this change is because whether something is an element (whether
+      HTML or JSX) or an autolink is ambiguous {'(Markdown normally treats `<svg:rect>`, `<xml:lang/>`, or `<svg:circle{...props}>` as links).'}
+
+      {/* Incorrect:
+
+      ```markdown
+      See <https://example.com> for more information
+      ``` */}
+
+      Correct:
+
+        See [example.com](https://example.com) for more information.
+
+      #### 7.2.4 Errors
+
+      Whereas all Markdown is valid, incorrect MDX will crash.
+      === end doc
+
+      Here are the components you have available, and how to use them:
+
+      === Begin components
+      <AI.React>{usageExamples}</AI.React>
+      === end components
+    </SystemMessage>
+    {children}
+  </ChatCompletion>;
+}

packages/docs/docs/changelog.md

@@ -1,6 +1,10 @@
 # Changelog
 
-## 0.5.14
+## 0.5.15
+
+- Add [`MdxChatCompletion`](./guides/mdx.md), so your model calls can now output [MDX](https://mdxjs.com/) using your components.
+
+## [0.5.14](https://github.com/fixie-ai/ai-jsx/commit/5971243)
 
 - Add [Llama2 support](./guides/models.md#llama2).
 

packages/docs/docs/guides/mdx.md

@@ -0,0 +1,81 @@
+# MDX Output
+
+By default, models emit text. If you ask them to, they'll emit markdown, which is a robust, easy way to provide more structured output to your user.
+
+To take it a step further, you can have the model emit [MDX](https://mdxjs.com/), which combines Markdown and your components:
+
+```mdx
+# Last year’s snowfall
+
+In {year}, the snowfall was above average.
+It was followed by a warm spring which caused
+flood conditions in many of the nearby rivers.
+
+<Chart year={year} color="#fcb32c" />
+```
+
+To do this, use the `MdxChatCompletion` component:
+
+```tsx
+<MdxChatCompletion usageExamples={usageExample}>
+  <UserMessage>Tell me a children's story. Summarize the key characters at the end.</UserMessage>
+</MdxChatCompletion>
+```
+
+[(See the examples project for a full working demo.)](https://github.com/fixie-ai/ai-jsx/blob/main/packages/examples/src/mdx.tsx)
+
+## Telling the model which components are available
+
+The API for `MdxChatCompletion` is the same as `ChatCompletion`, except it adds a `usageExamples` prop. That prop is an AI.JSX component that tells the model which components are available:
+
+<!-- prettier-ignore -->
+```tsx
+const usageExample = (
+  <>
+    Use a Card to display collected information to the user. The children can be markdown. Only use the card if you have
+    a logically-grouped set of information to show the user, in the context of a larger response. Generally, your entire
+    response should not be a card. A card takes optional header and footer props.
+
+    Example 1 of how you might use this
+    component: Here's the best candidate I found:
+    <Card header="Sam Smith">
+      **Skills**: React, TypeScript, Node.js **Location**: Seattle, WA **Years of experience**: 5 **Availability**:
+      Full-time
+    </Card>
+    {/* ... you may wish to add more examples */}
+  </>
+);
+
+<MdxChatCompletion usageExamples={usageExample}>
+  <UserMessage>Tell me a children's story. Summarize the key characters at the end.</UserMessage>
+</MdxChatCompletion>
+```
+
+In that example, `Card` may refer to a real component in scope, or you can just write it out as a string:
+
+<!-- prettier-ignore -->
+```tsx
+function Card({ header, footer, children }) {
+  return (
+    <div>
+      <div className="header">{header}</div>
+      <div className="content">{children}</div>
+      <div className="footer">{footer}</div>
+    </div>
+  );
+}
+
+const usageExample = (
+  <>
+    {/* Reference a component in scope */}
+    Here's how you use a Card: <Card>content</Card>
+
+    {/* Just write out a string */}
+    You can also use a ButtonGroup: {`<ButtonGroup labels=['Yes', 'No'] />`}
+  </>
+);
+```
+
+## Using the output
+
+The output will be MDX, as a string. It will not have any import statements. If you want to render it into your UI, it's up to you to parse/compile it in some way, using the [MDX APIs](https://mdxjs.com/).

packages/examples/.eslintrc.cjs

@@ -26,7 +26,7 @@ module.exports = {
     'no-use-before-define': 'off',
     '@typescript-eslint/no-use-before-define': ['error', { functions: false, variables: true }],
 
-    'no-trailing-spaces': 'warn',
+    'no-trailing-spaces': 'off',
     'no-else-return': ['warn', { allowElseIf: false }],
 
     // Disable style rules to let prettier own it

packages/examples/package.json

@@ -35,6 +35,7 @@
     "demo:multi-model-chat": "yarn build && node dist/multi-model-chat.js",
     "demo:simple-completion": "yarn build && node dist/simple-completion.js",
     "demo:getting-started": "yarn build && cp src/getting-started/data.json dist/getting-started/ && node dist/getting-started/index.js",
+    "demo:mdx-jit": "yarn build && node dist/mdx.js",
     "demo:inspector": "yarn build && node dist/inspector.js",
     "demo:router": "yarn build && node dist/router.js",
     "demo:replicate-llama2": "yarn build && node dist/replicate-llama2.js",