Fix memo bug, allow conversation rounds to run in parallel (#503)

This change:
- Fixes a bug where memoized streams sometimes duplicated their content
- Refactors `<Converse>` to allow the rounds to run in parallel when the
content allows

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

packages/ai-jsx/src/core/conversation.tsx

@@ -304,32 +304,30 @@ export async function renderToConversation(
  * ```
  *
  */
-export async function* Converse(
-  {
-    reply,
-    children,
-  }: {
-    reply: (messages: ConversationMessage[], fullConversation: ConversationMessage[]) => AI.Renderable;
-    children: AI.Node;
-  },
-  { render, memo, logger }: AI.ComponentContext
-): AI.RenderableStream {
-  yield AI.AppendOnlyStream;
-
-  const fullConversation = [] as ConversationMessage[];
-  let next = children;
-  while (true) {
-    const newMessages = await renderToConversation(next, render, logger);
-    if (newMessages.length === 0) {
-      break;
+export function Converse({
+  reply,
+  children,
+}: {
+  reply: (messages: ConversationMessage[], fullConversation: ConversationMessage[]) => AI.Renderable;
+  children: AI.Node;
+}) {
+  // Keep producing rounds until there's a round with no messages.
+  async function* ConversationRound(
+    { currentRound, history }: { currentRound: AI.Node; history: ConversationMessage[] },
+    { memo, render }: AI.ComponentContext
+  ) {
+    yield;
+    const currentRoundMessages = await renderToConversation(currentRound, render);
+    if (currentRoundMessages.length === 0) {
+      return;
     }
 
-    fullConversation.push(...newMessages);
-    next = memo(reply(newMessages, fullConversation.slice()));
-    yield next;
+    const newHistory = history.concat(currentRoundMessages);
+    const nextRound = memo(reply(currentRoundMessages, newHistory.slice()));
+    return [nextRound, <ConversationRound history={newHistory} currentRound={nextRound} />];
   }
 
-  return AI.AppendOnlyStream;
+  return <ConversationRound history={[]} currentRound={children} />;
 }
 
 /**

packages/ai-jsx/src/core/memoize.tsx

@@ -94,6 +94,7 @@ export function partialMemo(renderable: Renderable, id: number): Node | Renderab
       async *[Symbol.asyncIterator](): AsyncGenerator<Node | AppendOnlyStreamValue, Node | AppendOnlyStreamValue> {
         let index = 0;
         let isAppendOnly = false;
+        let didYieldSomething = false;
 
         while (true) {
           if (index < sink.length) {
@@ -102,6 +103,13 @@ export function partialMemo(renderable: Renderable, id: number): Node | Renderab
             while (index < sink.length) {
               let value = sink[index++];
               if (isAppendOnlyStreamValue(value)) {
+                if (!isAppendOnly && didYieldSomething && concatenatedNodes.length > 0) {
+                  // The stream is transitioning to append-only, but we previously yielded a value
+                  // that needs to be replaced before we start appending. Yield the replacement
+                  // value (`concatenatedNodes`) before we start appending.
+                  yield concatenatedNodes;
+                  concatenatedNodes = [];
+                }
                 isAppendOnly = true;
                 value = valueToAppend(value);
               }
@@ -119,6 +127,7 @@ export function partialMemo(renderable: Renderable, id: number): Node | Renderab
               return valueToYield;
             }
 
+            didYieldSomething = true;
             yield valueToYield;
             continue;
           }

packages/docs/docs/changelog.md

@@ -1,6 +1,15 @@
 # Changelog
 
-## 0.26.1
+## 0.27.1
+
+- Fix bug where memoized components could duplicate content
+- Refactor `<Converse>` to allow rounds to progress in parallel when content allows
+
+## [0.27.0](https://github.com/fixie-ai/ai-jsx/commit/83627e8d5d7bd86dd2fde505962af92bd25a02a1)
+
+- Add new `batchFrames` render option to coalesce ready frames
+
+## [0.26.1](https://github.com/fixie-ai/ai-jsx/commit/6f27bf8b5d1093e5523bd1214bdec2773182144c)
 
 - Fix `js-tiktoken` import that fails on 1.0.8.
 

packages/examples/test/core/memoize.tsx

@@ -157,6 +157,8 @@ it('works for streams that become append-only using a value', async () => {
 
 it('coalesces frames when there are multiple concurrent renders', async () => {
   async function* Component() {
+    yield 5;
+    yield 4;
     yield AI.AppendOnlyStream;
     yield 3;
     yield 2;
@@ -170,17 +172,22 @@ it('coalesces frames when there are multiple concurrent renders', async () => {
   const iterator1 = ctx.render(element)[Symbol.asyncIterator]();
   const iterator2 = ctx.render(element)[Symbol.asyncIterator]();
 
-  expect((await iterator1.next()).value).toBe('');
-  expect((await iterator2.next()).value).toBe('');
+  expect((await iterator1.next()).value).toBe('5');
+  expect((await iterator2.next()).value).toBe('5');
+
+  expect((await iterator1.next()).value).toBe('4');
+  expect((await iterator1.next()).value).toBe('4');
+  expect((await iterator1.next()).value).toBe('43');
+
+  expect((await iterator2.next()).value).toBe('4');
 
-  expect((await iterator1.next()).value).toBe('3');
-  expect((await iterator1.next()).value).toBe('32');
-  expect((await iterator1.next()).value).toBe('321');
-  expect((await iterator2.next()).value).toBe('321');
+  expect((await iterator1.next()).value).toBe('432');
+  expect((await iterator1.next()).value).toBe('4321');
+  expect((await iterator2.next()).value).toBe('4321');
 
-  expect(await iterator1.next()).toEqual({ value: '321LIFTOFF', done: true });
-  expect(await iterator2.next()).toEqual({ value: '321LIFTOFF', done: true });
+  expect(await iterator1.next()).toEqual({ value: '4321LIFTOFF', done: true });
+  expect(await iterator2.next()).toEqual({ value: '4321LIFTOFF', done: true });
 
   const iterator3 = ctx.render(element)[Symbol.asyncIterator]();
-  expect(await iterator3.next()).toEqual({ value: '321LIFTOFF', done: true });
+  expect(await iterator3.next()).toEqual({ value: '4321LIFTOFF', done: true });
 });