Merge branch 'development' into staging

Author: itsskofficial

src/server/main/chat/prompts.py

@@ -1,7 +1,7 @@
 STAGE_1_SYSTEM_PROMPT = """
 You are an expert Triage AI. You have two primary responsibilities:
 1.  Topic Change Detection: If the user mentions a topic that has not been discussed in the conversation history so far, set `topic_changed` to `true`. If the user is continuing a previously mentioned topic or asking a related question, set it to `false`.
-2.  Tool Selection: Based on the user's latest message and preceding relevant history/context, decide which tools are required to fulfill the request.
+2.  Tool Selection: Based on the user's latest message and preceding relevant history/context, decide which tools are required to fulfill the request. If the topic hasn't changed, keep the previous tools in your `tools` list.
 
 CRITICAL INSTRUCTIONS:
 - `topic_changed` (boolean): Set to `true` if the latest user message mentions a topic that has never been mentioned in the conversation history.

src/server/mcp_hub/gmail/main.py

@@ -73,11 +73,20 @@ async def _execute_tool(ctx: Context, func, *args, **kwargs) -> Dict[str, Any]:
 # --- Sync Tool Implementations ---
 
 def _send_email_sync(service, to: str, subject: str, body: str):
-    message_raw = base64.urlsafe_b64encode(MIMEText(body).as_bytes()).decode()
-    message_body = {"raw": message_raw, "to": to, "subject": subject}
+    msg = MIMEText(body)
+    msg["to"] = to
+    msg["subject"] = subject
+
+    # Optional: add From header to avoid confusion
+    msg["from"] = "me"
+
+    raw = base64.urlsafe_b64encode(msg.as_bytes()).decode()
+    message_body = {"raw": raw}
+
     service.users().messages().send(userId="me", body=message_body).execute()
     return {"message": "Email sent successfully."}
 
+
 def _reply_to_email_sync(service, message_id: str, body: str, reply_all: bool = False):
     original_msg = service.users().messages().get(userId="me", id=message_id, format="metadata", metadataHeaders=["subject", "from", "to", "cc", "message-id", "references"]).execute()
     headers = {h['name'].lower(): h['value'] for h in original_msg['payload']['headers']}
@@ -262,8 +271,11 @@ async def applyLabels(ctx: Context, message_id: str, label_ids: List[str]) -> Di
 async def createDraft(ctx: Context, to: str, subject: str, body: str) -> Dict[str, Any]:
     """Create a new draft email that can be edited before sending."""
     def _sync(service, to, subject, body):
-        message_raw = base64.urlsafe_b64encode(MIMEText(body).as_bytes()).decode()
-        message = {"message": {"raw": message_raw, "to": to, "subject": subject}}
+        msg = MIMEText(body)
+        msg["to"] = to
+        msg["subject"] = subject
+        message_raw = base64.urlsafe_b64encode(msg.as_bytes()).decode()
+        message = {"message": {"raw": message_raw}}
         draft = service.users().drafts().create(userId="me", body=message).execute()
         return {"draft_id": draft['id'], "message": "Draft created successfully."}
     return await _execute_tool(ctx, _sync, to=to, subject=subject, body=body)
@@ -371,8 +383,13 @@ async def removeLabels(ctx: Context, message_id: str, label_ids: List[str]) -> D
 async def updateDraft(ctx: Context, draft_id: str, to: Optional[str] = None, subject: Optional[str] = None, body: Optional[str] = None) -> Dict[str, Any]:
     """Update an existing draft email with new content."""
     def _sync(service, draft_id, to, subject, body):
-        message_raw = base64.urlsafe_b64encode(MIMEText(body).as_bytes()).decode()
-        message = {"message": {"raw": message_raw, "to": to, "subject": subject}}
+        msg = MIMEText(body)
+        if to:
+            msg["to"] = to
+        if subject:
+            msg["subject"] = subject
+        message_raw = base64.urlsafe_b64encode(msg.as_bytes()).decode()
+        message = {"message": {"raw": message_raw}}
         updated_draft = service.users().drafts().update(userId="me", id=draft_id, body=message).execute()
         return {"draft_id": updated_draft['id'], "message": "Draft updated."}
     return await _execute_tool(ctx, _sync, draft_id=draft_id, to=to, subject=subject, body=body)

src/server/workers/planner/llm.py

@@ -9,38 +9,6 @@
 
 logger = logging.getLogger(__name__)
 
-def run_agent_with_fallback(system_message: str, function_list: list, messages: list):
-    """
-    Initializes and runs a Qwen Assistant, trying a list of API keys in sequence if failures occur.
-    This function is a generator that yields the results from the successful agent run.
-    """
-    if not OPENAI_API_KEYS:
-        raise ValueError("No OpenAI API keys configured.")
-
-    errors = []
-    for i, key in enumerate(OPENAI_API_KEYS):
-        llm_cfg = {
-            'model': OPENAI_MODEL_NAME,
-            'model_server': OPENAI_API_BASE_URL,
-            'api_key': key,
-        }
-
-        try:
-            logger.info(f"Attempting to run agent with API key #{i+1}")
-            bot = Assistant(llm=llm_cfg, system_message=system_message, function_list=function_list or [])
-
-            yield from bot.run(messages=messages)
-            return # If the stream completes successfully, exit the generator.
-
-        except Exception as e:
-            error_message = f"Agent run with API key #{i+1} failed: {e}"
-            logger.warning(error_message, exc_info=True)
-            errors.append(error_message)
-            continue # Try the next key
-
-    # If the loop completes, all keys have failed
-    raise Exception(f"All OpenAI API keys failed. Errors: {errors}")
-
 def get_planner_agent(available_tools: dict, current_time_str: str, user_name: str, user_location: str, retrieved_context: dict = None):
     """Initializes and returns a Qwen Assistant agent for planning."""
 
@@ -64,23 +32,4 @@ def get_planner_agent(available_tools: dict, current_time_str: str, user_name: s
     return {
         "system_message": system_prompt,
         "function_list": []
-    }
-    
-def get_question_generator_agent(
-    original_context: dict,
-    available_tools_for_prompt: dict,
-    mcp_servers_for_agent: dict
-):
-    """Initializes a unified Qwen agent to verify context and generate clarifying questions."""
-    original_context_str = json.dumps(original_context, indent=2, default=str)
-
-    system_prompt = prompts.QUESTION_GENERATOR_SYSTEM_PROMPT.format(
-        original_context=original_context_str,
-    )
-    
-    tools_config = [{"mcpServers": mcp_servers_for_agent}]
-    
-    return {
-        "system_message": system_prompt,
-        "function_list": tools_config
-    }
+    }

src/server/workers/planner/prompts.py

@@ -71,32 +71,31 @@
 """
 
 QUESTION_GENERATOR_SYSTEM_PROMPT = """
-You are a highly intelligent context verification agent. Your primary function is to use a set of pre-selected, relevant tools to gather all necessary information about a task *before* deciding if you need to ask the user for clarification. Your goal is to be as autonomous as possible and only ask the user for information if it's a critical blocker that cannot be found with your tools.
-
-**Your Goal:**
-1.  Use the provided tools to search for context related to the task.
-2.  Analyze the gathered context.
-3.  If critical information is still missing, generate clarifying questions.
-4.  If you have enough information to create a plan, you will indicate that no questions are needed.
-
-**You have been given the following information:**
-**Original Context:** The raw information (e.g., email body) that triggered the task.
-```json
+You are a methodical Context Verification agent. Your sole purpose is to determine if enough information exists to fulfill a user's request. You will achieve this by first actively searching for information using your tools, and only then asking the user for clarification if necessary. You DO NOT perform the final task for the user.
+
+**Your Mandated Workflow:**
+
+**Step 1: Information Gathering (Tool Calls)**
+- Your first and primary action is to use the tools provided to you to find any missing information.
+- Analyze the user's request and the original context. Identify missing pieces of critical information (e.g., an email address, a document ID, a project name).
+- **You MUST call the appropriate tool to find this information.** For example, if the request is to 'email Sarthak', your first action MUST be to call `gpeople_server-search_contacts` with the query 'Sarthak'. If the request is about a document, use `gdrive_server-gdrive_search`.
+- If you believe you have enough information from the start, you can skip this step.
+
+**Step 2: Analysis & Final Output**
+- After you have exhausted your tool usage or determined no tools are needed, you must make a final decision.
+- **Scenario A: Sufficient Information:** If you have gathered all necessary information to proceed with a plan, your final output MUST be the following JSON object and nothing else:
+  `{{"clarifying_questions": []}}`
+- **Scenario B: Insufficient Information:** If critical information is still missing after you have tried to find it with your tools, your final output MUST be a JSON object containing a list of specific questions for the user. Follow this schema exactly:
+  `{{"clarifying_questions": ["What is the email address for Sarthak Karandikar?", "What should be the subject of this test email?"]}}`
+
+**CRITICAL RULES:**
+- Your response can only be one of two things: a tool call, or the final JSON object with `clarifying_questions`.
+- You are FORBIDDEN from performing the user's task (e.g., you cannot call `gmail_server-sendEmail`). Your job is only to verify context.
+- You are FORBIDDEN from outputting any JSON format other than the one specified for clarifying questions.
+- Do not include any text, explanations, or markdown formatting outside of your tool calls or the final JSON object.
+
+**Original Context Provided for this Task:**
 {original_context}
-```
-
-Output Requirements:
-Your response MUST be a single, valid JSON object that strictly adheres to the following schema. Do not include any other text or explanations.
 
-JSON Schema:
-```json
-{{
-"clarifying_questions": [
-"A clear, concise question for the user.",
-"Another question if needed."
-]
-}}
-```
-- If you have enough information to proceed with planning, return an empty list: {{"clarifying_questions": []}}.
-- If you need more information, populate the list with your questions.
+ONLY RETURN THE JSON OBJECT WITH CLARIFYING QUESTIONS OR A TOOL CALL. NEVER RETURN A PLAN OR A TEXTUAL RESPONSE. NEVER TRY TO TALK TO THE USER. NEVER TRY TO DIRECTLY PERFORM THE TASK. YOUR ROLE IS ONLY TO PERFORM CONTEXT VERIFICATION AND CHECK IF YOU HAVE THE RELEVANT INFORMATION. IF YOU DON'T HAVE THE INFORMATION, RETURN THE JSON OBJECT. NEVER TRY TO SEND EMAILS. NEVER TRY TO CREATE DOCUMENTS. NEVER TRY TO PERFORM ANY TASKS. 
 """

src/server/workers/tasks.py

@@ -18,14 +18,14 @@
 from main.llm import run_agent_with_fallback as run_main_agent_with_fallback
 from main.db import MongoManager
 from workers.celery_app import celery_app
-from workers.planner.llm import get_planner_agent, get_question_generator_agent # noqa: E501
+from workers.planner.llm import get_planner_agent # noqa: E501
+from workers.planner.prompts import QUESTION_GENERATOR_SYSTEM_PROMPT
 from workers.proactive.main import run_proactive_pipeline_logic
 from workers.planner.db import PlannerMongoManager, get_all_mcp_descriptions # noqa: E501
 from workers.memory_agent_utils import get_memory_qwen_agent, get_db_manager as get_memory_db_manager # noqa: E501
 from workers.executor.tasks import execute_task_plan
 from main.vector_db import get_conversation_summaries_collection
 from main.chat.prompts import STAGE_1_SYSTEM_PROMPT
-from workers.planner.llm import run_agent_with_fallback as run_worker_agent_with_fallback
 from mcp_hub.memory.utils import cud_memory, initialize_embedding_model, initialize_agents
 from workers.utils.text_utils import clean_llm_output
 
@@ -76,8 +76,7 @@ async def _select_relevant_tools(query: str, available_tools_map: Dict[str, str]
         return []
 
     try:
-        tools_description = "\n".join(f"- `{name}`: {desc}" for name, desc in available_tools_map.items())
-        prompt = f"User Query: \"{query}\"\n\nAvailable External Tools (for selection):\n{tools_description}"
+        prompt = f"The user is trying to perform the following task: \"{query}\" Choose the relevant tools needed to complete the task, as well as any tools where important information or context can be found related to the task. \n (For example, if the user is asking to perform a task using Gmail, you should definitely include gmail in the selected tools, but also include gpeople which can be used to find relevant contacts.)"
 
         messages = [{'role': 'user', 'content': prompt}]
 
@@ -91,11 +90,18 @@ def _run_selector_sync():
             return final_content_str
 
         final_content_str = await asyncio.to_thread(_run_selector_sync)
-        selected_tools = JsonExtractor.extract_valid_json(final_content_str)
-        if isinstance(selected_tools, list):
-            logger.info(f"Tool selector identified relevant tools for context search: {selected_tools}")
-            return selected_tools
-        return []
+        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", [])
+
+            # 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
     except Exception as e:
         logger.error(f"Error during tool selection for context search: {e}", exc_info=True)
         return list(available_tools_map.keys())
@@ -368,21 +374,26 @@ async def get_clarifying_questions(user_id: str, task_description: str, topics:
 
     logger.info(f"Context Verifier for user {user_id} will use tools: {list(mcp_servers_for_agent.keys())}")
 
-    agent_config = get_question_generator_agent(
-        original_context=original_context,
-        available_tools_for_prompt=available_tools_for_prompt,
-        mcp_servers_for_agent=mcp_servers_for_agent
+    """Initializes a unified Qwen agent to verify context and generate clarifying questions."""
+    original_context_str = json.dumps(original_context, indent=2, default=str)
+
+    system_prompt = QUESTION_GENERATOR_SYSTEM_PROMPT.format(
+        original_context=original_context_str,
     )
+    
+    tools_config = [{"mcpServers": mcp_servers_for_agent}]
 
-    user_prompt = f"Based on the task '{task_description}' and the provided context, please use your tools to find relevant information and then determine if any clarifying questions are necessary."
+    user_prompt = f"User's task request: '{task_description}'"
     messages = [{'role': 'user', 'content': user_prompt}]
 
     final_response_str = ""
-    for chunk in run_worker_agent_with_fallback(system_message=agent_config["system_message"], function_list=agent_config["function_list"], messages=messages):
+    for chunk in run_main_agent_with_fallback(system_message=system_prompt, function_list=tools_config, messages=messages):
         if isinstance(chunk, list) and chunk and chunk[-1].get("role") == "assistant":
             final_response_str = chunk[-1].get("content", "")
 
+    print ("RAW RESPONSE FROM QUESTION GENERATOR:", final_response_str)
     response_data = JsonExtractor.extract_valid_json(clean_llm_output(final_response_str))
+    print ("PARSED RESPONSE DATA:", response_data)
     if response_data and isinstance(response_data.get("clarifying_questions"), list):
         return response_data["clarifying_questions"]
     else:
@@ -531,7 +542,7 @@ async def async_generate_plan(task_id: str, user_id: str):
         messages = [{'role': 'user', 'content': user_prompt_content}]
 
         final_response_str = ""
-        for chunk in run_worker_agent_with_fallback(system_message=agent_config["system_message"], function_list=agent_config["function_list"], messages=messages):
+        for chunk in run_main_agent_with_fallback(system_message=agent_config["system_message"], function_list=agent_config["function_list"], messages=messages):
             if isinstance(chunk, list) and chunk and chunk[-1].get("role") == "assistant":
                 final_response_str = chunk[-1].get("content", "")