feat: Final error display and step duration (#88)


app/api/vision-agent/route.ts

@@ -161,8 +161,8 @@ export const POST = withLogging(
           // default to dev apikey
           apikey:
             process.env.LND_TIER === 'production'
-              ? 'land_sk_nMnUf8xiJJUjyw1l5QaIJJ4ZyrvPthzVmPAIG7TtJY7F9CW6lu'
-              : 'land_sk_DKeoYtaZZrYqJ9TMMiXe4BIQgJcZ0s3XAoB0JT3jv73FFqnr6k',
+              ? 'land_sk_nMnUf8xiJJUjyw1l5QaIJJ4ZyrvPthzVmPAIG7TtJY7F9CW6lu' // prod key
+              : 'land_sk_DKeoYtaZZrYqJ9TMMiXe4BIQgJcZ0s3XAoB0JT3jv73FFqnr6k', // dev key
         },
         body: formData,
       },

components/ChatInterface.tsx

@@ -25,12 +25,12 @@ const ChatInterface: React.FC<ChatInterfaceProps> = ({ chat, userId }) => {
         className="pl-4 peer absolute right-0 inset-y-0 hidden translate-x-full data-[state=open]:translate-x-0 z-30 duration-300 ease-in-out xl:flex flex-col items-start xl:w-1/2 h-full dark:bg-zinc-950 overflow-auto"
       >
         {messageCodeResult?.payload && (
-          <Card className="w-full">
+          <Card className="size-full overflow-auto">
             <CodeResultDisplay codeResult={messageCodeResult.payload} />
           </Card>
         )}
       </div>
-      <div className="w-full flex justify-center overflow-auto pr-0 animate-in duration-300 ease-in-out peer-[[data-state=open]]:xl:pr-[50%]">
+      <div className="w-full flex justify-center pr-0 animate-in duration-300 ease-in-out peer-[[data-state=open]]:xl:pr-[50%]">
         <ChatList chat={chat} userId={userId} />
       </div>
     </div>

components/CodeResultDisplay.tsx

@@ -21,7 +21,6 @@ import {
   CarouselNext,
   CarouselPrevious,
 } from './ui/carousel';
-import Link from 'next/link';
 
 export interface CodeResultDisplayProps {}
 

components/chat/ChatMessage.tsx

@@ -11,7 +11,11 @@ import {
   IconGlowingDot,
 } from '@/components/ui/Icons';
 import { MessageUI } from '@/lib/types';
-import { ChunkBody, CodeResult, formatStreamLogs } from '@/lib/utils/content';
+import {
+  WIPChunkBodyGroup,
+  CodeResult,
+  formatStreamLogs,
+} from '@/lib/utils/content';
 import {
   Table,
   TableBody,
@@ -43,10 +47,12 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
 }) => {
   const [messageId, setMessageId] = useAtom(selectedMessageId);
   const { id, mediaUrl, prompt, response, result } = message;
-  const [formattedSections, codeResult] = useMemo(
+  const [formattedSections, finalResult, finalError] = useMemo(
     () => formatStreamLogs(response ?? wipAssistantMessage?.content),
     [response, wipAssistantMessage?.content],
   );
+  // prioritize the result from the message over the WIP message
+  const codeResult = result?.payload ?? finalResult;
   return (
     <div
       className={cn(
@@ -103,7 +109,10 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
                         {ChunkStatusToIconDict[section.status]}
                       </TableCell>
                       <TableCell className="font-medium">
-                        <ChunkTypeToText chunk={section} />
+                        <ChunkTypeToText
+                          useTimer={!codeResult && !finalError}
+                          chunk={section}
+                        />
                       </TableCell>
                       <TableCell className="text-right">
                         <ChunkPayloadAction payload={section.payload} />
@@ -113,11 +122,30 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
                 </TableBody>
               </Table>
               {codeResult && (
-                <div className="xl:hidden">
-                  <CodeResultDisplay codeResult={codeResult} />
-                </div>
+                <>
+                  <div className="xl:hidden">
+                    <CodeResultDisplay codeResult={codeResult} />
+                  </div>
+                  <p>✨ Coding complete</p>
+                </>
+              )}
+              {!codeResult && finalError && (
+                <>
+                  <p>❌ Coding failed </p>
+                  <div>
+                    <CodeBlock
+                      language="error"
+                      value={
+                        finalError.name +
+                        '\n' +
+                        finalError.value +
+                        '\n' +
+                        finalError.traceback_raw.join('\n')
+                      }
+                    />
+                  </div>
+                </>
               )}
-              {codeResult && <p>✨ Coding complete</p>}
             </div>
           </div>
         </>
@@ -134,42 +162,54 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
   );
 };
 
-const ChunkStatusToIconDict: Record<ChunkBody['status'], React.ReactElement> = {
+const ChunkStatusToIconDict: Record<
+  WIPChunkBodyGroup['status'],
+  React.ReactElement
+> = {
   started: <IconGlowingDot className="bg-yellow-500/80" />,
   completed: <IconCheckCircle className="text-green-500" />,
   running: <IconGlowingDot className="bg-teal-500/80" />,
   failed: <IconCrossCircle className="text-red-500" />,
 };
 const ChunkTypeToText: React.FC<{
-  chunk: ChunkBody;
-}> = ({ chunk }) => {
-  const { status, type } = chunk;
+  chunk: WIPChunkBodyGroup;
+  useTimer: boolean;
+}> = ({ chunk, useTimer }) => {
+  const { status, type, timestamp, duration } = chunk;
 
-  const [seconds, setSeconds] = useState(0);
-  const isExecution = type === 'code' && status === 'running';
+  const [mSeconds, setMSeconds] = useState(0);
+  const isExecuting = !['completed', 'failed'].includes(status);
 
   useEffect(() => {
-    if (isExecution) {
+    if (isExecuting && timestamp && !useTimer) {
       const timerId = setInterval(() => {
-        setSeconds(prevSeconds => Math.round((prevSeconds + 0.2) * 10) / 10);
+        setMSeconds(Date.now() - Date.parse(timestamp));
       }, 200);
       return () => clearInterval(timerId);
     }
-  }, [isExecution]);
+  }, [isExecuting, timestamp, useTimer]);
+
+  const displayMs = isExecuting && useTimer ? mSeconds : duration;
+
+  const durationDisplay = displayMs
+    ? `(${Math.round(displayMs / 100) / 10}s)`
+    : '';
 
-  if (type === 'plans') return <p>Creating instructions</p>;
-  if (type === 'tools') return <p>Retrieving tools</p>;
-  if (type === 'code' && status === 'started') return <p>Generating code</p>;
-  if (isExecution) return <p>Executing code ({seconds}s)</p>;
+  if (type === 'plans') return <p>Creating instructions {durationDisplay}</p>;
+  if (type === 'tools') return <p>Retrieving tools {durationDisplay}</p>;
+  if (type === 'code' && status === 'started')
+    return <p>Generating code {durationDisplay}</p>;
+  if (type === 'code' && status === 'running')
+    return <p>Executing code {durationDisplay}</p>;
   if (type === 'code' && status === 'completed')
-    return <p>Code execution success ({seconds}s)</p>;
+    return <p>Code execution success {durationDisplay}</p>;
   if (type === 'code' && status === 'failed')
-    return <p>Code execution failure ({seconds}s)</p>;
+    return <p>Code execution failure {durationDisplay}</p>;
 
   return null;
 };
 const ChunkPayloadAction: React.FC<{
-  payload: ChunkBody['payload'];
+  payload: WIPChunkBodyGroup['payload'];
 }> = ({ payload }) => {
   if (!payload) return null;
   if (Array.isArray(payload)) {

components/chat/Composer.tsx

@@ -111,6 +111,7 @@ const Composer = forwardRef<ComposerRef, ComposerProps>(
               <Button
                 size="icon"
                 variant="ghost"
+                disabled={finalDisabled}
                 className="size-4"
                 onClick={() => setLocalMediaUrl(undefined)}
               >

lib/db/functions.ts

@@ -42,6 +42,12 @@ export async function dbFindOrCreateUser(
 
   // If the user doesn't exist, create it
   if (user) {
+    if (!user.avatar && avatar) {
+      await prisma.user.update({
+        where: { email: email },
+        data: { avatar: avatar },
+      });
+    }
     return user;
   } else {
     return prisma.user.create({

lib/utils/content.ts

@@ -1,4 +1,5 @@
 import toast from 'react-hot-toast';
+import { ResultPayload } from '../types';
 const ANSWERS_PREFIX = 'answers';
 
 export const generateAnswersImageMarkdown = (index: number, url: string) => {
@@ -15,80 +16,34 @@ export const cleanAnswerMessage = (content: string) => {
   return content.replace(/!\[answers.*?\.png\)/g, '');
 };
 
-type PlansBody =
-  | {
-      type: 'plans';
-      status: 'started';
-    }
-  | {
-      type: 'plans';
-      status: 'completed';
-      payload: Array<Record<string, string>>;
-    };
-
-type ToolsBody =
-  | {
-      type: 'tools';
-      status: 'started';
-    }
-  | {
-      type: 'tools';
-      status: 'completed';
-      payload: Array<Record<string, string>>;
-    };
-
-type CodeBody =
-  | {
-      type: 'code';
-      status: 'started';
-    }
-  | {
-      type: 'code';
-      status: 'running';
-      payload: {
-        code: string;
-        test: string;
-      };
-    }
-  | {
-      type: 'code';
-      status: 'completed' | 'failed';
-      payload: {
-        code: string;
-        test: string;
-        result: string;
-      };
-    };
-
-// this will return if self_reflection flag is true
-type ReflectionBody =
-  | {
-      type: 'self_reflection';
-      status: 'started';
-    }
-  | {
-      type: 'self_reflection';
-      status: 'completed' | 'failed';
-      payload: { feedback: string; success: boolean };
-    };
-
-
 export type CodeResult = {
   code: string;
   test: string;
   result: string;
 };
 
-export type ChunkBody =
-  | {
-      type: 'plans' | 'tools' | 'code' | 'final_code' | 'final_error';
-      status: 'started' | 'completed' | 'failed' | 'running';
-      timestamp: string;
-      payload:
-        | Array<Record<string, string>> // PlansBody | ToolsBody
-        | CodeResult; // CodeBody
-    }
-  | PrismaJson.FinalChatResult;
+const WIPLogTypes = ['plans', 'tools', 'code'];
+const AllLogTypes = [
+  'plans',
+  'tools',
+  'code',
+  'final_code',
+  'final_error',
+] as const;
+
+export type ChunkBody = {
+  type: (typeof AllLogTypes)[number];
+  status: 'started' | 'completed' | 'failed' | 'running';
+  timestamp?: string;
+  payload:
+    | Array<Record<string, string>> // PlansBody | ToolsBody
+    | CodeResult // CodeBody
+    | ResultPayload['error']; // ErrorBody
+};
+
+export type WIPChunkBodyGroup = ChunkBody & {
+  duration?: number;
+};
 
 /**
  * Formats the stream logs and returns an array of grouped sections.
@@ -98,7 +53,7 @@ export type ChunkBody =
  */
 export const formatStreamLogs = (
   content: string | null | undefined,
-): [ChunkBody[], CodeResult?] => {
+): [WIPChunkBodyGroup[], CodeResult?, ResultPayload['error']?] => {
   if (!content) return [[], undefined];
   const streamLogs = content.split('\n').filter(log => !!log);
 
@@ -124,21 +79,33 @@ export const formatStreamLogs = (
 
   // Merge consecutive logs of the same type to the latest status
   const groupedSections = parsedStreamLogs.reduce((acc, curr) => {
+    const lastGroup = acc[acc.length - 1];
     if (
       acc.length > 0 &&
-      acc[acc.length - 1].type === curr.type &&
+      lastGroup.type === curr.type &&
       curr.status !== 'started'
     ) {
-      acc[acc.length - 1] = curr;
+      acc[acc.length - 1] = {
+        ...curr,
+        // always use the timestamp of the first log
+        timestamp: lastGroup?.timestamp,
+        // duration is the difference between the last log and the first log
+        duration:
+          lastGroup?.timestamp && curr.timestamp
+            ? Date.parse(curr.timestamp) - Date.parse(lastGroup.timestamp)
+            : undefined,
+      };
     } else {
       acc.push(curr);
     }
     return acc;
-  }, [] as ChunkBody[]);
+  }, [] as WIPChunkBodyGroup[]);
 
   return [
-    groupedSections.filter(section => section.type !== 'final_code'),
+    groupedSections.filter(section => WIPLogTypes.includes(section.type)),
     groupedSections.find(section => section.type === 'final_code')
       ?.payload as CodeResult,
+    groupedSections.find(section => section.type === 'final_error')
+      ?.payload as ResultPayload['error'],
   ];
 };