dao-ai 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

dao_ai/cli.py

@@ -7,7 +7,7 @@ import sys
 import traceback
 from argparse import ArgumentParser, Namespace
 from pathlib import Path
-from typing import Optional, Sequence
+from typing import Any, Optional, Sequence
 
 from dotenv import find_dotenv, load_dotenv
 from loguru import logger
@@ -114,6 +114,7 @@ Examples:
   dao-ai validate -c config/model_config.yaml            # Validate a specific configuration file
   dao-ai graph -o architecture.png -c my_config.yaml -v  # Generate visual graph with verbose output
   dao-ai chat -c config/retail.yaml --custom-input store_num=87887  # Start interactive chat session
+  dao-ai list-mcp-tools -c config/mcp_config.yaml --apply-filters  # List filtered MCP tools only
   dao-ai validate                                        # Validate with detailed logging
   dao-ai bundle --deploy                                 # Deploy the DAO AI asset bundle
         """,
@@ -309,6 +310,53 @@ Examples:
         help="Path to the model configuration file to validate",
     )
 
+    # List MCP tools command
+    list_mcp_parser: ArgumentParser = subparsers.add_parser(
+        "list-mcp-tools",
+        help="List available MCP tools from configuration",
+        description="""
+List all available MCP tools from the configured MCP servers.
+This command shows:
+- All MCP servers/functions in the configuration
+- Available tools from each server
+- Full descriptions for each tool (no truncation)
+- Tool parameters in readable format (type, required/optional, descriptions)
+- Which tools are included/excluded based on filters
+- Filter patterns (include_tools, exclude_tools)
+
+Use this command to:
+- Discover available tools before configuring agents
+- Review tool descriptions and parameter schemas
+- Debug tool filtering configuration
+- Verify MCP server connectivity
+
+Options:
+- Use --apply-filters to only show tools that will be loaded (hides excluded tools)
+- Without --apply-filters, see all available tools with include/exclude status
+
+Note: Schemas are displayed in a concise, readable format instead of verbose JSON
+        """,
+        epilog="""Examples:
+  dao-ai list-mcp-tools -c config/model_config.yaml
+  dao-ai list-mcp-tools -c config/model_config.yaml --apply-filters
+        """,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    list_mcp_parser.add_argument(
+        "-c",
+        "--config",
+        type=str,
+        default="./config/model_config.yaml",
+        required=False,
+        metavar="FILE",
+        help="Path to the model configuration file (default: ./config/model_config.yaml)",
+    )
+    list_mcp_parser.add_argument(
+        "--apply-filters",
+        action="store_true",
+        help="Only show tools that pass include/exclude filters (hide excluded tools)",
+    )
+
     chat_parser: ArgumentParser = subparsers.add_parser(
         "chat",
         help="Interactive chat with the DAO AI system",
@@ -704,6 +752,275 @@ def handle_validate_command(options: Namespace) -> None:
         sys.exit(1)
 
 
+def _format_schema_pretty(schema: dict[str, Any], indent: int = 0) -> str:
+    """
+    Format a JSON schema in a more readable, concise format.
+
+    Args:
+        schema: The JSON schema to format
+        indent: Current indentation level
+
+    Returns:
+        Pretty-formatted schema string
+    """
+    if not schema:
+        return ""
+
+    lines: list[str] = []
+    indent_str = "  " * indent
+
+    # Get required fields
+    required_fields = set(schema.get("required", []))
+
+    # Handle object type with properties
+    if schema.get("type") == "object" and "properties" in schema:
+        properties = schema["properties"]
+
+        for prop_name, prop_schema in properties.items():
+            is_required = prop_name in required_fields
+            req_marker = " (required)" if is_required else " (optional)"
+
+            prop_type = prop_schema.get("type", "any")
+            prop_desc = prop_schema.get("description", "")
+
+            # Handle different types
+            if prop_type == "array":
+                items = prop_schema.get("items", {})
+                item_type = items.get("type", "any")
+                type_str = f"array<{item_type}>"
+            elif prop_type == "object":
+                type_str = "object"
+            else:
+                type_str = prop_type
+
+            # Format enum values if present
+            if "enum" in prop_schema:
+                enum_values = ", ".join(str(v) for v in prop_schema["enum"])
+                type_str = f"{type_str} (one of: {enum_values})"
+
+            # Build the line
+            line = f"{indent_str}{prop_name}: {type_str}{req_marker}"
+            if prop_desc:
+                line += f"\n{indent_str}  └─ {prop_desc}"
+
+            lines.append(line)
+
+            # Recursively handle nested objects
+            if prop_type == "object" and "properties" in prop_schema:
+                nested = _format_schema_pretty(prop_schema, indent + 1)
+                if nested:
+                    lines.append(nested)
+
+    # Handle simple types without properties
+    elif "type" in schema:
+        schema_type = schema["type"]
+        if schema.get("description"):
+            lines.append(f"{indent_str}Type: {schema_type}")
+            lines.append(f"{indent_str}└─ {schema['description']}")
+        else:
+            lines.append(f"{indent_str}Type: {schema_type}")
+
+    return "\n".join(lines)
+
+
+def handle_list_mcp_tools_command(options: Namespace) -> None:
+    """
+    List available MCP tools from configuration.
+
+    Shows all MCP servers and their available tools, indicating which
+    are included/excluded based on filter configuration.
+    """
+    logger.debug(f"Listing MCP tools from configuration: {options.config}")
+
+    try:
+        from dao_ai.config import McpFunctionModel
+        from dao_ai.tools.mcp import MCPToolInfo, _matches_pattern, list_mcp_tools
+
+        # Load configuration
+        config: AppConfig = AppConfig.from_file(options.config)
+
+        # Find all MCP tools in configuration
+        mcp_tools_config: list[tuple[str, McpFunctionModel]] = []
+        if config.tools:
+            for tool_name, tool_model in config.tools.items():
+                logger.debug(
+                    f"Checking tool: {tool_name}, function type: {type(tool_model.function)}"
+                )
+                if tool_model.function and isinstance(
+                    tool_model.function, McpFunctionModel
+                ):
+                    mcp_tools_config.append((tool_name, tool_model.function))
+
+        if not mcp_tools_config:
+            logger.warning("No MCP tools found in configuration")
+            print("\n⚠️  No MCP tools configured in this file.")
+            print(f"   Configuration: {options.config}")
+            print(
+                "\nTo add MCP tools, define them in the 'tools' section with 'type: mcp'"
+            )
+            sys.exit(0)
+
+        # Collect all results first (aggregate before displaying)
+        results: list[dict[str, Any]] = []
+        for tool_name, mcp_function in mcp_tools_config:
+            result = {
+                "tool_name": tool_name,
+                "mcp_function": mcp_function,
+                "error": None,
+                "all_tools": [],
+                "included_tools": [],
+                "excluded_tools": [],
+            }
+
+            try:
+                logger.info(f"Connecting to MCP server: {mcp_function.mcp_url}")
+
+                # Get all available tools (unfiltered)
+                all_tools: list[MCPToolInfo] = list_mcp_tools(
+                    mcp_function, apply_filters=False
+                )
+
+                # Get filtered tools (what will actually be loaded)
+                filtered_tools: list[MCPToolInfo] = list_mcp_tools(
+                    mcp_function, apply_filters=True
+                )
+
+                included_names = {t.name for t in filtered_tools}
+
+                # Categorize tools
+                for tool in sorted(all_tools, key=lambda t: t.name):
+                    if tool.name in included_names:
+                        result["included_tools"].append(tool)
+                    else:
+                        # Determine why it was excluded
+                        reason = ""
+                        if mcp_function.exclude_tools:
+                            if _matches_pattern(tool.name, mcp_function.exclude_tools):
+                                matching_patterns = [
+                                    p
+                                    for p in mcp_function.exclude_tools
+                                    if _matches_pattern(tool.name, [p])
+                                ]
+                                reason = f" (matches exclude pattern: {', '.join(matching_patterns)})"
+                        if not reason and mcp_function.include_tools:
+                            reason = " (not in include list)"
+                        result["excluded_tools"].append((tool, reason))
+
+                result["all_tools"] = all_tools
+
+            except KeyboardInterrupt:
+                result["error"] = "Connection interrupted by user"
+                results.append(result)
+                break
+            except Exception as e:
+                logger.error(f"Failed to list tools from MCP server: {e}")
+                result["error"] = str(e)
+
+            results.append(result)
+
+        # Now display all results at once (no logging interleaving)
+        print(f"\n{'=' * 80}")
+        print("MCP TOOLS DISCOVERY")
+        print(f"Configuration: {options.config}")
+        print(f"{'=' * 80}\n")
+
+        for result in results:
+            tool_name = result["tool_name"]
+            mcp_function = result["mcp_function"]
+
+            print(f"📦 Tool: {tool_name}")
+            print(f"   Server: {mcp_function.mcp_url}")
+
+            # Show connection type
+            if mcp_function.connection:
+                print(f"   Connection: UC Connection '{mcp_function.connection.name}'")
+            else:
+                print(f"   Transport: {mcp_function.transport.value}")
+
+            # Show filters if configured
+            if mcp_function.include_tools or mcp_function.exclude_tools:
+                print("\n   Filters:")
+                if mcp_function.include_tools:
+                    print(f"     Include: {', '.join(mcp_function.include_tools)}")
+                if mcp_function.exclude_tools:
+                    print(f"     Exclude: {', '.join(mcp_function.exclude_tools)}")
+
+            # Check for errors
+            if result["error"]:
+                print(f"\n   ❌ Error: {result['error']}")
+                print("   Could not connect to MCP server")
+                if result["error"] != "Connection interrupted by user":
+                    print(
+                        "   Tip: Verify server URL, authentication, and network connectivity"
+                    )
+            else:
+                all_tools = result["all_tools"]
+                included_tools = result["included_tools"]
+                excluded_tools = result["excluded_tools"]
+
+                # Show stats based on --apply-filters flag
+                if options.apply_filters:
+                    # Simplified view: only show filtered tools count
+                    print(
+                        f"\n   Available Tools: {len(included_tools)} (after filters)"
+                    )
+                else:
+                    # Full view: show all, included, and excluded counts
+                    print(f"\n   Available Tools: {len(all_tools)} total")
+                    print(f"   ├─ ✓ Included: {len(included_tools)}")
+                    print(f"   └─ ✗ Excluded: {len(excluded_tools)}")
+
+                # Show included tools with FULL descriptions and schemas
+                if included_tools:
+                    if options.apply_filters:
+                        print(f"\n   Tools ({len(included_tools)}):")
+                    else:
+                        print(f"\n   ✓ Included Tools ({len(included_tools)}):")
+
+                    for tool in included_tools:
+                        print(f"\n     • {tool.name}")
+                        if tool.description:
+                            # Show full description (no truncation)
+                            print(f"       Description: {tool.description}")
+                        if tool.input_schema:
+                            # Pretty print schema in readable format
+                            print("       Parameters:")
+                            pretty_schema = _format_schema_pretty(
+                                tool.input_schema, indent=0
+                            )
+                            if pretty_schema:
+                                # Indent the schema for better readability
+                                for line in pretty_schema.split("\n"):
+                                    print(f"         {line}")
+                            else:
+                                print("         (none)")
+
+                # Show excluded tools only if NOT applying filters
+                if excluded_tools and not options.apply_filters:
+                    print(f"\n   ✗ Excluded Tools ({len(excluded_tools)}):")
+                    for tool, reason in excluded_tools:
+                        print(f"     • {tool.name}{reason}")
+
+            print(f"\n{'-' * 80}\n")
+
+        # Summary
+        print(f"{'=' * 80}")
+        print(f"Summary: Found {len(mcp_tools_config)} MCP server(s)")
+        print(f"{'=' * 80}\n")
+
+        sys.exit(0)
+
+    except FileNotFoundError:
+        logger.error(f"Configuration file not found: {options.config}")
+        print(f"\n❌ Error: Configuration file not found: {options.config}")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"Failed to list MCP tools: {e}")
+        logger.debug(traceback.format_exc())
+        print(f"\n❌ Error: {e}")
+        sys.exit(1)
+
+
 def setup_logging(verbosity: int) -> None:
     levels: dict[int, str] = {
         0: "ERROR",
@@ -925,6 +1242,8 @@ def main() -> None:
             handle_deploy_command(options)
         case "chat":
             handle_chat_command(options)
+        case "list-mcp-tools":
+            handle_list_mcp_tools_command(options)
         case _:
             logger.error(f"Unknown command: {options.command}")
             sys.exit(1)

dao_ai/config.py

@@ -24,6 +24,7 @@ from databricks.sdk.credentials_provider import (
     ModelServingUserCredentials,
 )
 from databricks.sdk.errors.platform import NotFound
+from databricks.sdk.service.apps import App
 from databricks.sdk.service.catalog import FunctionInfo, TableInfo
 from databricks.sdk.service.dashboards import GenieSpace
 from databricks.sdk.service.database import DatabaseInstance
@@ -147,7 +148,7 @@ class PrimitiveVariableModel(BaseModel, HasValue):
         return str(value)
 
     @model_validator(mode="after")
-    def validate_value(self) -> "PrimitiveVariableModel":
+    def validate_value(self) -> Self:
         if not isinstance(self.as_value(), (str, int, float, bool)):
             raise ValueError("Value must be a primitive type (str, int, float, bool)")
         return self
@@ -391,10 +392,17 @@ class PermissionModel(BaseModel):
 
 class SchemaModel(BaseModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    catalog_name: str
-    schema_name: str
+    catalog_name: AnyVariable
+    schema_name: AnyVariable
     permissions: Optional[list[PermissionModel]] = Field(default_factory=list)
 
+    @model_validator(mode="after")
+    def resolve_variables(self) -> Self:
+        """Resolve AnyVariable fields to their actual string values."""
+        self.catalog_name = value_of(self.catalog_name)
+        self.schema_name = value_of(self.schema_name)
+        return self
+
     @property
     def full_name(self) -> str:
         return f"{self.catalog_name}.{self.schema_name}"
@@ -408,9 +416,45 @@ class SchemaModel(BaseModel, HasFullName):
 
 
 class DatabricksAppModel(IsDatabricksResource, HasFullName):
+    """
+    Configuration for a Databricks App resource.
+    
+    The `name` is the unique instance name of the Databricks App within the workspace.
+    The `url` is dynamically retrieved from the workspace client by calling 
+    `apps.get(name)` and returning the app's URL.
+    
+    Example:
+        ```yaml
+        resources:
+          apps:
+            my_app:
+              name: my-databricks-app
+        ```
+    """
+    
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
-    url: str
+    """The unique instance name of the Databricks App in the workspace."""
+
+    @property
+    def url(self) -> str:
+        """
+        Retrieve the URL of the Databricks App from the workspace.
+        
+        Returns:
+            The URL of the deployed Databricks App.
+            
+        Raises:
+            RuntimeError: If the app is not found or URL is not available.
+        """
+        app: App = self.workspace_client.apps.get(self.name)
+        if app.url is None:
+            raise RuntimeError(
+                f"Databricks App '{self.name}' does not have a URL. "
+                "The app may not be deployed yet."
+            )
+        return app.url
+
 
     @property
     def full_name(self) -> str:
@@ -432,7 +476,7 @@ class TableModel(IsDatabricksResource, HasFullName):
     name: Optional[str] = None
 
     @model_validator(mode="after")
-    def validate_name_or_schema_required(self) -> "TableModel":
+    def validate_name_or_schema_required(self) -> Self:
         if not self.name and not self.schema_model:
             raise ValueError(
                 "Either 'name' or 'schema_model' must be provided for TableModel"
@@ -998,7 +1042,7 @@ class VolumePathModel(BaseModel, HasFullName):
     path: Optional[str] = None
 
     @model_validator(mode="after")
-    def validate_path_or_volume(self) -> "VolumePathModel":
+    def validate_path_or_volume(self) -> Self:
         if not self.volume and not self.path:
             raise ValueError("Either 'volume' or 'path' must be provided")
         return self
@@ -1840,11 +1884,32 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
     headers: dict[str, AnyVariable] = Field(default_factory=dict)
     args: list[str] = Field(default_factory=list)
     # MCP-specific fields
+    app: Optional[DatabricksAppModel] = None
     connection: Optional[ConnectionModel] = None
     functions: Optional[SchemaModel] = None
     genie_room: Optional[GenieRoomModel] = None
     sql: Optional[bool] = None
     vector_search: Optional[VectorStoreModel] = None
+    # Tool filtering
+    include_tools: Optional[list[str]] = Field(
+        default=None,
+        description=(
+            "Optional list of tool names or glob patterns to include from the MCP server. "
+            "If specified, only tools matching these patterns will be loaded. "
+            "Supports glob patterns: * (any chars), ? (single char), [abc] (char set). "
+            "Examples: ['execute_query', 'list_*', 'get_?_data']"
+        ),
+    )
+    exclude_tools: Optional[list[str]] = Field(
+        default=None,
+        description=(
+            "Optional list of tool names or glob patterns to exclude from the MCP server. "
+            "Tools matching these patterns will not be loaded. "
+            "Takes precedence over include_tools. "
+            "Supports glob patterns: * (any chars), ? (single char), [abc] (char set). "
+            "Examples: ['drop_*', 'delete_*', 'execute_ddl']"
+        ),
+    )
 
     @property
     def api_scopes(self) -> Sequence[str]:
@@ -1907,6 +1972,7 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
 
         Returns the URL based on the configured source:
         - If url is set, returns it directly
+        - If app is set, retrieves URL from Databricks App via workspace client
         - If connection is set, constructs URL from connection
         - If genie_room is set, constructs Genie MCP URL
         - If sql is set, constructs DBSQL MCP URL (serverless)
@@ -1919,6 +1985,7 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
         - Vector Search: https://{host}/api/2.0/mcp/vector-search/{catalog}/{schema}
         - UC Functions: https://{host}/api/2.0/mcp/functions/{catalog}/{schema}
         - Connection: https://{host}/api/2.0/mcp/external/{connection_name}
+        - Databricks App: Retrieved dynamically from workspace
         """
         # Direct URL provided
         if self.url:
@@ -1940,6 +2007,10 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
         # DBSQL MCP server (serverless, workspace-level)
         if self.sql:
             return f"{workspace_host}/api/2.0/mcp/sql"
+        
+        # Databricks App
+        if self.app:
+            return self.app.url
 
         # Vector Search
         if self.vector_search:
@@ -1950,33 +2021,35 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
                 raise ValueError(
                     "vector_search must have an index with a schema (catalog/schema) configured"
                 )
-            catalog: str = self.vector_search.index.schema_model.catalog_name
-            schema: str = self.vector_search.index.schema_model.schema_name
+            catalog: str = value_of(self.vector_search.index.schema_model.catalog_name)
+            schema: str = value_of(self.vector_search.index.schema_model.schema_name)
             return f"{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}"
 
         # UC Functions MCP server
         if self.functions:
-            catalog: str = self.functions.catalog_name
-            schema: str = self.functions.schema_name
+            catalog: str = value_of(self.functions.catalog_name)
+            schema: str = value_of(self.functions.schema_name)
             return f"{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}"
 
         raise ValueError(
-            "No URL source configured. Provide one of: url, connection, genie_room, "
+            "No URL source configured. Provide one of: url, app, connection, genie_room, "
             "sql, vector_search, or functions"
         )
 
     @field_serializer("transport")
-    def serialize_transport(self, value) -> str:
+    def serialize_transport(self, value: TransportType) -> str:
+        """Serialize transport enum to string."""
         if isinstance(value, TransportType):
             return value.value
         return str(value)
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self) -> "McpFunctionModel":
+    def validate_mutually_exclusive(self) -> Self:
         """Validate that exactly one URL source is provided."""
         # Count how many URL sources are provided
         url_sources: list[tuple[str, Any]] = [
             ("url", self.url),
+            ("app", self.app),
             ("connection", self.connection),
             ("genie_room", self.genie_room),
             ("sql", self.sql),
@@ -1992,13 +2065,13 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
             if len(provided_sources) == 0:
                 raise ValueError(
                     "For STREAMABLE_HTTP transport, exactly one of the following must be provided: "
-                    "url, connection, genie_room, sql, vector_search, or functions"
+                    "url, app, connection, genie_room, sql, vector_search, or functions"
                 )
             if len(provided_sources) > 1:
                 raise ValueError(
                     f"For STREAMABLE_HTTP transport, only one URL source can be provided. "
                     f"Found: {', '.join(provided_sources)}. "
-                    f"Please provide only one of: url, connection, genie_room, sql, vector_search, or functions"
+                    f"Please provide only one of: url, app, connection, genie_room, sql, vector_search, or functions"
                 )
 
         if self.transport == TransportType.STDIO:
@@ -2010,14 +2083,41 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
         return self
 
     @model_validator(mode="after")
-    def update_url(self) -> "McpFunctionModel":
-        self.url = value_of(self.url)
+    def update_url(self) -> Self:
+        """Resolve AnyVariable to concrete value for URL."""
+        if self.url is not None:
+            resolved_value: Any = value_of(self.url)
+            # Cast to string since URL must be a string
+            self.url = str(resolved_value) if resolved_value else None
         return self
 
     @model_validator(mode="after")
-    def update_headers(self) -> "McpFunctionModel":
+    def update_headers(self) -> Self:
+        """Resolve AnyVariable to concrete values for headers."""
         for key, value in self.headers.items():
-            self.headers[key] = value_of(value)
+            resolved_value: Any = value_of(value)
+            # Headers must be strings
+            self.headers[key] = str(resolved_value) if resolved_value else ""
+        return self
+
+    @model_validator(mode="after")
+    def validate_tool_filters(self) -> Self:
+        """Validate tool filter configuration."""
+        from loguru import logger
+
+        # Warn if both are empty lists (explicit but pointless)
+        if self.include_tools is not None and len(self.include_tools) == 0:
+            logger.warning(
+                "include_tools is empty list - no tools will be loaded. "
+                "Remove field to load all tools."
+            )
+
+        if self.exclude_tools is not None and len(self.exclude_tools) == 0:
+            logger.warning(
+                "exclude_tools is empty list - has no effect. "
+                "Remove field or add patterns."
+            )
+
         return self
 
     def as_tools(self, **kwargs: Any) -> Sequence[RunnableLike]:

dao_ai/middleware/__init__.py

@@ -6,6 +6,7 @@ from langchain.agents.middleware import (
     ClearToolUsesEdit,
     ContextEditingMiddleware,
     HumanInTheLoopMiddleware,
+    LLMToolSelectorMiddleware,
     ModelCallLimitMiddleware,
     ModelRetryMiddleware,
     PIIMiddleware,
@@ -82,6 +83,7 @@ from dao_ai.middleware.summarization import (
 )
 from dao_ai.middleware.tool_call_limit import create_tool_call_limit_middleware
 from dao_ai.middleware.tool_retry import create_tool_retry_middleware
+from dao_ai.middleware.tool_selector import create_llm_tool_selector_middleware
 
 __all__ = [
     # Base class (from LangChain)
@@ -105,6 +107,7 @@ __all__ = [
     "ModelCallLimitMiddleware",
     "ToolRetryMiddleware",
     "ModelRetryMiddleware",
+    "LLMToolSelectorMiddleware",
     "ContextEditingMiddleware",
     "ClearToolUsesEdit",
     "PIIMiddleware",
@@ -150,6 +153,8 @@ __all__ = [
     "create_model_call_limit_middleware",
     "create_tool_retry_middleware",
     "create_model_retry_middleware",
+    # Tool selection middleware factory functions
+    "create_llm_tool_selector_middleware",
     # Context editing middleware factory functions
     "create_context_editing_middleware",
     "create_clear_tool_uses_edit",