khoj 1.16.1.dev15__py3-none-any.whl → 1.16.1.dev47__py3-none-any.whl

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}"

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()

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)

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}"

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.
-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.
 """.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.
-    """.strip()
+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.
+""".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 will be provided past questions(Q) and answers(A) for context.
+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.
+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 will be provided past questions(Q) and answers(A) for context.
+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(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}
 

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
 
 

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(

khoj/routers/api.py

@@ -6,7 +6,6 @@ import os
 import threading
 import time
 import uuid
-from random import random
 from typing import Any, Callable, List, Optional, Union
 
 import cron_descriptor
@@ -37,6 +36,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,
@@ -298,11 +298,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
@@ -313,11 +315,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):
@@ -335,6 +338,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,
@@ -372,7 +376,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(
@@ -391,7 +396,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)