dao-ai 0.1.10__py3-none-any.whl → 0.1.11__py3-none-any.whl

dao_ai/apps/resources.py

@@ -38,6 +38,8 @@ from databricks.sdk.service.apps import (
     AppResource,
     AppResourceDatabase,
     AppResourceDatabaseDatabasePermission,
+    AppResourceExperiment,
+    AppResourceExperimentExperimentPermission,
     AppResourceGenieSpace,
     AppResourceGenieSpaceGenieSpacePermission,
     AppResourceSecret,
@@ -518,7 +520,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 +533,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 +544,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 +697,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.
@@ -926,6 +968,7 @@ def generate_app_yaml(
     env_vars: list[dict[str, str]] = [
         {"name": "MLFLOW_TRACKING_URI", "value": "databricks"},
         {"name": "MLFLOW_REGISTRY_URI", "value": "databricks-uc"},
+        {"name": "MLFLOW_EXPERIMENT_ID", "valueFrom": "experiment"},
         {"name": "DAO_AI_CONFIG_PATH", "value": "model_config.yaml"},
     ]
 

dao_ai/providers/databricks.py

@@ -558,6 +558,16 @@ 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:
@@ -618,13 +628,13 @@ 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.dist-info → dao_ai-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.10
+Version: 0.1.11
 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.dist-info → dao_ai-0.1.11.dist-info}/RECORD

@@ -16,7 +16,7 @@ dao_ai/vector_search.py,sha256=8d3xROg9zSIYNXjRRl6rSexsJTlufjRl5Fy1ZA8daKA,4019
 dao_ai/apps/__init__.py,sha256=RLuhZf4gQ4pemwKDz1183aXib8UfaRhwfKvRx68GRlM,661
 dao_ai/apps/handlers.py,sha256=NN81uDV2hy83zT-kY36mxyBgCQIHBJApX4bnUceGB8k,2614
 dao_ai/apps/model_serving.py,sha256=XLt3_0pGSRceMK6YtOrND9Jnh7mKLPCtwjVDLIaptQU,847
-dao_ai/apps/resources.py,sha256=7dZiRLk3xsnCkatFqUiGJAK0Rg-zKjNwEm55fnOAXBc,36522
+dao_ai/apps/resources.py,sha256=kAXqYx-Xwba1SaltPZIAtB2xoIHxH98Q41b8YGPLfI0,38029
 dao_ai/apps/server.py,sha256=D3R3J1svtpxnpDjoM-oxg66dMDI8USgiQHPRvRHc7oQ,1276
 dao_ai/genie/__init__.py,sha256=vdEyGhrt6L8GlK75SyYvTnl8QpHKDCJC5hJKLg4DesQ,1063
 dao_ai/genie/core.py,sha256=HPKbocvhnnw_PkQwfoq5bpgQmL9lZyyS6_goTJL8yiY,1073
@@ -53,7 +53,7 @@ dao_ai/orchestration/supervisor.py,sha256=alKMEEo9G5LhdpMvTVdAMel234cZj5_MguWl4w
 dao_ai/orchestration/swarm.py,sha256=8tp1eGmsQqqWpaDcjPoJckddPWohZdmmN0RGRJ_xzOA,9198
 dao_ai/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dao_ai/providers/base.py,sha256=cJGo3UjUTPgS91dv38ePOHwQQtYhIa84ebb167CBXjk,2111
-dao_ai/providers/databricks.py,sha256=s4zYVvODtJNKE7SGXZpA6Vw-mlnNIQeF-JKc7Mx9IHE,70254
+dao_ai/providers/databricks.py,sha256=EZokRfbScW5K2h8EOCQ95XjoXt981ySslQk0d0DJmeI,70687
 dao_ai/tools/__init__.py,sha256=NfRpAKds_taHbx6gzLPWgtPXve-YpwzkoOAUflwxceM,1734
 dao_ai/tools/agent.py,sha256=plIWALywRjaDSnot13nYehBsrHRpBUpsVZakoGeajOE,1858
 dao_ai/tools/core.py,sha256=bRIN3BZhRQX8-Kpu3HPomliodyskCqjxynQmYbk6Vjs,3783
@@ -68,8 +68,8 @@ dao_ai/tools/sql.py,sha256=tKd1gjpLuKdQDyfmyYYtMiNRHDW6MGRbdEVaeqyB8Ok,7632
 dao_ai/tools/time.py,sha256=tufJniwivq29y0LIffbgeBTIDE6VgrLpmVf8Qr90qjw,9224
 dao_ai/tools/unity_catalog.py,sha256=AjQfW7bvV8NurqDLIyntYRv2eJuTwNdbvex1L5CRjOk,15534
 dao_ai/tools/vector_search.py,sha256=oe2uBwl2TfeJIXPpwiS6Rmz7wcHczSxNyqS9P3hE6co,14542
-dao_ai-0.1.10.dist-info/METADATA,sha256=lAelxDWOeBYOCdZhW1EAnu3PvY2EfoxocA8Mmy_bWII,16698
-dao_ai-0.1.10.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.1.10.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.1.10.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.1.10.dist-info/RECORD,,
+dao_ai-0.1.11.dist-info/METADATA,sha256=uSlkpK84xgQlBOkLAzdOCycWxUC1ovzrrrU_uCVMOd0,16698
+dao_ai-0.1.11.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.1.11.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.1.11.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.1.11.dist-info/RECORD,,