Merge branch 'master' into multi-image-chat-and-vision-for-gemini

2026-03-03 21:29:08 +00:00 · 2024-10-22 18:29:44 -07:00
parent e8fb79a369 0dad4212fa
commit 6c393800cc
43 changed files with 1080 additions and 147 deletions
--- a/src/interface/web/app/chat/chat.module.css
+++ b/src/interface/web/app/chat/chat.module.css
@@ -79,7 +79,7 @@ div.titleBar {
 div.chatBoxBody {
    display: grid;
    height: 100%;
-    width: 70%;
+    width: 95%;
    margin: auto;
 }

--- a/src/interface/web/app/chat/layout.tsx
+++ b/src/interface/web/app/chat/layout.tsx
@@ -47,7 +47,14 @@ export default function RootLayout({
                       child-src 'none';
                       object-src 'none';"
            ></meta>
-            <body className={inter.className}>{children}</body>
+            <body className={inter.className}>
+                {children}
+                <script
+                    dangerouslySetInnerHTML={{
+                        __html: `window.EXCALIDRAW_ASSET_PATH = 'https://assets.khoj.dev/@excalidraw/excalidraw/dist/';`,
+                    }}
+                />
+            </body>
        </html>
    );
 }
--- a/src/interface/web/app/chat/page.tsx
+++ b/src/interface/web/app/chat/page.tsx
@@ -41,6 +41,8 @@ function ChatBodyData(props: ChatBodyDataProps) {
    const setQueryToProcess = props.setQueryToProcess;
    const onConversationIdChange = props.onConversationIdChange;

+    const chatHistoryCustomClassName = props.isMobileWidth ? "w-full" : "w-4/6";
+
    useEffect(() => {
        if (images.length > 0) {
            const encodedImages = images.map((image) => encodeURIComponent(image));
@@ -105,10 +107,11 @@ function ChatBodyData(props: ChatBodyDataProps) {
                    setAgent={setAgentMetadata}
                    pendingMessage={processingMessage ? message : ""}
                    incomingMessages={props.streamedMessages}
+                    customClassName={chatHistoryCustomClassName}
                />
            </div>
            <div
-                className={`${styles.inputBox} p-1 md:px-2 shadow-md bg-background align-middle items-center justify-center dark:bg-neutral-700 dark:border-0 dark:shadow-sm rounded-t-2xl rounded-b-none md:rounded-xl h-fit`}
+                className={`${styles.inputBox} p-1 md:px-2 shadow-md bg-background align-middle items-center justify-center dark:bg-neutral-700 dark:border-0 dark:shadow-sm rounded-t-2xl rounded-b-none md:rounded-xl h-fit ${chatHistoryCustomClassName} mr-auto ml-auto`}
            >
                <ChatInputArea
                    agentColor={agentMetadata?.color}
--- a/src/interface/web/app/common/chatFunctions.ts
+++ b/src/interface/web/app/common/chatFunctions.ts
@@ -5,10 +5,10 @@ export interface RawReferenceData {
    onlineContext?: OnlineContext;
 }

-export interface ResponseWithReferences {
-    context?: Context[];
-    online?: OnlineContext;
-    response?: string;
+export interface ResponseWithIntent {
+    intentType: string;
+    response: string;
+    inferredQueries?: string[];
 }

 interface MessageChunk {
@@ -49,10 +49,14 @@ export function convertMessageChunkToJson(chunk: string): MessageChunk {
 function handleJsonResponse(chunkData: any) {
    const jsonData = chunkData as any;
    if (jsonData.image || jsonData.detail) {
-        let responseWithReference = handleImageResponse(chunkData, true);
-        if (responseWithReference.response) return responseWithReference.response;
+        let responseWithIntent = handleImageResponse(chunkData, true);
+        return responseWithIntent;
    } else if (jsonData.response) {
-        return jsonData.response;
+        return {
+            response: jsonData.response,
+            intentType: "",
+            inferredQueries: [],
+        };
    } else {
        throw new Error("Invalid JSON response");
    }
@@ -80,8 +84,18 @@ export function processMessageChunk(
        return { context, onlineContext };
    } else if (chunk.type === "message") {
        const chunkData = chunk.data;
+        // Here, handle if the response is a JSON response with an image, but the intentType is excalidraw
        if (chunkData !== null && typeof chunkData === "object") {
-            currentMessage.rawResponse += handleJsonResponse(chunkData);
+            let responseWithIntent = handleJsonResponse(chunkData);
+
+            if (responseWithIntent.intentType && responseWithIntent.intentType === "excalidraw") {
+                currentMessage.rawResponse = responseWithIntent.response;
+            } else {
+                currentMessage.rawResponse += responseWithIntent.response;
+            }
+
+            currentMessage.intentType = responseWithIntent.intentType;
+            currentMessage.inferredQueries = responseWithIntent.inferredQueries;
        } else if (
            typeof chunkData === "string" &&
            chunkData.trim()?.startsWith("{") &&
@@ -89,7 +103,10 @@ export function processMessageChunk(
        ) {
            try {
                const jsonData = JSON.parse(chunkData.trim());
-                currentMessage.rawResponse += handleJsonResponse(jsonData);
+                let responseWithIntent = handleJsonResponse(jsonData);
+                currentMessage.rawResponse += responseWithIntent.response;
+                currentMessage.intentType = responseWithIntent.intentType;
+                currentMessage.inferredQueries = responseWithIntent.inferredQueries;
            } catch (e) {
                currentMessage.rawResponse += JSON.stringify(chunkData);
            }
@@ -111,42 +128,26 @@ export function processMessageChunk(
    return { context, onlineContext };
 }

-export function handleImageResponse(imageJson: any, liveStream: boolean): ResponseWithReferences {
+export function handleImageResponse(imageJson: any, liveStream: boolean): ResponseWithIntent {
    let rawResponse = "";

    if (imageJson.image) {
-        const inferredQuery = imageJson.inferredQueries?.[0] ?? "generated image";
-
-        // If response has image field, response is a generated image.
-        if (imageJson.intentType === "text-to-image") {
-            rawResponse += `![generated_image](data:image/png;base64,${imageJson.image})`;
-        } else if (imageJson.intentType === "text-to-image2") {
-            rawResponse += `![generated_image](${imageJson.image})`;
-        } else if (imageJson.intentType === "text-to-image-v3") {
-            rawResponse = `![](data:image/webp;base64,${imageJson.image})`;
-        }
-        if (inferredQuery && !liveStream) {
-            rawResponse += `\n\n${inferredQuery}`;
-        }
+        // If response has image field, response may be a generated image
+        rawResponse = imageJson.image;
    }

-    let reference: ResponseWithReferences = {};
+    let responseWithIntent: ResponseWithIntent = {
+        intentType: imageJson.intentType,
+        response: rawResponse,
+        inferredQueries: imageJson.inferredQueries,
+    };

-    if (imageJson.context && imageJson.context.length > 0) {
-        const rawReferenceAsJson = imageJson.context;
-        if (rawReferenceAsJson instanceof Array) {
-            reference.context = rawReferenceAsJson;
-        } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
-            reference.online = rawReferenceAsJson;
-        }
-    }
    if (imageJson.detail) {
        // The detail field contains the improved image prompt
        rawResponse += imageJson.detail;
    }

-    reference.response = rawResponse;
-    return reference;
+    return responseWithIntent;
 }

 export function modifyFileFilterForConversation(
--- a/src/interface/web/app/common/iconUtils.tsx
+++ b/src/interface/web/app/common/iconUtils.tsx
@@ -48,6 +48,7 @@ import {
    Oven,
    Gavel,
    Broadcast,
+    KeyReturn,
 } from "@phosphor-icons/react";
 import { Markdown, OrgMode, Pdf, Word } from "@/app/components/logo/fileLogo";

@@ -193,6 +194,10 @@ export function getIconForSlashCommand(command: string, customClassName: string
    }

    if (command.includes("default")) {
+        return <KeyReturn className={className} />;
+    }
+
+    if (command.includes("diagram")) {
        return <Shapes className={className} />;
    }

--- a/src/interface/web/app/components/chatHistory/chatHistory.module.css
+++ b/src/interface/web/app/components/chatHistory/chatHistory.module.css
@@ -2,12 +2,7 @@ div.chatHistory {
    display: flex;
    flex-direction: column;
    height: 100%;
-}
-
-div.chatLayout {
-    height: 80vh;
-    overflow-y: auto;
-    margin: 0 auto;
+    margin: auto;
 }

 div.agentIndicator a {
--- a/src/interface/web/app/components/chatHistory/chatHistory.tsx
+++ b/src/interface/web/app/components/chatHistory/chatHistory.tsx
@@ -37,6 +37,7 @@ interface ChatHistoryProps {
    pendingMessage?: string;
    publicConversationSlug?: string;
    setAgent: (agent: AgentData) => void;
+    customClassName?: string;
 }

 function constructTrainOfThought(
@@ -255,7 +256,7 @@ export default function ChatHistory(props: ChatHistoryProps) {
    return (
        <ScrollArea className={`h-[80vh] relative`} ref={scrollAreaRef}>
            <div>
-                <div className={styles.chatHistory}>
+                <div className={`${styles.chatHistory} ${props.customClassName}`}>
                    <div ref={sentinelRef} style={{ height: "1px" }}>
                        {fetchingData && (
                            <InlineLoading message="Loading Conversation" className="opacity-50" />
@@ -322,6 +323,12 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                            by: "khoj",
                                            automationId: "",
                                            rawQuery: message.rawQuery,
+                                            intent: {
+                                                type: message.intentType || "",
+                                                query: message.rawQuery,
+                                                "memory-type": "",
+                                                "inferred-queries": message.inferredQueries || [],
+                                            },
                                        }}
                                        customClassName="fullHistory"
                                        borderLeftColor={`${data?.agent?.color}-500`}
@@ -365,18 +372,20 @@ export default function ChatHistory(props: ChatHistoryProps) {
                        </div>
                    )}
                </div>
-                {!isNearBottom && (
-                    <button
-                        title="Scroll to bottom"
-                        className="absolute bottom-4 right-5 bg-white dark:bg-[hsl(var(--background))] text-neutral-500 dark:text-white p-2 rounded-full shadow-xl"
-                        onClick={() => {
-                            scrollToBottom();
-                            setIsNearBottom(true);
-                        }}
-                    >
-                        <ArrowDown size={24} />
-                    </button>
-                )}
+                <div className={`${props.customClassName} fixed bottom-[15%] z-10`}>
+                    {!isNearBottom && (
+                        <button
+                            title="Scroll to bottom"
+                            className="absolute bottom-0 right-0 bg-white dark:bg-[hsl(var(--background))] text-neutral-500 dark:text-white p-2 rounded-full shadow-xl"
+                            onClick={() => {
+                                scrollToBottom();
+                                setIsNearBottom(true);
+                            }}
+                        >
+                            <ArrowDown size={24} />
+                        </button>
+                    )}
+                </div>
            </div>
        </ScrollArea>
    );
--- a/src/interface/web/app/components/chatMessage/chatMessage.tsx
+++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx
@@ -26,6 +26,7 @@ import {
    Palette,
    ClipboardText,
    Check,
+    Shapes,
 } from "@phosphor-icons/react";

 import DOMPurify from "dompurify";
@@ -35,6 +36,7 @@ import { AgentData } from "@/app/agents/page";

 import renderMathInElement from "katex/contrib/auto-render";
 import "katex/dist/katex.min.css";
+import ExcalidrawComponent from "../excalidraw/excalidraw";

 const md = new markdownIt({
    html: true,
@@ -127,6 +129,8 @@ export interface StreamMessage {
    timestamp: string;
    agent?: AgentData;
    images?: string[];
+    intentType?: string;
+    inferredQueries?: string[];
 }

 export interface ChatHistoryData {
@@ -251,6 +255,10 @@ function chooseIconFromHeader(header: string, iconColor: string) {
        return <Aperture className={`${classNames}`} />;
    }

+    if (compareHeader.includes("diagram")) {
+        return <Shapes className={`${classNames}`} />;
+    }
+
    if (compareHeader.includes("paint")) {
        return <Palette className={`${classNames}`} />;
    }
@@ -282,6 +290,7 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
    const [markdownRendered, setMarkdownRendered] = useState<string>("");
    const [isPlaying, setIsPlaying] = useState<boolean>(false);
    const [interrupted, setInterrupted] = useState<boolean>(false);
+    const [excalidrawData, setExcalidrawData] = useState<string>("");

    const interruptedRef = useRef<boolean>(false);
    const messageRef = useRef<HTMLDivElement>(null);
@@ -320,6 +329,11 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
    useEffect(() => {
        let message = props.chatMessage.message;

+        if (props.chatMessage.intent && props.chatMessage.intent.type == "excalidraw") {
+            message = props.chatMessage.intent["inferred-queries"][0];
+            setExcalidrawData(props.chatMessage.message);
+        }
+
        // Replace LaTeX delimiters with placeholders
        message = message
            .replace(/\\\(/g, "LEFTPAREN")
@@ -337,22 +351,27 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
            message = `<div class="${styles.imagesContainer}">${imagesInMd}</div>${message}`;
        }

-        if (props.chatMessage.intent && props.chatMessage.intent.type == "text-to-image") {
-            message = `![generated image](data:image/png;base64,${message})`;
-        } else if (props.chatMessage.intent && props.chatMessage.intent.type == "text-to-image2") {
-            message = `![generated image](${message})`;
-        } else if (
-            props.chatMessage.intent &&
-            props.chatMessage.intent.type == "text-to-image-v3"
-        ) {
-            message = `![generated image](data:image/webp;base64,${message})`;
-        }
-        if (
-            props.chatMessage.intent &&
-            props.chatMessage.intent.type.includes("text-to-image") &&
-            props.chatMessage.intent["inferred-queries"]?.length > 0
-        ) {
-            message += `\n\n${props.chatMessage.intent["inferred-queries"][0]}`;
+        const intentTypeHandlers = {
+            "text-to-image": (msg: string) => `![generated image](data:image/png;base64,${msg})`,
+            "text-to-image2": (msg: string) => `![generated image](${msg})`,
+            "text-to-image-v3": (msg: string) =>
+                `![generated image](data:image/webp;base64,${msg})`,
+            excalidraw: (msg: string) => {
+                return msg;
+            },
+        };
+
+        if (props.chatMessage.intent) {
+            const { type, "inferred-queries": inferredQueries } = props.chatMessage.intent;
+
+            console.log("intent type", type);
+            if (type in intentTypeHandlers) {
+                message = intentTypeHandlers[type as keyof typeof intentTypeHandlers](message);
+            }
+
+            if (type.includes("text-to-image") && inferredQueries?.length > 0) {
+                message += `\n\n${inferredQueries[0]}`;
+            }
        }

        setTextRendered(message);
@@ -559,6 +578,7 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
                    className={styles.chatMessage}
                    dangerouslySetInnerHTML={{ __html: markdownRendered }}
                />
+                {excalidrawData && <ExcalidrawComponent data={excalidrawData} />}
            </div>
            <div className={styles.teaserReferencesContainer}>
                <TeaserReferencesSection
--- a/src/interface/web/app/components/excalidraw/excalidraw.tsx
+++ b/src/interface/web/app/components/excalidraw/excalidraw.tsx
@@ -0,0 +1,24 @@
+"use client";
+
+import dynamic from "next/dynamic";
+import { Suspense } from "react";
+import Loading from "../../components/loading/loading";
+
+// Since client components get prerenderd on server as well hence importing
+// the excalidraw stuff dynamically with ssr false
+
+const ExcalidrawWrapper = dynamic(() => import("./excalidrawWrapper").then((mod) => mod.default), {
+    ssr: false,
+});
+
+interface ExcalidrawComponentProps {
+    data: any;
+}
+
+export default function ExcalidrawComponent(props: ExcalidrawComponentProps) {
+    return (
+        <Suspense fallback={<Loading />}>
+            <ExcalidrawWrapper data={props.data} />
+        </Suspense>
+    );
+}
--- a/src/interface/web/app/components/excalidraw/excalidrawWrapper.tsx
+++ b/src/interface/web/app/components/excalidraw/excalidrawWrapper.tsx
@@ -0,0 +1,149 @@
+"use client";
+
+import { useState, useEffect } from "react";
+
+import dynamic from "next/dynamic";
+
+import { ExcalidrawProps } from "@excalidraw/excalidraw/types/types";
+import { ExcalidrawElement } from "@excalidraw/excalidraw/types/element/types";
+import { ExcalidrawElementSkeleton } from "@excalidraw/excalidraw/types/data/transform";
+
+const Excalidraw = dynamic<ExcalidrawProps>(
+    async () => (await import("@excalidraw/excalidraw")).Excalidraw,
+    {
+        ssr: false,
+    },
+);
+
+import { convertToExcalidrawElements } from "@excalidraw/excalidraw";
+
+import { Button } from "@/components/ui/button";
+
+import { ArrowsInSimple, ArrowsOutSimple } from "@phosphor-icons/react";
+
+interface ExcalidrawWrapperProps {
+    data: ExcalidrawElementSkeleton[];
+}
+
+export default function ExcalidrawWrapper(props: ExcalidrawWrapperProps) {
+    const [excalidrawElements, setExcalidrawElements] = useState<ExcalidrawElement[]>([]);
+    const [expanded, setExpanded] = useState<boolean>(false);
+
+    const isValidExcalidrawElement = (element: ExcalidrawElementSkeleton): boolean => {
+        return (
+            element.x !== undefined &&
+            element.y !== undefined &&
+            element.id !== undefined &&
+            element.type !== undefined
+        );
+    };
+
+    useEffect(() => {
+        if (expanded) {
+            onkeydown = (e) => {
+                if (e.key === "Escape") {
+                    setExpanded(false);
+                    // Trigger a resize event to make Excalidraw adjust its size
+                    window.dispatchEvent(new Event("resize"));
+                }
+            };
+        } else {
+            onkeydown = null;
+        }
+    }, [expanded]);
+
+    useEffect(() => {
+        // Do some basic validation
+        const basicValidSkeletons: ExcalidrawElementSkeleton[] = [];
+
+        for (const element of props.data) {
+            if (isValidExcalidrawElement(element as ExcalidrawElementSkeleton)) {
+                basicValidSkeletons.push(element as ExcalidrawElementSkeleton);
+            }
+        }
+
+        const validSkeletons: ExcalidrawElementSkeleton[] = [];
+        for (const element of basicValidSkeletons) {
+            if (element.type === "frame") {
+                continue;
+            }
+            if (element.type === "arrow") {
+                const start = basicValidSkeletons.find((child) => child.id === element.start?.id);
+                const end = basicValidSkeletons.find((child) => child.id === element.end?.id);
+                if (start && end) {
+                    validSkeletons.push(element);
+                }
+            } else {
+                validSkeletons.push(element);
+            }
+        }
+
+        for (const element of basicValidSkeletons) {
+            if (element.type === "frame") {
+                const children = element.children?.map((childId) => {
+                    return validSkeletons.find((child) => child.id === childId);
+                });
+                // Get the valid children, filter out any undefined values
+                const validChildrenIds: readonly string[] = children
+                    ?.map((child) => child?.id)
+                    .filter((id) => id !== undefined) as string[];
+
+                if (validChildrenIds === undefined || validChildrenIds.length === 0) {
+                    continue;
+                }
+
+                validSkeletons.push({
+                    ...element,
+                    children: validChildrenIds,
+                });
+            }
+        }
+
+        const elements = convertToExcalidrawElements(validSkeletons);
+        setExcalidrawElements(elements);
+    }, []);
+
+    return (
+        <div className="relative">
+            <div
+                className={`${expanded ? "fixed inset-0 bg-black bg-opacity-50 backdrop-blur-sm z-50 flex items-center justify-center" : ""}`}
+            >
+                <Button
+                    onClick={() => {
+                        setExpanded(!expanded);
+                        // Trigger a resize event to make Excalidraw adjust its size
+                        window.dispatchEvent(new Event("resize"));
+                    }}
+                    variant={"outline"}
+                    className={`${expanded ? "absolute top-2 left-2 z-[60]" : ""}`}
+                >
+                    {expanded ? (
+                        <ArrowsInSimple className="h-4 w-4" />
+                    ) : (
+                        <ArrowsOutSimple className="h-4 w-4" />
+                    )}
+                </Button>
+                <div
+                    className={`
+                    ${expanded ? "w-[80vw] h-[80vh]" : "w-full h-[500px]"}
+                    bg-white overflow-hidden rounded-lg relative
+                `}
+                >
+                    <Excalidraw
+                        initialData={{
+                            elements: excalidrawElements,
+                            appState: { zenModeEnabled: true },
+                            scrollToContent: true,
+                        }}
+                        // TODO - Create a common function to detect if the theme is dark?
+                        theme={localStorage.getItem("theme") === "dark" ? "dark" : "light"}
+                        validateEmbeddable={true}
+                        renderTopRightUI={(isMobile, appState) => {
+                            return <></>;
+                        }}
+                    />
+                </div>
+            </div>
+        </div>
+    );
+}
--- a/src/interface/web/app/share/chat/layout.tsx
+++ b/src/interface/web/app/share/chat/layout.tsx
@@ -27,7 +27,14 @@ export default function RootLayout({
                       child-src 'none';
                       object-src 'none';"
            ></meta>
-            <body className={inter.className}>{children}</body>
+            <body className={inter.className}>
+                {children}
+                <script
+                    dangerouslySetInnerHTML={{
+                        __html: `window.EXCALIDRAW_ASSET_PATH = 'https://assets.khoj.dev/@excalidraw/excalidraw/dist/';`,
+                    }}
+                />
+            </body>
        </html>
    );
 }
--- a/src/interface/web/app/share/chat/page.tsx
+++ b/src/interface/web/app/share/chat/page.tsx
@@ -39,6 +39,8 @@ function ChatBodyData(props: ChatBodyDataProps) {
    const setQueryToProcess = props.setQueryToProcess;
    const streamedMessages = props.streamedMessages;

+    const chatHistoryCustomClassName = props.isMobileWidth ? "w-full" : "w-4/6";
+
    useEffect(() => {
        if (images.length > 0) {
            const encodedImages = images.map((image) => encodeURIComponent(image));
@@ -96,10 +98,11 @@ function ChatBodyData(props: ChatBodyDataProps) {
                    setTitle={props.setTitle}
                    pendingMessage={processingMessage ? message : ""}
                    incomingMessages={props.streamedMessages}
+                    customClassName={chatHistoryCustomClassName}
                />
            </div>
            <div
-                className={`${styles.inputBox} p-1 md:px-2 shadow-md bg-background align-middle items-center justify-center dark:bg-neutral-700 dark:border-0 dark:shadow-sm rounded-t-2xl rounded-b-none md:rounded-xl`}
+                className={`${styles.inputBox} p-1 md:px-2 shadow-md bg-background align-middle items-center justify-center dark:bg-neutral-700 dark:border-0 dark:shadow-sm rounded-t-2xl rounded-b-none md:rounded-xl h-fit ${chatHistoryCustomClassName} mr-auto ml-auto`}
            >
                <ChatInputArea
                    isLoggedIn={props.isLoggedIn}
@@ -293,6 +296,19 @@ export default function SharedChat() {

            <div className={styles.chatBox}>
                <div className={styles.chatBoxBody}>
+                    {!isMobileWidth && title && (
+                        <div
+                            className={`${styles.chatTitleWrapper} text-nowrap text-ellipsis overflow-hidden max-w-screen-md grid items-top font-bold mr-8 pt-6 col-auto h-fit`}
+                        >
+                            {title && (
+                                <h2
+                                    className={`text-lg text-ellipsis whitespace-nowrap overflow-x-hidden`}
+                                >
+                                    {title}
+                                </h2>
+                            )}
+                        </div>
+                    )}
                    <Suspense fallback={<Loading />}>
                        <ChatBodyData
                            conversationId={conversationId}
--- a/src/interface/web/app/share/chat/sharedChat.module.css
+++ b/src/interface/web/app/share/chat/sharedChat.module.css
@@ -75,7 +75,7 @@ div.titleBar {
 div.chatBoxBody {
    display: grid;
    height: 100%;
-    width: 70%;
+    width: 95%;
    margin: auto;
 }