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

dao_ai/providers/databricks.py

@@ -1,6 +1,5 @@
 import base64
 import uuid
-from importlib.metadata import version
 from pathlib import Path
 from typing import Any, Callable, Final, Sequence
 
@@ -32,6 +31,7 @@ from mlflow import MlflowClient
 from mlflow.entities import Experiment
 from mlflow.entities.model_registry import PromptVersion
 from mlflow.entities.model_registry.model_version import ModelVersion
+from mlflow.genai.prompts import load_prompt
 from mlflow.models.auth_policy import AuthPolicy, SystemAuthPolicy, UserAuthPolicy
 from mlflow.models.model import ModelInfo
 from mlflow.models.resources import (
@@ -46,6 +46,7 @@ from dao_ai.config import (
     AppConfig,
     ConnectionModel,
     DatabaseModel,
+    DatabricksAppModel,
     DatasetModel,
     FunctionModel,
     GenieRoomModel,
@@ -65,9 +66,11 @@ from dao_ai.config import (
 from dao_ai.models import get_latest_model_version
 from dao_ai.providers.base import ServiceProvider
 from dao_ai.utils import (
+    dao_ai_version,
     get_installed_packages,
     is_installed,
     is_lib_provided,
+    normalize_host,
     normalize_name,
 )
 from dao_ai.vector_search import endpoint_exists, index_exists
@@ -89,15 +92,18 @@ def _workspace_client(
     Create a WorkspaceClient instance with the provided parameters.
     If no parameters are provided, it will use the default configuration.
     """
-    if client_id and client_secret and workspace_host:
+    # Normalize the workspace host to ensure it has https:// scheme
+    normalized_host = normalize_host(workspace_host)
+
+    if client_id and client_secret and normalized_host:
         return WorkspaceClient(
-            host=workspace_host,
+            host=normalized_host,
             client_id=client_id,
             client_secret=client_secret,
             auth_type="oauth-m2m",
         )
     elif pat:
-        return WorkspaceClient(host=workspace_host, token=pat, auth_type="pat")
+        return WorkspaceClient(host=normalized_host, token=pat, auth_type="pat")
     else:
         return WorkspaceClient()
 
@@ -112,15 +118,18 @@ def _vector_search_client(
     Create a VectorSearchClient instance with the provided parameters.
     If no parameters are provided, it will use the default configuration.
     """
-    if client_id and client_secret and workspace_host:
+    # Normalize the workspace host to ensure it has https:// scheme
+    normalized_host = normalize_host(workspace_host)
+
+    if client_id and client_secret and normalized_host:
         return VectorSearchClient(
-            workspace_url=workspace_host,
+            workspace_url=normalized_host,
             service_principal_client_id=client_id,
             service_principal_client_secret=client_secret,
         )
-    elif pat and workspace_host:
+    elif pat and normalized_host:
         return VectorSearchClient(
-            workspace_url=workspace_host,
+            workspace_url=normalized_host,
             personal_access_token=pat,
         )
     else:
@@ -172,15 +181,17 @@ class DatabricksProvider(ServiceProvider):
         experiment: Experiment | None = mlflow.get_experiment_by_name(experiment_name)
         if experiment is None:
             experiment_id: str = mlflow.create_experiment(name=experiment_name)
-            logger.info(
-                f"Created new experiment: {experiment_name} (ID: {experiment_id})"
+            logger.success(
+                "Created new MLflow experiment",
+                experiment_name=experiment_name,
+                experiment_id=experiment_id,
             )
             experiment = mlflow.get_experiment(experiment_id)
         return experiment
 
     def create_token(self) -> str:
         current_user: User = self.w.current_user.me()
-        logger.debug(f"Authenticated to Databricks as {current_user}")
+        logger.debug("Authenticated to Databricks", user=str(current_user))
         headers: dict[str, str] = self.w.config.authenticate()
         token: str = headers["Authorization"].replace("Bearer ", "")
         return token
@@ -192,17 +203,24 @@ class DatabricksProvider(ServiceProvider):
             secret_response: GetSecretResponse = self.w.secrets.get_secret(
                 secret_scope, secret_key
             )
-            logger.debug(f"Retrieved secret {secret_key} from scope {secret_scope}")
+            logger.trace(
+                "Retrieved secret", secret_key=secret_key, secret_scope=secret_scope
+            )
             encoded_secret: str = secret_response.value
             decoded_secret: str = base64.b64decode(encoded_secret).decode("utf-8")
             return decoded_secret
         except NotFound:
             logger.warning(
-                f"Secret {secret_key} not found in scope {secret_scope}, using default value"
+                "Secret not found, using default value",
+                secret_key=secret_key,
+                secret_scope=secret_scope,
             )
         except Exception as e:
             logger.error(
-                f"Error retrieving secret {secret_key} from scope {secret_scope}: {e}"
+                "Error retrieving secret",
+                secret_key=secret_key,
+                secret_scope=secret_scope,
+                error=str(e),
             )
 
         return default_value
@@ -211,9 +229,18 @@ class DatabricksProvider(ServiceProvider):
         self,
         config: AppConfig,
     ) -> ModelInfo:
-        logger.debug("Creating agent...")
+        logger.info("Creating agent")
         mlflow.set_registry_uri("databricks-uc")
 
+        # Set up experiment for proper tracking
+        experiment: Experiment = self.get_or_create_experiment(config)
+        mlflow.set_experiment(experiment_id=experiment.experiment_id)
+        logger.debug(
+            "Using MLflow experiment",
+            experiment_name=experiment.name,
+            experiment_id=experiment.experiment_id,
+        )
+
         llms: Sequence[LLMModel] = list(config.resources.llms.values())
         vector_indexes: Sequence[IndexModel] = list(
             config.resources.vector_stores.values()
@@ -231,6 +258,7 @@ class DatabricksProvider(ServiceProvider):
         )
         databases: Sequence[DatabaseModel] = list(config.resources.databases.values())
         volumes: Sequence[VolumeModel] = list(config.resources.volumes.values())
+        apps: Sequence[DatabricksAppModel] = list(config.resources.apps.values())
 
         resources: Sequence[IsDatabricksResource] = (
             llms
@@ -242,6 +270,7 @@ class DatabricksProvider(ServiceProvider):
             + connections
             + databases
             + volumes
+            + apps
         )
 
         # Flatten all resources from all models into a single list
@@ -255,12 +284,16 @@ class DatabricksProvider(ServiceProvider):
             for resource in r.as_resources()
             if not r.on_behalf_of_user
         ]
-        logger.debug(f"system_resources: {[r.name for r in system_resources]}")
+        logger.trace(
+            "System resources identified",
+            count=len(system_resources),
+            resources=[r.name for r in system_resources],
+        )
 
         system_auth_policy: SystemAuthPolicy = SystemAuthPolicy(
             resources=system_resources
         )
-        logger.debug(f"system_auth_policy: {system_auth_policy}")
+        logger.trace("System auth policy created", policy=str(system_auth_policy))
 
         api_scopes: Sequence[str] = list(
             set(
@@ -272,15 +305,19 @@ class DatabricksProvider(ServiceProvider):
                 ]
             )
         )
-        logger.debug(f"api_scopes: {api_scopes}")
+        logger.trace("API scopes identified", scopes=api_scopes)
 
         user_auth_policy: UserAuthPolicy = UserAuthPolicy(api_scopes=api_scopes)
-        logger.debug(f"user_auth_policy: {user_auth_policy}")
+        logger.trace("User auth policy created", policy=str(user_auth_policy))
 
         auth_policy: AuthPolicy = AuthPolicy(
             system_auth_policy=system_auth_policy, user_auth_policy=user_auth_policy
         )
-        logger.debug(f"auth_policy: {auth_policy}")
+        logger.debug(
+            "Auth policy created",
+            has_system_auth=system_auth_policy is not None,
+            has_user_auth=user_auth_policy is not None,
+        )
 
         code_paths: list[str] = config.app.code_paths
         for path in code_paths:
@@ -296,7 +333,7 @@ class DatabricksProvider(ServiceProvider):
         if is_installed():
             if not is_lib_provided("dao-ai", pip_requirements):
                 pip_requirements += [
-                    f"dao-ai=={version('dao-ai')}",
+                    f"dao-ai=={dao_ai_version()}",
                 ]
         else:
             src_path: Path = model_root_path.parent
@@ -307,27 +344,52 @@ class DatabricksProvider(ServiceProvider):
 
             pip_requirements += get_installed_packages()
 
-        logger.debug(f"pip_requirements: {pip_requirements}")
-        logger.debug(f"code_paths: {code_paths}")
+        logger.trace("Pip requirements prepared", count=len(pip_requirements))
+        logger.trace("Code paths prepared", count=len(code_paths))
 
         run_name: str = normalize_name(config.app.name)
-        logger.debug(f"run_name: {run_name}")
-        logger.debug(f"model_path: {model_path.as_posix()}")
+        logger.debug(
+            "Agent run configuration",
+            run_name=run_name,
+            model_path=model_path.as_posix(),
+        )
 
         input_example: dict[str, Any] = None
         if config.app.input_example:
             input_example = config.app.input_example.model_dump()
 
-        logger.debug(f"input_example: {input_example}")
+        logger.trace("Input example configured", has_example=input_example is not None)
+
+        # Create conda environment with configured Python version
+        # This allows deploying from environments with different Python versions
+        # (e.g., Databricks Apps with Python 3.11 can deploy to Model Serving with 3.12)
+        target_python_version: str = config.app.python_version
+        logger.debug("Target Python version configured", version=target_python_version)
+
+        conda_env: dict[str, Any] = {
+            "name": "mlflow-env",
+            "channels": ["conda-forge"],
+            "dependencies": [
+                f"python={target_python_version}",
+                "pip",
+                {"pip": list(pip_requirements)},
+            ],
+        }
+        logger.trace(
+            "Conda environment configured",
+            python_version=target_python_version,
+            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(by_alias=True),
+                model_config=config.model_dump(mode="json", by_alias=True),
                 name="agent",
-                pip_requirements=pip_requirements,
+                conda_env=conda_env,
                 input_example=input_example,
                 # resources=all_resources,
                 auth_policy=auth_policy,
@@ -338,15 +400,26 @@ class DatabricksProvider(ServiceProvider):
         model_version: ModelVersion = mlflow.register_model(
             name=registered_model_name, model_uri=logged_agent_info.model_uri
         )
-        logger.debug(
-            f"Registered model: {registered_model_name} with version: {model_version.version}"
+        logger.success(
+            "Model registered",
+            model_name=registered_model_name,
+            version=model_version.version,
         )
 
         client: MlflowClient = MlflowClient()
 
+        # Set tags on the model version
+        client.set_model_version_tag(
+            name=registered_model_name,
+            version=model_version.version,
+            key="dao_ai",
+            value=dao_ai_version(),
+        )
+        logger.trace("Set dao_ai tag on model version", version=model_version.version)
+
         client.set_registered_model_alias(
             name=registered_model_name,
-            alias="Current",
+            alias="Champion",
             version=model_version.version,
         )
 
@@ -359,12 +432,15 @@ class DatabricksProvider(ServiceProvider):
             aliased_model: ModelVersion = client.get_model_version_by_alias(
                 registered_model_name, config.app.alias
             )
-            logger.debug(
-                f"Model {registered_model_name} aliased to {config.app.alias} with version: {aliased_model.version}"
+            logger.info(
+                "Model aliased",
+                model_name=registered_model_name,
+                alias=config.app.alias,
+                version=aliased_model.version,
             )
 
     def deploy_agent(self, config: AppConfig) -> None:
-        logger.debug("Deploying agent...")
+        logger.info("Deploying agent", endpoint_name=config.app.endpoint_name)
         mlflow.set_registry_uri("databricks-uc")
 
         endpoint_name: str = config.app.endpoint_name
@@ -372,7 +448,10 @@ class DatabricksProvider(ServiceProvider):
         scale_to_zero: bool = config.app.scale_to_zero
         environment_vars: dict[str, str] = config.app.environment_vars
         workload_size: str = config.app.workload_size
-        tags: dict[str, str] = config.app.tags
+        tags: dict[str, str] = config.app.tags.copy() if config.app.tags else {}
+
+        # Add dao_ai framework tag
+        tags["dao_ai"] = dao_ai_version()
 
         latest_version: int = get_latest_model_version(registered_model_name)
 
@@ -382,12 +461,10 @@ class DatabricksProvider(ServiceProvider):
             agents.get_deployments(endpoint_name)
             endpoint_exists = True
             logger.debug(
-                f"Endpoint {endpoint_name} already exists, updating without tags to avoid conflicts..."
+                "Endpoint already exists, updating", endpoint_name=endpoint_name
             )
         except Exception:
-            logger.debug(
-                f"Endpoint {endpoint_name} doesn't exist, creating new with tags..."
-            )
+            logger.debug("Creating new endpoint", endpoint_name=endpoint_name)
 
         # Deploy - skip tags for existing endpoints to avoid conflicts
         agents.deploy(
@@ -403,8 +480,11 @@ class DatabricksProvider(ServiceProvider):
         registered_model_name: str = config.app.registered_model.full_name
         permissions: Sequence[dict[str, Any]] = config.app.permissions
 
-        logger.debug(registered_model_name)
-        logger.debug(permissions)
+        logger.debug(
+            "Configuring model permissions",
+            model_name=registered_model_name,
+            permissions_count=len(permissions),
+        )
 
         for permission in permissions:
             principals: Sequence[str] = permission.principals
@@ -424,7 +504,7 @@ class DatabricksProvider(ServiceProvider):
         try:
             catalog_info = self.w.catalogs.get(name=schema.catalog_name)
         except NotFound:
-            logger.debug(f"Creating catalog: {schema.catalog_name}")
+            logger.info("Creating catalog", catalog_name=schema.catalog_name)
             catalog_info = self.w.catalogs.create(name=schema.catalog_name)
         return catalog_info
 
@@ -434,7 +514,7 @@ class DatabricksProvider(ServiceProvider):
         try:
             schema_info = self.w.schemas.get(full_name=schema.full_name)
         except NotFound:
-            logger.debug(f"Creating schema: {schema.full_name}")
+            logger.info("Creating schema", schema_name=schema.full_name)
             schema_info = self.w.schemas.create(
                 name=schema.schema_name, catalog_name=catalog_info.name
             )
@@ -446,7 +526,7 @@ class DatabricksProvider(ServiceProvider):
         try:
             volume_info = self.w.volumes.read(name=volume.full_name)
         except NotFound:
-            logger.debug(f"Creating volume: {volume.full_name}")
+            logger.info("Creating volume", volume_name=volume.full_name)
             volume_info = self.w.volumes.create(
                 catalog_name=schema_info.catalog_name,
                 schema_name=schema_info.name,
@@ -457,7 +537,7 @@ class DatabricksProvider(ServiceProvider):
 
     def create_path(self, volume_path: VolumePathModel) -> Path:
         path: Path = volume_path.full_name
-        logger.info(f"Creating volume path: {path}")
+        logger.info("Creating volume path", path=str(path))
         self.w.files.create_directory(path)
         return path
 
@@ -498,11 +578,12 @@ class DatabricksProvider(ServiceProvider):
 
         if ddl:
             ddl_path: Path = Path(ddl)
-            logger.debug(f"Executing DDL from: {ddl_path}")
+            logger.debug("Executing DDL", ddl_path=str(ddl_path))
             statements: Sequence[str] = sqlparse.parse(ddl_path.read_text())
             for statement in statements:
-                logger.debug(statement)
-                logger.debug(f"args: {args}")
+                logger.trace(
+                    "Executing DDL statement", statement=str(statement)[:100], args=args
+                )
                 spark.sql(
                     str(statement),
                     args=args,
@@ -511,20 +592,23 @@ class DatabricksProvider(ServiceProvider):
         if data:
             data_path: Path = Path(data)
             if format == "sql":
-                logger.debug(f"Executing SQL from: {data_path}")
+                logger.debug("Executing SQL from file", data_path=str(data_path))
                 data_statements: Sequence[str] = sqlparse.parse(data_path.read_text())
                 for statement in data_statements:
-                    logger.debug(statement)
-                    logger.debug(f"args: {args}")
+                    logger.trace(
+                        "Executing SQL statement",
+                        statement=str(statement)[:100],
+                        args=args,
+                    )
                     spark.sql(
                         str(statement),
                         args=args,
                     )
             else:
-                logger.debug(f"Writing to: {table}")
+                logger.debug("Writing dataset to table", table=table)
                 if not data_path.is_absolute():
                     data_path = current_dir / data_path
-                logger.debug(f"Data path: {data_path.as_posix()}")
+                logger.trace("Data path resolved", path=data_path.as_posix())
                 if format == "excel":
                     pdf = pd.read_excel(data_path.as_posix())
                     df = spark.createDataFrame(pdf, schema=dataset.table_schema)
@@ -548,13 +632,17 @@ class DatabricksProvider(ServiceProvider):
                 verbose=True,
             )
 
-        logger.debug(f"Endpoint named {vector_store.endpoint.name} is ready.")
+        logger.success(
+            "Vector search endpoint ready", endpoint_name=vector_store.endpoint.name
+        )
 
         if not index_exists(
             self.vsc, vector_store.endpoint.name, vector_store.index.full_name
         ):
-            logger.debug(
-                f"Creating index {vector_store.index.full_name} on endpoint {vector_store.endpoint.name}..."
+            logger.info(
+                "Creating vector search index",
+                index_name=vector_store.index.full_name,
+                endpoint_name=vector_store.endpoint.name,
             )
             self.vsc.create_delta_sync_index_and_wait(
                 endpoint_name=vector_store.endpoint.name,
@@ -568,7 +656,8 @@ class DatabricksProvider(ServiceProvider):
             )
         else:
             logger.debug(
-                f"Index {vector_store.index.full_name} already exists, checking status and syncing..."
+                "Vector search index already exists, checking status",
+                index_name=vector_store.index.full_name,
             )
             index = self.vsc.get_index(
                 vector_store.endpoint.name, vector_store.index.full_name
@@ -591,54 +680,61 @@ class DatabricksProvider(ServiceProvider):
 
                     if pipeline_status in [
                         "COMPLETED",
+                        "ONLINE",
                         "FAILED",
                         "CANCELED",
                         "ONLINE_PIPELINE_FAILED",
                     ]:
-                        logger.debug(
-                            f"Index is ready to sync (status: {pipeline_status})"
-                        )
+                        logger.debug("Index ready to sync", status=pipeline_status)
                         break
                     elif pipeline_status in [
                         "WAITING_FOR_RESOURCES",
                         "PROVISIONING",
                         "INITIALIZING",
                         "INDEXING",
-                        "ONLINE",
                     ]:
-                        logger.debug(
-                            f"Index not ready yet (status: {pipeline_status}), waiting {wait_interval} seconds..."
+                        logger.trace(
+                            "Index not ready, waiting",
+                            status=pipeline_status,
+                            wait_seconds=wait_interval,
                         )
                         time.sleep(wait_interval)
                         elapsed += wait_interval
                     else:
                         logger.warning(
-                            f"Unknown pipeline status: {pipeline_status}, attempting sync anyway"
+                            "Unknown pipeline status, attempting sync",
+                            status=pipeline_status,
                         )
                         break
                 except Exception as status_error:
                     logger.warning(
-                        f"Could not check index status: {status_error}, attempting sync anyway"
+                        "Could not check index status, attempting sync",
+                        error=str(status_error),
                     )
                     break
 
             if elapsed >= max_wait_time:
                 logger.warning(
-                    f"Timed out waiting for index to be ready after {max_wait_time} seconds"
+                    "Timed out waiting for index to be ready",
+                    max_wait_seconds=max_wait_time,
                 )
 
             # Now attempt to sync
             try:
                 index.sync()
-                logger.debug("Index sync completed successfully")
+                logger.success("Index sync completed")
             except Exception as sync_error:
                 if "not ready to sync yet" in str(sync_error).lower():
-                    logger.warning(f"Index still not ready to sync: {sync_error}")
+                    logger.warning(
+                        "Index still not ready to sync", error=str(sync_error)
+                    )
                 else:
                     raise sync_error
 
-        logger.debug(
-            f"index {vector_store.index.full_name} on table {vector_store.source_table.full_name} is ready"
+        logger.success(
+            "Vector search index ready",
+            index_name=vector_store.index.full_name,
+            source_table=vector_store.source_table.full_name,
         )
 
     def get_vector_index(self, vector_store: VectorStoreModel) -> None:
@@ -674,12 +770,16 @@ class DatabricksProvider(ServiceProvider):
             # sql = sql.replace("{catalog_name}", schema.catalog_name)
             # sql = sql.replace("{schema_name}", schema.schema_name)
 
-            logger.info(function.name)
-            logger.info(sql)
+            logger.info("Creating SQL function", function_name=function.name)
+            logger.trace("SQL function body", sql=sql[:200])
             _: FunctionInfo = self.dfs.create_function(sql_function_body=sql)
 
             if unity_catalog_function.test:
-                logger.info(unity_catalog_function.test.parameters)
+                logger.debug(
+                    "Testing function",
+                    function_name=function.full_name,
+                    parameters=unity_catalog_function.test.parameters,
+                )
 
                 result: FunctionExecutionResult = self.dfs.execute_function(
                     function_name=function.full_name,
@@ -687,37 +787,50 @@ class DatabricksProvider(ServiceProvider):
                 )
 
                 if result.error:
-                    logger.error(result.error)
+                    logger.error(
+                        "Function test failed",
+                        function_name=function.full_name,
+                        error=result.error,
+                    )
                 else:
-                    logger.info(f"Function {function.full_name} executed successfully.")
-                    logger.info(f"Result: {result}")
+                    logger.success(
+                        "Function test passed", function_name=function.full_name
+                    )
+                    logger.debug("Function test result", result=str(result))
 
     def find_columns(self, table_model: TableModel) -> Sequence[str]:
-        logger.debug(f"Finding columns for table: {table_model.full_name}")
+        logger.trace("Finding columns for table", table=table_model.full_name)
         table_info: TableInfo = self.w.tables.get(full_name=table_model.full_name)
         columns: Sequence[ColumnInfo] = table_info.columns
         column_names: Sequence[str] = [c.name for c in columns]
-        logger.debug(f"Columns found: {column_names}")
+        logger.debug(
+            "Columns found",
+            table=table_model.full_name,
+            columns_count=len(column_names),
+        )
         return column_names
 
     def find_primary_key(self, table_model: TableModel) -> Sequence[str] | None:
-        logger.debug(f"Finding primary key for table: {table_model.full_name}")
+        logger.trace("Finding primary key for table", table=table_model.full_name)
         primary_keys: Sequence[str] | None = None
         table_info: TableInfo = self.w.tables.get(full_name=table_model.full_name)
         constraints: Sequence[TableConstraint] = table_info.table_constraints
         primary_key_constraint: PrimaryKeyConstraint | None = next(
-            c.primary_key_constraint for c in constraints if c.primary_key_constraint
+            (c.primary_key_constraint for c in constraints if c.primary_key_constraint),
+            None,
         )
         if primary_key_constraint:
             primary_keys = primary_key_constraint.child_columns
 
-        logger.debug(f"Primary key for table {table_model.full_name}: {primary_keys}")
+        logger.debug(
+            "Primary key found", table=table_model.full_name, primary_keys=primary_keys
+        )
         return primary_keys
 
     def find_vector_search_endpoint(
         self, predicate: Callable[[dict[str, Any]], bool]
     ) -> str | None:
-        logger.debug("Finding vector search endpoint...")
+        logger.trace("Finding vector search endpoint")
         endpoint_name: str | None = None
         vector_search_endpoints: Sequence[dict[str, Any]] = (
             self.vsc.list_endpoints().get("endpoints", [])
@@ -726,11 +839,13 @@ class DatabricksProvider(ServiceProvider):
             if predicate(endpoint):
                 endpoint_name = endpoint["name"]
                 break
-        logger.debug(f"Vector search endpoint found: {endpoint_name}")
+        logger.debug("Vector search endpoint found", endpoint_name=endpoint_name)
         return endpoint_name
 
     def find_endpoint_for_index(self, index_model: IndexModel) -> str | None:
-        logger.debug(f"Finding vector search index: {index_model.full_name}")
+        logger.trace(
+            "Finding endpoint for vector search index", index_name=index_model.full_name
+        )
         all_endpoints: Sequence[dict[str, Any]] = self.vsc.list_endpoints().get(
             "endpoints", []
         )
@@ -740,14 +855,99 @@ class DatabricksProvider(ServiceProvider):
             endpoint_name: str = endpoint["name"]
             indexes = self.vsc.list_indexes(name=endpoint_name)
             vector_indexes: Sequence[dict[str, Any]] = indexes.get("vector_indexes", [])
-            logger.trace(f"Endpoint: {endpoint_name}, vector_indexes: {vector_indexes}")
+            logger.trace(
+                "Checking endpoint for indexes",
+                endpoint_name=endpoint_name,
+                indexes_count=len(vector_indexes),
+            )
             index_names = [vector_index["name"] for vector_index in vector_indexes]
             if index_name in index_names:
                 found_endpoint_name = endpoint_name
                 break
-        logger.debug(f"Vector search index found: {found_endpoint_name}")
+        logger.debug(
+            "Vector search index endpoint found",
+            index_name=index_model.full_name,
+            endpoint_name=found_endpoint_name,
+        )
         return found_endpoint_name
 
+    def _wait_for_database_available(
+        self,
+        workspace_client: WorkspaceClient,
+        instance_name: str,
+        max_wait_time: int = 600,
+        wait_interval: int = 10,
+    ) -> None:
+        """
+        Wait for a database instance to become AVAILABLE.
+
+        Args:
+            workspace_client: The Databricks workspace client
+            instance_name: Name of the database instance to wait for
+            max_wait_time: Maximum time to wait in seconds (default: 600 = 10 minutes)
+            wait_interval: Time between status checks in seconds (default: 10)
+
+        Raises:
+            TimeoutError: If the database doesn't become AVAILABLE within max_wait_time
+            RuntimeError: If the database enters a failed or deleted state
+        """
+        import time
+        from typing import Any
+
+        logger.info(
+            "Waiting for database instance to become AVAILABLE",
+            instance_name=instance_name,
+        )
+        elapsed: int = 0
+
+        while elapsed < max_wait_time:
+            try:
+                current_instance: Any = workspace_client.database.get_database_instance(
+                    name=instance_name
+                )
+                current_state: str = current_instance.state
+                logger.trace(
+                    "Database instance state checked",
+                    instance_name=instance_name,
+                    state=current_state,
+                )
+
+                if current_state == "AVAILABLE":
+                    logger.success(
+                        "Database instance is now AVAILABLE",
+                        instance_name=instance_name,
+                    )
+                    return
+                elif current_state in ["STARTING", "UPDATING", "PROVISIONING"]:
+                    logger.trace(
+                        "Database instance not ready, waiting",
+                        instance_name=instance_name,
+                        state=current_state,
+                        wait_seconds=wait_interval,
+                    )
+                    time.sleep(wait_interval)
+                    elapsed += wait_interval
+                elif current_state in ["STOPPED", "DELETING", "FAILED"]:
+                    raise RuntimeError(
+                        f"Database instance {instance_name} entered unexpected state: {current_state}"
+                    )
+                else:
+                    logger.warning(
+                        "Unknown database state, continuing to wait",
+                        instance_name=instance_name,
+                        state=current_state,
+                    )
+                    time.sleep(wait_interval)
+                    elapsed += wait_interval
+            except NotFound:
+                raise RuntimeError(
+                    f"Database instance {instance_name} was deleted while waiting for it to become AVAILABLE"
+                )
+
+        raise TimeoutError(
+            f"Timed out waiting for database instance {instance_name} to become AVAILABLE after {max_wait_time} seconds"
+        )
+
     def create_lakebase(self, database: DatabaseModel) -> None:
         """
         Create a Lakebase database instance using the Databricks workspace client.
@@ -778,13 +978,17 @@ class DatabricksProvider(ServiceProvider):
 
             if existing_instance:
                 logger.debug(
-                    f"Database instance {database.instance_name} already exists with state: {existing_instance.state}"
+                    "Database instance already exists",
+                    instance_name=database.instance_name,
+                    state=existing_instance.state,
                 )
 
                 # Check if database is in an intermediate state
                 if existing_instance.state in ["STARTING", "UPDATING"]:
                     logger.info(
-                        f"Database instance {database.instance_name} is in {existing_instance.state} state, waiting for it to become AVAILABLE..."
+                        "Database instance in intermediate state, waiting",
+                        instance_name=database.instance_name,
+                        state=existing_instance.state,
                     )
 
                     # Wait for database to reach a stable state
@@ -800,65 +1004,87 @@ class DatabricksProvider(ServiceProvider):
                                 )
                             )
                             current_state: str = current_instance.state
-                            logger.debug(f"Database instance state: {current_state}")
+                            logger.trace(
+                                "Checking database instance state",
+                                instance_name=database.instance_name,
+                                state=current_state,
+                            )
 
                             if current_state == "AVAILABLE":
-                                logger.info(
-                                    f"Database instance {database.instance_name} is now AVAILABLE"
+                                logger.success(
+                                    "Database instance is now AVAILABLE",
+                                    instance_name=database.instance_name,
                                 )
                                 break
                             elif current_state in ["STARTING", "UPDATING"]:
-                                logger.debug(
-                                    f"Database instance still in {current_state} state, waiting {wait_interval} seconds..."
+                                logger.trace(
+                                    "Database instance not ready, waiting",
+                                    instance_name=database.instance_name,
+                                    state=current_state,
+                                    wait_seconds=wait_interval,
                                 )
                                 time.sleep(wait_interval)
                                 elapsed += wait_interval
                             elif current_state in ["STOPPED", "DELETING"]:
                                 logger.warning(
-                                    f"Database instance {database.instance_name} is in unexpected state: {current_state}"
+                                    "Database instance in unexpected state",
+                                    instance_name=database.instance_name,
+                                    state=current_state,
                                 )
                                 break
                             else:
                                 logger.warning(
-                                    f"Unknown database state: {current_state}, proceeding anyway"
+                                    "Unknown database state, proceeding",
+                                    instance_name=database.instance_name,
+                                    state=current_state,
                                 )
                                 break
                         except NotFound:
                             logger.warning(
-                                f"Database instance {database.instance_name} no longer exists, will attempt to recreate"
+                                "Database instance no longer exists, will recreate",
+                                instance_name=database.instance_name,
                             )
                             break
                         except Exception as state_error:
                             logger.warning(
-                                f"Could not check database state: {state_error}, proceeding anyway"
+                                "Could not check database state, proceeding",
+                                instance_name=database.instance_name,
+                                error=str(state_error),
                             )
                             break
 
                     if elapsed >= max_wait_time:
                         logger.warning(
-                            f"Timed out waiting for database instance {database.instance_name} to become AVAILABLE after {max_wait_time} seconds"
+                            "Timed out waiting for database to become AVAILABLE",
+                            instance_name=database.instance_name,
+                            max_wait_seconds=max_wait_time,
                         )
 
                 elif existing_instance.state == "AVAILABLE":
                     logger.info(
-                        f"Database instance {database.instance_name} already exists and is AVAILABLE"
+                        "Database instance already exists and is AVAILABLE",
+                        instance_name=database.instance_name,
                     )
                     return
                 elif existing_instance.state in ["STOPPED", "DELETING"]:
                     logger.warning(
-                        f"Database instance {database.instance_name} is in {existing_instance.state} state"
+                        "Database instance in terminal state",
+                        instance_name=database.instance_name,
+                        state=existing_instance.state,
                     )
                     return
                 else:
                     logger.info(
-                        f"Database instance {database.instance_name} already exists with state: {existing_instance.state}"
+                        "Database instance already exists",
+                        instance_name=database.instance_name,
+                        state=existing_instance.state,
                     )
                     return
 
         except NotFound:
             # Database doesn't exist, proceed with creation
-            logger.debug(
-                f"Database instance {database.instance_name} not found, creating new instance..."
+            logger.info(
+                "Creating new database instance", instance_name=database.instance_name
             )
 
             try:
@@ -878,10 +1104,17 @@ class DatabricksProvider(ServiceProvider):
                 workspace_client.database.create_database_instance(
                     database_instance=database_instance
                 )
-                logger.info(
-                    f"Successfully created database instance: {database.instance_name}"
+                logger.success(
+                    "Database instance created successfully",
+                    instance_name=database.instance_name,
                 )
 
+                # Wait for the newly created database to become AVAILABLE
+                self._wait_for_database_available(
+                    workspace_client, database.instance_name
+                )
+                return
+
             except Exception as create_error:
                 error_msg: str = str(create_error)
 
@@ -891,13 +1124,20 @@ class DatabricksProvider(ServiceProvider):
                     or "RESOURCE_ALREADY_EXISTS" in error_msg
                 ):
                     logger.info(
-                        f"Database instance {database.instance_name} was created concurrently by another process"
+                        "Database instance was created concurrently",
+                        instance_name=database.instance_name,
+                    )
+                    # Still need to wait for the database to become AVAILABLE
+                    self._wait_for_database_available(
+                        workspace_client, database.instance_name
                     )
                     return
                 else:
                     # Re-raise unexpected errors
                     logger.error(
-                        f"Error creating database instance {database.instance_name}: {create_error}"
+                        "Error creating database instance",
+                        instance_name=database.instance_name,
+                        error=str(create_error),
                     )
                     raise
 
@@ -911,12 +1151,15 @@ class DatabricksProvider(ServiceProvider):
                 or "RESOURCE_ALREADY_EXISTS" in error_msg
             ):
                 logger.info(
-                    f"Database instance {database.instance_name} already exists (detected via exception)"
+                    "Database instance already exists (detected via exception)",
+                    instance_name=database.instance_name,
                 )
                 return
             else:
                 logger.error(
-                    f"Unexpected error while handling database {database.instance_name}: {e}"
+                    "Unexpected error while handling database",
+                    instance_name=database.instance_name,
+                    error=str(e),
                 )
                 raise
 
@@ -924,7 +1167,9 @@ class DatabricksProvider(ServiceProvider):
         """
         Ask Databricks to mint a fresh DB credential for this instance.
         """
-        logger.debug(f"Generating password for lakebase instance: {instance_name}")
+        logger.trace(
+            "Generating password for lakebase instance", instance_name=instance_name
+        )
         w: WorkspaceClient = self.w
         cred: DatabaseCredential = w.database.generate_database_credential(
             request_id=str(uuid.uuid4()),
@@ -960,7 +1205,8 @@ class DatabricksProvider(ServiceProvider):
         # Validate that client_id is provided
         if not database.client_id:
             logger.warning(
-                f"client_id is required to create instance role for database {database.instance_name}"
+                "client_id required to create instance role",
+                instance_name=database.instance_name,
             )
             return
 
@@ -970,7 +1216,10 @@ class DatabricksProvider(ServiceProvider):
         instance_name: str = database.instance_name
 
         logger.debug(
-            f"Creating instance role '{role_name}' for database {instance_name} with principal {client_id}"
+            "Creating instance role",
+            role_name=role_name,
+            instance_name=instance_name,
+            principal=client_id,
         )
 
         try:
@@ -981,13 +1230,15 @@ class DatabricksProvider(ServiceProvider):
                     name=role_name,
                 )
                 logger.info(
-                    f"Instance role '{role_name}' already exists for database {instance_name}"
+                    "Instance role already exists",
+                    role_name=role_name,
+                    instance_name=instance_name,
                 )
                 return
             except NotFound:
                 # Role doesn't exist, proceed with creation
                 logger.debug(
-                    f"Instance role '{role_name}' not found, creating new role..."
+                    "Instance role not found, creating new role", role_name=role_name
                 )
 
             # Create the database instance role
@@ -1003,8 +1254,10 @@ class DatabricksProvider(ServiceProvider):
                 database_instance_role=role,
             )
 
-            logger.info(
-                f"Successfully created instance role '{role_name}' for database {instance_name}"
+            logger.success(
+                "Instance role created successfully",
+                role_name=role_name,
+                instance_name=instance_name,
             )
 
         except Exception as e:
@@ -1016,88 +1269,308 @@ class DatabricksProvider(ServiceProvider):
                 or "RESOURCE_ALREADY_EXISTS" in error_msg
             ):
                 logger.info(
-                    f"Instance role '{role_name}' was created concurrently for database {instance_name}"
+                    "Instance role was created concurrently",
+                    role_name=role_name,
+                    instance_name=instance_name,
                 )
                 return
 
             # Re-raise unexpected errors
             logger.error(
-                f"Error creating instance role '{role_name}' for database {instance_name}: {e}"
+                "Error creating instance role",
+                role_name=role_name,
+                instance_name=instance_name,
+                error=str(e),
             )
             raise
 
-    def get_prompt(self, prompt_model: PromptModel) -> str:
-        """Load prompt from MLflow Prompt Registry or fall back to default_template."""
-        prompt_name: str = prompt_model.full_name
+    def get_prompt(self, prompt_model: PromptModel) -> PromptVersion:
+        """
+        Load prompt from MLflow Prompt Registry with fallback logic.
+
+        If an explicit version or alias is specified in the prompt_model, uses that directly.
+        Otherwise, tries to load prompts in this order:
+        1. champion alias
+        2. latest alias
+        3. default alias
+        4. Register default_template if provided (only if register_to_registry=True)
+        5. Use default_template directly (fallback)
+
+        The auto_register field controls whether the default_template is automatically
+        synced to the prompt registry:
+        - If True (default): Auto-registers/updates the default_template in the registry
+        - If False: Never registers, but can still load existing prompts from registry
+                   or use default_template directly as a local-only prompt
 
-        # Build prompt URI based on alias, version, or default to latest
-        if prompt_model.alias:
-            prompt_uri = f"prompts:/{prompt_name}@{prompt_model.alias}"
-        elif prompt_model.version:
-            prompt_uri = f"prompts:/{prompt_name}/{prompt_model.version}"
-        else:
-            prompt_uri = f"prompts:/{prompt_name}@latest"
+        Args:
+            prompt_model: The prompt model configuration
 
-        try:
-            from mlflow.genai.prompts import Prompt
+        Returns:
+            PromptVersion: The loaded prompt version
 
-            prompt_obj: Prompt = mlflow.genai.load_prompt(prompt_uri)
-            return prompt_obj.to_single_brace_format()
+        Raises:
+            ValueError: If no prompt can be loaded from any source
+        """
 
-        except Exception as e:
-            logger.warning(f"Failed to load prompt '{prompt_name}' from registry: {e}")
+        prompt_name: str = prompt_model.full_name
 
-            if prompt_model.default_template:
-                logger.info(f"Using default_template for '{prompt_name}'")
-                self._sync_default_template_to_registry(
-                    prompt_name, prompt_model.default_template, prompt_model.description
+        # If explicit version or alias is specified, use it directly
+        if prompt_model.version or prompt_model.alias:
+            try:
+                prompt_version: PromptVersion = prompt_model.as_prompt()
+                version_or_alias = (
+                    f"version {prompt_model.version}"
+                    if prompt_model.version
+                    else f"alias {prompt_model.alias}"
                 )
-                return prompt_model.default_template
+                logger.debug(
+                    "Loaded prompt with explicit version/alias",
+                    prompt_name=prompt_name,
+                    version_or_alias=version_or_alias,
+                )
+                return prompt_version
+            except Exception as e:
+                version_or_alias = (
+                    f"version {prompt_model.version}"
+                    if prompt_model.version
+                    else f"alias {prompt_model.alias}"
+                )
+                logger.warning(
+                    "Failed to load prompt with explicit version/alias",
+                    prompt_name=prompt_name,
+                    version_or_alias=version_or_alias,
+                    error=str(e),
+                )
+                # Fall through to try other methods
 
-            raise ValueError(
-                f"Prompt '{prompt_name}' not found in registry and no default_template provided"
-            ) from e
+        # Try to load in priority order: champion → default (with sync check)
+        logger.trace(
+            "Trying prompt fallback order",
+            prompt_name=prompt_name,
+            order="champion → default",
+        )
 
-    def _sync_default_template_to_registry(
-        self, prompt_name: str, default_template: str, description: str | None = None
-    ) -> None:
-        """Register default_template to prompt registry under 'default' alias if changed."""
-        try:
-            # Check if default alias already has the same template
+        # First, sync default alias if template has changed (even if champion exists)
+        # Only do this if auto_register is True
+        if prompt_model.default_template and prompt_model.auto_register:
             try:
-                logger.debug(f"Loading prompt '{prompt_name}' from registry...")
-                existing: PromptVersion = mlflow.genai.load_prompt(
-                    f"prompts:/{prompt_name}@default"
-                )
+                # Try to load existing default
+                existing_default = load_prompt(f"prompts:/{prompt_name}@default")
+
+                # Check if champion exists and if it matches default
+                champion_matches_default = False
+                try:
+                    existing_champion = load_prompt(f"prompts:/{prompt_name}@champion")
+                    champion_matches_default = (
+                        existing_champion.version == existing_default.version
+                    )
+                    status = (
+                        "tracking" if champion_matches_default else "pinned separately"
+                    )
+                    logger.trace(
+                        "Champion vs default version",
+                        prompt_name=prompt_name,
+                        champion_version=existing_champion.version,
+                        default_version=existing_default.version,
+                        status=status,
+                    )
+                except Exception:
+                    # No champion exists
+                    logger.trace("No champion alias found", prompt_name=prompt_name)
+
+                # Check if default_template differs from existing default
                 if (
-                    existing.to_single_brace_format().strip()
-                    == default_template.strip()
+                    existing_default.template.strip()
+                    != prompt_model.default_template.strip()
                 ):
-                    logger.debug(f"Prompt '{prompt_name}' is already up-to-date")
-                    return  # Already up-to-date
-            except Exception:
-                logger.debug(
-                    f"Default alias for prompt '{prompt_name}' doesn't exist yet"
+                    logger.info(
+                        "Default template changed, registering new version",
+                        prompt_name=prompt_name,
+                    )
+
+                    # Only update champion if it was pointing to the old default
+                    if champion_matches_default:
+                        logger.info(
+                            "Champion was tracking default, will update to new version",
+                            prompt_name=prompt_name,
+                            old_version=existing_default.version,
+                        )
+                        set_champion = True
+                    else:
+                        logger.info(
+                            "Champion is pinned separately, preserving it",
+                            prompt_name=prompt_name,
+                        )
+                        set_champion = False
+
+                    self._register_default_template(
+                        prompt_name,
+                        prompt_model.default_template,
+                        prompt_model.description,
+                        set_champion=set_champion,
+                    )
+            except Exception as e:
+                # No default exists yet, register it
+                logger.trace(
+                    "No default alias found", prompt_name=prompt_name, error=str(e)
+                )
+                logger.info(
+                    "Registering default template as default alias",
+                    prompt_name=prompt_name,
+                )
+                # First registration - set both default and champion
+                self._register_default_template(
+                    prompt_name,
+                    prompt_model.default_template,
+                    prompt_model.description,
+                    set_champion=True,
+                )
+        elif prompt_model.default_template and not prompt_model.auto_register:
+            logger.trace(
+                "Prompt has auto_register=False, skipping registration",
+                prompt_name=prompt_name,
+            )
+
+        # 1. Try champion alias (highest priority for execution)
+        try:
+            prompt_version = load_prompt(f"prompts:/{prompt_name}@champion")
+            logger.info("Loaded prompt from champion alias", prompt_name=prompt_name)
+            return prompt_version
+        except Exception as e:
+            logger.trace(
+                "Champion alias not found", prompt_name=prompt_name, error=str(e)
+            )
+
+        # 2. Try default alias (already synced above)
+        if prompt_model.default_template:
+            try:
+                prompt_version = load_prompt(f"prompts:/{prompt_name}@default")
+                logger.info("Loaded prompt from default alias", prompt_name=prompt_name)
+                return prompt_version
+            except Exception as e:
+                # Should not happen since we just registered it above, but handle anyway
+                logger.trace(
+                    "Default alias not found", prompt_name=prompt_name, error=str(e)
                 )
 
-            # Register new version and set as default alias
+        # 3. Try latest alias as final fallback
+        try:
+            prompt_version = load_prompt(f"prompts:/{prompt_name}@latest")
+            logger.info("Loaded prompt from latest alias", prompt_name=prompt_name)
+            return prompt_version
+        except Exception as e:
+            logger.trace(
+                "Latest alias not found", prompt_name=prompt_name, error=str(e)
+            )
+
+        # 4. Final fallback: use default_template directly if available
+        if prompt_model.default_template:
+            logger.warning(
+                "Could not load prompt from registry, using default_template directly",
+                prompt_name=prompt_name,
+            )
+            return PromptVersion(
+                name=prompt_name,
+                version=1,
+                template=prompt_model.default_template,
+                tags={"dao_ai": dao_ai_version()},
+            )
+
+        raise ValueError(
+            f"Prompt '{prompt_name}' not found in registry "
+            "(tried champion, default, latest aliases) "
+            "and no default_template provided"
+        )
+
+    def _register_default_template(
+        self,
+        prompt_name: str,
+        default_template: str,
+        description: str | None = None,
+        set_champion: bool = True,
+    ) -> PromptVersion:
+        """Register default_template as a new prompt version.
+
+        Registers the template and sets the 'default' alias.
+        Optionally sets 'champion' alias if no champion exists.
+
+        Args:
+            prompt_name: Full name of the prompt
+            default_template: The template content
+            description: Optional description for commit message
+            set_champion: Whether to also set champion alias (default: True)
+
+        If registration fails (e.g., in Model Serving with restricted permissions),
+        logs the error and raises.
+        """
+        logger.info(
+            "Registering default template",
+            prompt_name=prompt_name,
+            set_champion=set_champion,
+        )
+
+        try:
             commit_message = description or "Auto-synced from default_template"
-            prompt_version: PromptVersion = mlflow.genai.register_prompt(
+            prompt_version = mlflow.genai.register_prompt(
                 name=prompt_name,
                 template=default_template,
                 commit_message=commit_message,
+                tags={"dao_ai": dao_ai_version()},
             )
 
-            logger.debug(f"Setting default alias for prompt '{prompt_name}'")
-            mlflow.genai.set_prompt_alias(
-                name=prompt_name,
-                alias="default",
-                version=prompt_version.version,
-            )
+            # Always set default alias
+            try:
+                logger.debug(
+                    "Setting default alias",
+                    prompt_name=prompt_name,
+                    version=prompt_version.version,
+                )
+                mlflow.genai.set_prompt_alias(
+                    name=prompt_name, alias="default", version=prompt_version.version
+                )
+                logger.success(
+                    "Set default alias for prompt",
+                    prompt_name=prompt_name,
+                    version=prompt_version.version,
+                )
+            except Exception as alias_error:
+                logger.warning(
+                    "Could not set default alias",
+                    prompt_name=prompt_name,
+                    error=str(alias_error),
+                )
 
-            logger.info(
-                f"Synced prompt '{prompt_name}' v{prompt_version.version} to registry"
-            )
+            # Optionally set champion alias (only if no champion exists or explicitly requested)
+            if set_champion:
+                try:
+                    mlflow.genai.set_prompt_alias(
+                        name=prompt_name,
+                        alias="champion",
+                        version=prompt_version.version,
+                    )
+                    logger.success(
+                        "Set champion alias for prompt",
+                        prompt_name=prompt_name,
+                        version=prompt_version.version,
+                    )
+                except Exception as alias_error:
+                    logger.warning(
+                        "Could not set champion alias",
+                        prompt_name=prompt_name,
+                        error=str(alias_error),
+                    )
 
-        except Exception as e:
-            logger.warning(f"Failed to sync '{prompt_name}' to registry: {e}")
+            return prompt_version
+
+        except Exception as reg_error:
+            logger.error(
+                "Failed to register prompt - please register from notebook with write permissions",
+                prompt_name=prompt_name,
+                error=str(reg_error),
+            )
+            return PromptVersion(
+                name=prompt_name,
+                version=1,
+                template=default_template,
+                tags={"dao_ai": dao_ai_version()},
+            )