khoj 1.16.1.dev25__py3-none-any.whl → 1.17.1.dev216__py3-none-any.whl

khoj/routers/api_chat.py

@@ -1,41 +1,36 @@
+import asyncio
 import json
 import logging
-import math
+import time
 from datetime import datetime
+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 (
     ConversationAdapters,
-    DataStoreAdapters,
     EntryAdapters,
     FileObjectAdapters,
     PublicConversationAdapters,
     aget_user_name,
 )
 from khoj.database.models import KhojUser
-from khoj.processor.conversation.prompts import (
-    help_message,
-    no_entries_found,
-    no_notes_found,
-)
+from khoj.processor.conversation.prompts import help_message, no_entries_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,
@@ -58,7 +53,7 @@ from khoj.utils.helpers import (
     get_device,
     is_none_or_empty,
 )
-from khoj.utils.rawconfig import FilterRequest, LocationData
+from khoj.utils.rawconfig import FileFilterRequest, FilesFilterRequest, LocationData
 
 # Initialize Router
 logger = logging.getLogger(__name__)
@@ -92,68 +87,36 @@ def get_file_filter(request: Request, conversation_id: str) -> Response:
     return Response(content=json.dumps(file_filters), media_type="application/json", status_code=200)
 
 
-class FactCheckerStoreDataFormat(BaseModel):
-    factToVerify: str
-    response: str
-    references: Any
-    childReferences: List[Any]
-    runId: str
-    modelUsed: Dict[str, Any]
-
-
-class FactCheckerStoreData(BaseModel):
-    runId: str
-    storeData: FactCheckerStoreDataFormat
-
-
-@api_chat.post("/store/factchecker", response_class=Response)
+@api_chat.delete("/conversation/file-filters/bulk", response_class=Response)
 @requires(["authenticated"])
-async def store_factchecker(request: Request, common: CommonQueryParams, data: FactCheckerStoreData):
-    user = request.user.object
-
-    update_telemetry_state(
-        request=request,
-        telemetry_type="api",
-        api="store_factchecker",
-        **common.__dict__,
-    )
-    fact_checker_key = f"factchecker_{data.runId}"
-    await DataStoreAdapters.astore_data(data.storeData.model_dump_json(), fact_checker_key, user, private=False)
-    return Response(content=json.dumps({"status": "ok"}), media_type="application/json", status_code=200)
-
-
-@api_chat.get("/store/factchecker", response_class=Response)
-async def get_factchecker(request: Request, common: CommonQueryParams, runId: str):
-    update_telemetry_state(
-        request=request,
-        telemetry_type="api",
-        api="read_factchecker",
-        **common.__dict__,
-    )
+def remove_files_filter(request: Request, filter: FilesFilterRequest) -> Response:
+    conversation_id = int(filter.conversation_id)
+    files_filter = filter.filenames
+    file_filters = ConversationAdapters.remove_files_from_filter(request.user.object, conversation_id, files_filter)
+    return Response(content=json.dumps(file_filters), media_type="application/json", status_code=200)
 
-    fact_checker_key = f"factchecker_{runId}"
 
-    data = await DataStoreAdapters.aretrieve_public_data(fact_checker_key)
-    if data is None:
-        return Response(status_code=404)
-    return Response(content=json.dumps(data.value), media_type="application/json", status_code=200)
+@api_chat.post("/conversation/file-filters/bulk", response_class=Response)
+@requires(["authenticated"])
+def add_files_filter(request: Request, filter: FilesFilterRequest):
+    try:
+        conversation_id = int(filter.conversation_id)
+        files_filter = filter.filenames
+        file_filters = ConversationAdapters.add_files_to_filter(request.user.object, conversation_id, files_filter)
+        return Response(content=json.dumps(file_filters), media_type="application/json", status_code=200)
+    except Exception as e:
+        logger.error(f"Error adding file filter {filter.filename}: {e}", exc_info=True)
+        raise HTTPException(status_code=422, detail=str(e))
 
 
 @api_chat.post("/conversation/file-filters", response_class=Response)
 @requires(["authenticated"])
-def add_file_filter(request: Request, filter: FilterRequest):
+def add_file_filter(request: Request, filter: FileFilterRequest):
     try:
-        conversation = ConversationAdapters.get_conversation_by_user(
-            request.user.object, conversation_id=int(filter.conversation_id)
-        )
-        file_list = EntryAdapters.get_all_filenames_by_source(request.user.object, "computer")
-        if filter.filename in file_list and filter.filename not in conversation.file_filters:
-            conversation.file_filters.append(filter.filename)
-            conversation.save()
-        # remove files from conversation.file_filters that are not in file_list
-        conversation.file_filters = [file for file in conversation.file_filters if file in file_list]
-        conversation.save()
-        return Response(content=json.dumps(conversation.file_filters), media_type="application/json", status_code=200)
+        conversation_id = int(filter.conversation_id)
+        files_filter = [filter.filename]
+        file_filters = ConversationAdapters.add_files_to_filter(request.user.object, conversation_id, files_filter)
+        return Response(content=json.dumps(file_filters), media_type="application/json", status_code=200)
     except Exception as e:
         logger.error(f"Error adding file filter {filter.filename}: {e}", exc_info=True)
         raise HTTPException(status_code=422, detail=str(e))
@@ -161,18 +124,11 @@ def add_file_filter(request: Request, filter: FilterRequest):
 
 @api_chat.delete("/conversation/file-filters", response_class=Response)
 @requires(["authenticated"])
-def remove_file_filter(request: Request, filter: FilterRequest) -> Response:
-    conversation = ConversationAdapters.get_conversation_by_user(
-        request.user.object, conversation_id=int(filter.conversation_id)
-    )
-    if filter.filename in conversation.file_filters:
-        conversation.file_filters.remove(filter.filename)
-    conversation.save()
-    # remove files from conversation.file_filters that are not in file_list
-    file_list = EntryAdapters.get_all_filenames_by_source(request.user.object, "computer")
-    conversation.file_filters = [file for file in conversation.file_filters if file in file_list]
-    conversation.save()
-    return Response(content=json.dumps(conversation.file_filters), media_type="application/json", status_code=200)
+def remove_file_filter(request: Request, filter: FileFilterRequest) -> Response:
+    conversation_id = int(filter.conversation_id)
+    files_filter = [filter.filename]
+    file_filters = ConversationAdapters.remove_files_from_filter(request.user.object, conversation_id, files_filter)
+    return Response(content=json.dumps(file_filters), media_type="application/json", status_code=200)
 
 
 class FeedbackData(BaseModel):
@@ -195,10 +151,10 @@ async def text_to_speech(
     common: CommonQueryParams,
     text: str,
     rate_limiter_per_minute=Depends(
-        ApiUserRateLimiter(requests=5, subscribed_requests=20, window=60, slug="chat_minute")
+        ApiUserRateLimiter(requests=20, subscribed_requests=20, window=60, slug="chat_minute")
     ),
     rate_limiter_per_day=Depends(
-        ApiUserRateLimiter(requests=5, subscribed_requests=300, window=60 * 60 * 24, slug="chat_day")
+        ApiUserRateLimiter(requests=50, subscribed_requests=300, window=60 * 60 * 24, slug="chat_day")
     ),
 ) -> Response:
     voice_model = await ConversationAdapters.aget_voice_model_config(request.user.object)
@@ -252,6 +208,9 @@ def chat_history(
             "name": conversation.agent.name,
             "avatar": conversation.agent.avatar,
             "isCreator": conversation.agent.creator == user,
+            "color": conversation.agent.style_color,
+            "icon": conversation.agent.style_icon,
+            "persona": conversation.agent.personality,
         }
 
     meta_log = conversation.conversation_log
@@ -306,13 +265,21 @@ def get_shared_chat(
             "name": conversation.agent.name,
             "avatar": conversation.agent.avatar,
             "isCreator": conversation.agent.creator == user,
+            "color": conversation.agent.style_color,
+            "icon": conversation.agent.style_icon,
+            "persona": conversation.agent.personality,
         }
 
     meta_log = conversation.conversation_log
+    scrubbed_title = conversation.title if conversation.title else conversation.slug
+
+    if scrubbed_title:
+        scrubbed_title = scrubbed_title.replace("-", " ")
+
     meta_log.update(
         {
             "conversation_id": conversation.id,
-            "slug": conversation.title if conversation.title else conversation.slug,
+            "slug": scrubbed_title,
             "agent": agent_metadata,
         }
     )
@@ -328,7 +295,7 @@ def get_shared_chat(
     update_telemetry_state(
         request=request,
         telemetry_type="api",
-        api="public_conversation_history",
+        api="chat_history",
         **common.__dict__,
     )
 
@@ -370,7 +337,7 @@ def fork_public_conversation(
     public_conversation = PublicConversationAdapters.get_public_conversation_by_slug(public_conversation_slug)
 
     # Duplicate Public Conversation to User's Private Conversation
-    ConversationAdapters.create_conversation_from_public_conversation(
+    new_conversation = ConversationAdapters.create_conversation_from_public_conversation(
         user, public_conversation, request.user.client_app
     )
 
@@ -386,7 +353,16 @@ def fork_public_conversation(
 
     redirect_uri = str(request.app.url_path_for("chat_page"))
 
-    return Response(status_code=200, content=json.dumps({"status": "ok", "next_url": redirect_uri}))
+    return Response(
+        status_code=200,
+        content=json.dumps(
+            {
+                "status": "ok",
+                "next_url": redirect_uri,
+                "conversation_id": new_conversation.id,
+            }
+        ),
+    )
 
 
 @api_chat.post("/share")
@@ -427,15 +403,30 @@ def duplicate_chat_history_public_conversation(
 def chat_sessions(
     request: Request,
     common: CommonQueryParams,
+    recent: Optional[bool] = False,
 ):
     user = request.user.object
 
     # Load Conversation Sessions
-    sessions = ConversationAdapters.get_conversation_sessions(user, request.user.client_app).values_list(
-        "id", "slug", "title"
+    conversations = ConversationAdapters.get_conversation_sessions(user, request.user.client_app)
+    if recent:
+        conversations = conversations[:8]
+
+    sessions = conversations.values_list(
+        "id", "slug", "title", "agent__slug", "agent__name", "agent__avatar", "created_at", "updated_at"
     )
 
-    session_values = [{"conversation_id": session[0], "slug": session[2] or session[1]} for session in sessions]
+    session_values = [
+        {
+            "conversation_id": session[0],
+            "slug": session[2] or session[1],
+            "agent_name": session[4],
+            "agent_avatar": session[5],
+            "created": session[6].strftime("%Y-%m-%d %H:%M:%S"),
+            "updated": session[7].strftime("%Y-%m-%d %H:%M:%S"),
+        }
+        for session in sessions
+    ]
 
     update_telemetry_state(
         request=request,
@@ -477,7 +468,6 @@ async def create_chat_session(
 
 
 @api_chat.get("/options", response_class=Response)
-@requires(["authenticated"])
 async def chat_options(
     request: Request,
     common: CommonQueryParams,
@@ -526,141 +516,140 @@ async def set_conversation_title(
     )
 
 
-@api_chat.websocket("/ws")
-async def websocket_endpoint(
-    websocket: WebSocket,
-    conversation_id: int,
+@api_chat.get("")
+async def chat(
+    request: Request,
+    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_event(event_type: ChatEvent, data: str | dict):
+            nonlocal connection_alive, ttft
+            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,
+            )
 
-    async def send_status_update(message: str):
-        nonlocal connection_alive
-        if not connection_alive:
-            return
+        conversation_commands = [get_conversation_command(query=q, any_references=True)]
 
-        status_packet = {
-            "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):
-        nonlocal connection_alive
-        if not connection_alive:
-            return
-        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:
+        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 = {
-            "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
-    )
+        await is_ready_to_chat(user)
 
-    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
+        user_name = await aget_user_name(user)
+        location = None
+        if city or region or country:
+            location = LocationData(city=city, region=region, country=country)
 
         if is_query_empty(q):
-            await send_message("start_llm_response")
-            await send_message(
-                "It seems like your query is incomplete. Could you please provide more details or specify what you need help with?"
-            )
-            await send_message("end_llm_response")
-            continue
+            async for result in send_llm_response("Please ask your query to get started."):
+                yield result
+            return
 
         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}")
 
         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 (
@@ -676,28 +665,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)
-                        continue
+                        async for result in send_llm_response(response_log):
+                            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,
@@ -705,16 +703,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(
-                request=websocket,
-                telemetry_type="api",
-                api="chat",
-                metadata={"conversation_command": conversation_commands[0].value},
-            )
-            continue
+            return
 
         custom_filters = []
         if conversation_commands == [ConversationCommand.Help]:
@@ -724,8 +716,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)
-                continue
+                async for result in send_llm_response(formatted_help):
+                    yield result
+                return
             # Adding specification to search online specifically on khoj.dev pages.
             custom_filters.append("site:khoj.dev")
             conversation_commands.append(ConversationCommand.Online)
@@ -733,14 +726,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(
-                    f"Unable to create automation. Ensure the automation doesn't already exist."
-                )
-                continue
+                error_message = f"Unable to create automation. Ensure the automation doesn't already exist."
+                async for result in send_llm_response(error_message):
+                    yield result
+                return
 
             llm_response = construct_automation_created_message(automation, crontime, query_to_run, subject)
             await sync_to_async(save_to_conversation_log)(
@@ -750,57 +743,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(
-                client=websocket.user.client_app,
-                user_agent=websocket.headers.get("user-agent"),
-                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
+            async for result in send_llm_response(llm_response):
+                yield result
+            return
 
-        compiled_references, inferred_queries, defiltered_query = await extract_references_and_questions(
-            websocket, meta_log, q, 7, 0.18, conversation_id, conversation_commands, location, send_status_update
-        )
+        # Gather Context
+        ## 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()}")
-            continue
+            async for result in send_llm_response(f"{no_entries_found.format()}"):
+                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(
-                    defiltered_query, meta_log, location, send_status_update, custom_filters
-                )
+                async for result in search_online(
+                    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")
-                await send_complete_llm_response(
-                    f"Error searching online: {e}. Attempting to respond without online results"
-                )
-                continue
+                error_message = f"Error searching online: {e}. Attempting to respond without online results"
+                logger.warning(error_message)
+                async for result in send_llm_response(error_message):
+                    yield result
+                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):
@@ -810,38 +824,52 @@ async def websocket_endpoint(
 
                     for webpage in direct_web_pages[query]["webpages"]:
                         webpages.append(webpage["link"])
-
-                await send_status_update(f"**📚 Read web pages**: {webpages}")
+                async for result in send_event(ChatEvent.STATUS, 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(
-                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(
+            async for result in 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))
-                continue
+                async for result in send_llm_response(json.dumps(content_obj)):
+                    yield result
+                return
 
             await sync_to_async(save_to_conversation_log)(
                 q,
@@ -851,17 +879,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
-
-            await send_complete_llm_response(json.dumps(content_obj))
-            continue
+            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_status_update(f"**💭 Generating a well-informed response**")
+        ## Generate Text Output
+        async for result in send_event(ChatEvent.STATUS, f"**Generating a well-informed response**"):
+            yield result
         llm_response, chat_metadata = await agenerate_chat_response(
             defiltered_query,
             meta_log,
@@ -871,310 +905,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(
-            request=websocket,
-            telemetry_type="api",
-            api="chat",
-            metadata=chat_metadata,
-        )
+        continue_stream = True
         iterator = AsyncIteratorWrapper(llm_response)
-
-        await send_message("start_llm_response")
-
         async for item in iterator:
             if item is None:
-                break
-            if connection_alive:
-                try:
-                    await send_message(f"{item}")
-                except ConnectionClosedOK:
-                    connection_alive = False
-                    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:
+                async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""):
+                    yield result
+                logger.debug("Finished streaming response")
+                return
+            if not connection_alive or not continue_stream:
+                continue
             try:
-                file_object = await FileObjectAdapters.async_get_file_objects_by_name(user, file_filters[0])
-                if len(file_object) == 0:
-                    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)
+                async for result in send_event(ChatEvent.MESSAGE, f"{item}"):
+                    yield result
             except Exception as e:
-                logger.error(f"Error summarizing file for {user.email}: {e}")
-                llm_response = "Error summarizing file."
-        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)
+                continue_stream = False
+                logger.info(f"User {user} disconnected. Emitting rest of responses to clear thread: {e}")
 
+    ## Stream Text Response
     if stream:
-        return StreamingResponse(llm_response, media_type="text/event-stream", status_code=200)
-
-    iterator = AsyncIteratorWrapper(llm_response)
-
-    # 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)
+        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
+
+        return Response(content=json.dumps(response_obj), media_type="application/json", status_code=200)