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

universal_mcp/agents/codeact0/state.py

@@ -1,12 +1,36 @@
-from typing import Any
+from typing import Annotated, Any
 
-from langgraph.graph import MessagesState
+from langgraph.prebuilt.chat_agent_executor import AgentState
 
 
-class CodeActState(MessagesState):
+def _enqueue(left: list, right: list) -> list:
+    """Treat left as a FIFO queue, append new items from right (preserve order),
+    keep items unique, and cap total size to 20 (drop oldest items)."""
+    max_size = 30
+    preferred_size = 20
+    if len(right) > preferred_size:
+        preferred_size = min(max_size, len(right))
+    queue = list(left or [])
+
+    for item in right[:preferred_size] or []:
+        if item in queue:
+            queue.remove(item)
+        queue.append(item)
+
+    if len(queue) > preferred_size:
+        queue = queue[-preferred_size:]
+
+    return queue
+
+
+class CodeActState(AgentState):
     """State for CodeAct agent."""
 
     context: dict[str, Any]
     """Dictionary containing the execution context with available tools and variables."""
     add_context: dict[str, Any]
     """Dictionary containing the additional context (functions, classes, imports) to be added to the execution context."""
+    playbook_mode: str | None
+    """State for the playbook agent."""
+    selected_tool_ids: Annotated[list[str], _enqueue]
+    """Queue for tools exported from registry"""

universal_mcp/agents/codeact0/tools.py

@@ -0,0 +1,180 @@
+import asyncio
+from collections import defaultdict
+from typing import Any
+
+from langchain_core.tools import tool
+from universal_mcp.tools.registry import ToolRegistry
+
+MAX_LENGHT=100
+
+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
+
+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."""
+    return
+
+
+
+def create_meta_tools(tool_registry: ToolRegistry) -> dict[str, Any]:
+    """Create the meta tools for searching and loading tools"""
+
+    @tool
+    async def search_functions(queries: list[str]) -> str:
+        """Search for relevant functions given list of queries.
+        Each single query should be atomic (doable with a single function).
+        For tasks requiring multiple functions, add separate queries for each subtask"""
+        try:
+            # Fetch all connections
+            connections = await tool_registry.list_connected_apps()
+            connected_apps = {connection["app_id"] for connection in connections}
+
+            # Use defaultdict to avoid key existence checks
+            app_tools = defaultdict(list)
+
+            # Process all queries concurrently
+            search_tasks = []
+            for query in queries:
+                search_tasks.append(_search_query_tools(query))
+
+            query_results = await asyncio.gather(*search_tasks)
+
+            # Aggregate results with limit per app
+            for tools_list in query_results:
+                for tool in tools_list:
+                    app = tool["id"].split("__")[0]
+                    if len(app_tools[app]) < MAX_LENGHT:
+                        cleaned_desc = tool["description"].split("Context:")[0].strip()
+                        app_tools[app].append(f"{tool['id']}: {cleaned_desc}")
+
+            # Build result string efficiently
+            result_parts = []
+            for app, tools in app_tools.items():
+                app_status = "connected" if app in connected_apps else "NOT connected"
+                result_parts.append(f"Tools from {app} (status: {app_status} by user):")
+                for tool in tools:
+                    result_parts.append(f" - {tool}")
+                result_parts.append("")  # Empty line between apps
+
+            result_parts.append("Call load_functions to select the required functions only.")
+            return "\n".join(result_parts)
+
+        except Exception as e:
+            return f"Error: {e}"
+
+    async def _search_query_tools(query: str) -> list[dict]:
+        """Helper function to search apps and tools for a single query."""
+        # Start both searches concurrently
+        tools_search_task = tool_registry.search_tools(query, limit=10)
+        apps_search_task = tool_registry.search_apps(query, limit=4)
+
+        # Wait for both to complete
+        tools_from_general_search, apps_list = await asyncio.gather(tools_search_task, apps_search_task)
+
+        # Create tasks for searching tools from each app
+        app_tool_tasks = [tool_registry.search_tools(query, limit=5, app_id=app["id"]) for app in apps_list]
+
+        # Wait for all app-specific tool searches to complete
+        app_tools_results = await asyncio.gather(*app_tool_tasks)
+
+        # Combine all results
+        tools_list = list(tools_from_general_search)
+        for app_tools in app_tools_results:
+            tools_list.extend(app_tools)
+
+        return tools_list
+
+    @tool
+    async def load_functions(tool_ids: list[str]) -> str:
+        """Load specific functions by their IDs for use in subsequent steps.
+
+        Args:
+            tool_ids: Function ids in the form 'app__function'. Example: 'google_mail__send_email'
+
+        Returns:
+            Confirmation message about loaded functions
+        """
+        return f"Successfully loaded {len(tool_ids)} functions: {tool_ids}"
+
+    @tool
+    async def web_search(query: str) -> list:
+        """Search the web for the given query and return structured search results.
+
+        Do not use for app-specific searches (for example, reddit or linkedin searches
+
+        Returns:
+            list: A list of up to 10 search result dictionaries, each containing:
+                - id (str): Unique identifier, typically the URL
+                - title (str): The title/headline of the search result
+                - url (str): The web URL of the result
+                - publishedDate (str): ISO 8601 formatted date (e.g., "2025-01-01T00:00:00.000Z")
+                - author (str): Author name (may be empty string)
+                - summary (str): Text summary/snippet of the content
+                - image (str): URL to associated image (if available)
+
+        Example:
+            results = await web_search(query="python programming")
+        """
+        response = await tool_registry.call_tool(
+            "exa__search_with_filters", {"query": query, "contents": {"summary": True}}
+        )
+        return response["results"]
+
+    return {"search_functions": search_functions, "load_functions": load_functions, "web_search": web_search}
+
+
+async def get_valid_tools(tool_ids: list[str], registry: ToolRegistry) -> tuple[list[str], list[str]]:
+    """For a given list of tool_ids, validates the tools and returns a list of links for the apps that have not been logged in"""
+    correct, incorrect = [], []
+    connections = await registry.list_connected_apps()
+    connected_apps = {connection["app_id"] for connection in connections}
+    unconnected = set()
+    unconnected_links = []
+    app_tool_list: dict[str, set[str]] = {}
+
+    # Group tool_ids by app for fewer registry calls
+    app_to_tools: dict[str, list[tuple[str, str]]] = {}
+    for tool_id in tool_ids:
+        if "__" not in tool_id:
+            incorrect.append(tool_id)
+            continue
+        app, tool_name = tool_id.split("__", 1)
+        app_to_tools.setdefault(app, []).append((tool_id, tool_name))
+
+    # Fetch all apps concurrently
+    async def fetch_tools(app: str):
+        try:
+            tools_dict = await registry.list_tools(app)
+            return app, {tool_unit["name"] for tool_unit in tools_dict}
+        except Exception:
+            return app, None
+
+    results = await asyncio.gather(*(fetch_tools(app) for app in app_to_tools))
+
+    # Build map of available tools per app
+    for app, tools in results:
+        if tools is not None:
+            app_tool_list[app] = tools
+
+    # Validate tool_ids
+    for app, tool_entries in app_to_tools.items():
+        available = app_tool_list.get(app)
+        if available is None:
+            incorrect.extend(tool_id for tool_id, _ in tool_entries)
+            continue
+        if app not in connected_apps and app not in unconnected:
+            unconnected.add(app)
+            text = registry.client.get_authorization_url(app)
+            start = text.find(":") + 1
+            end = text.find(". R", start)
+            url = text[start:end].strip()
+            markdown_link = f"[{app}]({url})"
+            unconnected_links.append(markdown_link)
+        for tool_id, tool_name in tool_entries:
+            if tool_name in available:
+                correct.append(tool_id)
+            else:
+                incorrect.append(tool_id)
+
+    return correct, unconnected_links

universal_mcp/agents/codeact0/utils.py

@@ -1,10 +1,12 @@
+import ast
+import importlib
 import re
 from collections.abc import Sequence
 from typing import Any
 
 from langchain_core.messages import BaseMessage
-from pydantic import ValidationError
-from requests import JSONDecodeError
+
+MAX_CHARS = 300
 
 
 def light_copy(data):
@@ -15,13 +17,13 @@ def light_copy(data):
         data: Either a dictionary with string keys, or a sequence of such dictionaries
 
     Returns:
-        A deep copy where all string values are truncated to 30 characters
+        A deep copy where all string values are truncated to MAX_CHARS characters
     """
 
     def truncate_string(value):
-        """Truncate string to 30 chars, preserve other types"""
-        if isinstance(value, str) and len(value) > 30:
-            return value[:30] + "..."
+        """Truncate string to MAX_CHARS chars, preserve other types"""
+        if isinstance(value, str) and len(value) > MAX_CHARS:
+            return value[:MAX_CHARS] + "..."
         return value
 
     def copy_dict(d):
@@ -75,44 +77,6 @@ def make_safe_function_name(name: str) -> str:
     return safe_name
 
 
-def filter_retry_on(exc: Exception) -> bool:
-    import httpx
-    import requests
-
-    if isinstance(
-        exc,
-        (
-            ConnectionError,
-            JSONDecodeError,
-            ValidationError,
-        ),
-    ):
-        return True
-    if isinstance(
-        exc,
-        (
-            ValueError,
-            TypeError,
-            ArithmeticError,
-            ImportError,
-            LookupError,
-            NameError,
-            SyntaxError,
-            RuntimeError,
-            ReferenceError,
-            StopIteration,
-            StopAsyncIteration,
-            OSError,
-        ),
-    ):
-        return False
-    if isinstance(exc, httpx.HTTPStatusError):
-        return 500 <= exc.response.status_code < 600
-    if isinstance(exc, requests.HTTPError):
-        return 500 <= exc.response.status_code < 600 if exc.response else True
-    return True
-
-
 def derive_context(code: str, context: dict[str, Any]) -> dict[str, Any]:
     """
     Derive context from code by extracting classes, functions, and import statements.
@@ -124,8 +88,6 @@ def derive_context(code: str, context: dict[str, Any]) -> dict[str, Any]:
     Returns:
         Updated context dictionary with extracted entities
     """
-    import ast
-    import re
 
     # Initialize context keys if they don't exist
     if "imports" not in context:
@@ -177,19 +139,15 @@ def derive_context(code: str, context: dict[str, Any]) -> dict[str, Any]:
                 if class_def not in context["classes"]:
                     context["classes"].append(class_def)
 
-        # Extract function definitions (only top-level functions, not class methods)
+        # Extract function definitions (including async)
         for node in ast.walk(tree):
-            if isinstance(node, ast.FunctionDef):
-                # Get the function definition as a string
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                 func_lines = code.split("\n")[node.lineno - 1 : node.end_lineno]
                 func_def = "\n".join(func_lines)
 
-                # Check if this is a top-level function by looking at indentation
-                # Top-level functions should start at column 0 (no indentation)
+                # Only top-level functions (col_offset == 0)
                 if node.col_offset == 0:
-                    # Clean up the function definition (remove leading/trailing whitespace)
                     func_def = func_def.strip()
-
                     if func_def not in context["functions"]:
                         context["functions"].append(func_def)
 
@@ -208,8 +166,8 @@ def derive_context(code: str, context: dict[str, Any]) -> dict[str, Any]:
                 if "from" in pattern:
                     module = match.group(1)
                     imports = match.group(2).split(",")
-                    for imp in imports:
-                        imp = imp.strip()
+                    for import_name in imports:
+                        imp = import_name.strip()
                         if " as " in imp:
                             name, alias = imp.split(" as ")
                             import_stmt = f"from {module} import {name.strip()} as {alias.strip()}"
@@ -278,8 +236,6 @@ def inject_context(
         namespace = inject_context(context, existing_ns)
         # namespace will contain: {'math': <math module>, 'data': [1, 2, 3], 'pandas': <module>, 'pd': <module>, 'numpy': <module>, 'np': <module>, 'MyClass': <class>, 'MC': <class>, 'my_function': <function>, ...}
     """
-    import importlib
-    from typing import Any
 
     # Start with existing namespace or create new one
     namespace: dict[str, Any] = existing_namespace.copy() if existing_namespace is not None else {}
@@ -292,7 +248,6 @@ def inject_context(
                 exec(import_statement, namespace)
             except Exception as e:
                 # If execution fails, try to extract module name and create placeholder
-                import re
 
                 # Handle different import patterns
                 import_match = re.search(r"import\s+(\w+)(?:\s+as\s+(\w+))?", import_statement)
@@ -323,7 +278,6 @@ def inject_context(
                 exec(class_definition, namespace)
             except Exception:
                 # If execution fails, try to extract class name and create placeholder
-                import re
 
                 class_match = re.search(r"class\s+(\w+)", class_definition)
                 if class_match:
@@ -332,14 +286,14 @@ def inject_context(
                     # Create a placeholder class
                     class PlaceholderClass:
                         def __init__(self, *args, **kwargs):
-                            raise NotImplementedError(f"Class '{class_name}' failed to load: {str(e)}")
+                            raise NotImplementedError("Class '{class_name}' failed to load")
 
                     namespace[class_name] = PlaceholderClass
                 else:
                     # If we can't extract class name, create a generic placeholder
                     class GenericPlaceholderClass:
                         def __init__(self, *args, **kwargs):
-                            raise NotImplementedError(f"Class definition failed to load: {str(e)}")
+                            raise NotImplementedError("Class definition failed to load")
 
                     namespace[f"class_{len(namespace)}"] = GenericPlaceholderClass
 
@@ -351,24 +305,84 @@ def inject_context(
                 exec(function_definition, namespace)
             except Exception:
                 # If execution fails, try to extract function name and create placeholder
-                import re
-
-                func_match = re.search(r"def\s+(\w+)", function_definition)
+                func_match = re.search(r"(async\s+)?def\s+(\w+)", function_definition)
                 if func_match:
-                    func_name = func_match.group(1)
+                    func_name = func_match.group(2)
+                    is_async = bool(func_match.group(1))
+
+                    if is_async:
+
+                        async def placeholder_func(*args, **kwargs):
+                            raise NotImplementedError(f"Async function '{func_name}' failed to load")
+                    else:
 
-                    # Create a placeholder function
-                    def placeholder_func(*args, **kwargs):
-                        raise NotImplementedError(f"Function '{func_name}' failed to load: {str(e)}")
+                        def placeholder_func(*args, **kwargs):
+                            raise NotImplementedError(f"Function '{func_name}' failed to load")
 
                     placeholder_func.__name__ = func_name
                     namespace[func_name] = placeholder_func
-                else:
-                    # If we can't extract function name, create a generic placeholder
-                    def generic_placeholder_func(*args, **kwargs):
-                        raise NotImplementedError(f"Function definition failed to load: {str(e)}")
-
-                    generic_placeholder_func.__name__ = f"func_{len(namespace)}"
-                    namespace[generic_placeholder_func.__name__] = generic_placeholder_func
 
     return namespace
+
+
+def schema_to_signature(schema: dict, func_name="my_function") -> str:
+    type_map = {
+        "integer": "int",
+        "string": "str",
+        "boolean": "bool",
+        "null": "None",
+    }
+
+    params = []
+    for name, meta in schema.items():
+        # figure out type
+        if "type" in meta:
+            typ = type_map.get(meta["type"], "Any")
+        elif "anyOf" in meta:
+            types = [type_map.get(t["type"], "Any") for t in meta["anyOf"]]
+            typ = " | ".join(set(types))
+        else:
+            typ = "Any"
+
+        default = meta.get("default", None)
+        default_repr = repr(default)
+
+        params.append(f"{name}: {typ} = {default_repr}")
+
+    # join into signature
+    param_str = ",\n    ".join(params)
+    return f"def {func_name}(\n    {param_str},\n):"
+
+
+def smart_truncate(
+    output: str, max_chars_full: int = 2000, max_lines_headtail: int = 20, summary_threshold: int = 10000
+) -> str:
+    """
+    Truncates or summarizes output intelligently to avoid filling the context too fast.
+
+    Args:
+        output (str): The string output from code execution.
+        max_chars_full (int): Max characters to keep full output.
+        max_lines_headtail (int): Number of lines to keep from head and tail for medium outputs.
+        summary_threshold (int): If truncated output exceeds this, hard-truncate.
+
+    Returns:
+        str: Truncated or summarized output.
+    """
+    if len(output) <= max_chars_full:
+        return output  # Small output, include fully
+
+    lines = output.splitlines()
+    if len(lines) <= 2 * max_lines_headtail:
+        return output  # Medium output, include fully
+
+    # Medium-large output: take head + tail
+    head = "\n".join(lines[:max_lines_headtail])
+    tail = "\n".join(lines[-max_lines_headtail:])
+    truncated = f"{head}\n... [truncated {len(lines) - 2 * max_lines_headtail} lines] ...\n{tail}"
+
+    # If still too big, cut to summary threshold
+    if len(truncated) > summary_threshold:
+        truncated = truncated[:summary_threshold] + "\n... [output truncated to fit context] ..."
+
+    return truncated

universal_mcp/agents/shared/__main__.py

@@ -0,0 +1,44 @@
+import asyncio
+
+from rich import print
+from universal_mcp.agentr.registry import AgentrRegistry
+from universal_mcp.logger import setup_logger
+
+from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.agents.shared.tool_node import build_tool_node_graph
+
+
+async def main():
+    """
+    An example of how to run the tool_node graph independently.
+    """
+    setup_logger(level="INFO")
+
+    user_input = "What are the topics of my meetings today from Google Calendar and who are the attendees? Give a 1-line context for each attendee using LinkedIn or web search."
+
+    print(f"▶️  User Task: [bold cyan]'{user_input}'[/bold cyan]\n")
+
+    llm = load_chat_model("azure/gpt-4.1", thinking=False)
+    registry = AgentrRegistry()
+
+    graph = build_tool_node_graph(llm=llm, registry=registry)
+
+    initial_state = {"original_task": user_input}
+
+    print("🚀 Invoking the tool selection graph...")
+    final_state = await graph.ainvoke(initial_state)
+
+    execution_plan = final_state.get("execution_plan")
+
+    print("\n[bold green]✅ Graph execution complete![/bold green]")
+    print("\n--- Final Execution Plan (Selected Tools) ---")
+    if execution_plan:
+        print(execution_plan)
+    else:
+        print("[bold red]No execution plan was created.[/bold red]")
+        if messages := final_state.get("messages"):
+            print(f"Final Message: {messages[-1].content}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())