Merge branch 'master' of github.com:khoj-ai/khoj into features/big-upgrade-chat-ux

.github/workflows/pypi.yml

@@ -27,7 +27,7 @@ jobs:
     permissions:
       id-token: write
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
@@ -36,16 +36,12 @@ jobs:
         with:
           python-version: '3.11'
 
-      - name: ⬇️ Install Application
+      - name: ⬇️ Install Server
         run: python -m pip install --upgrade pip && pip install --upgrade .
 
-      - name: Install the Next.js application
+      - name: ⬇️ Install Web Client
         run: |
           yarn install
-        working-directory: src/interface/web
-
-      - name: Build & export static Next.js app to Django static assets
-        run: |
           yarn ciexport
         working-directory: src/interface/web
 
@@ -56,7 +52,12 @@ jobs:
           export SOURCE_DATE_EPOCH=$(git log -1 --pretty=%ct)
           rm -rf dist
 
-          # Build PyPi Package
+          # Build PyPI Package: khoj
+          pipx run build
+
+          # Build legacy PyPI Package: khoj-assistant
+          sed -i.bak '/^name = "khoj"$/s//name = "khoj-assistant"/' pyproject.toml
+          rm pyproject.toml.bak
           pipx run build
 
       - name: 🌡️ Validate Python Package
@@ -66,11 +67,11 @@ jobs:
           pipx run twine check dist/*
 
       - name: ⏫ Upload Python Package Artifacts
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: khoj-assistant
+          name: khoj
-          path: dist/*.whl
+          path: dist/khoj-*.whl
 
-      - name: 📦 Publish Python Package to PyPI
+      - name: 📦 Publish Python Packages to PyPI
         if: startsWith(github.ref, 'refs/tags') || github.ref == 'refs/heads/master'
         uses: pypa/gh-action-pypi-publish@v1.8.14

Dockerfile

@@ -1,12 +1,14 @@
 # syntax=docker/dockerfile:1
 FROM ubuntu:jammy
-LABEL org.opencontainers.image.source https://github.com/khoj-ai/khoj
+LABEL homepage="https://khoj.dev"
+LABEL repository="https://github.com/khoj-ai/khoj"
+LABEL org.opencontainers.image.source="https://github.com/khoj-ai/khoj"
 
 # Install System Dependencies
 RUN apt update -y && apt -y install python3-pip swig curl
 
 # Install Node.js and Yarn
-RUN curl -sL https://deb.nodesource.com/setup_22.x | bash -
+RUN curl -sL https://deb.nodesource.com/setup_20.x | bash -
 RUN apt -y install nodejs
 RUN curl -sS https://dl.yarnpkg.com/debian/pubkey.gpg | apt-key add -
 RUN echo "deb https://dl.yarnpkg.com/debian/ stable main" | tee /etc/apt/sources.list.d/yarn.list
@@ -31,7 +33,7 @@ ENV PYTHONPATH=/app/src:$PYTHONPATH
 
 # Go to the directory src/interface/web and export the built Next.js assets
 WORKDIR /app/src/interface/web
-RUN bash -c "yarn install && yarn ciexport"
+RUN bash -c "yarn cache clean && yarn install --verbose && yarn ciexport"
 WORKDIR /app
 
 # Run the Application

README.md

@@ -4,7 +4,7 @@
 
 [![test](https://github.com/khoj-ai/khoj/actions/workflows/test.yml/badge.svg)](https://github.com/khoj-ai/khoj/actions/workflows/test.yml)
 [![dockerize](https://github.com/khoj-ai/khoj/actions/workflows/dockerize.yml/badge.svg)](https://github.com/khoj-ai/khoj/pkgs/container/khoj)
-[![pypi](https://github.com/khoj-ai/khoj/actions/workflows/pypi.yml/badge.svg)](https://pypi.org/project/khoj-assistant/)
+[![pypi](https://github.com/khoj-ai/khoj/actions/workflows/pypi.yml/badge.svg)](https://pypi.org/project/khoj/)
 ![Discord](https://img.shields.io/discord/1112065956647284756?style=plastic&label=discord)
 
 </div>

documentation/docs/advanced/authentication.md

@@ -41,7 +41,7 @@ To set up your self-hosted Khoj with Google Auth, you need to create a project i
 To implement this, you'll need to:
 1. You must use the `python` package or build from source, because you'll need to install additional packages for the google auth libraries (`prod`). The syntax to install the right packages is
    ```
-   pip install khoj-assistant[prod]
+   pip install khoj[prod]
    ```
 2. [Create authorization credentials](https://developers.google.com/identity/sign-in/web/sign-in) for your application.
 3. Open your [Google cloud console](https://console.developers.google.com/apis/credentials) and create a configuration like below for the relevant `OAuth 2.0 Client IDs` project:

documentation/docs/contributing/development.mdx

@@ -229,7 +229,7 @@ The core code for the Obsidian plugin is under `src/interface/obsidian`. The fil
 4. Open the `khoj` folder in the file explorer that opens. You'll see a file called `main.js` in this folder. To test your changes, replace this file with the `main.js` file that was generated by the development server in the previous section.
 
 ## Create Khoj Release (Only for Maintainers)
-Follow the steps below to [release](https://github.com/debanjum/khoj/releases/) Khoj. This will create a stable release of Khoj on [Pypi](https://pypi.org/project/khoj-assistant/), [Melpa](https://stable.melpa.org/#%252Fkhoj) and [Obsidian](https://obsidian.md/plugins?id%253Dkhoj). It will also create desktop apps of Khoj and attach them to the latest release.
+Follow the steps below to [release](https://github.com/debanjum/khoj/releases/) Khoj. This will create a stable release of Khoj on [Pypi](https://pypi.org/project/khoj/), [Melpa](https://stable.melpa.org/#%252Fkhoj) and [Obsidian](https://obsidian.md/plugins?id%253Dkhoj). It will also create desktop apps of Khoj and attach them to the latest release.
 
 1. Create and tag release commit by running the bump_version script. The release commit sets version number in required metadata files.
   ```shell

documentation/docs/get-started/setup.mdx

@@ -105,10 +105,10 @@ Run the following command in your terminal to install the Khoj server.
     <TabItem value="macos" label="MacOS">
     ```shell
 # ARM/M1+ Machines
-MAKE_ARGS="-DLLAMA_METAL=on" python -m pip install khoj-assistant
+MAKE_ARGS="-DLLAMA_METAL=on" python -m pip install khoj
 
 # Intel Machines
-python -m pip install khoj-assistant
+python -m pip install khoj
     ```
     </TabItem>
     <TabItem value="win" label="Windows">
@@ -122,19 +122,19 @@ python -m pip install khoj-assistant
  $env:CMAKE_ARGS = "-DLLAMA_VULKAN=on"
 
  # 2. Install Khoj
- py -m pip install khoj-assistant
+ py -m pip install khoj
       ```
     </TabItem>
     <TabItem value="unix" label="Linux">
       ```shell
  # CPU
- python -m pip install khoj-assistant
+ python -m pip install khoj
  # NVIDIA (CUDA) GPU
- CMAKE_ARGS="DLLAMA_CUDA=on" FORCE_CMAKE=1 python -m pip install khoj-assistant
+ CMAKE_ARGS="DLLAMA_CUDA=on" FORCE_CMAKE=1 python -m pip install khoj
  # AMD (ROCm) GPU
- CMAKE_ARGS="-DLLAMA_HIPBLAS=on" FORCE_CMAKE=1 python -m pip install khoj-assistant
+ CMAKE_ARGS="-DLLAMA_HIPBLAS=on" FORCE_CMAKE=1 python -m pip install khoj
  # VULCAN GPU
- CMAKE_ARGS="-DLLAMA_VULKAN=on" FORCE_CMAKE=1 python -m pip install khoj-assistant
+ CMAKE_ARGS="-DLLAMA_VULKAN=on" FORCE_CMAKE=1 python -m pip install khoj
       ```
     </TabItem>
   </Tabs>
@@ -257,7 +257,7 @@ Set the host URL on your clients settings page to your Khoj server URL. By defau
   <Tabs groupId="environment">
     <TabItem value="localsetup" label="Local Setup">
       ```shell
-      pip install --upgrade khoj-assistant
+      pip install --upgrade khoj
       ```
       *Note: To upgrade to the latest pre-release version of the khoj server run below command*
     </TabItem>
@@ -285,7 +285,7 @@ Set the host URL on your clients settings page to your Khoj server URL. By defau
     <TabItem value="localsetup" label="Local Setup">
     ```shell
     # uninstall khoj server
-    pip uninstall khoj-assistant
+    pip uninstall khoj
 
     # delete khoj postgres db
     dropdb khoj -U postgres
@@ -318,13 +318,13 @@ Set the host URL on your clients settings page to your Khoj server URL. By defau
   1. Install [pipx](https://pypa.github.io/pipx/#install-pipx)
   2. Use `pipx` to install Khoj to avoid dependency conflicts with other python packages.
      ```shell
-     pipx install khoj-assistant
+     pipx install khoj
      ```
   3. Now start `khoj` using the standard steps described earlier
 
 
 #### Install fails while building Tokenizer dependency
-- **Details**: `pip install khoj-assistant` fails while building the `tokenizers` dependency. Complains about Rust.
+- **Details**: `pip install khoj` fails while building the `tokenizers` dependency. Complains about Rust.
 - **Fix**: Install Rust to build the tokenizers package. For example on Mac run:
     ```shell
     brew install rustup

manifest.json

@@ -1,7 +1,7 @@
 {
 	"id": "khoj",
 	"name": "Khoj",
-	"version": "1.16.0",
+	"version": "1.17.0",
 	"minAppVersion": "0.15.0",
 	"description": "An AI copilot for your Second Brain",
 	"author": "Khoj Inc.",

prod.Dockerfile

@@ -1,12 +1,12 @@
 FROM ubuntu:jammy
 
-LABEL org.opencontainers.image.source https://github.com/khoj-ai/khoj
+LABEL org.opencontainers.image.source="https://github.com/khoj-ai/khoj"
 
 # Install System Dependencies
 RUN apt update -y && apt -y install python3-pip libsqlite3-0 ffmpeg libsm6 libxext6 swig curl
 
 # Install Node.js and Yarn
-RUN curl -sL https://deb.nodesource.com/setup_22.x | bash -
+RUN curl -sL https://deb.nodesource.com/setup_20.x | bash -
 RUN apt -y install nodejs
 RUN curl -sS https://dl.yarnpkg.com/debian/pubkey.gpg | apt-key add -
 RUN echo "deb https://dl.yarnpkg.com/debian/ stable main" | tee /etc/apt/sources.list.d/yarn.list
@@ -29,7 +29,7 @@ ENV PYTHONPATH=/app/src:$PYTHONPATH
 
 # Go to the directory src/interface/web and export the built Next.js assets
 WORKDIR /app/src/interface/web
-RUN bash -c "yarn install && yarn ciexport"
+RUN bash -c "yarn cache clean && yarn install --verbose && yarn ciexport"
 WORKDIR /app
 
 # Run the Application

pyproject.toml

@@ -3,7 +3,7 @@ requires = ["hatchling", "hatch-vcs"]
 build-backend = "hatchling.build"
 
 [project]
-name = "khoj-assistant"
+name = "khoj"
 description = "An AI copilot for your Second Brain"
 readme = "README.md"
 license = "AGPL-3.0-or-later"
@@ -27,7 +27,6 @@ classifiers = [
     "License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)",
     "Operating System :: OS Independent",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
@@ -67,7 +66,7 @@ dependencies = [
     "pymupdf >= 1.23.5",
     "django == 5.0.7",
     "authlib == 1.2.1",
-    "llama-cpp-python == 0.2.76",
+    "llama-cpp-python == 0.2.82",
     "itsdangerous == 2.1.2",
     "httpx == 0.25.0",
     "pgvector == 0.2.4",
@@ -110,7 +109,7 @@ prod = [
     "resend == 1.0.1",
 ]
 dev = [
-    "khoj-assistant[prod]",
+    "khoj[prod]",
     "pytest >= 7.1.2",
     "pytest-xdist[psutil]",
     "pytest-django == 4.5.2",

src/interface/desktop/chat.html

@@ -61,6 +61,14 @@
         let city = null;
         let countryName = null;
         let timezone = null;
+        let chatMessageState = {
+            newResponseTextEl: null,
+            newResponseEl: null,
+            loadingEllipsis: null,
+            references: {},
+            rawResponse: "",
+            isVoice: false,
+        }
 
         fetch("https://ipapi.co/json")
             .then(response => response.json())
@@ -75,10 +83,9 @@
                 return;
             });
 
-        async function chat() {
+        async function chat(isVoice=false) {
-            // Extract required fields for search from form
+            // Extract chat message from chat input form
             let query = document.getElementById("chat-input").value.trim();
-            let resultsCount = localStorage.getItem("khojResultsCount") || 5;
             console.log(`Query: ${query}`);
 
             // Short circuit on empty query
@@ -106,9 +113,6 @@
                 await refreshChatSessionsPanel();
             }
 
-            // Generate backend API URL to execute query
-            let chatApi = `${hostURL}/api/chat?q=${encodeURIComponent(query)}&n=${resultsCount}&client=web&stream=true&conversation_id=${conversationID}&region=${region}&city=${city}&country=${countryName}&timezone=${timezone}`;
-
             let newResponseEl = document.createElement("div");
             newResponseEl.classList.add("chat-message", "khoj");
             newResponseEl.attributes["data-meta"] = "🏮 Khoj at " + formatDate(new Date());
@@ -119,25 +123,7 @@
             newResponseEl.appendChild(newResponseTextEl);
 
             // Temporary status message to indicate that Khoj is thinking
-            let loadingEllipsis = document.createElement("div");
+            let loadingEllipsis = createLoadingEllipsis();
-            loadingEllipsis.classList.add("lds-ellipsis");
-
-            let firstEllipsis = document.createElement("div");
-            firstEllipsis.classList.add("lds-ellipsis-item");
-
-            let secondEllipsis = document.createElement("div");
-            secondEllipsis.classList.add("lds-ellipsis-item");
-
-            let thirdEllipsis = document.createElement("div");
-            thirdEllipsis.classList.add("lds-ellipsis-item");
-
-            let fourthEllipsis = document.createElement("div");
-            fourthEllipsis.classList.add("lds-ellipsis-item");
-
-            loadingEllipsis.appendChild(firstEllipsis);
-            loadingEllipsis.appendChild(secondEllipsis);
-            loadingEllipsis.appendChild(thirdEllipsis);
-            loadingEllipsis.appendChild(fourthEllipsis);
 
             newResponseTextEl.appendChild(loadingEllipsis);
             document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
@@ -148,107 +134,36 @@
             let chatInput = document.getElementById("chat-input");
             chatInput.classList.remove("option-enabled");
 
+            // Setup chat message state
+            chatMessageState = {
+                newResponseTextEl,
+                newResponseEl,
+                loadingEllipsis,
+                references: {},
+                rawResponse: "",
+                rawQuery: query,
+                isVoice: isVoice,
+            }
+
             // Call Khoj chat API
-            let response = await fetch(chatApi, { headers });
+            let chatApi = `${hostURL}/api/chat?q=${encodeURIComponent(query)}&conversation_id=${conversationID}&stream=true&client=desktop`;
-            let rawResponse = "";
+            chatApi += (!!region && !!city && !!countryName && !!timezone)
-            let references = null;
+                ? `&region=${region}&city=${city}&country=${countryName}&timezone=${timezone}`
-            const contentType = response.headers.get("content-type");
+                : '';
 
-            if (contentType === "application/json") {
+            const response = await fetch(chatApi, { headers });
-                // Handle JSON response
-                try {
-                    const responseAsJson = await response.json();
-                    if (responseAsJson.image) {
-                        // If response has image field, response is a generated image.
-                        if (responseAsJson.intentType === "text-to-image") {
-                            rawResponse += `![${query}](data:image/png;base64,${responseAsJson.image})`;
-                        } else if (responseAsJson.intentType === "text-to-image2") {
-                            rawResponse += `![${query}](${responseAsJson.image})`;
-                        } else if (responseAsJson.intentType === "text-to-image-v3") {
-                            rawResponse += `![${query}](data:image/webp;base64,${responseAsJson.image})`;
-                        }
-                        const inferredQueries = responseAsJson.inferredQueries?.[0];
-                        if (inferredQueries) {
-                            rawResponse += `\n\n**Inferred Query**:\n\n${inferredQueries}`;
-                        }
-                    }
-                    if (responseAsJson.context) {
-                        const rawReferenceAsJson = responseAsJson.context;
-                        references = createReferenceSection(rawReferenceAsJson);
-                    }
-                    if (responseAsJson.detail) {
-                        // If response has detail field, response is an error message.
-                        rawResponse += responseAsJson.detail;
-                    }
-                } catch (error) {
-                    // If the chunk is not a JSON object, just display it as is
-                    rawResponse += chunk;
-                } finally {
-                    newResponseTextEl.innerHTML = "";
-                    newResponseTextEl.appendChild(formatHTMLMessage(rawResponse));
 
-                    if (references != null) {
+            try {
-                        newResponseTextEl.appendChild(references);
+                if (!response.ok) throw new Error(response.statusText);
-                    }
+                if (!response.body) throw new Error("Response body is empty");
-
+                // Stream and render chat response
-                    document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
+                await readChatStream(response);
-                    document.getElementById("chat-input").removeAttribute("disabled");
+            } catch (err) {
-                }
+                console.error(`Khoj chat response failed with\n${err}`);
-            } else {
+                if (chatMessageState.newResponseEl.getElementsByClassName("lds-ellipsis").length > 0 && chatMessageState.loadingEllipsis)
-                // Handle streamed response of type text/event-stream or text/plain
+                    chatMessageState.newResponseTextEl.removeChild(chatMessageState.loadingEllipsis);
-                const reader = response.body.getReader();
+                let errorMsg = "Sorry, unable to get response from Khoj backend ❤️‍🩹. Retry or contact developers for help at <a href=mailto:'team@khoj.dev'>team@khoj.dev</a> or <a href='https://discord.gg/BDgyabRM6e'>on Discord</a>";
-                const decoder = new TextDecoder();
+                newResponseTextEl.textContent = errorMsg;
-                let references = {};
-
-                readStream();
-
-                function readStream() {
-                    reader.read().then(({ done, value }) => {
-                        if (done) {
-                            // Append any references after all the data has been streamed
-                            if (references != {}) {
-                                newResponseTextEl.appendChild(createReferenceSection(references));
-                            }
-                            document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
-                            document.getElementById("chat-input").removeAttribute("disabled");
-                            return;
-                        }
-
-                        // Decode message chunk from stream
-                        const chunk = decoder.decode(value, { stream: true });
-
-                        if (chunk.includes("### compiled references:")) {
-                            const additionalResponse = chunk.split("### compiled references:")[0];
-                            rawResponse += additionalResponse;
-                            newResponseTextEl.innerHTML = "";
-                            newResponseTextEl.appendChild(formatHTMLMessage(rawResponse));
-
-                            const rawReference = chunk.split("### compiled references:")[1];
-                            const rawReferenceAsJson = JSON.parse(rawReference);
-                            if (rawReferenceAsJson instanceof Array) {
-                                references["notes"] = rawReferenceAsJson;
-                            } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
-                                references["online"] = rawReferenceAsJson;
-                            }
-                            readStream();
-                        } else {
-                            // Display response from Khoj
-                            if (newResponseTextEl.getElementsByClassName("lds-ellipsis").length > 0) {
-                                newResponseTextEl.removeChild(loadingEllipsis);
-                            }
-
-                            // If the chunk is not a JSON object, just display it as is
-                            rawResponse += chunk;
-                            newResponseTextEl.innerHTML = "";
-                            newResponseTextEl.appendChild(formatHTMLMessage(rawResponse));
-
-                            readStream();
-                        }
-
-                        // Scroll to bottom of chat window as chat response is streamed
-                        document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
-                    });
-                }
             }
         }
 

src/interface/desktop/chatutils.js

@@ -364,3 +364,194 @@ function createReferenceSection(references, createLinkerSection=false) {
 
     return referencesDiv;
 }
+
+function createLoadingEllipsis() {
+    let loadingEllipsis = document.createElement("div");
+    loadingEllipsis.classList.add("lds-ellipsis");
+
+    let firstEllipsis = document.createElement("div");
+    firstEllipsis.classList.add("lds-ellipsis-item");
+
+    let secondEllipsis = document.createElement("div");
+    secondEllipsis.classList.add("lds-ellipsis-item");
+
+    let thirdEllipsis = document.createElement("div");
+    thirdEllipsis.classList.add("lds-ellipsis-item");
+
+    let fourthEllipsis = document.createElement("div");
+    fourthEllipsis.classList.add("lds-ellipsis-item");
+
+    loadingEllipsis.appendChild(firstEllipsis);
+    loadingEllipsis.appendChild(secondEllipsis);
+    loadingEllipsis.appendChild(thirdEllipsis);
+    loadingEllipsis.appendChild(fourthEllipsis);
+
+    return loadingEllipsis;
+}
+
+function handleStreamResponse(newResponseElement, rawResponse, rawQuery, loadingEllipsis, replace=true) {
+    if (!newResponseElement) return;
+    // Remove loading ellipsis if it exists
+    if (newResponseElement.getElementsByClassName("lds-ellipsis").length > 0 && loadingEllipsis)
+        newResponseElement.removeChild(loadingEllipsis);
+    // Clear the response element if replace is true
+    if (replace) newResponseElement.innerHTML = "";
+
+    // Append response to the response element
+    newResponseElement.appendChild(formatHTMLMessage(rawResponse, false, replace, rawQuery));
+
+    // Append loading ellipsis if it exists
+    if (!replace && loadingEllipsis) newResponseElement.appendChild(loadingEllipsis);
+    // Scroll to bottom of chat view
+    document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
+}
+
+function handleImageResponse(imageJson, 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) {
+            rawResponse += `\n\n**Inferred Query**:\n\n${inferredQuery}`;
+        }
+    }
+
+    // If response has detail field, response is an error message.
+    if (imageJson.detail) rawResponse += imageJson.detail;
+
+    return rawResponse;
+}
+
+function finalizeChatBodyResponse(references, newResponseElement) {
+    if (!!newResponseElement && references != null && Object.keys(references).length > 0) {
+        newResponseElement.appendChild(createReferenceSection(references));
+    }
+    document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
+    document.getElementById("chat-input")?.removeAttribute("disabled");
+}
+
+function convertMessageChunkToJson(rawChunk) {
+    // Split the chunk into lines
+    if (rawChunk?.startsWith("{") && rawChunk?.endsWith("}")) {
+        try {
+            let jsonChunk = JSON.parse(rawChunk);
+            if (!jsonChunk.type)
+                jsonChunk = {type: 'message', data: jsonChunk};
+            return jsonChunk;
+        } catch (e) {
+            return {type: 'message', data: rawChunk};
+        }
+    } else if (rawChunk.length > 0) {
+        return {type: 'message', data: rawChunk};
+    }
+}
+
+function processMessageChunk(rawChunk) {
+    const chunk = convertMessageChunkToJson(rawChunk);
+    console.debug("Chunk:", chunk);
+    if (!chunk || !chunk.type) return;
+    if (chunk.type ==='status') {
+        console.log(`status: ${chunk.data}`);
+        const statusMessage = chunk.data;
+        handleStreamResponse(chatMessageState.newResponseTextEl, statusMessage, chatMessageState.rawQuery, chatMessageState.loadingEllipsis, false);
+    } else if (chunk.type === 'start_llm_response') {
+        console.log("Started streaming", new Date());
+    } else if (chunk.type === 'end_llm_response') {
+        console.log("Stopped streaming", new Date());
+
+        // Automatically respond with voice if the subscribed user has sent voice message
+        if (chatMessageState.isVoice && "{{ is_active }}" == "True")
+            textToSpeech(chatMessageState.rawResponse);
+
+        // Append any references after all the data has been streamed
+        finalizeChatBodyResponse(chatMessageState.references, chatMessageState.newResponseTextEl);
+
+        const liveQuery = chatMessageState.rawQuery;
+        // Reset variables
+        chatMessageState = {
+            newResponseTextEl: null,
+            newResponseEl: null,
+            loadingEllipsis: null,
+            references: {},
+            rawResponse: "",
+            rawQuery: liveQuery,
+            isVoice: false,
+        }
+    } else if (chunk.type === "references") {
+        chatMessageState.references = {"notes": chunk.data.context, "online": chunk.data.onlineContext};
+    } else if (chunk.type === 'message') {
+        const chunkData = chunk.data;
+        if (typeof chunkData === 'object' && chunkData !== null) {
+            // If chunkData is already a JSON object
+            handleJsonResponse(chunkData);
+        } else if (typeof chunkData  === 'string' && chunkData.trim()?.startsWith("{") && chunkData.trim()?.endsWith("}")) {
+            // Try process chunk data as if it is a JSON object
+            try {
+                const jsonData = JSON.parse(chunkData.trim());
+                handleJsonResponse(jsonData);
+            } catch (e) {
+                chatMessageState.rawResponse += chunkData;
+                handleStreamResponse(chatMessageState.newResponseTextEl, chatMessageState.rawResponse, chatMessageState.rawQuery, chatMessageState.loadingEllipsis);
+            }
+        } else {
+            chatMessageState.rawResponse += chunkData;
+            handleStreamResponse(chatMessageState.newResponseTextEl, chatMessageState.rawResponse, chatMessageState.rawQuery, chatMessageState.loadingEllipsis);
+        }
+    }
+}
+
+function handleJsonResponse(jsonData) {
+    if (jsonData.image || jsonData.detail) {
+        chatMessageState.rawResponse = handleImageResponse(jsonData, chatMessageState.rawResponse);
+    } else if (jsonData.response) {
+        chatMessageState.rawResponse = jsonData.response;
+    }
+
+    if (chatMessageState.newResponseTextEl) {
+        chatMessageState.newResponseTextEl.innerHTML = "";
+        chatMessageState.newResponseTextEl.appendChild(formatHTMLMessage(chatMessageState.rawResponse));
+    }
+}
+
+async function readChatStream(response) {
+    if (!response.body) return;
+    const reader = response.body.getReader();
+    const decoder = new TextDecoder();
+    const eventDelimiter = '␃🔚␗';
+    let buffer = '';
+
+    while (true) {
+        const { value, done } = await reader.read();
+        // If the stream is done
+        if (done) {
+            // Process the last chunk
+            processMessageChunk(buffer);
+            buffer = '';
+            break;
+        }
+
+        // Read chunk from stream and append it to the buffer
+        const chunk = decoder.decode(value, { stream: true });
+        console.debug("Raw Chunk:", chunk)
+        // Start buffering chunks until complete event is received
+        buffer += chunk;
+
+        // Once the buffer contains a complete event
+        let newEventIndex;
+        while ((newEventIndex = buffer.indexOf(eventDelimiter)) !== -1) {
+            // Extract the event from the buffer
+            const event = buffer.slice(0, newEventIndex);
+            buffer = buffer.slice(newEventIndex + eventDelimiter.length);
+
+            // Process the event
+            if (event) processMessageChunk(event);
+        }
+    }
+}

src/interface/desktop/package.json

@@ -1,6 +1,6 @@
 {
   "name": "Khoj",
-  "version": "1.16.0",
+  "version": "1.17.0",
   "description": "An AI copilot for your Second Brain",
   "author": "Saba Imran, Debanjum Singh Solanky <team@khoj.dev>",
   "license": "GPL-3.0-or-later",

src/interface/desktop/shortcut.html

@@ -346,7 +346,7 @@
             inp.focus();
         }
 
-        async function chat() {
+        async function chat(isVoice=false) {
             //set chat body to empty
             let chatBody = document.getElementById("chat-body");
             chatBody.innerHTML = "";
@@ -375,9 +375,6 @@
                 chat_body.dataset.conversationId = conversationID;
             }
 
-            // Generate backend API URL to execute query
-            let chatApi = `${hostURL}/api/chat?q=${encodeURIComponent(query)}&n=${resultsCount}&client=web&stream=true&conversation_id=${conversationID}&region=${region}&city=${city}&country=${countryName}&timezone=${timezone}`;
-
             let newResponseEl = document.createElement("div");
             newResponseEl.classList.add("chat-message", "khoj");
             newResponseEl.attributes["data-meta"] = "🏮 Khoj at " + formatDate(new Date());
@@ -388,128 +385,41 @@
             newResponseEl.appendChild(newResponseTextEl);
 
             // Temporary status message to indicate that Khoj is thinking
-            let loadingEllipsis = document.createElement("div");
+            let loadingEllipsis = createLoadingEllipsis();
-            loadingEllipsis.classList.add("lds-ellipsis");
-
-            let firstEllipsis = document.createElement("div");
-            firstEllipsis.classList.add("lds-ellipsis-item");
-
-            let secondEllipsis = document.createElement("div");
-            secondEllipsis.classList.add("lds-ellipsis-item");
-
-            let thirdEllipsis = document.createElement("div");
-            thirdEllipsis.classList.add("lds-ellipsis-item");
-
-            let fourthEllipsis = document.createElement("div");
-            fourthEllipsis.classList.add("lds-ellipsis-item");
-
-            loadingEllipsis.appendChild(firstEllipsis);
-            loadingEllipsis.appendChild(secondEllipsis);
-            loadingEllipsis.appendChild(thirdEllipsis);
-            loadingEllipsis.appendChild(fourthEllipsis);
-
-            newResponseTextEl.appendChild(loadingEllipsis);
             document.body.scrollTop = document.getElementById("chat-body").scrollHeight;
 
-            // Call Khoj chat API
-            let response = await fetch(chatApi, { headers });
-            let rawResponse = "";
-            let references = null;
-            const contentType = response.headers.get("content-type");
             toggleLoading();
-            if (contentType === "application/json") {
-                // Handle JSON response
-                try {
-                    const responseAsJson = await response.json();
-                    if (responseAsJson.image) {
-                        // If response has image field, response is a generated image.
-                        if (responseAsJson.intentType === "text-to-image") {
-                            rawResponse += `![${query}](data:image/png;base64,${responseAsJson.image})`;
-                        } else if (responseAsJson.intentType === "text-to-image2") {
-                            rawResponse += `![${query}](${responseAsJson.image})`;
-                        } else if (responseAsJson.intentType === "text-to-image-v3") {
-                            rawResponse += `![${query}](data:image/webp;base64,${responseAsJson.image})`;
-                        }
-                        const inferredQueries = responseAsJson.inferredQueries?.[0];
-                        if (inferredQueries) {
-                            rawResponse += `\n\n**Inferred Query**:\n\n${inferredQueries}`;
-                        }
-                    }
-                    if (responseAsJson.context) {
-                        const rawReferenceAsJson = responseAsJson.context;
-                        references = createReferenceSection(rawReferenceAsJson, createLinkerSection=true);
-                    }
-                    if (responseAsJson.detail) {
-                        // If response has detail field, response is an error message.
-                        rawResponse += responseAsJson.detail;
-                    }
-                } catch (error) {
-                    // If the chunk is not a JSON object, just display it as is
-                    rawResponse += chunk;
-                } finally {
-                    newResponseTextEl.innerHTML = "";
-                    newResponseTextEl.appendChild(formatHTMLMessage(rawResponse));
 
-                    if (references != null) {
+            // Setup chat message state
-                        newResponseTextEl.appendChild(references);
+            chatMessageState = {
-                    }
+                newResponseTextEl,
+                newResponseEl,
+                loadingEllipsis,
+                references: {},
+                rawResponse: "",
+                rawQuery: query,
+                isVoice: isVoice,
+            }
 
-                    document.body.scrollTop = document.getElementById("chat-body").scrollHeight;
+            // Construct API URL to execute chat query
-                }
+            let chatApi = `${hostURL}/api/chat?q=${encodeURIComponent(query)}&conversation_id=${conversationID}&stream=true&client=desktop`;
-            } else {
+            chatApi += (!!region && !!city && !!countryName && !!timezone)
-                // Handle streamed response of type text/event-stream or text/plain
+                ? `&region=${region}&city=${city}&country=${countryName}&timezone=${timezone}`
-                const reader = response.body.getReader();
+                : '';
-                const decoder = new TextDecoder();
-                let references = {};
 
-                readStream();
+            const response = await fetch(chatApi, { headers });
 
-                function readStream() {
+            try {
-                    reader.read().then(({ done, value }) => {
+                if (!response.ok) throw new Error(response.statusText);
-                        if (done) {
+                if (!response.body) throw new Error("Response body is empty");
-                            // Append any references after all the data has been streamed
+                // Stream and render chat response
-                            if (references != {}) {
+                await readChatStream(response);
-                                newResponseTextEl.appendChild(createReferenceSection(references, createLinkerSection=true));
+            } catch (err) {
-                            }
+                console.error(`Khoj chat response failed with\n${err}`);
-                            document.body.scrollTop = document.getElementById("chat-body").scrollHeight;
+                if (chatMessageState.newResponseEl.getElementsByClassName("lds-ellipsis").length > 0 && chatMessageState.loadingEllipsis)
-                            return;
+                    chatMessageState.newResponseTextEl.removeChild(chatMessageState.loadingEllipsis);
-                        }
+                let errorMsg = "Sorry, unable to get response from Khoj backend ❤️‍🩹. Retry or contact developers for help at <a href=mailto:'team@khoj.dev'>team@khoj.dev</a> or <a href='https://discord.gg/BDgyabRM6e'>on Discord</a>";
-
+                newResponseTextEl.textContent = errorMsg;
-                        // Decode message chunk from stream
-                        const chunk = decoder.decode(value, { stream: true });
-
-                        if (chunk.includes("### compiled references:")) {
-                            const additionalResponse = chunk.split("### compiled references:")[0];
-                            rawResponse += additionalResponse;
-                            newResponseTextEl.innerHTML = "";
-                            newResponseTextEl.appendChild(formatHTMLMessage(rawResponse));
-
-                            const rawReference = chunk.split("### compiled references:")[1];
-                            const rawReferenceAsJson = JSON.parse(rawReference);
-                            if (rawReferenceAsJson instanceof Array) {
-                                references["notes"] = rawReferenceAsJson;
-                            } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
-                                references["online"] = rawReferenceAsJson;
-                            }
-                            readStream();
-                        } else {
-                            // Display response from Khoj
-                            if (newResponseTextEl.getElementsByClassName("lds-ellipsis").length > 0) {
-                                newResponseTextEl.removeChild(loadingEllipsis);
-                            }
-
-                            // If the chunk is not a JSON object, just display it as is
-                            rawResponse += chunk;
-                            newResponseTextEl.innerHTML = "";
-                            newResponseTextEl.appendChild(formatHTMLMessage(rawResponse));
-
-                            readStream();
-                        }
-
-                        // Scroll to bottom of chat window as chat response is streamed
-                        document.body.scrollTop = document.getElementById("chat-body").scrollHeight;
-                    });
-                }
             }
             document.body.scrollTop = document.getElementById("chat-body").scrollHeight;
         }

src/interface/desktop/utils.js

@@ -34,8 +34,8 @@ function toggleNavMenu() {
 document.addEventListener('click', function(event) {
     let menu = document.getElementById("khoj-nav-menu");
     let menuContainer = document.getElementById("khoj-nav-menu-container");
-    let isClickOnMenu = menuContainer.contains(event.target) || menuContainer === event.target;
+    let isClickOnMenu = menuContainer?.contains(event.target) || menuContainer === event.target;
-    if (isClickOnMenu === false && menu.classList.contains("show")) {
+    if (menu && isClickOnMenu === false && menu.classList.contains("show")) {
         menu.classList.remove("show");
     }
 });

src/interface/emacs/khoj.el

@@ -6,7 +6,7 @@
 ;;         Saba Imran <saba@khoj.dev>
 ;; Description: An AI copilot for your Second Brain
 ;; Keywords: search, chat, org-mode, outlines, markdown, pdf, image
-;; Version: 1.16.0
+;; Version: 1.17.0
 ;; Package-Requires: ((emacs "27.1") (transient "0.3.0") (dash "2.19.1"))
 ;; URL: https://github.com/khoj-ai/khoj/tree/master/src/interface/emacs
 
@@ -283,9 +283,9 @@ Auto invokes setup steps on calling main entrypoint."
     (if (/= (apply #'call-process khoj-server-python-command
                    nil t nil
                    "-m" "pip" "install" "--upgrade"
-                   '("khoj-assistant"))
+                   '("khoj"))
             0)
-        (message "khoj.el: Failed to install Khoj server. Please install it manually using pip install `khoj-assistant'.\n%s" (buffer-string))
+        (message "khoj.el: Failed to install Khoj server. Please install it manually using pip install `khoj'.\n%s" (buffer-string))
       (message "khoj.el: Installed and upgraded Khoj server version: %s" (khoj--server-get-version)))))
 
 (defun khoj--server-start ()

src/interface/obsidian/manifest.json

@@ -1,7 +1,7 @@
 {
 	"id": "khoj",
 	"name": "Khoj",
-	"version": "1.16.0",
+	"version": "1.17.0",
 	"minAppVersion": "0.15.0",
 	"description": "An AI copilot for your Second Brain",
 	"author": "Khoj Inc.",

src/interface/obsidian/package.json

@@ -1,6 +1,6 @@
 {
     "name": "Khoj",
-    "version": "1.16.0",
+    "version": "1.17.0",
     "description": "An AI copilot for your Second Brain",
     "author": "Debanjum Singh Solanky, Saba Imran <team@khoj.dev>",
     "license": "GPL-3.0-or-later",

src/interface/obsidian/src/chat_view.ts

@@ -12,6 +12,25 @@ export interface ChatJsonResult {
     inferredQueries?: string[];
 }
 
+interface ChunkResult {
+    objects: string[];
+    remainder: string;
+}
+
+interface MessageChunk {
+    type: string;
+    data: any;
+}
+
+interface ChatMessageState {
+    newResponseTextEl: HTMLElement | null;
+    newResponseEl: HTMLElement | null;
+    loadingEllipsis: HTMLElement | null;
+    references: any;
+    rawResponse: string;
+    rawQuery: string;
+    isVoice: boolean;
+}
 
 interface Location {
     region: string;
@@ -26,6 +45,7 @@ export class KhojChatView extends KhojPaneView {
     waitingForLocation: boolean;
     location: Location;
     keyPressTimeout: NodeJS.Timeout | null = null;
+    chatMessageState: ChatMessageState;
 
     constructor(leaf: WorkspaceLeaf, setting: KhojSetting) {
         super(leaf, setting);
@@ -409,16 +429,15 @@ export class KhojChatView extends KhojPaneView {
         message = DOMPurify.sanitize(message);
 
         // Convert the message to html, sanitize the message html and render it to the real DOM
-        let chat_message_body_text_el = this.contentEl.createDiv();
+        let chatMessageBodyTextEl = this.contentEl.createDiv();
-        chat_message_body_text_el.className = "chat-message-text-response";
+        chatMessageBodyTextEl.innerHTML = this.markdownTextToSanitizedHtml(message, this);
-        chat_message_body_text_el.innerHTML = this.markdownTextToSanitizedHtml(message, this);
 
         // Add a copy button to each chat message, if it doesn't already exist
         if (willReplace === true) {
-            this.renderActionButtons(message, chat_message_body_text_el);
+            this.renderActionButtons(message, chatMessageBodyTextEl);
         }
 
-        return chat_message_body_text_el;
+        return chatMessageBodyTextEl;
     }
 
     markdownTextToSanitizedHtml(markdownText: string, component: ItemView): string {
@@ -502,23 +521,23 @@ export class KhojChatView extends KhojPaneView {
                 class: `khoj-chat-message ${sender}`
             },
         })
-        let chat_message_body_el = chatMessageEl.createDiv();
+        let chatMessageBodyEl = chatMessageEl.createDiv();
-        chat_message_body_el.addClasses(["khoj-chat-message-text", sender]);
+        chatMessageBodyEl.addClasses(["khoj-chat-message-text", sender]);
-        let chat_message_body_text_el = chat_message_body_el.createDiv();
+        let chatMessageBodyTextEl = chatMessageBodyEl.createDiv();
 
         // Sanitize the markdown to render
         message = DOMPurify.sanitize(message);
 
         if (raw) {
-            chat_message_body_text_el.innerHTML = message;
+            chatMessageBodyTextEl.innerHTML = message;
         } else {
             // @ts-ignore
-            chat_message_body_text_el.innerHTML = this.markdownTextToSanitizedHtml(message, this);
+            chatMessageBodyTextEl.innerHTML = this.markdownTextToSanitizedHtml(message, this);
         }
 
         // Add action buttons to each chat message element
         if (willReplace === true) {
-            this.renderActionButtons(message, chat_message_body_text_el);
+            this.renderActionButtons(message, chatMessageBodyTextEl);
         }
 
         // Remove user-select: none property to make text selectable
@@ -531,42 +550,38 @@ export class KhojChatView extends KhojPaneView {
     }
 
     createKhojResponseDiv(dt?: Date): HTMLDivElement {
-        let message_time = this.formatDate(dt ?? new Date());
+        let messageTime = this.formatDate(dt ?? new Date());
 
         // Append message to conversation history HTML element.
         // The chat logs should display above the message input box to follow standard UI semantics
-        let chat_body_el = this.contentEl.getElementsByClassName("khoj-chat-body")[0];
+        let chatBodyEl = this.contentEl.getElementsByClassName("khoj-chat-body")[0];
-        let chat_message_el = chat_body_el.createDiv({
+        let chatMessageEl = chatBodyEl.createDiv({
             attr: {
-                "data-meta": `🏮 Khoj at ${message_time}`,
+                "data-meta": `🏮 Khoj at ${messageTime}`,
                 class: `khoj-chat-message khoj`
             },
-        }).createDiv({
+        })
-            attr: {
-                class: `khoj-chat-message-text khoj`
-            },
-        }).createDiv();
 
         // Scroll to bottom after inserting chat messages
         this.scrollChatToBottom();
 
-        return chat_message_el;
+        return chatMessageEl;
     }
 
     async renderIncrementalMessage(htmlElement: HTMLDivElement, additionalMessage: string) {
-        this.result += additionalMessage;
+        this.chatMessageState.rawResponse += additionalMessage;
         htmlElement.innerHTML = "";
         // Sanitize the markdown to render
-        this.result = DOMPurify.sanitize(this.result);
+        this.chatMessageState.rawResponse = DOMPurify.sanitize(this.chatMessageState.rawResponse);
         // @ts-ignore
-        htmlElement.innerHTML = this.markdownTextToSanitizedHtml(this.result, this);
+        htmlElement.innerHTML = this.markdownTextToSanitizedHtml(this.chatMessageState.rawResponse, this);
         // Render action buttons for the message
-        this.renderActionButtons(this.result, htmlElement);
+        this.renderActionButtons(this.chatMessageState.rawResponse, htmlElement);
         // Scroll to bottom of modal, till the send message input box
         this.scrollChatToBottom();
     }
 
-    renderActionButtons(message: string, chat_message_body_text_el: HTMLElement) {
+    renderActionButtons(message: string, chatMessageBodyTextEl: HTMLElement) {
         let copyButton = this.contentEl.createEl('button');
         copyButton.classList.add("chat-action-button");
         copyButton.title = "Copy Message to Clipboard";
@@ -593,10 +608,10 @@ export class KhojChatView extends KhojPaneView {
         }
 
         // Append buttons to parent element
-        chat_message_body_text_el.append(copyButton, pasteToFile);
+        chatMessageBodyTextEl.append(copyButton, pasteToFile);
 
         if (speechButton) {
-            chat_message_body_text_el.append(speechButton);
+            chatMessageBodyTextEl.append(speechButton);
         }
     }
 
@@ -854,35 +869,122 @@ export class KhojChatView extends KhojPaneView {
         return true;
     }
 
-    async readChatStream(response: Response, responseElement: HTMLDivElement, isVoice: boolean = false): Promise<void> {
+    convertMessageChunkToJson(rawChunk: string): MessageChunk {
+        if (rawChunk?.startsWith("{") && rawChunk?.endsWith("}")) {
+            try {
+                let jsonChunk = JSON.parse(rawChunk);
+                if (!jsonChunk.type)
+                    jsonChunk = {type: 'message', data: jsonChunk};
+                return jsonChunk;
+            } catch (e) {
+                return {type: 'message', data: rawChunk};
+            }
+        } else if (rawChunk.length > 0) {
+            return {type: 'message', data: rawChunk};
+        }
+        return {type: '', data: ''};
+    }
+
+    processMessageChunk(rawChunk: string): void {
+        const chunk = this.convertMessageChunkToJson(rawChunk);
+        console.debug("Chunk:", chunk);
+        if (!chunk || !chunk.type) return;
+        if (chunk.type === 'status') {
+            console.log(`status: ${chunk.data}`);
+            const statusMessage = chunk.data;
+            this.handleStreamResponse(this.chatMessageState.newResponseTextEl, statusMessage, this.chatMessageState.loadingEllipsis, false);
+        } else if (chunk.type === 'start_llm_response') {
+            console.log("Started streaming", new Date());
+        } else if (chunk.type === 'end_llm_response') {
+            console.log("Stopped streaming", new Date());
+
+            // Automatically respond with voice if the subscribed user has sent voice message
+            if (this.chatMessageState.isVoice && this.setting.userInfo?.is_active)
+                this.textToSpeech(this.chatMessageState.rawResponse);
+
+            // Append any references after all the data has been streamed
+            this.finalizeChatBodyResponse(this.chatMessageState.references, this.chatMessageState.newResponseTextEl);
+
+            const liveQuery = this.chatMessageState.rawQuery;
+            // Reset variables
+            this.chatMessageState = {
+                newResponseTextEl: null,
+                newResponseEl: null,
+                loadingEllipsis: null,
+                references: {},
+                rawResponse: "",
+                rawQuery: liveQuery,
+                isVoice: false,
+            };
+        } else if (chunk.type === "references") {
+            this.chatMessageState.references = {"notes": chunk.data.context, "online": chunk.data.onlineContext};
+        } else if (chunk.type === 'message') {
+            const chunkData = chunk.data;
+            if (typeof chunkData === 'object' && chunkData !== null) {
+                // If chunkData is already a JSON object
+                this.handleJsonResponse(chunkData);
+            } else if (typeof chunkData === 'string' && chunkData.trim()?.startsWith("{") && chunkData.trim()?.endsWith("}")) {
+                // Try process chunk data as if it is a JSON object
+                try {
+                    const jsonData = JSON.parse(chunkData.trim());
+                    this.handleJsonResponse(jsonData);
+                } catch (e) {
+                    this.chatMessageState.rawResponse += chunkData;
+                    this.handleStreamResponse(this.chatMessageState.newResponseTextEl, this.chatMessageState.rawResponse, this.chatMessageState.loadingEllipsis);
+                }
+            } else {
+                this.chatMessageState.rawResponse += chunkData;
+                this.handleStreamResponse(this.chatMessageState.newResponseTextEl, this.chatMessageState.rawResponse, this.chatMessageState.loadingEllipsis);
+            }
+        }
+    }
+
+    handleJsonResponse(jsonData: any): void {
+        if (jsonData.image || jsonData.detail) {
+            this.chatMessageState.rawResponse = this.handleImageResponse(jsonData, this.chatMessageState.rawResponse);
+        } else if (jsonData.response) {
+            this.chatMessageState.rawResponse = jsonData.response;
+        }
+
+        if (this.chatMessageState.newResponseTextEl) {
+            this.chatMessageState.newResponseTextEl.innerHTML = "";
+            this.chatMessageState.newResponseTextEl.appendChild(this.formatHTMLMessage(this.chatMessageState.rawResponse));
+        }
+    }
+
+    async readChatStream(response: Response): Promise<void> {
         // Exit if response body is empty
         if (response.body == null) return;
 
         const reader = response.body.getReader();
         const decoder = new TextDecoder();
+        const eventDelimiter = '␃🔚␗';
+        let buffer = '';
 
         while (true) {
             const { value, done } = await reader.read();
 
             if (done) {
-                // Automatically respond with voice if the subscribed user has sent voice message
+                this.processMessageChunk(buffer);
-                if (isVoice && this.setting.userInfo?.is_active) this.textToSpeech(this.result);
+                buffer = '';
                 // Break if the stream is done
                 break;
             }
 
-            let responseText = decoder.decode(value);
+            const chunk = decoder.decode(value, { stream: true });
-            if (responseText.includes("### compiled references:")) {
+            console.debug("Raw Chunk:", chunk)
-                // Render any references used to generate the response
+            // Start buffering chunks until complete event is received
-                const [additionalResponse, rawReference] = responseText.split("### compiled references:", 2);
+            buffer += chunk;
-                await this.renderIncrementalMessage(responseElement, additionalResponse);
 
-                const rawReferenceAsJson = JSON.parse(rawReference);
+            // Once the buffer contains a complete event
-                let references = this.extractReferences(rawReferenceAsJson);
+            let newEventIndex;
-                responseElement.appendChild(this.createReferenceSection(references));
+            while ((newEventIndex = buffer.indexOf(eventDelimiter)) !== -1) {
-            } else {
+                // Extract the event from the buffer
-                // Render incremental chat response
+                const event = buffer.slice(0, newEventIndex);
-                await this.renderIncrementalMessage(responseElement, responseText);
+                buffer = buffer.slice(newEventIndex + eventDelimiter.length);
+
+                // Process the event
+                if (event) this.processMessageChunk(event);
             }
         }
     }
@@ -895,83 +997,59 @@ export class KhojChatView extends KhojPaneView {
         let chatBodyEl = this.contentEl.getElementsByClassName("khoj-chat-body")[0] as HTMLElement;
         this.renderMessage(chatBodyEl, query, "you");
 
-        let conversationID = chatBodyEl.dataset.conversationId;
+        let conversationId = chatBodyEl.dataset.conversationId;
-        if (!conversationID) {
+        if (!conversationId) {
             let chatUrl = `${this.setting.khojUrl}/api/chat/sessions?client=obsidian`;
             let response = await fetch(chatUrl, {
                 method: "POST",
                 headers: { "Authorization": `Bearer ${this.setting.khojApiKey}` },
             });
             let data = await response.json();
-            conversationID = data.conversation_id;
+            conversationId = data.conversation_id;
-            chatBodyEl.dataset.conversationId = conversationID;
+            chatBodyEl.dataset.conversationId = conversationId;
         }
 
         // Get chat response from Khoj backend
         let encodedQuery = encodeURIComponent(query);
-        let chatUrl = `${this.setting.khojUrl}/api/chat?q=${encodedQuery}&n=${this.setting.resultsCount}&client=obsidian&stream=true&region=${this.location.region}&city=${this.location.city}&country=${this.location.countryName}&timezone=${this.location.timezone}`;
+        let chatUrl = `${this.setting.khojUrl}/api/chat?q=${encodedQuery}&conversation_id=${conversationId}&n=${this.setting.resultsCount}&stream=true&client=obsidian`;
-        let responseElement = this.createKhojResponseDiv();
+        if (!!this.location) chatUrl += `&region=${this.location.region}&city=${this.location.city}&country=${this.location.countryName}&timezone=${this.location.timezone}`;
+
+        let newResponseEl = this.createKhojResponseDiv();
+        let newResponseTextEl = newResponseEl.createDiv();
+        newResponseTextEl.classList.add("khoj-chat-message-text", "khoj");
 
         // Temporary status message to indicate that Khoj is thinking
-        this.result = "";
         let loadingEllipsis = this.createLoadingEllipse();
-        responseElement.appendChild(loadingEllipsis);
+        newResponseTextEl.appendChild(loadingEllipsis);
+
+        // Set chat message state
+        this.chatMessageState = {
+            newResponseEl: newResponseEl,
+            newResponseTextEl: newResponseTextEl,
+            loadingEllipsis: loadingEllipsis,
+            references: {},
+            rawQuery: query,
+            rawResponse: "",
+            isVoice: isVoice,
+        };
 
         let response = await fetch(chatUrl, {
             method: "GET",
             headers: {
-                "Content-Type": "text/event-stream",
+                "Content-Type": "text/plain",
                 "Authorization": `Bearer ${this.setting.khojApiKey}`,
             },
         })
 
         try {
-            if (response.body === null) {
+            if (response.body === null) throw new Error("Response body is null");
-                throw new Error("Response body is null");
-            }
 
-            // Clear loading status message
+            // Stream and render chat response
-            if (responseElement.getElementsByClassName("lds-ellipsis").length > 0 && loadingEllipsis) {
+            await this.readChatStream(response);
-                responseElement.removeChild(loadingEllipsis);
-            }
-
-            // Reset collated chat result to empty string
-            this.result = "";
-            responseElement.innerHTML = "";
-            if (response.headers.get("content-type") === "application/json") {
-                let responseText = ""
-                try {
-                    const responseAsJson = await response.json() as ChatJsonResult;
-                    if (responseAsJson.image) {
-                        // If response has image field, response is a generated image.
-                        if (responseAsJson.intentType === "text-to-image") {
-                            responseText += `![${query}](data:image/png;base64,${responseAsJson.image})`;
-                        } else if (responseAsJson.intentType === "text-to-image2") {
-                            responseText += `![${query}](${responseAsJson.image})`;
-                        } else if (responseAsJson.intentType === "text-to-image-v3") {
-                            responseText += `![${query}](data:image/webp;base64,${responseAsJson.image})`;
-                        }
-                        const inferredQuery = responseAsJson.inferredQueries?.[0];
-                        if (inferredQuery) {
-                            responseText += `\n\n**Inferred Query**:\n\n${inferredQuery}`;
-                        }
-                    } else if (responseAsJson.detail) {
-                        responseText = responseAsJson.detail;
-                    }
-                } catch (error) {
-                    // If the chunk is not a JSON object, just display it as is
-                    responseText = await response.text();
-                } finally {
-                    await this.renderIncrementalMessage(responseElement, responseText);
-                }
-            } else {
-                // Stream and render chat response
-                await this.readChatStream(response, responseElement, isVoice);
-            }
         } catch (err) {
-            console.log(`Khoj chat response failed with\n${err}`);
+            console.error(`Khoj chat response failed with\n${err}`);
             let errorMsg = "Sorry, unable to get response from Khoj backend ❤️‍🩹. Retry or contact developers for help at <a href=mailto:'team@khoj.dev'>team@khoj.dev</a> or <a href='https://discord.gg/BDgyabRM6e'>on Discord</a>";
-            responseElement.innerHTML = errorMsg
+            newResponseTextEl.textContent = errorMsg;
         }
     }
 
@@ -1196,30 +1274,21 @@ export class KhojChatView extends KhojPaneView {
 
     handleStreamResponse(newResponseElement: HTMLElement | null, rawResponse: string, loadingEllipsis: HTMLElement | null, replace = true) {
         if (!newResponseElement) return;
-        if (newResponseElement.getElementsByClassName("lds-ellipsis").length > 0 && loadingEllipsis) {
+        // Remove loading ellipsis if it exists
+        if (newResponseElement.getElementsByClassName("lds-ellipsis").length > 0 && loadingEllipsis)
             newResponseElement.removeChild(loadingEllipsis);
-        }
+        // Clear the response element if replace is true
-        if (replace) {
+        if (replace) newResponseElement.innerHTML = "";
-            newResponseElement.innerHTML = "";
+
-        }
+        // Append response to the response element
         newResponseElement.appendChild(this.formatHTMLMessage(rawResponse, false, replace));
+
+        // Append loading ellipsis if it exists
+        if (!replace && loadingEllipsis) newResponseElement.appendChild(loadingEllipsis);
+        // Scroll to bottom of chat view
         this.scrollChatToBottom();
     }
 
-    handleCompiledReferences(rawResponseElement: HTMLElement | null, chunk: string, references: any, rawResponse: string) {
-        if (!rawResponseElement || !chunk) return { rawResponse, references };
-
-        const [additionalResponse, rawReference] = chunk.split("### compiled references:", 2);
-        rawResponse += additionalResponse;
-        rawResponseElement.innerHTML = "";
-        rawResponseElement.appendChild(this.formatHTMLMessage(rawResponse));
-
-        const rawReferenceAsJson = JSON.parse(rawReference);
-        references = this.extractReferences(rawReferenceAsJson);
-
-        return { rawResponse, references };
-    }
-
     handleImageResponse(imageJson: any, rawResponse: string) {
         if (imageJson.image) {
             const inferredQuery = imageJson.inferredQueries?.[0] ?? "generated image";
@@ -1236,33 +1305,10 @@ export class KhojChatView extends KhojPaneView {
                 rawResponse += `\n\n**Inferred Query**:\n\n${inferredQuery}`;
             }
         }
-        let references = {};
+        // If response has detail field, response is an error message.
-        if (imageJson.context && imageJson.context.length > 0) {
+        if (imageJson.detail) rawResponse += imageJson.detail;
-            references = this.extractReferences(imageJson.context);
-        }
-        if (imageJson.detail) {
-            // If response has detail field, response is an error message.
-            rawResponse += imageJson.detail;
-        }
-        return { rawResponse, references };
-    }
 
-    extractReferences(rawReferenceAsJson: any): object {
+        return rawResponse;
-        let references: any = {};
-        if (rawReferenceAsJson instanceof Array) {
-            references["notes"] = rawReferenceAsJson;
-        } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
-            references["online"] = rawReferenceAsJson;
-        }
-        return references;
-    }
-
-    addMessageToChatBody(rawResponse: string, newResponseElement: HTMLElement | null, references: any) {
-        if (!newResponseElement) return;
-        newResponseElement.innerHTML = "";
-        newResponseElement.appendChild(this.formatHTMLMessage(rawResponse));
-
-        this.finalizeChatBodyResponse(references, newResponseElement);
     }
 
     finalizeChatBodyResponse(references: object, newResponseElement: HTMLElement | null) {

src/interface/obsidian/src/search_modal.ts

@@ -1,6 +1,6 @@
 import { App, SuggestModal, request, MarkdownRenderer, Instruction, Platform } from 'obsidian';
 import { KhojSetting } from 'src/settings';
-import { createNoteAndCloseModal, getLinkToEntry } from 'src/utils';
+import { supportedBinaryFileTypes, createNoteAndCloseModal, getFileFromPath, getLinkToEntry, supportedImageFilesTypes } from 'src/utils';
 
 export interface SearchResult {
     entry: string;
@@ -112,28 +112,41 @@ export class KhojSearchModal extends SuggestModal<SearchResult> {
         let os_path_separator = result.file.includes('\\') ? '\\' : '/';
         let filename = result.file.split(os_path_separator).pop();
 
-        // Remove YAML frontmatter when rendering string
-        result.entry = result.entry.replace(/---[\n\r][\s\S]*---[\n\r]/, '');
-
-        // Truncate search results to lines_to_render
-        let entry_snipped_indicator = result.entry.split('\n').length > lines_to_render ? ' **...**' : '';
-        let snipped_entry = result.entry.split('\n').slice(0, lines_to_render).join('\n');
-
         // Show filename of each search result for context
         el.createEl("div",{ cls: 'khoj-result-file' }).setText(filename ?? "");
         let result_el = el.createEl("div", { cls: 'khoj-result-entry' })
 
+        let resultToRender = "";
+        let fileExtension = filename?.split(".").pop() ?? "";
+        if (supportedImageFilesTypes.includes(fileExtension) && filename) {
+            let linkToEntry: string = filename;
+            let imageFiles = this.app.vault.getFiles().filter(file => supportedImageFilesTypes.includes(fileExtension));
+            // Find vault file of chosen search result
+            let fileInVault = getFileFromPath(imageFiles, result.file);
+            if (fileInVault)
+                linkToEntry = this.app.vault.getResourcePath(fileInVault);
+
+            resultToRender = `![](${linkToEntry})`;
+        } else {
+            // Remove YAML frontmatter when rendering string
+            result.entry = result.entry.replace(/---[\n\r][\s\S]*---[\n\r]/, '');
+
+            // Truncate search results to lines_to_render
+            let entry_snipped_indicator = result.entry.split('\n').length > lines_to_render ? ' **...**' : '';
+            let snipped_entry = result.entry.split('\n').slice(0, lines_to_render).join('\n');
+            resultToRender = `${snipped_entry}${entry_snipped_indicator}`;
+        }
         // @ts-ignore
-        MarkdownRenderer.renderMarkdown(snipped_entry + entry_snipped_indicator, result_el, result.file, null);
+        MarkdownRenderer.renderMarkdown(resultToRender, result_el, result.file, null);
     }
 
     async onChooseSuggestion(result: SearchResult, _: MouseEvent | KeyboardEvent) {
-        // Get all markdown and PDF files in vault
+        // Get all markdown, pdf and image files in vault
         const mdFiles = this.app.vault.getMarkdownFiles();
-        const pdfFiles = this.app.vault.getFiles().filter(file => file.extension === 'pdf');
+        const binaryFiles = this.app.vault.getFiles().filter(file => supportedBinaryFileTypes.includes(file.extension));
 
         // Find, Open vault file at heading of chosen search result
-        let linkToEntry = getLinkToEntry(mdFiles.concat(pdfFiles), result.file, result.entry);
+        let linkToEntry = getLinkToEntry(mdFiles.concat(binaryFiles), result.file, result.entry);
         if (linkToEntry) this.app.workspace.openLinkText(linkToEntry, '');
     }
 }

src/interface/obsidian/src/settings.ts

@@ -10,7 +10,6 @@ export interface UserInfo {
     email?: string;
 }
 
-
 export interface KhojSetting {
     resultsCount: number;
     khojUrl: string;

src/interface/obsidian/src/utils.ts

@@ -48,11 +48,14 @@ function filenameToMimeType (filename: TFile): string {
     }
 }
 
+export const supportedImageFilesTypes = ['png', 'jpg', 'jpeg'];
+export const supportedBinaryFileTypes = ['pdf'].concat(supportedImageFilesTypes);
+export const supportedFileTypes = ['md', 'markdown'].concat(supportedBinaryFileTypes);
+
 export async function updateContentIndex(vault: Vault, setting: KhojSetting, lastSync: Map<TFile, number>, regenerate: boolean = false): Promise<Map<TFile, number>> {
     // Get all markdown, pdf files in the vault
     console.log(`Khoj: Updating Khoj content index...`)
-    const files = vault.getFiles().filter(file => file.extension === 'md' || file.extension === 'markdown' || file.extension === 'pdf');
+    const files = vault.getFiles().filter(file => supportedFileTypes.includes(file.extension));
-    const binaryFileTypes = ['pdf']
     let countOfFilesToIndex = 0;
     let countOfFilesToDelete = 0;
     lastSync = lastSync.size > 0 ? lastSync : new Map<TFile, number>();
@@ -66,7 +69,7 @@ export async function updateContentIndex(vault: Vault, setting: KhojSetting, las
         }
 
         countOfFilesToIndex++;
-        const encoding = binaryFileTypes.includes(file.extension) ? "binary" : "utf8";
+        const encoding = supportedBinaryFileTypes.includes(file.extension) ? "binary" : "utf8";
         const mimeType = fileExtensionToMimeType(file.extension) + (encoding === "utf8" ? "; charset=UTF-8" : "");
         const fileContent = encoding == 'binary' ? await vault.readBinary(file) : await vault.read(file);
         fileData.push({blob: new Blob([fileContent], { type: mimeType }), path: file.path});
@@ -354,7 +357,7 @@ export function pasteTextAtCursor(text: string | undefined) {
     }
 }
 
-export function getLinkToEntry(sourceFiles: TFile[], chosenFile: string, chosenEntry: string): string | undefined {
+export function getFileFromPath(sourceFiles: TFile[], chosenFile: string): TFile | undefined {
     // Find the vault file matching file of chosen file, entry
     let fileMatch = sourceFiles
         // Sort by descending length of path
@@ -363,6 +366,12 @@ export function getLinkToEntry(sourceFiles: TFile[], chosenFile: string, chosenE
         // The first match is the best file match across OS
         // e.g. Khoj server on Linux, Obsidian vault on Android
         .find(file => chosenFile.replace(/\\/g, "/").endsWith(file.path))
+    return fileMatch;
+}
+
+export function getLinkToEntry(sourceFiles: TFile[], chosenFile: string, chosenEntry: string): string | undefined {
+    // Find the vault file matching file of chosen file, entry
+    let fileMatch = getFileFromPath(sourceFiles, chosenFile);
 
     // Return link to vault file at heading of chosen search result
     if (fileMatch) {

src/interface/obsidian/styles.css

@@ -85,6 +85,12 @@ If your plugin does not need CSS, delete this file.
     margin-left: auto;
     white-space: pre-line;
 }
+/* Override white-space for ul, ol, li under khoj-chat-message-text.khoj */
+.khoj-chat-message-text.khoj ul,
+.khoj-chat-message-text.khoj ol,
+.khoj-chat-message-text.khoj li {
+    white-space: normal;
+}
 /* add left protrusion to khoj chat bubble */
 .khoj-chat-message-text.khoj:after {
     content: '';

src/interface/obsidian/versions.json

@@ -53,5 +53,6 @@
 	"1.13.0": "0.15.0",
 	"1.14.0": "0.15.0",
 	"1.15.0": "0.15.0",
-	"1.16.0": "0.15.0"
+	"1.16.0": "0.15.0",
+	"1.17.0": "0.15.0"
 }

src/khoj/database/adapters/__init__.py

@@ -559,7 +559,7 @@ class AgentAdapters:
         if default_conversation_config is None:
             logger.info("No default conversation config found, skipping default agent creation")
             return None
-        default_personality = prompts.personality.format(current_date="placeholder")
+        default_personality = prompts.personality.format(current_date="placeholder", day_of_week="placeholder")
 
         agent = Agent.objects.filter(name=AgentAdapters.DEFAULT_AGENT_NAME).first()
 
@@ -684,19 +684,18 @@ class ConversationAdapters:
     async def aget_conversation_by_user(
         user: KhojUser, client_application: ClientApplication = None, conversation_id: int = None, title: str = None
     ) -> Optional[Conversation]:
+        query = Conversation.objects.filter(user=user, client=client_application).prefetch_related("agent")
+
         if conversation_id:
-            return await Conversation.objects.filter(user=user, client=client_application, id=conversation_id).afirst()
+            return await query.filter(id=conversation_id).afirst()
         elif title:
-            return await Conversation.objects.filter(user=user, client=client_application, title=title).afirst()
+            return await query.filter(title=title).afirst()
-        else:
-            conversation = Conversation.objects.filter(user=user, client=client_application).order_by("-updated_at")
 
-        if await conversation.aexists():
+        conversation = await query.order_by("-updated_at").afirst()
-            return await conversation.prefetch_related("agent").afirst()
 
-        return await (
+        return conversation or await Conversation.objects.prefetch_related("agent").acreate(
-            Conversation.objects.filter(user=user, client=client_application).order_by("-updated_at").afirst()
+            user=user, client=client_application
-        ) or await Conversation.objects.acreate(user=user, client=client_application)
+        )
 
     @staticmethod
     async def adelete_conversation_by_user(

src/khoj/interface/web/chat.html

@@ -74,14 +74,13 @@ To get started, just start typing below. You can also type / to see a list of co
                 }, 1000);
             });
         }
-        var websocket = null;
+
         let region = null;
         let city = null;
         let countryName = null;
         let timezone = null;
         let waitingForLocation = true;
-
+        let chatMessageState = {
-        let websocketState = {
             newResponseTextEl: null,
             newResponseEl: null,
             loadingEllipsis: null,
@@ -105,7 +104,7 @@ To get started, just start typing below. You can also type / to see a list of co
             .finally(() => {
                 console.debug("Region:", region, "City:", city, "Country:", countryName, "Timezone:", timezone);
                 waitingForLocation = false;
-                setupWebSocket();
+                initMessageState();
             });
 
         function formatDate(date) {
@@ -599,13 +598,8 @@ To get started, just start typing below. You can also type / to see a list of co
         }
 
         async function chat(isVoice=false) {
-            if (websocket) {
+            // Extract chat message from chat input form
-                sendMessageViaWebSocket(isVoice);
+            var query = document.getElementById("chat-input").value.trim();
-                return;
-            }
-
-            let query = document.getElementById("chat-input").value.trim();
-            let resultsCount = localStorage.getItem("khojResultsCount") || 5;
             console.log(`Query: ${query}`);
 
             // Short circuit on empty query
@@ -624,31 +618,30 @@ To get started, just start typing below. You can also type / to see a list of co
             document.getElementById("chat-input").value = "";
             autoResize();
             document.getElementById("chat-input").setAttribute("disabled", "disabled");
-            let chat_body = document.getElementById("chat-body");
-
-            let conversationID = chat_body.dataset.conversationId;
 
+            let chatBody = document.getElementById("chat-body");
+            let conversationID = chatBody.dataset.conversationId;
             if (!conversationID) {
-                let response = await fetch('/api/chat/sessions', { method: "POST" });
+                let response = await fetch(`${hostURL}/api/chat/sessions`, { method: "POST" });
                 let data = await response.json();
                 conversationID = data.conversation_id;
-                chat_body.dataset.conversationId = conversationID;
+                chatBody.dataset.conversationId = conversationID;
-                refreshChatSessionsPanel();
+                await refreshChatSessionsPanel();
             }
 
-            let new_response = document.createElement("div");
+            let newResponseEl = document.createElement("div");
-            new_response.classList.add("chat-message", "khoj");
+            newResponseEl.classList.add("chat-message", "khoj");
-            new_response.attributes["data-meta"] = "🏮 Khoj at " + formatDate(new Date());
+            newResponseEl.attributes["data-meta"] = "🏮 Khoj at " + formatDate(new Date());
-            chat_body.appendChild(new_response);
+            chatBody.appendChild(newResponseEl);
 
-            let newResponseText = document.createElement("div");
+            let newResponseTextEl = document.createElement("div");
-            newResponseText.classList.add("chat-message-text", "khoj");
+            newResponseTextEl.classList.add("chat-message-text", "khoj");
-            new_response.appendChild(newResponseText);
+            newResponseEl.appendChild(newResponseTextEl);
 
             // Temporary status message to indicate that Khoj is thinking
             let loadingEllipsis = createLoadingEllipse();
 
-            newResponseText.appendChild(loadingEllipsis);
+            newResponseTextEl.appendChild(loadingEllipsis);
             document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
 
             let chatTooltip = document.getElementById("chat-tooltip");
@@ -657,65 +650,38 @@ To get started, just start typing below. You can also type / to see a list of co
             let chatInput = document.getElementById("chat-input");
             chatInput.classList.remove("option-enabled");
 
-            // Generate backend API URL to execute query
+            // Setup chat message state
-            let url = `/api/chat?q=${encodeURIComponent(query)}&n=${resultsCount}&client=web&stream=true&conversation_id=${conversationID}&region=${region}&city=${city}&country=${countryName}&timezone=${timezone}`;
+            chatMessageState = {
-
+                newResponseTextEl,
-            // Call specified Khoj API
+                newResponseEl,
-            let response = await fetch(url);
+                loadingEllipsis,
-            let rawResponse = "";
+                references: {},
-            let references = null;
+                rawResponse: "",
-            const contentType = response.headers.get("content-type");
+                rawQuery: query,
-
+                isVoice: isVoice,
-            if (contentType === "application/json") {
-                // Handle JSON response
-                try {
-                    const responseAsJson = await response.json();
-                    if (responseAsJson.image || responseAsJson.detail) {
-                        ({rawResponse, references } = handleImageResponse(responseAsJson, rawResponse));
-                    } else {
-                        rawResponse = responseAsJson.response;
-                    }
-                } catch (error) {
-                    // If the chunk is not a JSON object, just display it as is
-                    rawResponse += chunk;
-                } finally {
-                    addMessageToChatBody(rawResponse, newResponseText, references);
-                }
-            } else {
-                // Handle streamed response of type text/event-stream or text/plain
-                const reader = response.body.getReader();
-                const decoder = new TextDecoder();
-                let references = {};
-
-                readStream();
-
-                function readStream() {
-                    reader.read().then(({ done, value }) => {
-                        if (done) {
-                            // Append any references after all the data has been streamed
-                            finalizeChatBodyResponse(references, newResponseText);
-                            return;
-                        }
-
-                        // Decode message chunk from stream
-                        const chunk = decoder.decode(value, { stream: true });
-
-                        if (chunk.includes("### compiled references:")) {
-                            ({ rawResponse, references } = handleCompiledReferences(newResponseText, chunk, references, rawResponse));
-                            readStream();
-                        } else {
-                            // If the chunk is not a JSON object, just display it as is
-                            rawResponse += chunk;
-                            handleStreamResponse(newResponseText, rawResponse, query, loadingEllipsis);
-                            readStream();
-                        }
-                    });
-
-                    // Scroll to bottom of chat window as chat response is streamed
-                    document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
-                };
             }
-        };
+
+            // Call Khoj chat API
+            let chatApi = `/api/chat?q=${encodeURIComponent(query)}&conversation_id=${conversationID}&stream=true&client=web`;
+            chatApi += (!!region && !!city && !!countryName && !!timezone)
+                ? `&region=${region}&city=${city}&country=${countryName}&timezone=${timezone}`
+                : '';
+
+            const response = await fetch(chatApi);
+
+            try {
+                if (!response.ok) throw new Error(response.statusText);
+                if (!response.body) throw new Error("Response body is empty");
+                // Stream and render chat response
+                await readChatStream(response);
+            } catch (err) {
+                console.error(`Khoj chat response failed with\n${err}`);
+                if (chatMessageState.newResponseEl.getElementsByClassName("lds-ellipsis").length > 0 && chatMessageState.loadingEllipsis)
+                    chatMessageState.newResponseTextEl.removeChild(chatMessageState.loadingEllipsis);
+                let errorMsg = "Sorry, unable to get response from Khoj backend ❤️‍🩹. Retry or contact developers for help at <a href=mailto:'team@khoj.dev'>team@khoj.dev</a> or <a href='https://discord.gg/BDgyabRM6e'>on Discord</a>";
+                newResponseTextEl.innerHTML = errorMsg;
+            }
+        }
 
         function createLoadingEllipse() {
             // Temporary status message to indicate that Khoj is thinking
@@ -743,32 +709,22 @@ To get started, just start typing below. You can also type / to see a list of co
         }
 
         function handleStreamResponse(newResponseElement, rawResponse, rawQuery, loadingEllipsis, replace=true) {
-            if (newResponseElement.getElementsByClassName("lds-ellipsis").length > 0 && loadingEllipsis) {
+            if (!newResponseElement) return;
+            // Remove loading ellipsis if it exists
+            if (newResponseElement.getElementsByClassName("lds-ellipsis").length > 0 && loadingEllipsis)
                 newResponseElement.removeChild(loadingEllipsis);
-            }
+            // Clear the response element if replace is true
-            if (replace) {
+            if (replace) newResponseElement.innerHTML = "";
-                newResponseElement.innerHTML = "";
+
-            }
+            // Append response to the response element
             newResponseElement.appendChild(formatHTMLMessage(rawResponse, false, replace, rawQuery));
+
+            // Append loading ellipsis if it exists
+            if (!replace && loadingEllipsis) newResponseElement.appendChild(loadingEllipsis);
+            // Scroll to bottom of chat view
             document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
         }
 
-        function handleCompiledReferences(rawResponseElement, chunk, references, rawResponse) {
-            const additionalResponse = chunk.split("### compiled references:")[0];
-            rawResponse += additionalResponse;
-            rawResponseElement.innerHTML = "";
-            rawResponseElement.appendChild(formatHTMLMessage(rawResponse));
-
-            const rawReference = chunk.split("### compiled references:")[1];
-            const rawReferenceAsJson = JSON.parse(rawReference);
-            if (rawReferenceAsJson instanceof Array) {
-                references["notes"] = rawReferenceAsJson;
-            } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
-                references["online"] = rawReferenceAsJson;
-            }
-            return { rawResponse, references };
-        }
-
         function handleImageResponse(imageJson, rawResponse) {
             if (imageJson.image) {
                 const inferredQuery = imageJson.inferredQueries?.[0] ?? "generated image";
@@ -785,35 +741,139 @@ To get started, just start typing below. You can also type / to see a list of co
                     rawResponse += `\n\n**Inferred Query**:\n\n${inferredQuery}`;
                 }
             }
-            let references = {};
-            if (imageJson.context && imageJson.context.length > 0) {
-                const rawReferenceAsJson = imageJson.context;
-                if (rawReferenceAsJson instanceof Array) {
-                    references["notes"] = rawReferenceAsJson;
-                } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
-                    references["online"] = rawReferenceAsJson;
-                }
-            }
-            if (imageJson.detail) {
-                // If response has detail field, response is an error message.
-                rawResponse += imageJson.detail;
-            }
-            return { rawResponse, references };
-        }
 
-        function addMessageToChatBody(rawResponse, newResponseElement, references) {
+            // If response has detail field, response is an error message.
-            newResponseElement.innerHTML = "";
+            if (imageJson.detail) rawResponse += imageJson.detail;
-            newResponseElement.appendChild(formatHTMLMessage(rawResponse));
 
-            finalizeChatBodyResponse(references, newResponseElement);
+            return rawResponse;
         }
 
         function finalizeChatBodyResponse(references, newResponseElement) {
-            if (references != null && Object.keys(references).length > 0) {
+            if (!!newResponseElement && references != null && Object.keys(references).length > 0) {
                 newResponseElement.appendChild(createReferenceSection(references));
             }
             document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
-            document.getElementById("chat-input").removeAttribute("disabled");
+            document.getElementById("chat-input")?.removeAttribute("disabled");
+        }
+
+        function convertMessageChunkToJson(rawChunk) {
+            // Split the chunk into lines
+            console.debug("Raw Event:", rawChunk);
+            if (rawChunk?.startsWith("{") && rawChunk?.endsWith("}")) {
+                try {
+                    let jsonChunk = JSON.parse(rawChunk);
+                    if (!jsonChunk.type)
+                        jsonChunk = {type: 'message', data: jsonChunk};
+                    return jsonChunk;
+                } catch (e) {
+                    return {type: 'message', data: rawChunk};
+                }
+            } else if (rawChunk.length > 0) {
+                return {type: 'message', data: rawChunk};
+            }
+        }
+
+        function processMessageChunk(rawChunk) {
+            const chunk = convertMessageChunkToJson(rawChunk);
+            console.debug("Json Event:", chunk);
+            if (!chunk || !chunk.type) return;
+            if (chunk.type ==='status') {
+                console.log(`status: ${chunk.data}`);
+                const statusMessage = chunk.data;
+                handleStreamResponse(chatMessageState.newResponseTextEl, statusMessage, chatMessageState.rawQuery, chatMessageState.loadingEllipsis, false);
+            } else if (chunk.type === 'start_llm_response') {
+                console.log("Started streaming", new Date());
+            } else if (chunk.type === 'end_llm_response') {
+                console.log("Stopped streaming", new Date());
+
+                // Automatically respond with voice if the subscribed user has sent voice message
+                if (chatMessageState.isVoice && "{{ is_active }}" == "True")
+                    textToSpeech(chatMessageState.rawResponse);
+
+                // Append any references after all the data has been streamed
+                finalizeChatBodyResponse(chatMessageState.references, chatMessageState.newResponseTextEl);
+
+                const liveQuery = chatMessageState.rawQuery;
+                // Reset variables
+                chatMessageState = {
+                    newResponseTextEl: null,
+                    newResponseEl: null,
+                    loadingEllipsis: null,
+                    references: {},
+                    rawResponse: "",
+                    rawQuery: liveQuery,
+                    isVoice: false,
+                }
+            } else if (chunk.type === "references") {
+                chatMessageState.references = {"notes": chunk.data.context, "online": chunk.data.onlineContext};
+            } else if (chunk.type === 'message') {
+                const chunkData = chunk.data;
+                if (typeof chunkData === 'object' && chunkData !== null) {
+                    // If chunkData is already a JSON object
+                    handleJsonResponse(chunkData);
+                } else if (typeof chunkData  === 'string' && chunkData.trim()?.startsWith("{") && chunkData.trim()?.endsWith("}")) {
+                    // Try process chunk data as if it is a JSON object
+                    try {
+                        const jsonData = JSON.parse(chunkData.trim());
+                        handleJsonResponse(jsonData);
+                    } catch (e) {
+                        chatMessageState.rawResponse += chunkData;
+                        handleStreamResponse(chatMessageState.newResponseTextEl, chatMessageState.rawResponse, chatMessageState.rawQuery, chatMessageState.loadingEllipsis);
+                    }
+                } else {
+                    chatMessageState.rawResponse += chunkData;
+                    handleStreamResponse(chatMessageState.newResponseTextEl, chatMessageState.rawResponse, chatMessageState.rawQuery, chatMessageState.loadingEllipsis);
+                }
+            }
+        }
+
+        function handleJsonResponse(jsonData) {
+            if (jsonData.image || jsonData.detail) {
+                chatMessageState.rawResponse = handleImageResponse(jsonData, chatMessageState.rawResponse);
+            } else if (jsonData.response) {
+                chatMessageState.rawResponse = jsonData.response;
+            }
+
+            if (chatMessageState.newResponseTextEl) {
+                chatMessageState.newResponseTextEl.innerHTML = "";
+                chatMessageState.newResponseTextEl.appendChild(formatHTMLMessage(chatMessageState.rawResponse));
+            }
+        }
+
+        async function readChatStream(response) {
+            if (!response.body) return;
+            const reader = response.body.getReader();
+            const decoder = new TextDecoder();
+            const eventDelimiter = '␃🔚␗';
+            let buffer = '';
+
+            while (true) {
+                const { value, done } = await reader.read();
+                // If the stream is done
+                if (done) {
+                    // Process the last chunk
+                    processMessageChunk(buffer);
+                    buffer = '';
+                    break;
+                }
+
+                // Read chunk from stream and append it to the buffer
+                const chunk = decoder.decode(value, { stream: true });
+                console.debug("Raw Chunk:", chunk)
+                // Start buffering chunks until complete event is received
+                buffer += chunk;
+
+                // Once the buffer contains a complete event
+                let newEventIndex;
+                while ((newEventIndex = buffer.indexOf(eventDelimiter)) !== -1) {
+                    // Extract the event from the buffer
+                    const event = buffer.slice(0, newEventIndex);
+                    buffer = buffer.slice(newEventIndex + eventDelimiter.length);
+
+                    // Process the event
+                    if (event) processMessageChunk(event);
+                }
+            }
         }
 
         function incrementalChat(event) {
@@ -1069,17 +1129,13 @@ To get started, just start typing below. You can also type / to see a list of co
 
         window.onload = loadChat;
 
-        function setupWebSocket(isVoice=false) {
+        function initMessageState(isVoice=false) {
-            let chatBody = document.getElementById("chat-body");
-            let wsProtocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-            let webSocketUrl = `${wsProtocol}//${window.location.host}/api/chat/ws`;
-
             if (waitingForLocation) {
                 console.debug("Waiting for location data to be fetched. Will setup WebSocket once location data is available.");
                 return;
             }
 
-            websocketState = {
+            chatMessageState = {
                 newResponseTextEl: null,
                 newResponseEl: null,
                 loadingEllipsis: null,
@@ -1088,174 +1144,8 @@ To get started, just start typing below. You can also type / to see a list of co
                 rawQuery: "",
                 isVoice: isVoice,
             }
-
-            if (chatBody.dataset.conversationId) {
-                webSocketUrl += `?conversation_id=${chatBody.dataset.conversationId}`;
-                webSocketUrl += (!!region && !!city && !!countryName) && !!timezone ? `&region=${region}&city=${city}&country=${countryName}&timezone=${timezone}` : '';
-
-                websocket = new WebSocket(webSocketUrl);
-                websocket.onmessage = function(event) {
-
-                    // Get the last element in the chat-body
-                    let chunk = event.data;
-                    if (chunk == "start_llm_response") {
-                        console.log("Started streaming", new Date());
-                    } else if (chunk == "end_llm_response") {
-                        console.log("Stopped streaming", new Date());
-
-                        // Automatically respond with voice if the subscribed user has sent voice message
-                        if (websocketState.isVoice && "{{ is_active }}" == "True")
-                            textToSpeech(websocketState.rawResponse);
-
-                        // Append any references after all the data has been streamed
-                        finalizeChatBodyResponse(websocketState.references, websocketState.newResponseTextEl);
-
-                        const liveQuery = websocketState.rawQuery;
-                        // Reset variables
-                        websocketState = {
-                            newResponseTextEl: null,
-                            newResponseEl: null,
-                            loadingEllipsis: null,
-                            references: {},
-                            rawResponse: "",
-                            rawQuery: liveQuery,
-                            isVoice: false,
-                        }
-                    } else {
-                        try {
-                            if (chunk.includes("application/json"))
-                            {
-                                chunk = JSON.parse(chunk);
-                            }
-                        } catch (error) {
-                            // If the chunk is not a JSON object, continue.
-                        }
-
-                        const contentType = chunk["content-type"]
-
-                        if (contentType === "application/json") {
-                            // Handle JSON response
-                            try {
-                                if (chunk.image || chunk.detail) {
-                                    ({rawResponse, references } = handleImageResponse(chunk, websocketState.rawResponse));
-                                    websocketState.rawResponse = rawResponse;
-                                    websocketState.references = references;
-                                } else if (chunk.type == "status") {
-                                    handleStreamResponse(websocketState.newResponseTextEl, chunk.message, websocketState.rawQuery, null, false);
-                                } else if (chunk.type == "rate_limit") {
-                                    handleStreamResponse(websocketState.newResponseTextEl, chunk.message, websocketState.rawQuery, websocketState.loadingEllipsis, true);
-                                } else {
-                                    rawResponse = chunk.response;
-                                }
-                            } catch (error) {
-                                // If the chunk is not a JSON object, just display it as is
-                                websocketState.rawResponse += chunk;
-                            } finally {
-                                if (chunk.type != "status" && chunk.type != "rate_limit") {
-                                    addMessageToChatBody(websocketState.rawResponse, websocketState.newResponseTextEl, websocketState.references);
-                                }
-                            }
-                        } else {
-
-                            // Handle streamed response of type text/event-stream or text/plain
-                            if (chunk && chunk.includes("### compiled references:")) {
-                                ({ rawResponse, references } = handleCompiledReferences(websocketState.newResponseTextEl, chunk, websocketState.references, websocketState.rawResponse));
-                                websocketState.rawResponse = rawResponse;
-                                websocketState.references = references;
-                            } else {
-                                // If the chunk is not a JSON object, just display it as is
-                                websocketState.rawResponse += chunk;
-                                if (websocketState.newResponseTextEl) {
-                                    handleStreamResponse(websocketState.newResponseTextEl, websocketState.rawResponse, websocketState.rawQuery, websocketState.loadingEllipsis);
-                                }
-                            }
-
-                            // Scroll to bottom of chat window as chat response is streamed
-                            document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
-                        };
-                    }
-                }
-            };
-            websocket.onclose = function(event) {
-                websocket = null;
-                console.log("WebSocket is closed now.");
-                let setupWebSocketButton = document.createElement("button");
-                setupWebSocketButton.textContent = "Reconnect to Server";
-                setupWebSocketButton.onclick = setupWebSocket;
-                let statusDotIcon = document.getElementById("connection-status-icon");
-                statusDotIcon.style.backgroundColor = "red";
-                let statusDotText = document.getElementById("connection-status-text");
-                statusDotText.innerHTML = "";
-                statusDotText.style.marginTop = "5px";
-                statusDotText.appendChild(setupWebSocketButton);
-            }
-            websocket.onerror = function(event) {
-                console.log("WebSocket error observed:", event);
-            }
-
-            websocket.onopen = function(event) {
-                console.log("WebSocket is open now.")
-                let statusDotIcon = document.getElementById("connection-status-icon");
-                statusDotIcon.style.backgroundColor = "green";
-                let statusDotText = document.getElementById("connection-status-text");
-                statusDotText.textContent = "Connected to Server";
-            }
         }
 
-        function sendMessageViaWebSocket(isVoice=false) {
-            let chatBody = document.getElementById("chat-body");
-
-            var query = document.getElementById("chat-input").value.trim();
-            console.log(`Query: ${query}`);
-
-            if (userMessages.length >= 10) {
-                userMessages.shift();
-            }
-            userMessages.push(query);
-            resetUserMessageIndex();
-
-            // Add message by user to chat body
-            renderMessage(query, "you");
-            document.getElementById("chat-input").value = "";
-            autoResize();
-            document.getElementById("chat-input").setAttribute("disabled", "disabled");
-
-            let newResponseEl = document.createElement("div");
-            newResponseEl.classList.add("chat-message", "khoj");
-            newResponseEl.attributes["data-meta"] = "🏮 Khoj at " + formatDate(new Date());
-            chatBody.appendChild(newResponseEl);
-
-            let newResponseTextEl = document.createElement("div");
-            newResponseTextEl.classList.add("chat-message-text", "khoj");
-            newResponseEl.appendChild(newResponseTextEl);
-
-            // Temporary status message to indicate that Khoj is thinking
-            let loadingEllipsis = createLoadingEllipse();
-
-            newResponseTextEl.appendChild(loadingEllipsis);
-            document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
-
-            let chatTooltip = document.getElementById("chat-tooltip");
-            chatTooltip.style.display = "none";
-
-            let chatInput = document.getElementById("chat-input");
-            chatInput.classList.remove("option-enabled");
-
-            // Call specified Khoj API
-            websocket.send(query);
-            let rawResponse = "";
-            let references = {};
-
-            websocketState = {
-                newResponseTextEl,
-                newResponseEl,
-                loadingEllipsis,
-                references,
-                rawResponse,
-                rawQuery: query,
-                isVoice: isVoice,
-            }
-        }
         var userMessages = [];
         var userMessageIndex = -1;
         function loadChat() {
@@ -1265,7 +1155,7 @@ To get started, just start typing below. You can also type / to see a list of co
             let chatHistoryUrl = `/api/chat/history?client=web`;
             if (chatBody.dataset.conversationId) {
                 chatHistoryUrl += `&conversation_id=${chatBody.dataset.conversationId}`;
-                setupWebSocket();
+                initMessageState();
                 loadFileFiltersFromConversation();
             }
 
@@ -1305,7 +1195,7 @@ To get started, just start typing below. You can also type / to see a list of co
                     let chatBody = document.getElementById("chat-body");
                     chatBody.dataset.conversationId = response.conversation_id;
                     loadFileFiltersFromConversation();
-                    setupWebSocket();
+                    initMessageState();
                     chatBody.dataset.conversationTitle = response.slug || `New conversation 🌱`;
 
                     let agentMetadata = response.agent;

src/khoj/main.py

@@ -206,7 +206,7 @@ def set_state(args):
     state.host = args.host
     state.port = args.port
     state.anonymous_mode = args.anonymous_mode
-    state.khoj_version = version("khoj-assistant")
+    state.khoj_version = version("khoj")
     state.chat_on_gpu = args.chat_on_gpu
 
 

src/khoj/processor/conversation/anthropic/anthropic_chat.py

@@ -36,7 +36,7 @@ def extract_questions_anthropic(
     # Extract Past User Message and Inferred Questions from Conversation Log
     chat_history = "".join(
         [
-            f'Q: {chat["intent"]["query"]}\nKhoj: {{"queries": {chat["intent"].get("inferred-queries") or list([chat["intent"]["query"]])}}}\nA: {chat["message"]}\n\n'
+            f'User: {chat["intent"]["query"]}\nAssistant: {{"queries": {chat["intent"].get("inferred-queries") or list([chat["intent"]["query"]])}}}\nA: {chat["message"]}\n\n'
             for chat in conversation_log.get("chat", [])[-4:]
             if chat["by"] == "khoj" and "text-to-image" not in chat["intent"].get("type")
         ]
@@ -135,17 +135,23 @@ def converse_anthropic(
     Converse with user using Anthropic's Claude
     """
     # Initialize Variables
-    current_date = datetime.now().strftime("%Y-%m-%d")
+    current_date = datetime.now()
     compiled_references = "\n\n".join({f"# {item}" for item in references})
 
     conversation_primer = prompts.query_prompt.format(query=user_query)
 
     if agent and agent.personality:
         system_prompt = prompts.custom_personality.format(
-            name=agent.name, bio=agent.personality, current_date=current_date
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
         )
     else:
-        system_prompt = prompts.personality.format(current_date=current_date)
+        system_prompt = prompts.personality.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
 
     if location_data:
         location = f"{location_data.city}, {location_data.region}, {location_data.country}"

src/khoj/processor/conversation/offline/chat_model.py

@@ -55,6 +55,7 @@ def extract_questions_offline(
                 chat_history += f"Q: {chat['intent']['query']}\n"
                 chat_history += f"Khoj: {chat['message']}\n\n"
 
+    # Get dates relative to today for prompt creation
     today = datetime.today()
     yesterday = (today - timedelta(days=1)).strftime("%Y-%m-%d")
     last_year = today.year - 1
@@ -62,11 +63,13 @@ def extract_questions_offline(
         query=text,
         chat_history=chat_history,
         current_date=today.strftime("%Y-%m-%d"),
+        day_of_week=today.strftime("%A"),
         yesterday_date=yesterday,
         last_year=last_year,
         this_year=today.year,
         location=location,
     )
+
     messages = generate_chatml_messages_with_context(
         example_questions, model_name=model, loaded_model=offline_chat_model, max_prompt_size=max_prompt_size
     )
@@ -74,7 +77,7 @@ def extract_questions_offline(
     state.chat_lock.acquire()
     try:
         response = send_message_to_model_offline(
-            messages, loaded_model=offline_chat_model, max_prompt_size=max_prompt_size
+            messages, loaded_model=offline_chat_model, model=model, max_prompt_size=max_prompt_size
         )
     finally:
         state.chat_lock.release()
@@ -96,7 +99,7 @@ def extract_questions_offline(
     except:
         logger.warning(f"Llama returned invalid JSON. Falling back to using user message as search query.\n{response}")
         return all_questions
-    logger.debug(f"Extracted Questions by Llama: {questions}")
+    logger.debug(f"Questions extracted by {model}: {questions}")
     return questions
 
 
@@ -144,14 +147,20 @@ def converse_offline(
     offline_chat_model = loaded_model or download_model(model, max_tokens=max_prompt_size)
     compiled_references_message = "\n\n".join({f"{item['compiled']}" for item in references})
 
-    current_date = datetime.now().strftime("%Y-%m-%d")
+    current_date = datetime.now()
 
     if agent and agent.personality:
         system_prompt = prompts.custom_system_prompt_offline_chat.format(
-            name=agent.name, bio=agent.personality, current_date=current_date
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
         )
     else:
-        system_prompt = prompts.system_prompt_offline_chat.format(current_date=current_date)
+        system_prompt = prompts.system_prompt_offline_chat.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
 
     conversation_primer = prompts.query_prompt.format(query=user_query)
 
@@ -177,9 +186,9 @@ def converse_offline(
             if online_results[result].get("webpages"):
                 simplified_online_results[result] = online_results[result]["webpages"]
 
-        conversation_primer = f"{prompts.online_search_conversation.format(online_results=str(simplified_online_results))}\n{conversation_primer}"
+        conversation_primer = f"{prompts.online_search_conversation_offline.format(online_results=str(simplified_online_results))}\n{conversation_primer}"
     if not is_none_or_empty(compiled_references_message):
-        conversation_primer = f"{prompts.notes_conversation_offline.format(references=compiled_references_message)}\n{conversation_primer}"
+        conversation_primer = f"{prompts.notes_conversation_offline.format(references=compiled_references_message)}\n\n{conversation_primer}"
 
     # Setup Prompt with Primer or Conversation History
     messages = generate_chatml_messages_with_context(
@@ -192,6 +201,9 @@ def converse_offline(
         tokenizer_name=tokenizer_name,
     )
 
+    truncated_messages = "\n".join({f"{message.content[:70]}..." for message in messages})
+    logger.debug(f"Conversation Context for {model}: {truncated_messages}")
+
     g = ThreadedGenerator(references, online_results, completion_func=completion_func)
     t = Thread(target=llm_thread, args=(g, messages, offline_chat_model, max_prompt_size))
     t.start()

src/khoj/processor/conversation/offline/utils.py

@@ -24,6 +24,8 @@ def download_model(repo_id: str, filename: str = "*Q4_K_M.gguf", max_tokens: int
     # Add chat format if known
     if "llama-3" in repo_id.lower():
         kwargs["chat_format"] = "llama-3"
+    elif "gemma-2" in repo_id.lower():
+        kwargs["chat_format"] = "gemma"
 
     # Check if the model is already downloaded
     model_path = load_model_from_cache(repo_id, filename)

src/khoj/processor/conversation/openai/gpt.py

@@ -125,17 +125,23 @@ def converse(
     Converse with user using OpenAI's ChatGPT
     """
     # Initialize Variables
-    current_date = datetime.now().strftime("%Y-%m-%d")
+    current_date = datetime.now()
     compiled_references = "\n\n".join({f"# {item['compiled']}" for item in references})
 
     conversation_primer = prompts.query_prompt.format(query=user_query)
 
     if agent and agent.personality:
         system_prompt = prompts.custom_personality.format(
-            name=agent.name, bio=agent.personality, current_date=current_date
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
         )
     else:
-        system_prompt = prompts.personality.format(current_date=current_date)
+        system_prompt = prompts.personality.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
 
     if location_data:
         location = f"{location_data.city}, {location_data.region}, {location_data.country}"

src/khoj/processor/conversation/prompts.py

@@ -19,8 +19,8 @@ You were created by Khoj Inc. with the following capabilities:
 - Sometimes the user will share personal information that needs to be remembered, like an account ID or a residential address. These can be acknowledged with a simple "Got it" or "Okay".
 - Provide inline references to quotes from the user's notes or any web pages you refer to in your responses in markdown format. For example, "The farmer had ten sheep. [1](https://example.com)". *ALWAYS CITE YOUR SOURCES AND PROVIDE REFERENCES*. Add them inline to directly support your claim.
 
-Note: More information about you, the company or Khoj apps for download can be found at https://khoj.dev.
+Note: More information about you, the company or Khoj apps can be found at https://khoj.dev.
-Today is {current_date} in UTC.
+Today is {day_of_week}, {current_date} in UTC.
 """.strip()
 )
 
@@ -39,7 +39,7 @@ You were created by Khoj Inc. with the following capabilities:
 - Ask crisp follow-up questions to get additional context, when the answer cannot be inferred from the provided notes or past conversations.
 - Sometimes the user will share personal information that needs to be remembered, like an account ID or a residential address. These can be acknowledged with a simple "Got it" or "Okay".
 
-Today is {current_date} in UTC.
+Today is {day_of_week}, {current_date} in UTC.
 
 Instructions:\n{bio}
 """.strip()
@@ -79,10 +79,12 @@ You are Khoj, a smart, inquisitive and helpful personal assistant.
 - Use your general knowledge and past conversation with the user as context to inform your responses.
 - If you do not know the answer, say 'I don't know.'
 - Think step-by-step and ask questions to get the necessary information to answer the user's question.
+- Ask crisp follow-up questions to get additional context, when the answer cannot be inferred from the provided information or past conversations.
 - Do not print verbatim Notes unless necessary.
 
-Today is {current_date} in UTC.
+Note: More information about you, the company or Khoj apps can be found at https://khoj.dev.
-    """.strip()
+Today is {day_of_week}, {current_date} in UTC.
+""".strip()
 )
 
 custom_system_prompt_offline_chat = PromptTemplate.from_template(
@@ -91,12 +93,14 @@ You are {name}, a personal agent on Khoj.
 - Use your general knowledge and past conversation with the user as context to inform your responses.
 - If you do not know the answer, say 'I don't know.'
 - Think step-by-step and ask questions to get the necessary information to answer the user's question.
+- Ask crisp follow-up questions to get additional context, when the answer cannot be inferred from the provided information or past conversations.
 - Do not print verbatim Notes unless necessary.
 
-Today is {current_date} in UTC.
+Note: More information about you, the company or Khoj apps can be found at https://khoj.dev.
+Today is {day_of_week}, {current_date} in UTC.
 
 Instructions:\n{bio}
-    """.strip()
+""".strip()
 )
 
 ## Notes Conversation
@@ -106,13 +110,15 @@ notes_conversation = PromptTemplate.from_template(
 Use my personal notes and our past conversations to inform your response.
 Ask crisp follow-up questions to get additional context, when a helpful response cannot be provided from the provided notes or past conversations.
 
-Notes:
+User's Notes:
 {references}
 """.strip()
 )
 
 notes_conversation_offline = PromptTemplate.from_template(
     """
+Use my personal notes and our past conversations to inform your response.
+
 User's Notes:
 {references}
 """.strip()
@@ -174,6 +180,15 @@ Information from the internet:
 """.strip()
 )
 
+online_search_conversation_offline = PromptTemplate.from_template(
+    """
+Use this up-to-date information from the internet to inform your response.
+
+Information from the internet:
+{online_results}
+""".strip()
+)
+
 ## Query prompt
 ## --
 query_prompt = PromptTemplate.from_template(
@@ -186,15 +201,16 @@ Query: {query}""".strip()
 ## --
 extract_questions_offline = PromptTemplate.from_template(
     """
-You are Khoj, an extremely smart and helpful search assistant with the ability to retrieve information from the user's notes. Construct search queries to retrieve relevant information to answer the user's question.
+You are Khoj, an extremely smart and helpful search assistant with the ability to retrieve information from the user's notes. Disregard online search requests.
-- You will be provided past questions(Q) and answers(A) for context.
+Construct search queries to retrieve relevant information to answer the user's question.
+- You will be provided past questions(Q) and answers(Khoj) for context.
 - Try to be as specific as possible. Instead of saying "they" or "it" or "he", use proper nouns like name of the person or thing you are referring to.
 - Add as much context from the previous questions and answers as required into your search queries.
 - Break messages into multiple search queries when required to retrieve the relevant information.
 - Add date filters to your search queries from questions and answers when required to retrieve the relevant information.
 - Share relevant search queries as a JSON list of strings. Do not say anything else.
 
-Current Date: {current_date}
+Current Date: {day_of_week}, {current_date}
 User's Location: {location}
 
 Examples:
@@ -232,7 +248,8 @@ Q: {query}
 
 extract_questions = PromptTemplate.from_template(
     """
-You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes. Disregard online search requests. Construct search queries to retrieve relevant information to answer the user's question.
+You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes. Disregard online search requests.
+Construct search queries to retrieve relevant information to answer the user's question.
 - You will be provided past questions(Q) and answers(A) for context.
 - Add as much context from the previous questions and answers as required into your search queries.
 - Break messages into multiple search queries when required to retrieve the relevant information.
@@ -282,8 +299,9 @@ Khoj:
 
 extract_questions_anthropic_system_prompt = PromptTemplate.from_template(
     """
-You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes. Disregard online search requests. Construct search queries to retrieve relevant information to answer the user's question.
+You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes. Disregard online search requests.
-- You will be provided past questions(Q) and answers(A) for context.
+Construct search queries to retrieve relevant information to answer the user's question.
+- You will be provided past questions(User), extracted queries(Assistant) and answers(A) for context.
 - Add as much context from the previous questions and answers as required into your search queries.
 - Break messages into multiple search queries when required to retrieve the relevant information.
 - Add date filters to your search queries from questions and answers when required to retrieve the relevant information.
@@ -297,15 +315,19 @@ Here are some examples of how you can construct search queries to answer the use
 
 User: How was my trip to Cambodia?
 Assistant: {{"queries": ["How was my trip to Cambodia?"]}}
+A: The trip was amazing. You went to the Angkor Wat temple and it was beautiful.
 
 User: What national parks did I go to last year?
 Assistant: {{"queries": ["National park I visited in {last_new_year} dt>='{last_new_year_date}' dt<'{current_new_year_date}'"]}}
+A: You visited the Grand Canyon and Yellowstone National Park in {last_new_year}.
 
 User: How can you help me?
 Assistant: {{"queries": ["Social relationships", "Physical and mental health", "Education and career", "Personal life goals and habits"]}}
+A: I can help you live healthier and happier across work and personal life
 
 User: Who all did I meet here yesterday?
 Assistant: {{"queries": ["Met in {location} on {yesterday_date} dt>='{yesterday_date}' dt<'{current_date}'"]}}
+A: Yesterday's note mentions your visit to your local beach with Ram and Shyam.
 """.strip()
 )
 
@@ -319,7 +341,11 @@ Assistant:
 """.strip()
 )
 
-system_prompt_extract_relevant_information = """As a professional analyst, create a comprehensive report of the most relevant information from a web page in response to a user's query. The text provided is directly from within the web page. The report you create should be multiple paragraphs, and it should represent the content of the website. Tell the user exactly what the website says in response to their query, while adhering to these guidelines:
+system_prompt_extract_relevant_information = """
+As a professional analyst, create a comprehensive report of the most relevant information from a web page in response to a user's query.
+The text provided is directly from within the web page.
+The report you create should be multiple paragraphs, and it should represent the content of the website.
+Tell the user exactly what the website says in response to their query, while adhering to these guidelines:
 
 1. Answer the user's query as specifically as possible. Include many supporting details from the website.
 2. Craft a report that is detailed, thorough, in-depth, and complex, while maintaining clarity.
@@ -340,7 +366,11 @@ Collate only relevant information from the website to answer the target query.
 """.strip()
 )
 
-system_prompt_extract_relevant_summary = """As a professional analyst, create a comprehensive report of the most relevant information from the document in response to a user's query. The text provided is directly from within the document. The report you create should be multiple paragraphs, and it should represent the content of the document. Tell the user exactly what the document says in response to their query, while adhering to these guidelines:
+system_prompt_extract_relevant_summary = """
+As a professional analyst, create a comprehensive report of the most relevant information from the document in response to a user's query.
+The text provided is directly from within the document.
+The report you create should be multiple paragraphs, and it should represent the content of the document.
+Tell the user exactly what the document says in response to their query, while adhering to these guidelines:
 
 1. Answer the user's query as specifically as possible. Include many supporting details from the document.
 2. Craft a report that is detailed, thorough, in-depth, and complex, while maintaining clarity.
@@ -363,11 +393,13 @@ Collate only relevant information from the document to answer the target query.
 
 pick_relevant_output_mode = PromptTemplate.from_template(
     """
-You are Khoj, an excellent analyst for selecting the correct way to respond to a user's query. You have access to a limited set of modes for your response. You can only use one of these modes.
+You are Khoj, an excellent analyst for selecting the correct way to respond to a user's query.
+You have access to a limited set of modes for your response.
+You can only use one of these modes.
 
 {modes}
 
-Here are some example responses:
+Here are some examples:
 
 Example:
 Chat History:
@@ -383,7 +415,7 @@ User: I'm having trouble deciding which laptop to get. I want something with at
 AI: I can help with that. I see online that there is a new model of the Dell XPS 15 that meets your requirements.
 
 Q: What are the specs of the new Dell XPS 15?
-Khoj: default
+Khoj: text
 
 Example:
 Chat History:
@@ -391,7 +423,7 @@ User: Where did I go on my last vacation?
 AI: You went to Jordan and visited Petra, the Dead Sea, and Wadi Rum.
 
 Q: Remind me who did I go with on that trip?
-Khoj: default
+Khoj: text
 
 Example:
 Chat History:
@@ -399,7 +431,7 @@ User: How's the weather outside? Current Location: Bali, Indonesia
 AI: It's currently 28°C and partly cloudy in Bali.
 
 Q: Share a painting using the weather for Bali every morning.
-Khoj: reminder
+Khoj: automation
 
 Now it's your turn to pick the mode you would like to use to answer the user's question. Provide your response as a string.
 
@@ -422,7 +454,7 @@ Which of the data sources listed below you would use to answer the user's questi
 
 {tools}
 
-Here are some example responses:
+Here are some examples:
 
 Example:
 Chat History:
@@ -533,10 +565,10 @@ You are Khoj, an advanced google search assistant. You are tasked with construct
 - Break messages into multiple search queries when required to retrieve the relevant information.
 - Use site: google search operators when appropriate
 - You have access to the the whole internet to retrieve information.
-- Official, up-to-date information about you, Khoj, is available at site:khoj.dev
+- Official, up-to-date information about you, Khoj, is available at site:khoj.dev, github or pypi.
 
 What Google searches, if any, will you need to perform to answer the user's question?
-Provide search queries as a list of strings in a JSON object.
+Provide search queries as a list of strings in a JSON object. Do not wrap the json in a codeblock.
 Current Date: {current_date}
 User's Location: {location}
 
@@ -589,7 +621,6 @@ Q: How many oranges would fit in NASA's Saturn V rocket?
 Khoj: {{"queries": ["volume of an orange", "volume of saturn v rocket"]}}
 
 Now it's your turn to construct Google search queries to answer the user's question. Provide them as a list of strings in a JSON object. Do not say anything else.
-Now it's your turn to construct a search query for Google to answer the user's question.
 History:
 {chat_history}
 

src/khoj/processor/conversation/utils.py

@@ -62,10 +62,6 @@ class ThreadedGenerator:
         self.queue.put(data)
 
     def close(self):
-        if self.compiled_references and len(self.compiled_references) > 0:
-            self.queue.put(f"### compiled references:{json.dumps(self.compiled_references)}")
-        if self.online_results and len(self.online_results) > 0:
-            self.queue.put(f"### compiled references:{json.dumps(self.online_results)}")
         self.queue.put(StopIteration)
 
 
@@ -186,7 +182,7 @@ def generate_chatml_messages_with_context(
 
 def truncate_messages(
     messages: list[ChatMessage],
-    max_prompt_size,
+    max_prompt_size: int,
     model_name: str,
     loaded_model: Optional[Llama] = None,
     tokenizer_name=None,
@@ -232,7 +228,8 @@ def truncate_messages(
     tokens = sum([len(encoder.encode(message.content)) for message in messages if type(message.content) == str])
 
     # Drop older messages until under max supported prompt size by model
-    while (tokens + system_message_tokens) > max_prompt_size and len(messages) > 1:
+    # Reserves 4 tokens to demarcate each message (e.g <|im_start|>user, <|im_end|>, <|endoftext|> etc.)
+    while (tokens + system_message_tokens + 4 * len(messages)) > max_prompt_size and len(messages) > 1:
         messages.pop()
         tokens = sum([len(encoder.encode(message.content)) for message in messages if type(message.content) == str])
 
@@ -254,6 +251,8 @@ def truncate_messages(
             f"Truncate current message to fit within max prompt size of {max_prompt_size} supported by {model_name} model:\n {truncated_message}"
         )
 
+    if system_message:
+        system_message.role = "user" if "gemma-2" in model_name else "system"
     return messages + [system_message] if system_message else messages
 
 

src/khoj/processor/tools/online_search.py

@@ -11,6 +11,7 @@ from bs4 import BeautifulSoup
 from markdownify import markdownify
 
 from khoj.routers.helpers import (
+    ChatEvent,
     extract_relevant_info,
     generate_online_subqueries,
     infer_webpage_urls,
@@ -56,7 +57,8 @@ async def search_online(
     query += " ".join(custom_filters)
     if not is_internet_connected():
         logger.warn("Cannot search online as not connected to internet")
-        return {}
+        yield {}
+        return
 
     # Breakdown the query into subqueries to get the correct answer
     subqueries = await generate_online_subqueries(query, conversation_history, location)
@@ -66,7 +68,8 @@ async def search_online(
         logger.info(f"🌐 Searching the Internet for {list(subqueries)}")
         if send_status_func:
             subqueries_str = "\n- " + "\n- ".join(list(subqueries))
-            await send_status_func(f"**🌐 Searching the Internet for**: {subqueries_str}")
+            async for event in send_status_func(f"**🌐 Searching the Internet for**: {subqueries_str}"):
+                yield {ChatEvent.STATUS: event}
 
     with timer(f"Internet searches for {list(subqueries)} took", logger):
         search_func = search_with_google if SERPER_DEV_API_KEY else search_with_jina
@@ -89,7 +92,8 @@ async def search_online(
         logger.info(f"🌐👀 Reading web pages at: {list(webpage_links)}")
         if send_status_func:
             webpage_links_str = "\n- " + "\n- ".join(list(webpage_links))
-            await send_status_func(f"**📖 Reading web pages**: {webpage_links_str}")
+            async for event in send_status_func(f"**📖 Reading web pages**: {webpage_links_str}"):
+                yield {ChatEvent.STATUS: event}
     tasks = [read_webpage_and_extract_content(subquery, link, content) for link, subquery, content in webpages]
     results = await asyncio.gather(*tasks)
 
@@ -98,7 +102,7 @@ async def search_online(
         if webpage_extract is not None:
             response_dict[subquery]["webpages"] = {"link": url, "snippet": webpage_extract}
 
-    return response_dict
+    yield response_dict
 
 
 async def search_with_google(query: str) -> Tuple[str, Dict[str, List[Dict]]]:
@@ -127,13 +131,15 @@ async def read_webpages(
     "Infer web pages to read from the query and extract relevant information from them"
     logger.info(f"Inferring web pages to read")
     if send_status_func:
-        await send_status_func(f"**🧐 Inferring web pages to read**")
+        async for event in send_status_func(f"**🧐 Inferring web pages to read**"):
+            yield {ChatEvent.STATUS: event}
     urls = await infer_webpage_urls(query, conversation_history, location)
 
     logger.info(f"Reading web pages at: {urls}")
     if send_status_func:
         webpage_links_str = "\n- " + "\n- ".join(list(urls))
-        await send_status_func(f"**📖 Reading web pages**: {webpage_links_str}")
+        async for event in send_status_func(f"**📖 Reading web pages**: {webpage_links_str}"):
+            yield {ChatEvent.STATUS: event}
     tasks = [read_webpage_and_extract_content(query, url) for url in urls]
     results = await asyncio.gather(*tasks)
 
@@ -141,7 +147,7 @@ async def read_webpages(
     response[query]["webpages"] = [
         {"query": q, "link": url, "snippet": web_extract} for q, web_extract, url in results if web_extract is not None
     ]
-    return response
+    yield response
 
 
 async def read_webpage_and_extract_content(

src/khoj/routers/api.py

@@ -37,6 +37,7 @@ from khoj.processor.conversation.openai.gpt import extract_questions
 from khoj.processor.conversation.openai.whisper import transcribe_audio
 from khoj.routers.helpers import (
     ApiUserRateLimiter,
+    ChatEvent,
     CommonQueryParams,
     ConversationCommandRateLimiter,
     acreate_title_from_query,
@@ -342,11 +343,13 @@ async def extract_references_and_questions(
         not ConversationCommand.Notes in conversation_commands
         and not ConversationCommand.Default in conversation_commands
     ):
-        return compiled_references, inferred_queries, q
+        yield compiled_references, inferred_queries, q
+        return
 
     if not await sync_to_async(EntryAdapters.user_has_entries)(user=user):
         logger.debug("No documents in knowledge base. Use a Khoj client to sync and chat with your docs.")
-        return compiled_references, inferred_queries, q
+        yield compiled_references, inferred_queries, q
+        return
 
     # Extract filter terms from user message
     defiltered_query = q
@@ -357,11 +360,12 @@ async def extract_references_and_questions(
 
     if not conversation:
         logger.error(f"Conversation with id {conversation_id} not found.")
-        return compiled_references, inferred_queries, defiltered_query
+        yield compiled_references, inferred_queries, defiltered_query
+        return
 
     filters_in_query += " ".join([f'file:"{filter}"' for filter in conversation.file_filters])
     using_offline_chat = False
-    print(f"Filters in query: {filters_in_query}")
+    logger.debug(f"Filters in query: {filters_in_query}")
 
     # Infer search queries from user message
     with timer("Extracting search queries took", logger):
@@ -379,6 +383,7 @@ async def extract_references_and_questions(
 
             inferred_queries = extract_questions_offline(
                 defiltered_query,
+                model=chat_model,
                 loaded_model=loaded_model,
                 conversation_log=meta_log,
                 should_extract_questions=True,
@@ -416,7 +421,8 @@ async def extract_references_and_questions(
         logger.info(f"🔍 Searching knowledge base with queries: {inferred_queries}")
         if send_status_func:
             inferred_queries_str = "\n- " + "\n- ".join(inferred_queries)
-            await send_status_func(f"**Searching Documents for:** {inferred_queries_str}")
+            async for event in send_status_func(f"**Searching Documents for:** {inferred_queries_str}"):
+                yield {ChatEvent.STATUS: event}
         for query in inferred_queries:
             n_items = min(n, 3) if using_offline_chat else n
             search_results.extend(
@@ -435,7 +441,7 @@ async def extract_references_and_questions(
             {"compiled": item.additional["compiled"], "file": item.additional["file"]} for item in search_results
         ]
 
-    return compiled_references, inferred_queries, defiltered_query
+    yield compiled_references, inferred_queries, defiltered_query
 
 
 @api.get("/health", response_class=Response)

src/khoj/routers/api_chat.py

@@ -1,17 +1,17 @@
+import asyncio
 import json
 import logging
-import math
+import time
 from datetime import datetime
-from typing import Dict, Optional
+from functools import partial
+from typing import Any, Dict, List, Optional
 from urllib.parse import unquote
 
 from asgiref.sync import sync_to_async
-from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
+from fastapi import APIRouter, Depends, HTTPException, Request
 from fastapi.requests import Request
 from fastapi.responses import Response, StreamingResponse
 from starlette.authentication import requires
-from starlette.websockets import WebSocketDisconnect
-from websockets import ConnectionClosedOK
 
 from khoj.app.settings import ALLOWED_HOSTS
 from khoj.database.adapters import (
@@ -22,19 +22,15 @@ from khoj.database.adapters import (
     aget_user_name,
 )
 from khoj.database.models import KhojUser
-from khoj.processor.conversation.prompts import (
+from khoj.processor.conversation.prompts import help_message, no_entries_found
-    help_message,
-    no_entries_found,
-    no_notes_found,
-)
 from khoj.processor.conversation.utils import save_to_conversation_log
 from khoj.processor.speech.text_to_speech import generate_text_to_speech
 from khoj.processor.tools.online_search import read_webpages, search_online
 from khoj.routers.api import extract_references_and_questions
 from khoj.routers.helpers import (
     ApiUserRateLimiter,
+    ChatEvent,
     CommonQueryParams,
-    CommonQueryParamsClass,
     ConversationCommandRateLimiter,
     agenerate_chat_response,
     aget_relevant_information_sources,
@@ -519,141 +515,142 @@ async def set_conversation_title(
     )
 
 
-@api_chat.websocket("/ws")
+@api_chat.get("")
-async def websocket_endpoint(
+async def chat(
-    websocket: WebSocket,
+    request: Request,
-    conversation_id: int,
+    common: CommonQueryParams,
+    q: str,
+    n: int = 7,
+    d: float = 0.18,
+    stream: Optional[bool] = False,
+    title: Optional[str] = None,
+    conversation_id: Optional[int] = None,
     city: Optional[str] = None,
     region: Optional[str] = None,
     country: Optional[str] = None,
     timezone: Optional[str] = None,
+    rate_limiter_per_minute=Depends(
+        ApiUserRateLimiter(requests=5, subscribed_requests=60, window=60, slug="chat_minute")
+    ),
+    rate_limiter_per_day=Depends(
+        ApiUserRateLimiter(requests=5, subscribed_requests=600, window=60 * 60 * 24, slug="chat_day")
+    ),
 ):
-    connection_alive = True
+    async def event_generator(q: str):
+        start_time = time.perf_counter()
+        ttft = None
+        chat_metadata: dict = {}
+        connection_alive = True
+        user: KhojUser = request.user.object
+        event_delimiter = "␃🔚␗"
+        q = unquote(q)
 
-    async def send_status_update(message: str):
+        async def send_event(event_type: ChatEvent, data: str | dict):
-        nonlocal connection_alive
+            nonlocal connection_alive, ttft
-        if not connection_alive:
+            if not connection_alive or await request.is_disconnected():
+                connection_alive = False
+                logger.warn(f"User {user} disconnected from {common.client} client")
+                return
+            try:
+                if event_type == ChatEvent.END_LLM_RESPONSE:
+                    collect_telemetry()
+                if event_type == ChatEvent.START_LLM_RESPONSE:
+                    ttft = time.perf_counter() - start_time
+                if event_type == ChatEvent.MESSAGE:
+                    yield data
+                elif event_type == ChatEvent.REFERENCES or stream:
+                    yield json.dumps({"type": event_type.value, "data": data}, ensure_ascii=False)
+            except asyncio.CancelledError as e:
+                connection_alive = False
+                logger.warn(f"User {user} disconnected from {common.client} client: {e}")
+                return
+            except Exception as e:
+                connection_alive = False
+                logger.error(f"Failed to stream chat API response to {user} on {common.client}: {e}", exc_info=True)
+                return
+            finally:
+                if stream:
+                    yield event_delimiter
+
+        async def send_llm_response(response: str):
+            async for result in send_event(ChatEvent.START_LLM_RESPONSE, ""):
+                yield result
+            async for result in send_event(ChatEvent.MESSAGE, response):
+                yield result
+            async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""):
+                yield result
+
+        def collect_telemetry():
+            # Gather chat response telemetry
+            nonlocal chat_metadata
+            latency = time.perf_counter() - start_time
+            cmd_set = set([cmd.value for cmd in conversation_commands])
+            chat_metadata = chat_metadata or {}
+            chat_metadata["conversation_command"] = cmd_set
+            chat_metadata["agent"] = conversation.agent.slug if conversation.agent else None
+            chat_metadata["latency"] = f"{latency:.3f}"
+            chat_metadata["ttft_latency"] = f"{ttft:.3f}"
+
+            logger.info(f"Chat response time to first token: {ttft:.3f} seconds")
+            logger.info(f"Chat response total time: {latency:.3f} seconds")
+            update_telemetry_state(
+                request=request,
+                telemetry_type="api",
+                api="chat",
+                client=request.user.client_app,
+                user_agent=request.headers.get("user-agent"),
+                host=request.headers.get("host"),
+                metadata=chat_metadata,
+            )
+
+        conversation = await ConversationAdapters.aget_conversation_by_user(
+            user, client_application=request.user.client_app, conversation_id=conversation_id, title=title
+        )
+        if not conversation:
+            async for result in send_llm_response(f"Conversation {conversation_id} not found"):
+                yield result
             return
 
-        status_packet = {
+        await is_ready_to_chat(user)
-            "type": "status",
-            "message": message,
-            "content-type": "application/json",
-        }
-        try:
-            await websocket.send_text(json.dumps(status_packet))
-        except ConnectionClosedOK:
-            connection_alive = False
-            logger.info(f"User {user} disconnected web socket. Emitting rest of responses to clear thread")
 
-    async def send_complete_llm_response(llm_response: str):
+        user_name = await aget_user_name(user)
-        nonlocal connection_alive
+        location = None
-        if not connection_alive:
+        if city or region or country:
-            return
+            location = LocationData(city=city, region=region, country=country)
-        try:
-            await websocket.send_text("start_llm_response")
-            await websocket.send_text(llm_response)
-            await websocket.send_text("end_llm_response")
-        except ConnectionClosedOK:
-            connection_alive = False
-            logger.info(f"User {user} disconnected web socket. Emitting rest of responses to clear thread")
-
-    async def send_message(message: str):
-        nonlocal connection_alive
-        if not connection_alive:
-            return
-        try:
-            await websocket.send_text(message)
-        except ConnectionClosedOK:
-            connection_alive = False
-            logger.info(f"User {user} disconnected web socket. Emitting rest of responses to clear thread")
-
-    async def send_rate_limit_message(message: str):
-        nonlocal connection_alive
-        if not connection_alive:
-            return
-
-        status_packet = {
-            "type": "rate_limit",
-            "message": message,
-            "content-type": "application/json",
-        }
-        try:
-            await websocket.send_text(json.dumps(status_packet))
-        except ConnectionClosedOK:
-            connection_alive = False
-            logger.info(f"User {user} disconnected web socket. Emitting rest of responses to clear thread")
-
-    user: KhojUser = websocket.user.object
-    conversation = await ConversationAdapters.aget_conversation_by_user(
-        user, client_application=websocket.user.client_app, conversation_id=conversation_id
-    )
-
-    hourly_limiter = ApiUserRateLimiter(requests=5, subscribed_requests=60, window=60, slug="chat_minute")
-
-    daily_limiter = ApiUserRateLimiter(requests=5, subscribed_requests=600, window=60 * 60 * 24, slug="chat_day")
-
-    await is_ready_to_chat(user)
-
-    user_name = await aget_user_name(user)
-
-    location = None
-
-    if city or region or country:
-        location = LocationData(city=city, region=region, country=country)
-
-    await websocket.accept()
-    while connection_alive:
-        try:
-            if conversation:
-                await sync_to_async(conversation.refresh_from_db)(fields=["conversation_log"])
-            q = await websocket.receive_text()
-
-            # Refresh these because the connection to the database might have been closed
-            await conversation.arefresh_from_db()
-
-        except WebSocketDisconnect:
-            logger.debug(f"User {user} disconnected web socket")
-            break
-
-        try:
-            await sync_to_async(hourly_limiter)(websocket)
-            await sync_to_async(daily_limiter)(websocket)
-        except HTTPException as e:
-            await send_rate_limit_message(e.detail)
-            break
 
         if is_query_empty(q):
-            await send_message("start_llm_response")
+            async for result in send_llm_response("Please ask your query to get started."):
-            await send_message(
+                yield result
-                "It seems like your query is incomplete. Could you please provide more details or specify what you need help with?"
+            return
-            )
-            await send_message("end_llm_response")
-            continue
 
         user_message_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         conversation_commands = [get_conversation_command(query=q, any_references=True)]
 
-        await send_status_update(f"**Understanding Query**: {q}")
+        async for result in send_event(ChatEvent.STATUS, f"**Understanding Query**: {q}"):
+            yield result
 
         meta_log = conversation.conversation_log
         is_automated_task = conversation_commands == [ConversationCommand.AutomatedTask]
-        used_slash_summarize = conversation_commands == [ConversationCommand.Summarize]
 
         if conversation_commands == [ConversationCommand.Default] or is_automated_task:
             conversation_commands = await aget_relevant_information_sources(q, meta_log, is_automated_task)
             conversation_commands_str = ", ".join([cmd.value for cmd in conversation_commands])
-            await send_status_update(f"**Chose Data Sources to Search:** {conversation_commands_str}")
+            async for result in send_event(
+                ChatEvent.STATUS, f"**Chose Data Sources to Search:** {conversation_commands_str}"
+            ):
+                yield result
 
             mode = await aget_relevant_output_modes(q, meta_log, is_automated_task)
-            await send_status_update(f"**Decided Response Mode:** {mode.value}")
+            async for result in send_event(ChatEvent.STATUS, f"**Decided Response Mode:** {mode.value}"):
+                yield result
             if mode not in conversation_commands:
                 conversation_commands.append(mode)
 
         for cmd in conversation_commands:
-            await conversation_command_rate_limiter.update_and_check_if_valid(websocket, cmd)
+            await conversation_command_rate_limiter.update_and_check_if_valid(request, cmd)
             q = q.replace(f"/{cmd.value}", "").strip()
 
+        used_slash_summarize = conversation_commands == [ConversationCommand.Summarize]
         file_filters = conversation.file_filters if conversation else []
         # Skip trying to summarize if
         if (
@@ -669,28 +666,37 @@ async def websocket_endpoint(
             response_log = ""
             if len(file_filters) == 0:
                 response_log = "No files selected for summarization. Please add files using the section on the left."
-                await send_complete_llm_response(response_log)
+                async for result in send_llm_response(response_log):
+                    yield result
             elif len(file_filters) > 1:
                 response_log = "Only one file can be selected for summarization."
-                await send_complete_llm_response(response_log)
+                async for result in send_llm_response(response_log):
+                    yield result
             else:
                 try:
                     file_object = await FileObjectAdapters.async_get_file_objects_by_name(user, file_filters[0])
                     if len(file_object) == 0:
                         response_log = "Sorry, we couldn't find the full text of this file. Please re-upload the document and try again."
-                        await send_complete_llm_response(response_log)
+                        async for result in send_llm_response(response_log):
-                        continue
+                            yield result
+                        return
                     contextual_data = " ".join([file.raw_text for file in file_object])
                     if not q:
                         q = "Create a general summary of the file"
-                    await send_status_update(f"**Constructing Summary Using:** {file_object[0].file_name}")
+                    async for result in send_event(
+                        ChatEvent.STATUS, f"**Constructing Summary Using:** {file_object[0].file_name}"
+                    ):
+                        yield result
+
                     response = await extract_relevant_summary(q, contextual_data)
                     response_log = str(response)
-                    await send_complete_llm_response(response_log)
+                    async for result in send_llm_response(response_log):
+                        yield result
                 except Exception as e:
                     response_log = "Error summarizing file."
                     logger.error(f"Error summarizing file for {user.email}: {e}", exc_info=True)
-                    await send_complete_llm_response(response_log)
+                    async for result in send_llm_response(response_log):
+                        yield result
             await sync_to_async(save_to_conversation_log)(
                 q,
                 response_log,
@@ -698,16 +704,10 @@ async def websocket_endpoint(
                 meta_log,
                 user_message_time,
                 intent_type="summarize",
-                client_application=websocket.user.client_app,
+                client_application=request.user.client_app,
                 conversation_id=conversation_id,
             )
-            update_telemetry_state(
+            return
-                request=websocket,
-                telemetry_type="api",
-                api="chat",
-                metadata={"conversation_command": conversation_commands[0].value},
-            )
-            continue
 
         custom_filters = []
         if conversation_commands == [ConversationCommand.Help]:
@@ -717,8 +717,9 @@ async def websocket_endpoint(
                     conversation_config = await ConversationAdapters.aget_default_conversation_config()
                 model_type = conversation_config.model_type
                 formatted_help = help_message.format(model=model_type, version=state.khoj_version, device=get_device())
-                await send_complete_llm_response(formatted_help)
+                async for result in send_llm_response(formatted_help):
-                continue
+                    yield result
+                return
             # Adding specification to search online specifically on khoj.dev pages.
             custom_filters.append("site:khoj.dev")
             conversation_commands.append(ConversationCommand.Online)
@@ -726,14 +727,14 @@ async def websocket_endpoint(
         if ConversationCommand.Automation in conversation_commands:
             try:
                 automation, crontime, query_to_run, subject = await create_automation(
-                    q, timezone, user, websocket.url, meta_log
+                    q, timezone, user, request.url, meta_log
                 )
             except Exception as e:
                 logger.error(f"Error scheduling task {q} for {user.email}: {e}")
-                await send_complete_llm_response(
+                error_message = f"Unable to create automation. Ensure the automation doesn't already exist."
-                    f"Unable to create automation. Ensure the automation doesn't already exist."
+                async for result in send_llm_response(error_message):
-                )
+                    yield result
-                continue
+                return
 
             llm_response = construct_automation_created_message(automation, crontime, query_to_run, subject)
             await sync_to_async(save_to_conversation_log)(
@@ -743,57 +744,78 @@ async def websocket_endpoint(
                 meta_log,
                 user_message_time,
                 intent_type="automation",
-                client_application=websocket.user.client_app,
+                client_application=request.user.client_app,
                 conversation_id=conversation_id,
                 inferred_queries=[query_to_run],
                 automation_id=automation.id,
             )
-            common = CommonQueryParamsClass(
+            async for result in send_llm_response(llm_response):
-                client=websocket.user.client_app,
+                yield result
-                user_agent=websocket.headers.get("user-agent"),
+            return
-                host=websocket.headers.get("host"),
-            )
-            update_telemetry_state(
-                request=websocket,
-                telemetry_type="api",
-                api="chat",
-                **common.__dict__,
-            )
-            await send_complete_llm_response(llm_response)
-            continue
 
-        compiled_references, inferred_queries, defiltered_query = await extract_references_and_questions(
+        # Gather Context
-            websocket, meta_log, q, 7, 0.18, conversation_id, conversation_commands, location, send_status_update
+        ## Extract Document References
-        )
+        compiled_references, inferred_queries, defiltered_query = [], [], None
+        async for result in extract_references_and_questions(
+            request,
+            meta_log,
+            q,
+            (n or 7),
+            (d or 0.18),
+            conversation_id,
+            conversation_commands,
+            location,
+            partial(send_event, ChatEvent.STATUS),
+        ):
+            if isinstance(result, dict) and ChatEvent.STATUS in result:
+                yield result[ChatEvent.STATUS]
+            else:
+                compiled_references.extend(result[0])
+                inferred_queries.extend(result[1])
+                defiltered_query = result[2]
 
-        if compiled_references:
+        if not is_none_or_empty(compiled_references):
             headings = "\n- " + "\n- ".join(set([c.get("compiled", c).split("\n")[0] for c in compiled_references]))
-            await send_status_update(f"**Found Relevant Notes**: {headings}")
+            async for result in send_event(ChatEvent.STATUS, f"**Found Relevant Notes**: {headings}"):
+                yield result
 
         online_results: Dict = dict()
 
         if conversation_commands == [ConversationCommand.Notes] and not await EntryAdapters.auser_has_entries(user):
-            await send_complete_llm_response(f"{no_entries_found.format()}")
+            async for result in send_llm_response(f"{no_entries_found.format()}"):
-            continue
+                yield result
+            return
 
         if ConversationCommand.Notes in conversation_commands and is_none_or_empty(compiled_references):
             conversation_commands.remove(ConversationCommand.Notes)
 
+        ## Gather Online References
         if ConversationCommand.Online in conversation_commands:
             try:
-                online_results = await search_online(
+                async for result in search_online(
-                    defiltered_query, meta_log, location, send_status_update, custom_filters
+                    defiltered_query, meta_log, location, partial(send_event, ChatEvent.STATUS), custom_filters
-                )
+                ):
+                    if isinstance(result, dict) and ChatEvent.STATUS in result:
+                        yield result[ChatEvent.STATUS]
+                    else:
+                        online_results = result
             except ValueError as e:
-                logger.warning(f"Error searching online: {e}. Attempting to respond without online results")
+                error_message = f"Error searching online: {e}. Attempting to respond without online results"
-                await send_complete_llm_response(
+                logger.warning(error_message)
-                    f"Error searching online: {e}. Attempting to respond without online results"
+                async for result in send_llm_response(error_message):
-                )
+                    yield result
-                continue
+                return
 
+        ## Gather Webpage References
         if ConversationCommand.Webpage in conversation_commands:
             try:
-                direct_web_pages = await read_webpages(defiltered_query, meta_log, location, send_status_update)
+                async for result in read_webpages(
+                    defiltered_query, meta_log, location, partial(send_event, ChatEvent.STATUS)
+                ):
+                    if isinstance(result, dict) and ChatEvent.STATUS in result:
+                        yield result[ChatEvent.STATUS]
+                    else:
+                        direct_web_pages = result
                 webpages = []
                 for query in direct_web_pages:
                     if online_results.get(query):
@@ -803,38 +825,52 @@ async def websocket_endpoint(
 
                     for webpage in direct_web_pages[query]["webpages"]:
                         webpages.append(webpage["link"])
-
+                async for result in send_event(ChatEvent.STATUS, f"**Read web pages**: {webpages}"):
-                await send_status_update(f"**Read web pages**: {webpages}")
+                    yield result
             except ValueError as e:
                 logger.warning(
-                    f"Error directly reading webpages: {e}. Attempting to respond without online results", exc_info=True
+                    f"Error directly reading webpages: {e}. Attempting to respond without online results",
+                    exc_info=True,
                 )
 
+        ## Send Gathered References
+        async for result in send_event(
+            ChatEvent.REFERENCES,
+            {
+                "inferredQueries": inferred_queries,
+                "context": compiled_references,
+                "onlineContext": online_results,
+            },
+        ):
+            yield result
+
+        # Generate Output
+        ## Generate Image Output
         if ConversationCommand.Image in conversation_commands:
-            update_telemetry_state(
+            async for result in text_to_image(
-                request=websocket,
-                telemetry_type="api",
-                api="chat",
-                metadata={"conversation_command": conversation_commands[0].value},
-            )
-            image, status_code, improved_image_prompt, intent_type = await text_to_image(
                 q,
                 user,
                 meta_log,
                 location_data=location,
                 references=compiled_references,
                 online_results=online_results,
-                send_status_func=send_status_update,
+                send_status_func=partial(send_event, ChatEvent.STATUS),
-            )
+            ):
+                if isinstance(result, dict) and ChatEvent.STATUS in result:
+                    yield result[ChatEvent.STATUS]
+                else:
+                    image, status_code, improved_image_prompt, intent_type = result
+
             if image is None or status_code != 200:
                 content_obj = {
-                    "image": image,
+                    "content-type": "application/json",
                     "intentType": intent_type,
                     "detail": improved_image_prompt,
-                    "content-type": "application/json",
+                    "image": image,
                 }
-                await send_complete_llm_response(json.dumps(content_obj))
+                async for result in send_llm_response(json.dumps(content_obj)):
-                continue
+                    yield result
+                return
 
             await sync_to_async(save_to_conversation_log)(
                 q,
@@ -844,17 +880,23 @@ async def websocket_endpoint(
                 user_message_time,
                 intent_type=intent_type,
                 inferred_queries=[improved_image_prompt],
-                client_application=websocket.user.client_app,
+                client_application=request.user.client_app,
                 conversation_id=conversation_id,
                 compiled_references=compiled_references,
                 online_results=online_results,
             )
-            content_obj = {"image": image, "intentType": intent_type, "inferredQueries": [improved_image_prompt], "context": compiled_references, "content-type": "application/json", "online_results": online_results}  # type: ignore
+            content_obj = {
+                "intentType": intent_type,
+                "inferredQueries": [improved_image_prompt],
+                "image": image,
+            }
+            async for result in send_llm_response(json.dumps(content_obj)):
+                yield result
+            return
 
-            await send_complete_llm_response(json.dumps(content_obj))
+        ## Generate Text Output
-            continue
+        async for result in send_event(ChatEvent.STATUS, f"**Generating a well-informed response**"):
-
+            yield result
-        await send_status_update(f"**Generating a well-informed response**")
         llm_response, chat_metadata = await agenerate_chat_response(
             defiltered_query,
             meta_log,
@@ -864,310 +906,49 @@ async def websocket_endpoint(
             inferred_queries,
             conversation_commands,
             user,
-            websocket.user.client_app,
+            request.user.client_app,
             conversation_id,
             location,
             user_name,
         )
 
-        chat_metadata["agent"] = conversation.agent.slug if conversation.agent else None
+        # Send Response
+        async for result in send_event(ChatEvent.START_LLM_RESPONSE, ""):
+            yield result
 
-        update_telemetry_state(
+        continue_stream = True
-            request=websocket,
-            telemetry_type="api",
-            api="chat",
-            metadata=chat_metadata,
-        )
         iterator = AsyncIteratorWrapper(llm_response)
-
-        await send_message("start_llm_response")
-
         async for item in iterator:
             if item is None:
-                break
+                async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""):
-            if connection_alive:
+                    yield result
-                try:
+                logger.debug("Finished streaming response")
-                    await send_message(f"{item}")
+                return
-                except ConnectionClosedOK:
+            if not connection_alive or not continue_stream:
-                    connection_alive = False
+                continue
-                    logger.info(f"User {user} disconnected web socket. Emitting rest of responses to clear thread")
-
-        await send_message("end_llm_response")
-
-
-@api_chat.get("", response_class=Response)
-@requires(["authenticated"])
-async def chat(
-    request: Request,
-    common: CommonQueryParams,
-    q: str,
-    n: Optional[int] = 5,
-    d: Optional[float] = 0.22,
-    stream: Optional[bool] = False,
-    title: Optional[str] = None,
-    conversation_id: Optional[int] = None,
-    city: Optional[str] = None,
-    region: Optional[str] = None,
-    country: Optional[str] = None,
-    timezone: Optional[str] = None,
-    rate_limiter_per_minute=Depends(
-        ApiUserRateLimiter(requests=5, subscribed_requests=60, window=60, slug="chat_minute")
-    ),
-    rate_limiter_per_day=Depends(
-        ApiUserRateLimiter(requests=5, subscribed_requests=600, window=60 * 60 * 24, slug="chat_day")
-    ),
-) -> Response:
-    user: KhojUser = request.user.object
-    q = unquote(q)
-    if is_query_empty(q):
-        return Response(
-            content="It seems like your query is incomplete. Could you please provide more details or specify what you need help with?",
-            media_type="text/plain",
-            status_code=400,
-        )
-    user_message_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    logger.info(f"Chat request by {user.username}: {q}")
-
-    await is_ready_to_chat(user)
-    conversation_commands = [get_conversation_command(query=q, any_references=True)]
-
-    _custom_filters = []
-    if conversation_commands == [ConversationCommand.Help]:
-        help_str = "/" + ConversationCommand.Help
-        if q.strip() == help_str:
-            conversation_config = await ConversationAdapters.aget_user_conversation_config(user)
-            if conversation_config == None:
-                conversation_config = await ConversationAdapters.aget_default_conversation_config()
-            model_type = conversation_config.model_type
-            formatted_help = help_message.format(model=model_type, version=state.khoj_version, device=get_device())
-            return StreamingResponse(iter([formatted_help]), media_type="text/event-stream", status_code=200)
-        # Adding specification to search online specifically on khoj.dev pages.
-        _custom_filters.append("site:khoj.dev")
-        conversation_commands.append(ConversationCommand.Online)
-
-    conversation = await ConversationAdapters.aget_conversation_by_user(
-        user, request.user.client_app, conversation_id, title
-    )
-    conversation_id = conversation.id if conversation else None
-
-    if not conversation:
-        return Response(
-            content=f"No conversation found with requested id, title", media_type="text/plain", status_code=400
-        )
-    else:
-        meta_log = conversation.conversation_log
-
-    if ConversationCommand.Summarize in conversation_commands:
-        file_filters = conversation.file_filters
-        llm_response = ""
-        if len(file_filters) == 0:
-            llm_response = "No files selected for summarization. Please add files using the section on the left."
-        elif len(file_filters) > 1:
-            llm_response = "Only one file can be selected for summarization."
-        else:
             try:
-                file_object = await FileObjectAdapters.async_get_file_objects_by_name(user, file_filters[0])
+                async for result in send_event(ChatEvent.MESSAGE, f"{item}"):
-                if len(file_object) == 0:
+                    yield result
-                    llm_response = "Sorry, we couldn't find the full text of this file. Please re-upload the document and try again."
-                    return StreamingResponse(content=llm_response, media_type="text/event-stream", status_code=200)
-                contextual_data = " ".join([file.raw_text for file in file_object])
-                summarizeStr = "/" + ConversationCommand.Summarize
-                if q.strip() == summarizeStr:
-                    q = "Create a general summary of the file"
-                response = await extract_relevant_summary(q, contextual_data)
-                llm_response = str(response)
             except Exception as e:
-                logger.error(f"Error summarizing file for {user.email}: {e}")
+                continue_stream = False
-                llm_response = "Error summarizing file."
+                logger.info(f"User {user} disconnected. Emitting rest of responses to clear thread: {e}")
-        await sync_to_async(save_to_conversation_log)(
-            q,
-            llm_response,
-            user,
-            conversation.conversation_log,
-            user_message_time,
-            intent_type="summarize",
-            client_application=request.user.client_app,
-            conversation_id=conversation_id,
-        )
-        update_telemetry_state(
-            request=request,
-            telemetry_type="api",
-            api="chat",
-            metadata={"conversation_command": conversation_commands[0].value},
-            **common.__dict__,
-        )
-        return StreamingResponse(content=llm_response, media_type="text/event-stream", status_code=200)
-
-    is_automated_task = conversation_commands == [ConversationCommand.AutomatedTask]
-
-    if conversation_commands == [ConversationCommand.Default] or is_automated_task:
-        conversation_commands = await aget_relevant_information_sources(q, meta_log, is_automated_task)
-        mode = await aget_relevant_output_modes(q, meta_log, is_automated_task)
-        if mode not in conversation_commands:
-            conversation_commands.append(mode)
-
-    for cmd in conversation_commands:
-        await conversation_command_rate_limiter.update_and_check_if_valid(request, cmd)
-        q = q.replace(f"/{cmd.value}", "").strip()
-
-    location = None
-
-    if city or region or country:
-        location = LocationData(city=city, region=region, country=country)
-
-    user_name = await aget_user_name(user)
-
-    if ConversationCommand.Automation in conversation_commands:
-        try:
-            automation, crontime, query_to_run, subject = await create_automation(
-                q, timezone, user, request.url, meta_log
-            )
-        except Exception as e:
-            logger.error(f"Error creating automation {q} for {user.email}: {e}", exc_info=True)
-            return Response(
-                content=f"Unable to create automation. Ensure the automation doesn't already exist.",
-                media_type="text/plain",
-                status_code=500,
-            )
-
-        llm_response = construct_automation_created_message(automation, crontime, query_to_run, subject)
-        await sync_to_async(save_to_conversation_log)(
-            q,
-            llm_response,
-            user,
-            meta_log,
-            user_message_time,
-            intent_type="automation",
-            client_application=request.user.client_app,
-            conversation_id=conversation_id,
-            inferred_queries=[query_to_run],
-            automation_id=automation.id,
-        )
-
-        if stream:
-            return StreamingResponse(llm_response, media_type="text/event-stream", status_code=200)
-        else:
-            return Response(content=llm_response, media_type="text/plain", status_code=200)
-
-    compiled_references, inferred_queries, defiltered_query = await extract_references_and_questions(
-        request, meta_log, q, (n or 5), (d or math.inf), conversation_id, conversation_commands, location
-    )
-    online_results: Dict[str, Dict] = {}
-
-    if conversation_commands == [ConversationCommand.Notes] and not await EntryAdapters.auser_has_entries(user):
-        no_entries_found_format = no_entries_found.format()
-        if stream:
-            return StreamingResponse(iter([no_entries_found_format]), media_type="text/event-stream", status_code=200)
-        else:
-            response_obj = {"response": no_entries_found_format}
-            return Response(content=json.dumps(response_obj), media_type="text/plain", status_code=200)
-
-    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(compiled_references):
-        no_notes_found_format = no_notes_found.format()
-        if stream:
-            return StreamingResponse(iter([no_notes_found_format]), media_type="text/event-stream", status_code=200)
-        else:
-            response_obj = {"response": no_notes_found_format}
-            return Response(content=json.dumps(response_obj), media_type="text/plain", status_code=200)
-
-    if ConversationCommand.Notes in conversation_commands and is_none_or_empty(compiled_references):
-        conversation_commands.remove(ConversationCommand.Notes)
-
-    if ConversationCommand.Online in conversation_commands:
-        try:
-            online_results = await search_online(defiltered_query, meta_log, location, custom_filters=_custom_filters)
-        except ValueError as e:
-            logger.warning(f"Error searching online: {e}. Attempting to respond without online results")
-
-    if ConversationCommand.Webpage in conversation_commands:
-        try:
-            online_results = await read_webpages(defiltered_query, meta_log, location)
-        except ValueError as e:
-            logger.warning(
-                f"Error directly reading webpages: {e}. Attempting to respond without online results", exc_info=True
-            )
-
-    if ConversationCommand.Image in conversation_commands:
-        update_telemetry_state(
-            request=request,
-            telemetry_type="api",
-            api="chat",
-            metadata={"conversation_command": conversation_commands[0].value},
-            **common.__dict__,
-        )
-        image, status_code, improved_image_prompt, intent_type = await text_to_image(
-            q, user, meta_log, location_data=location, references=compiled_references, online_results=online_results
-        )
-        if image is None:
-            content_obj = {"image": image, "intentType": intent_type, "detail": improved_image_prompt}
-            return Response(content=json.dumps(content_obj), media_type="application/json", status_code=status_code)
-
-        await sync_to_async(save_to_conversation_log)(
-            q,
-            image,
-            user,
-            meta_log,
-            user_message_time,
-            intent_type=intent_type,
-            inferred_queries=[improved_image_prompt],
-            client_application=request.user.client_app,
-            conversation_id=conversation.id,
-            compiled_references=compiled_references,
-            online_results=online_results,
-        )
-        content_obj = {"image": image, "intentType": intent_type, "inferredQueries": [improved_image_prompt], "context": compiled_references, "online_results": online_results}  # type: ignore
-        return Response(content=json.dumps(content_obj), media_type="application/json", status_code=status_code)
-
-    # Get the (streamed) chat response from the LLM of choice.
-    llm_response, chat_metadata = await agenerate_chat_response(
-        defiltered_query,
-        meta_log,
-        conversation,
-        compiled_references,
-        online_results,
-        inferred_queries,
-        conversation_commands,
-        user,
-        request.user.client_app,
-        conversation.id,
-        location,
-        user_name,
-    )
-
-    cmd_set = set([cmd.value for cmd in conversation_commands])
-    chat_metadata["conversation_command"] = cmd_set
-    chat_metadata["agent"] = conversation.agent.slug if conversation.agent else None
-
-    update_telemetry_state(
-        request=request,
-        telemetry_type="api",
-        api="chat",
-        metadata=chat_metadata,
-        **common.__dict__,
-    )
-
-    if llm_response is None:
-        return Response(content=llm_response, media_type="text/plain", status_code=500)
 
+    ## Stream Text Response
     if stream:
-        return StreamingResponse(llm_response, media_type="text/event-stream", status_code=200)
+        return StreamingResponse(event_generator(q), media_type="text/plain")
+    ## Non-Streaming Text Response
+    else:
+        # Get the full response from the generator if the stream is not requested.
+        response_obj = {}
+        actual_response = ""
+        iterator = event_generator(q)
+        async for item in iterator:
+            try:
+                item_json = json.loads(item)
+                if "type" in item_json and item_json["type"] == ChatEvent.REFERENCES.value:
+                    response_obj = item_json["data"]
+            except:
+                actual_response += item
+        response_obj["response"] = actual_response
 
-    iterator = AsyncIteratorWrapper(llm_response)
+        return Response(content=json.dumps(response_obj), media_type="application/json", status_code=200)
-
-    # Get the full response from the generator if the stream is not requested.
-    aggregated_gpt_response = ""
-    async for item in iterator:
-        if item is None:
-            break
-        aggregated_gpt_response += item
-
-    actual_response = aggregated_gpt_response.split("### compiled references:")[0]
-
-    response_obj = {
-        "response": actual_response,
-        "inferredQueries": inferred_queries,
-        "context": compiled_references,
-        "online_results": online_results,
-    }
-
-    return Response(content=json.dumps(response_obj), media_type="application/json", status_code=200)

src/khoj/routers/auth.py

@@ -41,7 +41,7 @@ if not state.anonymous_mode:
         from google.auth.transport import requests as google_requests
         from google.oauth2 import id_token
     except ImportError:
-        missing_requirements += ["Install the Khoj production package with `pip install khoj-assistant[prod]`"]
+        missing_requirements += ["Install the Khoj production package with `pip install khoj[prod]`"]
     if not os.environ.get("RESEND_API_KEY") and (
         not os.environ.get("GOOGLE_CLIENT_ID") or not os.environ.get("GOOGLE_CLIENT_SECRET")
     ):

src/khoj/routers/helpers.py

@@ -9,6 +9,7 @@ import os
 import re
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timedelta, timezone
+from enum import Enum
 from functools import partial
 from random import random
 from typing import (
@@ -330,6 +331,7 @@ async def aget_relevant_output_modes(query: str, conversation_history: dict, is_
             # Check whether the tool exists as a valid ConversationCommand
             return ConversationCommand(response)
 
+        logger.error(f"Invalid output mode selected: {response}. Defaulting to text.")
         return ConversationCommand.Text
     except Exception:
         logger.error(f"Invalid response for determining relevant mode: {response}")
@@ -542,9 +544,6 @@ async def send_message_to_model_wrapper(
         chat_model_option or await ConversationAdapters.aget_default_conversation_config()
     )
 
-    if conversation_config is None:
-        raise HTTPException(status_code=500, detail="Contact the server administrator to set a default chat model.")
-
     chat_model = conversation_config.chat_model
     max_tokens = conversation_config.max_prompt_size
     tokenizer = conversation_config.tokenizer
@@ -778,7 +777,7 @@ async def text_to_image(
     references: List[Dict[str, Any]],
     online_results: Dict[str, Any],
     send_status_func: Optional[Callable] = None,
-) -> Tuple[Optional[str], int, Optional[str], str]:
+):
     status_code = 200
     image = None
     response = None
@@ -790,7 +789,8 @@ async def text_to_image(
         # If the user has not configured a text to image model, return an unsupported on server error
         status_code = 501
         message = "Failed to generate image. Setup image generation on the server."
-        return image_url or image, status_code, message, intent_type.value
+        yield image_url or image, status_code, message, intent_type.value
+        return
 
     text2image_model = text_to_image_config.model_name
     chat_history = ""
@@ -802,20 +802,21 @@ async def text_to_image(
             chat_history += f"Q: Prompt: {chat['intent']['query']}\n"
             chat_history += f"A: Improved Prompt: {chat['intent']['inferred-queries'][0]}\n"
 
-    with timer("Improve the original user query", logger):
+    if send_status_func:
-        if send_status_func:
+        async for event in send_status_func("**Enhancing the Painting Prompt**"):
-            await send_status_func("**Enhancing the Painting Prompt**")
+            yield {ChatEvent.STATUS: event}
-        improved_image_prompt = await generate_better_image_prompt(
+    improved_image_prompt = await generate_better_image_prompt(
-            message,
+        message,
-            chat_history,
+        chat_history,
-            location_data=location_data,
+        location_data=location_data,
-            note_references=references,
+        note_references=references,
-            online_results=online_results,
+        online_results=online_results,
-            model_type=text_to_image_config.model_type,
+        model_type=text_to_image_config.model_type,
-        )
+    )
 
     if send_status_func:
-        await send_status_func(f"**🖼️ Painting using Enhanced Prompt**:\n{improved_image_prompt}")
+        async for event in send_status_func(f"**🖼️ Painting using Enhanced Prompt**:\n{improved_image_prompt}"):
+            yield {ChatEvent.STATUS: event}
 
     if text_to_image_config.model_type == TextToImageModelConfig.ModelType.OPENAI:
         with timer("Generate image with OpenAI", logger):
@@ -840,12 +841,14 @@ async def text_to_image(
                     logger.error(f"Image Generation blocked by OpenAI: {e}")
                     status_code = e.status_code  # type: ignore
                     message = f"Image generation blocked by OpenAI: {e.message}"  # type: ignore
-                    return image_url or image, status_code, message, intent_type.value
+                    yield image_url or image, status_code, message, intent_type.value
+                    return
                 else:
                     logger.error(f"Image Generation failed with {e}", exc_info=True)
                     message = f"Image generation failed with OpenAI error: {e.message}"  # type: ignore
                     status_code = e.status_code  # type: ignore
-                    return image_url or image, status_code, message, intent_type.value
+                    yield image_url or image, status_code, message, intent_type.value
+                    return
 
     elif text_to_image_config.model_type == TextToImageModelConfig.ModelType.STABILITYAI:
         with timer("Generate image with Stability AI", logger):
@@ -867,7 +870,8 @@ async def text_to_image(
                 logger.error(f"Image Generation failed with {e}", exc_info=True)
                 message = f"Image generation failed with Stability AI error: {e}"
                 status_code = e.status_code  # type: ignore
-                return image_url or image, status_code, message, intent_type.value
+                yield image_url or image, status_code, message, intent_type.value
+                return
 
     with timer("Convert image to webp", logger):
         # Convert png to webp for faster loading
@@ -887,7 +891,7 @@ async def text_to_image(
         intent_type = ImageIntentType.TEXT_TO_IMAGE_V3
         image = base64.b64encode(webp_image_bytes).decode("utf-8")
 
-    return image_url or image, status_code, improved_image_prompt, intent_type.value
+    yield image_url or image, status_code, improved_image_prompt, intent_type.value
 
 
 class ApiUserRateLimiter:
@@ -1211,6 +1215,14 @@ Manage your automations [here](/automations).
     """.strip()
 
 
+class ChatEvent(Enum):
+    START_LLM_RESPONSE = "start_llm_response"
+    END_LLM_RESPONSE = "end_llm_response"
+    MESSAGE = "message"
+    REFERENCES = "references"
+    STATUS = "status"
+
+
 def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False):
     user_picture = request.session.get("user", {}).get("picture")
     is_active = has_required_scope(request, ["premium"])

src/khoj/routers/indexer.py

@@ -0,0 +1,349 @@
+import asyncio
+import logging
+from typing import Dict, Optional, Union
+
+from fastapi import APIRouter, Depends, Header, Request, Response, UploadFile
+from pydantic import BaseModel
+from starlette.authentication import requires
+
+from khoj.database.models import GithubConfig, KhojUser, NotionConfig
+from khoj.processor.content.docx.docx_to_entries import DocxToEntries
+from khoj.processor.content.github.github_to_entries import GithubToEntries
+from khoj.processor.content.images.image_to_entries import ImageToEntries
+from khoj.processor.content.markdown.markdown_to_entries import MarkdownToEntries
+from khoj.processor.content.notion.notion_to_entries import NotionToEntries
+from khoj.processor.content.org_mode.org_to_entries import OrgToEntries
+from khoj.processor.content.pdf.pdf_to_entries import PdfToEntries
+from khoj.processor.content.plaintext.plaintext_to_entries import PlaintextToEntries
+from khoj.routers.helpers import ApiIndexedDataLimiter, update_telemetry_state
+from khoj.search_type import text_search
+from khoj.utils import constants, state
+from khoj.utils.config import SearchModels
+from khoj.utils.helpers import LRU, get_file_type
+from khoj.utils.rawconfig import ContentConfig, FullConfig, SearchConfig
+from khoj.utils.yaml import save_config_to_file_updated_state
+
+logger = logging.getLogger(__name__)
+
+indexer = APIRouter()
+
+
+class File(BaseModel):
+    path: str
+    content: Union[str, bytes]
+
+
+class IndexBatchRequest(BaseModel):
+    files: list[File]
+
+
+class IndexerInput(BaseModel):
+    org: Optional[dict[str, str]] = None
+    markdown: Optional[dict[str, str]] = None
+    pdf: Optional[dict[str, bytes]] = None
+    plaintext: Optional[dict[str, str]] = None
+    image: Optional[dict[str, bytes]] = None
+    docx: Optional[dict[str, bytes]] = None
+
+
+@indexer.post("/update")
+@requires(["authenticated"])
+async def update(
+    request: Request,
+    files: list[UploadFile],
+    force: bool = False,
+    t: Optional[Union[state.SearchType, str]] = state.SearchType.All,
+    client: Optional[str] = None,
+    user_agent: Optional[str] = Header(None),
+    referer: Optional[str] = Header(None),
+    host: Optional[str] = Header(None),
+    indexed_data_limiter: ApiIndexedDataLimiter = Depends(
+        ApiIndexedDataLimiter(
+            incoming_entries_size_limit=10,
+            subscribed_incoming_entries_size_limit=75,
+            total_entries_size_limit=10,
+            subscribed_total_entries_size_limit=100,
+        )
+    ),
+):
+    user = request.user.object
+    index_files: Dict[str, Dict[str, str]] = {
+        "org": {},
+        "markdown": {},
+        "pdf": {},
+        "plaintext": {},
+        "image": {},
+        "docx": {},
+    }
+    try:
+        logger.info(f"📬 Updating content index via API call by {client} client")
+        for file in files:
+            file_content = file.file.read()
+            file_type, encoding = get_file_type(file.content_type, file_content)
+            if file_type in index_files:
+                index_files[file_type][file.filename] = file_content.decode(encoding) if encoding else file_content
+            else:
+                logger.warning(f"Skipped indexing unsupported file type sent by {client} client: {file.filename}")
+
+        indexer_input = IndexerInput(
+            org=index_files["org"],
+            markdown=index_files["markdown"],
+            pdf=index_files["pdf"],
+            plaintext=index_files["plaintext"],
+            image=index_files["image"],
+            docx=index_files["docx"],
+        )
+
+        if state.config == None:
+            logger.info("📬 Initializing content index on first run.")
+            default_full_config = FullConfig(
+                content_type=None,
+                search_type=SearchConfig.model_validate(constants.default_config["search-type"]),
+                processor=None,
+            )
+            state.config = default_full_config
+            default_content_config = ContentConfig(
+                org=None,
+                markdown=None,
+                pdf=None,
+                docx=None,
+                image=None,
+                github=None,
+                notion=None,
+                plaintext=None,
+            )
+            state.config.content_type = default_content_config
+            save_config_to_file_updated_state()
+            configure_search(state.search_models, state.config.search_type)
+
+        # Extract required fields from config
+        loop = asyncio.get_event_loop()
+        success = await loop.run_in_executor(
+            None,
+            configure_content,
+            indexer_input.model_dump(),
+            force,
+            t,
+            False,
+            user,
+        )
+        if not success:
+            raise RuntimeError("Failed to update content index")
+        logger.info(f"Finished processing batch indexing request")
+    except Exception as e:
+        logger.error(f"Failed to process batch indexing request: {e}", exc_info=True)
+        logger.error(
+            f'🚨 Failed to {"force " if force else ""}update {t} content index triggered via API call by {client} client: {e}',
+            exc_info=True,
+        )
+        return Response(content="Failed", status_code=500)
+
+    indexing_metadata = {
+        "num_org": len(index_files["org"]),
+        "num_markdown": len(index_files["markdown"]),
+        "num_pdf": len(index_files["pdf"]),
+        "num_plaintext": len(index_files["plaintext"]),
+        "num_image": len(index_files["image"]),
+        "num_docx": len(index_files["docx"]),
+    }
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="index/update",
+        client=client,
+        user_agent=user_agent,
+        referer=referer,
+        host=host,
+        metadata=indexing_metadata,
+    )
+
+    logger.info(f"📪 Content index updated via API call by {client} client")
+
+    indexed_filenames = ",".join(file for ctype in index_files for file in index_files[ctype]) or ""
+    return Response(content=indexed_filenames, status_code=200)
+
+
+def configure_search(search_models: SearchModels, search_config: Optional[SearchConfig]) -> Optional[SearchModels]:
+    # Run Validation Checks
+    if search_models is None:
+        search_models = SearchModels()
+
+    return search_models
+
+
+def configure_content(
+    files: Optional[dict[str, dict[str, str]]],
+    regenerate: bool = False,
+    t: Optional[state.SearchType] = state.SearchType.All,
+    full_corpus: bool = True,
+    user: KhojUser = None,
+) -> bool:
+    success = True
+    if t == None:
+        t = state.SearchType.All
+
+    if t is not None and t in [type.value for type in state.SearchType]:
+        t = state.SearchType(t)
+
+    if t is not None and not t.value in [type.value for type in state.SearchType]:
+        logger.warning(f"🚨 Invalid search type: {t}")
+        return False
+
+    search_type = t.value if t else None
+
+    no_documents = all([not files.get(file_type) for file_type in files])
+
+    if files is None:
+        logger.warning(f"🚨 No files to process for {search_type} search.")
+        return True
+
+    try:
+        # Initialize Org Notes Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Org.value) and files["org"]:
+            logger.info("🦄 Setting up search for orgmode notes")
+            # Extract Entries, Generate Notes Embeddings
+            text_search.setup(
+                OrgToEntries,
+                files.get("org"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup org: {e}", exc_info=True)
+        success = False
+
+    try:
+        # Initialize Markdown Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Markdown.value) and files[
+            "markdown"
+        ]:
+            logger.info("💎 Setting up search for markdown notes")
+            # Extract Entries, Generate Markdown Embeddings
+            text_search.setup(
+                MarkdownToEntries,
+                files.get("markdown"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup markdown: {e}", exc_info=True)
+        success = False
+
+    try:
+        # Initialize PDF Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Pdf.value) and files["pdf"]:
+            logger.info("🖨️ Setting up search for pdf")
+            # Extract Entries, Generate PDF Embeddings
+            text_search.setup(
+                PdfToEntries,
+                files.get("pdf"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup PDF: {e}", exc_info=True)
+        success = False
+
+    try:
+        # Initialize Plaintext Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Plaintext.value) and files[
+            "plaintext"
+        ]:
+            logger.info("📄 Setting up search for plaintext")
+            # Extract Entries, Generate Plaintext Embeddings
+            text_search.setup(
+                PlaintextToEntries,
+                files.get("plaintext"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup plaintext: {e}", exc_info=True)
+        success = False
+
+    try:
+        if no_documents:
+            github_config = GithubConfig.objects.filter(user=user).prefetch_related("githubrepoconfig").first()
+            if (
+                search_type == state.SearchType.All.value or search_type == state.SearchType.Github.value
+            ) and github_config is not None:
+                logger.info("🐙 Setting up search for github")
+                # Extract Entries, Generate Github Embeddings
+                text_search.setup(
+                    GithubToEntries,
+                    None,
+                    regenerate=regenerate,
+                    full_corpus=full_corpus,
+                    user=user,
+                    config=github_config,
+                )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup GitHub: {e}", exc_info=True)
+        success = False
+
+    try:
+        if no_documents:
+            # Initialize Notion Search
+            notion_config = NotionConfig.objects.filter(user=user).first()
+            if (
+                search_type == state.SearchType.All.value or search_type == state.SearchType.Notion.value
+            ) and notion_config:
+                logger.info("🔌 Setting up search for notion")
+                text_search.setup(
+                    NotionToEntries,
+                    None,
+                    regenerate=regenerate,
+                    full_corpus=full_corpus,
+                    user=user,
+                    config=notion_config,
+                )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup Notion: {e}", exc_info=True)
+        success = False
+
+    try:
+        # Initialize Image Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Image.value) and files[
+            "image"
+        ]:
+            logger.info("🖼️ Setting up search for images")
+            # Extract Entries, Generate Image Embeddings
+            text_search.setup(
+                ImageToEntries,
+                files.get("image"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup images: {e}", exc_info=True)
+        success = False
+    try:
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Docx.value) and files["docx"]:
+            logger.info("📄 Setting up search for docx")
+            text_search.setup(
+                DocxToEntries,
+                files.get("docx"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup docx: {e}", exc_info=True)
+        success = False
+
+    # Invalidate Query Cache
+    if user:
+        state.query_cache[user.uuid] = LRU()
+
+    return success

src/khoj/utils/cli.py

@@ -59,7 +59,7 @@ def cli(args=None):
     # Set default values for arguments
     args.chat_on_gpu = not args.disable_chat_on_gpu
 
-    args.version_no = version("khoj-assistant")
+    args.version_no = version("khoj")
     if args.version:
         # Show version of khoj installed and exit
         print(args.version_no)

src/khoj/utils/fs_syncer.py

@@ -22,7 +22,7 @@ magika = Magika()
 
 
 def collect_files(search_type: Optional[SearchType] = SearchType.All, user=None) -> dict:
-    files = {}
+    files: dict[str, dict] = {"docx": {}, "image": {}}
 
     if search_type == SearchType.All or search_type == SearchType.Org:
         org_config = LocalOrgConfig.objects.filter(user=user).first()

src/khoj/utils/helpers.py

@@ -259,7 +259,7 @@ def log_telemetry(
     # Populate telemetry data to log
     request_body = {
         "telemetry_type": telemetry_type,
-        "server_version": version("khoj-assistant"),
+        "server_version": version("khoj"),
         "os": platform.system(),
         "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
     }

tests/test_client.py

@@ -61,7 +61,7 @@ def test_search_with_invalid_content_type(client):
 @pytest.mark.django_db(transaction=True)
 def test_search_with_valid_content_type(client):
     headers = {"Authorization": "Bearer kk-secret"}
-    for content_type in ["all", "org", "markdown", "image", "pdf", "github", "notion", "plaintext", "docx"]:
+    for content_type in ["all", "org", "markdown", "image", "pdf", "github", "notion", "plaintext", "image", "docx"]:
         # Act
         response = client.get(f"/api/search?q=random&t={content_type}", headers=headers)
         # Assert
@@ -127,6 +127,8 @@ def test_index_update_big_files(client):
     # Arrange
     state.billing_enabled = True
     files = get_big_size_sample_files_data()
+
+    # Credential for the default_user, who is subscribed
     headers = {"Authorization": "Bearer kk-secret"}
 
     # Act
@@ -455,13 +457,13 @@ def test_user_no_data_returns_empty(client, sample_org_data, api_user3: KhojApiU
 
 @pytest.mark.skipif(os.getenv("OPENAI_API_KEY") is None, reason="requires OPENAI_API_KEY")
 @pytest.mark.django_db(transaction=True)
-def test_chat_with_unauthenticated_user(chat_client_with_auth, api_user2: KhojApiUser):
+async def test_chat_with_unauthenticated_user(chat_client_with_auth, api_user2: KhojApiUser):
     # Arrange
     headers = {"Authorization": f"Bearer {api_user2.token}"}
 
     # Act
-    auth_response = chat_client_with_auth.get(f'/api/chat?q="Hello!"&stream=true', headers=headers)
+    auth_response = chat_client_with_auth.get(f'/api/chat?q="Hello!"', headers=headers)
-    no_auth_response = chat_client_with_auth.get(f'/api/chat?q="Hello!"&stream=true')
+    no_auth_response = chat_client_with_auth.get(f'/api/chat?q="Hello!"')
 
     # Assert
     assert auth_response.status_code == 200
@@ -497,7 +499,8 @@ def get_sample_files_data():
 
 
 def get_big_size_sample_files_data():
-    big_text = "a" * (25 * 1024 * 1024)  # a string of approximately 25 MB
+    # a string of approximately 100 MB
+    big_text = "a" * (100 * 1024 * 1024)
     return [
         (
             "files",

tests/test_offline_chat_actors.py

@@ -286,7 +286,7 @@ def test_answer_from_chat_history_and_currently_retrieved_content(loaded_model):
     # Act
     response_gen = converse_offline(
         references=[
-            "Testatron was born on 1st April 1984 in Testville."
+            {"compiled": "Testatron was born on 1st April 1984 in Testville."}
         ],  # Assume context retrieved from notes for the user_query
         user_query="Where was I born?",
         conversation_log=populate_chat_history(message_list),
@@ -341,14 +341,22 @@ def test_answer_requires_current_date_awareness(loaded_model):
     "Chat actor should be able to answer questions relative to current date using provided notes"
     # Arrange
     context = [
-        f"""{datetime.now().strftime("%Y-%m-%d")} "Naco Taco" "Tacos for Dinner"
+        {
-Expenses:Food:Dining  10.00 USD""",
+            "compiled": f"""{datetime.now().strftime("%Y-%m-%d")} "Naco Taco" "Tacos for Dinner"
-        f"""{datetime.now().strftime("%Y-%m-%d")} "Sagar Ratna" "Dosa for Lunch"
+Expenses:Food:Dining  10.00 USD"""
-Expenses:Food:Dining  10.00 USD""",
+        },
-        f"""2020-04-01 "SuperMercado" "Bananas"
+        {
-Expenses:Food:Groceries  10.00 USD""",
+            "compiled": f"""{datetime.now().strftime("%Y-%m-%d")} "Sagar Ratna" "Dosa for Lunch"
-        f"""2020-01-01 "Naco Taco" "Burittos for Dinner"
+Expenses:Food:Dining  10.00 USD"""
-Expenses:Food:Dining  10.00 USD""",
+        },
+        {
+            "compiled": f"""2020-04-01 "SuperMercado" "Bananas"
+Expenses:Food:Groceries  10.00 USD"""
+        },
+        {
+            "compiled": f"""2020-01-01 "Naco Taco" "Burittos for Dinner"
+Expenses:Food:Dining  10.00 USD"""
+        },
     ]
 
     # Act
@@ -373,14 +381,22 @@ def test_answer_requires_date_aware_aggregation_across_provided_notes(loaded_mod
     "Chat actor should be able to answer questions that require date aware aggregation across multiple notes"
     # Arrange
     context = [
-        f"""# {datetime.now().strftime("%Y-%m-%d")} "Naco Taco" "Tacos for Dinner"
+        {
-Expenses:Food:Dining  10.00 USD""",
+            "compiled": f"""# {datetime.now().strftime("%Y-%m-%d")} "Naco Taco" "Tacos for Dinner"
-        f"""{datetime.now().strftime("%Y-%m-%d")} "Sagar Ratna" "Dosa for Lunch"
+Expenses:Food:Dining  10.00 USD"""
-Expenses:Food:Dining  10.00 USD""",
+        },
-        f"""2020-04-01 "SuperMercado" "Bananas"
+        {
-Expenses:Food:Groceries  10.00 USD""",
+            "compiled": f"""{datetime.now().strftime("%Y-%m-%d")} "Sagar Ratna" "Dosa for Lunch"
-        f"""2020-01-01 "Naco Taco" "Burittos for Dinner"
+Expenses:Food:Dining  10.00 USD"""
-Expenses:Food:Dining  10.00 USD""",
+        },
+        {
+            "compiled": f"""2020-04-01 "SuperMercado" "Bananas"
+Expenses:Food:Groceries  10.00 USD"""
+        },
+        {
+            "compiled": f"""2020-01-01 "Naco Taco" "Burittos for Dinner"
+Expenses:Food:Dining  10.00 USD"""
+        },
     ]
 
     # Act
@@ -430,12 +446,18 @@ def test_ask_for_clarification_if_not_enough_context_in_question(loaded_model):
     "Chat actor should ask for clarification if question cannot be answered unambiguously with the provided context"
     # Arrange
     context = [
-        f"""# Ramya
+        {
-My sister, Ramya, is married to Kali Devi. They have 2 kids, Ravi and Rani.""",
+            "compiled": f"""# Ramya
-        f"""# Fang
+My sister, Ramya, is married to Kali Devi. They have 2 kids, Ravi and Rani."""
-My sister, Fang Liu is married to Xi Li. They have 1 kid, Xiao Li.""",
+        },
-        f"""# Aiyla
+        {
-My sister, Aiyla is married to Tolga. They have 3 kids, Yildiz, Ali and Ahmet.""",
+            "compiled": f"""# Fang
+My sister, Fang Liu is married to Xi Li. They have 1 kid, Xiao Li."""
+        },
+        {
+            "compiled": f"""# Aiyla
+My sister, Aiyla is married to Tolga. They have 3 kids, Yildiz, Ali and Ahmet."""
+        },
     ]
 
     # Act
@@ -459,9 +481,9 @@ def test_agent_prompt_should_be_used(loaded_model, offline_agent):
     "Chat actor should ask be tuned to think like an accountant based on the agent definition"
     # Arrange
     context = [
-        f"""I went to the store and bought some bananas for 2.20""",
+        {"compiled": f"""I went to the store and bought some bananas for 2.20"""},
-        f"""I went to the store and bought some apples for 1.30""",
+        {"compiled": f"""I went to the store and bought some apples for 1.30"""},
-        f"""I went to the store and bought some oranges for 6.00""",
+        {"compiled": f"""I went to the store and bought some oranges for 6.00"""},
     ]
 
     # Act
@@ -499,7 +521,7 @@ def test_chat_does_not_exceed_prompt_size(loaded_model):
     "Ensure chat context and response together do not exceed max prompt size for the model"
     # Arrange
     prompt_size_exceeded_error = "ERROR: The prompt size exceeds the context window size and cannot be processed"
-    context = [" ".join([f"{number}" for number in range(2043)])]
+    context = [{"compiled": " ".join([f"{number}" for number in range(2043)])}]
 
     # Act
     response_gen = converse_offline(
@@ -530,7 +552,7 @@ def test_filter_questions():
 # ----------------------------------------------------------------------------------------------------
 @pytest.mark.anyio
 @pytest.mark.django_db(transaction=True)
-async def test_use_default_response_mode(client_offline_chat):
+async def test_use_text_response_mode(client_offline_chat):
     # Arrange
     user_query = "What's the latest in the Israel/Palestine conflict?"
 
@@ -538,7 +560,7 @@ async def test_use_default_response_mode(client_offline_chat):
     mode = await aget_relevant_output_modes(user_query, {})
 
     # Assert
-    assert mode.value == "default"
+    assert mode.value == "text"
 
 
 # ----------------------------------------------------------------------------------------------------

tests/test_offline_chat_director.py

@@ -45,7 +45,6 @@ def create_conversation(message_list, user, agent=None):
 
 # Tests
 # ----------------------------------------------------------------------------------------------------
-@pytest.mark.xfail(AssertionError, reason="Chat director not capable of answering this question yet")
 @pytest.mark.chatquality
 @pytest.mark.django_db(transaction=True)
 def test_offline_chat_with_no_chat_history_or_retrieved_content(client_offline_chat):
@@ -68,10 +67,8 @@ def test_chat_with_online_content(client_offline_chat):
     # Act
     q = "/online give me the link to paul graham's essay how to do great work"
     encoded_q = quote(q, safe="")
-    response = client_offline_chat.get(f"/api/chat?q={encoded_q}&stream=true")
+    response = client_offline_chat.get(f"/api/chat?q={encoded_q}")
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
-
-    response_message = response_message.split("### compiled references")[0]
 
     # Assert
     expected_responses = [
@@ -92,10 +89,8 @@ def test_chat_with_online_webpage_content(client_offline_chat):
     # Act
     q = "/online how many firefighters were involved in the great chicago fire and which year did it take place?"
     encoded_q = quote(q, safe="")
-    response = client_offline_chat.get(f"/api/chat?q={encoded_q}&stream=true")
+    response = client_offline_chat.get(f"/api/chat?q={encoded_q}")
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
-
-    response_message = response_message.split("### compiled references")[0]
 
     # Assert
     expected_responses = ["185", "1871", "horse"]
@@ -179,10 +174,6 @@ def test_answer_from_chat_history_and_previously_retrieved_content(client_offlin
 
 
 # ----------------------------------------------------------------------------------------------------
-@pytest.mark.xfail(
-    AssertionError,
-    reason="Chat director not capable of answering this question yet because it requires extract_questions",
-)
 @pytest.mark.chatquality
 @pytest.mark.django_db(transaction=True)
 def test_answer_from_chat_history_and_currently_retrieved_content(client_offline_chat, default_user2):

tests/test_openai_chat_director.py

@@ -49,8 +49,8 @@ def create_conversation(message_list, user, agent=None):
 @pytest.mark.django_db(transaction=True)
 def test_chat_with_no_chat_history_or_retrieved_content(chat_client):
     # Act
-    response = chat_client.get(f'/api/chat?q="Hello, my name is Testatron. Who are you?"&stream=true')
+    response = chat_client.get(f'/api/chat?q="Hello, my name is Testatron. Who are you?"')
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
 
     # Assert
     expected_responses = ["Khoj", "khoj"]
@@ -67,10 +67,8 @@ def test_chat_with_online_content(chat_client):
     # Act
     q = "/online give me the link to paul graham's essay how to do great work"
     encoded_q = quote(q, safe="")
-    response = chat_client.get(f"/api/chat?q={encoded_q}&stream=true")
+    response = chat_client.get(f"/api/chat?q={encoded_q}")
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
-
-    response_message = response_message.split("### compiled references")[0]
 
     # Assert
     expected_responses = [
@@ -91,10 +89,8 @@ def test_chat_with_online_webpage_content(chat_client):
     # Act
     q = "/online how many firefighters were involved in the great chicago fire and which year did it take place?"
     encoded_q = quote(q, safe="")
-    response = chat_client.get(f"/api/chat?q={encoded_q}&stream=true")
+    response = chat_client.get(f"/api/chat?q={encoded_q}")
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
-
-    response_message = response_message.split("### compiled references")[0]
 
     # Assert
     expected_responses = ["185", "1871", "horse"]
@@ -144,7 +140,7 @@ def test_answer_from_currently_retrieved_content(chat_client, default_user2: Kho
 
     # Act
     response = chat_client.get(f'/api/chat?q="Where was Xi Li born?"')
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
 
     # Assert
     assert response.status_code == 200
@@ -168,7 +164,7 @@ def test_answer_from_chat_history_and_previously_retrieved_content(chat_client_n
 
     # Act
     response = chat_client_no_background.get(f'/api/chat?q="Where was I born?"')
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
 
     # Assert
     assert response.status_code == 200
@@ -191,7 +187,7 @@ def test_answer_from_chat_history_and_currently_retrieved_content(chat_client, d
 
     # Act
     response = chat_client.get(f'/api/chat?q="Where was I born?"')
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
 
     # Assert
     assert response.status_code == 200
@@ -215,8 +211,8 @@ def test_no_answer_in_chat_history_or_retrieved_content(chat_client, default_use
     create_conversation(message_list, default_user2)
 
     # Act
-    response = chat_client.get(f'/api/chat?q="Where was I born?"&stream=true')
+    response = chat_client.get(f'/api/chat?q="Where was I born?"')
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
 
     # Assert
     expected_responses = [
@@ -226,6 +222,7 @@ def test_no_answer_in_chat_history_or_retrieved_content(chat_client, default_use
         "do not have",
         "don't have",
         "where were you born?",
+        "where you were born?",
     ]
 
     assert response.status_code == 200
@@ -280,8 +277,8 @@ def test_answer_not_known_using_notes_command(chat_client_no_background, default
     create_conversation(message_list, default_user2)
 
     # Act
-    response = chat_client_no_background.get(f"/api/chat?q={query}&stream=true")
+    response = chat_client_no_background.get(f"/api/chat?q={query}")
-    response_message = response.content.decode("utf-8")
+    response_message = response.json()["response"]
 
     # Assert
     assert response.status_code == 200
@@ -527,8 +524,8 @@ def test_answer_general_question_not_in_chat_history_or_retrieved_content(chat_c
     create_conversation(message_list, default_user2)
 
     # Act
-    response = chat_client.get(f'/api/chat?q="Write a haiku about unit testing. Do not say anything else."&stream=true')
+    response = chat_client.get(f'/api/chat?q="Write a haiku about unit testing. Do not say anything else.')
-    response_message = response.content.decode("utf-8").split("### compiled references")[0]
+    response_message = response.json()["response"]
 
     # Assert
     expected_responses = ["test", "Test"]
@@ -544,9 +541,8 @@ def test_answer_general_question_not_in_chat_history_or_retrieved_content(chat_c
 @pytest.mark.chatquality
 def test_ask_for_clarification_if_not_enough_context_in_question(chat_client_no_background):
     # Act
-
+    response = chat_client_no_background.get(f'/api/chat?q="What is the name of Namitas older son?"')
-    response = chat_client_no_background.get(f'/api/chat?q="What is the name of Namitas older son?"&stream=true')
+    response_message = response.json()["response"].lower()
-    response_message = response.content.decode("utf-8").split("### compiled references")[0].lower()
 
     # Assert
     expected_responses = [
@@ -658,8 +654,8 @@ def test_answer_in_chat_history_by_conversation_id_with_agent(
 def test_answer_requires_multiple_independent_searches(chat_client):
     "Chat director should be able to answer by doing multiple independent searches for required information"
     # Act
-    response = chat_client.get(f'/api/chat?q="Is Xi older than Namita? Just the older persons full name"&stream=true')
+    response = chat_client.get(f'/api/chat?q="Is Xi older than Namita? Just the older persons full name"')
-    response_message = response.content.decode("utf-8").split("### compiled references")[0].lower()
+    response_message = response.json()["response"].lower()
 
     # Assert
     expected_responses = ["he is older than namita", "xi is older than namita", "xi li is older than namita"]
@@ -683,8 +679,8 @@ def test_answer_using_file_filter(chat_client):
         'Is Xi older than Namita? Just say the older persons full name. file:"Namita.markdown" file:"Xi Li.markdown"'
     )
 
-    response = chat_client.get(f"/api/chat?q={query}&stream=true")
+    response = chat_client.get(f"/api/chat?q={query}")
-    response_message = response.content.decode("utf-8").split("### compiled references")[0].lower()
+    response_message = response.json()["response"].lower()
 
     # Assert
     expected_responses = ["he is older than namita", "xi is older than namita", "xi li is older than namita"]

versions.json

@@ -53,5 +53,6 @@
 	"1.13.0": "0.15.0",
 	"1.14.0": "0.15.0",
 	"1.15.0": "0.15.0",
-	"1.16.0": "0.15.0"
+	"1.16.0": "0.15.0",
+	"1.17.0": "0.15.0"
 }