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

dao_ai/app_server.py

@@ -0,0 +1,103 @@
+"""
+App server module for running dao-ai agents as Databricks Apps.
+
+This module provides the entry point for deploying dao-ai agents as Databricks Apps
+using MLflow's AgentServer. It follows the same pattern as agent_as_code.py but
+uses the AgentServer for the Databricks Apps runtime.
+
+Configuration Loading:
+    The config path is specified via the DAO_AI_CONFIG_PATH environment variable,
+    or defaults to model_config.yaml in the current directory.
+
+Usage:
+    # With environment variable
+    DAO_AI_CONFIG_PATH=/path/to/config.yaml python -m dao_ai.app_server
+
+    # With default model_config.yaml in current directory
+    python -m dao_ai.app_server
+"""
+
+import os
+from typing import AsyncGenerator
+
+import mlflow
+from dotenv import load_dotenv
+from mlflow.genai.agent_server import AgentServer, invoke, stream
+from mlflow.pyfunc import ResponsesAgent
+from mlflow.types.responses import (
+    ResponsesAgentRequest,
+    ResponsesAgentResponse,
+    ResponsesAgentStreamEvent,
+)
+
+from dao_ai.config import AppConfig
+from dao_ai.logging import configure_logging
+
+# Load environment variables from .env.local if it exists
+load_dotenv(dotenv_path=".env.local", override=True)
+
+# Configure MLflow
+mlflow.set_registry_uri("databricks-uc")
+mlflow.set_tracking_uri("databricks")
+mlflow.langchain.autolog()
+
+# Get config path from environment or use default
+config_path: str = os.environ.get("DAO_AI_CONFIG_PATH", "model_config.yaml")
+
+# Load configuration using AppConfig.from_file (consistent with CLI, notebook, builder)
+config: AppConfig = AppConfig.from_file(config_path)
+
+# Configure logging
+if config.app and config.app.log_level:
+    configure_logging(level=config.app.log_level)
+
+# Create the ResponsesAgent
+_responses_agent: ResponsesAgent = config.as_responses_agent()
+
+
+@invoke()
+def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
+    """
+    Handle non-streaming requests by delegating to the ResponsesAgent.
+
+    Args:
+        request: The incoming ResponsesAgentRequest
+
+    Returns:
+        ResponsesAgentResponse with the complete output
+    """
+    return _responses_agent.predict(request)
+
+
+@stream()
+def streaming(
+    request: ResponsesAgentRequest,
+) -> AsyncGenerator[ResponsesAgentStreamEvent, None]:
+    """
+    Handle streaming requests by delegating to the ResponsesAgent.
+
+    Args:
+        request: The incoming ResponsesAgentRequest
+
+    Yields:
+        ResponsesAgentStreamEvent objects as they are generated
+    """
+    # The predict_stream method returns a generator, convert to async generator
+    for event in _responses_agent.predict_stream(request):
+        yield event
+
+
+# Create the AgentServer instance
+agent_server = AgentServer("ResponsesAgent", enable_chat_proxy=True)
+
+# Define the app as a module level variable to enable multiple workers
+app = agent_server.app
+
+
+def main() -> None:
+    """Entry point for running the agent server."""
+    agent_server.run(app_import_string="dao_ai.app_server:app")
+
+
+if __name__ == "__main__":
+    main()

dao_ai/cli.py

@@ -309,6 +309,14 @@ Examples:
         metavar="FILE",
         help="Path to the model configuration file to validate",
     )
+    deploy_parser.add_argument(
+        "-t",
+        "--target",
+        type=str,
+        choices=["model_serving", "apps"],
+        default="model_serving",
+        help="Deployment target: 'model_serving' (default) or 'apps'",
+    )
 
     # List MCP tools command
     list_mcp_parser: ArgumentParser = subparsers.add_parser(
@@ -729,11 +737,17 @@ def handle_graph_command(options: Namespace) -> None:
 
 
 def handle_deploy_command(options: Namespace) -> None:
+    from dao_ai.config import DeploymentTarget
+
     logger.debug(f"Validating configuration from {options.config}...")
     try:
         config: AppConfig = AppConfig.from_file(options.config)
+
+        # Convert target string to enum
+        target: DeploymentTarget = DeploymentTarget(options.target)
+
         config.create_agent()
-        config.deploy_agent()
+        config.deploy_agent(target=target)
         sys.exit(0)
     except Exception as e:
         logger.error(f"Deployment failed: {e}")

dao_ai/config.py

@@ -208,7 +208,9 @@ class IsDatabricksResource(ABC, BaseModel):
     Authentication Options:
     ----------------------
     1. **On-Behalf-Of User (OBO)**: Set on_behalf_of_user=True to use the
-       calling user's identity via ModelServingUserCredentials.
+       calling user's identity. Implementation varies by deployment:
+       - Databricks Apps: Uses X-Forwarded-Access-Token from request headers
+       - Model Serving: Uses ModelServingUserCredentials
 
     2. **Service Principal (OAuth M2M)**: Provide service_principal or
        (client_id + client_secret + workspace_host) for service principal auth.
@@ -221,9 +223,17 @@ class IsDatabricksResource(ABC, BaseModel):
 
     Authentication Priority:
     1. OBO (on_behalf_of_user=True)
+       - Checks for forwarded headers (Databricks Apps)
+       - Falls back to ModelServingUserCredentials (Model Serving)
     2. Service Principal (client_id + client_secret + workspace_host)
     3. PAT (pat + workspace_host)
     4. Ambient/default authentication
+
+    Note: When on_behalf_of_user=True, the agent acts as the calling user regardless
+    of deployment target. In Databricks Apps, this uses X-Forwarded-Access-Token
+    automatically captured by MLflow AgentServer. In Model Serving, this uses
+    ModelServingUserCredentials. Forwarded headers are ONLY used when
+    on_behalf_of_user=True.
     """
 
     model_config = ConfigDict(use_enum_values=True)
@@ -235,9 +245,6 @@ class IsDatabricksResource(ABC, BaseModel):
     workspace_host: Optional[AnyVariable] = None
     pat: Optional[AnyVariable] = None
 
-    # Private attribute to cache the workspace client (lazy instantiation)
-    _workspace_client: Optional[WorkspaceClient] = PrivateAttr(default=None)
-
     @abstractmethod
     def as_resources(self) -> Sequence[DatabricksResource]: ...
 
@@ -273,32 +280,56 @@ class IsDatabricksResource(ABC, BaseModel):
         """
         Get a WorkspaceClient configured with the appropriate authentication.
 
-        The client is lazily instantiated on first access and cached for subsequent calls.
+        A new client is created on each access.
 
         Authentication priority:
-        1. If on_behalf_of_user is True, uses ModelServingUserCredentials (OBO)
-        2. If service principal credentials are configured (client_id, client_secret,
-           workspace_host), uses OAuth M2M
-        3. If PAT is configured, uses token authentication
-        4. Otherwise, uses default/ambient authentication
+        1. On-Behalf-Of User (on_behalf_of_user=True):
+           - Forwarded headers (Databricks Apps)
+           - ModelServingUserCredentials (Model Serving)
+        2. Service Principal (client_id + client_secret + workspace_host)
+        3. PAT (pat + workspace_host)
+        4. Ambient/default authentication
         """
-        # Return cached client if already instantiated
-        if self._workspace_client is not None:
-            return self._workspace_client
-
         from dao_ai.utils import normalize_host
 
         # Check for OBO first (highest priority)
         if self.on_behalf_of_user:
+            # NEW: In Databricks Apps, use forwarded headers for per-user auth
+            try:
+                from mlflow.genai.agent_server import get_request_headers
+
+                headers = get_request_headers()
+                forwarded_token = headers.get("x-forwarded-access-token")
+
+                if forwarded_token:
+                    forwarded_user = headers.get("x-forwarded-user", "unknown")
+                    logger.debug(
+                        f"Creating WorkspaceClient for {self.__class__.__name__} "
+                        f"with OBO using forwarded token from Databricks Apps",
+                        forwarded_user=forwarded_user,
+                    )
+                    # Use workspace_host if configured, otherwise SDK will auto-detect
+                    workspace_host_value: str | None = (
+                        normalize_host(value_of(self.workspace_host))
+                        if self.workspace_host
+                        else None
+                    )
+                    return WorkspaceClient(
+                        host=workspace_host_value,
+                        token=forwarded_token,
+                        auth_type="pat",
+                    )
+            except (ImportError, LookupError):
+                # mlflow not available or headers not set - fall through to Model Serving
+                pass
+
+            # Fall back to Model Serving OBO (existing behavior)
             credentials_strategy: CredentialsStrategy = ModelServingUserCredentials()
             logger.debug(
                 f"Creating WorkspaceClient for {self.__class__.__name__} "
-                f"with OBO credentials strategy"
-            )
-            self._workspace_client = WorkspaceClient(
-                credentials_strategy=credentials_strategy
+                f"with OBO credentials strategy (Model Serving)"
             )
-            return self._workspace_client
+            return WorkspaceClient(credentials_strategy=credentials_strategy)
 
         # Check for service principal credentials
         client_id_value: str | None = (
@@ -318,13 +349,12 @@ class IsDatabricksResource(ABC, BaseModel):
                 f"Creating WorkspaceClient for {self.__class__.__name__} with service principal: "
                 f"client_id={client_id_value}, host={workspace_host_value}"
             )
-            self._workspace_client = WorkspaceClient(
+            return WorkspaceClient(
                 host=workspace_host_value,
                 client_id=client_id_value,
                 client_secret=client_secret_value,
                 auth_type="oauth-m2m",
             )
-            return self._workspace_client
 
         # Check for PAT authentication
         pat_value: str | None = value_of(self.pat) if self.pat else None
@@ -332,20 +362,28 @@ class IsDatabricksResource(ABC, BaseModel):
             logger.debug(
                 f"Creating WorkspaceClient for {self.__class__.__name__} with PAT"
             )
-            self._workspace_client = WorkspaceClient(
+            return WorkspaceClient(
                 host=workspace_host_value,
                 token=pat_value,
                 auth_type="pat",
             )
-            return self._workspace_client
 
         # Default: use ambient authentication
         logger.debug(
             f"Creating WorkspaceClient for {self.__class__.__name__} "
             "with default/ambient authentication"
         )
-        self._workspace_client = WorkspaceClient()
-        return self._workspace_client
+        return WorkspaceClient()
+
+
+class DeploymentTarget(str, Enum):
+    """Target platform for agent deployment."""
+
+    MODEL_SERVING = "model_serving"
+    """Deploy to Databricks Model Serving endpoint."""
+
+    APPS = "apps"
+    """Deploy as a Databricks App."""
 
 
 class Privilege(str, Enum):
@@ -418,11 +456,11 @@ 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 
+    The `url` is dynamically retrieved from the workspace client by calling
     `apps.get(name)` and returning the app's URL.
-    
+
     Example:
         ```yaml
         resources:
@@ -431,7 +469,7 @@ class DatabricksAppModel(IsDatabricksResource, HasFullName):
               name: my-databricks-app
         ```
     """
-    
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
     """The unique instance name of the Databricks App in the workspace."""
@@ -440,10 +478,10 @@ class DatabricksAppModel(IsDatabricksResource, HasFullName):
     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.
         """
@@ -455,7 +493,6 @@ class DatabricksAppModel(IsDatabricksResource, HasFullName):
             )
         return app.url
 
-
     @property
     def full_name(self) -> str:
         return self.name
@@ -761,11 +798,20 @@ class FunctionModel(IsDatabricksResource, HasFullName):
 
 
 class WarehouseModel(IsDatabricksResource):
-    model_config = ConfigDict()
-    name: str
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    name: Optional[str] = None
     description: Optional[str] = None
     warehouse_id: AnyVariable
 
+    _warehouse_details: Optional[GetWarehouseResponse] = PrivateAttr(default=None)
+
+    def _get_warehouse_details(self) -> GetWarehouseResponse:
+        if self._warehouse_details is None:
+            self._warehouse_details = self.workspace_client.warehouses.get(
+                id=value_of(self.warehouse_id)
+            )
+        return self._warehouse_details
+
     @property
     def api_scopes(self) -> Sequence[str]:
         return [
@@ -786,10 +832,22 @@ class WarehouseModel(IsDatabricksResource):
         self.warehouse_id = value_of(self.warehouse_id)
         return self
 
+    @model_validator(mode="after")
+    def populate_name(self) -> Self:
+        """Populate name from warehouse details if not provided."""
+        if self.warehouse_id and not self.name:
+            try:
+                warehouse_details = self._get_warehouse_details()
+                if warehouse_details.name:
+                    self.name = warehouse_details.name
+            except Exception as e:
+                logger.debug(f"Could not fetch details from warehouse: {e}")
+        return self
+
 
 class GenieRoomModel(IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    name: str
+    name: Optional[str] = None
     description: Optional[str] = None
     space_id: AnyVariable
 
@@ -845,10 +903,6 @@ class GenieRoomModel(IsDatabricksResource):
                 pat=self.pat,
             )
 
-            # Share the cached workspace client if available
-            if self._workspace_client is not None:
-                warehouse_model._workspace_client = self._workspace_client
-
             return warehouse_model
         except Exception as e:
             logger.warning(
@@ -892,9 +946,6 @@ class GenieRoomModel(IsDatabricksResource):
                                 workspace_host=self.workspace_host,
                                 pat=self.pat,
                             )
-                            # Share the cached workspace client if available
-                            if self._workspace_client is not None:
-                                table_model._workspace_client = self._workspace_client
 
                             # Verify the table exists before adding
                             if not table_model.exists():
@@ -932,9 +983,6 @@ class GenieRoomModel(IsDatabricksResource):
                 workspace_host=self.workspace_host,
                 pat=self.pat,
             )
-            # Share the cached workspace client if available
-            if self._workspace_client is not None:
-                function_model._workspace_client = self._workspace_client
 
             # Verify the function exists before adding
             if not function_model.exists():
@@ -998,15 +1046,17 @@ class GenieRoomModel(IsDatabricksResource):
         return self
 
     @model_validator(mode="after")
-    def update_description_from_space(self) -> Self:
-        """Populate description from GenieSpace if not provided."""
-        if not self.description:
+    def populate_name_and_description(self) -> Self:
+        """Populate name and description from GenieSpace if not provided."""
+        if self.space_id and (not self.name or not self.description):
             try:
                 space_details = self._get_space_details()
-                if space_details.description:
+                if not self.name and space_details.title:
+                    self.name = space_details.title
+                if not self.description and space_details.description:
                     self.description = space_details.description
             except Exception as e:
-                logger.debug(f"Could not fetch description from Genie space: {e}")
+                logger.debug(f"Could not fetch details from Genie space: {e}")
         return self
 
 
@@ -2007,7 +2057,7 @@ 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
@@ -3233,6 +3283,7 @@ class ResourcesModel(BaseModel):
 
 class AppConfig(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    version: Optional[str] = None
     variables: dict[str, AnyVariable] = Field(default_factory=dict)
     service_principals: dict[str, ServicePrincipalModel] = Field(default_factory=dict)
     schemas: dict[str, SchemaModel] = Field(default_factory=dict)
@@ -3253,6 +3304,9 @@ class AppConfig(BaseModel):
     )
     providers: Optional[dict[type | str, Any]] = None
 
+    # Private attribute to track the source config file path (set by from_file)
+    _source_config_path: str | None = None
+
     @classmethod
     def from_file(cls, path: PathLike) -> "AppConfig":
         path = Path(path).as_posix()
@@ -3260,12 +3314,20 @@ class AppConfig(BaseModel):
         model_config: ModelConfig = ModelConfig(development_config=path)
         config: AppConfig = AppConfig(**model_config.to_dict())
 
+        # Store the source config path for later use (e.g., Apps deployment)
+        config._source_config_path = path
+
         config.initialize()
 
         atexit.register(config.shutdown)
 
         return config
 
+    @property
+    def source_config_path(self) -> str | None:
+        """Get the source config file path if loaded via from_file."""
+        return self._source_config_path
+
     def initialize(self) -> None:
         from dao_ai.hooks.core import create_hooks
         from dao_ai.logging import configure_logging
@@ -3336,6 +3398,7 @@ class AppConfig(BaseModel):
 
     def deploy_agent(
         self,
+        target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
         w: WorkspaceClient | None = None,
         vsc: "VectorSearchClient | None" = None,
         pat: str | None = None,
@@ -3343,6 +3406,18 @@ class AppConfig(BaseModel):
         client_secret: str | None = None,
         workspace_host: str | None = None,
     ) -> None:
+        """
+        Deploy the agent to the specified target.
+
+        Args:
+            target: The deployment target (MODEL_SERVING or APPS). Defaults to MODEL_SERVING.
+            w: Optional WorkspaceClient instance
+            vsc: Optional VectorSearchClient instance
+            pat: Optional personal access token for authentication
+            client_id: Optional client ID for service principal authentication
+            client_secret: Optional client secret for service principal authentication
+            workspace_host: Optional workspace host URL
+        """
         from dao_ai.providers.base import ServiceProvider
         from dao_ai.providers.databricks import DatabricksProvider
 
@@ -3354,7 +3429,7 @@ class AppConfig(BaseModel):
             client_secret=client_secret,
             workspace_host=workspace_host,
         )
-        provider.deploy_agent(self)
+        provider.deploy_agent(self, target=target)
 
     def find_agents(
         self, predicate: Callable[[AgentModel], bool] | None = None

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,
@@ -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,196 @@ 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
+        # Use pip to install dao-ai and run the app server
+        app_yaml_content: str = """command:
+  - /bin/bash
+  - -c
+  - |
+    pip install dao-ai && python -m dao_ai.app_server
+
+env:
+  - name: MLFLOW_TRACKING_URI
+    value: "databricks"
+  - name: MLFLOW_REGISTRY_URI
+    value: "databricks-uc"
+  - name: DAO_AI_CONFIG_PATH
+    value: "model_config.yaml"
+"""
+        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 uploaded", path=app_yaml_path)
+
+        # 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 get the app
+        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}",
+            )
+            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
+
+        # Deploy the app with source code
+        # The app will use the dao_ai.app_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

@@ -26,9 +26,9 @@ from loguru import logger
 from mcp.types import CallToolResult, TextContent, Tool
 
 from dao_ai.config import (
+    IsDatabricksResource,
     McpFunctionModel,
     TransportType,
-    value_of,
 )
 
 
@@ -143,12 +143,54 @@ def _should_include_tool(
     return True
 
 
+def _get_auth_resource(function: McpFunctionModel) -> IsDatabricksResource:
+    """
+    Get the IsDatabricksResource to use for authentication.
+
+    Follows a priority hierarchy:
+    1. Explicit resource with auth (app, connection, genie_room, vector_search, functions)
+    2. McpFunctionModel itself (which also inherits from IsDatabricksResource)
+
+    Returns the resource whose workspace_client should be used for authentication.
+    """
+    # Check each possible resource source in priority order
+    # These resources may have their own auth configured
+    if function.app:
+        return function.app
+    if function.connection:
+        return function.connection
+    if function.genie_room:
+        return function.genie_room
+    if function.vector_search:
+        return function.vector_search
+    if function.functions:
+        # SchemaModel doesn't have auth - fall through to McpFunctionModel
+        pass
+
+    # Fall back to McpFunctionModel itself (it inherits from IsDatabricksResource)
+    return function
+
+
 def _build_connection_config(
     function: McpFunctionModel,
 ) -> dict[str, Any]:
     """
     Build the connection configuration dictionary for MultiServerMCPClient.
 
+    Authentication Strategy:
+    -----------------------
+    For HTTP transport, authentication is handled consistently using
+    DatabricksOAuthClientProvider with the workspace_client from the appropriate
+    IsDatabricksResource. The auth resource is selected in this priority:
+
+    1. Nested resource (app, connection, genie_room, vector_search) if it has auth
+    2. McpFunctionModel itself (inherits from IsDatabricksResource)
+
+    This approach ensures:
+    - Consistent auth handling across all MCP sources
+    - Automatic token refresh for long-running connections
+    - Support for OBO, service principal, PAT, and ambient auth
+
     Args:
         function: The MCP function model configuration.
 
@@ -162,52 +204,30 @@ def _build_connection_config(
             "transport": function.transport.value,
         }
 
-    # For HTTP transport with UC Connection, use DatabricksOAuthClientProvider
-    if function.connection:
-        from databricks_mcp import DatabricksOAuthClientProvider
+    # For HTTP transport, use DatabricksOAuthClientProvider with unified auth
+    from databricks_mcp import DatabricksOAuthClientProvider
 
-        workspace_client = function.connection.workspace_client
-        auth_provider = DatabricksOAuthClientProvider(workspace_client)
+    # Get the resource to use for authentication
+    auth_resource = _get_auth_resource(function)
 
-        logger.trace(
-            "Using DatabricksOAuthClientProvider for authentication",
-            connection_name=function.connection.name,
-        )
-
-        return {
-            "url": function.mcp_url,
-            "transport": "http",
-            "auth": auth_provider,
-        }
-
-    # For HTTP transport with headers-based authentication
-    headers: dict[str, str] = {
-        key: str(value_of(val)) for key, val in function.headers.items()
-    }
+    # Get workspace client from the auth resource
+    workspace_client = auth_resource.workspace_client
+    auth_provider = DatabricksOAuthClientProvider(workspace_client)
 
-    if "Authorization" not in headers:
-        logger.trace("Generating fresh authentication token")
-
-        from dao_ai.providers.databricks import DatabricksProvider
-
-        try:
-            provider = DatabricksProvider(
-                workspace_host=value_of(function.workspace_host),
-                client_id=value_of(function.client_id),
-                client_secret=value_of(function.client_secret),
-                pat=value_of(function.pat),
-            )
-            headers["Authorization"] = f"Bearer {provider.create_token()}"
-            logger.trace("Generated fresh authentication token")
-        except Exception as e:
-            logger.error("Failed to create fresh token", error=str(e))
-    else:
-        logger.trace("Using existing authentication token")
+    # Log which resource is providing auth
+    resource_name = (
+        getattr(auth_resource, "name", None) or auth_resource.__class__.__name__
+    )
+    logger.trace(
+        "Using DatabricksOAuthClientProvider for authentication",
+        auth_resource=resource_name,
+        resource_type=auth_resource.__class__.__name__,
+    )
 
     return {
         "url": function.mcp_url,
         "transport": "http",
-        "headers": headers,
+        "auth": auth_provider,
     }
 
 

{dao_ai-0.1.7.dist-info → dao_ai-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.7
+Version: 0.1.9
 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

{dao_ai-0.1.7.dist-info → dao_ai-0.1.9.dist-info}/RECORD

@@ -1,8 +1,9 @@
 dao_ai/__init__.py,sha256=18P98ExEgUaJ1Byw440Ct1ty59v6nxyWtc5S6Uq2m9Q,1062
 dao_ai/agent_as_code.py,sha256=xIlLDpPVfmDVzLvbdY_V_CrC4Jvj2ItCWJ-NzdrszTo,538
+dao_ai/app_server.py,sha256=QKpl068z-s1gLF67dPW-3fT77i33t_Oab4_ugmxISWs,3010
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
-dao_ai/cli.py,sha256=7LGrVDRgSBpznr8c8EksAhzPW_8NJ9h4St3DSpx-0z4,48196
-dao_ai/config.py,sha256=GY-n2PtPg4pYtO46KYpHFIGkNkCDVoPA5S7sKLjWpVc,124699
+dao_ai/cli.py,sha256=57Kmmi0zgS92ACBTD-gH5hZzW6rPDkbdkRVlFjX4onQ,48604
+dao_ai/config.py,sha256=t7kXU7XjdMaCZ3G9Hn-O9NDOaTS_LaMXX6s5mdyv3dM,127944
 dao_ai/graph.py,sha256=1-uQlo7iXZQTT3uU8aYu0N5rnhw5_g_2YLwVsAs6M-U,1119
 dao_ai/logging.py,sha256=lYy4BmucCHvwW7aI3YQkQXKJtMvtTnPDu9Hnd7_O4oc,1556
 dao_ai/messages.py,sha256=4ZBzO4iFdktGSLrmhHzFjzMIt2tpaL-aQLHOQJysGnY,6959
@@ -10,7 +11,7 @@ dao_ai/models.py,sha256=AwzwTRTNZF-UOh59HsuXEgFk_YH6q6M-mERNDe64Z8k,81783
 dao_ai/nodes.py,sha256=7W6Ek6Uk9-pKa-H06nVCwuDllCrgX02IYy3rHtuL0aM,10777
 dao_ai/optimization.py,sha256=phK6t4wYmWPObCjGUBHdZzsaFXGhQOjhAek2bAEfwXo,22971
 dao_ai/prompts.py,sha256=4cz5bZ7cOzrjyQ8hMp-K4evK6cVYrkGrAGdUl8-KDEM,2784
-dao_ai/state.py,sha256=0wbbzfQmldkCu26gdTE5j0Rl-_pfilza-YIHPbSWlvI,6394
+dao_ai/state.py,sha256=ifDTAC7epdowk3Z1CP3Xqw4uH2dIxQEVF3C747dA8yI,6436
 dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dao_ai/utils.py,sha256=_Urd7Nj2VzrgPKf3NS4E6vt0lWRhEUddBqWN9BksqeE,11543
 dao_ai/vector_search.py,sha256=8d3xROg9zSIYNXjRRl6rSexsJTlufjRl5Fy1ZA8daKA,4019
@@ -48,14 +49,14 @@ dao_ai/orchestration/core.py,sha256=qoU7uMXBJCth-sqfu0jRE1L0GOn5H4LoZdRUY1Ib3DI,
 dao_ai/orchestration/supervisor.py,sha256=alKMEEo9G5LhdpMvTVdAMel234cZj5_MguWl4wFB7XQ,9873
 dao_ai/orchestration/swarm.py,sha256=8tp1eGmsQqqWpaDcjPoJckddPWohZdmmN0RGRJ_xzOA,9198
 dao_ai/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dao_ai/providers/base.py,sha256=-fjKypCOk28h6vioPfMj9YZSw_3Kcbi2nMuAyY7vX9k,1383
-dao_ai/providers/databricks.py,sha256=XxYkyoDYkwGV_Xg1IJBpGOl4d7U5HiFP4RtjjSLgenI,61437
+dao_ai/providers/base.py,sha256=cJGo3UjUTPgS91dv38ePOHwQQtYhIa84ebb167CBXjk,2111
+dao_ai/providers/databricks.py,sha256=X_VjzZogwiSlNpPBWP0iMCCXAvfFDRlbC4AZCHleb2A,68608
 dao_ai/tools/__init__.py,sha256=NfRpAKds_taHbx6gzLPWgtPXve-YpwzkoOAUflwxceM,1734
 dao_ai/tools/agent.py,sha256=plIWALywRjaDSnot13nYehBsrHRpBUpsVZakoGeajOE,1858
 dao_ai/tools/core.py,sha256=bRIN3BZhRQX8-Kpu3HPomliodyskCqjxynQmYbk6Vjs,3783
 dao_ai/tools/email.py,sha256=A3TsCoQgJR7UUWR0g45OPRGDpVoYwctFs1MOZMTt_d4,7389
 dao_ai/tools/genie.py,sha256=4e_5MeAe7kDzHbYeXuNPFbY5z8ci3ouj8l5254CZ2lA,8874
-dao_ai/tools/mcp.py,sha256=0OfP4b4skcjeF2rzkOLYqd65ti1Mj55N_l8VoQlH9qo,17818
+dao_ai/tools/mcp.py,sha256=ZNalYo2atZECatZjMT8w4mHEsaUZJQ_fsCjia7px1nc,18689
 dao_ai/tools/memory.py,sha256=lwObKimAand22Nq3Y63tsv-AXQ5SXUigN9PqRjoWKes,1836
 dao_ai/tools/python.py,sha256=jWFnZPni2sCdtd8D1CqXnZIPHnWkdK27bCJnBXpzhvo,1879
 dao_ai/tools/search.py,sha256=cJ3D9FKr1GAR6xz55dLtRkjtQsI0WRueGt9TPDFpOxc,433
@@ -64,8 +65,8 @@ dao_ai/tools/sql.py,sha256=tKd1gjpLuKdQDyfmyYYtMiNRHDW6MGRbdEVaeqyB8Ok,7632
 dao_ai/tools/time.py,sha256=tufJniwivq29y0LIffbgeBTIDE6VgrLpmVf8Qr90qjw,9224
 dao_ai/tools/unity_catalog.py,sha256=AjQfW7bvV8NurqDLIyntYRv2eJuTwNdbvex1L5CRjOk,15534
 dao_ai/tools/vector_search.py,sha256=oe2uBwl2TfeJIXPpwiS6Rmz7wcHczSxNyqS9P3hE6co,14542
-dao_ai-0.1.7.dist-info/METADATA,sha256=jzaENv6Ic9S-uds3qTC4Pu7chwzYG8y0zvTBni4VCW8,16685
-dao_ai-0.1.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.1.7.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.1.7.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.1.7.dist-info/RECORD,,
+dao_ai-0.1.9.dist-info/METADATA,sha256=_cVhNKdywyci5ZyrmMfkF9CiWveo4kdoj9OlaM_w-cg,16697
+dao_ai-0.1.9.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.1.9.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.1.9.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.1.9.dist-info/RECORD,,