From b5337176cd0d50c1691c2ae4bcc02e71a37b407c Mon Sep 17 00:00:00 2001
From: Peter Salas <peter@fixie.ai>
Date: Thu, 12 Oct 2023 14:46:39 -0700
Subject: [PATCH] Tweak <Sidekick> behavior (#400)

- Move user system message to the start of the prompt
- Make MDX formatting logic conditional on using MDX
---
 packages/ai-jsx/package.json                  |  2 +-
 .../sidekick/platform/conversation.tsx        | 18 +++++-----
 .../batteries/sidekick/platform/sidekick.tsx  | 35 +++++++++++++------
 .../sidekick/platform/system-message.tsx      |  8 +++--
 .../sidekick/platform/use-tools-eject.tsx     | 12 -------
 5 files changed, 39 insertions(+), 36 deletions(-)
 delete mode 100644 packages/ai-jsx/src/batteries/sidekick/platform/use-tools-eject.tsx

diff --git a/packages/ai-jsx/package.json b/packages/ai-jsx/package.json
index cde0dbd80..edde4e1c8 100644
--- 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.21.2",
+  "version": "0.22.0",
   "volta": {
     "extends": "../../package.json"
   },
diff --git a/packages/ai-jsx/src/batteries/sidekick/platform/conversation.tsx b/packages/ai-jsx/src/batteries/sidekick/platform/conversation.tsx
index e54305fe2..8e59f049c 100644
--- a/packages/ai-jsx/src/batteries/sidekick/platform/conversation.tsx
+++ b/packages/ai-jsx/src/batteries/sidekick/platform/conversation.tsx
@@ -15,6 +15,7 @@ import {
 import { LargeFunctionResponseWrapper, redactedFunctionTools } from './large-response-handler.js';
 import { ExecuteFunction, UseToolsProps } from '../../use-tools.js';
 import _ from 'lodash';
+import { SidekickOutputFormat } from './sidekick.js';
 
 /**
  * This function defines the shrinking policy. It's activated when the conversation history overflows the context
@@ -54,8 +55,8 @@ export function getShrinkableConversation(messages: ConversationMessage[], fullC
   });
 }
 
-export function present(conversationElement: ConversationMessage) {
-  if (conversationElement.type === 'assistant') {
+export function present(conversationElement: ConversationMessage, outputFormat: SidekickOutputFormat) {
+  if (conversationElement.type === 'assistant' && outputFormat === 'text/mdx') {
     return (
       <AssistantMessage>
         <LimitToValidMdx>{conversationElement.element}</LimitToValidMdx>
@@ -78,6 +79,7 @@ export function present(conversationElement: ConversationMessage) {
 export function getNextConversationStep(
   messages: ConversationMessage[],
   fullConversation: ConversationMessage[],
+  outputFormat: SidekickOutputFormat,
   tools?: UseToolsProps['tools']
 ) {
   const shrinkableConversation = getShrinkableConversation(messages, fullConversation);
@@ -118,14 +120,12 @@ export function getNextConversationStep(
      */
     case 'system':
     case 'user':
-    case 'functionResponse':
-      return (
-        <RepairMdxInConversation>
-          <ChatCompletion functionDefinitions={tools ? updatedTools : undefined}>
-            {shrinkableConversation}
-          </ChatCompletion>
-        </RepairMdxInConversation>
+    case 'functionResponse': {
+      const generation = (
+        <ChatCompletion functionDefinitions={tools ? updatedTools : undefined}>{shrinkableConversation}</ChatCompletion>
       );
+      return outputFormat === 'text/mdx' ? <RepairMdxInConversation>{generation}</RepairMdxInConversation> : generation;
+    }
     default:
       return null;
   }
diff --git a/packages/ai-jsx/src/batteries/sidekick/platform/sidekick.tsx b/packages/ai-jsx/src/batteries/sidekick/platform/sidekick.tsx
index f0eb2a25e..880abe174 100644
--- a/packages/ai-jsx/src/batteries/sidekick/platform/sidekick.tsx
+++ b/packages/ai-jsx/src/batteries/sidekick/platform/sidekick.tsx
@@ -1,14 +1,18 @@
-import { present } from './conversation.js';
-import { UseTools } from './use-tools-eject.js';
+import { getNextConversationStep, present } from './conversation.js';
 import { SidekickSystemMessage } from './system-message.js';
 import { UseToolsProps } from '../../use-tools.js';
 import * as AI from '../../../index.js';
-import { ConversationHistory, ShowConversation } from '../../../core/conversation.js';
+import { ConversationHistory, Converse, ShowConversation } from '../../../core/conversation.js';
 import { MergeExclusive } from 'type-fest';
 
 interface UniversalSidekickProps {
   tools?: UseToolsProps['tools'];
   systemMessage?: AI.Node;
+
+  /**
+   * The conversation to act on. If not specified, uses the <ConversationHistory /> component.
+   */
+  children?: AI.Node;
 }
 
 type OutputFormatSidekickProps = MergeExclusive<
@@ -77,21 +81,30 @@ type OutputFormatSidekickProps = MergeExclusive<
 
 export type SidekickProps = UniversalSidekickProps & OutputFormatSidekickProps;
 
+export type SidekickOutputFormat = Exclude<SidekickProps['outputFormat'], undefined>;
+
 export function Sidekick(props: SidekickProps) {
+  const outputFormat = props.outputFormat ?? 'text/mdx';
   return (
-    <ShowConversation present={present}>
-      <UseTools tools={props.tools ?? undefined} showSteps>
+    <ShowConversation present={(msg) => present(msg, outputFormat)}>
+      <Converse
+        reply={(messages, fullConversation) =>
+          getNextConversationStep(messages, fullConversation, outputFormat, props.tools)
+        }
+      >
+        {props.systemMessage}
         <SidekickSystemMessage
           timeZone="America/Los_Angeles"
-          includeNextStepsRecommendations={props.includeNextStepsRecommendations ?? true}
-          useCitationCard={props.useCitationCard ?? true}
-          outputFormat={props.outputFormat ?? 'text/mdx'}
+          includeNextStepsRecommendations={
+            outputFormat === 'text/mdx' && (props.includeNextStepsRecommendations ?? true)
+          }
+          useCitationCard={outputFormat === 'text/mdx' && (props.useCitationCard ?? true)}
+          outputFormat={outputFormat}
           userProvidedGenUIUsageExamples={props.genUIExamples}
           userProvidedGenUIComponentNames={props.genUIComponentNames}
         />
-        <ConversationHistory />
-        {props.systemMessage}
-      </UseTools>
+        {props.children ?? <ConversationHistory />}
+      </Converse>
     </ShowConversation>
   );
 }
diff --git a/packages/ai-jsx/src/batteries/sidekick/platform/system-message.tsx b/packages/ai-jsx/src/batteries/sidekick/platform/system-message.tsx
index b48a278a3..7e8c91799 100644
--- a/packages/ai-jsx/src/batteries/sidekick/platform/system-message.tsx
+++ b/packages/ai-jsx/src/batteries/sidekick/platform/system-message.tsx
@@ -2,10 +2,12 @@ import { SystemMessage } from '../../../core/conversation.js';
 import { MdxSystemMessage } from '../../../react/jit-ui/mdx.js';
 import { MdxUsageExamples } from './gen-ui.js';
 import { Node } from '../../../index.js';
-import { SidekickProps } from './sidekick.js';
+import { SidekickOutputFormat } from './sidekick.js';
 
-export interface SidekickSystemMessageProps
-  extends Pick<SidekickProps, 'outputFormat' | 'includeNextStepsRecommendations' | 'useCitationCard'> {
+export interface SidekickSystemMessageProps {
+  outputFormat: SidekickOutputFormat;
+  includeNextStepsRecommendations: boolean;
+  useCitationCard: boolean;
   timeZone: string;
   userProvidedGenUIUsageExamples?: Node;
   userProvidedGenUIComponentNames?: string[];
diff --git a/packages/ai-jsx/src/batteries/sidekick/platform/use-tools-eject.tsx b/packages/ai-jsx/src/batteries/sidekick/platform/use-tools-eject.tsx
deleted file mode 100644
index 7b7c955c8..000000000
--- a/packages/ai-jsx/src/batteries/sidekick/platform/use-tools-eject.tsx
+++ /dev/null
@@ -1,12 +0,0 @@
-import { getNextConversationStep } from './conversation.js';
-import { Converse } from '../../../core/conversation.js';
-import { UseToolsProps } from '../../use-tools.js';
-import { SetOptional } from 'type-fest';
-
-export function UseTools(props: SetOptional<UseToolsProps, 'tools'>) {
-  return (
-    <Converse reply={(messages, fullConversation) => getNextConversationStep(messages, fullConversation, props.tools)}>
-      {props.children}
-    </Converse>
-  );
-}