dao-ai 0.1.9__tar.gz → 0.1.10__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 (305) hide show
  1. {dao_ai-0.1.9 → dao_ai-0.1.10}/.gitignore +3 -0
  2. {dao_ai-0.1.9 → dao_ai-0.1.10}/PKG-INFO +1 -1
  3. {dao_ai-0.1.9 → dao_ai-0.1.10}/app.yaml +1 -1
  4. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/hardware_store_lakebase.yaml +16 -0
  5. {dao_ai-0.1.9 → dao_ai-0.1.10}/databricks.yaml.template +4 -0
  6. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/05_deploy_agent.py +21 -3
  7. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/99_scratchpad.py +1 -1
  8. {dao_ai-0.1.9 → dao_ai-0.1.10}/pyproject.toml +1 -1
  9. {dao_ai-0.1.9 → dao_ai-0.1.10}/requirements.txt +1 -0
  10. {dao_ai-0.1.9 → dao_ai-0.1.10}/schemas/model_config_schema.json +21 -0
  11. dao_ai-0.1.10/src/dao_ai/apps/__init__.py +24 -0
  12. dao_ai-0.1.9/src/dao_ai/app_server.py → dao_ai-0.1.10/src/dao_ai/apps/handlers.py +20 -39
  13. dao_ai-0.1.10/src/dao_ai/apps/model_serving.py +29 -0
  14. dao_ai-0.1.10/src/dao_ai/apps/resources.py +1029 -0
  15. dao_ai-0.1.10/src/dao_ai/apps/server.py +39 -0
  16. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/cli.py +51 -4
  17. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/config.py +34 -4
  18. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/memory/postgres.py +29 -4
  19. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/models.py +327 -370
  20. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/providers/databricks.py +52 -20
  21. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/mcp.py +165 -68
  22. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_databricks.py +109 -0
  23. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_input_output_structure.py +6 -4
  24. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_responses_agent_structured_output_unit.py +61 -19
  25. dao_ai-0.1.9/src/dao_ai/agent_as_code.py +0 -22
  26. {dao_ai-0.1.9 → dao_ai-0.1.10}/.python-version +0 -0
  27. {dao_ai-0.1.9 → dao_ai-0.1.10}/CHANGELOG.md +0 -0
  28. {dao_ai-0.1.9 → dao_ai-0.1.10}/CONTRIBUTING.md +0 -0
  29. {dao_ai-0.1.9 → dao_ai-0.1.10}/LICENSE +0 -0
  30. {dao_ai-0.1.9 → dao_ai-0.1.10}/Makefile +0 -0
  31. {dao_ai-0.1.9 → dao_ai-0.1.10}/README.md +0 -0
  32. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/01_getting_started/README.md +0 -0
  33. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/01_getting_started/minimal.yaml +0 -0
  34. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/02_mcp/README.md +0 -0
  35. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/02_mcp/custom_mcp.yaml +0 -0
  36. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/02_mcp/external_mcp.yaml +0 -0
  37. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/02_mcp/filtered_mcp.yaml +0 -0
  38. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/02_mcp/managed_mcp.yaml +0 -0
  39. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/02_mcp/slack_integration.yaml +0 -0
  40. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/03_reranking/README.md +0 -0
  41. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/03_reranking/vector_search_with_reranking.yaml +0 -0
  42. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/04_genie/README.md +0 -0
  43. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/04_genie/genie_basic.yaml +0 -0
  44. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/04_genie/genie_lru_cache.yaml +0 -0
  45. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/04_genie/genie_semantic_cache.yaml +0 -0
  46. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/04_genie/genie_with_conversation_id.yaml +0 -0
  47. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/05_memory/README.md +0 -0
  48. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/05_memory/conversation_summarization.yaml +0 -0
  49. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/05_memory/in_memory_basic.yaml +0 -0
  50. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/05_memory/lakebase_persistence.yaml +0 -0
  51. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/05_memory/postgres_persistence.yaml +0 -0
  52. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/06_on_behalf_of_user/README.md +0 -0
  53. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/06_on_behalf_of_user/obo_basic.yaml +0 -0
  54. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/07_human_in_the_loop/README.md +0 -0
  55. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/07_human_in_the_loop/human_in_the_loop.yaml +0 -0
  56. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/08_guardrails/README.md +0 -0
  57. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/08_guardrails/guardrails_basic.yaml +0 -0
  58. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/09_structured_output/README.md +0 -0
  59. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/09_structured_output/structured_output.yaml +0 -0
  60. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/10_agent_integrations/README.md +0 -0
  61. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/10_agent_integrations/agent_bricks.yaml +0 -0
  62. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/10_agent_integrations/kasal.yaml +0 -0
  63. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/11_prompt_engineering/README.md +0 -0
  64. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/11_prompt_engineering/prompt_optimization.yaml +0 -0
  65. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/11_prompt_engineering/prompt_registry.yaml +0 -0
  66. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/README.md +0 -0
  67. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/combined_middleware.yaml +0 -0
  68. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/context_management.yaml +0 -0
  69. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/custom_field_validation.yaml +0 -0
  70. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/limit_middleware.yaml +0 -0
  71. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/logging_middleware.yaml +0 -0
  72. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/pii_middleware.yaml +0 -0
  73. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/retry_middleware.yaml +0 -0
  74. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/12_middleware/tool_selector_middleware.yaml +0 -0
  75. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/13_orchestration/README.md +0 -0
  76. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/13_orchestration/supervisor_pattern.yaml +0 -0
  77. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/13_orchestration/swarm_pattern.yaml +0 -0
  78. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/14_basic_tools/README.md +0 -0
  79. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/14_basic_tools/sql_tool_example.yaml +0 -0
  80. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/README.md +0 -0
  81. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/brick_store.yaml +0 -0
  82. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/deep_research.yaml +0 -0
  83. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/executive_assistant.yaml +0 -0
  84. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/genie_and_genie_mcp.yaml +0 -0
  85. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/genie_vector_search_hybrid.yaml +0 -0
  86. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/hardware_store.yaml +0 -0
  87. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/hardware_store_swarm.yaml +0 -0
  88. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/quick_serve_restaurant.yaml +0 -0
  89. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/15_complete_applications/reservations_system.yaml +0 -0
  90. {dao_ai-0.1.9 → dao_ai-0.1.10}/config/examples/README.md +0 -0
  91. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/appointments.sql +0 -0
  92. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/appointments_data.sql +0 -0
  93. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/brand_rep_demo_data.sql +0 -0
  94. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/brand_rep_demo_queries.sql +0 -0
  95. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/brand_rep_demo_tables.sql +0 -0
  96. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/brand_rep_demo_validation.sql +0 -0
  97. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/customers.sql +0 -0
  98. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/customers_data.sql +0 -0
  99. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/dim_stores.sql +0 -0
  100. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/dim_stores_data.sql +0 -0
  101. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/employee_performance.sql +0 -0
  102. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/employee_performance_data.sql +0 -0
  103. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/employee_tasks.sql +0 -0
  104. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/employee_tasks_data.sql +0 -0
  105. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/inventory.sql +0 -0
  106. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/inventory_data.sql +0 -0
  107. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/managers.sql +0 -0
  108. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/managers_data.sql +0 -0
  109. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/product_data.sql +0 -0
  110. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/products.sql +0 -0
  111. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/dais2025/task_assignments.sql +0 -0
  112. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/hardware_store/inventory.snappy.parquet +0 -0
  113. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/hardware_store/inventory.sql +0 -0
  114. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/hardware_store/products.snappy.parquet +0 -0
  115. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/hardware_store/products.sql +0 -0
  116. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/quick_serve_restaurant/.gitkeep +0 -0
  117. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/quick_serve_restaurant/fulfil_item_orders.sql +0 -0
  118. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/quick_serve_restaurant/items_description.csv +0 -0
  119. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/quick_serve_restaurant/items_description.sql +0 -0
  120. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/quick_serve_restaurant/items_raw.csv +0 -0
  121. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/quick_serve_restaurant/items_raw.sql +0 -0
  122. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/quick_serve_restaurant/orders_raw.csv +0 -0
  123. {dao_ai-0.1.9 → dao_ai-0.1.10}/data/quick_serve_restaurant/orders_raw.sql +0 -0
  124. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/architecture.md +0 -0
  125. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/cli-reference.md +0 -0
  126. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/configuration-reference.md +0 -0
  127. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/contributing.md +0 -0
  128. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/examples.md +0 -0
  129. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/faq.md +0 -0
  130. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/hardware_store/README.md +0 -0
  131. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/hardware_store/retail_supervisor.png +0 -0
  132. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/hardware_store/retail_swarm.png +0 -0
  133. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/images/genie.png +0 -0
  134. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/key-capabilities.md +0 -0
  135. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/python-api.md +0 -0
  136. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/quick_serve_restaurant/.gitkeep +0 -0
  137. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/quick_serve_restaurant/quick-serve-restaurant.png +0 -0
  138. {dao_ai-0.1.9 → dao_ai-0.1.10}/docs/why-dao.md +0 -0
  139. {dao_ai-0.1.9 → dao_ai-0.1.10}/environment.yaml +0 -0
  140. {dao_ai-0.1.9 → dao_ai-0.1.10}/examples/dais2025/examples.yaml +0 -0
  141. {dao_ai-0.1.9 → dao_ai-0.1.10}/examples/deep_research/examples.yaml +0 -0
  142. {dao_ai-0.1.9 → dao_ai-0.1.10}/examples/executive_assistant/examples.yaml +0 -0
  143. {dao_ai-0.1.9 → dao_ai-0.1.10}/examples/hardware_store/examples.yaml +0 -0
  144. {dao_ai-0.1.9 → dao_ai-0.1.10}/examples/quick_serve_restaurant/.gitkeep +0 -0
  145. {dao_ai-0.1.9 → dao_ai-0.1.10}/examples/quick_serve_restaurant/examples.yaml +0 -0
  146. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/dais2025/extract_store_numbers.sql +0 -0
  147. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/dais2025/find_inventory_by_sku.sql +0 -0
  148. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/dais2025/find_inventory_by_upc.sql +0 -0
  149. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/dais2025/find_product_by_sku.sql +0 -0
  150. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/dais2025/find_product_by_upc.sql +0 -0
  151. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/dais2025/find_store_by_number.sql +0 -0
  152. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/dais2025/find_store_inventory_by_sku.sql +0 -0
  153. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/dais2025/find_store_inventory_by_upc.sql +0 -0
  154. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/hardware_store/find_inventory_by_sku.sql +0 -0
  155. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/hardware_store/find_inventory_by_upc.sql +0 -0
  156. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/hardware_store/find_product_by_sku.sql +0 -0
  157. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/hardware_store/find_product_by_upc.sql +0 -0
  158. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/hardware_store/find_store_inventory_by_sku.sql +0 -0
  159. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/hardware_store/find_store_inventory_by_upc.sql +0 -0
  160. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/quick_serve_restaurant/.gitkeep +0 -0
  161. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/quick_serve_restaurant/insert_coffee_order.sql +0 -0
  162. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/quick_serve_restaurant/lookup_items_by_descriptions.sql +0 -0
  163. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/quick_serve_restaurant/match_historical_item_order_by_date.sql +0 -0
  164. {dao_ai-0.1.9 → dao_ai-0.1.10}/functions/quick_serve_restaurant/match_item_by_description_and_price.sql +0 -0
  165. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/01_ingest_and_transform.py +0 -0
  166. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/02_provision_vector_search.py +0 -0
  167. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/03_provision_lakebase.py +0 -0
  168. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/04_unity_catalog_tools.py +0 -0
  169. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/06_generate_evaluation_data.py +0 -0
  170. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/07_run_evaluation.py +0 -0
  171. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/08_run_examples.py +0 -0
  172. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/09_evaluate_inferences.py +0 -0
  173. {dao_ai-0.1.9 → dao_ai-0.1.10}/notebooks/10_optimize_prompts.py +0 -0
  174. {dao_ai-0.1.9 → dao_ai-0.1.10}/schemas/bundle_config_schema.json +0 -0
  175. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/__init__.py +0 -0
  176. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/models.py +0 -0
  177. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/tools/__init__.py +0 -0
  178. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/tools/customer.py +0 -0
  179. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/tools/employee.py +0 -0
  180. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/tools/executive.py +0 -0
  181. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/tools/genie.py +0 -0
  182. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/tools/inventory.py +0 -0
  183. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/tools/models.py +0 -0
  184. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dais2025/tools/store.py +0 -0
  185. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/__init__.py +0 -0
  186. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/catalog.py +0 -0
  187. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/genie/__init__.py +0 -0
  188. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/genie/cache/__init__.py +0 -0
  189. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/genie/cache/base.py +0 -0
  190. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/genie/cache/core.py +0 -0
  191. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/genie/cache/lru.py +0 -0
  192. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/genie/cache/semantic.py +0 -0
  193. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/genie/core.py +0 -0
  194. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/graph.py +0 -0
  195. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/hooks/__init__.py +0 -0
  196. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/hooks/core.py +0 -0
  197. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/logging.py +0 -0
  198. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/memory/__init__.py +0 -0
  199. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/memory/base.py +0 -0
  200. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/memory/core.py +0 -0
  201. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/memory/databricks.py +0 -0
  202. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/messages.py +0 -0
  203. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/__init__.py +0 -0
  204. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/assertions.py +0 -0
  205. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/base.py +0 -0
  206. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/context_editing.py +0 -0
  207. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/core.py +0 -0
  208. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/guardrails.py +0 -0
  209. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/human_in_the_loop.py +0 -0
  210. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/message_validation.py +0 -0
  211. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/model_call_limit.py +0 -0
  212. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/model_retry.py +0 -0
  213. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/pii.py +0 -0
  214. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/summarization.py +0 -0
  215. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/tool_call_limit.py +0 -0
  216. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/tool_retry.py +0 -0
  217. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/middleware/tool_selector.py +0 -0
  218. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/nodes.py +0 -0
  219. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/optimization.py +0 -0
  220. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/orchestration/__init__.py +0 -0
  221. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/orchestration/core.py +0 -0
  222. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/orchestration/supervisor.py +0 -0
  223. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/orchestration/swarm.py +0 -0
  224. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/prompts.py +0 -0
  225. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/providers/__init__.py +0 -0
  226. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/providers/base.py +0 -0
  227. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/state.py +0 -0
  228. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/__init__.py +0 -0
  229. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/agent.py +0 -0
  230. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/core.py +0 -0
  231. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/email.py +0 -0
  232. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/genie.py +0 -0
  233. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/memory.py +0 -0
  234. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/python.py +0 -0
  235. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/search.py +0 -0
  236. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/slack.py +0 -0
  237. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/sql.py +0 -0
  238. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/time.py +0 -0
  239. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/unity_catalog.py +0 -0
  240. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/tools/vector_search.py +0 -0
  241. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/types.py +0 -0
  242. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/utils.py +0 -0
  243. {dao_ai-0.1.9 → dao_ai-0.1.10}/src/dao_ai/vector_search.py +0 -0
  244. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/config/test_model_config.yaml +0 -0
  245. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/conftest.py +0 -0
  246. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/middleware/test_context_editing.py +0 -0
  247. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/middleware/test_model_call_limit.py +0 -0
  248. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/middleware/test_model_retry.py +0 -0
  249. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/middleware/test_pii.py +0 -0
  250. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/middleware/test_tool_call_limit.py +0 -0
  251. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/middleware/test_tool_retry.py +0 -0
  252. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/middleware/test_tool_selector.py +0 -0
  253. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_agent_response_format.py +0 -0
  254. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_assertions_middleware.py +0 -0
  255. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_catalog.py +0 -0
  256. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_chat_history.py +0 -0
  257. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_config.py +0 -0
  258. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_function_parsing.py +0 -0
  259. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_genie.py +0 -0
  260. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_genie_conversation_ids_in_outputs.py +0 -0
  261. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_genie_databricks_integration.py +0 -0
  262. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_genie_room_model.py +0 -0
  263. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_guardrail_retry.py +0 -0
  264. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_hitl_config_model.py +0 -0
  265. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_hitl_responses_agent.py +0 -0
  266. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_hooks.py +0 -0
  267. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_human_in_the_loop.py +0 -0
  268. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_inference.py +0 -0
  269. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_inference_integration.py +0 -0
  270. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_interrupt_type.py +0 -0
  271. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_llm_interrupt_handling.py +0 -0
  272. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_mcp.py +0 -0
  273. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_mcp_filtering.py +0 -0
  274. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_mcp_filtering_integration.py +0 -0
  275. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_mcp_function_model.py +0 -0
  276. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_message_validation_middleware.py +0 -0
  277. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_messages.py +0 -0
  278. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_models.py +0 -0
  279. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_optimization.py +0 -0
  280. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_postgres_integration.py +0 -0
  281. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_prompt_optimizations.py +0 -0
  282. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_prompts.py +0 -0
  283. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_reranking.py +0 -0
  284. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_reranking_integration.py +0 -0
  285. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_resources_model_genie_integration.py +0 -0
  286. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_response_format.py +0 -0
  287. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_semantic_cache_context.py +0 -0
  288. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_sql_tool.py +0 -0
  289. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_sql_tool_integration.py +0 -0
  290. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_state.py +0 -0
  291. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_summarization_inference.py +0 -0
  292. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_swarm_middleware.py +0 -0
  293. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_tools.py +0 -0
  294. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_types.py +0 -0
  295. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_unity_catalog.py +0 -0
  296. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_utils.py +0 -0
  297. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_utils_type_from_fqn.py +0 -0
  298. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_vector_search.py +0 -0
  299. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/test_warehouse_model.py +0 -0
  300. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/dao_ai/weather_server_mcp.py +0 -0
  301. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/hardware_store/.gitkeep +0 -0
  302. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/hardware_store/test_graph.py +0 -0
  303. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/images/doritos_upc.png +0 -0
  304. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/images/lays_upc.png +0 -0
  305. {dao_ai-0.1.9 → dao_ai-0.1.10}/tests/quick_serve_restaurant/.gitkeep +0 -0
@@ -278,3 +278,6 @@ mlruns/
278
278
  databricks.yaml
279
279
 
280
280
  .cursor/
281
+
282
+ # Deployment directory for local source code
283
+ deploy/
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dao-ai
3
- Version: 0.1.9
3
+ Version: 0.1.10
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
@@ -1,4 +1,4 @@
1
- command: ["uv", "run", "python", "-m", "dao_ai.app_server"]
1
+ command: ["uv", "run", "python", "-m", "dao_ai.apps.server"]
2
2
  # Databricks Apps listen by default on port 8000
3
3
 
4
4
  env:
@@ -10,6 +10,20 @@
10
10
  # DATABRICKS SCHEMAS
11
11
  # =============================================================================
12
12
  # Define Unity Catalog schemas where data and functions will be stored
13
+
14
+
15
+ variables:
16
+ client_id: &client_id
17
+ options:
18
+ - env: RETAIL_AI_DATABRICKS_CLIENT_ID # Service principal client ID
19
+ - scope: retail_consumer_goods
20
+ secret: RETAIL_AI_DATABRICKS_CLIENT_ID
21
+ client_secret: &client_secret
22
+ options:
23
+ - env: RETAIL_AI_DATABRICKS_CLIENT_SECRET # Service principal secret
24
+ - scope: retail_consumer_goods
25
+ secret: RETAIL_AI_DATABRICKS_CLIENT_SECRET
26
+
13
27
  schemas:
14
28
  retail_schema: &retail_schema
15
29
  catalog_name: retail_consumer_goods # Unity Catalog name
@@ -134,6 +148,8 @@ resources:
134
148
  instance_name: "retail-consumer-goods" # Use instance_name for Databricks Lakebase
135
149
  description: "Database for agent memory and checkpoints"
136
150
 
151
+ client_id: *client_id
152
+ client_secret: *client_secret
137
153
  on_behalf_of_user: false
138
154
 
139
155
  # =============================================================================
@@ -16,6 +16,9 @@ variables:
16
16
  description: The job compute node type (cloud-specific, set per target)
17
17
  cloud:
18
18
  description: Cloud provider (azure, aws, gcp) - auto-detected or set per target
19
+ deployment_target:
20
+ description: Agent deployment target (model_serving or apps) - defaults to model_serving
21
+ default: model_serving
19
22
 
20
23
  resources:
21
24
  jobs:
@@ -83,6 +86,7 @@ resources:
83
86
  notebook_path: ./notebooks/05_deploy_agent.py
84
87
  base_parameters:
85
88
  config-path: ${var.config_path}
89
+ deployment-target: ${var.deployment_target}
86
90
  environment_key: dao-ai-env
87
91
  - task_key: run-evaluation
88
92
  depends_on:
@@ -28,16 +28,23 @@ def find_yaml_files_os_walk(base_path: str) -> Sequence[str]:
28
28
  # COMMAND ----------
29
29
 
30
30
  dbutils.widgets.text(name="config-path", defaultValue="")
31
+ dbutils.widgets.dropdown(
32
+ name="deployment-target",
33
+ choices=["", "model_serving", "apps"],
34
+ defaultValue="",
35
+ )
31
36
 
32
37
  config_files: Sequence[str] = find_yaml_files_os_walk("../config")
33
38
  dbutils.widgets.dropdown(name="config-paths", choices=config_files, defaultValue=next(iter(config_files), ""))
34
39
 
35
40
  config_path: str | None = dbutils.widgets.get("config-path") or None
36
41
  project_path: str = dbutils.widgets.get("config-paths") or None
42
+ deployment_target_str: str | None = dbutils.widgets.get("deployment-target") or None
37
43
 
38
44
  config_path: str = config_path or project_path
39
45
 
40
- print(config_path)
46
+ print(f"Config path: {config_path}")
47
+ print(f"Deployment target: {deployment_target_str or '(using config default)'}")
41
48
 
42
49
  # COMMAND ----------
43
50
 
@@ -93,10 +100,21 @@ nest_asyncio.apply()
93
100
 
94
101
  # COMMAND ----------
95
102
 
96
- from dao_ai.config import AppConfig
103
+ from dao_ai.config import AppConfig, DeploymentTarget
97
104
 
98
105
  config: AppConfig = AppConfig.from_file(path=config_path)
99
106
 
107
+ # Resolve deployment target from widget (hybrid resolution)
108
+ # If widget is empty/None, deploy_agent() will use config.app.deployment_target or default
109
+ deployment_target: DeploymentTarget | None = None
110
+ if deployment_target_str:
111
+ deployment_target = DeploymentTarget(deployment_target_str)
112
+ print(f"Using widget-specified deployment target: {deployment_target.value}")
113
+ elif config.app and config.app.deployment_target:
114
+ print(f"Using config file deployment target: {config.app.deployment_target.value}")
115
+ else:
116
+ print("Using default deployment target: model_serving")
117
+
100
118
  # COMMAND ----------
101
119
 
102
120
  config.display_graph()
@@ -107,4 +125,4 @@ config.create_agent()
107
125
 
108
126
  # COMMAND ----------
109
127
 
110
- config.deploy_agent()
128
+ config.deploy_agent(target=deployment_target)
@@ -618,7 +618,7 @@ result = reflection_app.invoke({"messages": example_query})
618
618
 
619
619
  # COMMAND ----------
620
620
 
621
- from agent_as_code import app
621
+ from dao_ai.apps.model_serving import app
622
622
 
623
623
 
624
624
  # COMMAND ----------
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "dao-ai"
7
- version = "0.1.9"
7
+ version = "0.1.10"
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" }
@@ -71,6 +71,7 @@ grandalf==0.8
71
71
  graphene==3.4.3
72
72
  graphql-core==3.2.7
73
73
  graphql-relay==3.2.0
74
+ greenlet==3.3.0
74
75
  grpcio==1.76.0
75
76
  grpcio-status==1.76.0
76
77
  gunicorn==23.0.0
@@ -395,6 +395,18 @@
395
395
  "default": "3.12",
396
396
  "description": "Python version for Model Serving deployment. Defaults to 3.12 which is supported by Databricks Model Serving. This allows deploying from environments with different Python versions (e.g., Databricks Apps with 3.11).",
397
397
  "title": "Python Version"
398
+ },
399
+ "deployment_target": {
400
+ "anyOf": [
401
+ {
402
+ "$ref": "#/$defs/DeploymentTarget"
403
+ },
404
+ {
405
+ "type": "null"
406
+ }
407
+ ],
408
+ "default": null,
409
+ "description": "Default deployment target. If not specified, defaults to MODEL_SERVING. Can be overridden via CLI --target flag. Options: 'model_serving' or 'apps'."
398
410
  }
399
411
  },
400
412
  "required": [
@@ -1472,6 +1484,15 @@
1472
1484
  "title": "DatasetModel",
1473
1485
  "type": "object"
1474
1486
  },
1487
+ "DeploymentTarget": {
1488
+ "description": "Target platform for agent deployment.",
1489
+ "enum": [
1490
+ "model_serving",
1491
+ "apps"
1492
+ ],
1493
+ "title": "DeploymentTarget",
1494
+ "type": "string"
1495
+ },
1475
1496
  "Entitlement": {
1476
1497
  "enum": [
1477
1498
  "CAN_MANAGE",
@@ -0,0 +1,24 @@
1
+ """
2
+ Databricks Apps deployment module for dao-ai.
3
+
4
+ This subpackage contains all modules related to deploying dao-ai agents
5
+ as Databricks Apps or Model Serving endpoints.
6
+
7
+ Modules:
8
+ handlers: MLflow AgentServer request handlers (@invoke, @stream)
9
+ server: Entry point for Databricks Apps deployment
10
+ resources: Databricks App resource configuration generation
11
+ model_serving: Entry point for Databricks Model Serving deployment
12
+ """
13
+
14
+ from dao_ai.apps.resources import (
15
+ generate_app_resources,
16
+ generate_app_yaml,
17
+ generate_sdk_resources,
18
+ )
19
+
20
+ __all__ = [
21
+ "generate_app_resources",
22
+ "generate_app_yaml",
23
+ "generate_sdk_resources",
24
+ ]
@@ -1,20 +1,12 @@
1
1
  """
2
- App server module for running dao-ai agents as Databricks Apps.
2
+ Agent request handlers for MLflow AgentServer.
3
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.
4
+ This module defines the invoke and stream handlers that are registered
5
+ with the MLflow AgentServer. These handlers delegate to the ResponsesAgent
6
+ created from the dao-ai configuration.
7
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
8
+ The handlers use async methods (apredict, apredict_stream) to be compatible
9
+ with both Databricks Model Serving and Databricks Apps environments.
18
10
  """
19
11
 
20
12
  import os
@@ -22,8 +14,7 @@ from typing import AsyncGenerator
22
14
 
23
15
  import mlflow
24
16
  from dotenv import load_dotenv
25
- from mlflow.genai.agent_server import AgentServer, invoke, stream
26
- from mlflow.pyfunc import ResponsesAgent
17
+ from mlflow.genai.agent_server import invoke, stream
27
18
  from mlflow.types.responses import (
28
19
  ResponsesAgentRequest,
29
20
  ResponsesAgentResponse,
@@ -32,6 +23,7 @@ from mlflow.types.responses import (
32
23
 
33
24
  from dao_ai.config import AppConfig
34
25
  from dao_ai.logging import configure_logging
26
+ from dao_ai.models import LanggraphResponsesAgent
35
27
 
36
28
  # Load environment variables from .env.local if it exists
37
29
  load_dotenv(dotenv_path=".env.local", override=True)
@@ -51,53 +43,42 @@ config: AppConfig = AppConfig.from_file(config_path)
51
43
  if config.app and config.app.log_level:
52
44
  configure_logging(level=config.app.log_level)
53
45
 
54
- # Create the ResponsesAgent
55
- _responses_agent: ResponsesAgent = config.as_responses_agent()
46
+ # Create the ResponsesAgent - cast to LanggraphResponsesAgent to access async methods
47
+ _responses_agent: LanggraphResponsesAgent = config.as_responses_agent() # type: ignore[assignment]
56
48
 
57
49
 
58
50
  @invoke()
59
- def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
51
+ async def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
60
52
  """
61
53
  Handle non-streaming requests by delegating to the ResponsesAgent.
62
54
 
55
+ Uses the async apredict() method for compatibility with both
56
+ Model Serving and Apps environments.
57
+
63
58
  Args:
64
59
  request: The incoming ResponsesAgentRequest
65
60
 
66
61
  Returns:
67
62
  ResponsesAgentResponse with the complete output
68
63
  """
69
- return _responses_agent.predict(request)
64
+ return await _responses_agent.apredict(request)
70
65
 
71
66
 
72
67
  @stream()
73
- def streaming(
68
+ async def streaming(
74
69
  request: ResponsesAgentRequest,
75
70
  ) -> AsyncGenerator[ResponsesAgentStreamEvent, None]:
76
71
  """
77
72
  Handle streaming requests by delegating to the ResponsesAgent.
78
73
 
74
+ Uses the async apredict_stream() method for compatibility with both
75
+ Model Serving and Apps environments.
76
+
79
77
  Args:
80
78
  request: The incoming ResponsesAgentRequest
81
79
 
82
80
  Yields:
83
81
  ResponsesAgentStreamEvent objects as they are generated
84
82
  """
85
- # The predict_stream method returns a generator, convert to async generator
86
- for event in _responses_agent.predict_stream(request):
83
+ async for event in _responses_agent.apredict_stream(request):
87
84
  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()
@@ -0,0 +1,29 @@
1
+ # Apply nest_asyncio FIRST before any other imports
2
+ # This allows dao-ai's async/sync patterns to work in Model Serving
3
+ # where there may already be an event loop running (e.g., notebook context)
4
+ import nest_asyncio
5
+
6
+ nest_asyncio.apply()
7
+
8
+ import mlflow # noqa: E402
9
+ from mlflow.models import ModelConfig # noqa: E402
10
+ from mlflow.pyfunc import ResponsesAgent # noqa: E402
11
+
12
+ from dao_ai.config import AppConfig # noqa: E402
13
+ from dao_ai.logging import configure_logging # noqa: E402
14
+
15
+ mlflow.set_registry_uri("databricks-uc")
16
+ mlflow.set_tracking_uri("databricks")
17
+
18
+ mlflow.langchain.autolog()
19
+
20
+ model_config: ModelConfig = ModelConfig()
21
+ config: AppConfig = AppConfig(**model_config.to_dict())
22
+
23
+ log_level: str = config.app.log_level
24
+
25
+ configure_logging(level=log_level)
26
+
27
+ app: ResponsesAgent = config.as_responses_agent()
28
+
29
+ mlflow.models.set_model(app)