From f00581403047137c30634ead299e546b094aeb13 Mon Sep 17 00:00:00 2001
From: Cole McCracken <colemccracken@gmail.com>
Date: Thu, 21 Dec 2023 10:25:28 -0800
Subject: [PATCH] Cole/def 5751 allow multiple textoutputs to render multiple
 text blocks (#74)

---
 .../library/piedpiper/spend_all_product.yaml  |  29 ++
 .../library/piedpiper/spend_per_product.yaml  |  29 --
 examples/next/src/app/page.tsx                |   9 +-
 packages/openassistants-react/package.json    |   2 +-
 .../src/components/chat-content.tsx           |  59 ++++
 .../src/components/chat-list.tsx              |  63 +---
 .../src/components/chat-message.tsx           | 269 +++++++++---------
 .../src/components/dataframe-table.tsx        |   4 +-
 .../src/components/function-form.tsx          |   9 +-
 .../src/components/plotly-vis.tsx             |   1 -
 10 files changed, 242 insertions(+), 232 deletions(-)
 create mode 100644 examples/fast-api-server/library/piedpiper/spend_all_product.yaml
 delete mode 100644 examples/fast-api-server/library/piedpiper/spend_per_product.yaml
 create mode 100644 packages/openassistants-react/src/components/chat-content.tsx

diff --git a/examples/fast-api-server/library/piedpiper/spend_all_product.yaml b/examples/fast-api-server/library/piedpiper/spend_all_product.yaml
new file mode 100644
index 0000000..0e8049c
--- /dev/null
+++ b/examples/fast-api-server/library/piedpiper/spend_all_product.yaml
@@ -0,0 +1,29 @@
+display_name: Spend all products
+description: |
+  Show the spend for all products (top 5)
+sample_questions:
+  - show spend for all products
+parameters:
+  json_schema:
+    type: object
+    properties: {}
+    required: []
+type: DuckDBQueryFunction
+dataset: 'dummy-data/'
+sqls:
+  - |
+    SELECT purchases.product, SUM(purchases.amount) as amount,
+    FROM "purchases.csv"
+    GROUP BY purchases.product
+    ORDER BY amount DESC
+    LIMIT 5
+visualizations:
+  - |
+    def create_chart_config(dataframes):
+      df = dataframes[0]
+      data = [{'x': df.iloc[:, 0].tolist(), 'y': df.iloc[:, 1].tolist(), 'type': 'bar'}]
+      layout = {'title': 'Spend per Product', 'xaxis': {'title': 'Product'}, 'yaxis': {'title': 'Total Spend'}}
+      result = {'data': data, 'layout': layout}
+
+      return result
+summarization: 'Describe the spend per product'
diff --git a/examples/fast-api-server/library/piedpiper/spend_per_product.yaml b/examples/fast-api-server/library/piedpiper/spend_per_product.yaml
deleted file mode 100644
index 654eb32..0000000
--- a/examples/fast-api-server/library/piedpiper/spend_per_product.yaml
+++ /dev/null
@@ -1,29 +0,0 @@
-display_name: Spend per product
-description: |
-  Show the spend per product (top 5)
-sample_questions:
-- show spend per product
-parameters:
-  json_schema:
-    type: object
-    properties: {}
-    required: []
-type: DuckDBQueryFunction
-dataset: "dummy-data/"
-sqls:
-- |
-  SELECT purchases.product, SUM(purchases.amount) as amount,
-  FROM "purchases.csv"
-  GROUP BY purchases.product
-  ORDER BY amount DESC
-  LIMIT 5
-visualizations:
-  - | 
-      def create_chart_config(dataframes):
-        df = dataframes[0]
-        data = [{'x': df.iloc[:, 0].tolist(), 'y': df.iloc[:, 1].tolist(), 'type': 'bar'}]
-        layout = {'title': 'Spend per Product', 'xaxis': {'title': 'Product'}, 'yaxis': {'title': 'Total Spend'}}
-        result = {'data': data, 'layout': layout}
-
-        return result
-summarization: 'Describe the spend per product'
diff --git a/examples/next/src/app/page.tsx b/examples/next/src/app/page.tsx
index 712949d..91d2e8c 100644
--- a/examples/next/src/app/page.tsx
+++ b/examples/next/src/app/page.tsx
@@ -8,9 +8,8 @@ export default function Home() {
       <div className="relative flex w-full h-full overflow-hidden">
         <div className="group w-full pl-0 animate-in duration-300 ease-in-out peer-[[data-state=open]]:lg:pl-[250px] peer-[[data-state=open]]:xl:pl-[300px]">
           <OpenAssistantsChat
-            api={`${
-              process.env.NEXT_PUBLIC_API_URL || 'http://localhost:8000'
-            }/v1alpha/assistants/hooli/chat`}
+            api={`${process.env.NEXT_PUBLIC_API_URL ||
+              'http://localhost:8000'}/v1alpha/assistants/hooli/chat`}
             name="Example Assistant"
             starterPrompts={[
               {
@@ -18,8 +17,8 @@ export default function Home() {
                 prompt: 'Show recent purchases',
               },
               {
-                title: 'Show spend per product',
-                prompt: 'Show spend per product',
+                title: 'Show spend all products',
+                prompt: 'Show spend all products',
               },
               {
                 title: 'What can I ask?',
diff --git a/packages/openassistants-react/package.json b/packages/openassistants-react/package.json
index fe0ceb9..45d24dc 100644
--- a/packages/openassistants-react/package.json
+++ b/packages/openassistants-react/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@definitive-io/openassistants-react",
-  "version": "0.1.4",
+  "version": "0.1.5",
   "license": "MIT",
   "module": "dist/openassistants-react.esm.js",
   "main": "dist/index.js",
diff --git a/packages/openassistants-react/src/components/chat-content.tsx b/packages/openassistants-react/src/components/chat-content.tsx
new file mode 100644
index 0000000..089bcc4
--- /dev/null
+++ b/packages/openassistants-react/src/components/chat-content.tsx
@@ -0,0 +1,59 @@
+import React from 'react';
+import { MemoizedReactMarkdown } from './markdown';
+import remarkGfm from 'remark-gfm';
+import { CodeBlock } from './ui';
+
+export interface ChatContentProps {
+  content: string;
+}
+
+export const ChatContent = ({ content }: ChatContentProps) => {
+  return (
+    <div>
+      <MemoizedReactMarkdown
+        className={`prose break-words prose-p:leading-relaxed prose-pre:p-0`}
+        remarkPlugins={[remarkGfm]}
+        components={{
+          p({ children }) {
+            return <p className="mb-2 last:mb-0">{children}</p>;
+          },
+          code({ node, inline, className, children, ...props }) {
+            if (children.length) {
+              if (children[0] == '▍') {
+                return (
+                  <span className="mt-1 animate-pulse cursor-default">▍</span>
+                );
+              }
+
+              children[0] = (children[0] as string).replace('`▍`', '▍');
+            }
+
+            const match = /language-(\w+)/.exec(className || '');
+
+            if (inline) {
+              return (
+                <code className={className} {...props}>
+                  {children}
+                </code>
+              );
+            }
+            return (
+              <div>
+                {
+                  <CodeBlock
+                    key={Math.random()}
+                    language={(match && match[1]) || ''}
+                    value={String(children).replace(/\n$/, '')}
+                    {...props}
+                  />
+                }
+              </div>
+            );
+          },
+        }}
+      >
+        {content}
+      </MemoizedReactMarkdown>
+    </div>
+  );
+};
diff --git a/packages/openassistants-react/src/components/chat-list.tsx b/packages/openassistants-react/src/components/chat-list.tsx
index b5d4ad3..3e0542e 100644
--- a/packages/openassistants-react/src/components/chat-list.tsx
+++ b/packages/openassistants-react/src/components/chat-list.tsx
@@ -33,62 +33,17 @@ export function OpenAssistantsChatList({
   return (
     <div className="relative mx-auto max-w-2xl px-4">
       {messages.map((message, index) => {
-        const outputs = 'outputs' in message ? message.outputs : null;
         return (
           <div key={index}>
-            {outputs &&
-              outputs.map((output, i) => {
-                console.log(
-                  `output.type=${output.type} output=${JSON.stringify(output)}`
-                );
-                return (
-                  <div key={`message_${index}_output_${i}`} className="mb-4">
-                    {output.type === 'dataframe' && (
-                      <DataframeTable
-                        key={`message_${index}_df_${i}`}
-                        title={output.title || ''}
-                        dataframe={output.dataframe}
-                      ></DataframeTable>
-                    )}
-                    {output.type === 'visualization' && (
-                      <RenderVisMessage
-                        id={`message-${index}-output-${i}`}
-                        data={output.visualization.data}
-                        layout={output.visualization.layout}
-                      />
-                    )}
-                  </div>
-                );
-              })}
-
-            {getContent(message) && (
-              <div key={index}>
-                <OpenAssistantsChatMessage
-                  key={index}
-                  message={message}
-                  onEdit={(m) => onEdit(index, m)}
-                  isLoading={isLoading}
-                />
-                {index < messages.length - 1 && <Separator className="my-4 " />}
-              </div>
-            )}
-
-            {message.role === 'assistant' && message.input_request && (
-              <FunctionForm
-                message={message}
-                onSubmit={(values) => {
-                  const newMessage = {
-                    role: 'user' as const,
-                    input_response: {
-                      name: message.input_request!.name,
-                      arguments: values,
-                    },
-                    content: '',
-                  };
-                  append(newMessage);
-                }}
-              ></FunctionForm>
-            )}
+            <OpenAssistantsChatMessage
+              key={index}
+              total={messages.length}
+              index={index}
+              message={message}
+              append={append}
+              onEdit={(m) => onEdit(index, m)}
+              isLoading={isLoading}
+            />
           </div>
         );
       })}
diff --git a/packages/openassistants-react/src/components/chat-message.tsx b/packages/openassistants-react/src/components/chat-message.tsx
index 7dbe111..eba474e 100644
--- a/packages/openassistants-react/src/components/chat-message.tsx
+++ b/packages/openassistants-react/src/components/chat-message.tsx
@@ -1,18 +1,21 @@
-import remarkGfm from 'remark-gfm';
-
 import { cn } from '../lib/utils';
-import { CodeBlock } from './ui/codeblock';
-import { MemoizedReactMarkdown } from './markdown';
 import { IconArrowRight, IconDefinitive, IconUser } from './ui/icons';
 import { ChatMessageAction, ChatMessageActions } from './chat-message-actions';
 import { useState } from 'react';
 import React from 'react';
 import { Message } from './chat-models';
-import { Button, Textarea } from './ui';
+import { Button, Separator, Textarea } from './ui';
+import { DataframeTable } from './dataframe-table';
+import { RenderVisMessage } from './plotly-vis';
+import { FunctionForm } from './function-form';
+import { ChatContent } from './chat-content';
 
 export interface DefinitiveChatMessageProps {
+  index: number;
+  total: number;
   message: Message;
   onEdit?: (message: Message) => Promise<void>;
+  append: (message?: Message) => Promise<void>;
   isLoading: boolean;
 }
 
@@ -36,24 +39,7 @@ const isCollapsable = (message: Message): boolean => {
   );
 };
 
-function formatCode(input: string): string {
-  const trimmed = input
-    .replace('BEGIN_CODE', '')
-    .replace('END_CODE', '')
-    .trim();
-  return trimmed
-    .split(';')
-    .map((line) => line.trim())
-    .join(';\n');
-}
-
 export const getContent = (message: Message): string => {
-  if (message.role === 'function') {
-    const textOutput =
-      message.outputs.find((output) => output.type === 'text') || {};
-    //@ts-ignore
-    return 'text' in textOutput ? (textOutput.text as string) : '';
-  }
   if (message.role === 'user') {
     if (message.content) {
       return message.content;
@@ -61,11 +47,6 @@ export const getContent = (message: Message): string => {
       return '';
     }
   }
-  if (message.role === 'assistant') {
-    if (message.function_call) {
-      return '';
-    }
-  }
   if (message.role === 'assistant' && message.content) {
     return message.content;
   }
@@ -73,14 +54,17 @@ export const getContent = (message: Message): string => {
 };
 
 export function OpenAssistantsChatMessage({
+  index,
+  total,
   message,
   onEdit,
   isLoading,
+  append,
   ...props
 }: DefinitiveChatMessageProps) {
   const [isEditing, setIsEditing] = useState(false);
+  const [submitted, setSubmitted] = React.useState(false);
   const [editedMessageContent, setEditedMessageContent] = useState<string>('');
-  const content = getContent(message);
   const supportedActions: ChatMessageAction[] = [];
   if (!isCollapsable(message) && !isEditing) {
     supportedActions.push('copy');
@@ -88,127 +72,136 @@ export function OpenAssistantsChatMessage({
   if (message.role === 'user' && !isEditing && onEdit) {
     supportedActions.push('edit');
   }
-  if (!content) {
+  const outputs = 'outputs' in message ? message.outputs : null;
+  // executed function call- ignore
+  if (message.role === 'assistant' && message.function_call) {
+    return <div></div>;
+  }
+  // completed input response from user- ignore
+  if (message.role === 'user' && message.input_response) {
     return <div></div>;
   }
   return (
-    <div
-      className={cn(
-        'group relative mb-4 flex items-start md:-ml-12 bg-highlight'
-      )}
-      {...props}
-    >
+    <div>
       <div
         className={cn(
-          'flex h-8 w-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-background'
+          'group relative flex items-start md:-ml-12 bg-highlight '
         )}
+        {...props}
       >
-        {message && getIcon(message)}
-      </div>
-      <div className="ml-4 flex-1 space-y-2 overflow-hidden px-1 text-primary">
-        {isEditing && message.role === 'user' && (
-          <div>
-            <Textarea
-              className="w-full p-2 resize-y"
-              value={editedMessageContent}
-              onChange={(e) => setEditedMessageContent(e.target.value)}
-              onKeyDown={(e) => {
-                if (e.key === 'Enter' && !e.shiftKey) {
-                  e.preventDefault();
-                  setIsEditing(false);
-                  onEdit &&
-                    onEdit({
-                      ...message,
-                      content: editedMessageContent,
-                    });
-                }
-              }}
-            />
-            <Button
-              variant="ghost"
-              onClick={() => {
-                setIsEditing(false);
-              }}
-            >
-              Cancel
-            </Button>
-            <Button
-              variant="ghost"
-              onClick={() => {
-                setIsEditing(false);
-                onEdit &&
-                  onEdit({
-                    ...message,
-                    content: editedMessageContent,
-                  });
-              }}
-            >
-              Save
-            </Button>
-          </div>
-        )}
-        {!isEditing && (
-          <div>
-            <MemoizedReactMarkdown
-              className={`prose break-words prose-p:leading-relaxed prose-pre:p-0`}
-              remarkPlugins={[remarkGfm]}
-              components={{
-                p({ children }) {
-                  return <p className="mb-2 last:mb-0">{children}</p>;
-                },
-                code({ node, inline, className, children, ...props }) {
-                  if (children.length) {
-                    if (children[0] == '▍') {
-                      return (
-                        <span className="mt-1 animate-pulse cursor-default">
-                          ▍
-                        </span>
-                      );
-                    }
-
-                    children[0] = (children[0] as string).replace('`▍`', '▍');
-                  }
-
-                  const match = /language-(\w+)/.exec(className || '');
-
-                  if (inline) {
-                    return (
-                      <code className={className} {...props}>
-                        {children}
-                      </code>
-                    );
-                  }
-                  return (
-                    <div>
-                      {
-                        <CodeBlock
-                          key={Math.random()}
-                          language={(match && match[1]) || ''}
-                          value={String(children).replace(/\n$/, '')}
-                          {...props}
-                        />
-                      }
-                    </div>
-                  );
-                },
-              }}
-            >
-              {content}
-            </MemoizedReactMarkdown>
-
-            {'content' in message && (
-              <ChatMessageActions
-                content={message.content}
-                actions={supportedActions}
-                onEditClicked={() => {
-                  setIsEditing(true);
-                  setEditedMessageContent(message.content || '');
+        <div
+          className={cn(
+            'flex h-8 w-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-background'
+          )}
+        >
+          {message && getIcon(message)}
+        </div>
+        <div className={cn('flex flex-col')}>
+          {message.role === 'assistant' && message.input_request && (
+            // add padding
+            <div className={cn('px-8')}>
+              <FunctionForm
+                message={message}
+                submitted={submitted}
+                onSubmit={(values) => {
+                  const newMessage = {
+                    role: 'user' as const,
+                    input_response: {
+                      name: message.input_request!.name,
+                      arguments: values,
+                    },
+                    content: '',
+                  };
+                  setSubmitted(true);
+                  append(newMessage);
                 }}
-              />
+              ></FunctionForm>
+            </div>
+          )}
+
+          <div className="ml-4 flex-1 space-y-2 overflow-hidden px-1 text-primary">
+            {isEditing && message.role === 'user' && (
+              <div>
+                <Textarea
+                  className="w-full p-2 resize-y"
+                  value={editedMessageContent}
+                  onChange={(e) => setEditedMessageContent(e.target.value)}
+                  onKeyDown={(e) => {
+                    if (e.key === 'Enter' && !e.shiftKey) {
+                      e.preventDefault();
+                      setIsEditing(false);
+                      onEdit &&
+                        onEdit({
+                          ...message,
+                          content: editedMessageContent,
+                        });
+                    }
+                  }}
+                />
+                <Button
+                  variant="ghost"
+                  onClick={() => {
+                    setIsEditing(false);
+                  }}
+                >
+                  Cancel
+                </Button>
+                <Button
+                  variant="ghost"
+                  onClick={() => {
+                    setIsEditing(false);
+                    onEdit &&
+                      onEdit({
+                        ...message,
+                        content: editedMessageContent,
+                      });
+                  }}
+                >
+                  Save
+                </Button>
+              </div>
+            )}
+            {outputs &&
+              outputs.map((output, i) => {
+                return (
+                  <div key={`message_${index}_output_${i}`} className="mb-4">
+                    {output.type === 'dataframe' && (
+                      <DataframeTable
+                        key={`message_${index}_df_${i}`}
+                        title={output.title || ''}
+                        dataframe={output.dataframe}
+                      ></DataframeTable>
+                    )}
+                    {output.type === 'visualization' && (
+                      <RenderVisMessage
+                        id={`message-${index}-output-${i}`}
+                        data={output.visualization.data}
+                        layout={output.visualization.layout}
+                      />
+                    )}
+                    {output.type === 'text' && !isEditing && (
+                      <ChatContent content={output.text} />
+                    )}
+                  </div>
+                );
+              })}
+            {'content' in message && !submitted && (
+              <div>
+                <ChatContent content={message.content || ''} />
+                <ChatMessageActions
+                  content={message.content}
+                  actions={supportedActions}
+                  onEditClicked={() => {
+                    setIsEditing(true);
+                    setEditedMessageContent(message.content || '');
+                  }}
+                />
+              </div>
             )}
           </div>
-        )}
+        </div>
       </div>
+      {index < total - 1 && <Separator className="my-4 " />}
     </div>
   );
 }
diff --git a/packages/openassistants-react/src/components/dataframe-table.tsx b/packages/openassistants-react/src/components/dataframe-table.tsx
index 7fca73f..421a898 100644
--- a/packages/openassistants-react/src/components/dataframe-table.tsx
+++ b/packages/openassistants-react/src/components/dataframe-table.tsx
@@ -22,7 +22,9 @@ export const DataframeTable = ({
       className="max-h-96 overflow-y-auto px-4 mx-auto overflow-x-auto"
       style={{ maxWidth: '90vw' }}
     >
-      <h2 className={cn('text-xl text-center text-primary mb-4')}>{title}</h2>
+      {title && (
+        <h2 className={cn('text-xl text-center text-primary mb-4')}>{title}</h2>
+      )}
       <Table
         className={cn(
           'w-full table-auto divide-y divide-gray-700 sm:table-fixed'
diff --git a/packages/openassistants-react/src/components/function-form.tsx b/packages/openassistants-react/src/components/function-form.tsx
index 094b63a..c6f7fa3 100644
--- a/packages/openassistants-react/src/components/function-form.tsx
+++ b/packages/openassistants-react/src/components/function-form.tsx
@@ -9,6 +9,7 @@ import { cn } from '../lib/utils';
 
 interface FuncionFormProps extends React.ComponentProps<'div'> {
   message: AssistantMessage;
+  submitted: boolean;
   onSubmit: (values: any) => void;
 }
 
@@ -44,8 +45,11 @@ const SelectWidget = (props: WidgetProps) => {
   );
 };
 
-export function FunctionForm({ message, onSubmit }: FuncionFormProps) {
-  const [submitted, setSubmitted] = React.useState(false);
+export function FunctionForm({
+  message,
+  submitted,
+  onSubmit,
+}: FuncionFormProps) {
   const json_schema = message?.input_request?.json_schema;
   const uiSchema = {};
   // expand with more widgets
@@ -120,7 +124,6 @@ export function FunctionForm({ message, onSubmit }: FuncionFormProps) {
           validator={validator}
           onSubmit={(form) => {
             onSubmit(form.formData);
-            setSubmitted(true);
           }}
           onError={() => console.log('errors')}
         />
diff --git a/packages/openassistants-react/src/components/plotly-vis.tsx b/packages/openassistants-react/src/components/plotly-vis.tsx
index 9431786..25569f9 100644
--- a/packages/openassistants-react/src/components/plotly-vis.tsx
+++ b/packages/openassistants-react/src/components/plotly-vis.tsx
@@ -14,7 +14,6 @@ export const RenderVisMessage = ({
 }: RenderVisMessageProps) => {
   const Plotly =
     typeof window !== 'undefined' ? require('plotly.js-dist-min') : null;
-  console.log('Plotly', Plotly);
   useEffect(() => {
     if (!data || !layout) {
       return;