Show executed code in web app chat message references

src/interface/web/app/chat/page.tsx

@@ -12,7 +12,12 @@ import { processMessageChunk } from "../common/chatFunctions";
 
 import "katex/dist/katex.min.css";
 
-import { Context, OnlineContext, StreamMessage } from "../components/chatMessage/chatMessage";
+import {
+    CodeContext,
+    Context,
+    OnlineContext,
+    StreamMessage,
+} from "../components/chatMessage/chatMessage";
 import { useIPLocationData, useIsMobileWidth, welcomeConsole } from "../common/utils";
 import ChatInputArea, { ChatOptions } from "../components/chatInputArea/chatInputArea";
 import { useAuthenticatedData } from "../common/auth";
@@ -167,6 +172,7 @@ export default function Chat() {
                 trainOfThought: [],
                 context: [],
                 onlineContext: {},
+                codeContext: {},
                 completed: false,
                 timestamp: new Date().toISOString(),
                 rawQuery: queryToProcess || "",
@@ -195,6 +201,7 @@ export default function Chat() {
         // Track context used for chat response
         let context: Context[] = [];
         let onlineContext: OnlineContext = {};
+        let codeContext: CodeContext = {};
 
         while (true) {
             const { done, value } = await reader.read();
@@ -221,11 +228,12 @@ export default function Chat() {
                     }
 
                     // Track context used for chat response. References are rendered at the end of the chat
-                    ({ context, onlineContext } = processMessageChunk(
+                    ({ context, onlineContext, codeContext } = processMessageChunk(
                         event,
                         currentMessage,
                         context,
                         onlineContext,
+                        codeContext,
                     ));
 
                     setMessages([...messages]);

src/interface/web/app/common/chatFunctions.ts

@@ -1,13 +1,20 @@
-import { Context, OnlineContext, StreamMessage } from "../components/chatMessage/chatMessage";
+import {
+    CodeContext,
+    Context,
+    OnlineContext,
+    StreamMessage,
+} from "../components/chatMessage/chatMessage";
 
 export interface RawReferenceData {
     context?: Context[];
     onlineContext?: OnlineContext;
+    codeContext?: CodeContext;
 }
 
 export interface ResponseWithReferences {
     context?: Context[];
     online?: OnlineContext;
+    codeContext?: CodeContext;
     response?: string;
 }
 
@@ -63,10 +70,11 @@ export function processMessageChunk(
     currentMessage: StreamMessage,
     context: Context[] = [],
     onlineContext: OnlineContext = {},
-): { context: Context[]; onlineContext: OnlineContext } {
+    codeContext: CodeContext = {},
+): { context: Context[]; onlineContext: OnlineContext; codeContext: CodeContext } {
     const chunk = convertMessageChunkToJson(rawChunk);
 
-    if (!currentMessage || !chunk || !chunk.type) return { context, onlineContext };
+    if (!currentMessage || !chunk || !chunk.type) return { context, onlineContext, codeContext };
 
     if (chunk.type === "status") {
         console.log(`status: ${chunk.data}`);
@@ -77,7 +85,8 @@ export function processMessageChunk(
 
         if (references.context) context = references.context;
         if (references.onlineContext) onlineContext = references.onlineContext;
-        return { context, onlineContext };
+        if (references.codeContext) codeContext = references.codeContext;
+        return { context, onlineContext, codeContext };
     } else if (chunk.type === "message") {
         const chunkData = chunk.data;
         if (chunkData !== null && typeof chunkData === "object") {
@@ -102,13 +111,14 @@ export function processMessageChunk(
         console.log(`Completed streaming: ${new Date()}`);
 
         // Append any references after all the data has been streamed
+        if (codeContext) currentMessage.codeContext = codeContext;
         if (onlineContext) currentMessage.onlineContext = onlineContext;
         if (context) currentMessage.context = context;
 
         // Mark current message streaming as completed
         currentMessage.completed = true;
     }
-    return { context, onlineContext };
+    return { context, onlineContext, codeContext };
 }
 
 export function handleImageResponse(imageJson: any, liveStream: boolean): ResponseWithReferences {

src/interface/web/app/components/chatHistory/chatHistory.tsx

@@ -295,6 +295,7 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                             message: message.rawQuery,
                                             context: [],
                                             onlineContext: {},
+                                            codeContext: {},
                                             created: message.timestamp,
                                             by: "you",
                                             automationId: "",
@@ -318,6 +319,7 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                             message: message.rawResponse,
                                             context: message.context,
                                             onlineContext: message.onlineContext,
+                                            codeContext: message.codeContext,
                                             created: message.timestamp,
                                             by: "khoj",
                                             automationId: "",
@@ -338,6 +340,7 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                 message: props.pendingMessage,
                                 context: [],
                                 onlineContext: {},
+                                codeContext: {},
                                 created: new Date().getTime().toString(),
                                 by: "you",
                                 automationId: "",

src/interface/web/app/components/chatMessage/chatMessage.tsx

@@ -97,6 +97,26 @@ export interface OnlineContextData {
     peopleAlsoAsk: PeopleAlsoAsk[];
 }
 
+export interface CodeContext {
+    [key: string]: CodeContextData;
+}
+
+export interface CodeContextData {
+    code: string;
+    results: {
+        success: boolean;
+        output_files: CodeContextFile[];
+        std_out: string;
+        std_err: string;
+        code_runtime: number;
+    };
+}
+
+export interface CodeContextFile {
+    filename: string;
+    b64_data: string;
+}
+
 interface Intent {
     type: string;
     query: string;
@@ -111,6 +131,7 @@ export interface SingleChatMessage {
     created: string;
     context: Context[];
     onlineContext: OnlineContext;
+    codeContext: CodeContext;
     rawQuery?: string;
     intent?: Intent;
     agent?: AgentData;
@@ -122,6 +143,7 @@ export interface StreamMessage {
     trainOfThought: string[];
     context: Context[];
     onlineContext: OnlineContext;
+    codeContext: CodeContext;
     completed: boolean;
     rawQuery: string;
     timestamp: string;
@@ -539,6 +561,7 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
     const allReferences = constructAllReferences(
         props.chatMessage.context,
         props.chatMessage.onlineContext,
+        props.chatMessage.codeContext,
     );
 
     return (
@@ -560,6 +583,7 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
                     isMobileWidth={props.isMobileWidth}
                     notesReferenceCardData={allReferences.notesReferenceCardData}
                     onlineReferenceCardData={allReferences.onlineReferenceCardData}
+                    codeReferenceCardData={allReferences.codeReferenceCardData}
                 />
             </div>
             <div className={styles.chatFooter}>

src/interface/web/app/components/referencePanel/referencePanel.tsx

@@ -11,7 +11,7 @@ const md = new markdownIt({
     typographer: true,
 });
 
-import { Context, WebPage, OnlineContext } from "../chatMessage/chatMessage";
+import { Context, WebPage, OnlineContext, CodeContext } from "../chatMessage/chatMessage";
 import { Card } from "@/components/ui/card";
 
 import {
@@ -94,11 +94,67 @@ function NotesContextReferenceCard(props: NotesContextReferenceCardProps) {
     );
 }
 
+interface CodeContextReferenceCardProps {
+    code: string;
+    output: string;
+    error: string;
+    showFullContent: boolean;
+}
+
+function CodeContextReferenceCard(props: CodeContextReferenceCardProps) {
+    const fileIcon = getIconFromFilename(".py", "w-6 h-6 text-muted-foreground inline-flex mr-2");
+    const snippet = DOMPurify.sanitize(props.code);
+    const [isHovering, setIsHovering] = useState(false);
+
+    return (
+        <>
+            <Popover open={isHovering && !props.showFullContent} onOpenChange={setIsHovering}>
+                <PopoverTrigger asChild>
+                    <Card
+                        onMouseEnter={() => setIsHovering(true)}
+                        onMouseLeave={() => setIsHovering(false)}
+                        className={`${props.showFullContent ? "w-auto" : "w-[200px]"} overflow-hidden break-words text-balance rounded-lg p-2 bg-muted border-none`}
+                    >
+                        <h3
+                            className={`${props.showFullContent ? "block" : "line-clamp-1"} text-muted-foreground}`}
+                        >
+                            {fileIcon}
+                            Code
+                        </h3>
+                        <p
+                            className={`${props.showFullContent ? "block" : "overflow-hidden line-clamp-2"}`}
+                        >
+                            {snippet}
+                        </p>
+                    </Card>
+                </PopoverTrigger>
+                <PopoverContent className="w-[400px] mx-2">
+                    <Card
+                        className={`w-auto overflow-hidden break-words text-balance rounded-lg p-2 border-none`}
+                    >
+                        <h3 className={`line-clamp-2 text-muted-foreground}`}>
+                            {fileIcon}
+                            Code
+                        </h3>
+                        <p className={`overflow-hidden line-clamp-3`}>{snippet}</p>
+                    </Card>
+                </PopoverContent>
+            </Popover>
+        </>
+    );
+}
+
 export interface ReferencePanelData {
     notesReferenceCardData: NotesContextReferenceData[];
     onlineReferenceCardData: OnlineReferenceData[];
 }
 
+export interface CodeReferenceData {
+    code: string;
+    output: string;
+    error: string;
+}
+
 interface OnlineReferenceData {
     title: string;
     description: string;
@@ -214,9 +270,27 @@ function GenericOnlineReferenceCard(props: OnlineReferenceCardProps) {
     );
 }
 
-export function constructAllReferences(contextData: Context[], onlineData: OnlineContext) {
+export function constructAllReferences(
+    contextData: Context[],
+    onlineData: OnlineContext,
+    codeContext: CodeContext,
+) {
     const onlineReferences: OnlineReferenceData[] = [];
     const contextReferences: NotesContextReferenceData[] = [];
+    const codeReferences: CodeReferenceData[] = [];
+
+    if (codeContext) {
+        for (const [key, value] of Object.entries(codeContext)) {
+            if (!value.results) {
+                continue;
+            }
+            codeReferences.push({
+                code: value.code,
+                output: value.results.std_out,
+                error: value.results.std_err,
+            });
+        }
+    }
 
     if (onlineData) {
         let localOnlineReferences = [];
@@ -298,12 +372,14 @@ export function constructAllReferences(contextData: Context[], onlineData: Onlin
     return {
         notesReferenceCardData: contextReferences,
         onlineReferenceCardData: onlineReferences,
+        codeReferenceCardData: codeReferences,
     };
 }
 
 export interface TeaserReferenceSectionProps {
     notesReferenceCardData: NotesContextReferenceData[];
     onlineReferenceCardData: OnlineReferenceData[];
+    codeReferenceCardData: CodeReferenceData[];
     isMobileWidth: boolean;
 }
 
@@ -315,16 +391,27 @@ export function TeaserReferencesSection(props: TeaserReferenceSectionProps) {
     }, [props.isMobileWidth]);
 
     const notesDataToShow = props.notesReferenceCardData.slice(0, numTeaserSlots);
+    const codeDataToShow = props.codeReferenceCardData.slice(
+        0,
+        numTeaserSlots - notesDataToShow.length,
+    );
     const onlineDataToShow =
-        notesDataToShow.length < numTeaserSlots
+        notesDataToShow.length + codeDataToShow.length < numTeaserSlots
-            ? props.onlineReferenceCardData.slice(0, numTeaserSlots - notesDataToShow.length)
+            ? props.onlineReferenceCardData.slice(
+                  0,
+                  numTeaserSlots - codeDataToShow.length - notesDataToShow.length,
+              )
             : [];
 
     const shouldShowShowMoreButton =
-        props.notesReferenceCardData.length > 0 || props.onlineReferenceCardData.length > 0;
+        props.notesReferenceCardData.length > 0 ||
+        props.codeReferenceCardData.length > 0 ||
+        props.onlineReferenceCardData.length > 0;
 
     const numReferences =
-        props.notesReferenceCardData.length + props.onlineReferenceCardData.length;
+        props.notesReferenceCardData.length +
+        props.codeReferenceCardData.length +
+        props.onlineReferenceCardData.length;
 
     if (numReferences === 0) {
         return null;
@@ -346,6 +433,15 @@ export function TeaserReferencesSection(props: TeaserReferenceSectionProps) {
                         />
                     );
                 })}
+                {codeDataToShow.map((code, index) => {
+                    return (
+                        <CodeContextReferenceCard
+                            showFullContent={false}
+                            {...code}
+                            key={`code-${index}`}
+                        />
+                    );
+                })}
                 {onlineDataToShow.map((online, index) => {
                     return (
                         <GenericOnlineReferenceCard

src/interface/web/app/factchecker/page.tsx

@@ -5,6 +5,7 @@ import { useAuthenticatedData } from "@/app/common/auth";
 import { useState, useEffect } from "react";
 
 import ChatMessage, {
+    CodeContext,
     Context,
     OnlineContext,
     OnlineContextData,
@@ -46,6 +47,7 @@ interface SupplementReferences {
 interface ResponseWithReferences {
     context?: Context[];
     online?: OnlineContext;
+    code?: CodeContext;
     response?: string;
 }
 
@@ -192,6 +194,7 @@ function ReferenceVerification(props: ReferenceVerificationProps) {
                     context: [],
                     created: new Date().toISOString(),
                     onlineContext: {},
+                    codeContext: {},
                 }}
                 isMobileWidth={isMobileWidth}
             />
@@ -622,6 +625,7 @@ export default function FactChecker() {
                                         context: [],
                                         created: new Date().toISOString(),
                                         onlineContext: {},
+                                        codeContext: {},
                                     }}
                                     isMobileWidth={isMobileWidth}
                                 />

src/interface/web/app/share/chat/page.tsx

@@ -164,6 +164,7 @@ export default function SharedChat() {
                 trainOfThought: [],
                 context: [],
                 onlineContext: {},
+                codeContext: {},
                 completed: false,
                 timestamp: new Date().toISOString(),
                 rawQuery: queryToProcess || "",