dao-ai 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

dao_ai/providers/databricks.py

@@ -181,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
@@ -201,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
@@ -220,14 +229,16 @@ 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(
-            f"Using experiment: {experiment.name} (ID: {experiment.experiment_id})"
+            "Using MLflow experiment",
+            experiment_name=experiment.name,
+            experiment_id=experiment.experiment_id,
         )
 
         llms: Sequence[LLMModel] = list(config.resources.llms.values())
@@ -273,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(
@@ -290,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:
@@ -325,24 +344,27 @@ 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(f"target_python_version: {target_python_version}")
+        logger.debug("Target Python version configured", version=target_python_version)
 
         conda_env: dict[str, Any] = {
             "name": "mlflow-env",
@@ -353,7 +375,11 @@ class DatabricksProvider(ServiceProvider):
                 {"pip": list(pip_requirements)},
             ],
         }
-        logger.debug(f"conda_env: {conda_env}")
+        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")
@@ -374,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()
@@ -387,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,
@@ -404,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
@@ -430,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(
@@ -451,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
@@ -472,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
 
@@ -482,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
             )
@@ -494,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,
@@ -505,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
 
@@ -546,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,
@@ -559,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)
@@ -589,6 +625,22 @@ class DatabricksProvider(ServiceProvider):
                 df.write.mode("overwrite").saveAsTable(table)
 
     def create_vector_store(self, vector_store: VectorStoreModel) -> None:
+        # Validate that this is a provisioning-mode config
+        if vector_store.source_table is None:
+            raise ValueError(
+                "Cannot create vector store: source_table is required for provisioning. "
+                "This VectorStoreModel appears to be configured for 'use existing index' mode. "
+                "To provision a new vector store, provide source_table and embedding_source_column."
+            )
+        if vector_store.embedding_source_column is None:
+            raise ValueError(
+                "Cannot create vector store: embedding_source_column is required for provisioning."
+            )
+        if vector_store.endpoint is None:
+            raise ValueError(
+                "Cannot create vector store: endpoint is required for provisioning."
+            )
+
         if not endpoint_exists(self.vsc, vector_store.endpoint.name):
             self.vsc.create_endpoint_and_wait(
                 name=vector_store.endpoint.name,
@@ -596,13 +648,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,
@@ -616,7 +672,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
@@ -639,54 +696,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:
@@ -722,12 +786,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,
@@ -735,37 +803,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", [])
@@ -774,11 +855,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", []
         )
@@ -788,12 +871,20 @@ 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(
@@ -820,7 +911,8 @@ class DatabricksProvider(ServiceProvider):
         from typing import Any
 
         logger.info(
-            f"Waiting for database instance {instance_name} to become AVAILABLE..."
+            "Waiting for database instance to become AVAILABLE",
+            instance_name=instance_name,
         )
         elapsed: int = 0
 
@@ -830,16 +922,24 @@ class DatabricksProvider(ServiceProvider):
                     name=instance_name
                 )
                 current_state: str = current_instance.state
-                logger.debug(
-                    f"Database instance {instance_name} state: {current_state}"
+                logger.trace(
+                    "Database instance state checked",
+                    instance_name=instance_name,
+                    state=current_state,
                 )
 
                 if current_state == "AVAILABLE":
-                    logger.info(f"Database instance {instance_name} is now AVAILABLE")
+                    logger.success(
+                        "Database instance is now AVAILABLE",
+                        instance_name=instance_name,
+                    )
                     return
                 elif current_state in ["STARTING", "UPDATING", "PROVISIONING"]:
-                    logger.debug(
-                        f"Database instance still in {current_state} state, waiting {wait_interval} seconds..."
+                    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
@@ -849,7 +949,9 @@ class DatabricksProvider(ServiceProvider):
                     )
                 else:
                     logger.warning(
-                        f"Unknown database state: {current_state}, continuing to wait..."
+                        "Unknown database state, continuing to wait",
+                        instance_name=instance_name,
+                        state=current_state,
                     )
                     time.sleep(wait_interval)
                     elapsed += wait_interval
@@ -892,13 +994,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
@@ -914,65 +1020,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:
@@ -992,8 +1120,9 @@ 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
@@ -1011,7 +1140,8 @@ 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(
@@ -1021,7 +1151,9 @@ class DatabricksProvider(ServiceProvider):
                 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
 
@@ -1035,12 +1167,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
 
@@ -1048,7 +1183,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()),
@@ -1084,7 +1221,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
 
@@ -1094,7 +1232,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:
@@ -1105,13 +1246,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
@@ -1127,8 +1270,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:
@@ -1140,13 +1285,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
 
@@ -1159,7 +1309,14 @@ class DatabricksProvider(ServiceProvider):
         1. champion alias
         2. latest alias
         3. default alias
-        4. Register default_template if provided
+        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
@@ -1177,25 +1334,41 @@ class DatabricksProvider(ServiceProvider):
         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 try other methods
 
         # Try to load in priority order: champion → default (with sync check)
-        logger.debug(
-            f"Trying fallback order for '{prompt_name}': champion → default (with auto-sync)"
+        logger.trace(
+            "Trying prompt fallback order",
+            prompt_name=prompt_name,
+            order="champion → default",
         )
 
         # First, sync default alias if template has changed (even if champion exists)
-        if prompt_model.default_template:
+        # Only do this if auto_register is True
+        if prompt_model.default_template and prompt_model.auto_register:
             try:
                 # Try to load existing default
                 existing_default = load_prompt(f"prompts:/{prompt_name}@default")
@@ -1207,13 +1380,19 @@ class DatabricksProvider(ServiceProvider):
                     champion_matches_default = (
                         existing_champion.version == existing_default.version
                     )
-                    logger.debug(
-                        f"Champion v{existing_champion.version} vs Default v{existing_default.version}: "
-                        f"{'tracking' if champion_matches_default else 'pinned separately'}"
+                    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.debug(f"No champion alias found for '{prompt_name}'")
+                    logger.trace("No champion alias found", prompt_name=prompt_name)
 
                 # Check if default_template differs from existing default
                 if (
@@ -1221,19 +1400,23 @@ class DatabricksProvider(ServiceProvider):
                     != prompt_model.default_template.strip()
                 ):
                     logger.info(
-                        f"Default template for '{prompt_name}' has changed, "
-                        "registering new version with default alias"
+                        "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(
-                            f"Champion was tracking default (v{existing_default.version}), "
-                            "will update champion to new default version"
+                            "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")
+                        logger.info(
+                            "Champion is pinned separately, preserving it",
+                            prompt_name=prompt_name,
+                        )
                         set_champion = False
 
                     self._register_default_template(
@@ -1244,9 +1427,12 @@ class DatabricksProvider(ServiceProvider):
                     )
             except Exception as e:
                 # No default exists yet, register it
-                logger.debug(f"No default alias found for '{prompt_name}': {e}")
+                logger.trace(
+                    "No default alias found", prompt_name=prompt_name, error=str(e)
+                )
                 logger.info(
-                    f"Registering default_template for '{prompt_name}' as default alias"
+                    "Registering default template as default alias",
+                    prompt_name=prompt_name,
                 )
                 # First registration - set both default and champion
                 self._register_default_template(
@@ -1255,40 +1441,49 @@ class DatabricksProvider(ServiceProvider):
                     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(
-                f"Loaded prompt '{prompt_name}' from champion alias (default was synced separately)"
-            )
+            logger.info("Loaded prompt from champion alias", prompt_name=prompt_name)
             return prompt_version
         except Exception as e:
-            logger.debug(f"Champion alias not found for '{prompt_name}': {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(f"Loaded prompt '{prompt_name}' from default alias")
+                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.debug(f"Default alias not found for '{prompt_name}': {e}")
+                logger.trace(
+                    "Default alias not found", prompt_name=prompt_name, error=str(e)
+                )
 
         # 3. Try latest alias as final fallback
         try:
             prompt_version = load_prompt(f"prompts:/{prompt_name}@latest")
-            logger.info(f"Loaded prompt '{prompt_name}' from latest alias")
+            logger.info("Loaded prompt from latest alias", prompt_name=prompt_name)
             return prompt_version
         except Exception as e:
-            logger.debug(f"Latest alias not found for '{prompt_name}': {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_name}' from registry. "
-                "Using default_template directly (likely in test environment)"
+                "Could not load prompt from registry, using default_template directly",
+                prompt_name=prompt_name,
             )
             return PromptVersion(
                 name=prompt_name,
@@ -1325,8 +1520,9 @@ class DatabricksProvider(ServiceProvider):
         logs the error and raises.
         """
         logger.info(
-            f"Registering default_template for '{prompt_name}' "
-            f"(set_champion={set_champion})"
+            "Registering default template",
+            prompt_name=prompt_name,
+            set_champion=set_champion,
         )
 
         try:
@@ -1340,15 +1536,24 @@ class DatabricksProvider(ServiceProvider):
 
             # 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.info(
-                    f"Set default alias for '{prompt_name}' v{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(
-                    f"Could not set default alias for '{prompt_name}': {alias_error}"
+                    "Could not set default alias",
+                    prompt_name=prompt_name,
+                    error=str(alias_error),
                 )
 
             # Optionally set champion alias (only if no champion exists or explicitly requested)
@@ -1359,20 +1564,25 @@ class DatabricksProvider(ServiceProvider):
                         alias="champion",
                         version=prompt_version.version,
                     )
-                    logger.info(
-                        f"Set champion alias for '{prompt_name}' v{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(
-                        f"Could not set champion alias for '{prompt_name}': {alias_error}"
+                        "Could not set champion alias",
+                        prompt_name=prompt_name,
+                        error=str(alias_error),
                     )
 
             return prompt_version
 
         except Exception as reg_error:
             logger.error(
-                f"Failed to register prompt '{prompt_name}': {reg_error}. "
-                f"Please register the prompt from a notebook with write permissions before deployment."
+                "Failed to register prompt - please register from notebook with write permissions",
+                prompt_name=prompt_name,
+                error=str(reg_error),
             )
             return PromptVersion(
                 name=prompt_name,