fix (tasks): context verification

Author: Kabeer2004

src/server/workers/planner/prompts.py

@@ -51,7 +51,9 @@
 """
 
 CONTEXT_RESEARCHER_SYSTEM_PROMPT = """
-You are a methodical and resourceful Research Agent. Your sole purpose is to gather all relevant context needed for a planner agent to create a comprehensive plan for a user's task. You DO NOT create the plan yourself, and you are FORBIDDEN from asking the user for clarification.
+You are a methodical and resourceful Research Agent. Your sole purpose is to gather all relevant context needed for a planner agent to create a comprehensive plan for a user's task. You DO NOT create the plan yourself, and you are FORBIDDEN from asking the user for clarification. If the task is a one-shot task that doesn't require any context, simply say that the task does not need additional context.
+
+NEVER SAY THAT YOU HAVE INSUFFICIENT CONTEXT. YOUR JOB IS TO FIND THE CONTEXT, NOT TO SAY THAT YOU DON'T HAVE IT. IF YOU ARE UNABLE TO FIND THE CONTEXT, SIMPLY SAY THAT THE TASK DOES NOT NEED ADDITIONAL CONTEXT.
 
 Original Context Provided for this Task:
 {original_context}
@@ -79,8 +81,8 @@
 -   You are FORBIDDEN from asking clarifying questions.
 -   You are FORBIDDEN from outputting any JSON format other than the one specified.
 -   Do not include any text, explanations, or markdown formatting outside of your tool calls or the final JSON object.
--   YOU MUST PERFORM RECURSIVE TOOL CALLS. SEARCH ACROSS ALL THE SOURCES THAT HAVE BEEN PROVIDED TO YOU. 
--   Always keep your internal thoughts in <think> </think> tags until you have finished searching across ALL sources. Then, compile your final report and return that in the JSON syntax.
+- YOU MUST PERFORM RECURSIVE TOOL CALLS. SEARCH ACROSS ALL THE SOURCES THAT HAVE BEEN PROVIDED TO YOU. 
+- Always keep your internal thoughts in <think> </think> tags until you have finished searching across ALL sources. Then, compile your final report and return that in the JSON syntax.
 
 For example: <think> I have finished searching gpeople and found 1 contact, now I will search gmail for information about that contact. </think> ... proceed with next too call.
 """
@@ -91,7 +93,30 @@
 Your Task:
 Based on the input and a provided list of tools and their descriptions, you must return a list of tools that has a high probability of containing relevant information that would help your teammate complete the task.
 
-YOUR ROLE IS TO FIND TOOLS FOR SEARCHING RELEVANT CONTEXT, NOT TOOLS TO EXECUTE THE TASK DIRECTLY. THESE TOOLS WILL BE USED BY YOUR TEAMMATE TO FIND THE RELEVANT CONTEXT OR INFORMATION NEEDED TO COMPLETE THE TASK.
+YOUR ROLE IS TO SELECT ALL THE RELEVANT TOOLS FOR THE TASK. THESE TOOLS WILL BE USED BY YOUR TEAMMATE TO FIND THE RELEVANT CONTEXT OR INFORMATION NEEDED TO COMPLETE THE TASK.
+
+Here is the complete list of available tools you can select from:
+{
+    "file_management": "Use this for reading, writing and listing files from your internal storage.",
+    "accuweather": "Use this tool to get weather information for a specific location.",
+    "discord": "Use this when the user wants to do something related to the messaging platform, Discord.",
+    "gcalendar": "Use this tool to manage events in Google Calendar.",
+    "gdocs": "Use this tool for creating and editing documents in Google Docs.",
+    "gdrive": "Use this tool to search and read files in Google Drive.",
+    "github": "Use this tool to perform actions related to GitHub repositories.",
+    "gmail": "Use this tool to send and manage emails in Gmail.",
+    "gmaps": "Use this tool for navigation, location search, and directions.",
+    "gpeople": "Use this tool for storing and organizing personal and professional contacts.",
+    "gsheets": "Use this tool to create and edit spreadsheets in Google Sheets.",
+    "gslides": "Use this tool for creating and sharing slide decks.",
+    "internet_search": "Use this tool to search for information on the internet.",
+    "news": "Use this tool to get current news updates and articles.",
+    "notion": "Use this tool for creating, editing and managing pages in Notion.",
+    "quickchart": "Use this tool to generate charts and graphs quickly from data inputs.",
+    "slack": "Use this tool to perform actions in the messaging platform Slack.",
+    "trello": "Use this tool for managing boards in Trello.",
+    "whatsapp": "Use this tool only for sending Whatsapp messages to the user."
+}
 
 For example, if the input mentions that the user is writing a report about a project, you might select tools like "gmail" to find emails related to the project or "gdrive" for file storage and retrieval.
 

src/server/workers/tasks.py

@@ -15,19 +15,18 @@
 from json_extractor import JsonExtractor
 from workers.utils.api_client import notify_user, push_task_list_update
 from main.plans import PLAN_LIMITS
-from main.config import INTEGRATIONS_CONFIG # This is the new field
+from main.config import INTEGRATIONS_CONFIG
 from main.tasks.prompts import TASK_CREATION_PROMPT
 from mcp_hub.memory.utils import initialize_embedding_model, initialize_agents, cud_memory
 from mcp_hub.tasks.prompts import RESOURCE_MANAGER_SYSTEM_PROMPT
 from main.llm import run_agent as run_main_agent, LLMProviderDownError
 from main.db import MongoManager
 from workers.celery_app import celery_app
-from workers.planner.llm import get_planner_agent # noqa: E501
+from workers.planner.llm import get_planner_agent
 from workers.planner.prompts import CONTEXT_RESEARCHER_SYSTEM_PROMPT, TOOL_SELECTOR_SYSTEM_PROMPT
-from workers.planner.db import PlannerMongoManager, get_all_mcp_descriptions # noqa: E501
+from workers.planner.db import PlannerMongoManager, get_all_mcp_descriptions
 from workers.executor.tasks import execute_task_plan, run_single_item_worker, aggregate_results_callback
 from main.vector_db import get_conversation_summaries_collection
-from main.chat.prompts import STAGE_1_SYSTEM_PROMPT
 from mcp_hub.tasks.prompts import ITEM_EXTRACTOR_SYSTEM_PROMPT
 from workers.utils.text_utils import clean_llm_output
 
@@ -76,15 +75,14 @@ async def _select_relevant_tools(query: str, available_tools_map: Dict[str, str]
     """
     if not available_tools_map:
         return []
-
     try:
-        prompt = f"The user is trying to perform the following task: \"{query}\""
+        prompt = f"User Query: \"{query}\""
 
-        messages = [{'role':'system', 'content': TOOL_SELECTOR_SYSTEM_PROMPT}, {'role': 'user', 'content': prompt}]
+        messages = [{'role': 'user', 'content': prompt}]
 
         def _run_selector_sync():
             final_content_str = ""
-            for chunk in run_main_agent(system_message=STAGE_1_SYSTEM_PROMPT, function_list=[], messages=messages):
+            for chunk in run_main_agent(system_message=TOOL_SELECTOR_SYSTEM_PROMPT, function_list=[], messages=messages):
                 if isinstance(chunk, list) and chunk:
                     last_message = chunk[-1]
                     if last_message.get("role") == "assistant" and isinstance(last_message.get("content"), str):
@@ -93,19 +91,17 @@ def _run_selector_sync():
 
         final_content_str = await asyncio.to_thread(_run_selector_sync)
         cleaned_output = clean_llm_output(final_content_str)
-        parsed_output = JsonExtractor.extract_valid_json(cleaned_output)
-        selected_tools = []
-        if isinstance(parsed_output, dict) and "topic_changed" in parsed_output and "tools" in parsed_output:
-            selected_tools = parsed_output.get("tools", [])
+        selected_tools = JsonExtractor.extract_valid_json(cleaned_output)
 
-            # Separate into connected and disconnected
-            connected_tools_selected = [tool for tool in selected_tools if tool in available_tools_map]
-            
-            selected_tools = connected_tools_selected
-        
-        return selected_tools
+        if isinstance(selected_tools, list):
+            # Filter the LLM's selection to ensure they are valid and available tools
+            valid_selected_tools = [tool for tool in selected_tools if tool in available_tools_map]
+            logger.info(f"Unified Search Tool Selector identified relevant tools: {valid_selected_tools}")
+            return valid_selected_tools
+        return []
     except Exception as e:
         logger.error(f"Error during tool selection for context search: {e}", exc_info=True)
+        # Fallback to all connected tools on error
         return list(available_tools_map.keys())
 
 # Helper to run async code in Celery's sync context
@@ -371,7 +367,7 @@ async def async_refine_and_plan_ai_task(task_id: str, user_id: str):
             user_timezone = ZoneInfo("UTC")
         current_time_str = datetime.datetime.now(user_timezone).strftime('%Y-%m-%d %H:%M:%S %Z')
 
-        system_prompt = TASK_CREATION_PROMPT.format( # noqa
+        system_prompt = TASK_CREATION_PROMPT.format(
             user_name=user_name,
             user_timezone=user_timezone_str,
             current_time=current_time_str
@@ -443,7 +439,7 @@ async def async_process_change_request(task_id: str, user_id: str, user_message:
     db_manager = PlannerMongoManager()
     try:
         # 1. Fetch the task and its full history
-        task = await db_manager.get_task(task_id) # This is the new field
+        task = await db_manager.get_task(task_id)
         if not task:
             logger.error(f"Task Change Request: Task {task_id} not found.")
             return
@@ -458,7 +454,7 @@ async def async_process_change_request(task_id: str, user_id: str, user_message:
 
         # 3. Update task status and chat history in DB
         await db_manager.update_task_field(task_id, {
-            "chat_history": chat_history, # This is the new field
+            "chat_history": chat_history,
             "status": "planning" # Revert to planning to re-evaluate
         })
 
@@ -744,7 +740,7 @@ async def async_refine_task_details(task_id: str):
         user_timezone = ZoneInfo(user_timezone_str)
         current_time_str = datetime.datetime.now(user_timezone).strftime('%Y-%m-%d %H:%M:%S %Z')
 
-        system_prompt = TASK_CREATION_PROMPT.format( # noqa
+        system_prompt = TASK_CREATION_PROMPT.format(
             user_name=user_name,
             user_timezone=user_timezone_str,
             current_time=current_time_str
@@ -777,7 +773,7 @@ def execute_triggered_task(user_id: str, source: str, event_type: str, event_dat
     logger.info(f"Checking for triggered tasks for user '{user_id}' from source '{source}' event '{event_type}'.")
     run_async(async_execute_triggered_task(user_id, source, event_type, event_data))
 
-def _event_matches_filter(event_data: Dict[str, Any], task_filter: Dict[str, Any], source: str) -> bool: # noqa
+def _event_matches_filter(event_data: Dict[str, Any], task_filter: Dict[str, Any], source: str) -> bool:
     """
     Checks if an event's data matches the conditions defined in a task's filter.
     Supports complex, MongoDB-like query syntax including $or, $and, $not,