Update the iterative chat director prompt to generalize across chat models

These prompts work across o1 and standard openai model. Works with
anthropic and google models as well

src/khoj/processor/conversation/prompts.py

@@ -484,45 +484,49 @@ Khoj:
 
 plan_function_execution = PromptTemplate.from_template(
     """
-You are Khoj, a smart, methodical researcher. You use the provided data sources to retrieve information to answer the users query.
-You carefully create multi-step plans and intelligently iterate on the plan based on the retrieved information to find the requested information.
+You are Khoj, a smart, methodical researcher agent. Use the provided tool AIs to answer my query.
+Create a multi-step plan and intelligently iterate on the plan based on the retrieved information to find the requested information.
 {personality_context}
-- Use the data sources provided below, one at a time, if you need to find more information. Their output will be shown to you in the next iteration.
-- You are allowed upto {max_iterations} iterations to use these data sources to answer the user's question
-- If you have enough information to answer the question, then exit execution by returning an empty response. E.g., {{}}
-- Ensure the query contains enough context to retrieve relevant information from the data sources.
-- Break down the problem into smaller steps. Some examples are provided below assuming you have access to the notes and online data sources:
-  - If the user asks for the population of their hometown
-    1. Try look up their hometown in their notes
+
+# Instructions
+- Ask detailed queries to the tool AIs provided below, one at a time, to discover required information or run calculations. Their response will be shown to you in the next iteration.
+- Break down your discovery and research process into independent, self-contained steps that can be executed sequentially.
+- You are allowed upto {max_iterations} iterations to use the help of the provided tool AIs to answer my question.
+- When you have the required information return an empty JSON object. E.g., {{}}
+
+# Examples
+Assuming you can search my notes and the internet.
+- When I ask for the population of my hometown
+  1. Try look up my hometown in my notes
   2. Only then try find the population of the city online.
-  - If the user asks for their computer's specs
-    1. Try find the computer model in their notes
-    2. Now look up their computer models spec online
-  - If the user asks what clothes to carry for their upcoming trip
-    1. Find the itinerary of their upcoming trip in their notes
+- When I ask for my computer's specs
+  1. Try find my computer model in my notes
+  2. Now look up my computer model's spec online
+- When I ask what clothes to carry for my upcoming trip
+  1. Find the itinerary of my upcoming trip in my notes
   2. Next find the weather forecast at the destination online
-    3. Then find if they mention what clothes they own in their notes
+  3. Then find if I mentioned what clothes I own in my notes
 
-Background Context:
+# Background Context
 - Current Date: {day_of_week}, {current_date}
-- User's Location: {location}
-- {username}
+- My Location: {location}
+- My {username}
 
-Which of the data sources listed below you would use to answer the user's question? You **only** have access to the following data sources:
+# Available Tool AIs
+Which of the tool AIs listed below would you use to answer my question? You **only** have access to the following tool AIs:
 
 {tools}
 
-Provide the data source and associated query in a JSON object. Do not say anything else.
-
-Previous Iterations:
+# Previous Iterations
 {previous_iterations}
 
-Response format:
-{{"data_source": "<tool_name>", "query": "<your_new_query>"}}
-
-Chat History:
+# Chat History:
 {chat_history}
 
+Return the next tool AI to use and the query to ask it. Your response should always be a valid JSON object. Do not say anything else.
+Response format:
+{{"scratchpad": "<your_scratchpad_to_reason_about_which_tool_to_use>", "tool": "<name_of_tool_ai>", "query": "<your_query_for_the_tool_ai>"}}
+
 User: {query}
 Khoj:
 """.strip()
@@ -530,11 +534,10 @@ Khoj:
 
 previous_iteration = PromptTemplate.from_template(
     """
-# Iteration {index}:
-# ---
-- data_source: {data_source}
+## Iteration {index}:
+- tool: {tool}
 - query: {query}
-- summary: {summary}
+- result: {result}
 """
 )
 

src/khoj/processor/conversation/utils.py

@@ -82,14 +82,14 @@ class ThreadedGenerator:
 class InformationCollectionIteration:
     def __init__(
         self,
-        data_source: str,
+        tool: str,
         query: str,
         context: Dict[str, Dict] = None,
         onlineContext: dict = None,
         codeContext: dict = None,
         summarizedResult: str = None,
     ):
-        self.data_source = data_source
+        self.tool = tool
         self.query = query
         self.context = context
         self.onlineContext = onlineContext
@@ -103,9 +103,9 @@ def construct_iteration_history(
     previous_iterations_history = ""
     for idx, iteration in enumerate(previous_iterations):
         iteration_data = previous_iteration_prompt.format(
+            tool=iteration.tool,
             query=iteration.query,
-            data_source=iteration.data_source,
-            summary=iteration.summarizedResult,
+            result=iteration.summarizedResult,
             index=idx + 1,
         )
 

src/khoj/routers/research.py

@@ -100,20 +100,21 @@ async def apick_next_tool(
         response = response.strip()
         response = remove_json_codeblock(response)
         response = json.loads(response)
-        suggested_data_source = response.get("data_source", None)
-        suggested_query = response.get("query", None)
+        selected_tool = response.get("tool", None)
+        generated_query = response.get("query", None)
+        scratchpad = response.get("scratchpad", None)
 
         logger.info(f"Response for determining relevant tools: {response}")
 
         return InformationCollectionIteration(
-            data_source=suggested_data_source,
-            query=suggested_query,
+            tool=selected_tool,
+            query=generated_query,
         )
 
     except Exception as e:
         logger.error(f"Invalid response for determining relevant tools: {response}. {e}", exc_info=True)
         return InformationCollectionIteration(
-            data_source=None,
+            tool=None,
             query=None,
         )
 
@@ -155,7 +156,7 @@ async def execute_information_collection(
             previous_iterations_history,
             MAX_ITERATIONS,
         )
-        if this_iteration.data_source == ConversationCommand.Notes:
+        if this_iteration.tool == ConversationCommand.Notes:
             ## Extract Document References
             compiled_references, inferred_queries, defiltered_query = [], [], None
             async for result in extract_references_and_questions(
@@ -190,7 +191,7 @@ async def execute_information_collection(
                 # TODO Get correct type for compiled across research notes extraction
                 logger.error(f"Error extracting references: {e}", exc_info=True)
 
-        elif this_iteration.data_source == ConversationCommand.Online:
+        elif this_iteration.tool == ConversationCommand.Online:
             async for result in search_online(
                 this_iteration.query,
                 conversation_history,
@@ -209,7 +210,7 @@ async def execute_information_collection(
                     online_results: Dict[str, Dict] = result  # type: ignore
                     this_iteration.onlineContext = online_results
 
-        elif this_iteration.data_source == ConversationCommand.Webpage:
+        elif this_iteration.tool == ConversationCommand.Webpage:
             try:
                 async for result in read_webpages(
                     this_iteration.query,
@@ -239,7 +240,7 @@ async def execute_information_collection(
             except Exception as e:
                 logger.error(f"Error reading webpages: {e}", exc_info=True)
 
-        elif this_iteration.data_source == ConversationCommand.Code:
+        elif this_iteration.tool == ConversationCommand.Code:
             try:
                 async for result in run_code(
                     this_iteration.query,

src/khoj/utils/helpers.py

@@ -350,8 +350,8 @@ tool_descriptions_for_llm = {
 
 function_calling_description_for_llm = {
     ConversationCommand.Notes: "To search the user's personal knowledge base. Especially helpful if the question expects context from the user's notes or documents.",
-    ConversationCommand.Online: "To search for the latest, up-to-date information from the internet.",
-    ConversationCommand.Webpage: "To use if the user has directly provided the webpage urls or you are certain of the webpage urls to read.",
+    ConversationCommand.Online: "To search the internet for the latest, up-to-date information.",
+    ConversationCommand.Webpage: "To read a webpage url for detailed information from the internet.",
     ConversationCommand.Code: "To run Python code in a Pyodide sandbox with no network access. Helpful when need to parse information, run complex calculations, create documents and charts for user. Matplotlib, bs4, pandas, numpy, etc. are available.",
 }