dao-ai 0.1.8__tar.gz → 0.1.9__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (301) hide show
  1. {dao_ai-0.1.8 → dao_ai-0.1.9}/PKG-INFO +2 -2
  2. dao_ai-0.1.9/app.yaml +12 -0
  3. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/02_mcp/filtered_mcp.yaml +50 -23
  4. {dao_ai-0.1.8 → dao_ai-0.1.9}/pyproject.toml +2 -2
  5. {dao_ai-0.1.8 → dao_ai-0.1.9}/requirements.txt +10 -0
  6. {dao_ai-0.1.8 → dao_ai-0.1.9}/schemas/model_config_schema.json +12 -0
  7. dao_ai-0.1.9/src/dao_ai/app_server.py +103 -0
  8. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/cli.py +15 -1
  9. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/config.py +89 -36
  10. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/providers/base.py +28 -2
  11. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/providers/databricks.py +205 -3
  12. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/state.py +1 -0
  13. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_databricks.py +377 -20
  14. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_genie_room_model.py +0 -6
  15. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_mcp.py +25 -15
  16. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_mcp_filtering_integration.py +215 -111
  17. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_mcp_function_model.py +45 -21
  18. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_sql_tool.py +72 -31
  19. {dao_ai-0.1.8 → dao_ai-0.1.9}/.gitignore +0 -0
  20. {dao_ai-0.1.8 → dao_ai-0.1.9}/.python-version +0 -0
  21. {dao_ai-0.1.8 → dao_ai-0.1.9}/CHANGELOG.md +0 -0
  22. {dao_ai-0.1.8 → dao_ai-0.1.9}/CONTRIBUTING.md +0 -0
  23. {dao_ai-0.1.8 → dao_ai-0.1.9}/LICENSE +0 -0
  24. {dao_ai-0.1.8 → dao_ai-0.1.9}/Makefile +0 -0
  25. {dao_ai-0.1.8 → dao_ai-0.1.9}/README.md +0 -0
  26. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/01_getting_started/README.md +0 -0
  27. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/01_getting_started/minimal.yaml +0 -0
  28. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/02_mcp/README.md +0 -0
  29. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/02_mcp/custom_mcp.yaml +0 -0
  30. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/02_mcp/external_mcp.yaml +0 -0
  31. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/02_mcp/managed_mcp.yaml +0 -0
  32. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/02_mcp/slack_integration.yaml +0 -0
  33. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/03_reranking/README.md +0 -0
  34. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/03_reranking/vector_search_with_reranking.yaml +0 -0
  35. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/04_genie/README.md +0 -0
  36. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/04_genie/genie_basic.yaml +0 -0
  37. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/04_genie/genie_lru_cache.yaml +0 -0
  38. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/04_genie/genie_semantic_cache.yaml +0 -0
  39. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/04_genie/genie_with_conversation_id.yaml +0 -0
  40. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/05_memory/README.md +0 -0
  41. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/05_memory/conversation_summarization.yaml +0 -0
  42. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/05_memory/in_memory_basic.yaml +0 -0
  43. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/05_memory/lakebase_persistence.yaml +0 -0
  44. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/05_memory/postgres_persistence.yaml +0 -0
  45. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/06_on_behalf_of_user/README.md +0 -0
  46. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/06_on_behalf_of_user/obo_basic.yaml +0 -0
  47. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/07_human_in_the_loop/README.md +0 -0
  48. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/07_human_in_the_loop/human_in_the_loop.yaml +0 -0
  49. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/08_guardrails/README.md +0 -0
  50. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/08_guardrails/guardrails_basic.yaml +0 -0
  51. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/09_structured_output/README.md +0 -0
  52. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/09_structured_output/structured_output.yaml +0 -0
  53. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/10_agent_integrations/README.md +0 -0
  54. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/10_agent_integrations/agent_bricks.yaml +0 -0
  55. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/10_agent_integrations/kasal.yaml +0 -0
  56. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/11_prompt_engineering/README.md +0 -0
  57. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/11_prompt_engineering/prompt_optimization.yaml +0 -0
  58. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/11_prompt_engineering/prompt_registry.yaml +0 -0
  59. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/README.md +0 -0
  60. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/combined_middleware.yaml +0 -0
  61. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/context_management.yaml +0 -0
  62. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/custom_field_validation.yaml +0 -0
  63. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/limit_middleware.yaml +0 -0
  64. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/logging_middleware.yaml +0 -0
  65. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/pii_middleware.yaml +0 -0
  66. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/retry_middleware.yaml +0 -0
  67. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/12_middleware/tool_selector_middleware.yaml +0 -0
  68. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/13_orchestration/README.md +0 -0
  69. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/13_orchestration/supervisor_pattern.yaml +0 -0
  70. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/13_orchestration/swarm_pattern.yaml +0 -0
  71. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/14_basic_tools/README.md +0 -0
  72. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/14_basic_tools/sql_tool_example.yaml +0 -0
  73. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/README.md +0 -0
  74. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/brick_store.yaml +0 -0
  75. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/deep_research.yaml +0 -0
  76. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/executive_assistant.yaml +0 -0
  77. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/genie_and_genie_mcp.yaml +0 -0
  78. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/genie_vector_search_hybrid.yaml +0 -0
  79. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/hardware_store.yaml +0 -0
  80. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/hardware_store_lakebase.yaml +0 -0
  81. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/hardware_store_swarm.yaml +0 -0
  82. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/quick_serve_restaurant.yaml +0 -0
  83. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/15_complete_applications/reservations_system.yaml +0 -0
  84. {dao_ai-0.1.8 → dao_ai-0.1.9}/config/examples/README.md +0 -0
  85. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/appointments.sql +0 -0
  86. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/appointments_data.sql +0 -0
  87. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/brand_rep_demo_data.sql +0 -0
  88. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/brand_rep_demo_queries.sql +0 -0
  89. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/brand_rep_demo_tables.sql +0 -0
  90. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/brand_rep_demo_validation.sql +0 -0
  91. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/customers.sql +0 -0
  92. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/customers_data.sql +0 -0
  93. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/dim_stores.sql +0 -0
  94. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/dim_stores_data.sql +0 -0
  95. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/employee_performance.sql +0 -0
  96. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/employee_performance_data.sql +0 -0
  97. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/employee_tasks.sql +0 -0
  98. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/employee_tasks_data.sql +0 -0
  99. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/inventory.sql +0 -0
  100. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/inventory_data.sql +0 -0
  101. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/managers.sql +0 -0
  102. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/managers_data.sql +0 -0
  103. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/product_data.sql +0 -0
  104. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/products.sql +0 -0
  105. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/dais2025/task_assignments.sql +0 -0
  106. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/hardware_store/inventory.snappy.parquet +0 -0
  107. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/hardware_store/inventory.sql +0 -0
  108. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/hardware_store/products.snappy.parquet +0 -0
  109. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/hardware_store/products.sql +0 -0
  110. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/quick_serve_restaurant/.gitkeep +0 -0
  111. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/quick_serve_restaurant/fulfil_item_orders.sql +0 -0
  112. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/quick_serve_restaurant/items_description.csv +0 -0
  113. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/quick_serve_restaurant/items_description.sql +0 -0
  114. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/quick_serve_restaurant/items_raw.csv +0 -0
  115. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/quick_serve_restaurant/items_raw.sql +0 -0
  116. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/quick_serve_restaurant/orders_raw.csv +0 -0
  117. {dao_ai-0.1.8 → dao_ai-0.1.9}/data/quick_serve_restaurant/orders_raw.sql +0 -0
  118. {dao_ai-0.1.8 → dao_ai-0.1.9}/databricks.yaml.template +0 -0
  119. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/architecture.md +0 -0
  120. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/cli-reference.md +0 -0
  121. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/configuration-reference.md +0 -0
  122. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/contributing.md +0 -0
  123. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/examples.md +0 -0
  124. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/faq.md +0 -0
  125. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/hardware_store/README.md +0 -0
  126. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/hardware_store/retail_supervisor.png +0 -0
  127. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/hardware_store/retail_swarm.png +0 -0
  128. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/images/genie.png +0 -0
  129. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/key-capabilities.md +0 -0
  130. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/python-api.md +0 -0
  131. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/quick_serve_restaurant/.gitkeep +0 -0
  132. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/quick_serve_restaurant/quick-serve-restaurant.png +0 -0
  133. {dao_ai-0.1.8 → dao_ai-0.1.9}/docs/why-dao.md +0 -0
  134. {dao_ai-0.1.8 → dao_ai-0.1.9}/environment.yaml +0 -0
  135. {dao_ai-0.1.8 → dao_ai-0.1.9}/examples/dais2025/examples.yaml +0 -0
  136. {dao_ai-0.1.8 → dao_ai-0.1.9}/examples/deep_research/examples.yaml +0 -0
  137. {dao_ai-0.1.8 → dao_ai-0.1.9}/examples/executive_assistant/examples.yaml +0 -0
  138. {dao_ai-0.1.8 → dao_ai-0.1.9}/examples/hardware_store/examples.yaml +0 -0
  139. {dao_ai-0.1.8 → dao_ai-0.1.9}/examples/quick_serve_restaurant/.gitkeep +0 -0
  140. {dao_ai-0.1.8 → dao_ai-0.1.9}/examples/quick_serve_restaurant/examples.yaml +0 -0
  141. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/dais2025/extract_store_numbers.sql +0 -0
  142. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/dais2025/find_inventory_by_sku.sql +0 -0
  143. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/dais2025/find_inventory_by_upc.sql +0 -0
  144. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/dais2025/find_product_by_sku.sql +0 -0
  145. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/dais2025/find_product_by_upc.sql +0 -0
  146. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/dais2025/find_store_by_number.sql +0 -0
  147. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/dais2025/find_store_inventory_by_sku.sql +0 -0
  148. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/dais2025/find_store_inventory_by_upc.sql +0 -0
  149. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/hardware_store/find_inventory_by_sku.sql +0 -0
  150. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/hardware_store/find_inventory_by_upc.sql +0 -0
  151. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/hardware_store/find_product_by_sku.sql +0 -0
  152. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/hardware_store/find_product_by_upc.sql +0 -0
  153. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/hardware_store/find_store_inventory_by_sku.sql +0 -0
  154. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/hardware_store/find_store_inventory_by_upc.sql +0 -0
  155. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/quick_serve_restaurant/.gitkeep +0 -0
  156. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/quick_serve_restaurant/insert_coffee_order.sql +0 -0
  157. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/quick_serve_restaurant/lookup_items_by_descriptions.sql +0 -0
  158. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/quick_serve_restaurant/match_historical_item_order_by_date.sql +0 -0
  159. {dao_ai-0.1.8 → dao_ai-0.1.9}/functions/quick_serve_restaurant/match_item_by_description_and_price.sql +0 -0
  160. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/01_ingest_and_transform.py +0 -0
  161. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/02_provision_vector_search.py +0 -0
  162. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/03_provision_lakebase.py +0 -0
  163. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/04_unity_catalog_tools.py +0 -0
  164. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/05_deploy_agent.py +0 -0
  165. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/06_generate_evaluation_data.py +0 -0
  166. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/07_run_evaluation.py +0 -0
  167. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/08_run_examples.py +0 -0
  168. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/09_evaluate_inferences.py +0 -0
  169. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/10_optimize_prompts.py +0 -0
  170. {dao_ai-0.1.8 → dao_ai-0.1.9}/notebooks/99_scratchpad.py +0 -0
  171. {dao_ai-0.1.8 → dao_ai-0.1.9}/schemas/bundle_config_schema.json +0 -0
  172. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/__init__.py +0 -0
  173. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/models.py +0 -0
  174. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/tools/__init__.py +0 -0
  175. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/tools/customer.py +0 -0
  176. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/tools/employee.py +0 -0
  177. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/tools/executive.py +0 -0
  178. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/tools/genie.py +0 -0
  179. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/tools/inventory.py +0 -0
  180. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/tools/models.py +0 -0
  181. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dais2025/tools/store.py +0 -0
  182. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/__init__.py +0 -0
  183. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/agent_as_code.py +0 -0
  184. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/catalog.py +0 -0
  185. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/genie/__init__.py +0 -0
  186. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/genie/cache/__init__.py +0 -0
  187. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/genie/cache/base.py +0 -0
  188. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/genie/cache/core.py +0 -0
  189. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/genie/cache/lru.py +0 -0
  190. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/genie/cache/semantic.py +0 -0
  191. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/genie/core.py +0 -0
  192. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/graph.py +0 -0
  193. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/hooks/__init__.py +0 -0
  194. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/hooks/core.py +0 -0
  195. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/logging.py +0 -0
  196. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/memory/__init__.py +0 -0
  197. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/memory/base.py +0 -0
  198. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/memory/core.py +0 -0
  199. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/memory/databricks.py +0 -0
  200. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/memory/postgres.py +0 -0
  201. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/messages.py +0 -0
  202. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/__init__.py +0 -0
  203. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/assertions.py +0 -0
  204. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/base.py +0 -0
  205. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/context_editing.py +0 -0
  206. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/core.py +0 -0
  207. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/guardrails.py +0 -0
  208. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/human_in_the_loop.py +0 -0
  209. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/message_validation.py +0 -0
  210. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/model_call_limit.py +0 -0
  211. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/model_retry.py +0 -0
  212. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/pii.py +0 -0
  213. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/summarization.py +0 -0
  214. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/tool_call_limit.py +0 -0
  215. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/tool_retry.py +0 -0
  216. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/middleware/tool_selector.py +0 -0
  217. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/models.py +0 -0
  218. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/nodes.py +0 -0
  219. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/optimization.py +0 -0
  220. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/orchestration/__init__.py +0 -0
  221. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/orchestration/core.py +0 -0
  222. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/orchestration/supervisor.py +0 -0
  223. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/orchestration/swarm.py +0 -0
  224. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/prompts.py +0 -0
  225. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/providers/__init__.py +0 -0
  226. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/__init__.py +0 -0
  227. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/agent.py +0 -0
  228. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/core.py +0 -0
  229. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/email.py +0 -0
  230. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/genie.py +0 -0
  231. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/mcp.py +0 -0
  232. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/memory.py +0 -0
  233. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/python.py +0 -0
  234. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/search.py +0 -0
  235. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/slack.py +0 -0
  236. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/sql.py +0 -0
  237. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/time.py +0 -0
  238. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/unity_catalog.py +0 -0
  239. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/tools/vector_search.py +0 -0
  240. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/types.py +0 -0
  241. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/utils.py +0 -0
  242. {dao_ai-0.1.8 → dao_ai-0.1.9}/src/dao_ai/vector_search.py +0 -0
  243. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/config/test_model_config.yaml +0 -0
  244. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/conftest.py +0 -0
  245. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/middleware/test_context_editing.py +0 -0
  246. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/middleware/test_model_call_limit.py +0 -0
  247. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/middleware/test_model_retry.py +0 -0
  248. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/middleware/test_pii.py +0 -0
  249. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/middleware/test_tool_call_limit.py +0 -0
  250. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/middleware/test_tool_retry.py +0 -0
  251. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/middleware/test_tool_selector.py +0 -0
  252. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_agent_response_format.py +0 -0
  253. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_assertions_middleware.py +0 -0
  254. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_catalog.py +0 -0
  255. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_chat_history.py +0 -0
  256. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_config.py +0 -0
  257. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_function_parsing.py +0 -0
  258. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_genie.py +0 -0
  259. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_genie_conversation_ids_in_outputs.py +0 -0
  260. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_genie_databricks_integration.py +0 -0
  261. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_guardrail_retry.py +0 -0
  262. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_hitl_config_model.py +0 -0
  263. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_hitl_responses_agent.py +0 -0
  264. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_hooks.py +0 -0
  265. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_human_in_the_loop.py +0 -0
  266. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_inference.py +0 -0
  267. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_inference_integration.py +0 -0
  268. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_input_output_structure.py +0 -0
  269. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_interrupt_type.py +0 -0
  270. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_llm_interrupt_handling.py +0 -0
  271. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_mcp_filtering.py +0 -0
  272. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_message_validation_middleware.py +0 -0
  273. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_messages.py +0 -0
  274. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_models.py +0 -0
  275. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_optimization.py +0 -0
  276. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_postgres_integration.py +0 -0
  277. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_prompt_optimizations.py +0 -0
  278. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_prompts.py +0 -0
  279. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_reranking.py +0 -0
  280. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_reranking_integration.py +0 -0
  281. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_resources_model_genie_integration.py +0 -0
  282. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_response_format.py +0 -0
  283. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_responses_agent_structured_output_unit.py +0 -0
  284. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_semantic_cache_context.py +0 -0
  285. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_sql_tool_integration.py +0 -0
  286. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_state.py +0 -0
  287. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_summarization_inference.py +0 -0
  288. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_swarm_middleware.py +0 -0
  289. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_tools.py +0 -0
  290. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_types.py +0 -0
  291. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_unity_catalog.py +0 -0
  292. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_utils.py +0 -0
  293. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_utils_type_from_fqn.py +0 -0
  294. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_vector_search.py +0 -0
  295. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/test_warehouse_model.py +0 -0
  296. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/dao_ai/weather_server_mcp.py +0 -0
  297. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/hardware_store/.gitkeep +0 -0
  298. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/hardware_store/test_graph.py +0 -0
  299. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/images/doritos_upc.png +0 -0
  300. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/images/lays_upc.png +0 -0
  301. {dao_ai-0.1.8 → dao_ai-0.1.9}/tests/quick_serve_restaurant/.gitkeep +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dao-ai
3
- Version: 0.1.8
3
+ Version: 0.1.9
4
4
  Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
5
5
  Project-URL: Homepage, https://github.com/natefleming/dao-ai
6
6
  Project-URL: Documentation, https://natefleming.github.io/dao-ai
@@ -43,7 +43,7 @@ Requires-Dist: langgraph>=1.0.5
43
43
  Requires-Dist: langmem>=0.0.30
44
44
  Requires-Dist: loguru>=0.7.3
45
45
  Requires-Dist: mcp>=1.24.0
46
- Requires-Dist: mlflow>=3.8.1
46
+ Requires-Dist: mlflow[databricks]>=3.8.1
47
47
  Requires-Dist: nest-asyncio>=1.6.0
48
48
  Requires-Dist: openevals>=0.1.3
49
49
  Requires-Dist: openpyxl>=3.1.5
dao_ai-0.1.9/app.yaml ADDED
@@ -0,0 +1,12 @@
1
+ command: ["uv", "run", "python", "-m", "dao_ai.app_server"]
2
+ # Databricks Apps listen by default on port 8000
3
+
4
+ env:
5
+ - name: MLFLOW_TRACKING_URI
6
+ value: "databricks"
7
+ - name: MLFLOW_REGISTRY_URI
8
+ value: "databricks-uc"
9
+ # Path to the dao-ai configuration file
10
+ # This can be set to a workspace path or local file path
11
+ - name: DAO_AI_CONFIG_PATH
12
+ value: "model_config.yaml"
@@ -85,8 +85,8 @@ tools:
85
85
  - execute_query # Exact match
86
86
  - list_tables # Exact match
87
87
  - describe_table # Exact match
88
- - get_* # Pattern: all getters
89
- - show_* # Pattern: all show operations
88
+ - "get_*" # Pattern: all getters
89
+ - "show_*" # Pattern: all show operations
90
90
 
91
91
  # ---------------------------------------------------------------------------
92
92
  # Example 2: Exclude Dangerous Tools (Denylist)
@@ -104,11 +104,11 @@ tools:
104
104
  workspace_host: *workspace_host
105
105
  # Load all tools EXCEPT these dangerous ones
106
106
  exclude_tools:
107
- - drop_* # Block all drop operations
108
- - delete_* # Block all delete operations
109
- - truncate_* # Block all truncate operations
107
+ - "drop_*" # Block all drop operations
108
+ - "delete_*" # Block all delete operations
109
+ - "truncate_*" # Block all truncate operations
110
110
  - execute_ddl # Block DDL execution
111
- - alter_* # Block all alter operations
111
+ - "alter_*" # Block all alter operations
112
112
 
113
113
  # ---------------------------------------------------------------------------
114
114
  # Example 3: Hybrid Filtering (Include + Exclude)
@@ -126,14 +126,14 @@ tools:
126
126
  workspace_host: *workspace_host
127
127
  # Start with these categories
128
128
  include_tools:
129
- - query_* # All query functions
130
- - get_* # All getter functions
131
- - list_* # All list functions
129
+ - "query_*" # All query functions
130
+ - "get_*" # All getter functions
131
+ - "list_*" # All list functions
132
132
  # But exclude sensitive ones
133
133
  exclude_tools:
134
- - *_sensitive # Exclude anything with "_sensitive"
135
- - *_admin # Exclude admin functions
136
- - get_secret_* # Exclude secret getters
134
+ - "*_sensitive" # Exclude anything with "_sensitive"
135
+ - "*_admin" # Exclude admin functions
136
+ - "get_secret_*" # Exclude secret getters
137
137
 
138
138
  # ---------------------------------------------------------------------------
139
139
  # Example 4: Pattern-Based Inclusion
@@ -151,10 +151,10 @@ tools:
151
151
  workspace_host: *workspace_host
152
152
  # Only read operations with patterns
153
153
  include_tools:
154
- - query_* # All queries
155
- - list_* # All lists
156
- - describe_* # All describe operations
157
- - show_* # All show operations
154
+ - "query_*" # All queries
155
+ - "list_*" # All lists
156
+ - "describe_*" # All describe operations
157
+ - "show_*" # All show operations
158
158
 
159
159
  # ---------------------------------------------------------------------------
160
160
  # Example 5: Maximum Security (Very Restrictive)
@@ -192,8 +192,8 @@ tools:
192
192
  workspace_host: *workspace_host
193
193
  # Allow everything except the really dangerous stuff
194
194
  exclude_tools:
195
- - drop_* # Can't drop anything
196
- - truncate_* # Can't truncate
195
+ - "drop_*" # Can't drop anything
196
+ - "truncate_*" # Can't truncate
197
197
  - execute_ddl # Can't run arbitrary DDL
198
198
 
199
199
  # =============================================================================
@@ -295,17 +295,44 @@ agents:
295
295
  # APPLICATION CONFIGURATION
296
296
  # =============================================================================
297
297
 
298
- app_name: filtered_mcp_example
299
- entry_agent: safe_sql_agent
298
+ app:
299
+ name: filtered_mcp_example # Application name
300
+ log_level: INFO # Logging level (DEBUG, INFO, WARNING, ERROR)
301
+ registered_model: # MLflow registered model configuration
302
+ schema: *retail_schema # Schema where model will be registered
303
+ name: filtered_mcp_example # Model name in MLflow registry
304
+ agents: # List of agents included in the system
305
+ - *safe_sql_agent # Safe SQL agent (read-only)
306
+ - *analyst_agent # Data analyst agent
307
+ - *dev_agent # Development agent
308
+ - *secure_agent # High-security agent
309
+ - *functions_agent # Functions agent
310
+ orchestration: # Agent orchestration configuration
311
+ swarm: # Supervisor orchestration pattern
312
+ model: *default_llm # LLM for routing decisions
313
+ # Agents will be automatically routed based on their capabilities
300
314
 
301
315
  # =============================================================================
302
316
  # USAGE NOTES
303
317
  # =============================================================================
304
318
  #
319
+ # Deployment:
320
+ # databricks apps deploy dao-ai-filtered-mcp --source-dir . --config-file config/examples/02_mcp/filtered_mcp.yaml
321
+ #
322
+ # The app uses swarm orchestration to automatically route user queries to the
323
+ # most appropriate agent based on their capabilities and tool access:
324
+ #
325
+ # - safe_sql_agent: Read-only SQL queries (allowlist of 5 safe operations)
326
+ # - analyst_agent: Data analysis queries (query_*, list_*, describe_*, show_*)
327
+ # - dev_agent: Development queries (all tools except dangerous DDL operations)
328
+ # - secure_agent: Maximum security (only 3 specific tools)
329
+ # - functions_agent: UC Functions access (filtered to exclude sensitive/admin)
330
+ #
305
331
  # Testing Filters:
306
- # 1. Run with safe_sql_agent - try to call drop_table (should not be available)
307
- # 2. Run with dev_agent - verify you can query but not drop
308
- # 3. Run with secure_agent - verify only 3 tools are available
332
+ # 1. Ask to query data Routes to analyst_agent or safe_sql_agent
333
+ # 2. Ask to drop a table Routes to dev_agent, which will refuse (blocked)
334
+ # 3. Ask in secure mode Routes to secure_agent (only 3 tools available)
335
+ # 4. Ask to call a function → Routes to functions_agent (sensitive functions blocked)
309
336
  #
310
337
  # Pattern Examples:
311
338
  # - "query_*" matches: query_sales, query_inventory, query_anything
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "dao-ai"
7
- version = "0.1.8"
7
+ version = "0.1.9"
8
8
  description = "DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML."
9
9
  readme = "README.md"
10
10
  license = { text = "MIT" }
@@ -61,7 +61,7 @@ dependencies = [
61
61
  "langmem>=0.0.30",
62
62
  "loguru>=0.7.3",
63
63
  "mcp>=1.24.0",
64
- "mlflow>=3.8.1",
64
+ "mlflow[databricks]>=3.8.1",
65
65
  "nest-asyncio>=1.6.0",
66
66
  "openevals>=0.1.3",
67
67
  "openpyxl>=3.1.5",
@@ -9,6 +9,9 @@ anthropic==0.75.0
9
9
  anyio==4.12.0
10
10
  asyncer==0.0.8
11
11
  attrs==25.4.0
12
+ azure-core==1.37.0
13
+ azure-storage-blob==12.28.0
14
+ azure-storage-file-datalake==12.23.0
12
15
  backoff==2.2.1
13
16
  blinker==1.9.0
14
17
  boto3==1.42.16
@@ -57,7 +60,12 @@ fsspec==2025.10.0
57
60
  gepa==0.0.22
58
61
  gitdb==4.0.12
59
62
  gitpython==3.1.45
63
+ google-api-core==2.28.1
60
64
  google-auth==2.45.0
65
+ google-cloud-core==2.5.0
66
+ google-cloud-storage==3.7.0
67
+ google-crc32c==1.8.0
68
+ google-resumable-media==2.8.0
61
69
  googleapis-common-protos==1.72.0
62
70
  grandalf==0.8
63
71
  graphene==3.4.3
@@ -81,6 +89,7 @@ identify==2.6.15
81
89
  idna==3.11
82
90
  importlib-metadata==8.7.1
83
91
  iniconfig==2.3.0
92
+ isodate==0.7.2
84
93
  itsdangerous==2.2.0
85
94
  jinja2==3.1.6
86
95
  jiter==0.12.0
@@ -151,6 +160,7 @@ pluggy==1.6.0
151
160
  pre-commit==4.5.1
152
161
  primp==0.15.0
153
162
  propcache==0.4.1
163
+ proto-plus==1.27.0
154
164
  protobuf==6.33.2
155
165
  psycopg==3.3.2
156
166
  psycopg-binary==3.3.2
@@ -5150,6 +5150,18 @@
5150
5150
  },
5151
5151
  "additionalProperties": false,
5152
5152
  "properties": {
5153
+ "version": {
5154
+ "anyOf": [
5155
+ {
5156
+ "type": "string"
5157
+ },
5158
+ {
5159
+ "type": "null"
5160
+ }
5161
+ ],
5162
+ "default": null,
5163
+ "title": "Version"
5164
+ },
5153
5165
  "variables": {
5154
5166
  "additionalProperties": {
5155
5167
  "anyOf": [
@@ -0,0 +1,103 @@
1
+ """
2
+ App server module for running dao-ai agents as Databricks Apps.
3
+
4
+ This module provides the entry point for deploying dao-ai agents as Databricks Apps
5
+ using MLflow's AgentServer. It follows the same pattern as agent_as_code.py but
6
+ uses the AgentServer for the Databricks Apps runtime.
7
+
8
+ Configuration Loading:
9
+ The config path is specified via the DAO_AI_CONFIG_PATH environment variable,
10
+ or defaults to model_config.yaml in the current directory.
11
+
12
+ Usage:
13
+ # With environment variable
14
+ DAO_AI_CONFIG_PATH=/path/to/config.yaml python -m dao_ai.app_server
15
+
16
+ # With default model_config.yaml in current directory
17
+ python -m dao_ai.app_server
18
+ """
19
+
20
+ import os
21
+ from typing import AsyncGenerator
22
+
23
+ import mlflow
24
+ from dotenv import load_dotenv
25
+ from mlflow.genai.agent_server import AgentServer, invoke, stream
26
+ from mlflow.pyfunc import ResponsesAgent
27
+ from mlflow.types.responses import (
28
+ ResponsesAgentRequest,
29
+ ResponsesAgentResponse,
30
+ ResponsesAgentStreamEvent,
31
+ )
32
+
33
+ from dao_ai.config import AppConfig
34
+ from dao_ai.logging import configure_logging
35
+
36
+ # Load environment variables from .env.local if it exists
37
+ load_dotenv(dotenv_path=".env.local", override=True)
38
+
39
+ # Configure MLflow
40
+ mlflow.set_registry_uri("databricks-uc")
41
+ mlflow.set_tracking_uri("databricks")
42
+ mlflow.langchain.autolog()
43
+
44
+ # Get config path from environment or use default
45
+ config_path: str = os.environ.get("DAO_AI_CONFIG_PATH", "model_config.yaml")
46
+
47
+ # Load configuration using AppConfig.from_file (consistent with CLI, notebook, builder)
48
+ config: AppConfig = AppConfig.from_file(config_path)
49
+
50
+ # Configure logging
51
+ if config.app and config.app.log_level:
52
+ configure_logging(level=config.app.log_level)
53
+
54
+ # Create the ResponsesAgent
55
+ _responses_agent: ResponsesAgent = config.as_responses_agent()
56
+
57
+
58
+ @invoke()
59
+ def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
60
+ """
61
+ Handle non-streaming requests by delegating to the ResponsesAgent.
62
+
63
+ Args:
64
+ request: The incoming ResponsesAgentRequest
65
+
66
+ Returns:
67
+ ResponsesAgentResponse with the complete output
68
+ """
69
+ return _responses_agent.predict(request)
70
+
71
+
72
+ @stream()
73
+ def streaming(
74
+ request: ResponsesAgentRequest,
75
+ ) -> AsyncGenerator[ResponsesAgentStreamEvent, None]:
76
+ """
77
+ Handle streaming requests by delegating to the ResponsesAgent.
78
+
79
+ Args:
80
+ request: The incoming ResponsesAgentRequest
81
+
82
+ Yields:
83
+ ResponsesAgentStreamEvent objects as they are generated
84
+ """
85
+ # The predict_stream method returns a generator, convert to async generator
86
+ for event in _responses_agent.predict_stream(request):
87
+ yield event
88
+
89
+
90
+ # Create the AgentServer instance
91
+ agent_server = AgentServer("ResponsesAgent", enable_chat_proxy=True)
92
+
93
+ # Define the app as a module level variable to enable multiple workers
94
+ app = agent_server.app
95
+
96
+
97
+ def main() -> None:
98
+ """Entry point for running the agent server."""
99
+ agent_server.run(app_import_string="dao_ai.app_server:app")
100
+
101
+
102
+ if __name__ == "__main__":
103
+ main()
@@ -309,6 +309,14 @@ Examples:
309
309
  metavar="FILE",
310
310
  help="Path to the model configuration file to validate",
311
311
  )
312
+ deploy_parser.add_argument(
313
+ "-t",
314
+ "--target",
315
+ type=str,
316
+ choices=["model_serving", "apps"],
317
+ default="model_serving",
318
+ help="Deployment target: 'model_serving' (default) or 'apps'",
319
+ )
312
320
 
313
321
  # List MCP tools command
314
322
  list_mcp_parser: ArgumentParser = subparsers.add_parser(
@@ -729,11 +737,17 @@ def handle_graph_command(options: Namespace) -> None:
729
737
 
730
738
 
731
739
  def handle_deploy_command(options: Namespace) -> None:
740
+ from dao_ai.config import DeploymentTarget
741
+
732
742
  logger.debug(f"Validating configuration from {options.config}...")
733
743
  try:
734
744
  config: AppConfig = AppConfig.from_file(options.config)
745
+
746
+ # Convert target string to enum
747
+ target: DeploymentTarget = DeploymentTarget(options.target)
748
+
735
749
  config.create_agent()
736
- config.deploy_agent()
750
+ config.deploy_agent(target=target)
737
751
  sys.exit(0)
738
752
  except Exception as e:
739
753
  logger.error(f"Deployment failed: {e}")
@@ -208,7 +208,9 @@ class IsDatabricksResource(ABC, BaseModel):
208
208
  Authentication Options:
209
209
  ----------------------
210
210
  1. **On-Behalf-Of User (OBO)**: Set on_behalf_of_user=True to use the
211
- calling user's identity via ModelServingUserCredentials.
211
+ calling user's identity. Implementation varies by deployment:
212
+ - Databricks Apps: Uses X-Forwarded-Access-Token from request headers
213
+ - Model Serving: Uses ModelServingUserCredentials
212
214
 
213
215
  2. **Service Principal (OAuth M2M)**: Provide service_principal or
214
216
  (client_id + client_secret + workspace_host) for service principal auth.
@@ -221,9 +223,17 @@ class IsDatabricksResource(ABC, BaseModel):
221
223
 
222
224
  Authentication Priority:
223
225
  1. OBO (on_behalf_of_user=True)
226
+ - Checks for forwarded headers (Databricks Apps)
227
+ - Falls back to ModelServingUserCredentials (Model Serving)
224
228
  2. Service Principal (client_id + client_secret + workspace_host)
225
229
  3. PAT (pat + workspace_host)
226
230
  4. Ambient/default authentication
231
+
232
+ Note: When on_behalf_of_user=True, the agent acts as the calling user regardless
233
+ of deployment target. In Databricks Apps, this uses X-Forwarded-Access-Token
234
+ automatically captured by MLflow AgentServer. In Model Serving, this uses
235
+ ModelServingUserCredentials. Forwarded headers are ONLY used when
236
+ on_behalf_of_user=True.
227
237
  """
228
238
 
229
239
  model_config = ConfigDict(use_enum_values=True)
@@ -235,9 +245,6 @@ class IsDatabricksResource(ABC, BaseModel):
235
245
  workspace_host: Optional[AnyVariable] = None
236
246
  pat: Optional[AnyVariable] = None
237
247
 
238
- # Private attribute to cache the workspace client (lazy instantiation)
239
- _workspace_client: Optional[WorkspaceClient] = PrivateAttr(default=None)
240
-
241
248
  @abstractmethod
242
249
  def as_resources(self) -> Sequence[DatabricksResource]: ...
243
250
 
@@ -273,32 +280,56 @@ class IsDatabricksResource(ABC, BaseModel):
273
280
  """
274
281
  Get a WorkspaceClient configured with the appropriate authentication.
275
282
 
276
- The client is lazily instantiated on first access and cached for subsequent calls.
283
+ A new client is created on each access.
277
284
 
278
285
  Authentication priority:
279
- 1. If on_behalf_of_user is True, uses ModelServingUserCredentials (OBO)
280
- 2. If service principal credentials are configured (client_id, client_secret,
281
- workspace_host), uses OAuth M2M
282
- 3. If PAT is configured, uses token authentication
283
- 4. Otherwise, uses default/ambient authentication
286
+ 1. On-Behalf-Of User (on_behalf_of_user=True):
287
+ - Forwarded headers (Databricks Apps)
288
+ - ModelServingUserCredentials (Model Serving)
289
+ 2. Service Principal (client_id + client_secret + workspace_host)
290
+ 3. PAT (pat + workspace_host)
291
+ 4. Ambient/default authentication
284
292
  """
285
- # Return cached client if already instantiated
286
- if self._workspace_client is not None:
287
- return self._workspace_client
288
-
289
293
  from dao_ai.utils import normalize_host
290
294
 
291
295
  # Check for OBO first (highest priority)
292
296
  if self.on_behalf_of_user:
297
+ # NEW: In Databricks Apps, use forwarded headers for per-user auth
298
+ try:
299
+ from mlflow.genai.agent_server import get_request_headers
300
+
301
+ headers = get_request_headers()
302
+ forwarded_token = headers.get("x-forwarded-access-token")
303
+
304
+ if forwarded_token:
305
+ forwarded_user = headers.get("x-forwarded-user", "unknown")
306
+ logger.debug(
307
+ f"Creating WorkspaceClient for {self.__class__.__name__} "
308
+ f"with OBO using forwarded token from Databricks Apps",
309
+ forwarded_user=forwarded_user,
310
+ )
311
+ # Use workspace_host if configured, otherwise SDK will auto-detect
312
+ workspace_host_value: str | None = (
313
+ normalize_host(value_of(self.workspace_host))
314
+ if self.workspace_host
315
+ else None
316
+ )
317
+ return WorkspaceClient(
318
+ host=workspace_host_value,
319
+ token=forwarded_token,
320
+ auth_type="pat",
321
+ )
322
+ except (ImportError, LookupError):
323
+ # mlflow not available or headers not set - fall through to Model Serving
324
+ pass
325
+
326
+ # Fall back to Model Serving OBO (existing behavior)
293
327
  credentials_strategy: CredentialsStrategy = ModelServingUserCredentials()
294
328
  logger.debug(
295
329
  f"Creating WorkspaceClient for {self.__class__.__name__} "
296
- f"with OBO credentials strategy"
297
- )
298
- self._workspace_client = WorkspaceClient(
299
- credentials_strategy=credentials_strategy
330
+ f"with OBO credentials strategy (Model Serving)"
300
331
  )
301
- return self._workspace_client
332
+ return WorkspaceClient(credentials_strategy=credentials_strategy)
302
333
 
303
334
  # Check for service principal credentials
304
335
  client_id_value: str | None = (
@@ -318,13 +349,12 @@ class IsDatabricksResource(ABC, BaseModel):
318
349
  f"Creating WorkspaceClient for {self.__class__.__name__} with service principal: "
319
350
  f"client_id={client_id_value}, host={workspace_host_value}"
320
351
  )
321
- self._workspace_client = WorkspaceClient(
352
+ return WorkspaceClient(
322
353
  host=workspace_host_value,
323
354
  client_id=client_id_value,
324
355
  client_secret=client_secret_value,
325
356
  auth_type="oauth-m2m",
326
357
  )
327
- return self._workspace_client
328
358
 
329
359
  # Check for PAT authentication
330
360
  pat_value: str | None = value_of(self.pat) if self.pat else None
@@ -332,20 +362,28 @@ class IsDatabricksResource(ABC, BaseModel):
332
362
  logger.debug(
333
363
  f"Creating WorkspaceClient for {self.__class__.__name__} with PAT"
334
364
  )
335
- self._workspace_client = WorkspaceClient(
365
+ return WorkspaceClient(
336
366
  host=workspace_host_value,
337
367
  token=pat_value,
338
368
  auth_type="pat",
339
369
  )
340
- return self._workspace_client
341
370
 
342
371
  # Default: use ambient authentication
343
372
  logger.debug(
344
373
  f"Creating WorkspaceClient for {self.__class__.__name__} "
345
374
  "with default/ambient authentication"
346
375
  )
347
- self._workspace_client = WorkspaceClient()
348
- return self._workspace_client
376
+ return WorkspaceClient()
377
+
378
+
379
+ class DeploymentTarget(str, Enum):
380
+ """Target platform for agent deployment."""
381
+
382
+ MODEL_SERVING = "model_serving"
383
+ """Deploy to Databricks Model Serving endpoint."""
384
+
385
+ APPS = "apps"
386
+ """Deploy as a Databricks App."""
349
387
 
350
388
 
351
389
  class Privilege(str, Enum):
@@ -865,10 +903,6 @@ class GenieRoomModel(IsDatabricksResource):
865
903
  pat=self.pat,
866
904
  )
867
905
 
868
- # Share the cached workspace client if available
869
- if self._workspace_client is not None:
870
- warehouse_model._workspace_client = self._workspace_client
871
-
872
906
  return warehouse_model
873
907
  except Exception as e:
874
908
  logger.warning(
@@ -912,9 +946,6 @@ class GenieRoomModel(IsDatabricksResource):
912
946
  workspace_host=self.workspace_host,
913
947
  pat=self.pat,
914
948
  )
915
- # Share the cached workspace client if available
916
- if self._workspace_client is not None:
917
- table_model._workspace_client = self._workspace_client
918
949
 
919
950
  # Verify the table exists before adding
920
951
  if not table_model.exists():
@@ -952,9 +983,6 @@ class GenieRoomModel(IsDatabricksResource):
952
983
  workspace_host=self.workspace_host,
953
984
  pat=self.pat,
954
985
  )
955
- # Share the cached workspace client if available
956
- if self._workspace_client is not None:
957
- function_model._workspace_client = self._workspace_client
958
986
 
959
987
  # Verify the function exists before adding
960
988
  if not function_model.exists():
@@ -3255,6 +3283,7 @@ class ResourcesModel(BaseModel):
3255
3283
 
3256
3284
  class AppConfig(BaseModel):
3257
3285
  model_config = ConfigDict(use_enum_values=True, extra="forbid")
3286
+ version: Optional[str] = None
3258
3287
  variables: dict[str, AnyVariable] = Field(default_factory=dict)
3259
3288
  service_principals: dict[str, ServicePrincipalModel] = Field(default_factory=dict)
3260
3289
  schemas: dict[str, SchemaModel] = Field(default_factory=dict)
@@ -3275,6 +3304,9 @@ class AppConfig(BaseModel):
3275
3304
  )
3276
3305
  providers: Optional[dict[type | str, Any]] = None
3277
3306
 
3307
+ # Private attribute to track the source config file path (set by from_file)
3308
+ _source_config_path: str | None = None
3309
+
3278
3310
  @classmethod
3279
3311
  def from_file(cls, path: PathLike) -> "AppConfig":
3280
3312
  path = Path(path).as_posix()
@@ -3282,12 +3314,20 @@ class AppConfig(BaseModel):
3282
3314
  model_config: ModelConfig = ModelConfig(development_config=path)
3283
3315
  config: AppConfig = AppConfig(**model_config.to_dict())
3284
3316
 
3317
+ # Store the source config path for later use (e.g., Apps deployment)
3318
+ config._source_config_path = path
3319
+
3285
3320
  config.initialize()
3286
3321
 
3287
3322
  atexit.register(config.shutdown)
3288
3323
 
3289
3324
  return config
3290
3325
 
3326
+ @property
3327
+ def source_config_path(self) -> str | None:
3328
+ """Get the source config file path if loaded via from_file."""
3329
+ return self._source_config_path
3330
+
3291
3331
  def initialize(self) -> None:
3292
3332
  from dao_ai.hooks.core import create_hooks
3293
3333
  from dao_ai.logging import configure_logging
@@ -3358,6 +3398,7 @@ class AppConfig(BaseModel):
3358
3398
 
3359
3399
  def deploy_agent(
3360
3400
  self,
3401
+ target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
3361
3402
  w: WorkspaceClient | None = None,
3362
3403
  vsc: "VectorSearchClient | None" = None,
3363
3404
  pat: str | None = None,
@@ -3365,6 +3406,18 @@ class AppConfig(BaseModel):
3365
3406
  client_secret: str | None = None,
3366
3407
  workspace_host: str | None = None,
3367
3408
  ) -> None:
3409
+ """
3410
+ Deploy the agent to the specified target.
3411
+
3412
+ Args:
3413
+ target: The deployment target (MODEL_SERVING or APPS). Defaults to MODEL_SERVING.
3414
+ w: Optional WorkspaceClient instance
3415
+ vsc: Optional VectorSearchClient instance
3416
+ pat: Optional personal access token for authentication
3417
+ client_id: Optional client ID for service principal authentication
3418
+ client_secret: Optional client secret for service principal authentication
3419
+ workspace_host: Optional workspace host URL
3420
+ """
3368
3421
  from dao_ai.providers.base import ServiceProvider
3369
3422
  from dao_ai.providers.databricks import DatabricksProvider
3370
3423
 
@@ -3376,7 +3429,7 @@ class AppConfig(BaseModel):
3376
3429
  client_secret=client_secret,
3377
3430
  workspace_host=workspace_host,
3378
3431
  )
3379
- provider.deploy_agent(self)
3432
+ provider.deploy_agent(self, target=target)
3380
3433
 
3381
3434
  def find_agents(
3382
3435
  self, predicate: Callable[[AgentModel], bool] | None = None
@@ -1,15 +1,19 @@
1
1
  from abc import ABC, abstractmethod
2
- from typing import Any, Sequence
2
+ from typing import TYPE_CHECKING, Any, Sequence
3
3
 
4
4
  from dao_ai.config import (
5
5
  AppModel,
6
6
  DatasetModel,
7
+ DeploymentTarget,
7
8
  SchemaModel,
8
9
  UnityCatalogFunctionSqlModel,
9
10
  VectorStoreModel,
10
11
  VolumeModel,
11
12
  )
12
13
 
14
+ if TYPE_CHECKING:
15
+ from dao_ai.config import AppConfig
16
+
13
17
 
14
18
  class ServiceProvider(ABC):
15
19
  @abstractmethod
@@ -52,4 +56,26 @@ class ServiceProvider(ABC):
52
56
  ) -> Any: ...
53
57
 
54
58
  @abstractmethod
55
- def deploy_agent(self, config: AppModel) -> Any: ...
59
+ def deploy_model_serving_agent(self, config: "AppConfig") -> Any:
60
+ """Deploy agent to Databricks Model Serving endpoint."""
61
+ ...
62
+
63
+ @abstractmethod
64
+ def deploy_apps_agent(self, config: "AppConfig") -> Any:
65
+ """Deploy agent as a Databricks App."""
66
+ ...
67
+
68
+ @abstractmethod
69
+ def deploy_agent(
70
+ self,
71
+ config: "AppConfig",
72
+ target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
73
+ ) -> Any:
74
+ """
75
+ Deploy agent to the specified target.
76
+
77
+ Args:
78
+ config: The AppConfig containing deployment configuration
79
+ target: The deployment target (MODEL_SERVING or APPS)
80
+ """
81
+ ...