Hydrate MDX (#220)

![image](https://github.com/fixie-ai/ai-jsx/assets/829827/ee36980d-555f-400d-9f91-0e0d747c4006)


With this approach, we compile MDX on the client.

Pros:
* It's easy to get the components (e.g. `Card`) in scope – no need to
pass them up and down to the server.
* Easy for the client to otherwise customize the Markdown rendering as
they would like.
* Maintains support for the text streaming model, allowing us to stay
closer to the Vercel `useChat` paved path.
* Enables other AI.JSX components to render UI simply by emitting MDX.
(Of course, we need to ensure those emitted components are in scope at
compile time.)
* I suspect that the MDX compiler/runtime may be published in a way that
causes trouble for CJS importers. By keeping the MDX usage to
"user-land", or at least on the client, if there is a problem here, we
minimize the impact.

Potential objections
* Performance of doing `O(count of stream chunks)` compiles – I don't
think this will be noticeable, particularly in context of LLM response
times.

Future work
* Devising a scheme for the user to interact with the components – e.g.
if the user clicks a button or fills out a form, how do we communicate
that to the AI?
* AI does a decent but not amazing job of adhering to the spec – for
instance, all my few-shot examples say `Badge` needs a `color` prop
(e.g. `<Badge color="yellow">In progress</Badge>`), but the model does
not always do that. There may be more prompt engineering work we can do
here.
* The demo itself in the nextjs project isn't super compelling. I would
rather focus that demo energy on HS.
* Sometimes the model still emits \`\`\`mdx blocks wrapping large parts
of its response.
* Sometimes the model uses `<details>` and `<summary>` in a way that
causes the compiler to produce `<details><p><summary>,` which is
invalid.

I think this approach will work reasonable well for RSC / Architecture
4, with some modification.

Once we align on this approach, I'll add docs.

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.7.3",
+  "version": "0.8.0",
   "volta": {
     "extends": "../../package.json"
   },

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

@@ -119,5 +119,12 @@ export function collectComponents(node: React.ReactNode | AI.Node) {
     }
   }
   collectComponentsRec(node, true);
-  return Object.fromEntries(Array.from(reactComponents).map((c) => [reactComponentName(c), c]));
+  return Object.fromEntries(
+    Array.from(reactComponents)
+      /* Filter out symbols, e.g. React.Fragment.
+        `usageExample` will often be passed as a Fragment, and we don't want to tell the AI to output fragments.
+      */
+      .filter((component) => typeof component !== 'symbol')
+      .map((c) => [reactComponentName(c), c])
+  );
 }

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

@@ -1,16 +1,28 @@
+/** @jsxImportSource ai-jsx/react */
+
 import * as AI from '../core.js';
-import { ChatCompletion, SystemMessage } from '../../core/completion.js';
+import { SystemMessage } from '../../core/completion.js';
 import React from 'react';
 import { collectComponents } from '../completion.js';
 
 /**
- * Use GPT-4 with this.
+ * A completion component that emits [MDX](https://mdxjs.com/).
+ *
+ * By default, the result streamed out of this component will sometimes be unparsable, as the model emits a partial value.
+ * (For instance, if the model is emitting the string `foo <Bar />`, and
+ * it streams out `foo <Ba`, that's not parsable.)
+ *
+ * You'll get better results with this if you use GPT-4.
+ *
+ * Use `usageExamples` to teach the model how to use your components.
+ *
+ * @see https://docs.ai-jsx.com/guides/mdx
+ * @see https://github.com/fixie-ai/ai-jsx/blob/main/packages/examples/src/mdx.tsx
  */
-export function MdxChatCompletion({ children, usageExamples }: { children: AI.Node; usageExamples: React.ReactNode }) {
+export function MdxSystemMessage({ usageExamples }: { usageExamples: React.ReactNode }) {
   const components = collectComponents(usageExamples);
   /* prettier-ignore */
-  return <ChatCompletion>
-    <SystemMessage>
+  return <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:
@@ -26,7 +38,7 @@ export function MdxChatCompletion({ children, usageExamples }: { children: AI.No
 
       === Begin example
       {`
-        Here is some markdown text 
+        Here is some markdown text
         <MyComponent id="123" />
 
         # Here is more markdown text
@@ -109,7 +121,5 @@ export function MdxChatCompletion({ children, usageExamples }: { children: AI.No
       === Begin components
       <AI.React>{usageExamples}</AI.React>
       === end components
-    </SystemMessage>
-    {children}
-  </ChatCompletion>;
+    </SystemMessage>;
 }

packages/docs/docs/changelog.md

@@ -1,6 +1,10 @@
 # Changelog
 
-## 0.7.3
+## 0.8.0
+
+- Move `MdxChatCompletion` to be `MdxSystemMessage`. You can now put this `SystemMessage` in any `ChatCompletion` to prompt the model to give MDX output.
+
+## [0.7.3](https://github.com/fixie-ai/ai-jsx/commit/670ea52647138052cb116cbc56b6cc4bb49512a0)
 
 - Update readme.
 

packages/examples/src/mdx.tsx

@@ -1,14 +1,13 @@
 /** @jsxImportSource ai-jsx/react */
 import * as AI from 'ai-jsx';
-import { SystemMessage, UserMessage } from 'ai-jsx/core/completion';
-// import { showInspector } from 'ai-jsx/core/inspector';
-import { MdxChatCompletion } from 'ai-jsx/react/jit-ui/mdx';
+import { SystemMessage, UserMessage, ChatCompletion } from 'ai-jsx/core/completion';
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+import { showInspector } from 'ai-jsx/core/inspector';
+import { MdxSystemMessage } from 'ai-jsx/react/jit-ui/mdx';
 import { JsonChatCompletion } from 'ai-jsx/batteries/constrained-output';
 import z from 'zod';
 
 import { OpenAI } from 'ai-jsx/lib/openai';
-import { PinoLogger } from 'ai-jsx/core/log';
-import { pino } from 'pino';
 
 /* eslint-disable @typescript-eslint/no-unused-vars */
 function Card({ header, footer, children }: { header?: string; footer?: string; children: string }) {
@@ -115,7 +114,11 @@ function QuestionAndAnswer({ children }: { children: AI.Node }, { memo }: AI.Com
       <OpenAI chatModel="gpt-4">
         Q: {question}
         {'\n'}
-        A: <MdxChatCompletion usageExamples={usageExample}>{question}</MdxChatCompletion>
+        A:{' '}
+        <ChatCompletion>
+          <MdxSystemMessage usageExamples={usageExample} />
+          <UserMessage>{question}</UserMessage>
+        </ChatCompletion>
         {'\n\n'}
       </OpenAI>
     </>
@@ -166,19 +169,8 @@ export function App() {
 
 // showInspector(<App />);
 
-const logger = pino({
-  name: 'ai-jsx',
-  level: process.env.loglevel ?? 'trace',
-  transport: {
-    target: 'pino-pretty',
-    options: {
-      colorize: true,
-    },
-  },
-});
-
 let lastValue = '';
-const rendering = AI.createRenderContext({ logger: new PinoLogger(logger) }).render(<App />, { appendOnly: true });
+const rendering = AI.createRenderContext().render(<App />, { appendOnly: true });
 for await (const frame of rendering) {
   process.stdout.write(frame.slice(lastValue.length));
   lastValue = frame;

packages/nextjs-demo/package.json

@@ -13,6 +13,8 @@
   "dependencies": {
     "@headlessui/react": "^1.7.15",
     "@heroicons/react": "^2.0.18",
+    "@mdx-js/mdx": "^2.3.0",
+    "@mdx-js/react": "^2.3.0",
     "@octokit/graphql": "^5.0.6",
     "@tailwindcss/forms": "^0.5.3",
     "@types/node": "20.2.5",
@@ -29,6 +31,7 @@
     "postcss": "8.4.24",
     "react": "18.2.0",
     "react-dom": "18.2.0",
+    "remark-gfm": "^3.0.1",
     "tailwindcss": "3.3.2",
     "typescript": "^5.1.3"
   },

packages/nextjs-demo/src/app/500.tsx

@@ -0,0 +1,7 @@
+export default function Failure() {
+  return (
+    <div className="flex items-center justify-center w-screen h-screen">
+      <p>500</p>
+    </div>
+  );
+}

packages/nextjs-demo/src/app/building-blocks/api/route.tsx

@@ -0,0 +1,67 @@
+/** @jsxImportSource ai-jsx/react */
+import { NextRequest } from 'next/server';
+import { UserMessage, ChatCompletion } from 'ai-jsx/core/completion';
+import * as BuildingBlocks from '@/components/BuildingBlocks';
+const { Card, ButtonGroup, Badge, Toggle } = BuildingBlocks;
+import { MdxSystemMessage } from 'ai-jsx/react/jit-ui/mdx';
+import { toTextStream } from 'ai-jsx/stream';
+import { Message, StreamingTextResponse } from 'ai';
+import _ from 'lodash';
+import { OpenAI } from 'ai-jsx/lib/openai';
+
+function BuildingBlocksAI({ query }: { query: string }) {
+  const usageExamples = (
+    <>
+      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>
+      Example 2 of how you might use this component:
+      <Card header="Your Ferry Booking" footer="Reservation held for 20 minutes">
+        **Leaves** at 4:15p and **arrives** at 6:20p.
+      </Card>
+      Example 3 of how you might use this component (using with surrounding markdown): Sure, I'd be happy to help you
+      find a car wash. Here are some options:
+      <Card header="AutoWorld">$50 for a quick car wash.</Card>
+      <Card header="Big Joel Big Trucks">$155 for a detailing</Card>
+      <Card header="Small Joel Small Trucks">$10 for some guy to spray your car with a hose.</Card>
+      Example 4 of how you might use this component, after writing out a report on economics: ... and that concludes the
+      report on economics.
+      <Card header="Primary Points">
+        * Price is determined by supply and demand * Setting price floors or ceilings cause deadweight loss. *
+        Interfering with the natural price can also cause shortages.
+      </Card>
+      Use a button group when the user needs to make a choice. A ButtonGroup requires a labels prop. Example 1 of how
+      you might use this component:
+      <ButtonGroup labels={['Yes', 'No']} />
+      Example 2 of how you might use this component (using with surrounding markdown): The system is configured. How
+      would you like to proceed?
+      <ButtonGroup labels={['Deploy to prod', 'Deploy to staging', 'Cancel']} />
+      Use a badge to indicate status:
+      <Badge color="yellow">In progress</Badge>
+      <Badge color="green">Complete</Badge>
+      Use a toggle to let users enable/disable an option:
+      <Toggle title="Use rocket fuel" subtitle="($7 surcharge)" />
+    </>
+  );
+
+  return (
+    <OpenAI chatModel="gpt-4">
+      <ChatCompletion>
+        <MdxSystemMessage usageExamples={usageExamples} />
+        <UserMessage>{query}</UserMessage>
+      </ChatCompletion>
+    </OpenAI>
+  );
+}
+
+export async function POST(request: NextRequest) {
+  const { messages } = await request.json();
+  const lastMessage = _.last(messages) as Message;
+
+  return new StreamingTextResponse(toTextStream(<BuildingBlocksAI query={lastMessage.content} />));
+}

packages/nextjs-demo/src/app/building-blocks/loading.tsx

@@ -0,0 +1,7 @@
+export default function Loading() {
+  return (
+    <div className="flex items-center justify-center w-screen h-screen">
+      <p>Loading...</p>
+    </div>
+  );
+}

packages/nextjs-demo/src/app/building-blocks/page.tsx

@@ -0,0 +1,17 @@
+import InputPrompt from '@/components/InputPrompt';
+import ResultContainer from '@/components/ResultContainer';
+import BuildingBlocks from '@/components/BuildingBlocksGenerator';
+
+export default function BuildingBlocksPage({ searchParams }: { searchParams: any }) {
+  const defaultValue =
+    'Summarize this JSON blob for me, using a Card: {"reservation":{"reservationId":"1234567890","passengerName":"John Doe","flightNumber":"ABC123","origin":"Los Angeles","destination":"New York","departureDate":"2022-01-01","departureTime":"09:00","arrivalDate":"2022-01-01","arrivalTime":"15:00"}}. Also use a Badge. And give me some other github-flavored markdown, using a table, <details>, and strikethrough.';
+  const query = searchParams.q ?? defaultValue;
+  return (
+    <div>
+      <ResultContainer title="Building Blocks" description="In this demo, the AI can use building block UI components">
+        <InputPrompt label="Ask me anything..." defaultValue={query} />
+      </ResultContainer>
+      <BuildingBlocks topic={query} />
+    </div>
+  );
+}