Observability + make streaming demo look better (#168)

Last night, I noticed the streaming demo wouldn't display any results until the whole thing had finished. I wanted to make the demo have an immediately-streaming text result, following by the UI, but that didn't work because everything appeared at the end. This PR includes a number of fixes I needed to make that work, as well as observability I added along the way. Gif of the UI changes: https://github.com/fixie-ai/ai-jsx/assets/829827/055baf67-b63c-4a31-be5d-e2ad7bc48401 To push the UI further, I think we'll want to do more engine-level stuff, like add the ability for AI to pass props. What I did here is pretty hacky.
fixie-ai · Jun 30, 2023 · e2735fd · e2735fd · vercel · Jun 30, 2023
1 parent 183f4bb
commit e2735fd
Show file tree

Hide file tree

Showing 22 changed files with 538 additions and 65 deletions.
diff --git a/packages/ai-jsx/package.json b/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.9",
+  "version": "0.5.10",
   "volta": {
     "extends": "../../package.json"
   },

diff --git a/packages/ai-jsx/src/core/core.ts b/packages/ai-jsx/src/core/core.ts
@@ -367,10 +367,20 @@ async function* renderStream(
     const logImpl = renderingContext.getContext(LoggerContext);
     const renderId = uuidv4();
     try {
-      return yield* renderingContext.render(
+      /**
+       * This approach is pretty noisy because there are many internal components about which the users don't care.
+       * For instance, if the user writes <ChatCompletion>, that'll generate a bunch of internal helpers to
+       * locate + call the model provider, etc.
+       *
+       * To get around this, maybe we want components to be able to choose the loglevel used for their rendering.
+       */
+      logImpl.log('debug', renderable, renderId, 'Start rendering element');
+      const finalResult = yield* renderingContext.render(
         renderable.render(renderingContext, new BoundLogger(logImpl, renderId, renderable)),
         recursiveRenderOpts
       );
+      logImpl.log('debug', renderable, renderId, { finalResult }, 'Finished rendering element');
+      return finalResult;
     } catch (ex) {
       logImpl.logException(renderable, renderId, ex);
       throw ex;
@@ -385,6 +395,8 @@ async function* renderStream(
     while (true) {
       const next = await iterator.next();
       if (next.value === AppendOnlyStream) {
+        // TODO: I'd like to emit a log here indicating that an element has chosen to AppendOnlyStream,
+        // but I'm not sure what the best way is to know which element/renderId produced `renderable`.
         isAppendOnlyStream = true;
       } else if (isAppendOnlyStream) {
         const renderResult = context.render(next.value, recursiveRenderOpts);
@@ -428,7 +440,7 @@ export function createRenderContext(opts?: { logger?: LogImplementation }) {
   });
 }
 
-function createRenderContextInternal(render: StreamRenderer, userContext: Record<symbol, any>): RenderContext {
+function createRenderContextInternal(renderStream: StreamRenderer, userContext: Record<symbol, any>): RenderContext {
   const context: RenderContext = {
     render: <TFinal extends string | PartiallyRendered[], TIntermediate>(
       renderable: Renderable,
@@ -441,7 +453,7 @@ function createRenderContextInternal(render: StreamRenderer, userContext: Record
       const generator = (async function* () {
         // eslint-disable-next-line @typescript-eslint/prefer-nullish-coalescing
         const shouldStop = (opts?.stop || (() => false)) as ElementPredicate;
-        const generatorToWrap = render(context, renderable, shouldStop, Boolean(opts?.appendOnly));
+        const generatorToWrap = renderStream(context, renderable, shouldStop, Boolean(opts?.appendOnly));
         while (true) {
           const next = await generatorToWrap.next();
           const value = opts?.stop ? (next.value as TFinal) : (next.value.join('') as TFinal);
@@ -520,10 +532,13 @@ function createRenderContextInternal(render: StreamRenderer, userContext: Record
       return defaultValue;
     },
 
-    wrapRender: (getRender) => createRenderContextInternal(getRender(render), userContext),
+    wrapRender: (getRenderStream) => createRenderContextInternal(getRenderStream(renderStream), userContext),
 
     [pushContextSymbol]: (contextReference, value) =>
-      createRenderContextInternal(render, { ...userContext, [contextReference[contextKey].userContextSymbol]: value }),
+      createRenderContextInternal(renderStream, {
+        ...userContext,
+        [contextReference[contextKey].userContextSymbol]: value,
+      }),
   };
 
   return context;

diff --git a/packages/ai-jsx/src/experimental/next/index.ts b/packages/ai-jsx/src/experimental/next/index.ts
@@ -160,8 +160,12 @@ export const jsx = asJsxBoundary(function jsx(
 });
 export const JSX = jsx;
 
-export function toReactStream(componentMap: ComponentMap<any>, renderable: AI.Renderable): Response {
-  const renderResult = AI.createRenderContext().render(renderable, {
+export function toReactStream(
+  componentMap: ComponentMap<any>,
+  renderable: AI.Renderable,
+  renderContextOpts?: Parameters<typeof AI.createRenderContext>[0]
+): Response {
+  const renderResult = AI.createRenderContext(renderContextOpts).render(renderable, {
     stop: (e) => boundaryElements.some((special) => special.tag === e.tag),
     map: (x) => x,
   });

diff --git a/packages/ai-jsx/src/lib/openai.tsx b/packages/ai-jsx/src/lib/openai.tsx
@@ -448,14 +448,10 @@ export async function* OpenAIChatModel(
  * @param size The size of the image to generate. Defaults to `512x512`.
  * @returns URL(s) to the generated image, wrapped in {@link Image} component(s).
  */
-export async function DalleImageGen(
+export async function* DalleImageGen(
   { numSamples = 1, size = '512x512', children }: ImageGenPropsWithChildren,
   { render, getContext, logger }: AI.ComponentContext
 ) {
-  const prompt = await render(children);
-
-  const openai = getContext(openAiClientContext);
-
   let sizeEnum;
   switch (size) {
     case '256x256':
@@ -475,6 +471,19 @@ export async function DalleImageGen(
       );
   }
 
+  // Consider emitting http://via.placeholder.com/256x256 instead.
+  yield (
+    <Image
+      url={`http://via.placeholder.com/${size}`}
+      prompt="placeholder while real results renderes"
+      modelName="placeholder.com"
+    />
+  );
+
+  const prompt = await render(children);
+
+  const openai = getContext(openAiClientContext);
+
   const imageRequest = {
     prompt,
     n: numSamples,

diff --git a/packages/ai-jsx/src/react/completion.tsx b/packages/ai-jsx/src/react/completion.tsx
@@ -9,10 +9,11 @@ function reactComponentName(component: React.JSXElementConstructor<any> | string
   return typeof component === 'string' ? component : component.name;
 }
 
-export async function UICompletion(
+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)) {

diff --git a/packages/create-react-app-demo/package.json b/packages/create-react-app-demo/package.json
@@ -13,7 +13,7 @@
     "@testing-library/react": "^13.4.0",
     "@testing-library/user-event": "^13.5.0",
     "@types/http-proxy": "^1.17.11",
-    "ai-jsx": "0.5.9",
+    "ai-jsx": "0.5.10",
     "babel-jest": "^27.4.2",
     "babel-loader": "^8.2.3",
     "babel-plugin-named-asset-import": "^0.3.8",

diff --git a/packages/docs/docs/changelog.md b/packages/docs/docs/changelog.md
@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.5.10
+
+- Update logging to log the output of every component.
+- Update [`UseTools`](./api/modules/batteries_use_tools.md) to use [OpenAI function calls](https://openai.com/blog/function-calling-and-other-api-updates) if you're using a model that supports them.
+
 ## [0.5.9](https://github.com/fixie-ai/ai-jsx/commit/92b6e0f28580fbd9b8fb62072d8c13e28b14d9fe)
 
 - [Add Anthropic support.](./guides/models.md).

diff --git a/packages/docs/docs/guides/rules-of-jsx.md b/packages/docs/docs/guides/rules-of-jsx.md
@@ -39,7 +39,6 @@ function App({ query }) {
     <ChatCompletion>
       <SystemMessage>
         Answer customer questions based on their data: <CustomerData />
-        Here's data about our company: <OrgData />
       </SystemMessage>
       <UserMessage>{query}</UserMessage>
     </ChatCompletion>
@@ -50,14 +49,48 @@ async function CustomerData() {
   const accountId = await getCustomerAccount();
   return isLegacyAccount(accountId) ? fetchLegacy() : fetchModern();
 }
+```
+
+:::caution Edge case
+
+Imagine you have an slow async component, which is used as a sibling of faster components:
+
+```tsx
+async function Slow() {
+  await new Promise((resolve) => setTimeout(resolve, 4000));
+  return 'slow result';
+}
+
+async function Fast() {
+  await Promise.resolve();
+  return 'fast result';
+}
+
+const app = (
+  <>
+    <Fast />
+    <Slow />
+  </>
+);
+```
+
+Surprisingly, you won't get any results streamed out of `Fast` until `Slow` completes.
 
-function* OrgData() {
-  yield firstData;
-  yield secondData;
-  yield thirdData;
+To solve this, return an intermediate value from `Slow`:
+
+```tsx
+async function* Slow() {
+  // highlight-next-line
+  yield '';
+
+  await new Promise((resolve) => setTimeout(resolve, 4000));
+  return 'slow result';
 }
 ```
 
+This is not ideal and we plan to improve it in the future.
+:::
+
 ### Append-only generators
 
 If your component is a generator, the default behavior is that each `yield`ed value replaces the previous value. For instance, imagine you have an image generation API like Midjourney that returns a series of image URLs showing the image render in progress:

diff --git a/packages/examples/package.json b/packages/examples/package.json
@@ -29,6 +29,7 @@
     "demo:wandb": "yarn build && node dist/wandb.js",
     "demo:helloworld": "yarn build && node dist/helloworld.js",
     "demo:inline-chat": "yarn build && node dist/inline-chat.js",
+    "demo:stream-ui": "yarn build && node dist/stream-ui.js",
     "demo:inline-completion": "yarn build && node dist/inline-completion.js",
     "demo:simple-chat": "yarn build && node dist/simple-chat.js",
     "demo:multi-model-chat": "yarn build && node dist/multi-model-chat.js",
@@ -59,7 +60,7 @@
     "@opentelemetry/sdk-metrics": "^1.13.0",
     "@opentelemetry/sdk-node": "^0.39.1",
     "@wandb/sdk": "^0.5.1",
-    "ai-jsx": "0.5.9",
+    "ai-jsx": "0.5.10",
     "axios": "^1.4.0",
     "csv-stringify": "^6.4.0",
     "globby": "^13.1.4",

diff --git a/packages/examples/src/logger.tsx b/packages/examples/src/logger.tsx
@@ -20,6 +20,13 @@ class ConsoleLogger extends LogImplementation {
   }
 }
 
+async function* Slow({ delay }: { delay: number }): AI.RenderableStream {
+  yield AI.AppendOnlyStream;
+  yield `first ${delay}`;
+  await new Promise((resolve) => setTimeout(resolve, delay));
+  return ` second ${delay}`;
+}
+
 function CharacterGenerator() {
   const inlineCompletion = (prompt: AI.Node) => (
     <Completion stop={['"']} temperature={1.0}>
@@ -28,14 +35,18 @@ function CharacterGenerator() {
   );
 
   return (
-    <Inline>
-      The following is a character profile for an RPG game in JSON format:{'\n'}
-      {'{'}
-      {'\n  '}"class": "{inlineCompletion}",
-      {'\n  '}"name": "{inlineCompletion}",
-      {'\n  '}"mantra": "{inlineCompletion}"{'\n'}
-      {'}'}
-    </Inline>
+    <>
+      <Slow delay={1000} />
+      <Slow delay={2000} />
+      <Inline>
+        The following is a character profile for an RPG game in JSON format:{'\n'}
+        {'{'}
+        {'\n  '}"class": "{inlineCompletion}",
+        {'\n  '}"name": "{inlineCompletion}",
+        {'\n  '}"mantra": "{inlineCompletion}"{'\n'}
+        {'}'}
+      </Inline>
+    </>
   );
 }
 

diff --git a/packages/examples/src/stream-ui.tsx b/packages/examples/src/stream-ui.tsx
@@ -0,0 +1,62 @@
+/**
+ * To make this demo work, comment out the `import 'server-only'` line in `ai-jsx/experimental/next`.
+ */
+
+/** @jsxImportSource ai-jsx/react */
+import * as AI from 'ai-jsx/experimental/next';
+import { memo } from 'ai-jsx/core/memoize';
+import { ChatCompletion, UserMessage } from 'ai-jsx/core/completion';
+import { makeComponentMap } from 'ai-jsx/react/map';
+import { pino } from 'pino';
+import { PinoLogger } from 'ai-jsx/core/log';
+
+function App() {
+  const chatCompletion = memo(
+    <ChatCompletion temperature={1}>
+      <UserMessage>List five dog names</UserMessage>
+    </ChatCompletion>
+  );
+
+  return (
+    <>
+      {chatCompletion}
+      <Slow delay={2000} />
+    </>
+  );
+}
+
+async function Slow({ delay }: { delay: number }) {
+  // By default, this demo will show that the tree stream waits for `Slow` to complete before rendering anything.
+  // If we `yield ''`, the problem is solved.
+
+  // yield ''
+  await new Promise((resolve) => setTimeout(resolve, delay));
+  return ` returned after ${delay}`;
+}
+
+const pinoStdoutLogger = pino({
+  name: 'ai-jsx',
+  level: process.env.loglevel ?? 'trace',
+  transport: {
+    target: 'pino-pretty',
+    options: {
+      colorize: true,
+    },
+  },
+});
+
+const decoder = new TextDecoder();
+
+const response = await AI.toReactStream(makeComponentMap({}), <App />, {
+  logger: new PinoLogger(pinoStdoutLogger),
+});
+const body = await response.body;
+const reader = body.getReader();
+// eslint-disable-next-line no-constant-condition
+while (true) {
+  const { done, value } = await reader.read();
+  if (done) {
+    break;
+  }
+  process.stdout.write(decoder.decode(value));
+}
diff --git a/packages/nextjs-demo/package.json b/packages/nextjs-demo/package.json
@@ -11,12 +11,15 @@
     "lint:fix": "yarn lint --fix"
   },
   "dependencies": {
+    "@headlessui/react": "^1.7.15",
+    "@heroicons/react": "^2.0.18",
     "@octokit/graphql": "^5.0.6",
+    "@tailwindcss/forms": "^0.5.3",
     "@types/node": "20.2.5",
     "@types/react": "18.2.8",
     "@types/react-dom": "18.2.4",
     "ai": "^2.1.8",
-    "ai-jsx": "0.5.9",
+    "ai-jsx": "0.5.10",
     "autoprefixer": "10.4.14",
     "classnames": "^2.3.2",
     "eslint": "8.42.0",

diff --git a/packages/nextjs-demo/src/app/layout.tsx b/packages/nextjs-demo/src/app/layout.tsx
@@ -8,12 +8,20 @@ export const metadata = {
   description: 'A framework for AI-native UIs',
 };
 
+const isScreenshare = false;
+
 export default function RootLayout({ children }: { children: React.ReactNode }) {
   return (
     <html lang="en">
       <body>
-        <NavBar />
-        <About />
+        {/* This is intentionally a constant value. */}
+        {/* eslint-disable-next-line @typescript-eslint/no-unnecessary-condition */}
+        {!isScreenshare && (
+          <>
+            <NavBar />
+            <About />
+          </>
+        )}
         <main className="flex min-h-screen flex-col items-start px-24 py-6">{children}</main>
         <div className="hidden">
           <h1 className="text-2xl"></h1>