universal-mcp-agents 0.1.22__py3-none-any.whl → 0.1.23__py3-none-any.whl

universal_mcp/agents/codeact0/tools.py

@@ -1,22 +1,21 @@
 import asyncio
-import json
+import base64
 from collections import defaultdict
+from pathlib import Path
 from typing import Annotated, Any
 
 from langchain_core.tools import tool
-from loguru import logger
 from pydantic import Field
+from universal_mcp.agentr.client import AgentrClient
 from universal_mcp.agentr.registry import AgentrRegistry
+from universal_mcp.applications.markitdown.app import MarkitdownApp
 from universal_mcp.types import ToolFormat
 
-
-def enter_playbook_mode():
-    """Call this function to enter playbook mode. Playbook mode is when the user wants to store a repeated task as a script with some inputs for the future."""
-    return
+from universal_mcp.agents.codeact0.prompts import build_tool_definitions
 
 
-def exit_playbook_mode():
-    """Call this function to exit playbook mode. Playbook mode is when the user wants to store a repeated task as a script with some inputs for the future."""
+def enter_agent_builder_mode():
+    """Call this function to enter agent builder mode. Agent builder mode is when the user wants to store a repeated task as a script with some inputs for the future."""
     return
 
 
@@ -26,189 +25,323 @@ def create_meta_tools(tool_registry: AgentrRegistry) -> dict[str, Any]:
     @tool
     async def search_functions(
         queries: Annotated[
-            list[str] | str | None,
-            Field(description="A single query or a list of queries to search for relevant functions"),
+            list[list[str]] | None,
+            Field(
+                description="A list of query lists. Each inner list contains one or more search terms that will be used together to find relevant tools."
+            ),
         ] = None,
-        app_id: Annotated[
-            str | None,
-            Field(description="The ID or common name of a specific application to search within"),
+        app_ids: Annotated[
+            list[str] | None,
+            Field(description="The ID or list of IDs (common names) of specific applications to search within."),
         ] = None,
     ) -> str:
         """
-        Searches for relevant functions across applications based on queries and/or a specific app.
-        This function operates in three modes:
-
-        1.  **Global Search (provide `queries` only):**
-            - Use when the user wants to perform an action without specifying an application.
-            - The system will search across all available functions.
-            - Example: For "how can I create a presentation?", call with queries=["create presentation"].
-
-        2.  **App Discovery (provide `app_id` only):**
-            - Use when the user asks about the capabilities of a specific application.
-            - The `app_id` can be the common name of the app (e.g., "Gmail", "Google Drive").
-            - This will return all available functions for that application, up to a limit.
-            - Example: For "what can you do with Gmail?", call with app_id="Gmail".
-
-        3.  **Scoped Search (provide `queries` AND `app_id`):**
-            - Use when the user wants to perform an action within a specific application.
-            - This performs a targeted search only within the specified app's functions.
-            - Example: For "how do I find an email in Gmail?", call with queries=["find email"], app_id="Gmail".
+        Searches for relevant functions based on queries and/or applications. This function
+        operates in three powerful modes with support for multi-query searches:
+
+        1.  **Global Search** (`queries` only as List[List[str]]):
+            - Searches all functions across all applications.
+            - Supports multiple independent searches in parallel.
+            - Each inner list represents a separate search query.
+
+            Examples:
+            - Single global search:
+              `search_functions(queries=[["create presentation"]])`
+
+            - Multiple independent global searches:
+              `search_functions(queries=[["send email"], ["schedule meeting"]])`
+
+            - Multi-term search for comprehensive results:
+              `search_functions(queries=[["send email", "draft email", "compose email"]])`
+
+        2.  **App Discovery** (`app_ids` only as List[str]):
+            - Returns ALL available functions for one or more specific applications.
+            - Use this to explore the complete capability set of an application.
+
+            Examples:
+            - Single app discovery:
+              `search_functions(app_ids=["Gmail"])`
+
+            - Multiple app discovery:
+              `search_functions(app_ids=["Gmail", "Google Calendar", "Slack"])`
+
+        3.  **Scoped Search** (`queries` as List[List[str]] and `app_ids` as List[str]):
+            - Performs targeted searches within specific applications in parallel.
+            - The number of app_ids must match the number of inner query lists.
+            - Each query list is searched within its corresponding app_id.
+            - Supports multiple search terms per app for comprehensive discovery.
+
+            Examples:
+            - Basic scoped search (one query per app):
+              `search_functions(queries=[["find email"], ["share file"]], app_ids=["Gmail", "Google_Drive"])`
+
+            - Multi-term scoped search (multiple queries per app):
+              `search_functions(
+                  queries=[
+                      ["send email", "draft email", "compose email", "reply to email"],
+                      ["create event", "schedule meeting", "find free time"],
+                      ["upload file", "share file", "create folder", "search files"]
+                  ],
+                  app_ids=["Gmail", "Google Calendar", "Google_Drive"]
+              )`
+
+            - Mixed complexity (some apps with single query, others with multiple):
+              `search_functions(
+                  queries=[
+                      ["list messages"],
+                      ["create event", "delete event", "update event"]
+                  ],
+                  app_ids=["Gmail", "Google Calendar"]
+              )`
+
+        **Pro Tips:**
+        - Use multiple search terms in a single query list to cast a wider net and discover related functionality
+        - Multi-term searches are more efficient than separate calls
+        - Scoped searches return more focused results than global searches
+        - The function returns connection status for each app (connected vs NOT connected)
+        - All searches within a single call execute in parallel for maximum efficiency
+
+        **Parameters:**
+        - `queries` (List[List[str]], optional): A list of query lists. Each inner list contains one or more
+          search terms that will be used together to find relevant tools.
+        - `app_ids` (List[str], optional): A list of application IDs to search within or discover.
+
+        **Returns:**
+        - A structured response containing:
+          - Matched tools with their descriptions
+          - Connection status for each app
+          - Recommendations for which tools to load next
         """
-        if isinstance(queries, str):  # Handle JSON string input
-            try:
-                queries = json.loads(queries)
-            except json.JSONDecodeError:
-                # If it's a single query as a string, convert to list
-                queries = [queries] if queries else None
+        registry = tool_registry
 
-        if not queries and not app_id:
-            raise ValueError("You must provide 'queries', an 'app_id', or both.")
+        TOOL_THRESHOLD = 0.75
+        APP_THRESHOLD = 0.7
 
-        registry = tool_registry
-        connections = await registry.list_connected_apps()
-        connected_app_ids = {connection["app_id"] for connection in connections}
+        # --- Helper Functions for Different Search Modes ---
 
-        canonical_app_id = None
-        found_tools_result = []
-        THRESHOLD = 0.8
-
-        if app_id:
-            relevant_apps = await registry.search_apps(query=app_id, distance_threshold=THRESHOLD)
-            if not relevant_apps:
-                return {
-                    "found_tools": [],
-                    "message": f"Search failed. Application '{app_id}' was not found.",
-                }
-            canonical_app_id = relevant_apps[0]["id"]
-
-        if canonical_app_id and not queries:
-            all_app_tools = await registry.search_tools(query="", app_id=canonical_app_id, limit=20)
-
-            tool_list = []
-            for tool in all_app_tools:
-                cleaned_description = tool.get("description", "").split("Context:")[0].strip()
-                tool_list.append({"id": tool["id"], "description": cleaned_description})
-
-            found_tools_result.append(
-                {
-                    "app_id": canonical_app_id,
-                    "connection_status": "connected" if canonical_app_id in connected_app_ids else "not_connected",
-                    "tools": tool_list,
-                }
-            )
+        async def _handle_global_search(queries: list[str]) -> list[list[dict[str, Any]]]:
+            """Performs a broad search across all apps to find relevant tools and apps."""
+            # 1. Perform initial broad searches for tools and apps concurrently.
+            initial_tool_tasks = [registry.search_tools(query=q, distance_threshold=TOOL_THRESHOLD) for q in queries]
+            app_search_tasks = [registry.search_apps(query=q, distance_threshold=APP_THRESHOLD) for q in queries]
 
-        else:
-            query_results = []
-            prioritized_app_id_list = []
+            initial_tool_results, app_search_results = await asyncio.gather(
+                asyncio.gather(*initial_tool_tasks), asyncio.gather(*app_search_tasks)
+            )
 
-            if canonical_app_id:
-                prioritized_app_id_list = [canonical_app_id]
-            else:
-                # 1. Perform an initial broad search for tools.
-                initial_tool_search_tasks = [registry.search_tools(query=q, distance_threshold=THRESHOLD) for q in queries]
-                initial_tool_results = await asyncio.gather(*initial_tool_search_tasks)
-
-                # 2. Search for relevant apps.
-                app_search_tasks = [registry.search_apps(query=q, distance_threshold=THRESHOLD) for q in queries]
-                app_search_results = await asyncio.gather(*app_search_tasks)
-
-                # 3. Create a prioritized list of app IDs for the final search.
-                # Apps found via search_apps are considered higher priority and come first.
-                app_ids_from_apps = {app["id"] for result_list in app_search_results for app in result_list}
-                # Use a list to maintain order.
-                prioritized_app_id_list.extend(list(app_ids_from_apps))
-
-                # Add app_ids from the initial tool search, ensuring no duplicates.
-                app_ids_from_tools = {tool["app_id"] for result_list in initial_tool_results for tool in result_list}
-
-                for tool_app_id in app_ids_from_tools:
-                    if tool_app_id not in app_ids_from_apps:
-                        prioritized_app_id_list.append(tool_app_id)
-
-            # 4. Perform the final, comprehensive tool search across the prioritized list of apps.
-            if prioritized_app_id_list:
-                # print(f"Prioritized app IDs for final search: {prioritized_app_id_list}")
-                final_tool_search_tasks = []
-                for app_id_to_search in prioritized_app_id_list:
-                    for query in queries:
-                        final_tool_search_tasks.append(
-                            registry.search_tools(query=query, app_id=app_id_to_search, distance_threshold=THRESHOLD)
-                        )
-                query_results = await asyncio.gather(*final_tool_search_tasks)
-
-            # 5. Aggregate all found tools for easy lookup.
+            # 2. Create a prioritized list of app IDs for the final search.
+            app_ids_from_apps = {app["id"] for result_list in app_search_results for app in result_list}
+            prioritized_app_id_list = list(app_ids_from_apps)
+
+            app_ids_from_tools = {tool["app_id"] for result_list in initial_tool_results for tool in result_list}
+            for tool_app_id in app_ids_from_tools:
+                if tool_app_id not in app_ids_from_apps:
+                    prioritized_app_id_list.append(tool_app_id)
+
+            if not prioritized_app_id_list:
+                return []
+
+            # 3. Perform the final, comprehensive tool search across the prioritized apps.
+            final_tool_search_tasks = [
+                registry.search_tools(query=query, app_id=app_id_to_search, distance_threshold=TOOL_THRESHOLD)
+                for app_id_to_search in prioritized_app_id_list
+                for query in queries
+            ]
+            return await asyncio.gather(*final_tool_search_tasks)
+
+        async def _handle_scoped_search(app_ids: list[str], queries: list[list[str]]) -> list[list[dict[str, Any]]]:
+            """Performs targeted searches for specific queries within specific applications."""
+            if len(app_ids) != len(queries):
+                raise ValueError("The number of app_ids must match the number of query lists.")
+
+            tasks = []
+            for app_id, query_list in zip(app_ids, queries):
+                for query in query_list:
+                    # Create a search task for each query in the list for the corresponding app
+                    tasks.append(registry.search_tools(query=query, app_id=app_id, distance_threshold=TOOL_THRESHOLD))
+
+            return await asyncio.gather(*tasks)
+
+        async def _handle_app_discovery(app_ids: list[str]) -> list[list[dict[str, Any]]]:
+            """Fetches all tools for a list of applications."""
+            tasks = [registry.search_tools(query="", app_id=app_id, limit=20) for app_id in app_ids]
+            return await asyncio.gather(*tasks)
+
+        # --- Helper Functions for Structuring and Formatting Results ---
+
+        def _format_response(structured_results: list[dict[str, Any]]) -> str:
+            """Builds the final, user-facing formatted string response from structured data."""
+            if not structured_results:
+                return "No relevant functions were found."
+
+            result_parts = []
+            apps_in_results = {app["app_id"] for app in structured_results}
+            connected_apps_in_results = {
+                app["app_id"] for app in structured_results if app["connection_status"] == "connected"
+            }
+
+            for app in structured_results:
+                app_id = app["app_id"]
+                app_status = "connected" if app["connection_status"] == "connected" else "NOT connected"
+                result_parts.append(f"Tools from {app_id} (status: {app_status} by user):")
+
+                for tool in app["tools"]:
+                    result_parts.append(f" - {tool['id']}: {tool['description']}")
+                result_parts.append("")  # Empty line for readability
+
+            # Add summary connection status messages
+            if not connected_apps_in_results and len(apps_in_results) > 1:
+                result_parts.append(
+                    "Connection Status: None of the apps in the results are connected. "
+                    "You must ask the user to choose the application."
+                )
+            elif len(connected_apps_in_results) > 1:
+                connected_list = ", ".join(sorted(list(connected_apps_in_results)))
+                result_parts.append(
+                    f"Connection Status: Multiple apps are connected ({connected_list}). "
+                    "You must ask the user to select which application they want to use."
+                )
+
+            result_parts.append("Call load_functions to select the required functions only.")
+            if 0 <= len(connected_apps_in_results) < len(apps_in_results):
+                result_parts.append(
+                    "Unconnected app functions can also be loaded if asked for by the user, they will generate a connection link"
+                    "but prefer connected ones. Ask the user to choose the app if none of the "
+                    "relevant apps are connected."
+                )
+
+            return "\n".join(result_parts)
+
+        def _structure_tool_results(
+            raw_tool_lists: list[list[dict[str, Any]]], connected_app_ids: set[str]
+        ) -> list[dict[str, Any]]:
+            """
+            Converts raw search results into a structured format, handling duplicates,
+            cleaning descriptions, and adding connection status.
+            """
             aggregated_tools = defaultdict(dict)
-            for tool_list in query_results:
+            # Use a list to maintain the order of apps as they are found.
+            ordered_app_ids = []
+
+            for tool_list in raw_tool_lists:
                 for tool in tool_list:
-                    app_id_from_tool = tool.get("app_id", "unknown")
+                    app_id = tool.get("app_id", "unknown")
                     tool_id = tool.get("id")
-                    if not tool_id or tool_id in aggregated_tools[app_id_from_tool]:
+
+                    if not tool_id:
                         continue
-                    cleaned_description = tool.get("description", "").split("Context:")[0].strip()
-                    aggregated_tools[app_id_from_tool][tool_id] = {
-                        "id": tool_id,
-                        "description": cleaned_description,
-                    }
-
-            # 6. Build the final results list, respecting the prioritized app order.
-            for app_id_from_list in prioritized_app_id_list:
-                if app_id_from_list in aggregated_tools and aggregated_tools[app_id_from_list]:
+
+                    if app_id not in aggregated_tools:
+                        ordered_app_ids.append(app_id)
+
+                    if tool_id not in aggregated_tools[app_id]:
+                        aggregated_tools[app_id][tool_id] = {
+                            "id": tool_id,
+                            "description": _clean_tool_description(tool.get("description", "")),
+                        }
+
+            # Build the final results list respecting the discovery order.
+            found_tools_result = []
+            for app_id in ordered_app_ids:
+                if app_id in aggregated_tools and aggregated_tools[app_id]:
                     found_tools_result.append(
                         {
-                            "app_id": app_id_from_list,
-                            "connection_status": "connected"
-                            if app_id_from_list in connected_app_ids
-                            else "not_connected",
-                            "tools": list(aggregated_tools[app_id_from_list].values()),
+                            "app_id": app_id,
+                            "connection_status": "connected" if app_id in connected_app_ids else "not_connected",
+                            "tools": list(aggregated_tools[app_id].values()),
                         }
                     )
+            return found_tools_result
 
-        # Build result string efficiently
-        result_parts = []
-        apps_in_results = {app["app_id"] for app in found_tools_result}
-        connected_apps_in_results = apps_in_results.intersection(connected_app_ids)
-
-        for app in found_tools_result:
-            app_id = app["app_id"]
-            connection_status = app["connection_status"]
-            tools = app["tools"]
-
-            app_status = "connected" if connection_status == "connected" else "NOT connected"
-            result_parts.append(f"Tools from {app_id} (status: {app_status} by user):")
-
-            for tool in tools:
-                tool_id = tool["id"]
-                description = tool["description"]
-                result_parts.append(f" - {tool_id}: {description}")
-            result_parts.append("")  # Empty line between apps
-
-        # Add connection status information
-        if len(connected_apps_in_results) == 0 and len(apps_in_results) > 0:
-            result_parts.append(
-                "Connection Status: None of the apps in the results are connected. You must ask the user to choose the application."
-            )
-        elif len(connected_apps_in_results) > 1:
-            connected_list = ", ".join(connected_apps_in_results)
-            result_parts.append(
-                f"Connection Status: Multiple apps are connected ({connected_list}). You must ask the user to select which application they want to use."
+        def _clean_tool_description(description: str) -> str:
+            """Consistently formats tool descriptions by removing implementation details."""
+            return description.split("Context:")[0].strip()
+
+        # Main Function Logic
+
+        if not queries and not app_ids:
+            raise ValueError("You must provide 'queries', 'app_ids', or both.")
+
+        # --- Initialization and Input Normalization ---
+        connections = await registry.list_connected_apps()
+        connected_app_ids = {connection["app_id"] for connection in connections}
+
+        canonical_app_ids = []
+        if app_ids:
+            # Concurrently search for all provided app names
+            app_search_tasks = [
+                registry.search_apps(query=app_name, distance_threshold=APP_THRESHOLD) for app_name in app_ids
+            ]
+            app_search_results = await asyncio.gather(*app_search_tasks)
+
+            # Process results and build the list of canonical IDs, handling not found errors
+            for app_name, result_list in zip(app_ids, app_search_results):
+                if not result_list:
+                    raise ValueError(f"Application '{app_name}' could not be found.")
+                # Assume the first result is the correct one
+                canonical_app_ids.append(result_list[0]["id"])
+
+        # --- Mode Dispatching ---
+        raw_results = []
+
+        if canonical_app_ids and queries:
+            raw_results = await _handle_scoped_search(canonical_app_ids, queries)
+        elif canonical_app_ids:
+            raw_results = await _handle_app_discovery(canonical_app_ids)
+        elif queries:
+            # Flatten list of lists to list of strings for global search
+            flat_queries = (
+                [q for sublist in queries for q in sublist] if queries and not isinstance(queries[0], str) else queries
             )
+            raw_results = await _handle_global_search(flat_queries)
 
-        result_parts.append("Call load_functions to select the required functions only.")
-        return "\n".join(result_parts)
+        # --- Structuring and Formatting ---
+        structured_data = _structure_tool_results(raw_results, connected_app_ids)
+        return _format_response(structured_data)
 
     @tool
     async def load_functions(tool_ids: list[str]) -> str:
-        """Load specific functions by their IDs for use in subsequent steps.
+        """
+        Loads specified functions and returns their Python signatures and docstrings.
+        This makes the functions available for use inside the 'execute_ipython_cell' tool.
+        The agent MUST use the returned information to understand how to call the functions correctly.
 
         Args:
-            tool_ids: Function ids in the form 'app__function'. Example: 'google_mail__send_email'
+            tool_ids: A list of function IDs in the format 'app__function'. Example: ['google_mail__send_email']
 
         Returns:
-            Confirmation message about loaded functions
+            A string containing the signatures and docstrings of the successfully loaded functions,
+            ready for the agent to use in its code.
         """
-        return f"Successfully loaded {len(tool_ids)} functions: {tool_ids}"
+        if not tool_ids:
+            return "No tool IDs provided to load."
+
+        # Step 1: Validate which tools are usable and get login links for others.
+        valid_tools, unconnected_links = await get_valid_tools(tool_ids=tool_ids, registry=tool_registry)
+
+        if not valid_tools:
+            response_string = "Error: None of the provided tool IDs could be validated or loaded."
+            return response_string, {}, [], ""
+
+        # Step 2: Export the schemas of the valid tools.
+        await tool_registry.load_tools(valid_tools)
+        exported_tools = await tool_registry.export_tools(
+            valid_tools, ToolFormat.NATIVE
+        )  # Get definition for only the new tools
+
+        # Step 3: Build the informational string for the agent.
+        tool_definitions, new_tools_context = build_tool_definitions(exported_tools)
+
+        result_parts = [
+            f"Successfully loaded {len(exported_tools)} functions. They are now available for use inside `execute_ipython_cell`:",
+            "\n".join(tool_definitions),
+        ]
+
+        response_string = "\n\n".join(result_parts)
+        unconnected_links = "\n".join(unconnected_links)
+
+        return response_string, new_tools_context, valid_tools, unconnected_links
 
-    @tool
     async def web_search(query: str) -> dict:
         """
         Get an LLM answer to a question informed by Exa search results. Useful when you need information from a wide range of real-time sources on the web. Do not use this when you need to access contents of a specific webpage.
@@ -234,7 +367,95 @@ def create_meta_tools(tool_registry: AgentrRegistry) -> dict[str, Any]:
             "citations": response.get("citations", []),
         }
 
-    return {"search_functions": search_functions, "load_functions": load_functions, "web_search": web_search}
+    async def read_file(uri: str) -> str:
+        """
+        Asynchronously reads a local file or uri and returns the content as a markdown string.
+
+        This tool aims to extract the main text content from various sources.
+        It automatically prepends 'file://' to the input string if it appears
+        to be a local path without a specified scheme (like http, https, data, file).
+
+        Args:
+            uri (str): The URI pointing to the resource or a local file path.
+                       Supported schemes:
+                       - http:// or https:// (Web pages, feeds, APIs)
+                       - file:// (Local or accessible network files)
+                       - data: (Embedded data)
+
+        Returns:
+            A string containing the markdown representation of the content at the specified URI
+
+        Raises:
+            ValueError: If the URI is invalid, empty, or uses an unsupported scheme
+                        after automatic prefixing.
+
+        Tags:
+            convert, markdown, async, uri, transform, document, important
+        """
+        markitdown = MarkitdownApp()
+        response = await markitdown.convert_to_markdown(uri)
+        return response
+
+    def save_file(file_name: str, content: str) -> dict:
+        """
+        Saves a file to the local filesystem.
+
+        Args:
+            file_name (str): The name of the file to save.
+            content (str): The content to save to the file.
+
+        Returns:
+            dict: A dictionary containing the result of the save operation with the following fields:
+                - status (str): "success" if the save succeeded, "error" otherwise.
+                - message (str): A message returned by the server, typically indicating success or providing error details.
+        """
+        with Path(file_name).open("w") as f:
+            f.write(content)
+
+        return {
+            "status": "success",
+            "message": f"File {file_name} saved successfully",
+            "file_path": Path(file_name).absolute(),
+        }
+
+    def upload_file(file_name: str, mime_type: str, base64_data: str) -> dict:
+        """
+        Uploads a file to the server via the AgentrClient.
+
+        Args:
+            file_name (str): The name of the file to upload.
+            mime_type (str): The MIME type of the file.
+            base64_data (str): The file content encoded as a base64 string.
+
+        Returns:
+            dict: A dictionary containing the result of the upload operation with the following fields:
+                - status (str): "success" if the upload succeeded, "error" otherwise.
+                - message (str): A message returned by the server, typically indicating success or providing error details.
+                - signed_url (str or None): The signed URL to access the uploaded file if successful, None otherwise.
+        """
+        client: AgentrClient = tool_registry.client
+        bytes_data = base64.b64decode(base64_data)
+        response = client._upload_file(file_name, mime_type, bytes_data)
+        if response.get("status") != "success":
+            return {
+                "status": "error",
+                "message": response.get("message"),
+                "signed_url": None,
+            }
+        return {
+            "status": "success",
+            "message": response.get("message"),
+            "signed_url": response.get("signed_url"),
+        }
+
+    return {
+        "search_functions": search_functions,
+        "load_functions": load_functions,
+        "web_search": web_search,
+        "read_file": read_file,
+        "upload_file": upload_file,
+        "save_file": save_file,
+    }
 
 
 async def get_valid_tools(tool_ids: list[str], registry: AgentrRegistry) -> tuple[list[str], list[str]]:
@@ -282,7 +503,7 @@ async def get_valid_tools(tool_ids: list[str], registry: AgentrRegistry) -> tupl
             start = text.find(":") + 1
             end = text.find(". R", start)
             url = text[start:end].strip()
-            markdown_link = f"[{app}]({url})"
+            markdown_link = f"![Connect to {app.capitalize()}]({url})"
             unconnected_links.append(markdown_link)
         for tool_id, tool_name in tool_entries:
             if tool_name in available: