Improve data source, output mode selection

- Set output mode to single string. Specify output schema in prompt
  - Both thesee should encourage model to select only 1 output mode
    instead of encouraging it in prompt too many times
  - Output schema should also improve schema following in general
- Standardize variable, func name of io selector for readability
- Fix chat actors to test the io selector chat actor
- Make chat actor return sources, output separately for better
  disambiguation, at least during tests, for now

tests/test_offline_chat_director.py

@@ -7,7 +7,7 @@ from freezegun import freeze_time
 from khoj.database.models import Agent, Entry, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.utils import message_to_log
-from khoj.routers.helpers import aget_relevant_tools_to_execute
+from khoj.routers.helpers import aget_data_sources_and_output_format
 from tests.helpers import ConversationFactory
 
 SKIP_TESTS = True
@@ -735,7 +735,7 @@ async def test_get_correct_tools_online(client_offline_chat):
     user_query = "What's the weather in Patagonia this week?"
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, {}, is_task=False)
+    tools = await aget_data_sources_and_output_format(user_query, {}, is_task=False)
 
     # Assert
     tools = [tool.value for tool in tools]
@@ -750,7 +750,7 @@ async def test_get_correct_tools_notes(client_offline_chat):
     user_query = "Where did I go for my first battleship training?"
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, {}, is_task=False)
+    tools = await aget_data_sources_and_output_format(user_query, {}, is_task=False)
 
     # Assert
     tools = [tool.value for tool in tools]
@@ -765,7 +765,7 @@ async def test_get_correct_tools_online_or_general_and_notes(client_offline_chat
     user_query = "What's the highest point in Patagonia and have I been there?"
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, {}, is_task=False)
+    tools = await aget_data_sources_and_output_format(user_query, {}, is_task=False)
 
     # Assert
     tools = [tool.value for tool in tools]
@@ -782,7 +782,7 @@ async def test_get_correct_tools_general(client_offline_chat):
     user_query = "How many noble gases are there?"
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, {}, is_task=False)
+    tools = await aget_data_sources_and_output_format(user_query, {}, is_task=False)
 
     # Assert
     tools = [tool.value for tool in tools]
@@ -806,7 +806,7 @@ async def test_get_correct_tools_with_chat_history(client_offline_chat, default_
     chat_history = create_conversation(chat_log, default_user2)
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, chat_history, is_task=False)
+    tools = await aget_data_sources_and_output_format(user_query, chat_history, is_task=False)
 
     # Assert
     tools = [tool.value for tool in tools]

tests/test_openai_chat_actors.py

@@ -1,4 +1,3 @@
-import os
 from datetime import datetime
 
 import freezegun
@@ -8,7 +7,7 @@ from freezegun import freeze_time
 from khoj.processor.conversation.openai.gpt import converse, extract_questions
 from khoj.processor.conversation.utils import message_to_log
 from khoj.routers.helpers import (
-    aget_relevant_tools_to_execute,
+    aget_data_sources_and_output_format,
     generate_online_subqueries,
     infer_webpage_urls,
     schedule_query,
@@ -529,19 +528,36 @@ async def test_websearch_khoj_website_for_info_about_khoj(chat_client, default_u
 @pytest.mark.parametrize(
     "user_query, expected_conversation_commands",
     [
-        ("Where did I learn to swim?", [ConversationCommand.Notes]),
-        ("Where is the nearest hospital?", [ConversationCommand.Online]),
-        ("Summarize the wikipedia page on the history of the internet", [ConversationCommand.Webpage]),
+        (
+            "Where did I learn to swim?",
+            {"sources": [ConversationCommand.Notes], "output": ConversationCommand.Text},
+        ),
+        (
+            "Where is the nearest hospital?",
+            {"sources": [ConversationCommand.Online], "output": ConversationCommand.Text},
+        ),
+        (
+            "Summarize the wikipedia page on the history of the internet",
+            {"sources": [ConversationCommand.Webpage], "output": ConversationCommand.Text},
+        ),
+        (
+            "Make a painting incorporating my past diving experiences",
+            {"sources": [ConversationCommand.Notes], "output": ConversationCommand.Image},
+        ),
+        (
+            "Create a chart of the weather over the next 7 days in Timbuktu",
+            {"sources": [ConversationCommand.Online, ConversationCommand.Code], "output": ConversationCommand.Text},
+        ),
     ],
 )
 async def test_select_data_sources_actor_chooses_to_search_notes(
-    chat_client, user_query, expected_conversation_commands
+    chat_client, user_query, expected_conversation_commands, default_user2
 ):
     # Act
-    conversation_commands = await aget_relevant_tools_to_execute(user_query, {}, False, False)
+    selected_conversation_commands = await aget_data_sources_and_output_format(user_query, {}, False, default_user2)
 
     # Assert
-    assert set(expected_conversation_commands) == set(conversation_commands)
+    assert expected_conversation_commands == selected_conversation_commands
 
 
 # ----------------------------------------------------------------------------------------------------

tests/test_openai_chat_director.py

@@ -8,7 +8,7 @@ from freezegun import freeze_time
 from khoj.database.models import Agent, Entry, KhojUser, LocalPdfConfig
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.utils import message_to_log
-from khoj.routers.helpers import aget_relevant_tools_to_execute
+from khoj.routers.helpers import aget_data_sources_and_output_format
 from tests.helpers import ConversationFactory
 
 # Initialize variables for tests
@@ -719,7 +719,7 @@ async def test_get_correct_tools_online(chat_client):
     user_query = "What's the weather in Patagonia this week?"
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, {}, False, False)
+    tools = await aget_data_sources_and_output_format(user_query, {}, False, False)
 
     # Assert
     tools = [tool.value for tool in tools]
@@ -734,7 +734,7 @@ async def test_get_correct_tools_notes(chat_client):
     user_query = "Where did I go for my first battleship training?"
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, {}, False, False)
+    tools = await aget_data_sources_and_output_format(user_query, {}, False, False)
 
     # Assert
     tools = [tool.value for tool in tools]
@@ -749,7 +749,7 @@ async def test_get_correct_tools_online_or_general_and_notes(chat_client):
     user_query = "What's the highest point in Patagonia and have I been there?"
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, {}, False, False)
+    tools = await aget_data_sources_and_output_format(user_query, {}, False, False)
 
     # Assert
     tools = [tool.value for tool in tools]
@@ -766,7 +766,7 @@ async def test_get_correct_tools_general(chat_client):
     user_query = "How many noble gases are there?"
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, {}, False, False)
+    tools = await aget_data_sources_and_output_format(user_query, {}, False, False)
 
     # Assert
     tools = [tool.value for tool in tools]
@@ -790,7 +790,7 @@ async def test_get_correct_tools_with_chat_history(chat_client):
     chat_history = generate_history(chat_log)
 
     # Act
-    tools = await aget_relevant_tools_to_execute(user_query, chat_history, False, False)
+    tools = await aget_data_sources_and_output_format(user_query, chat_history, False, False)
 
     # Assert
     tools = [tool.value for tool in tools]