dao-ai 0.1.10__tar.gz → 0.1.12__tar.gz

{dao_ai-0.1.10 → dao_ai-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.10
+Version: 0.1.12
 Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
 Project-URL: Homepage, https://github.com/natefleming/dao-ai
 Project-URL: Documentation, https://natefleming.github.io/dao-ai
@@ -28,7 +28,7 @@ Requires-Python: >=3.11
 Requires-Dist: databricks-agents>=1.9.0
 Requires-Dist: databricks-langchain[memory]>=0.12.1
 Requires-Dist: databricks-mcp>=0.5.0
-Requires-Dist: databricks-sdk[openai]>=0.76.0
+Requires-Dist: databricks-sdk[openai]>=0.77.0
 Requires-Dist: ddgs>=9.10.0
 Requires-Dist: dspy>=2.6.27
 Requires-Dist: flashrank>=0.2.10

{dao_ai-0.1.10 → dao_ai-0.1.12}/config/examples/07_human_in_the_loop/human_in_the_loop.yaml

@@ -26,6 +26,18 @@
   # - Composite: { type: composite, options: [${VAR1}, ${VAR2}, default] }
   # - Primitives: Just use the value directly
 
+variables:
+  client_id: &client_id
+    options:
+      - env: RETAIL_AI_DATABRICKS_CLIENT_ID      # Service principal client ID
+      - scope: retail_consumer_goods
+        secret: RETAIL_AI_DATABRICKS_CLIENT_ID
+  client_secret: &client_secret
+    options:
+      - env: RETAIL_AI_DATABRICKS_CLIENT_SECRET  # Service principal secret
+      - scope: retail_consumer_goods
+        secret: RETAIL_AI_DATABRICKS_CLIENT_SECRET
+
 schemas:
   retail_schema: &retail_schema
     catalog_name: nfleming                    
@@ -42,6 +54,8 @@ resources:
     lakebase_db: &lakebase_db
       name: retail-consumer-goods
       instance_name: retail-consumer-goods
+      client_id: *client_id
+      client_secret: *client_secret
 
 # =============================================================================
 # TOOLS WITH HUMAN-IN-THE-LOOP CONFIGURATIONS

{dao_ai-0.1.10 → dao_ai-0.1.12}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dao-ai"
-version = "0.1.10"
+version = "0.1.12"
 description = "DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML."
 readme = "README.md"
 license = { text = "MIT" }
@@ -46,7 +46,7 @@ dependencies = [
     "databricks-agents>=1.9.0",
     "databricks-langchain[memory]>=0.12.1",
     "databricks-mcp>=0.5.0",
-    "databricks-sdk[openai]>=0.76.0",
+    "databricks-sdk[openai]>=0.77.0",
     "ddgs>=9.10.0",
     "dspy>=2.6.27",
     "flashrank>=0.2.10",

{dao_ai-0.1.10 → dao_ai-0.1.12}/requirements.txt

@@ -33,7 +33,7 @@ databricks-agents==1.9.0
 databricks-ai-bridge==0.11.0
 databricks-langchain==0.12.1
 databricks-mcp==0.5.1
-databricks-sdk==0.76.0
+databricks-sdk==0.77.0
 dataclasses-json==0.6.7
 datasets==4.4.2
 ddgs==9.10.0

{dao_ai-0.1.10 → dao_ai-0.1.12}/src/dao_ai/apps/handlers.py

@@ -34,7 +34,7 @@ mlflow.set_tracking_uri("databricks")
 mlflow.langchain.autolog()
 
 # Get config path from environment or use default
-config_path: str = os.environ.get("DAO_AI_CONFIG_PATH", "model_config.yaml")
+config_path: str = os.environ.get("DAO_AI_CONFIG_PATH", "dao_ai.yaml")
 
 # Load configuration using AppConfig.from_file (consistent with CLI, notebook, builder)
 config: AppConfig = AppConfig.from_file(config_path)

{dao_ai-0.1.10 → dao_ai-0.1.12}/src/dao_ai/apps/resources.py

@@ -38,6 +38,8 @@ from databricks.sdk.service.apps import (
     AppResource,
     AppResourceDatabase,
     AppResourceDatabaseDatabasePermission,
+    AppResourceExperiment,
+    AppResourceExperimentExperimentPermission,
     AppResourceGenieSpace,
     AppResourceGenieSpaceGenieSpacePermission,
     AppResourceSecret,
@@ -308,6 +310,25 @@ def _extract_secrets_from_config(config: AppConfig) -> list[dict[str, Any]]:
         A list of secret resource dictionaries with unique scope/key pairs
     """
     secrets: dict[tuple[str, str], dict[str, Any]] = {}
+    used_names: set[str] = set()
+
+    def get_unique_resource_name(base_name: str) -> str:
+        """Generate a unique resource name, adding suffix if needed."""
+        sanitized = _sanitize_resource_name(base_name)
+        if sanitized not in used_names:
+            used_names.add(sanitized)
+            return sanitized
+        # Name collision - add numeric suffix
+        counter = 1
+        while True:
+            # Leave room for suffix (e.g., "_1", "_2", etc.)
+            suffix = f"_{counter}"
+            max_base_len = 30 - len(suffix)
+            candidate = sanitized[:max_base_len] + suffix
+            if candidate not in used_names:
+                used_names.add(candidate)
+                return candidate
+            counter += 1
 
     def extract_from_value(value: Any, path: str = "") -> None:
         """Recursively extract secrets from any value."""
@@ -315,9 +336,10 @@ def _extract_secrets_from_config(config: AppConfig) -> list[dict[str, Any]]:
             secret_key = (value.scope, value.secret)
             if secret_key not in secrets:
                 # Create a unique name for the secret resource
-                resource_name = f"{value.scope}_{value.secret}".replace(
-                    "-", "_"
-                ).replace("/", "_")
+                base_name = f"{value.scope}_{value.secret}".replace("-", "_").replace(
+                    "/", "_"
+                )
+                resource_name = get_unique_resource_name(base_name)
                 secrets[secret_key] = {
                     "name": resource_name,
                     "type": "secret",
@@ -325,7 +347,9 @@ def _extract_secrets_from_config(config: AppConfig) -> list[dict[str, Any]]:
                     "key": value.secret,
                     "permissions": [{"level": "READ"}],
                 }
-                logger.debug(f"Found secret: {value.scope}/{value.secret} at {path}")
+                logger.debug(
+                    f"Found secret: {value.scope}/{value.secret} at {path} -> resource: {resource_name}"
+                )
         elif isinstance(value, dict):
             for k, v in value.items():
                 extract_from_value(v, f"{path}.{k}" if path else k)
@@ -518,7 +542,10 @@ def _sanitize_resource_name(name: str) -> str:
     return sanitized
 
 
-def generate_sdk_resources(config: AppConfig) -> list[AppResource]:
+def generate_sdk_resources(
+    config: AppConfig,
+    experiment_id: str | None = None,
+) -> list[AppResource]:
     """
     Generate Databricks SDK AppResource objects from an AppConfig.
 
@@ -528,6 +555,9 @@ def generate_sdk_resources(config: AppConfig) -> list[AppResource]:
 
     Args:
         config: The AppConfig containing resource definitions
+        experiment_id: Optional MLflow experiment ID to add as a resource.
+            When provided, the experiment is added with CAN_EDIT permission,
+            allowing the app to log traces and runs.
 
     Returns:
         A list of AppResource objects for the Databricks SDK
@@ -536,13 +566,17 @@ def generate_sdk_resources(config: AppConfig) -> list[AppResource]:
         >>> from databricks.sdk import WorkspaceClient
         >>> from databricks.sdk.service.apps import App
         >>> config = AppConfig.from_file("model_config.yaml")
-        >>> resources = generate_sdk_resources(config)
+        >>> resources = generate_sdk_resources(config, experiment_id="12345")
         >>> w = WorkspaceClient()
         >>> app = App(name="my-app", resources=resources)
         >>> w.apps.create_and_wait(app=app)
     """
     resources: list[AppResource] = []
 
+    # Add experiment resource if provided
+    if experiment_id:
+        resources.append(_extract_sdk_experiment_resource(experiment_id))
+
     if config.resources is None:
         logger.debug("No resources defined in config")
         return resources
@@ -685,6 +719,36 @@ def _extract_sdk_volume_resources(
     return resources
 
 
+def _extract_sdk_experiment_resource(
+    experiment_id: str,
+    resource_name: str = "experiment",
+) -> AppResource:
+    """Create SDK AppResource for MLflow experiment.
+
+    This allows the Databricks App to log traces and runs to the specified
+    MLflow experiment. The experiment ID is exposed via the MLFLOW_EXPERIMENT_ID
+    environment variable using valueFrom: experiment in app.yaml.
+
+    Args:
+        experiment_id: The MLflow experiment ID
+        resource_name: The resource key name (default: "experiment")
+
+    Returns:
+        An AppResource for the MLflow experiment
+    """
+    resource = AppResource(
+        name=resource_name,
+        experiment=AppResourceExperiment(
+            experiment_id=experiment_id,
+            permission=AppResourceExperimentExperimentPermission.CAN_EDIT,
+        ),
+    )
+    logger.debug(
+        f"Extracted SDK experiment resource: {resource_name} -> {experiment_id}"
+    )
+    return resource
+
+
 def _extract_sdk_secrets_from_config(config: AppConfig) -> list[AppResource]:
     """
     Extract SDK AppResource objects for all secrets referenced in the config.
@@ -700,6 +764,25 @@ def _extract_sdk_secrets_from_config(config: AppConfig) -> list[AppResource]:
         A list of AppResource objects for secrets
     """
     secrets: dict[tuple[str, str], AppResource] = {}
+    used_names: set[str] = set()
+
+    def get_unique_resource_name(base_name: str) -> str:
+        """Generate a unique resource name, adding suffix if needed."""
+        sanitized = _sanitize_resource_name(base_name)
+        if sanitized not in used_names:
+            used_names.add(sanitized)
+            return sanitized
+        # Name collision - add numeric suffix
+        counter = 1
+        while True:
+            # Leave room for suffix (e.g., "_1", "_2", etc.)
+            suffix = f"_{counter}"
+            max_base_len = 30 - len(suffix)
+            candidate = sanitized[:max_base_len] + suffix
+            if candidate not in used_names:
+                used_names.add(candidate)
+                return candidate
+            counter += 1
 
     def extract_from_value(value: Any) -> None:
         """Recursively extract secrets from any value."""
@@ -707,10 +790,10 @@ def _extract_sdk_secrets_from_config(config: AppConfig) -> list[AppResource]:
             secret_key = (value.scope, value.secret)
             if secret_key not in secrets:
                 # Create a unique name for the secret resource
-                resource_name = f"{value.scope}_{value.secret}".replace(
-                    "-", "_"
-                ).replace("/", "_")
-                resource_name = _sanitize_resource_name(resource_name)
+                base_name = f"{value.scope}_{value.secret}".replace("-", "_").replace(
+                    "/", "_"
+                )
+                resource_name = get_unique_resource_name(base_name)
 
                 resource = AppResource(
                     name=resource_name,
@@ -722,7 +805,7 @@ def _extract_sdk_secrets_from_config(config: AppConfig) -> list[AppResource]:
                 )
                 secrets[secret_key] = resource
                 logger.debug(
-                    f"Found secret for SDK resource: {value.scope}/{value.secret}"
+                    f"Found secret for SDK resource: {value.scope}/{value.secret} -> resource: {resource_name}"
                 )
         elif isinstance(value, dict):
             for v in value.values():
@@ -926,12 +1009,22 @@ def generate_app_yaml(
     env_vars: list[dict[str, str]] = [
         {"name": "MLFLOW_TRACKING_URI", "value": "databricks"},
         {"name": "MLFLOW_REGISTRY_URI", "value": "databricks-uc"},
-        {"name": "DAO_AI_CONFIG_PATH", "value": "model_config.yaml"},
+        {"name": "MLFLOW_EXPERIMENT_ID", "valueFrom": "experiment"},
+        {"name": "DAO_AI_CONFIG_PATH", "value": "dao_ai.yaml"},
     ]
 
     # Extract environment variables from config.app.environment_vars
     config_env_vars = _extract_env_vars_from_config(config)
 
+    # Environment variables that are automatically provided by Databricks Apps
+    # and should not be included in app.yaml
+    platform_provided_env_vars = {"DATABRICKS_HOST"}
+
+    # Filter out platform-provided env vars from config
+    config_env_vars = [
+        e for e in config_env_vars if e["name"] not in platform_provided_env_vars
+    ]
+
     # Merge config env vars, avoiding duplicates (config takes precedence)
     base_env_names = {e["name"] for e in env_vars}
     for config_env in config_env_vars:

{dao_ai-0.1.10 → dao_ai-0.1.12}/src/dao_ai/apps/server.py

@@ -7,13 +7,13 @@ uses the AgentServer for the Databricks Apps runtime.
 
 Configuration Loading:
     The config path is specified via the DAO_AI_CONFIG_PATH environment variable,
-    or defaults to model_config.yaml in the current directory.
+    or defaults to dao_ai.yaml in the current directory.
 
 Usage:
     # With environment variable
     DAO_AI_CONFIG_PATH=/path/to/config.yaml python -m dao_ai.apps.server
 
-    # With default model_config.yaml in current directory
+    # With default dao_ai.yaml in current directory
     python -m dao_ai.apps.server
 """
 

{dao_ai-0.1.10 → dao_ai-0.1.12}/src/dao_ai/providers/databricks.py

@@ -152,25 +152,77 @@ class DatabricksProvider(ServiceProvider):
         client_secret: str | None = None,
         workspace_host: str | None = None,
     ) -> None:
-        if w is None:
-            w = _workspace_client(
-                pat=pat,
-                client_id=client_id,
-                client_secret=client_secret,
-                workspace_host=workspace_host,
+        # Store credentials for lazy initialization
+        self._pat = pat
+        self._client_id = client_id
+        self._client_secret = client_secret
+        self._workspace_host = workspace_host
+
+        # Lazy initialization for WorkspaceClient
+        self._w: WorkspaceClient | None = w
+        self._w_initialized = w is not None
+
+        # Lazy initialization for VectorSearchClient - only create when needed
+        # This avoids authentication errors in Databricks Apps where VSC
+        # requires explicit credentials but the platform uses ambient auth
+        self._vsc: VectorSearchClient | None = vsc
+        self._vsc_initialized = vsc is not None
+
+        # Lazy initialization for DatabricksFunctionClient
+        self._dfs: DatabricksFunctionClient | None = dfs
+        self._dfs_initialized = dfs is not None
+
+    @property
+    def w(self) -> WorkspaceClient:
+        """Lazy initialization of WorkspaceClient."""
+        if not self._w_initialized:
+            self._w = _workspace_client(
+                pat=self._pat,
+                client_id=self._client_id,
+                client_secret=self._client_secret,
+                workspace_host=self._workspace_host,
             )
-        if vsc is None:
-            vsc = _vector_search_client(
-                pat=pat,
-                client_id=client_id,
-                client_secret=client_secret,
-                workspace_host=workspace_host,
+            self._w_initialized = True
+        return self._w  # type: ignore[return-value]
+
+    @w.setter
+    def w(self, value: WorkspaceClient) -> None:
+        """Set WorkspaceClient and mark as initialized."""
+        self._w = value
+        self._w_initialized = True
+
+    @property
+    def vsc(self) -> VectorSearchClient:
+        """Lazy initialization of VectorSearchClient."""
+        if not self._vsc_initialized:
+            self._vsc = _vector_search_client(
+                pat=self._pat,
+                client_id=self._client_id,
+                client_secret=self._client_secret,
+                workspace_host=self._workspace_host,
             )
-        if dfs is None:
-            dfs = _function_client(w=w)
-        self.w = w
-        self.vsc = vsc
-        self.dfs = dfs
+            self._vsc_initialized = True
+        return self._vsc  # type: ignore[return-value]
+
+    @vsc.setter
+    def vsc(self, value: VectorSearchClient) -> None:
+        """Set VectorSearchClient and mark as initialized."""
+        self._vsc = value
+        self._vsc_initialized = True
+
+    @property
+    def dfs(self) -> DatabricksFunctionClient:
+        """Lazy initialization of DatabricksFunctionClient."""
+        if not self._dfs_initialized:
+            self._dfs = _function_client(w=self.w)
+            self._dfs_initialized = True
+        return self._dfs  # type: ignore[return-value]
+
+    @dfs.setter
+    def dfs(self, value: DatabricksFunctionClient) -> None:
+        """Set DatabricksFunctionClient and mark as initialized."""
+        self._dfs = value
+        self._dfs_initialized = True
 
     def experiment_name(self, config: AppConfig) -> str:
         current_user: User = self.w.current_user.me()
@@ -558,11 +610,21 @@ class DatabricksProvider(ServiceProvider):
 
         logger.info("Using workspace source path", source_path=source_path)
 
+        # Get or create experiment for this app (for tracing and tracking)
+        from mlflow.entities import Experiment
+
+        experiment: Experiment = self.get_or_create_experiment(config)
+        logger.info(
+            "Using MLflow experiment for app",
+            experiment_name=experiment.name,
+            experiment_id=experiment.experiment_id,
+        )
+
         # Upload the configuration file to the workspace
         source_config_path: str | None = config.source_config_path
         if source_config_path:
             # Read the config file and upload to workspace
-            config_file_name: str = "model_config.yaml"
+            config_file_name: str = "dao_ai.yaml"
             workspace_config_path: str = f"{source_path}/{config_file_name}"
 
             logger.info(
@@ -593,7 +655,7 @@ class DatabricksProvider(ServiceProvider):
             logger.warning(
                 "No source config path available. "
                 "Ensure DAO_AI_CONFIG_PATH is set in the app environment or "
-                "model_config.yaml exists in the app source directory."
+                "dao_ai.yaml exists in the app source directory."
             )
 
         # Generate and upload app.yaml with dynamically discovered resources
@@ -618,13 +680,15 @@ class DatabricksProvider(ServiceProvider):
         )
         logger.info("app.yaml with resources uploaded", path=app_yaml_path)
 
-        # Generate SDK resources from the config
+        # Generate SDK resources from the config (including experiment)
         from dao_ai.apps.resources import (
             generate_sdk_resources,
             generate_user_api_scopes,
         )
 
-        sdk_resources = generate_sdk_resources(config)
+        sdk_resources = generate_sdk_resources(
+            config, experiment_id=experiment.experiment_id
+        )
         if sdk_resources:
             logger.info(
                 "Discovered app resources from config",

{dao_ai-0.1.10 → dao_ai-0.1.12}/tests/dao_ai/test_databricks.py

@@ -684,22 +684,28 @@ def test_deploy_apps_agent_creates_new_app():
         # Mock current user
         provider.w.current_user.me.return_value = mock_user
 
-        # Simulate app doesn't exist
-        provider.w.apps.get.side_effect = NotFound("App not found")
-        provider.w.apps.create_and_wait.return_value = mock_created_app
-        provider.w.apps.deploy_and_wait.return_value = mock_deployment
+        # Mock MLflow experiment
+        mock_experiment = MagicMock()
+        mock_experiment.experiment_id = "12345"
+        with patch.object(
+            provider, "get_or_create_experiment", return_value=mock_experiment
+        ):
+            # Simulate app doesn't exist
+            provider.w.apps.get.side_effect = NotFound("App not found")
+            provider.w.apps.create_and_wait.return_value = mock_created_app
+            provider.w.apps.deploy_and_wait.return_value = mock_deployment
 
-        provider.deploy_apps_agent(mock_config)
+            provider.deploy_apps_agent(mock_config)
 
-        # Verify create_and_wait was called with an App object
-        provider.w.apps.create_and_wait.assert_called_once()
-        call_args = provider.w.apps.create_and_wait.call_args
-        app_arg = call_args.kwargs.get("app")
-        assert app_arg is not None
-        assert app_arg.name == "test-app"  # Normalized: underscores become dashes
-        assert app_arg.description == "Test app description"
-        # Verify deploy_and_wait was called
-        provider.w.apps.deploy_and_wait.assert_called_once()
+            # Verify create_and_wait was called with an App object
+            provider.w.apps.create_and_wait.assert_called_once()
+            call_args = provider.w.apps.create_and_wait.call_args
+            app_arg = call_args.kwargs.get("app")
+            assert app_arg is not None
+            assert app_arg.name == "test-app"  # Normalized: underscores become dashes
+            assert app_arg.description == "Test app description"
+            # Verify deploy_and_wait was called
+            provider.w.apps.deploy_and_wait.assert_called_once()
 
 
 @pytest.mark.unit
@@ -746,16 +752,22 @@ def test_deploy_apps_agent_updates_existing_app():
         # Mock current user
         provider.w.current_user.me.return_value = mock_user
 
-        # Simulate app already exists
-        provider.w.apps.get.return_value = mock_existing_app
-        provider.w.apps.deploy_and_wait.return_value = mock_deployment
+        # Mock MLflow experiment
+        mock_experiment = MagicMock()
+        mock_experiment.experiment_id = "12345"
+        with patch.object(
+            provider, "get_or_create_experiment", return_value=mock_experiment
+        ):
+            # Simulate app already exists
+            provider.w.apps.get.return_value = mock_existing_app
+            provider.w.apps.deploy_and_wait.return_value = mock_deployment
 
-        provider.deploy_apps_agent(mock_config)
+            provider.deploy_apps_agent(mock_config)
 
-        # Verify create_and_wait was NOT called (app already exists)
-        provider.w.apps.create_and_wait.assert_not_called()
-        # Verify deploy_and_wait was called
-        provider.w.apps.deploy_and_wait.assert_called_once()
+            # Verify create_and_wait was NOT called (app already exists)
+            provider.w.apps.create_and_wait.assert_not_called()
+            # Verify deploy_and_wait was called
+            provider.w.apps.deploy_and_wait.assert_called_once()
 
 
 @pytest.mark.unit