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

dao_ai/apps/server.py

@@ -0,0 +1,39 @@
+"""
+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 model_serving.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.apps.server
+
+    # With default model_config.yaml in current directory
+    python -m dao_ai.apps.server
+"""
+
+from mlflow.genai.agent_server import AgentServer
+
+# Import the agent handlers to register the invoke and stream decorators
+# This MUST happen before creating the AgentServer instance
+import dao_ai.apps.handlers  # noqa: E402, F401
+
+# 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.apps.server:app")
+
+
+if __name__ == "__main__":
+    main()

dao_ai/cli.py

@@ -285,6 +285,15 @@ Examples:
         action="store_true",
         help="Perform a dry run without executing the deployment or run commands",
     )
+    bundle_parser.add_argument(
+        "--deployment-target",
+        type=str,
+        choices=["model_serving", "apps"],
+        default=None,
+        help="Agent deployment target: 'model_serving' or 'apps'. "
+        "If not specified, uses app.deployment_target from config file, "
+        "or defaults to 'model_serving'. Passed to the deploy notebook.",
+    )
 
     # Deploy command
     deploy_parser: ArgumentParser = subparsers.add_parser(
@@ -309,6 +318,16 @@ 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=None,
+        help="Deployment target: 'model_serving' or 'apps'. "
+        "If not specified, uses app.deployment_target from config file, "
+        "or defaults to 'model_serving'.",
+    )
 
     # List MCP tools command
     list_mcp_parser: ArgumentParser = subparsers.add_parser(
@@ -729,11 +748,28 @@ 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)
+
+        # Hybrid target resolution:
+        # 1. CLI --target takes precedence
+        # 2. Fall back to config.app.deployment_target
+        # 3. Default to MODEL_SERVING (handled in deploy_agent)
+        target: DeploymentTarget | None = None
+        if options.target is not None:
+            target = DeploymentTarget(options.target)
+            logger.info(f"Using CLI-specified deployment target: {target.value}")
+        elif config.app is not None and config.app.deployment_target is not None:
+            target = config.app.deployment_target
+            logger.info(f"Using config file deployment target: {target.value}")
+        else:
+            logger.info("No deployment target specified, defaulting to model_serving")
+
         config.create_agent()
-        config.deploy_agent()
+        config.deploy_agent(target=target)
         sys.exit(0)
     except Exception as e:
         logger.error(f"Deployment failed: {e}")
@@ -1083,6 +1119,7 @@ def run_databricks_command(
     target: Optional[str] = None,
     cloud: Optional[str] = None,
     dry_run: bool = False,
+    deployment_target: Optional[str] = None,
 ) -> None:
     """Execute a databricks CLI command with optional profile, target, and cloud.
 
@@ -1093,6 +1130,8 @@ def run_databricks_command(
         target: Optional bundle target name (if not provided, auto-generated from app name and cloud)
         cloud: Optional cloud provider ('azure', 'aws', 'gcp'). Auto-detected if not specified.
         dry_run: If True, print the command without executing
+        deployment_target: Optional agent deployment target ('model_serving' or 'apps').
+            Passed to the deploy notebook via bundle variable.
     """
     config_path = Path(config) if config else None
 
@@ -1148,6 +1187,24 @@ def run_databricks_command(
 
         cmd.append(f'--var="config_path={relative_config}"')
 
+    # Add deployment_target variable for notebooks (hybrid resolution)
+    # Priority: CLI arg > config file > default (model_serving)
+    resolved_deployment_target: str = "model_serving"
+    if deployment_target is not None:
+        resolved_deployment_target = deployment_target
+        logger.debug(
+            f"Using CLI-specified deployment target: {resolved_deployment_target}"
+        )
+    elif app_config and app_config.app and app_config.app.deployment_target:
+        resolved_deployment_target = app_config.app.deployment_target.value
+        logger.debug(
+            f"Using config file deployment target: {resolved_deployment_target}"
+        )
+    else:
+        logger.debug("Using default deployment target: model_serving")
+
+    cmd.append(f'--var="deployment_target={resolved_deployment_target}"')
+
     logger.debug(f"Executing command: {' '.join(cmd)}")
 
     if dry_run:
@@ -1190,6 +1247,7 @@ def handle_bundle_command(options: Namespace) -> None:
     target: Optional[str] = options.target
     cloud: Optional[str] = options.cloud
     dry_run: bool = options.dry_run
+    deployment_target: Optional[str] = options.deployment_target
 
     if options.deploy:
         logger.info("Deploying DAO AI asset bundle...")
@@ -1200,6 +1258,7 @@ def handle_bundle_command(options: Namespace) -> None:
             target=target,
             cloud=cloud,
             dry_run=dry_run,
+            deployment_target=deployment_target,
         )
     if options.run:
         logger.info("Running DAO AI system with current configuration...")
@@ -1211,6 +1270,7 @@ def handle_bundle_command(options: Namespace) -> None:
             target=target,
             cloud=cloud,
             dry_run=dry_run,
+            deployment_target=deployment_target,
         )
     if options.destroy:
         logger.info("Destroying DAO AI system with current configuration...")
@@ -1221,6 +1281,7 @@ def handle_bundle_command(options: Namespace) -> None:
             target=target,
             cloud=cloud,
             dry_run=dry_run,
+            deployment_target=deployment_target,
         )
     else:
         logger.warning("No action specified. Use --deploy, --run or --destroy flags.")

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 = (
@@ -313,18 +344,24 @@ class IsDatabricksResource(ABC, BaseModel):
             else None
         )
 
-        if client_id_value and client_secret_value and workspace_host_value:
+        if client_id_value and client_secret_value:
+            # If workspace_host is not provided, check DATABRICKS_HOST env var first,
+            # then fall back to WorkspaceClient().config.host
+            if not workspace_host_value:
+                workspace_host_value = os.getenv("DATABRICKS_HOST")
+                if not workspace_host_value:
+                    workspace_host_value = WorkspaceClient().config.host
+
             logger.debug(
                 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 +369,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):
@@ -865,10 +910,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(
@@ -912,9 +953,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():
@@ -952,9 +990,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():
@@ -2775,6 +2810,11 @@ class AppModel(BaseModel):
         "which is supported by Databricks Model Serving. This allows deploying from "
         "environments with different Python versions (e.g., Databricks Apps with 3.11).",
     )
+    deployment_target: Optional[DeploymentTarget] = Field(
+        default=None,
+        description="Default deployment target. If not specified, defaults to MODEL_SERVING. "
+        "Can be overridden via CLI --target flag. Options: 'model_serving' or 'apps'.",
+    )
 
     @model_validator(mode="after")
     def set_databricks_env_vars(self) -> Self:
@@ -3255,6 +3295,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)
@@ -3275,6 +3316,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()
@@ -3282,12 +3326,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
@@ -3358,6 +3410,7 @@ class AppConfig(BaseModel):
 
     def deploy_agent(
         self,
+        target: DeploymentTarget | None = None,
         w: WorkspaceClient | None = None,
         vsc: "VectorSearchClient | None" = None,
         pat: str | None = None,
@@ -3365,9 +3418,39 @@ class AppConfig(BaseModel):
         client_secret: str | None = None,
         workspace_host: str | None = None,
     ) -> None:
+        """
+        Deploy the agent to the specified target.
+
+        Target resolution follows this priority:
+        1. Explicit `target` parameter (if provided)
+        2. `app.deployment_target` from config file (if set)
+        3. Default: MODEL_SERVING
+
+        Args:
+            target: The deployment target (MODEL_SERVING or APPS). If None, uses
+                config.app.deployment_target or 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
 
+        # Resolve target using hybrid logic:
+        # 1. Explicit parameter takes precedence
+        # 2. Fall back to config.app.deployment_target
+        # 3. Default to MODEL_SERVING
+        resolved_target: DeploymentTarget
+        if target is not None:
+            resolved_target = target
+        elif self.app is not None and self.app.deployment_target is not None:
+            resolved_target = self.app.deployment_target
+        else:
+            resolved_target = DeploymentTarget.MODEL_SERVING
+
         provider: ServiceProvider = DatabricksProvider(
             w=w,
             vsc=vsc,
@@ -3376,7 +3459,7 @@ class AppConfig(BaseModel):
             client_secret=client_secret,
             workspace_host=workspace_host,
         )
-        provider.deploy_agent(self)
+        provider.deploy_agent(self, target=resolved_target)
 
     def find_agents(
         self, predicate: Callable[[AgentModel], bool] | None = None

dao_ai/memory/postgres.py

@@ -178,7 +178,20 @@ class AsyncPostgresStoreManager(StoreManagerBase):
     def _setup(self):
         if self._setup_complete:
             return
-        asyncio.run(self._async_setup())
+        try:
+            # Check if we're already in an async context
+            asyncio.get_running_loop()
+            # If we get here, we're in an async context - raise to caller
+            raise RuntimeError(
+                "Cannot call sync _setup() from async context. "
+                "Use await _async_setup() instead."
+            )
+        except RuntimeError as e:
+            if "no running event loop" in str(e).lower():
+                # No event loop running - safe to use asyncio.run()
+                asyncio.run(self._async_setup())
+            else:
+                raise
 
     async def _async_setup(self):
         if self._setup_complete:
@@ -237,13 +250,25 @@ class AsyncPostgresCheckpointerManager(CheckpointManagerBase):
 
     def _setup(self):
         """
-        Run the async setup. Works in both sync and async contexts when nest_asyncio is applied.
+        Run the async setup. For async contexts, use await _async_setup() directly.
         """
         if self._setup_complete:
             return
 
-        # With nest_asyncio applied in notebooks, asyncio.run() works everywhere
-        asyncio.run(self._async_setup())
+        try:
+            # Check if we're already in an async context
+            asyncio.get_running_loop()
+            # If we get here, we're in an async context - raise to caller
+            raise RuntimeError(
+                "Cannot call sync _setup() from async context. "
+                "Use await _async_setup() instead."
+            )
+        except RuntimeError as e:
+            if "no running event loop" in str(e).lower():
+                # No event loop running - safe to use asyncio.run()
+                asyncio.run(self._async_setup())
+            else:
+                raise
 
     async def _async_setup(self):
         """