dao-ai 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

dao_ai/providers/base.py

@@ -1,15 +1,19 @@
 from abc import ABC, abstractmethod
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from dao_ai.config import (
     AppModel,
     DatasetModel,
+    DeploymentTarget,
     SchemaModel,
     UnityCatalogFunctionSqlModel,
     VectorStoreModel,
     VolumeModel,
 )
 
+if TYPE_CHECKING:
+    from dao_ai.config import AppConfig
+
 
 class ServiceProvider(ABC):
     @abstractmethod
@@ -52,4 +56,26 @@ class ServiceProvider(ABC):
     ) -> Any: ...
 
     @abstractmethod
-    def deploy_agent(self, config: AppModel) -> Any: ...
+    def deploy_model_serving_agent(self, config: "AppConfig") -> Any:
+        """Deploy agent to Databricks Model Serving endpoint."""
+        ...
+
+    @abstractmethod
+    def deploy_apps_agent(self, config: "AppConfig") -> Any:
+        """Deploy agent as a Databricks App."""
+        ...
+
+    @abstractmethod
+    def deploy_agent(
+        self,
+        config: "AppConfig",
+        target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
+    ) -> Any:
+        """
+        Deploy agent to the specified target.
+
+        Args:
+            config: The AppConfig containing deployment configuration
+            target: The deployment target (MODEL_SERVING or APPS)
+        """
+        ...

dao_ai/providers/databricks.py

@@ -23,7 +23,7 @@ from databricks.sdk.service.catalog import (
 )
 from databricks.sdk.service.database import DatabaseCredential
 from databricks.sdk.service.iam import User
-from databricks.sdk.service.workspace import GetSecretResponse
+from databricks.sdk.service.workspace import GetSecretResponse, ImportFormat
 from databricks.vector_search.client import VectorSearchClient
 from databricks.vector_search.index import VectorSearchIndex
 from loguru import logger
@@ -48,6 +48,7 @@ from dao_ai.config import (
     DatabaseModel,
     DatabricksAppModel,
     DatasetModel,
+    DeploymentTarget,
     FunctionModel,
     GenieRoomModel,
     HasFullName,
@@ -326,7 +327,7 @@ class DatabricksProvider(ServiceProvider):
                 raise FileNotFoundError(f"Code path does not exist: {path}")
 
         model_root_path: Path = Path(dao_ai.__file__).parent
-        model_path: Path = model_root_path / "agent_as_code.py"
+        model_path: Path = model_root_path / "apps" / "model_serving.py"
 
         pip_requirements: Sequence[str] = config.app.pip_requirements
 
@@ -439,8 +440,19 @@ class DatabricksProvider(ServiceProvider):
                 version=aliased_model.version,
             )
 
-    def deploy_agent(self, config: AppConfig) -> None:
-        logger.info("Deploying agent", endpoint_name=config.app.endpoint_name)
+    def deploy_model_serving_agent(self, config: AppConfig) -> None:
+        """
+        Deploy agent to Databricks Model Serving endpoint.
+
+        This is the original deployment method that creates/updates a Model Serving
+        endpoint with the registered model.
+
+        Args:
+            config: The AppConfig containing deployment configuration
+        """
+        logger.info(
+            "Deploying agent to Model Serving", endpoint_name=config.app.endpoint_name
+        )
         mlflow.set_registry_uri("databricks-uc")
 
         endpoint_name: str = config.app.endpoint_name
@@ -499,6 +511,228 @@ class DatabricksProvider(ServiceProvider):
                     permission_level=PermissionLevel[entitlement],
                 )
 
+    def deploy_apps_agent(self, config: AppConfig) -> None:
+        """
+        Deploy agent as a Databricks App.
+
+        This method creates or updates a Databricks App that serves the agent
+        using the app_server module.
+
+        The deployment process:
+        1. Determine the workspace source path for the app
+        2. Upload the configuration file to the workspace
+        3. Create the app if it doesn't exist
+        4. Deploy the app
+
+        Args:
+            config: The AppConfig containing deployment configuration
+
+        Note:
+            The config file must be loaded via AppConfig.from_file() so that
+            the source_config_path is available for upload.
+        """
+        import io
+
+        from databricks.sdk.service.apps import (
+            App,
+            AppDeployment,
+            AppDeploymentMode,
+            AppDeploymentState,
+        )
+
+        # Normalize app name: lowercase, replace underscores with dashes
+        raw_name: str = config.app.name
+        app_name: str = raw_name.lower().replace("_", "-")
+        if app_name != raw_name:
+            logger.info(
+                "Normalized app name for Databricks Apps",
+                original=raw_name,
+                normalized=app_name,
+            )
+        logger.info("Deploying agent to Databricks Apps", app_name=app_name)
+
+        # Use convention-based workspace path: /Workspace/Users/{user}/apps/{app_name}
+        current_user: User = self.w.current_user.me()
+        user_name: str = current_user.user_name or "default"
+        source_path: str = f"/Workspace/Users/{user_name}/apps/{app_name}"
+
+        logger.info("Using workspace source path", source_path=source_path)
+
+        # Upload the configuration file to the workspace
+        source_config_path: str | None = config.source_config_path
+        if source_config_path:
+            # Read the config file and upload to workspace
+            config_file_name: str = "model_config.yaml"
+            workspace_config_path: str = f"{source_path}/{config_file_name}"
+
+            logger.info(
+                "Uploading config file to workspace",
+                source=source_config_path,
+                destination=workspace_config_path,
+            )
+
+            # Read the source config file
+            with open(source_config_path, "rb") as f:
+                config_content: bytes = f.read()
+
+            # Create the directory if it doesn't exist and upload the file
+            try:
+                self.w.workspace.mkdirs(source_path)
+            except Exception as e:
+                logger.debug(f"Directory may already exist: {e}")
+
+            # Upload the config file
+            self.w.workspace.upload(
+                path=workspace_config_path,
+                content=io.BytesIO(config_content),
+                format=ImportFormat.AUTO,
+                overwrite=True,
+            )
+            logger.info("Config file uploaded", path=workspace_config_path)
+        else:
+            logger.warning(
+                "No source config path available. "
+                "Ensure DAO_AI_CONFIG_PATH is set in the app environment or "
+                "model_config.yaml exists in the app source directory."
+            )
+
+        # Generate and upload app.yaml with dynamically discovered resources
+        from dao_ai.apps.resources import generate_app_yaml
+
+        app_yaml_content: str = generate_app_yaml(
+            config,
+            command=[
+                "/bin/bash",
+                "-c",
+                "pip install dao-ai && python -m dao_ai.apps.server",
+            ],
+            include_resources=True,
+        )
+
+        app_yaml_path: str = f"{source_path}/app.yaml"
+        self.w.workspace.upload(
+            path=app_yaml_path,
+            content=io.BytesIO(app_yaml_content.encode("utf-8")),
+            format=ImportFormat.AUTO,
+            overwrite=True,
+        )
+        logger.info("app.yaml with resources uploaded", path=app_yaml_path)
+
+        # Generate SDK resources from the config
+        from dao_ai.apps.resources import (
+            generate_sdk_resources,
+            generate_user_api_scopes,
+        )
+
+        sdk_resources = generate_sdk_resources(config)
+        if sdk_resources:
+            logger.info(
+                "Discovered app resources from config",
+                resource_count=len(sdk_resources),
+                resources=[r.name for r in sdk_resources],
+            )
+
+        # Generate user API scopes for on-behalf-of-user resources
+        user_api_scopes = generate_user_api_scopes(config)
+        if user_api_scopes:
+            logger.info(
+                "Discovered user API scopes for OBO resources",
+                scopes=user_api_scopes,
+            )
+
+        # Check if app exists
+        app_exists: bool = False
+        try:
+            existing_app: App = self.w.apps.get(name=app_name)
+            app_exists = True
+            logger.debug("App already exists, updating", app_name=app_name)
+        except NotFound:
+            logger.debug("Creating new app", app_name=app_name)
+
+        # Create or update the app with resources and user_api_scopes
+        if not app_exists:
+            logger.info("Creating Databricks App", app_name=app_name)
+            app_spec = App(
+                name=app_name,
+                description=config.app.description or f"DAO AI Agent: {app_name}",
+                resources=sdk_resources if sdk_resources else None,
+                user_api_scopes=user_api_scopes if user_api_scopes else None,
+            )
+            app: App = self.w.apps.create_and_wait(app=app_spec)
+            logger.info("App created", app_name=app.name, app_url=app.url)
+        else:
+            app = existing_app
+            # Update resources and scopes on existing app
+            if sdk_resources or user_api_scopes:
+                logger.info("Updating app resources and scopes", app_name=app_name)
+                updated_app = App(
+                    name=app_name,
+                    description=config.app.description or app.description,
+                    resources=sdk_resources if sdk_resources else None,
+                    user_api_scopes=user_api_scopes if user_api_scopes else None,
+                )
+                app = self.w.apps.update(name=app_name, app=updated_app)
+                logger.info("App resources and scopes updated", app_name=app_name)
+
+        # Deploy the app with source code
+        # The app will use the dao_ai.apps.server module as the entry point
+        logger.info("Deploying app", app_name=app_name)
+
+        # Create deployment configuration
+        app_deployment = AppDeployment(
+            mode=AppDeploymentMode.SNAPSHOT,
+            source_code_path=source_path,
+        )
+
+        # Deploy the app
+        deployment: AppDeployment = self.w.apps.deploy_and_wait(
+            app_name=app_name,
+            app_deployment=app_deployment,
+        )
+
+        if (
+            deployment.status
+            and deployment.status.state == AppDeploymentState.SUCCEEDED
+        ):
+            logger.info(
+                "App deployed successfully",
+                app_name=app_name,
+                deployment_id=deployment.deployment_id,
+                app_url=app.url if app else None,
+            )
+        else:
+            status_message: str = (
+                deployment.status.message if deployment.status else "Unknown error"
+            )
+            logger.error(
+                "App deployment failed",
+                app_name=app_name,
+                status=status_message,
+            )
+            raise RuntimeError(f"App deployment failed: {status_message}")
+
+    def deploy_agent(
+        self,
+        config: AppConfig,
+        target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
+    ) -> None:
+        """
+        Deploy agent to the specified target.
+
+        This is the main deployment method that routes to the appropriate
+        deployment implementation based on the target.
+
+        Args:
+            config: The AppConfig containing deployment configuration
+            target: The deployment target (MODEL_SERVING or APPS)
+        """
+        if target == DeploymentTarget.MODEL_SERVING:
+            self.deploy_model_serving_agent(config)
+        elif target == DeploymentTarget.APPS:
+            self.deploy_apps_agent(config)
+        else:
+            raise ValueError(f"Unknown deployment target: {target}")
+
     def create_catalog(self, schema: SchemaModel) -> CatalogInfo:
         catalog_info: CatalogInfo
         try:

dao_ai/state.py

@@ -164,6 +164,7 @@ class Context(BaseModel):
 
     user_id: str | None = None
     thread_id: str | None = None
+    headers: dict[str, Any] | None = None
 
     @classmethod
     def from_runnable_config(cls, config: dict[str, Any]) -> "Context":

dao_ai/tools/mcp.py

@@ -261,12 +261,12 @@ def _extract_text_content(result: CallToolResult) -> str:
     return "\n".join(text_parts)
 
 
-def _fetch_tools_from_server(function: McpFunctionModel) -> list[Tool]:
+async def _afetch_tools_from_server(function: McpFunctionModel) -> list[Tool]:
     """
-    Fetch raw MCP tools from the server.
+    Async version: Fetch raw MCP tools from the server.
 
-    This is the core async operation that connects to the MCP server
-    and retrieves the list of available tools.
+    This is the primary async implementation that handles the actual MCP connection
+    and tool listing. It's used by both alist_mcp_tools() and acreate_mcp_tools().
 
     Args:
         function: The MCP function model configuration.
@@ -280,14 +280,10 @@ def _fetch_tools_from_server(function: McpFunctionModel) -> list[Tool]:
     connection_config = _build_connection_config(function)
     client = MultiServerMCPClient({"mcp_function": connection_config})
 
-    async def _list_tools_async() -> list[Tool]:
-        """Async helper to list tools from MCP server."""
+    try:
         async with client.session("mcp_function") as session:
             result = await session.list_tools()
             return result.tools if hasattr(result, "tools") else list(result)
-
-    try:
-        return asyncio.run(_list_tools_async())
     except Exception as e:
         if function.connection:
             logger.error(
@@ -312,57 +308,48 @@ def _fetch_tools_from_server(function: McpFunctionModel) -> list[Tool]:
             ) from e
 
 
-def list_mcp_tools(
-    function: McpFunctionModel,
-    apply_filters: bool = True,
-) -> list[MCPToolInfo]:
+def _fetch_tools_from_server(function: McpFunctionModel) -> list[Tool]:
     """
-    List available tools from an MCP server.
-
-    This function connects to an MCP server and returns information about
-    all available tools. It's designed for:
-    - Tool discovery and exploration
-    - UI-based tool selection (e.g., in DAO AI Builder)
-    - Debugging and validation of MCP configurations
+    Sync wrapper: Fetch raw MCP tools from the server.
 
-    The returned MCPToolInfo objects contain all information needed to
-    display tools in a UI and allow users to select which tools to use.
+    For async contexts, use _afetch_tools_from_server() directly.
 
     Args:
-        function: The MCP function model configuration containing:
-            - Connection details (url, connection, headers, etc.)
-            - Optional filtering (include_tools, exclude_tools)
-        apply_filters: Whether to apply include_tools/exclude_tools filters.
-            Set to False to get the complete list of available tools
-            regardless of filter configuration. Default True.
+        function: The MCP function model configuration.
 
     Returns:
-        List of MCPToolInfo objects describing available tools.
-        Each contains name, description, and input_schema.
+        List of raw MCP Tool objects from the server.
 
     Raises:
         RuntimeError: If connection to MCP server fails.
+    """
+    return asyncio.run(_afetch_tools_from_server(function))
 
-    Example:
-        # List all tools from a DBSQL MCP server
-        from dao_ai.config import McpFunctionModel
-        from dao_ai.tools.mcp import list_mcp_tools
 
-        function = McpFunctionModel(sql=True)
-        tools = list_mcp_tools(function)
+async def alist_mcp_tools(
+    function: McpFunctionModel,
+    apply_filters: bool = True,
+) -> list[MCPToolInfo]:
+    """
+    Async version: List available tools from an MCP server.
+
+    This is the primary async implementation for tool discovery.
+    For sync contexts, use list_mcp_tools() instead.
 
-        for tool in tools:
-            print(f"{tool.name}: {tool.description}")
+    Args:
+        function: The MCP function model configuration.
+        apply_filters: Whether to apply include_tools/exclude_tools filters.
 
-        # Get unfiltered list (ignore include_tools/exclude_tools)
-        all_tools = list_mcp_tools(function, apply_filters=False)
+    Returns:
+        List of MCPToolInfo objects describing available tools.
 
-    Note:
-        For creating executable LangChain tools, use create_mcp_tools() instead.
-        This function is for discovery/display purposes only.
+    Raises:
+        RuntimeError: If connection to MCP server fails.
     """
     mcp_url = function.mcp_url
-    logger.debug("Listing MCP tools", mcp_url=mcp_url, apply_filters=apply_filters)
+    logger.debug(
+        "Listing MCP tools (async)", mcp_url=mcp_url, apply_filters=apply_filters
+    )
 
     # Log connection type
     if function.connection:
@@ -378,8 +365,8 @@ def list_mcp_tools(
             mcp_url=mcp_url,
         )
 
-    # Fetch tools from server
-    mcp_tools: list[Tool] = _fetch_tools_from_server(function)
+    # Fetch tools from server (async)
+    mcp_tools: list[Tool] = await _afetch_tools_from_server(function)
 
     # Log discovered tools
     logger.info(
@@ -433,45 +420,155 @@ def list_mcp_tools(
     return tool_infos
 
 
-def create_mcp_tools(
+def list_mcp_tools(
     function: McpFunctionModel,
-) -> Sequence[RunnableLike]:
+    apply_filters: bool = True,
+) -> list[MCPToolInfo]:
     """
-    Create executable LangChain tools for invoking Databricks MCP functions.
+    Sync wrapper: List available tools from an MCP server.
+
+    For async contexts, use alist_mcp_tools() directly.
+
+    Args:
+        function: The MCP function model configuration.
+        apply_filters: Whether to apply include_tools/exclude_tools filters.
 
-    Supports both direct MCP connections and UC Connection-based MCP access.
-    Uses manual tool wrappers to ensure response format compatibility with
-    Databricks APIs (which reject extra fields in tool results).
+    Returns:
+        List of MCPToolInfo objects describing available tools.
+
+    Raises:
+        RuntimeError: If connection to MCP server fails.
+    """
+    return asyncio.run(alist_mcp_tools(function, apply_filters))
 
-    This function:
-    1. Fetches available tools from the MCP server
-    2. Applies include_tools/exclude_tools filters
-    3. Wraps each tool for LangChain agent execution
 
-    For tool discovery without creating executable tools, use list_mcp_tools().
+async def acreate_mcp_tools(
+    function: McpFunctionModel,
+) -> Sequence[RunnableLike]:
+    """
+    Async version: Create executable LangChain tools for invoking Databricks MCP functions.
 
-    Based on: https://docs.databricks.com/aws/en/generative-ai/mcp/external-mcp
+    This is the primary async implementation. For sync contexts, use create_mcp_tools().
 
     Args:
-        function: The MCP function model configuration containing:
-            - Connection details (url, connection, headers, etc.)
-            - Optional filtering (include_tools, exclude_tools)
+        function: The MCP function model configuration.
 
     Returns:
         A sequence of LangChain tools that can be used by agents.
 
     Raises:
         RuntimeError: If connection to MCP server fails.
+    """
+    mcp_url = function.mcp_url
+    logger.debug("Creating MCP tools (async)", mcp_url=mcp_url)
+
+    # Fetch tools from server (async)
+    mcp_tools: list[Tool] = await _afetch_tools_from_server(function)
+
+    # Log discovered tools
+    logger.info(
+        "Discovered MCP tools from server",
+        tools_count=len(mcp_tools),
+        tool_names=[t.name for t in mcp_tools],
+        mcp_url=mcp_url,
+    )
+
+    # Apply filtering if configured
+    if function.include_tools or function.exclude_tools:
+        original_count = len(mcp_tools)
+        mcp_tools = [
+            tool
+            for tool in mcp_tools
+            if _should_include_tool(
+                tool.name,
+                function.include_tools,
+                function.exclude_tools,
+            )
+        ]
+        filtered_count = original_count - len(mcp_tools)
+
+        logger.info(
+            "Filtered MCP tools",
+            original_count=original_count,
+            filtered_count=filtered_count,
+            final_count=len(mcp_tools),
+            include_patterns=function.include_tools,
+            exclude_patterns=function.exclude_tools,
+        )
+
+    # Log final tool list
+    for mcp_tool in mcp_tools:
+        logger.debug(
+            "MCP tool available",
+            tool_name=mcp_tool.name,
+            tool_description=(
+                mcp_tool.description[:100] if mcp_tool.description else None
+            ),
+        )
+
+    def _create_tool_wrapper(mcp_tool: Tool) -> RunnableLike:
+        """
+        Create a LangChain tool wrapper for an MCP tool.
+        """
+
+        @create_tool(
+            mcp_tool.name,
+            description=mcp_tool.description or f"MCP tool: {mcp_tool.name}",
+            args_schema=mcp_tool.inputSchema,
+        )
+        async def tool_wrapper(**kwargs: Any) -> str:
+            """Execute MCP tool with fresh session."""
+            logger.trace("Invoking MCP tool", tool_name=mcp_tool.name, args=kwargs)
+
+            invocation_client = MultiServerMCPClient(
+                {"mcp_function": _build_connection_config(function)}
+            )
+
+            try:
+                async with invocation_client.session("mcp_function") as session:
+                    result: CallToolResult = await session.call_tool(
+                        mcp_tool.name, kwargs
+                    )
+
+                    text_result = _extract_text_content(result)
+
+                    logger.trace(
+                        "MCP tool completed",
+                        tool_name=mcp_tool.name,
+                        result_length=len(text_result),
+                    )
+
+                    return text_result
+
+            except Exception as e:
+                logger.error(
+                    "MCP tool failed",
+                    tool_name=mcp_tool.name,
+                    error=str(e),
+                )
+                raise
+
+        return tool_wrapper
+
+    return [_create_tool_wrapper(tool) for tool in mcp_tools]
 
-    Example:
-        from dao_ai.config import McpFunctionModel
-        from dao_ai.tools.mcp import create_mcp_tools
 
-        function = McpFunctionModel(sql=True)
-        tools = create_mcp_tools(function)
+def create_mcp_tools(
+    function: McpFunctionModel,
+) -> Sequence[RunnableLike]:
+    """
+    Sync wrapper: Create executable LangChain tools for invoking Databricks MCP functions.
 
-        # Use tools in an agent
-        agent = create_agent(model=model, tools=tools)
+    For async contexts, use acreate_mcp_tools() directly.
+
+    Args:
+        function: The MCP function model configuration.
+
+    Returns:
+        A sequence of LangChain tools that can be used by agents.
+
+    Raises:
+        RuntimeError: If connection to MCP server fails.
     """
     mcp_url = function.mcp_url
     logger.debug("Creating MCP tools", mcp_url=mcp_url)

{dao_ai-0.1.8.dist-info → dao_ai-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.8
+Version: 0.1.10
 Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
 Project-URL: Homepage, https://github.com/natefleming/dao-ai
 Project-URL: Documentation, https://natefleming.github.io/dao-ai
@@ -43,7 +43,7 @@ Requires-Dist: langgraph>=1.0.5
 Requires-Dist: langmem>=0.0.30
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: mcp>=1.24.0
-Requires-Dist: mlflow>=3.8.1
+Requires-Dist: mlflow[databricks]>=3.8.1
 Requires-Dist: nest-asyncio>=1.6.0
 Requires-Dist: openevals>=0.1.3
 Requires-Dist: openpyxl>=3.1.5