universal-mcp-agents 0.1.13__py3-none-any.whl → 0.1.15__py3-none-any.whl

universal_mcp/agents/shared/prompts.py

@@ -1,132 +1,83 @@
-TASK_DECOMPOSITION_PROMPT = """
-You are an expert planner. Your goal is to consolidate a complex user request into the minimum number of high-level sub-tasks required. Each sub-task should correspond to a major, consolidated action within a single target application.
+TOOL_SEARCH_QUERIES_PROMPT = """
+You are an expert at breaking down a complex user task into a list of simple, atomic search queries for finding tools. Your goal is to generate a list of queries that will cover all aspects of the user's request.
 
 **CORE PRINCIPLES:**
-1.  **App-Centric Grouping:** Group all related actions for a single application into ONE sub-task.
-2.  **Focus on Data Handoffs:** A good decomposition often involves one sub-task to *retrieve* information and a subsequent sub-task to *use* that information.
-3.  **Assume Internal Capabilities:** Do NOT create sub-tasks for abstract cognitive work like 'summarize' or 'analyze'.
-4.  **Simplify Single Actions:** If the user's task is already a single, simple action, the output should be a single sub-task that concisely describes that action. Do not make it identical to the user's input.
-5.  **General purpose sub tasks:** You also need to realise that these subtasks are going to be used to search for tools and apps. And the names and description of these tools and apps are going to be general in nature so the sub tasks should not be too specific. The task which you will get may be specific in nature but the sub taks must be general.
+1.  **Deconstruct the Task:** Analyze the user's request and identify the distinct actions or sub-tasks required.
+2.  **Include Application Context:** If the user mentions a specific application (e.g., Gmail, Google Docs, Exa), include it in the query.
+3.  **Focus on the Action:** Each query must describe a general capability. It should combine the core action (verb) and the general type of object it acts on (e.g., "create document", "get pull requests", "web search").
+4.  **STRIP SPECIFIC DETAILS:** **This is critical.** Do NOT include specific data, parameters, names, or details from the user's prompt in your queries. Your goal is to find a general tool, not to run the specific command.
+
 **--- EXAMPLES ---**
 
 **EXAMPLE 1:**
 - **User Task:** "Create a Google Doc summarizing the last 5 merged pull requests in my GitHub repo universal-mcp/universal-mcp."
-- **CORRECT DECOMPOSITION:**
-    - "Fetch the last 5 merged pull requests from the GitHub repository 'universal-mcp/universal-mcp'."
-    - "Create a new Google Doc containing the summary of the pull requests."
+- **CORRECT QUERIES:**
+    - "github get pull requests from repository"
+    - "google docs create document"
+    - "google docs append text to document"
+- **INCORRECT QUERIES:**
+    - "github get pull requests from universal-mcp/universal-mcp" (Contains specific repo name)
+    - "google docs create 'summary' document" (Contains specific document title)
+
 
 **EXAMPLE 2:**
-- **User Task:** "Find the best restaurants in Goa using perplexity web search."
-- **CORRECT DECOMPOSITION:**
-    - "Perform a web search using Perplexity to find the best restaurants in Goa."
+- **User Task:** "Find the best restaurants in Goa using exa web search, then email the list to my friend at test@example.com."
+- **CORRECT QUERIES:**
+    - "exa web search"
+    - "send email"
+- **INCORRECT QUERIES:**
+    - "exa search for best restaurants in Goa" (Contains specific search details)
+    - "email list to test@example.com" (Contains specific parameters)
+
+**EXAMPLE 3:**
+- **User Task:** "add an event to my google calendar at 2pm called 'Walk in the park'?"
+- **CORRECT QUERIES:**
+    - "google calendar create calendar event"
+- **INCORRECT QUERIES:**
+    - "google calendar create event 'Walk in the park' at 2pm" (Contains specific event details)
 
 **--- YOUR TASK ---**
 
 **USER TASK:**
 "{task}"
 
-**YOUR DECOMPOSITION (as a list of strings):**
-"""
-
-
-APP_SEARCH_QUERY_PROMPT = """
-You are an expert at selecting an application to perform a specific sub-task. Your goal is to generate a concise query for an app search engine.
-
-Analyze the current sub-task in the context of the original user goal and the ENTIRE PLAN so far.
-
-**CORE INSTRUCTION:** If any application already used in the plan is capable of performing the current sub-task, your query MUST BE the name of that application to ensure continuity and efficiency. Otherwise, generate a concise query for the category of application needed.
-
-**--- EXAMPLES ---**
-
-**EXAMPLE 1: Reusing an app from two steps ago**
-- **Original User Task:** "Find my latest order confirmation in Gmail, search for reviews of the main product on perplexity, and then send an email to ankit@agentr.dev telling about the reviews"
-- **Plan So Far:**
-  - The sub-task 'Find order confirmation in Gmail' was assigned to app 'google_mail'.
-  - The sub-task 'Search for product reviews on perplexity' was assigned to app 'perplexity'.
-- **Current Sub-task:** "send an email to ankit@agentr.dev"
-- **CORRECT QUERY:** "google_mail"
-
-**EXAMPLE 2: First Step (No previous context)**
-- **Original User Task:** "Find the best restaurants in Goa."
-- **Plan So Far:** None. This is the first step.
-- **Current Sub-task:** "Perform a web search to find the best restaurants in Goa."
-- **CORRECT QUERY:** "web search"
-
-**--- YOUR TASK ---**
-
-**Original User Task:**
-"{original_task}"
-
-**Plan So Far:**
-{plan_context}
-
-**Current Sub-task:**
-"{sub_task}"
-
-**YOUR CONCISE APP SEARCH QUERY:**
+**YOUR SEARCH QUERIES (as a list of strings):**
 """
 
-
-TOOL_SEARCH_QUERY_PROMPT = """
-You are an expert at summarizing the core *action* of a sub-task into a concise query for finding a tool. This query should ignore any application names.
+APP_SELECTION_PROMPT = """
+You are an AI assistant that selects the most appropriate applications (apps) from a list to accomplish a user's task.
 
 **INSTRUCTIONS:**
-1.  Focus only on the verb or action being performed in the sub-task.
-2.  Include key entities related to the action.
-3.  Do NOT include the names of applications (e.g., "Perplexity", "Gmail").
-4.  You also need to realise that this query is going to be used to search for tools in a particular app. And the names and description of these tools are going to be general in nature so the query should not be too specific. The sub task which you will get may be specific in nature but the query must be general.
-
-**EXAMPLES:**
-- **Sub-task:** "Perform a web search using Perplexity to find the best restaurants in Goa."
-- **Query:** "web search for restaurants"
-
-- **Sub-task:** "Fetch all marketing emails received from Gmail in the last 7 days."
-- **Query:** "get emails by date"
-
-- **Sub-task:** "Create a new Google Doc and append a summary."
-- **Query:** "create document, append text"
-
-**SUB-TASK:**
-"{sub_task}"
-
-**YOUR CONCISE TOOL SEARCH QUERY:**
-"""
-
-REVISE_DECOMPOSITION_PROMPT = """
-You are an expert planner who revises plans that have failed. Your previous attempt to break down a task resulted in a sub-task that could not be matched with any available tools.
-
-**INSTRUCTIONS:**
-1.  Analyze the original user task and the failed sub-task.
-2.  Generate a NEW, alternative decomposition of the original task.
-3.  This new plan should try to achieve the same overall goal but with different, perhaps broader or more combined, sub-tasks to increase the chance of finding a suitable tool.
+1.  Carefully review the original user task to understand the complete goal.
+2.  Examine the list of available apps, their IDs, and their descriptions.
+3.  Select ALL app IDs that are necessary to complete the entire task.
+4.  If the user's task mentions a specific app, you MUST select it.
+5.  If no apps are a good fit, return an empty list.
 
 **ORIGINAL USER TASK:**
 "{task}"
 
-**FAILED SUB-TASK FROM PREVIOUS PLAN:**
-"{failed_sub_task}"
+**AVAILABLE APPS:**
+{app_candidates}
 
-**YOUR NEW, REVISED DECOMPOSITION (as a list of strings):**
+**YOUR SELECTED APP ID(s) (as a list of strings):**
 """
 
-
 TOOL_SELECTION_PROMPT = """
-You are an AI assistant that selects the most appropriate tool(s) from a list to accomplish a specific sub-task.
+You are an AI assistant that selects the most appropriate tool(s) from a list to accomplish a user's overall task.
 
 **INSTRUCTIONS:**
-1.  Carefully review the sub-task to understand the required action.
-2.  Examine the list of available tools and their descriptions.
-3.  Select the best tool ID that matches the sub-task. You are encouraged to select multiple tools if there are multiple tools with similar capabilties
-or names. It is always good to have more tools than having insufficent tools.
-4.  If no tool is a good fit, return an empty list.
-5.  Only return the tool IDs.
-6. You should understand that the sub task maybe specific in nature but the tools are made to be general purpose and therefore the tool_candidates you will get will be very general purpose but that should not stop you from selecting the tools as these tools will be given to a very smart agent who will be able to use these tools for the specific sub-taks
+1.  Carefully review the original user task to understand the complete goal.
+2.  Examine the list of available tools, their IDs, and their descriptions. These tools have been found using a search based on the task.
+3.  Select all tool IDs that are necessary to complete the entire task. It is critical to select all tools needed for a multi-step task.
+4.  If no tools are a good fit for the task, return an empty list.
+5.  Only return the tool IDs. The tools are general purpose, but you are smart enough to see how they can be used for the specific task.
 
-**SUB-TASK:**
-"{sub_task}"
+**ORIGINAL USER TASK:**
+"{task}"
 
 **AVAILABLE TOOLS:**
 {tool_candidates}
 
-**YOUR SELECTED TOOL ID(s):**
+**YOUR SELECTED TOOL ID(s) (as a list of strings):**
 """

universal_mcp/agents/shared/tool_node.py

@@ -1,227 +1,211 @@
+import asyncio
+from collections import defaultdict
 from typing import Annotated, TypedDict
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, AnyMessage
 from langgraph.graph import END, StateGraph
 from langgraph.graph.message import add_messages
+from langgraph.types import Command
 from loguru import logger
 from pydantic import BaseModel, Field
 from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolConfig
 
 from universal_mcp.agents.shared.prompts import (
-    APP_SEARCH_QUERY_PROMPT,
-    REVISE_DECOMPOSITION_PROMPT,
-    TASK_DECOMPOSITION_PROMPT,
-    TOOL_SEARCH_QUERY_PROMPT,
+    APP_SELECTION_PROMPT,
+    TOOL_SEARCH_QUERIES_PROMPT,
     TOOL_SELECTION_PROMPT,
 )
 
-MAX_DECOMPOSITION_ATTEMPTS = 2
+MAX_RETRIES = 1
 
-# --- Pydantic Models for Structured LLM Outputs ---
 
+class SearchQueries(BaseModel):
+    queries: list[str] = Field(description="A list of search queries for finding tools.")
 
-class TaskDecomposition(BaseModel):
-    sub_tasks: list[str] = Field(description="A list of sub-task descriptions.")
 
-
-class SearchQuery(BaseModel):
-    query: str = Field(description="A concise search query.")
+class AppSelection(BaseModel):
+    app_ids: list[str] = Field(description="The IDs of the selected applications.")
 
 
 class ToolSelection(BaseModel):
     tool_ids: list[str] = Field(description="The IDs of the selected tools.")
 
 
-# --- LangGraph Agent State ---
+class AgentState(TypedDict):
+    """The central state of our agent graph."""
 
+    original_task: str
+    queries: list[str]
+    candidate_tools: list[dict]
+    execution_plan: ToolConfig
+    messages: Annotated[list[AnyMessage], add_messages]
+    retry_count: int
 
-class SubTask(TypedDict, total=False):
-    """Represents a single step in the execution plan."""
 
-    task: str
-    status: str  # "pending", "success", "failed"
-    app_id: str
-    tool_ids: list[str]
-    reasoning: str
+def build_tool_node_graph(llm: BaseChatModel, registry: ToolRegistry) -> StateGraph:
+    """Builds a workflow for tool selection with a retry mechanism."""
 
+    async def _search_for_tools(state: AgentState) -> Command:
+        """
+        Performs a hierarchical search:
+        1. Generates search queries for the task.
+        2. Searches for candidate *applications*.
+        3. Uses an LLM to select the most relevant applications.
+        4. Searches for tools only within the selected applications.
+        If any step fails, it can trigger a retry.
+        """
+        task = state["original_task"]
+        logger.info(f"Starting hierarchical tool search for task: '{task}'")
+
+        prompt = TOOL_SEARCH_QUERIES_PROMPT.format(task=task)
+        response = await llm.with_structured_output(SearchQueries).ainvoke(prompt)
+        queries = response.queries
+        logger.info(f"Generated search queries: {queries}")
+
+        if not queries:
+            logger.error("LLM failed to generate any search queries.")
+            return Command(
+                update={"messages": [AIMessage(content="I could not understand the task to search for tools.")]},
+                goto="handle_failure",
+            )
 
-class AgentState(TypedDict):
-    """The central state of our agent graph."""
+        # Always store queries for potential retry
+        update_state = {"queries": queries}
 
-    original_task: str
-    decomposition_attempts: int
-    failed_sub_task_info: str  # To inform re-decomposition
-    sub_tasks: list[SubTask]
-    execution_plan: list[SubTask]
-    messages: Annotated[list[AnyMessage], add_messages]
+        app_search_tasks = [registry.search_apps(query, distance_threshold=0.7) for query in queries]
+        app_results = await asyncio.gather(*app_search_tasks)
+        unique_apps = {app["id"]: app for app_list in app_results for app in app_list}
 
+        if not unique_apps:
+            logger.warning(f"No applications found for queries: {queries}. Triggering retry.")
+            return Command(update=update_state, goto="general_search_and_select")
 
-# --- Graph Builder ---
+        logger.info(f"Found {len(unique_apps)} candidate applications.")
 
+        app_candidates_str = "\n - ".join([f"{app['id']}: {app['description']}" for app in unique_apps.values()])
+        app_selection_prompt = APP_SELECTION_PROMPT.format(task=task, app_candidates=app_candidates_str)
+        app_selection_response = await llm.with_structured_output(AppSelection).ainvoke(app_selection_prompt)
+        selected_app_ids = app_selection_response.app_ids
 
-def build_tool_node_graph(llm: BaseChatModel, registry: ToolRegistry) -> StateGraph:
-    """Builds the adaptive LangGraph workflow for tool selection."""
+        if not selected_app_ids:
+            logger.warning("LLM did not select any applications from the candidate list. Triggering retry.")
+            return Command(update=update_state, goto="general_search_and_select")
 
-    async def _decompose_task(state: AgentState) -> AgentState:
-        """Decomposes the main task or revises a failed decomposition."""
-        attempts = state.get("decomposition_attempts", 0)
-        task = state["original_task"]
-        failed_info = state.get("failed_sub_task_info")
+        logger.success(f"Selected {len(selected_app_ids)} applications: {selected_app_ids}")
 
-        if attempts > 0 and failed_info:
-            logger.warning(f"Revising decomposition. Attempt {attempts + 1}.")
-            prompt = REVISE_DECOMPOSITION_PROMPT.format(task=task, failed_sub_task=failed_info)
-        else:
-            logger.info("Performing initial task decomposition.")
-            prompt = TASK_DECOMPOSITION_PROMPT.format(task=task)
-
-        response = await llm.with_structured_output(TaskDecomposition).ainvoke(prompt)
-        sub_tasks = [{"task": sub_task_str, "status": "pending"} for sub_task_str in response.sub_tasks]
-
-        return {
-            "sub_tasks": sub_tasks,
-            "decomposition_attempts": attempts + 1,
-            "messages": [AIMessage(content=f"New plan created with {len(sub_tasks)} steps.")],
-        }
-
-    async def _resolve_sub_tasks(state: AgentState) -> AgentState:
-        """Iterates through sub-tasks, providing full plan context to the app selection prompt."""
-        sub_tasks = state["sub_tasks"]
-        original_task = state["original_task"]
-        current_plan = []
-
-        for i, sub_task in enumerate(sub_tasks):
-            task_desc = sub_task["task"]
-            logger.info(f"Resolving sub-task: '{task_desc}'")
-
-            # 1. Build the FULL context string from the entire plan so far
-            if not current_plan:
-                plan_context_str = "None. This is the first step."
-            else:
-                context_lines = [
-                    f"- The sub-task '{step['task']}' was assigned to app '{step['app_id']}'." for step in current_plan
-                ]
-                plan_context_str = "\n".join(context_lines)
-
-            # 2. Generate the App-specific query using the NEW full-context prompt
-            app_query_prompt = APP_SEARCH_QUERY_PROMPT.format(
-                original_task=original_task, plan_context=plan_context_str, sub_task=task_desc
-            )
-            app_query_response = await llm.with_structured_output(SearchQuery).ainvoke(app_query_prompt)
-            app_search_query = app_query_response.query
-            logger.info(f"Generated context-aware app search query: '{app_search_query}'")
-
-            # 3. Search for candidate apps (the rest of the logic is the same)
-            candidate_apps = await registry.search_apps(query=app_search_query, limit=5)
-            if not candidate_apps:
-                logger.error(f"No apps found for query '{app_search_query}' from sub-task: '{task_desc}'")
-                return {"failed_sub_task_info": task_desc, "sub_tasks": []}
-
-            # 4. Generate Action-specific query for finding the tool
-            tool_query_prompt = TOOL_SEARCH_QUERY_PROMPT.format(sub_task=task_desc)
-            tool_query_response = await llm.with_structured_output(SearchQuery).ainvoke(tool_query_prompt)
-            tool_search_query = tool_query_response.query
-            logger.info(f"Generated tool search query: '{tool_search_query}'")
-
-            # 5. Find a suitable tool within the candidate apps
-            tool_found = False
-            for app in candidate_apps:
-                app_id = app["id"]
-                logger.info(f"Searching for tools in app '{app_id}' with query '{tool_search_query}'...")
-
-                found_tools = await registry.search_tools(query=tool_search_query, app_id=app_id, limit=5)
-                if not found_tools:
-                    continue
-
-                tool_candidates_str = "\n - ".join([f"{tool['name']}: {tool['description']}" for tool in found_tools])
-                selection_prompt = TOOL_SELECTION_PROMPT.format(sub_task=task_desc, tool_candidates=tool_candidates_str)
-                selection_response = await llm.with_structured_output(ToolSelection).ainvoke(selection_prompt)
-
-                if selection_response.tool_ids:
-                    logger.success(f"Found and selected tool(s) {selection_response.tool_ids} in app '{app_id}'.")
-                    sub_task.update(
-                        {
-                            "status": "success",
-                            "app_id": app_id,
-                            "tool_ids": selection_response.tool_ids,
-                            "reasoning": f"Selected tool(s) {selection_response.tool_ids} from app '{app_id}' for sub-task.",
-                        }
-                    )
-                    current_plan.append(sub_task)
-                    tool_found = True
-                    break
-
-            if not tool_found:
-                logger.error(f"Could not find any suitable tool for sub-task: '{task_desc}'")
-                return {"failed_sub_task_info": task_desc, "sub_tasks": []}
-
-        return {"execution_plan": current_plan, "sub_tasks": []}
-
-    def _handle_planning_failure(state: AgentState) -> AgentState:
-        """Handles the case where all decomposition attempts have failed."""
-        logger.error("Maximum decomposition attempts reached. Planning failed.")
-        return {
-            "messages": [
-                AIMessage(
-                    content="I am unable to create a complete plan for this task with the available tools. Please try rephrasing your request."
-                )
-            ]
-        }
+        tool_search_tasks = [
+            registry.search_tools(task, app_id=app_id, distance_threshold=0.8) for app_id in selected_app_ids
+        ]
+        tool_results = await asyncio.gather(*tool_search_tasks)
+        candidate_tools = [tool for tool_list in tool_results for tool in tool_list]
 
-    def _consolidate_plan(state: AgentState) -> AgentState:
+        if not candidate_tools:
+            logger.warning(f"No tools found within the selected applications: {selected_app_ids}. Triggering retry.")
+            return Command(update=update_state, goto="general_search_and_select")
+
+        logger.success(f"Found {len(candidate_tools)} candidate tools from selected apps.")
+        update_state["candidate_tools"] = candidate_tools
+        return Command(update=update_state, goto="select_tools_for_plan")
+
+    async def _general_search_and_select(state: AgentState) -> Command:
         """
-        NEW: Merges steps in the execution plan that use the same app_id.
-        It combines their tool_ids into a single unique list.
+        A retry node that performs a general tool search without app filters.
         """
-        logger.info("Consolidating final execution plan.")
-        plan = state["execution_plan"]
-        merged_apps: dict[str, SubTask] = {}
-
-        for step in plan:
-            app_id = step["app_id"]
-            if app_id not in merged_apps:
-                # Store the first occurrence of this app
-                merged_apps[app_id] = step.copy()
-                merged_apps[app_id]["tool_ids"] = set(step["tool_ids"])
-            else:
-                # If app already seen, just update its set of tool_ids
-                merged_apps[app_id]["tool_ids"].update(step["tool_ids"])
+        state["original_task"]
+        queries = state["queries"]
+        retry_count = state.get("retry_count", 0)
+
+        if retry_count >= MAX_RETRIES:
+            logger.error("Max retries reached. Failing the planning process.")
+            return Command(
+                update={
+                    "messages": [AIMessage(content="I could not find any relevant tools after extensive searching.")]
+                },
+                goto="handle_failure",
+            )
 
-        # Convert the merged dictionary back to a list of SubTasks
-        final_plan = []
-        for app_id, step_data in merged_apps.items():
-            step_data["tool_ids"] = sorted(list(step_data["tool_ids"]))
-            final_plan.append(step_data)
+        logger.info(f"--- RETRY {retry_count + 1}/{MAX_RETRIES} ---")
+        logger.info("Performing a general tool search without app filters.")
+
+        general_search_tasks = [registry.search_tools(query, distance_threshold=0.85) for query in queries]
+        tool_results = await asyncio.gather(*general_search_tasks)
+
+        unique_tools = {tool["id"]: tool for tool_list in tool_results for tool in tool_list}
+        candidate_tools = list(unique_tools.values())
+
+        if not candidate_tools:
+            logger.error("General search (retry) also failed to find any tools.")
+            return Command(
+                update={
+                    "messages": [
+                        AIMessage(content="I could not find any tools for your request, even with a broader search.")
+                    ],
+                    "retry_count": retry_count + 1,
+                },
+                goto="handle_failure",
+            )
 
-        return {"execution_plan": final_plan}
+        logger.success(f"General search found {len(candidate_tools)} candidate tools.")
+        return Command(
+            update={"candidate_tools": candidate_tools, "retry_count": retry_count + 1},
+            goto="select_tools_for_plan",
+        )
 
-    # --- Graph Definition ---
+    async def _select_tools_for_plan(state: AgentState) -> Command:
+        """Selects the best tools from the candidates and builds the final execution plan."""
+        task = state["original_task"]
+        candidate_tools = state["candidate_tools"]
+        retry_count = state.get("retry_count", 0)
+        logger.info("Starting tool selection from candidate list.")
+
+        tool_candidates_str = "\n - ".join([f"{tool['id']}: {tool['description']}" for tool in candidate_tools])
+        prompt = TOOL_SELECTION_PROMPT.format(task=task, tool_candidates=tool_candidates_str)
+        response = await llm.with_structured_output(ToolSelection).ainvoke(prompt)
+        selected_tool_ids = response.tool_ids
+
+        if not selected_tool_ids:
+            if retry_count >= MAX_RETRIES:
+                logger.error("LLM did not select any tools, even after a retry. Failing.")
+                return Command(
+                    update={
+                        "messages": [AIMessage(content="I found potential tools, but could not create a final plan.")]
+                    },
+                    goto="handle_failure",
+                )
+            else:
+                logger.warning(
+                    "LLM did not select any tools from the current candidate list. Triggering general search."
+                )
+                return Command(goto="general_search_and_select")
 
-    workflow = StateGraph(AgentState)
+        logger.success(f"Selected {len(selected_tool_ids)} tools for the final plan: {selected_tool_ids}")
 
-    workflow.add_node("decompose_task", _decompose_task)
-    workflow.add_node("resolve_sub_tasks", _resolve_sub_tasks)
-    workflow.add_node("consolidate_plan", _consolidate_plan)  # NEW NODE
-    workflow.add_node("handle_planning_failure", _handle_planning_failure)
+        final_plan = defaultdict(list)
+        for tool_id in selected_tool_ids:
+            if "__" in tool_id:
+                app_id, tool_name = tool_id.split("__", 1)
+                final_plan[app_id].append(tool_name)
 
-    workflow.set_entry_point("decompose_task")
+        sorted_final_plan = {app_id: sorted(tools) for app_id, tools in final_plan.items()}
+        return Command(update={"execution_plan": sorted_final_plan}, goto=END)
 
-    def should_continue(state: AgentState):
-        if not state.get("sub_tasks"):  # Resolution failed or succeeded
-            if state.get("execution_plan"):
-                return "consolidate_plan"  # MODIFIED: Go to consolidate on success
-            elif state["decomposition_attempts"] >= MAX_DECOMPOSITION_ATTEMPTS:
-                return "handle_planning_failure"
-            else:
-                return "decompose_task"  # Re-try decomposition
+    def _handle_planning_failure(state: AgentState) -> Command:
+        """Handles cases where tool search or selection fails by logging the final error message."""
+        if messages := state.get("messages"):
+            last_message = messages[-1].content
+            logger.error(f"Planning failed. Final message: {last_message}")
         else:
-            return "resolve_sub_tasks"
+            logger.error("Planning failed with no specific message.")
+        return Command(goto=END)
 
-    workflow.add_conditional_edges("decompose_task", lambda s: "resolve_sub_tasks")
-    workflow.add_conditional_edges("resolve_sub_tasks", should_continue)
+    workflow = StateGraph(AgentState)
+    workflow.add_node("search_for_tools", _search_for_tools)
+    workflow.add_node("general_search_and_select", _general_search_and_select)
+    workflow.add_node("select_tools_for_plan", _select_tools_for_plan)
+    workflow.add_node("handle_failure", _handle_planning_failure)
 
-    workflow.add_edge("consolidate_plan", END)  # NEW EDGE
-    workflow.add_edge("handle_planning_failure", END)
+    workflow.set_entry_point("search_for_tools")
 
     return workflow.compile()