feat: Code result image display (#76)




components/CodeResultDisplay.tsx

@@ -29,13 +29,18 @@ const CodeResultDisplay: React.FC<{
         results: detail.results,
         stderr: detail.logs.stderr,
         stdout: detail.logs.stdout,
+        error: detail.error,
       };
     } catch {
       return {};
     }
   };
 
-  const { results, stderr, stdout } = getDetail();
+  const { results = [], stderr, stdout, error } = getDetail();
+
+  const imageResults = results?.filter(_ => !!_.png).map(_ => _.png);
+  const videoResults = results?.filter(_ => !!_.mp4).map(_ => _.mp4);
+  const finalResult = results?.find(_ => _.is_main_result)?.text;
 
   return (
     <div className="rounded-lg overflow-hidden relative max-w-5xl">
@@ -50,7 +55,7 @@ const CodeResultDisplay: React.FC<{
             </DialogTrigger>
             <DialogContent className="max-w-5xl">
               <DialogHeader>
-                <DialogTitle>Test Code</DialogTitle>
+                <DialogTitle>Test code</DialogTitle>
               </DialogHeader>
               <CodeBlock language="python" value={test} />
             </DialogContent>
@@ -75,47 +80,68 @@ const CodeResultDisplay: React.FC<{
           <CodeBlock language="print" value={stdout.join('').trim()} />
         </>
       )}
-      {Array.isArray(results) && !!results.length && (
+      {!!error && (
         <>
           <Separator />
-          {results.map((result, index) => {
-            if (result.png) {
-              return (
-                <Img
-                  key={'png' + index}
-                  src={result.png}
-                  alt={'answer-image'}
-                  quality={100}
-                  sizes="(min-width: 66em) 15vw,
-                        (min-width: 44em) 20vw,
-                        100vw"
-                />
-              );
-            } else if (result.mp4) {
-              return (
-                <video
-                  key={'mp4' + index}
-                  src={result.mp4}
-                  controls
-                  width={500}
-                  height={500}
-                />
-              );
-            } else if (result.text) {
-              return (
-                <CodeBlock
-                  key={'text' + index}
-                  language="output"
-                  value={result.text}
-                />
-              );
-            } else {
-              return null;
+          <CodeBlock
+            language="error"
+            value={
+              error.name +
+              '\n' +
+              error.value +
+              '\n' +
+              error.traceback_raw.join('\n')
             }
-          })}
+          />
         </>
       )}
-      <Separator />
+      {!!imageResults.length && (
+        <div className="p-4 text-xs lowercase bg-zinc-900 space-y-4 border-t border-muted">
+          <p>image output</p>
+          <div className="flex flex-row space-x-4 overflow-auto">
+            {imageResults.map((png, index) => (
+              <Dialog key={'png' + index}>
+                <DialogTrigger asChild>
+                  <Img
+                    key={'png' + index}
+                    src={png!}
+                    width={200}
+                    alt="result-image"
+                    className="cursor-zoom-in"
+                  />
+                </DialogTrigger>
+                <DialogContent className="max-w-5xl">
+                  <Img
+                    src={png!}
+                    width={1200}
+                    height={800}
+                    alt="result-image"
+                    quality={100}
+                  />
+                </DialogContent>
+              </Dialog>
+            ))}
+          </div>
+        </div>
+      )}
+      {!!videoResults.length && (
+        <div className="p-4 text-xs lowercase bg-zinc-900 space-y-4">
+          <p>video output</p>
+          <div className="flex flex-row space-x-4 overflow-auto">
+            {videoResults.map((mp4, index) => (
+              <Dialog key={'png' + index}>
+                <DialogTrigger asChild>
+                  <video src={mp4} controls width={400} height={400} />
+                </DialogTrigger>
+                <DialogContent className="max-w-5xl">
+                  <video src={mp4} controls width={400} height={400} />
+                </DialogContent>
+              </Dialog>
+            ))}
+          </div>
+        </div>
+      )}
+      {!!finalResult && <CodeBlock language="output" value={finalResult} />}
     </div>
   );
 };

components/chat/ChatList.tsx

@@ -57,10 +57,11 @@ const ChatList: React.FC<ChatListProps> = ({ chat }) => {
       <div className="overflow-auto h-full p-4 z-10" ref={messagesRef}>
         {dbMessages.map((message, index) => (
           <ChatMessage
-            key={index}
+            key={message.id}
             message={message}
             wipAssistantMessage={
-              isLoading && lastMessage.role === 'assistant'
+              lastMessage.role === 'assistant' &&
+              index === dbMessages.length - 1
                 ? lastMessage
                 : undefined
             }
@@ -78,10 +79,12 @@ const ChatList: React.FC<ChatListProps> = ({ chat }) => {
           initMediaUrl={dbMessages[dbMessages.length - 1]?.mediaUrl}
           isLoading={isLoading}
           onSubmit={async ({ input, mediaUrl: newMediaUrl }) => {
-            append({
+            const messageInput = {
               prompt: input,
               mediaUrl: newMediaUrl,
-            });
+            };
+            const resp = await dbPostCreateMessage(id, messageInput);
+            append(resp);
           }}
         />
       </div>

components/chat/ChatMessage.tsx

@@ -67,48 +67,59 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
           {mediaUrl && (
             <>
               {mediaUrl?.endsWith('.mp4') ? (
-                <video src={mediaUrl} controls width={500} height={500} />
+                <video src={mediaUrl} controls width={400} height={400} />
               ) : (
-                <Img src={mediaUrl} alt={mediaUrl} quality={100} width={300} />
+                <Dialog>
+                  <DialogTrigger asChild>
+                    <Img src={mediaUrl} alt={mediaUrl} width={300} />
+                  </DialogTrigger>
+                  <DialogContent className="max-w-5xl">
+                    <Img src={mediaUrl} alt={mediaUrl} quality={100} />
+                  </DialogContent>
+                </Dialog>
               )}
             </>
           )}
         </div>
       </div>
-      <Separator className="bg-primary/30 my-4" />
-      <div className="flex">
-        <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-primary text-primary-foreground">
-          <IconLandingAI />
-        </div>
-        <div className="flex-1 px-1 space-y-4 ml-4 overflow-hidden">
-          <Table className="w-[400px]">
-            <TableBody>
-              {formattedSections.map(section => (
-                <TableRow
-                  className="border-primary/50 h-[56px]"
-                  key={section.type}
-                >
-                  <TableCell className="text-center text-webkit-center">
-                    {ChunkStatusToIconDict[section.status]}
-                  </TableCell>
-                  <TableCell className="font-medium">
-                    {ChunkTypeToTextDict[section.type]}
-                  </TableCell>
-                  <TableCell className="text-right">
-                    <ChunkPayloadAction payload={section.payload} />
-                  </TableCell>
-                </TableRow>
-              ))}
-            </TableBody>
-          </Table>
-          {codeResult && (
-            <div className="xl:hidden">
-              <CodeResultDisplay codeResult={codeResult} />
+      {!!formattedSections.length && (
+        <>
+          <Separator className="bg-primary/30 my-4" />
+          <div className="flex">
+            <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-primary text-primary-foreground">
+              <IconLandingAI />
             </div>
-          )}
-          {codeResult && <p>✨ Coding complete</p>}
-        </div>
-      </div>
+            <div className="flex-1 px-1 space-y-4 ml-4 overflow-hidden">
+              <Table className="w-[400px]">
+                <TableBody>
+                  {formattedSections.map(section => (
+                    <TableRow
+                      className="border-primary/50 h-[56px]"
+                      key={section.type}
+                    >
+                      <TableCell className="text-center text-webkit-center">
+                        {ChunkStatusToIconDict[section.status]}
+                      </TableCell>
+                      <TableCell className="font-medium">
+                        {ChunkTypeToTextDict[section.type]}
+                      </TableCell>
+                      <TableCell className="text-right">
+                        <ChunkPayloadAction payload={section.payload} />
+                      </TableCell>
+                    </TableRow>
+                  ))}
+                </TableBody>
+              </Table>
+              {codeResult && (
+                <div className="xl:hidden">
+                  <CodeResultDisplay codeResult={codeResult} />
+                </div>
+              )}
+              {codeResult && <p>✨ Coding complete</p>}
+            </div>
+          </div>
+        </>
+      )}
     </div>
   );
 };

components/chat/Composer.tsx

@@ -98,7 +98,7 @@ const Composer = forwardRef<ComposerRef, ComposerProps>(
                     <p>{mediaName ?? 'unnamed_media'}</p>
                   </div>
                 </TooltipTrigger>
-                <TooltipContent sideOffset={12}>
+                <TooltipContent sideOffset={12} className="max-w-2xl">
                   <Img
                     src={localMediaUrl}
                     className="m-1"

components/ui/CodeBlock.tsx

@@ -45,12 +45,16 @@ export const programmingLanguages: languageMap = {
   sql: '.sql',
   html: '.html',
   css: '.css',
+  // custom titles
   print: '.txt',
+  error: 'vim',
   // add more file extensions here, make sure the key is same as language prop in CodeBlock.tsx component
 };
 
 const customSyntax: languageMap = {
   print: 'vim',
+  output: 'vim',
+  error: 'vim',
 };
 
 export const generateRandomString = (length: number, lowercase = false) => {
@@ -99,7 +103,7 @@ const CodeBlock: FC<Props> = memo(({ language, value }) => {
   };
   return (
     <div className="relative w-full codeblock bg-zinc-900 overflow-hidden">
-      <div className="flex items-center justify-between w-full pl-8 pr-4 pt-2 text-zinc-100">
+      <div className="flex items-center justify-between w-full px-4 pt-2 text-zinc-100">
         <span className="text-xs lowercase">{language}</span>
         <div className="flex items-center space-x-1">
           <Button
@@ -131,7 +135,7 @@ const CodeBlock: FC<Props> = memo(({ language, value }) => {
           margin: 0,
           width: '100%',
           background: 'transparent',
-          padding: '0.5rem 1rem 1.5rem 1rem',
+          padding: '0.5rem 1rem 1rem 1rem',
         }}
         lineNumberStyle={{
           userSelect: 'none',

components/ui/Dialog.tsx

@@ -21,7 +21,7 @@ const DialogOverlay = React.forwardRef<
   <DialogPrimitive.Overlay
     ref={ref}
     className={cn(
-      'fixed inset-0 z-50 bg-black/80  data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0',
+      'fixed inset-0 z-50 bg-black/80 data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0',
       className,
     )}
     {...props}
@@ -38,7 +38,7 @@ const DialogContent: React.ForwardRefExoticComponent<
     <DialogPrimitive.Content
       ref={ref}
       className={cn(
-        'fixed left-1/2 top-1/2 z-50 grid w-full max-w-lg -translate-x-1/2 -translate-y-1/2 gap-4 border bg-background p-6 shadow-lg duration-200 data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[state=closed]:slide-out-to-left-1/2 data-[state=closed]:slide-out-to-top-[48%] data-[state=open]:slide-in-from-left-1/2 data-[state=open]:slide-in-from-top-[48%] sm:rounded-lg',
+        'fixed left-1/2 top-1/2 z-50 grid w-full max-w-lg max-h-[1000px] overflow-auto -translate-x-1/2 -translate-y-1/2 gap-4 border bg-background p-6 shadow-lg duration-200 data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[state=closed]:slide-out-to-left-1/2 data-[state=closed]:slide-out-to-top-[48%] data-[state=open]:slide-in-from-left-1/2 data-[state=open]:slide-in-from-top-[48%] sm:rounded-lg',
         className,
       )}
       {...props}

lib/hooks/useVisionAgent.ts

@@ -12,6 +12,7 @@ import {
 } from '../utils/message';
 import { useSetAtom } from 'jotai';
 import { selectedMessageId } from '@/state/chat';
+import { Message } from '@prisma/client';
 
 const useVisionAgent = (chat: ChatWithMessages) => {
   const { messages: dbMessages, id, mediaUrl } = chat;
@@ -66,16 +67,15 @@ const useVisionAgent = (chat: ChatWithMessages) => {
 
   return {
     messages: messages as MessageUI[],
-    append: async (messageInput: MessageUserInput) => {
-      const resp = await dbPostCreateMessage(id, messageInput);
-      currMediaUrl.current = messageInput.mediaUrl;
-      currMessageId.current = resp.id;
+    append: (message: Message) => {
+      currMediaUrl.current = message.mediaUrl;
+      currMessageId.current = message.id;
       append({
         id,
         role: 'user',
-        content: messageInput.prompt,
+        content: message.prompt,
         // @ts-ignore valid when setting sendExtraMessageFields
-        mediaUrl: messageInput.mediaUrl,
+        mediaUrl: message.mediaUrl,
       });
     },
     reload,

lib/types.ts

@@ -30,4 +30,9 @@ export type ResultPayload = {
     text: string;
     is_main_result: boolean;
   }>;
+  error: {
+    name: string;
+    value: string;
+    traceback_raw: string[];
+  };
 };

lib/utils/content.ts

@@ -1,16 +1,4 @@
 import toast from 'react-hot-toast';
-import { Message } from 'ai';
-
-const PAIRS: Record<string, string> = {
-  '┍': '┑',
-  '┝': '┥',
-  '├': '┤',
-  '┕': '┙',
-};
-
-const MIDDLE_STARTER = '┝';
-const MIDDLE_SEPARATOR = '┿';
-
 const ANSWERS_PREFIX = 'answers';
 
 export const generateAnswersImageMarkdown = (index: number, url: string) => {
@@ -27,56 +15,6 @@ export const cleanAnswerMessage = (content: string) => {
   return content.replace(/!\[answers.*?\.png\)/g, '');
 };
 
-const generateJSONArrayMarkdown = (
-  message: string,
-  payload: Array<Record<string, string | boolean>>,
-) => {
-  if (payload.length === 0) return '';
-  const keys = Object.keys(payload[0]);
-  message += '\n';
-  message += '| ' + keys.join(' | ') + ' |' + '\n';
-  message += new Array(keys.length + 1).fill('|').join(' :- ') + '\n';
-  payload.forEach((obj: any) => {
-    message +=
-      '| ' +
-      keys
-        .map(key => {
-          if (key === 'documentation') {
-            const doc = `\`\`\`\n${obj[key]}\n\`\`\`\n`;
-            return `<button data-details=${JSON.stringify(encodeURI(doc))}>Show</button>`;
-          } else {
-            return obj[key];
-          }
-        })
-        .join(' | ') +
-      ' |' +
-      '\n';
-  });
-  message += '\n';
-  return message;
-};
-
-const generateCodeExecutionMarkdown = (
-  message: string,
-  payload: {
-    code: string;
-    test: string;
-    result?: string;
-  },
-) => {
-  let Details = 'Code: \n';
-  Details += `\`\`\`python\n${payload.code}\n\`\`\`\n`;
-  Details += 'Test: \n';
-  Details += `\`\`\`python\n${payload.test}\n\`\`\`\n`;
-  if (payload.result) {
-    Details += 'Execution result: \n';
-    Details += `\`\`\`python\n${payload.result}\n\`\`\`\n`;
-  }
-  message += `<button data-details=${JSON.stringify(encodeURI(Details))}>View details</button> \n`;
-
-  return message;
-};
-
 type PlansBody =
   | {
       type: 'plans';
@@ -183,26 +121,6 @@ const getMessageTitle = (json: MessageBody) => {
       throw 'Not supported type';
   }
 };
-const parseLine = (json: MessageBody) => {
-  const title = getMessageTitle(json);
-  if (json.status === 'started') {
-    return title;
-  }
-  switch (json.type) {
-    case 'plans':
-      return generateJSONArrayMarkdown(title, json.payload);
-    case 'tools':
-      return generateJSONArrayMarkdown(title, json.payload);
-    case 'code':
-      return generateCodeExecutionMarkdown(title, json.payload);
-    case 'self_reflection':
-      return generateJSONArrayMarkdown(title, [json.payload]);
-    case 'final_code':
-      return '';
-    default:
-      throw 'Not supported type';
-  }
-};
 
 export type CodeResult = {
   code: string;
@@ -254,7 +172,11 @@ export const formatStreamLogs = (
 
   // Merge consecutive logs of the same type to the latest status
   const groupedSections = parsedStreamLogs.reduce((acc, curr) => {
-    if (acc.length > 0 && acc[acc.length - 1].type === curr.type) {
+    if (
+      acc.length > 0 &&
+      acc[acc.length - 1].type === curr.type &&
+      curr.status !== 'started'
+    ) {
       acc[acc.length - 1] = curr;
     } else {
       acc.push(curr);