Change ChatCompletion to emit AssistantMessage (#218)

This changes the various `ChatCompletion` implementations to emit
`<AssistantMessage>`s rather than loose text. This allows callers to
easily separate `<AssistantMessage>` and `<FunctionCall>` outputs.

Preserving the append-only nature of the stream requires non-trivial
gymnastics to split the iterator consumption across multiple components,
but this technique also allows for the removal of the
`experimental_streamFunctionCallOnly` flag.

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

packages/ai-jsx/src/batteries/constrained-output.tsx

@@ -11,6 +11,7 @@ import { zodToJsonSchema } from 'zod-to-json-schema';
 import {
   AssistantMessage,
   ChatCompletion,
+  FunctionCall,
   ModelPropsWithChildren,
   SystemMessage,
   UserMessage,
@@ -318,7 +319,6 @@ export async function* JsonChatCompletionFunctionCall(
 
   const childrenWithCompletion = (
     <ChatCompletion
-      experimental_streamFunctionCallOnly
       {...props}
       functionDefinitions={{
         print: {
@@ -336,28 +336,44 @@ export async function* JsonChatCompletionFunctionCall(
     </ChatCompletion>
   );
 
-  const frames = render(childrenWithCompletion);
+  const frames = render(childrenWithCompletion, { stop: (e) => e.tag === FunctionCall, map: (e) => e });
   for await (const frame of frames) {
-    const object = JSON.parse(frame).arguments;
+    const functionCall = frame.find((e) => AI.isElement(e) && e.tag === FunctionCall) as
+      | AI.Element<AI.PropsOfComponent<typeof FunctionCall>>
+      | undefined;
+    if (!functionCall) {
+      continue;
+    }
+
+    const jsonResult = functionCall.props.args;
     try {
       for (const validator of validatorsAndSchema) {
-        validator(object);
+        validator(jsonResult);
       }
     } catch (e: any) {
       continue;
     }
-    yield JSON.stringify(object);
+    yield JSON.stringify(jsonResult);
   }
-  const object = JSON.parse(await frames).arguments;
+
+  const functionCall = (await frames).find((e) => AI.isElement(e) && e.tag === FunctionCall) as
+    | AI.Element<AI.PropsOfComponent<typeof FunctionCall>>
+    | undefined;
+
+  if (functionCall === undefined) {
+    return null;
+  }
+
+  const jsonResult = functionCall.props.args;
   try {
     for (const validator of validatorsAndSchema) {
-      validator(object);
+      validator(jsonResult);
     }
   } catch (e: any) {
     throw new CompletionError('The model did not produce a valid JSON object', 'runtime', {
-      output: JSON.stringify(object),
+      output: JSON.stringify(jsonResult),
       validationError: e.message,
     });
   }
-  return JSON.stringify(object);
+  return JSON.stringify(jsonResult);
 }

packages/ai-jsx/src/batteries/use-tools.tsx

@@ -221,7 +221,7 @@ export async function* UseTools(props: UseToolsProps, { render }: RenderContext)
 /** @hidden */
 export async function* UseToolsFunctionCall(
   props: UseToolsProps,
-  { render, memo }: ComponentContext
+  { render, memo, logger }: ComponentContext
 ): RenderableStream {
   yield AppendOnlyStream;
 
@@ -233,35 +233,30 @@ export async function* UseToolsFunctionCall(
       yield modelResponse;
     }
 
-    const renderResult = await render(modelResponse, { stop: (el) => el.tag == FunctionCall });
+    const renderResult = await render(modelResponse, {
+      stop: (el) => el.tag === AssistantMessage || el.tag == FunctionCall,
+    });
     let functionCallElement: Element<any> | null = null;
 
-    let currentString = '';
     for (const element of renderResult) {
-      if (typeof element === 'string') {
-        // Model has generated a string response. Record it.
-        currentString += element;
-      } else if (isElement(element) && element.tag === FunctionCall) {
-        if (currentString.trim() !== '') {
-          conversation.push(<AssistantMessage>{currentString}</AssistantMessage>);
-          currentString = '';
-        }
+      if (isElement(element)) {
+        conversation.push(memo(element));
 
-        // Model has generated a function call.
-        if (functionCallElement) {
-          throw new AIJSXError(
-            `ChatCompletion returned 2 function calls at the same time ${renderResult.join(', ')}`,
-            ErrorCode.ModelOutputCouldNotBeParsedForTool,
-            'runtime'
-          );
+        if (element.tag === FunctionCall) {
+          // Model has generated a function call.
+          if (functionCallElement) {
+            throw new AIJSXError(
+              `ChatCompletion returned 2 function calls at the same time ${renderResult.join(', ')}`,
+              ErrorCode.ModelOutputCouldNotBeParsedForTool,
+              'runtime'
+            );
+          }
+          functionCallElement = element;
         }
-        conversation.push(memo(element));
-        functionCallElement = element;
       } else {
-        throw new AIJSXError(
-          `Unexpected result from render ${renderResult.join(', ')}`,
-          ErrorCode.ModelOutputCouldNotBeParsedForTool,
-          'runtime'
+        logger.debug(
+          { text: element },
+          '<ChatCompletion> emitted something other than <AssistantMessage> or <FunctionCall>, which is unexpected.'
         );
       }
     }

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

@@ -299,8 +299,16 @@ export function ConversationHistory({ messages }: { messages: ChatCompletionResp
  *    ==> "That would be 83,076."
  * ```
  */
-export function FunctionCall({ name, args }: { name: string; args: Record<string, string | number | boolean | null> }) {
-  return `Call function ${name} with ${JSON.stringify(args)}`;
+export function FunctionCall({
+  name,
+  partial,
+  args,
+}: {
+  name: string;
+  partial?: boolean;
+  args: Record<string, string | number | boolean | null>;
+}) {
+  return `Call function ${name} with ${partial ? '(incomplete) ' : ''}${JSON.stringify(args)}`;
 }
 
 /**

packages/ai-jsx/src/lib/anthropic.tsx

@@ -93,7 +93,7 @@ interface AnthropicChatModelProps extends ModelPropsWithChildren {
 }
 export async function* AnthropicChatModel(
   props: AnthropicChatModelProps,
-  { render, getContext, logger }: AI.ComponentContext
+  { render, getContext, logger, memo }: AI.ComponentContext
 ): AI.RenderableStream {
   if ('functionDefinitions' in props) {
     throw new AIJSXError(
@@ -186,22 +186,34 @@ export async function* AnthropicChatModel(
     }
     throw err;
   }
-  let resultSoFar = '';
-  let isFirstResponse = true;
-  for await (const completion of response) {
-    let text = completion.completion;
-    if (isFirstResponse && text.length > 0) {
-      isFirstResponse = false;
-      if (text.startsWith(' ')) {
-        text = text.slice(1);
-      }
-    }
-    resultSoFar += text;
-    logger.trace({ completion }, 'Got Anthropic stream event');
-    yield text;
-  }
 
-  logger.debug({ completion: resultSoFar }, 'Anthropic completion finished');
+  // Embed the stream "within" an <AssistantMessage>, memoizing it to ensure it's only consumed once.
+  yield (
+    <AssistantMessage>
+      {memo(
+        (async function* (): AI.RenderableStream {
+          yield AI.AppendOnlyStream;
+          let resultSoFar = '';
+          let isFirstResponse = true;
+          for await (const completion of response) {
+            let text = completion.completion;
+            if (isFirstResponse && text.length > 0) {
+              isFirstResponse = false;
+              if (text.startsWith(' ')) {
+                text = text.slice(1);
+              }
+            }
+            resultSoFar += text;
+            logger.trace({ completion }, 'Got Anthropic stream event');
+            yield text;
+          }
+
+          logger.debug({ completion: resultSoFar }, 'Anthropic completion finished');
+          return AI.AppendOnlyStream;
+        })()
+      )}
+    </AssistantMessage>
+  );
 
   return AI.AppendOnlyStream;
 }