dao-ai 0.1.9__py3-none-any.whl → 0.1.11__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(
@@ -314,8 +323,10 @@ Examples:
         "--target",
         type=str,
         choices=["model_serving", "apps"],
-        default="model_serving",
-        help="Deployment target: 'model_serving' (default) or '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
@@ -743,8 +754,19 @@ def handle_deploy_command(options: Namespace) -> None:
     try:
         config: AppConfig = AppConfig.from_file(options.config)
 
-        # Convert target string to enum
-        target: DeploymentTarget = DeploymentTarget(options.target)
+        # 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(target=target)
@@ -1097,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.
 
@@ -1107,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
 
@@ -1162,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:
@@ -1204,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...")
@@ -1214,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...")
@@ -1225,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...")
@@ -1235,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

@@ -344,7 +344,14 @@ 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}"
@@ -2803,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:
@@ -3398,7 +3410,7 @@ class AppConfig(BaseModel):
 
     def deploy_agent(
         self,
-        target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
+        target: DeploymentTarget | None = None,
         w: WorkspaceClient | None = None,
         vsc: "VectorSearchClient | None" = None,
         pat: str | None = None,
@@ -3409,8 +3421,14 @@ class AppConfig(BaseModel):
         """
         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). Defaults to MODEL_SERVING.
+            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
@@ -3421,6 +3439,18 @@ class AppConfig(BaseModel):
         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,
@@ -3429,7 +3459,7 @@ class AppConfig(BaseModel):
             client_secret=client_secret,
             workspace_host=workspace_host,
         )
-        provider.deploy_agent(self, target=target)
+        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):
         """