Research Mode [Part 2]: Improve Prompts, Edit Chat Messages. Set LLM Seed for Reproducibility

docker-compose.yml

@@ -14,6 +14,10 @@ services:
       interval: 30s
       timeout: 10s
       retries: 5
+  sandbox:
+    image: ghcr.io/khoj-ai/terrarium:latest
+    ports:
+      - "8080:8080"
   server:
     depends_on:
       database:

documentation/docs/features/chat.md

@@ -43,3 +43,6 @@ Slash commands allows you to change what Khoj uses to respond to your query
 - **/image**: Generate an image in response to your query.
 - **/help**: Use /help to get all available commands and general information about Khoj
 - **/summarize**: Can be used to summarize 1 selected file filter for that conversation. Refer to [File Summarization](summarization) for details.
+- **/diagram**: Generate a diagram in response to your query. This is built on [Excalidraw](https://excalidraw.com/).
+- **/code**: Generate and run very simple Python code snippets. Refer to [Code Generation](code_generation) for details.
+- **/research**: Go deeper in a topic for more accurate, in-depth responses.

documentation/docs/features/code_execution.md

@@ -0,0 +1,30 @@
+---
+---
+
+# Code Execution
+
+Khoj can generate and run very simple Python code snippets as well. This is useful if you want to generate a plot, run a simple calculation, or do some basic data manipulation. LLMs by default aren't skilled at complex quantitative tasks. Code generation & execution can come in handy for such tasks.
+
+Just use `/code` in your chat command.
+
+### Setup (Self-Hosting)
+Run [Cohere's Terrarium](https://github.com/cohere-ai/cohere-terrarium) on your machine to enable code generation and execution.
+
+Check the [instructions](https://github.com/cohere-ai/cohere-terrarium?tab=readme-ov-file#development) for running from source.
+
+For running with Docker, you can use our [docker-compose.yml](https://github.com/khoj-ai/khoj/blob/master/docker-compose.yml), or start it manually like this:
+
+```bash
+docker pull ghcr.io/khoj-ai/terrarium:latest
+docker run -d -p 8080:8080 ghcr.io/khoj-ai/terrarium:latest
+```
+
+#### Verify
+Verify that it's running, by evaluating a simple Python expression:
+
+```bash
+curl -X POST -H "Content-Type: application/json" \
+--url http://localhost:8080 \
+--data-raw '{"code": "1 + 1"}' \
+--no-buffer
+```

pyproject.toml

@@ -87,7 +87,7 @@ dependencies = [
     "django_apscheduler == 0.6.2",
     "anthropic == 0.26.1",
     "docx2txt == 0.8",
-    "google-generativeai == 0.7.2"
+    "google-generativeai == 0.8.3"
 ]
 dynamic = ["version"]
 

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

@@ -29,6 +29,7 @@ interface ChatBodyDataProps {
     onConversationIdChange?: (conversationId: string) => void;
     setQueryToProcess: (query: string) => void;
     streamedMessages: StreamMessage[];
+    setStreamedMessages: (messages: StreamMessage[]) => void;
     setUploadedFiles: (files: string[]) => void;
     isMobileWidth?: boolean;
     isLoggedIn: boolean;
@@ -118,6 +119,7 @@ function ChatBodyData(props: ChatBodyDataProps) {
                     setAgent={setAgentMetadata}
                     pendingMessage={processingMessage ? message : ""}
                     incomingMessages={props.streamedMessages}
+                    setIncomingMessages={props.setStreamedMessages}
                     customClassName={chatHistoryCustomClassName}
                 />
             </div>
@@ -351,6 +353,7 @@ export default function Chat() {
                         <ChatBodyData
                             isLoggedIn={authenticatedData !== null}
                             streamedMessages={messages}
+                            setStreamedMessages={setMessages}
                             chatOptionsData={chatOptionsData}
                             setTitle={setTitle}
                             setQueryToProcess={setQueryToProcess}

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

@@ -11,6 +11,11 @@ export interface RawReferenceData {
     codeContext?: CodeContext;
 }
 
+export interface MessageMetadata {
+    conversationId: string;
+    turnId: string;
+}
+
 export interface ResponseWithIntent {
     intentType: string;
     response: string;
@@ -90,6 +95,9 @@ export function processMessageChunk(
         if (references.onlineContext) onlineContext = references.onlineContext;
         if (references.codeContext) codeContext = references.codeContext;
         return { context, onlineContext, codeContext };
+    } else if (chunk.type === "metadata") {
+        const messageMetadata = chunk.data as MessageMetadata;
+        currentMessage.turnId = messageMetadata.turnId;
     } 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

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

@@ -42,6 +42,13 @@ export function converColorToBgGradient(color: string) {
     return `${convertToBGGradientClass(color)} dark:border dark:border-neutral-700`;
 }
 
+export function convertColorToCaretClass(color: string | undefined) {
+    if (color && tailwindColors.includes(color)) {
+        return `caret-${color}-500`;
+    }
+    return `caret-orange-500`;
+}
+
 export function convertColorToRingClass(color: string | undefined) {
     if (color && tailwindColors.includes(color)) {
         return `focus-visible:ring-${color}-500`;

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

@@ -34,8 +34,9 @@ interface ChatHistory {
 interface ChatHistoryProps {
     conversationId: string;
     setTitle: (title: string) => void;
-    incomingMessages?: StreamMessage[];
     pendingMessage?: string;
+    incomingMessages?: StreamMessage[];
+    setIncomingMessages?: (incomingMessages: StreamMessage[]) => void;
     publicConversationSlug?: string;
     setAgent: (agent: AgentData) => void;
     customClassName?: string;
@@ -45,7 +46,7 @@ interface TrainOfThoughtComponentProps {
     trainOfThought: string[];
     lastMessage: boolean;
     agentColor: string;
-    key: string;
+    keyId: string;
     completed?: boolean;
 }
 
@@ -56,7 +57,7 @@ function TrainOfThoughtComponent(props: TrainOfThoughtComponentProps) {
     return (
         <div
             className={`${!collapsed ? styles.trainOfThought + " shadow-sm" : ""}`}
-            key={props.key}
+            key={props.keyId}
         >
             {!props.completed && <InlineLoading className="float-right" />}
             {props.completed &&
@@ -97,6 +98,7 @@ export default function ChatHistory(props: ChatHistoryProps) {
     const [data, setData] = useState<ChatHistoryData | null>(null);
     const [currentPage, setCurrentPage] = useState(0);
     const [hasMoreMessages, setHasMoreMessages] = useState(true);
+    const [currentTurnId, setCurrentTurnId] = useState<string | null>(null);
     const sentinelRef = useRef<HTMLDivElement | null>(null);
     const scrollAreaRef = useRef<HTMLDivElement | null>(null);
     const latestUserMessageRef = useRef<HTMLDivElement | null>(null);
@@ -177,6 +179,10 @@ export default function ChatHistory(props: ChatHistoryProps) {
             if (lastMessage && !lastMessage.completed) {
                 setIncompleteIncomingMessageIndex(props.incomingMessages.length - 1);
                 props.setTitle(lastMessage.rawQuery);
+                // Store the turnId when we get it
+                if (lastMessage.turnId) {
+                    setCurrentTurnId(lastMessage.turnId);
+                }
             }
         }
     }, [props.incomingMessages]);
@@ -278,6 +284,25 @@ export default function ChatHistory(props: ChatHistoryProps) {
         return data.agent?.persona;
     }
 
+    const handleDeleteMessage = (turnId?: string) => {
+        if (!turnId) return;
+
+        setData((prevData) => {
+            if (!prevData || !turnId) return prevData;
+            return {
+                ...prevData,
+                chat: prevData.chat.filter((msg) => msg.turnId !== turnId),
+            };
+        });
+
+        // Update incoming messages if they exist
+        if (props.incomingMessages && props.setIncomingMessages) {
+            props.setIncomingMessages(
+                props.incomingMessages.filter((msg) => msg.turnId !== turnId),
+            );
+        }
+    };
+
     if (!props.conversationId && !props.publicConversationSlug) {
         return null;
     }
@@ -293,6 +318,18 @@ export default function ChatHistory(props: ChatHistoryProps) {
                         data.chat &&
                         data.chat.map((chatMessage, index) => (
                             <>
+                                {chatMessage.trainOfThought && chatMessage.by === "khoj" && (
+                                    <TrainOfThoughtComponent
+                                        trainOfThought={chatMessage.trainOfThought?.map(
+                                            (train) => train.data,
+                                        )}
+                                        lastMessage={false}
+                                        agentColor={data?.agent?.color || "orange"}
+                                        key={`${index}trainOfThought`}
+                                        keyId={`${index}trainOfThought`}
+                                        completed={true}
+                                    />
+                                )}
                                 <ChatMessage
                                     key={`${index}fullHistory`}
                                     ref={
@@ -312,22 +349,14 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                     customClassName="fullHistory"
                                     borderLeftColor={`${data?.agent?.color}-500`}
                                     isLastMessage={index === data.chat.length - 1}
+                                    onDeleteMessage={handleDeleteMessage}
+                                    conversationId={props.conversationId}
                                 />
-                                {chatMessage.trainOfThought && chatMessage.by === "khoj" && (
-                                    <TrainOfThoughtComponent
-                                        trainOfThought={chatMessage.trainOfThought?.map(
-                                            (train) => train.data,
-                                        )}
-                                        lastMessage={false}
-                                        agentColor={data?.agent?.color || "orange"}
-                                        key={`${index}trainOfThought`}
-                                        completed={true}
-                                    />
-                                )}
                             </>
                         ))}
                     {props.incomingMessages &&
                         props.incomingMessages.map((message, index) => {
+                            const messageTurnId = message.turnId ?? currentTurnId ?? undefined;
                             return (
                                 <React.Fragment key={`incomingMessage${index}`}>
                                     <ChatMessage
@@ -342,16 +371,22 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                             by: "you",
                                             automationId: "",
                                             images: message.images,
+                                            conversationId: props.conversationId,
+                                            turnId: messageTurnId,
                                         }}
                                         customClassName="fullHistory"
                                         borderLeftColor={`${data?.agent?.color}-500`}
+                                        onDeleteMessage={handleDeleteMessage}
+                                        conversationId={props.conversationId}
+                                        turnId={messageTurnId}
                                     />
                                     {message.trainOfThought && (
                                         <TrainOfThoughtComponent
                                             trainOfThought={message.trainOfThought}
                                             lastMessage={index === incompleteIncomingMessageIndex}
                                             agentColor={data?.agent?.color || "orange"}
                                             key={`${index}trainOfThought`}
+                                            keyId={`${index}trainOfThought`}
                                             completed={message.completed}
                                         />
                                     )}
@@ -373,7 +408,12 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                                 "memory-type": "",
                                                 "inferred-queries": message.inferredQueries || [],
                                             },
+                                            conversationId: props.conversationId,
+                                            turnId: messageTurnId,
                                         }}
+                                        conversationId={props.conversationId}
+                                        turnId={messageTurnId}
+                                        onDeleteMessage={handleDeleteMessage}
                                         customClassName="fullHistory"
                                         borderLeftColor={`${data?.agent?.color}-500`}
                                         isLastMessage={true}
@@ -393,7 +433,11 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                 created: new Date().getTime().toString(),
                                 by: "you",
                                 automationId: "",
+                                conversationId: props.conversationId,
+                                turnId: undefined,
                             }}
+                            conversationId={props.conversationId}
+                            onDeleteMessage={handleDeleteMessage}
                             customClassName="fullHistory"
                             borderLeftColor={`${data?.agent?.color}-500`}
                             isLastMessage={true}

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

@@ -149,7 +149,7 @@ export const ChatInputArea = forwardRef<HTMLTextAreaElement, ChatInputProps>((pr
         }
 
         let messageToSend = message.trim();
-        if (useResearchMode) {
+        if (useResearchMode && !messageToSend.startsWith("/research")) {
             messageToSend = `/research ${messageToSend}`;
         }
 
@@ -398,7 +398,7 @@ export const ChatInputArea = forwardRef<HTMLTextAreaElement, ChatInputProps>((pr
                         <PopoverContent
                             onOpenAutoFocus={(e) => e.preventDefault()}
                             className={`${props.isMobileWidth ? "w-[100vw]" : "w-full"} rounded-md`}
-                            side="top"
+                            side="bottom"
                             align="center"
                             /* Offset below text area on home page (i.e where conversationId is unset) */
                             sideOffset={props.conversationId ? 0 : 80}
@@ -590,8 +590,8 @@ export const ChatInputArea = forwardRef<HTMLTextAreaElement, ChatInputProps>((pr
                             </Button>
                         </TooltipTrigger>
                         <TooltipContent className="text-xs">
-                            Research Mode allows you to get more deeply researched, detailed
-                            responses. Response times may be longer.
+                            (Experimental) Research Mode allows you to get more deeply researched,
+                            detailed responses. Response times may be longer.
                         </TooltipContent>
                     </Tooltip>
                 </TooltipProvider>

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

@@ -29,6 +29,7 @@ import {
     Check,
     Code,
     Shapes,
+    Trash,
 } from "@phosphor-icons/react";
 
 import DOMPurify from "dompurify";
@@ -146,6 +147,8 @@ export interface SingleChatMessage {
     intent?: Intent;
     agent?: AgentData;
     images?: string[];
+    conversationId: string;
+    turnId?: string;
 }
 
 export interface StreamMessage {
@@ -161,6 +164,7 @@ export interface StreamMessage {
     images?: string[];
     intentType?: string;
     inferredQueries?: string[];
+    turnId?: string;
 }
 
 export interface ChatHistoryData {
@@ -242,6 +246,9 @@ interface ChatMessageProps {
     borderLeftColor?: string;
     isLastMessage?: boolean;
     agent?: AgentData;
+    onDeleteMessage: (turnId?: string) => void;
+    conversationId: string;
+    turnId?: string;
 }
 
 interface TrainOfThoughtProps {
@@ -654,6 +661,27 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
         });
     }
 
+    const deleteMessage = async (message: SingleChatMessage) => {
+        const turnId = message.turnId || props.turnId;
+        const response = await fetch("/api/chat/conversation/message", {
+            method: "DELETE",
+            headers: {
+                "Content-Type": "application/json",
+            },
+            body: JSON.stringify({
+                conversation_id: props.conversationId,
+                turn_id: turnId,
+            }),
+        });
+
+        if (response.ok) {
+            // Update the UI after successful deletion
+            props.onDeleteMessage(turnId);
+        } else {
+            console.error("Failed to delete message");
+        }
+    };
+
     const allReferences = constructAllReferences(
         props.chatMessage.context,
         props.chatMessage.onlineContext,
@@ -716,6 +744,18 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
                                         />
                                     </button>
                                 ))}
+                            {props.chatMessage.turnId && (
+                                <button
+                                    title="Delete"
+                                    className={`${styles.deleteButton}`}
+                                    onClick={() => deleteMessage(props.chatMessage)}
+                                >
+                                    <Trash
+                                        alt="Delete Message"
+                                        className="hsl(var(--muted-foreground)) hover:text-red-500"
+                                    />
+                                </button>
+                            )}
                             <button
                                 title="Copy"
                                 className={`${styles.copyButton}`}