From 3f70d2f68571294c5557d667f96616a7ddf89b2c Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Mon, 18 Nov 2024 09:34:49 -0800
Subject: [PATCH] Add more graceful exception handling when tool selection
 doesn't work

---
 src/khoj/processor/conversation/prompts.py |  3 ++-
 src/khoj/routers/helpers.py                | 19 +++++++++----------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/khoj/processor/conversation/prompts.py b/src/khoj/processor/conversation/prompts.py
index 8d31ffcd..a1b3033a 100644
--- a/src/khoj/processor/conversation/prompts.py
+++ b/src/khoj/processor/conversation/prompts.py
@@ -634,6 +634,7 @@ You are Khoj, an extremely smart and helpful search assistant.
 - You can use the data sources listed below to collect more relevant information
 - You can select certain types of output to respond to the user's question. Select just one output type to answer the user's question
 - You can use any combination of these data sources and output types to answer the user's question
+- You can only select one output type to answer the user's question
 
 Which of the tools listed below you would use to answer the user's question? You **only** have access to the following:
 
@@ -685,7 +686,7 @@ AI: Learning to play the guitar is a great hobby. It can be a lot of fun and a g
 Q: Draw a painting of a guitar.
 Khoj: {{"source": ["general"], "output": ["image"]}}
 
-Now it's your turn to pick the data sources you would like to use to answer the user's question. Provide the data sources as a list of strings in a JSON object. Do not say anything else.
+Now it's your turn to pick the sources and output to answer the user's query. Respond with a JSON object, including both `source` and `output`. The values should be a list of strings. Do not say anything else.
 
 Chat History:
 {chat_history}
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index f653f279..109deca2 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -402,15 +402,14 @@ async def aget_relevant_tools_to_execute(
     try:
         response = clean_json(response)
         response = json.loads(response)
-        input_tools = [q.strip() for q in response["source"] if q.strip()]
-        if not isinstance(input_tools, list) or not input_tools or len(input_tools) == 0:
-            logger.error(f"Invalid response for determining relevant tools: {input_tools}")
-            return tool_options
 
-        output_modes = [q.strip() for q in response["output"] if q.strip()]
-        if not isinstance(output_modes, list) or not output_modes or len(output_modes) == 0:
-            logger.error(f"Invalid response for determining relevant output modes: {output_modes}")
-            return mode_options
+        input_tools = [q.strip() for q in response.get("source", []) if q.strip()]
+        output_modes = [q.strip() for q in response.get("output", ["text"]) if q.strip()]  # Default to text output
+
+        if not isinstance(input_tools, list) or not input_tools or len(input_tools) == 0:
+            raise ValueError(
+                f"Invalid response for determining relevant tools: {input_tools}. Raw Response: {response}"
+            )
 
         final_response = [] if not is_task else [ConversationCommand.AutomatedTask]
         for llm_suggested_tool in input_tools:
@@ -434,8 +433,8 @@ async def aget_relevant_tools_to_execute(
                 final_response = [ConversationCommand.Default, ConversationCommand.Text]
             else:
                 final_response = [ConversationCommand.General, ConversationCommand.Text]
-    except Exception:
-        logger.error(f"Invalid response for determining relevant tools: {response}")
+    except Exception as e:
+        logger.error(f"Invalid response for determining relevant tools: {response}. Error: {e}", exc_info=True)
         if len(agent_tools) == 0:
             final_response = [ConversationCommand.Default, ConversationCommand.Text]
         else: