dao-ai 0.0.28__py3-none-any.whl → 0.1.5__py3-none-any.whl

dao_ai/providers/databricks.py

@@ -1,5 +1,4 @@
 import base64
-import re
 import uuid
 from pathlib import Path
 from typing import Any, Callable, Final, Sequence
@@ -32,14 +31,12 @@ 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.datasets import EvaluationDataset, get_dataset
 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 (
     DatabricksResource,
 )
-from mlflow.pyfunc import ResponsesAgent
 from pyspark.sql import SparkSession
 from unitycatalog.ai.core.base import FunctionExecutionResult
 from unitycatalog.ai.core.databricks import DatabricksFunctionClient
@@ -49,6 +46,7 @@ from dao_ai.config import (
     AppConfig,
     ConnectionModel,
     DatabaseModel,
+    DatabricksAppModel,
     DatasetModel,
     FunctionModel,
     GenieRoomModel,
@@ -57,7 +55,6 @@ from dao_ai.config import (
     IsDatabricksResource,
     LLMModel,
     PromptModel,
-    PromptOptimizationModel,
     SchemaModel,
     TableModel,
     UnityCatalogFunctionSqlModel,
@@ -73,6 +70,7 @@ from dao_ai.utils import (
     get_installed_packages,
     is_installed,
     is_lib_provided,
+    normalize_host,
     normalize_name,
 )
 from dao_ai.vector_search import endpoint_exists, index_exists
@@ -94,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()
 
@@ -117,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:
@@ -177,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
@@ -197,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
@@ -216,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()
@@ -236,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
@@ -247,6 +270,7 @@ class DatabricksProvider(ServiceProvider):
             + connections
             + databases
             + volumes
+            + apps
         )
 
         # Flatten all resources from all models into a single list
@@ -260,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(
@@ -277,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:
@@ -312,18 +344,42 @@ 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")
@@ -333,7 +389,7 @@ class DatabricksProvider(ServiceProvider):
                 code_paths=code_paths,
                 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,
@@ -344,8 +400,10 @@ 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()
@@ -357,7 +415,7 @@ class DatabricksProvider(ServiceProvider):
             key="dao_ai",
             value=dao_ai_version(),
         )
-        logger.debug(f"Set dao_ai tag on model version {model_version.version}")
+        logger.trace("Set dao_ai tag on model version", version=model_version.version)
 
         client.set_registered_model_alias(
             name=registered_model_name,
@@ -374,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
@@ -400,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(
@@ -421,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
@@ -442,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
 
@@ -452,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
             )
@@ -464,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,
@@ -475,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
 
@@ -516,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,
@@ -529,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)
@@ -559,6 +625,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,
@@ -566,13 +643,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,
@@ -586,7 +667,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
@@ -609,54 +691,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:
@@ -692,12 +781,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,
@@ -705,37 +798,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", [])
@@ -744,11 +850,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", []
         )
@@ -758,14 +866,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.
@@ -796,13 +989,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
@@ -818,65 +1015,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:
@@ -896,10 +1115,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)
 
@@ -909,13 +1135,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
 
@@ -929,12 +1162,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
 
@@ -942,7 +1178,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()),
@@ -978,7 +1216,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
 
@@ -988,7 +1227,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:
@@ -999,13 +1241,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
@@ -1021,8 +1265,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:
@@ -1034,13 +1280,18 @@ 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
 
@@ -1050,9 +1301,17 @@ class DatabricksProvider(ServiceProvider):
 
         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 (if it exists)
-        2. latest alias (if it exists)
-        3. default_template (if provided)
+        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
 
         Args:
             prompt_model: The prompt model configuration
@@ -1063,542 +1322,266 @@ class DatabricksProvider(ServiceProvider):
         Raises:
             ValueError: If no prompt can be loaded from any source
         """
+
         prompt_name: str = prompt_model.full_name
 
-        # If explicit version or alias is specified, use it directly without fallback
+        # 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}"
+                )
                 logger.debug(
-                    f"Loaded prompt '{prompt_name}' with explicit "
-                    f"{'version ' + str(prompt_model.version) if prompt_model.version else 'alias ' + prompt_model.alias}"
+                    "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(
-                    f"Failed to load prompt '{prompt_name}' with explicit "
-                    f"{'version ' + str(prompt_model.version) if prompt_model.version else 'alias ' + prompt_model.alias}: {e}"
+                    "Failed to load prompt with explicit version/alias",
+                    prompt_name=prompt_name,
+                    version_or_alias=version_or_alias,
+                    error=str(e),
                 )
-                # Fall through to default_template if available
-        else:
-            # No explicit version/alias specified - check if default_template needs syncing first
-            logger.debug(
-                f"No explicit version/alias specified for '{prompt_name}', "
-                "checking if default_template needs syncing"
-            )
-
-            # If we have a default_template, check if it differs from what's in the registry
-            # This ensures we always sync config changes before returning any alias
-            if prompt_model.default_template:
-                try:
-                    default_uri: str = f"prompts:/{prompt_name}@default"
-                    default_version: PromptVersion = load_prompt(default_uri)
+                # Fall through to try other methods
 
-                    if (
-                        default_version.to_single_brace_format().strip()
-                        != prompt_model.default_template.strip()
-                    ):
-                        logger.info(
-                            f"Config default_template for '{prompt_name}' differs from registry, syncing..."
-                        )
-                        return self._sync_default_template_to_registry(
-                            prompt_name,
-                            prompt_model.default_template,
-                            prompt_model.description,
-                        )
-                except Exception as e:
-                    logger.debug(f"Could not check default alias for sync: {e}")
-
-            # Now try aliases in order: champion → latest → default
-            logger.debug(
-                f"Trying fallback order for '{prompt_name}': champion → latest → default"
-            )
-
-            # Try champion alias first
-            try:
-                champion_uri: str = f"prompts:/{prompt_name}@champion"
-                prompt_version: PromptVersion = load_prompt(champion_uri)
-                logger.info(f"Loaded prompt '{prompt_name}' from champion alias")
-                return prompt_version
-            except Exception as e:
-                logger.debug(f"Champion alias not found for '{prompt_name}': {e}")
-
-            # Try latest alias next
-            try:
-                latest_uri: str = f"prompts:/{prompt_name}@latest"
-                prompt_version: PromptVersion = load_prompt(latest_uri)
-                logger.info(f"Loaded prompt '{prompt_name}' from latest alias")
-                return prompt_version
-            except Exception as e:
-                logger.debug(f"Latest alias not found for '{prompt_name}': {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",
+        )
 
-            # Try default alias last
+        # 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:
-                default_uri: str = f"prompts:/{prompt_name}@default"
-                prompt_version: PromptVersion = load_prompt(default_uri)
-                logger.info(f"Loaded prompt '{prompt_name}' from default alias")
-                return prompt_version
-            except Exception as e:
-                logger.debug(f"Default alias not found for '{prompt_name}': {e}")
+                # Try to load existing default
+                existing_default = load_prompt(f"prompts:/{prompt_name}@default")
 
-        # Fall back to registering default_template if provided
-        if prompt_model.default_template:
-            logger.info(
-                f"Registering default_template for '{prompt_name}' "
-                "(no aliases found in registry)"
-            )
-            return self._sync_default_template_to_registry(
-                prompt_name, prompt_model.default_template, prompt_model.description
-            )
-
-        raise ValueError(
-            f"Prompt '{prompt_name}' not found in registry "
-            "(tried champion, latest, default aliases) and no default_template provided"
-        )
-
-    def _sync_default_template_to_registry(
-        self, prompt_name: str, default_template: str, description: str | None = None
-    ) -> PromptVersion:
-        """Register default_template to prompt registry under 'default' alias if changed."""
-        prompt_version: PromptVersion | None = None
+                # 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)
 
-        try:
-            # Check if default alias already has the same template
-            try:
-                logger.debug(f"Loading prompt '{prompt_name}' from registry...")
-                existing: PromptVersion = mlflow.genai.load_prompt(
-                    f"prompts:/{prompt_name}@default"
-                )
+                # 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")
+                    logger.info(
+                        "Default template changed, registering new version",
+                        prompt_name=prompt_name,
+                    )
 
-                    # Ensure the "latest" and "champion" aliases also exist and point to the same version
-                    # This handles prompts created before the fix that added these aliases
-                    try:
-                        latest_version: PromptVersion = mlflow.genai.load_prompt(
-                            f"prompts:/{prompt_name}@latest"
-                        )
-                        logger.debug(
-                            f"Latest alias already exists for '{prompt_name}' pointing to version {latest_version.version}"
-                        )
-                    except Exception:
+                    # Only update champion if it was pointing to the old default
+                    if champion_matches_default:
                         logger.info(
-                            f"Setting 'latest' alias for existing prompt '{prompt_name}' v{existing.version}"
+                            "Champion was tracking default, will update to new version",
+                            prompt_name=prompt_name,
+                            old_version=existing_default.version,
                         )
-                        mlflow.genai.set_prompt_alias(
-                            name=prompt_name,
-                            alias="latest",
-                            version=existing.version,
-                        )
-
-                    # Ensure champion alias exists for first-time deployments
-                    try:
-                        champion_version: PromptVersion = mlflow.genai.load_prompt(
-                            f"prompts:/{prompt_name}@champion"
-                        )
-                        logger.debug(
-                            f"Champion alias already exists for '{prompt_name}' pointing to version {champion_version.version}"
-                        )
-                    except Exception:
+                        set_champion = True
+                    else:
                         logger.info(
-                            f"Setting 'champion' alias for existing prompt '{prompt_name}' v{existing.version}"
-                        )
-                        mlflow.genai.set_prompt_alias(
-                            name=prompt_name,
-                            alias="champion",
-                            version=existing.version,
+                            "Champion is pinned separately, preserving it",
+                            prompt_name=prompt_name,
                         )
+                        set_champion = False
 
-                    return existing  # Already up-to-date, return existing version
-            except Exception:
-                logger.debug(
-                    f"Default alias for prompt '{prompt_name}' doesn't exist yet"
+                    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)
                 )
-
-            # Register new version and set as default alias
-            commit_message = description or "Auto-synced from default_template"
-            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, latest, and champion aliases for prompt '{prompt_name}'"
-            )
-            mlflow.genai.set_prompt_alias(
-                name=prompt_name,
-                alias="default",
-                version=prompt_version.version,
-            )
-            mlflow.genai.set_prompt_alias(
-                name=prompt_name,
-                alias="latest",
-                version=prompt_version.version,
-            )
-            mlflow.genai.set_prompt_alias(
-                name=prompt_name,
-                alias="champion",
-                version=prompt_version.version,
+                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,
             )
 
-            logger.info(
-                f"Synced prompt '{prompt_name}' v{prompt_version.version} to registry with 'default', 'latest', and 'champion' aliases"
-            )
+        # 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.error(f"Failed to sync '{prompt_name}' to registry: {e}")
-            raise ValueError(
-                f"Failed to sync prompt '{prompt_name}' to registry and unable to retrieve existing version"
-            ) from e
-
-    def optimize_prompt(self, optimization: PromptOptimizationModel) -> PromptModel:
-        """
-        Optimize a prompt using MLflow's prompt optimization (MLflow 3.5+).
-
-        This uses the MLflow GenAI optimize_prompts API with GepaPromptOptimizer as documented at:
-        https://mlflow.org/docs/latest/genai/prompt-registry/optimize-prompts/
+            logger.trace(
+                "Champion alias not found", prompt_name=prompt_name, error=str(e)
+            )
 
-        Args:
-            optimization: PromptOptimizationModel containing configuration
+        # 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)
+                )
 
-        Returns:
-            PromptModel: The optimized prompt with new URI
-        """
-        from mlflow.genai.optimize import GepaPromptOptimizer, optimize_prompts
-        from mlflow.genai.scorers import Correctness
-
-        from dao_ai.config import AgentModel, PromptModel
-
-        logger.info(f"Optimizing prompt: {optimization.name}")
-
-        # Get agent and prompt (prompt is guaranteed to be set by validator)
-        agent_model: AgentModel = optimization.agent
-        prompt: PromptModel = optimization.prompt  # type: ignore[assignment]
-        agent_model.prompt = prompt.uri
-
-        print(f"prompt={agent_model.prompt}")
-        # Log the prompt URI scheme being used
-        # Supports three schemes:
-        # 1. Specific version: "prompts:/qa/1" (when version is specified)
-        # 2. Alias: "prompts:/qa@champion" (when alias is specified)
-        # 3. Latest: "prompts:/qa@latest" (default when neither version nor alias specified)
-        prompt_uri: str = prompt.uri
-        logger.info(f"Using prompt URI for optimization: {prompt_uri}")
-
-        # Load the specific prompt version by URI for comparison
-        # Try to load the exact version specified, but if it doesn't exist,
-        # use get_prompt to create it from default_template
-        prompt_version: PromptVersion
+        # 3. Try latest alias as final fallback
         try:
-            prompt_version = load_prompt(prompt_uri)
-            logger.info(f"Successfully loaded prompt from registry: {prompt_uri}")
+            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(
-                f"Could not load prompt '{prompt_uri}' directly: {e}. "
-                "Attempting to create from default_template..."
+                "Could not load prompt from registry, using default_template directly",
+                prompt_name=prompt_name,
             )
-            # Use get_prompt which will create from default_template if needed
-            prompt_version = self.get_prompt(prompt)
-            logger.info(
-                f"Created/loaded prompt '{prompt.full_name}' (will optimize against this version)"
+            return PromptVersion(
+                name=prompt_name,
+                version=1,
+                template=prompt_model.default_template,
+                tags={"dao_ai": dao_ai_version()},
             )
 
-        # Load the evaluation dataset by name
-        logger.debug(f"Looking up dataset: {optimization.dataset}")
-        dataset: EvaluationDataset
-        if isinstance(optimization.dataset, str):
-            dataset = get_dataset(name=optimization.dataset)
-        else:
-            dataset = optimization.dataset.as_dataset()
-
-        # Set up reflection model for the optimizer
-        reflection_model_name: str
-        if optimization.reflection_model:
-            if isinstance(optimization.reflection_model, str):
-                reflection_model_name = optimization.reflection_model
-            else:
-                reflection_model_name = optimization.reflection_model.uri
-        else:
-            reflection_model_name = agent_model.model.uri
-        logger.debug(f"Using reflection model: {reflection_model_name}")
-
-        # Create the GepaPromptOptimizer
-        optimizer: GepaPromptOptimizer = GepaPromptOptimizer(
-            reflection_model=reflection_model_name,
-            max_metric_calls=optimization.num_candidates,
-            display_progress_bar=True,
+        raise ValueError(
+            f"Prompt '{prompt_name}' not found in registry "
+            "(tried champion, default, latest aliases) "
+            "and no default_template provided"
         )
 
-        # Set up scorer (judge model for evaluation)
-        scorer_model: str
-        if optimization.scorer_model:
-            if isinstance(optimization.scorer_model, str):
-                scorer_model = optimization.scorer_model
-            else:
-                scorer_model = optimization.scorer_model.uri
-        else:
-            scorer_model = agent_model.model.uri  # Use Databricks default
-        logger.debug(f"Using scorer with model: {scorer_model}")
-
-        scorers: list[Correctness] = [Correctness(model=scorer_model)]
-
-        # Use prompt_uri from line 1188 - already set to prompt.uri (configured URI)
-        # DO NOT overwrite with prompt_version.uri as that uses fallback logic
-        logger.debug(f"Optimizing prompt: {prompt_uri}")
-
-        agent: ResponsesAgent = agent_model.as_responses_agent()
-
-        # Create predict function that will be optimized
-        def predict_fn(**inputs: dict[str, Any]) -> str:
-            """Prediction function that uses the ResponsesAgent with ChatPayload.
-
-            The agent already has the prompt referenced/applied, so we just need to
-            convert the ChatPayload inputs to ResponsesAgentRequest format and call predict.
-
-            Args:
-                **inputs: Dictionary containing ChatPayload fields (messages/input, custom_inputs)
-
-            Returns:
-                str: The agent's response content
-            """
-            from mlflow.types.responses import (
-                ResponsesAgentRequest,
-                ResponsesAgentResponse,
-            )
-            from mlflow.types.responses_helpers import Message
-
-            from dao_ai.config import ChatPayload
-
-            # Verify agent is accessible (should be captured from outer scope)
-            if agent is None:
-                raise RuntimeError(
-                    "Agent object is not available in predict_fn. "
-                    "This may indicate a serialization issue with the ResponsesAgent."
-                )
-
-            # Convert inputs to ChatPayload
-            chat_payload: ChatPayload = ChatPayload(**inputs)
-
-            # Convert ChatPayload messages to MLflow Message format
-            mlflow_messages: list[Message] = [
-                Message(role=msg.role, content=msg.content)
-                for msg in chat_payload.messages
-            ]
-
-            # Create ResponsesAgentRequest
-            request: ResponsesAgentRequest = ResponsesAgentRequest(
-                input=mlflow_messages,
-                custom_inputs=chat_payload.custom_inputs,
-            )
+    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.
 
-            # Call the ResponsesAgent's predict method
-            response: ResponsesAgentResponse = agent.predict(request)
-
-            if response.output and len(response.output) > 0:
-                content = response.output[0].content
-                logger.debug(f"Response content type: {type(content)}")
-                logger.debug(f"Response content: {content}")
-
-                # Extract text from content using same logic as LanggraphResponsesAgent._extract_text_from_content
-                # Content can be:
-                # - A string (return as is)
-                # - A list of items with 'text' keys (extract and join)
-                # - Other types (try to get 'text' attribute or convert to string)
-                if isinstance(content, str):
-                    return content
-                elif isinstance(content, list):
-                    text_parts = []
-                    for content_item in content:
-                        if isinstance(content_item, str):
-                            text_parts.append(content_item)
-                        elif isinstance(content_item, dict) and "text" in content_item:
-                            text_parts.append(content_item["text"])
-                        elif hasattr(content_item, "text"):
-                            text_parts.append(content_item.text)
-                    return "".join(text_parts) if text_parts else str(content)
-                else:
-                    # Fallback for unknown types - try to extract text attribute
-                    return getattr(content, "text", str(content))
-            else:
-                return ""
+        Registers the template and sets the 'default' alias.
+        Optionally sets 'champion' alias if no champion exists.
 
-        # Set registry URI for Databricks Unity Catalog
-        mlflow.set_registry_uri("databricks-uc")
+        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)
 
-        # Run optimization with tracking disabled to prevent auto-registering all candidates
-        logger.info("Running prompt optimization with GepaPromptOptimizer...")
+        If registration fails (e.g., in Model Serving with restricted permissions),
+        logs the error and raises.
+        """
         logger.info(
-            f"Generating {optimization.num_candidates} candidate prompts for evaluation"
-        )
-
-        from mlflow.genai.optimize.types import (
-            PromptOptimizationResult,
-        )
-
-        result: PromptOptimizationResult = optimize_prompts(
-            predict_fn=predict_fn,
-            train_data=dataset,
-            prompt_uris=[prompt_uri],  # Use the configured URI (version/alias/latest)
-            optimizer=optimizer,
-            scorers=scorers,
-            enable_tracking=False,  # Don't auto-register all candidates
+            "Registering default template",
+            prompt_name=prompt_name,
+            set_champion=set_champion,
         )
 
-        # Log the optimization results
-        logger.info("Optimization complete!")
-        logger.info(f"Optimizer used: {result.optimizer_name}")
-
-        if result.optimized_prompts:
-            optimized_prompt_version: PromptVersion = result.optimized_prompts[0]
-
-            # Check if the optimized prompt is actually different from the original
-            original_template: str = prompt_version.to_single_brace_format().strip()
-            optimized_template: str = (
-                optimized_prompt_version.to_single_brace_format().strip()
-            )
-
-            # Normalize whitespace for more robust comparison
-            original_normalized: str = re.sub(r"\s+", " ", original_template).strip()
-            optimized_normalized: str = re.sub(r"\s+", " ", optimized_template).strip()
-
-            logger.debug(f"Original template length: {len(original_template)} chars")
-            logger.debug(f"Optimized template length: {len(optimized_template)} chars")
-            logger.debug(
-                f"Templates identical: {original_normalized == optimized_normalized}"
-            )
-
-            if original_normalized == optimized_normalized:
-                logger.info(
-                    f"Optimized prompt is identical to original for '{prompt.full_name}'. "
-                    "No new version will be registered."
-                )
-                return prompt
-
-            logger.info("Optimized prompt is DIFFERENT from original")
-            logger.info(
-                f"Original length: {len(original_template)}, Optimized length: {len(optimized_template)}"
-            )
-            logger.debug(
-                f"Original template (first 300 chars): {original_template[:300]}..."
-            )
-            logger.debug(
-                f"Optimized template (first 300 chars): {optimized_template[:300]}..."
+        try:
+            commit_message = description or "Auto-synced from default_template"
+            prompt_version = mlflow.genai.register_prompt(
+                name=prompt_name,
+                template=default_template,
+                commit_message=commit_message,
+                tags={"dao_ai": dao_ai_version()},
             )
 
-            # Check evaluation scores to determine if we should register the optimized prompt
-            should_register: bool = False
-            has_improvement: bool = False
-
-            if (
-                result.initial_eval_score is not None
-                and result.final_eval_score is not None
-            ):
-                logger.info("Evaluation scores:")
-                logger.info(f"  Initial score: {result.initial_eval_score}")
-                logger.info(f"  Final score: {result.final_eval_score}")
-
-                # Only register if there's improvement
-                if result.final_eval_score > result.initial_eval_score:
-                    improvement: float = (
-                        (result.final_eval_score - result.initial_eval_score)
-                        / result.initial_eval_score
-                    ) * 100
-                    logger.info(
-                        f"Optimized prompt improved by {improvement:.2f}% "
-                        f"(initial: {result.initial_eval_score}, final: {result.final_eval_score})"
-                    )
-                    should_register = True
-                    has_improvement = True
-                else:
-                    logger.info(
-                        f"Optimized prompt (score: {result.final_eval_score}) did NOT improve over baseline (score: {result.initial_eval_score}). "
-                        "No new version will be registered."
-                    )
-            else:
-                # No scores available - register anyway but warn
-                logger.warning(
-                    "No evaluation scores available to compare performance. "
-                    "Registering optimized prompt without performance validation."
-                )
-                should_register = True
-
-            if not should_register:
-                logger.info(
-                    f"Skipping registration for '{prompt.full_name}' (no improvement)"
-                )
-                return prompt
-
-            # Register the optimized prompt manually
+            # Always set default alias
             try:
-                logger.info(f"Registering optimized prompt '{prompt.full_name}'")
-                registered_version: PromptVersion = mlflow.genai.register_prompt(
-                    name=prompt.full_name,
-                    template=optimized_template,
-                    commit_message=f"Optimized for {agent_model.model.uri} using GepaPromptOptimizer",
-                    tags={
-                        "dao_ai": dao_ai_version(),
-                        "target_model": agent_model.model.uri,
-                    },
-                )
-                logger.info(
-                    f"Registered optimized prompt as version {registered_version.version}"
-                )
-
-                # Always set "latest" alias (represents most recently registered prompt)
-                logger.info(
-                    f"Setting 'latest' alias for optimized prompt '{prompt.full_name}' version {registered_version.version}"
+                logger.debug(
+                    "Setting default alias",
+                    prompt_name=prompt_name,
+                    version=prompt_version.version,
                 )
                 mlflow.genai.set_prompt_alias(
-                    name=prompt.full_name,
-                    alias="latest",
-                    version=registered_version.version,
+                    name=prompt_name, alias="default", version=prompt_version.version
                 )
-                logger.info(
-                    f"Successfully set 'latest' alias for '{prompt.full_name}' v{registered_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),
                 )
 
-                # If there's confirmed improvement, also set the "champion" alias
-                # (represents the prompt that should be used by deployed agents)
-                if has_improvement:
-                    logger.info(
-                        f"Setting 'champion' alias for improved prompt '{prompt.full_name}' version {registered_version.version}"
-                    )
+            # Optionally set champion alias (only if no champion exists or explicitly requested)
+            if set_champion:
+                try:
                     mlflow.genai.set_prompt_alias(
-                        name=prompt.full_name,
+                        name=prompt_name,
                         alias="champion",
-                        version=registered_version.version,
+                        version=prompt_version.version,
                     )
-                    logger.info(
-                        f"Successfully set 'champion' alias for '{prompt.full_name}' v{registered_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),
                     )
 
-                # Add target_model and dao_ai tags
-                tags: dict[str, Any] = prompt.tags.copy() if prompt.tags else {}
-                tags["target_model"] = agent_model.model.uri
-                tags["dao_ai"] = dao_ai_version()
-
-                # Return the optimized prompt with the appropriate alias
-                # Use "champion" if there was improvement, otherwise "latest"
-                result_alias: str = "champion" if has_improvement else "latest"
-                return PromptModel(
-                    name=prompt.name,
-                    schema=prompt.schema_model,
-                    description=f"Optimized version of {prompt.name} for {agent_model.model.uri}",
-                    alias=result_alias,
-                    tags=tags,
-                )
+            return prompt_version
 
-            except Exception as e:
-                logger.error(
-                    f"Failed to register optimized prompt '{prompt.full_name}': {e}"
-                )
-                return prompt
-        else:
-            logger.warning("No optimized prompts returned from optimization")
-            return prompt
+        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()},
+            )