dao-ai 0.1.2__py3-none-any.whl → 0.1.20__py3-none-any.whl

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,
@@ -151,25 +152,77 @@ class DatabricksProvider(ServiceProvider):
         client_secret: str | None = None,
         workspace_host: str | None = None,
     ) -> None:
-        if w is None:
-            w = _workspace_client(
-                pat=pat,
-                client_id=client_id,
-                client_secret=client_secret,
-                workspace_host=workspace_host,
+        # Store credentials for lazy initialization
+        self._pat = pat
+        self._client_id = client_id
+        self._client_secret = client_secret
+        self._workspace_host = workspace_host
+
+        # Lazy initialization for WorkspaceClient
+        self._w: WorkspaceClient | None = w
+        self._w_initialized = w is not None
+
+        # Lazy initialization for VectorSearchClient - only create when needed
+        # This avoids authentication errors in Databricks Apps where VSC
+        # requires explicit credentials but the platform uses ambient auth
+        self._vsc: VectorSearchClient | None = vsc
+        self._vsc_initialized = vsc is not None
+
+        # Lazy initialization for DatabricksFunctionClient
+        self._dfs: DatabricksFunctionClient | None = dfs
+        self._dfs_initialized = dfs is not None
+
+    @property
+    def w(self) -> WorkspaceClient:
+        """Lazy initialization of WorkspaceClient."""
+        if not self._w_initialized:
+            self._w = _workspace_client(
+                pat=self._pat,
+                client_id=self._client_id,
+                client_secret=self._client_secret,
+                workspace_host=self._workspace_host,
             )
-        if vsc is None:
-            vsc = _vector_search_client(
-                pat=pat,
-                client_id=client_id,
-                client_secret=client_secret,
-                workspace_host=workspace_host,
+            self._w_initialized = True
+        return self._w  # type: ignore[return-value]
+
+    @w.setter
+    def w(self, value: WorkspaceClient) -> None:
+        """Set WorkspaceClient and mark as initialized."""
+        self._w = value
+        self._w_initialized = True
+
+    @property
+    def vsc(self) -> VectorSearchClient:
+        """Lazy initialization of VectorSearchClient."""
+        if not self._vsc_initialized:
+            self._vsc = _vector_search_client(
+                pat=self._pat,
+                client_id=self._client_id,
+                client_secret=self._client_secret,
+                workspace_host=self._workspace_host,
             )
-        if dfs is None:
-            dfs = _function_client(w=w)
-        self.w = w
-        self.vsc = vsc
-        self.dfs = dfs
+            self._vsc_initialized = True
+        return self._vsc  # type: ignore[return-value]
+
+    @vsc.setter
+    def vsc(self, value: VectorSearchClient) -> None:
+        """Set VectorSearchClient and mark as initialized."""
+        self._vsc = value
+        self._vsc_initialized = True
+
+    @property
+    def dfs(self) -> DatabricksFunctionClient:
+        """Lazy initialization of DatabricksFunctionClient."""
+        if not self._dfs_initialized:
+            self._dfs = _function_client(w=self.w)
+            self._dfs_initialized = True
+        return self._dfs  # type: ignore[return-value]
+
+    @dfs.setter
+    def dfs(self, value: DatabricksFunctionClient) -> None:
+        """Set DatabricksFunctionClient and mark as initialized."""
+        self._dfs = value
+        self._dfs_initialized = True
 
     def experiment_name(self, config: AppConfig) -> str:
         current_user: User = self.w.current_user.me()
@@ -326,7 +379,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
 
@@ -344,6 +397,8 @@ class DatabricksProvider(ServiceProvider):
 
             pip_requirements += get_installed_packages()
 
+        code_paths = list(dict.fromkeys(code_paths))
+
         logger.trace("Pip requirements prepared", count=len(pip_requirements))
         logger.trace("Code paths prepared", count=len(code_paths))
 
@@ -381,19 +436,38 @@ class DatabricksProvider(ServiceProvider):
             pip_packages_count=len(pip_requirements),
         )
 
-        with mlflow.start_run(run_name=run_name):
-            mlflow.set_tag("type", "agent")
-            mlflow.set_tag("dao_ai", dao_ai_version())
-            logged_agent_info: ModelInfo = mlflow.pyfunc.log_model(
-                python_model=model_path.as_posix(),
-                code_paths=code_paths,
-                model_config=config.model_dump(mode="json", by_alias=True),
-                name="agent",
-                conda_env=conda_env,
-                input_example=input_example,
-                # resources=all_resources,
-                auth_policy=auth_policy,
+        # End any stale runs before starting to ensure clean state on retry
+        if mlflow.active_run():
+            logger.warning(
+                "Ending stale MLflow run before creating new agent",
+                run_id=mlflow.active_run().info.run_id,
             )
+            mlflow.end_run()
+
+        try:
+            with mlflow.start_run(run_name=run_name):
+                mlflow.set_tag("type", "agent")
+                mlflow.set_tag("dao_ai", dao_ai_version())
+                logged_agent_info: ModelInfo = mlflow.pyfunc.log_model(
+                    python_model=model_path.as_posix(),
+                    code_paths=code_paths,
+                    model_config=config.model_dump(mode="json", by_alias=True),
+                    name="agent",
+                    conda_env=conda_env,
+                    input_example=input_example,
+                    # resources=all_resources,
+                    auth_policy=auth_policy,
+                )
+        except Exception as e:
+            # Ensure run is ended on failure to prevent stale state on retry
+            if mlflow.active_run():
+                mlflow.end_run(status="FAILED")
+            logger.error(
+                "Failed to log model",
+                run_name=run_name,
+                error=str(e),
+            )
+            raise
 
         registered_model_name: str = config.app.registered_model.full_name
 
@@ -439,8 +513,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 +584,240 @@ 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)
+
+        # Get or create experiment for this app (for tracing and tracking)
+        from mlflow.entities import Experiment
+
+        experiment: Experiment = self.get_or_create_experiment(config)
+        logger.info(
+            "Using MLflow experiment for app",
+            experiment_name=experiment.name,
+            experiment_id=experiment.experiment_id,
+        )
+
+        # 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 = "dao_ai.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 "
+                "dao_ai.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 (including experiment)
+        from dao_ai.apps.resources import (
+            generate_sdk_resources,
+            generate_user_api_scopes,
+        )
+
+        sdk_resources = generate_sdk_resources(
+            config, experiment_id=experiment.experiment_id
+        )
+        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:
@@ -625,6 +944,17 @@ class DatabricksProvider(ServiceProvider):
                 df.write.mode("overwrite").saveAsTable(table)
 
     def create_vector_store(self, vector_store: VectorStoreModel) -> None:
+        """
+        Create a vector search index from a source table.
+
+        This method expects a VectorStoreModel in provisioning mode with all
+        required fields validated. Use VectorStoreModel.create() which handles
+        mode detection and validation.
+
+        Args:
+            vector_store: VectorStoreModel configured for provisioning
+        """
+        # Ensure endpoint exists
         if not endpoint_exists(self.vsc, vector_store.endpoint.name):
             self.vsc.create_endpoint_and_wait(
                 name=vector_store.endpoint.name,

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/__init__.py

@@ -1,10 +1,10 @@
-from dao_ai.genie.cache import LRUCacheService, SemanticCacheService
+from dao_ai.genie.cache import LRUCacheService, PostgresContextAwareGenieService
 from dao_ai.hooks.core import create_hooks
 from dao_ai.tools.agent import create_agent_endpoint_tool
 from dao_ai.tools.core import create_tools, say_hello_tool
 from dao_ai.tools.email import create_send_email_tool
 from dao_ai.tools.genie import create_genie_tool
-from dao_ai.tools.mcp import create_mcp_tools
+from dao_ai.tools.mcp import MCPToolInfo, create_mcp_tools, list_mcp_tools
 from dao_ai.tools.memory import create_search_memory_tool
 from dao_ai.tools.python import create_factory_tool, create_python_tool
 from dao_ai.tools.search import create_search_tool
@@ -30,6 +30,8 @@ __all__ = [
     "create_genie_tool",
     "create_hooks",
     "create_mcp_tools",
+    "list_mcp_tools",
+    "MCPToolInfo",
     "create_python_tool",
     "create_search_memory_tool",
     "create_search_tool",
@@ -42,8 +44,8 @@ __all__ = [
     "format_time_tool",
     "is_business_hours_tool",
     "LRUCacheService",
+    "PostgresContextAwareGenieService",
     "say_hello_tool",
-    "SemanticCacheService",
     "time_difference_tool",
     "time_in_timezone_tool",
     "time_until_tool",