dao-ai 0.0.31__py3-none-any.whl → 0.0.33__py3-none-any.whl

dao_ai/providers/databricks.py

@@ -332,6 +332,23 @@ class DatabricksProvider(ServiceProvider):
 
         logger.debug(f"input_example: {input_example}")
 
+        # 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}")
+
+        conda_env: dict[str, Any] = {
+            "name": "mlflow-env",
+            "channels": ["conda-forge"],
+            "dependencies": [
+                f"python={target_python_version}",
+                "pip",
+                {"pip": list(pip_requirements)},
+            ],
+        }
+        logger.debug(f"conda_env: {conda_env}")
+
         with mlflow.start_run(run_name=run_name):
             mlflow.set_tag("type", "agent")
             mlflow.set_tag("dao_ai", dao_ai_version())
@@ -340,7 +357,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,
@@ -773,6 +790,72 @@ class DatabricksProvider(ServiceProvider):
         logger.debug(f"Vector search index found: {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(
+            f"Waiting for database instance {instance_name} to become AVAILABLE..."
+        )
+        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.debug(
+                    f"Database instance {instance_name} state: {current_state}"
+                )
+
+                if current_state == "AVAILABLE":
+                    logger.info(f"Database instance {instance_name} is now AVAILABLE")
+                    return
+                elif current_state in ["STARTING", "UPDATING", "PROVISIONING"]:
+                    logger.debug(
+                        f"Database instance still in {current_state} state, waiting {wait_interval} seconds..."
+                    )
+                    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(
+                        f"Unknown database state: {current_state}, continuing to wait..."
+                    )
+                    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.
@@ -907,6 +990,12 @@ class DatabricksProvider(ServiceProvider):
                     f"Successfully created database instance: {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)
 
@@ -918,6 +1007,10 @@ class DatabricksProvider(ServiceProvider):
                     logger.info(
                         f"Database instance {database.instance_name} was created concurrently by another process"
                     )
+                    # 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
@@ -1057,9 +1150,10 @@ 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 version (max version number from search_prompt_versions)
+        3. default alias
+        4. Register default_template if provided
 
         Args:
             prompt_model: The prompt model configuration
@@ -1070,9 +1164,11 @@ class DatabricksProvider(ServiceProvider):
         Raises:
             ValueError: If no prompt can be loaded from any source
         """
+
         prompt_name: str = prompt_model.full_name
+        mlflow_client: MlflowClient = MlflowClient()
 
-        # 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()
@@ -1086,73 +1182,48 @@ class DatabricksProvider(ServiceProvider):
                     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}"
                 )
-                # 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)
-
-                    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}")
+                # Fall through to try other methods
 
-            # Now try aliases in order: champion → latest → default
-            logger.debug(
-                f"Trying fallback order for '{prompt_name}': champion → latest → default"
-            )
+        # Try to load in priority 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}")
+        # 1. Try champion alias
+        try:
+            prompt_version = load_prompt(f"prompts:/{prompt_name}@champion")
+            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}")
+        # 2. Try to get latest version by finding the max version number
+        try:
+            versions = mlflow_client.search_prompt_versions(
+                prompt_name, max_results=100
+            )
+            if versions:
+                latest = max(versions, key=lambda v: int(v.version))
+                logger.info(
+                    f"Loaded prompt '{prompt_name}' version {latest.version} (latest by max version)"
+                )
+                return latest
+        except Exception as e:
+            logger.debug(f"Failed to find latest version for '{prompt_name}': {e}")
 
-            # Try default alias last
-            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}")
+        # 3. Try default alias
+        try:
+            prompt_version = load_prompt(f"prompts:/{prompt_name}@default")
+            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}")
 
-        # Fall back to registering default_template if provided
+        # 4. Try to register default_template if provided
         if prompt_model.default_template:
             logger.info(
-                f"Registering default_template for '{prompt_name}' "
-                "(no aliases found in registry)"
+                f"No existing prompt found for '{prompt_name}', "
+                "attempting to register default_template"
             )
             return self._sync_default_template_to_registry(
                 prompt_name, prompt_model.default_template, prompt_model.description
@@ -1160,72 +1231,58 @@ class DatabricksProvider(ServiceProvider):
 
         raise ValueError(
             f"Prompt '{prompt_name}' not found in registry "
-            "(tried champion, latest, default aliases) and no default_template provided"
+            "(tried champion alias, latest version, default alias) "
+            "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
+        """Get the best available prompt version, or register default_template if possible.
+
+        Tries to load prompts in order: champion → latest (max version) → default.
+        If none found and we have write permissions, registers the default_template.
+        If registration fails (e.g., in Model Serving), logs the error and raises.
+        """
+        mlflow_client: MlflowClient = MlflowClient()
 
+        # Try to find an existing prompt version in priority order
+        # 1. Try champion alias
         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"
-                )
-                if (
-                    existing.to_single_brace_format().strip()
-                    == default_template.strip()
-                ):
-                    logger.debug(f"Prompt '{prompt_name}' is already up-to-date")
+            champion = mlflow.genai.load_prompt(f"prompts:/{prompt_name}@champion")
+            logger.info(f"Loaded prompt '{prompt_name}' from champion alias")
+            return champion
+        except Exception as e:
+            logger.debug(f"Champion alias not found for '{prompt_name}': {e}")
 
-                    # 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:
-                        logger.info(
-                            f"Setting 'latest' alias for existing prompt '{prompt_name}' v{existing.version}"
-                        )
-                        mlflow.genai.set_prompt_alias(
-                            name=prompt_name,
-                            alias="latest",
-                            version=existing.version,
-                        )
+        # 2. Try to get the latest version by finding the max version number
+        try:
+            versions = mlflow_client.search_prompt_versions(
+                prompt_name, max_results=100
+            )
+            if versions:
+                latest = max(versions, key=lambda v: int(v.version))
+                logger.info(
+                    f"Loaded prompt '{prompt_name}' version {latest.version} (latest by max version)"
+                )
+                return latest
+        except Exception as e:
+            logger.debug(f"Failed to search versions for '{prompt_name}': {e}")
 
-                    # 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:
-                        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,
-                        )
+        # 3. Try default alias
+        try:
+            default = mlflow.genai.load_prompt(f"prompts:/{prompt_name}@default")
+            logger.info(f"Loaded prompt '{prompt_name}' from default alias")
+            return default
+        except Exception as e:
+            logger.debug(f"Default alias not found for '{prompt_name}': {e}")
 
-                    return existing  # Already up-to-date, return existing version
-            except Exception:
-                logger.debug(
-                    f"Default alias for prompt '{prompt_name}' doesn't exist yet"
-                )
+        # No existing prompt found - try to register if we have a template
+        logger.info(
+            f"No existing prompt found for '{prompt_name}', attempting to register default_template"
+        )
 
-            # Register new version and set as default alias
+        try:
             commit_message = description or "Auto-synced from default_template"
             prompt_version = mlflow.genai.register_prompt(
                 name=prompt_name,
@@ -1234,35 +1291,36 @@ class DatabricksProvider(ServiceProvider):
                 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,
-            )
+            # Try to set aliases (may fail in restricted environments)
+            try:
+                mlflow.genai.set_prompt_alias(
+                    name=prompt_name, alias="default", version=prompt_version.version
+                )
+                mlflow.genai.set_prompt_alias(
+                    name=prompt_name, alias="champion", version=prompt_version.version
+                )
+                logger.info(
+                    f"Registered prompt '{prompt_name}' v{prompt_version.version} with aliases"
+                )
+            except Exception as alias_error:
+                logger.warning(
+                    f"Registered prompt '{prompt_name}' v{prompt_version.version} "
+                    f"but failed to set aliases: {alias_error}"
+                )
 
-            logger.info(
-                f"Synced prompt '{prompt_name}' v{prompt_version.version} to registry with 'default', 'latest', and 'champion' aliases"
-            )
             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
+        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."
+            )
+            return PromptVersion(
+                name=prompt_name,
+                version=1,
+                template=default_template,
+                tags={"dao_ai": dao_ai_version()},
+            )
 
     def optimize_prompt(self, optimization: PromptOptimizationModel) -> PromptModel:
         """

dao_ai/tools/core.py

@@ -35,7 +35,7 @@ def create_tools(tool_models: Sequence[ToolModel]) -> Sequence[RunnableLike]:
         if name in tools:
             logger.warning(f"Tools already registered for: {name}, skipping creation.")
             continue
-        registered_tools: Sequence[RunnableLike] = tool_registry.get(name)
+        registered_tools: Sequence[RunnableLike] | None = tool_registry.get(name)
         if registered_tools is None:
             logger.debug(f"Creating tools for: {name}...")
             function: AnyTool = tool_config.function