dao-ai 0.0.28__py3-none-any.whl → 0.1.2__py3-none-any.whl
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.
- dao_ai/__init__.py +29 -0
- dao_ai/agent_as_code.py +2 -5
- dao_ai/cli.py +245 -40
- dao_ai/config.py +1491 -370
- dao_ai/genie/__init__.py +38 -0
- dao_ai/genie/cache/__init__.py +43 -0
- dao_ai/genie/cache/base.py +72 -0
- dao_ai/genie/cache/core.py +79 -0
- dao_ai/genie/cache/lru.py +347 -0
- dao_ai/genie/cache/semantic.py +970 -0
- dao_ai/genie/core.py +35 -0
- dao_ai/graph.py +27 -253
- dao_ai/hooks/__init__.py +9 -6
- dao_ai/hooks/core.py +27 -195
- dao_ai/logging.py +56 -0
- dao_ai/memory/__init__.py +10 -0
- dao_ai/memory/core.py +65 -30
- dao_ai/memory/databricks.py +402 -0
- dao_ai/memory/postgres.py +79 -38
- dao_ai/messages.py +6 -4
- dao_ai/middleware/__init__.py +125 -0
- dao_ai/middleware/assertions.py +806 -0
- dao_ai/middleware/base.py +50 -0
- dao_ai/middleware/core.py +67 -0
- dao_ai/middleware/guardrails.py +420 -0
- dao_ai/middleware/human_in_the_loop.py +232 -0
- dao_ai/middleware/message_validation.py +586 -0
- dao_ai/middleware/summarization.py +197 -0
- dao_ai/models.py +1306 -114
- dao_ai/nodes.py +245 -159
- dao_ai/optimization.py +674 -0
- dao_ai/orchestration/__init__.py +52 -0
- dao_ai/orchestration/core.py +294 -0
- dao_ai/orchestration/supervisor.py +278 -0
- dao_ai/orchestration/swarm.py +271 -0
- dao_ai/prompts.py +128 -31
- dao_ai/providers/databricks.py +573 -601
- dao_ai/state.py +157 -21
- dao_ai/tools/__init__.py +13 -5
- dao_ai/tools/agent.py +1 -3
- dao_ai/tools/core.py +64 -11
- dao_ai/tools/email.py +232 -0
- dao_ai/tools/genie.py +144 -294
- dao_ai/tools/mcp.py +223 -155
- dao_ai/tools/memory.py +50 -0
- dao_ai/tools/python.py +9 -14
- dao_ai/tools/search.py +14 -0
- dao_ai/tools/slack.py +22 -10
- dao_ai/tools/sql.py +202 -0
- dao_ai/tools/time.py +30 -7
- dao_ai/tools/unity_catalog.py +165 -88
- dao_ai/tools/vector_search.py +331 -221
- dao_ai/utils.py +166 -20
- dao_ai-0.1.2.dist-info/METADATA +455 -0
- dao_ai-0.1.2.dist-info/RECORD +64 -0
- dao_ai/chat_models.py +0 -204
- dao_ai/guardrails.py +0 -112
- dao_ai/tools/human_in_the_loop.py +0 -100
- dao_ai-0.0.28.dist-info/METADATA +0 -1168
- dao_ai-0.0.28.dist-info/RECORD +0 -41
- {dao_ai-0.0.28.dist-info → dao_ai-0.1.2.dist-info}/WHEEL +0 -0
- {dao_ai-0.0.28.dist-info → dao_ai-0.1.2.dist-info}/entry_points.txt +0 -0
- {dao_ai-0.0.28.dist-info → dao_ai-0.1.2.dist-info}/licenses/LICENSE +0 -0
dao_ai-0.0.28.dist-info/METADATA
DELETED
|
@@ -1,1168 +0,0 @@
|
|
|
1
|
-
Metadata-Version: 2.4
|
|
2
|
-
Name: dao-ai
|
|
3
|
-
Version: 0.0.28
|
|
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
|
-
Project-URL: Homepage, https://github.com/natefleming/dao-ai
|
|
6
|
-
Project-URL: Documentation, https://natefleming.github.io/dao-ai
|
|
7
|
-
Project-URL: Repository, https://github.com/natefleming/dao-ai
|
|
8
|
-
Project-URL: Issues, https://github.com/natefleming/dao-ai/issues
|
|
9
|
-
Project-URL: Changelog, https://github.com/natefleming/dao-ai/blob/main/CHANGELOG.md
|
|
10
|
-
Author-email: Nate Fleming <nate.fleming@databricks.com>, Nate Fleming <nate.fleming@gmail.com>
|
|
11
|
-
Maintainer-email: Nate Fleming <nate.fleming@databricks.com>
|
|
12
|
-
License: MIT
|
|
13
|
-
License-File: LICENSE
|
|
14
|
-
Keywords: agents,ai,databricks,langchain,langgraph,llm,multi-agent,orchestration,vector-search,workflow
|
|
15
|
-
Classifier: Development Status :: 3 - Alpha
|
|
16
|
-
Classifier: Intended Audience :: Developers
|
|
17
|
-
Classifier: Intended Audience :: Science/Research
|
|
18
|
-
Classifier: License :: OSI Approved :: MIT License
|
|
19
|
-
Classifier: Operating System :: OS Independent
|
|
20
|
-
Classifier: Programming Language :: Python :: 3
|
|
21
|
-
Classifier: Programming Language :: Python :: 3.12
|
|
22
|
-
Classifier: Programming Language :: Python :: 3.13
|
|
23
|
-
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
24
|
-
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
25
|
-
Classifier: Topic :: System :: Distributed Computing
|
|
26
|
-
Requires-Python: >=3.12
|
|
27
|
-
Requires-Dist: databricks-agents>=1.7.0
|
|
28
|
-
Requires-Dist: databricks-langchain>=0.8.1
|
|
29
|
-
Requires-Dist: databricks-mcp>=0.3.0
|
|
30
|
-
Requires-Dist: databricks-sdk[openai]>=0.67.0
|
|
31
|
-
Requires-Dist: duckduckgo-search>=8.0.2
|
|
32
|
-
Requires-Dist: flashrank>=0.2.8
|
|
33
|
-
Requires-Dist: gepa>=0.0.17
|
|
34
|
-
Requires-Dist: grandalf>=0.8
|
|
35
|
-
Requires-Dist: langchain-mcp-adapters>=0.1.10
|
|
36
|
-
Requires-Dist: langchain-tavily>=0.2.11
|
|
37
|
-
Requires-Dist: langchain>=0.3.27
|
|
38
|
-
Requires-Dist: langgraph-checkpoint-postgres>=2.0.25
|
|
39
|
-
Requires-Dist: langgraph-supervisor>=0.0.29
|
|
40
|
-
Requires-Dist: langgraph-swarm>=0.0.14
|
|
41
|
-
Requires-Dist: langgraph>=0.6.10
|
|
42
|
-
Requires-Dist: langmem>=0.0.29
|
|
43
|
-
Requires-Dist: loguru>=0.7.3
|
|
44
|
-
Requires-Dist: mcp>=1.17.0
|
|
45
|
-
Requires-Dist: mlflow>=3.5.1
|
|
46
|
-
Requires-Dist: nest-asyncio>=1.6.0
|
|
47
|
-
Requires-Dist: openevals>=0.0.19
|
|
48
|
-
Requires-Dist: openpyxl>=3.1.5
|
|
49
|
-
Requires-Dist: psycopg[binary,pool]>=3.2.9
|
|
50
|
-
Requires-Dist: pydantic>=2.12.0
|
|
51
|
-
Requires-Dist: python-dotenv>=1.1.0
|
|
52
|
-
Requires-Dist: pyyaml>=6.0.2
|
|
53
|
-
Requires-Dist: rich>=14.0.0
|
|
54
|
-
Requires-Dist: scipy<=1.15
|
|
55
|
-
Requires-Dist: sqlparse>=0.5.3
|
|
56
|
-
Requires-Dist: tomli>=2.3.0
|
|
57
|
-
Requires-Dist: unitycatalog-ai[databricks]>=0.3.0
|
|
58
|
-
Provides-Extra: databricks
|
|
59
|
-
Requires-Dist: databricks-connect>=15.0.0; extra == 'databricks'
|
|
60
|
-
Requires-Dist: databricks-vectorsearch>=0.63; extra == 'databricks'
|
|
61
|
-
Requires-Dist: pyspark>=3.5.0; extra == 'databricks'
|
|
62
|
-
Provides-Extra: dev
|
|
63
|
-
Requires-Dist: mypy>=1.0.0; extra == 'dev'
|
|
64
|
-
Requires-Dist: pre-commit>=3.0.0; extra == 'dev'
|
|
65
|
-
Requires-Dist: pytest>=8.3.5; extra == 'dev'
|
|
66
|
-
Requires-Dist: ruff>=0.11.11; extra == 'dev'
|
|
67
|
-
Provides-Extra: docs
|
|
68
|
-
Requires-Dist: mkdocs-material>=9.0.0; extra == 'docs'
|
|
69
|
-
Requires-Dist: mkdocs>=1.5.0; extra == 'docs'
|
|
70
|
-
Requires-Dist: mkdocstrings[python]>=0.24.0; extra == 'docs'
|
|
71
|
-
Provides-Extra: test
|
|
72
|
-
Requires-Dist: pytest-cov>=4.0.0; extra == 'test'
|
|
73
|
-
Requires-Dist: pytest-mock>=3.10.0; extra == 'test'
|
|
74
|
-
Requires-Dist: pytest>=8.3.5; extra == 'test'
|
|
75
|
-
Description-Content-Type: text/markdown
|
|
76
|
-
|
|
77
|
-
# Declarative Agent Orchestration (DAO) Framework
|
|
78
|
-
|
|
79
|
-
A modular, multi-agent orchestration framework for building sophisticated AI workflows on Databricks. While this implementation provides a complete retail AI reference architecture, the framework is designed to support any domain or use case requiring agent coordination, tool integration, and dynamic configuration.
|
|
80
|
-
|
|
81
|
-
## Overview
|
|
82
|
-
|
|
83
|
-
This project implements a LangGraph-based multi-agent orchestration framework that can:
|
|
84
|
-
|
|
85
|
-
- **Route queries** to specialized agents based on content and context
|
|
86
|
-
- **Coordinate multiple AI agents** working together on complex tasks
|
|
87
|
-
- **Integrate diverse tools** including databases, APIs, vector search, and external services
|
|
88
|
-
- **Support flexible orchestration patterns** (supervisor, swarm, and custom workflows)
|
|
89
|
-
- **Provide dynamic configuration** through YAML-based agent and tool definitions
|
|
90
|
-
- **Enable domain-specific specialization** while maintaining a unified interface
|
|
91
|
-
|
|
92
|
-
**Retail Reference Implementation**: This repository includes a complete retail AI system demonstrating:
|
|
93
|
-
- Product inventory management and search
|
|
94
|
-
- Customer recommendation engines
|
|
95
|
-
- Order tracking and management
|
|
96
|
-
- Product classification and information retrieval
|
|
97
|
-
|
|
98
|
-
The system uses Databricks Vector Search, Unity Catalog, and LLMs to provide accurate, context-aware responses across any domain.
|
|
99
|
-
|
|
100
|
-
## Key Features
|
|
101
|
-
|
|
102
|
-
- **Multi-Modal Interface**: CLI commands and Python API for development and deployment
|
|
103
|
-
- **Agent Lifecycle Management**: Create, deploy, and monitor agents programmatically
|
|
104
|
-
- **Vector Search Integration**: Built-in support for Databricks Vector Search with retrieval tools
|
|
105
|
-
- **Configuration-Driven**: YAML-based configuration with validation and IDE support
|
|
106
|
-
- **MLflow Integration**: Automatic model packaging, versioning, and deployment
|
|
107
|
-
- **Monitoring & Evaluation**: Built-in assessment and monitoring capabilities
|
|
108
|
-
|
|
109
|
-
## Architecture
|
|
110
|
-
|
|
111
|
-
### Overview
|
|
112
|
-
|
|
113
|
-
The Multi-Agent AI system is built as a component-based agent architecture that routes queries to specialized agents based on the nature of the request. This approach enables domain-specific handling while maintaining a unified interface that can be adapted to any industry or use case.
|
|
114
|
-
|
|
115
|
-

|
|
116
|
-
|
|
117
|
-
### Core Components
|
|
118
|
-
|
|
119
|
-
#### Configuration Components
|
|
120
|
-
|
|
121
|
-
All components are defined from the provided [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml) using a modular approach:
|
|
122
|
-
|
|
123
|
-
- **Schemas**: Define database and catalog structures
|
|
124
|
-
- **Resources**: Configure infrastructure components like LLMs, vector stores, catalogs, warehouses, and databases
|
|
125
|
-
- **Tools**: Define functions that agents can use to perform tasks (dictionary-based with keys as tool names)
|
|
126
|
-
- **Agents**: Specialized AI assistants configured for specific domains (dictionary-based with keys as agent names)
|
|
127
|
-
- **Guardrails**: Quality control mechanisms to ensure accurate responses
|
|
128
|
-
- **Retrievers**: Configuration for vector search and retrieval
|
|
129
|
-
- **Evaluation**: Configuration for model evaluation and testing
|
|
130
|
-
- **Datasets**: Configuration for training and evaluation datasets
|
|
131
|
-
- **App**: Overall application configuration including orchestration and logging
|
|
132
|
-
|
|
133
|
-
#### Message Processing Flow
|
|
134
|
-
|
|
135
|
-
The system uses a LangGraph-based workflow with the following key nodes:
|
|
136
|
-
|
|
137
|
-
- **Message Validation**: Validates incoming requests (`message_validation_node`)
|
|
138
|
-
- **Agent Routing**: Routes messages to appropriate specialized agents using supervisor or swarm patterns
|
|
139
|
-
- **Agent Execution**: Processes requests using specialized agents with their configured tools
|
|
140
|
-
- **Response Generation**: Returns structured responses to users
|
|
141
|
-
|
|
142
|
-
#### Specialized Agents
|
|
143
|
-
|
|
144
|
-
Agents are dynamically configured from the provided [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml) file and can include:
|
|
145
|
-
- Custom LLM models and parameters
|
|
146
|
-
- Specific sets of available tools (Python functions, Unity Catalog functions, factory tools, MCP services)
|
|
147
|
-
- Domain-specific system prompts
|
|
148
|
-
- Guardrails for response quality
|
|
149
|
-
- Handoff prompts for agent coordination
|
|
150
|
-
|
|
151
|
-
### Technical Implementation
|
|
152
|
-
|
|
153
|
-
The system is implemented using:
|
|
154
|
-
|
|
155
|
-
- **LangGraph**: For workflow orchestration and state management
|
|
156
|
-
- **LangChain**: For LLM interactions and tool integration
|
|
157
|
-
- **MLflow**: For model tracking and deployment
|
|
158
|
-
- **Databricks**: LLM APIs, Vector Search, Unity Catalog, and Model Serving
|
|
159
|
-
- **Pydantic**: For configuration validation and schema management
|
|
160
|
-
|
|
161
|
-
## Prerequisites
|
|
162
|
-
|
|
163
|
-
- Python 3.12+
|
|
164
|
-
- Databricks workspace with access to:
|
|
165
|
-
- Unity Catalog
|
|
166
|
-
- Model Serving
|
|
167
|
-
- Vector Search
|
|
168
|
-
- Genie (optional)
|
|
169
|
-
- Databricks CLI configured with appropriate permissions
|
|
170
|
-
- Databricks model endpoints for LLMs and embeddings
|
|
171
|
-
|
|
172
|
-
## Setup
|
|
173
|
-
|
|
174
|
-
1. Clone this repository
|
|
175
|
-
2. Install dependencies:
|
|
176
|
-
|
|
177
|
-
```bash
|
|
178
|
-
# Create and activate a Python virtual environment
|
|
179
|
-
uv venv
|
|
180
|
-
source .venv/bin/activate # On Windows: .venv\Scripts\activate
|
|
181
|
-
|
|
182
|
-
# Install dependencies using Makefile
|
|
183
|
-
make install
|
|
184
|
-
```
|
|
185
|
-
|
|
186
|
-
3. Configure Databricks CLI with appropriate workspace access
|
|
187
|
-
|
|
188
|
-
## Quick Start
|
|
189
|
-
|
|
190
|
-
### Option 1: Using Python API (Recommended for Development)
|
|
191
|
-
|
|
192
|
-
```python
|
|
193
|
-
from retail_ai.config import AppConfig
|
|
194
|
-
|
|
195
|
-
# Load your configuration
|
|
196
|
-
config = AppConfig.from_file("config/hardware_store/supervisor_postgres.yaml")
|
|
197
|
-
|
|
198
|
-
# Create vector search infrastructure
|
|
199
|
-
for name, vector_store in config.resources.vector_stores.items():
|
|
200
|
-
vector_store.create()
|
|
201
|
-
|
|
202
|
-
# Create and deploy your agent
|
|
203
|
-
config.create_agent()
|
|
204
|
-
config.deploy_agent()
|
|
205
|
-
|
|
206
|
-
```
|
|
207
|
-
|
|
208
|
-
### Option 2: Using CLI Commands
|
|
209
|
-
|
|
210
|
-
```bash
|
|
211
|
-
# Validate configuration
|
|
212
|
-
dao-ai validate -c config/hardware_store/supervisor_postgres.yaml
|
|
213
|
-
|
|
214
|
-
# Generate workflow diagram
|
|
215
|
-
dao-ai graph -o architecture.png
|
|
216
|
-
|
|
217
|
-
# Deploy using Databricks Asset Bundles
|
|
218
|
-
dao-ai bundle --deploy --run
|
|
219
|
-
|
|
220
|
-
# Deploy using Databricks Asset Bundles with specific configuration
|
|
221
|
-
dao-ai -vvvv bundle --deploy --run --target dev --config config/hardware_store/supervisor_postgres.yaml --profile DEFAULT
|
|
222
|
-
```
|
|
223
|
-
|
|
224
|
-
See the [Python API](#python-api) section for detailed programmatic usage, or [Command Line Interface](#command-line-interface) for CLI usage.
|
|
225
|
-
|
|
226
|
-
## Command Line Interface
|
|
227
|
-
|
|
228
|
-
The framework includes a comprehensive CLI for managing, validating, and visualizing your multi-agent system:
|
|
229
|
-
|
|
230
|
-
### Schema Generation
|
|
231
|
-
Generate JSON schema for configuration validation and IDE autocompletion:
|
|
232
|
-
```bash
|
|
233
|
-
dao-ai schema > schema.json
|
|
234
|
-
```
|
|
235
|
-
|
|
236
|
-
### Configuration Validation
|
|
237
|
-
Validate your configuration file for syntax and semantic correctness:
|
|
238
|
-
```bash
|
|
239
|
-
# Validate default configuration (config/hardware_store/supervisor_postgres.yaml)
|
|
240
|
-
dao-ai validate
|
|
241
|
-
|
|
242
|
-
# Validate specific configuration file
|
|
243
|
-
dao-ai validate -c config/production.yaml
|
|
244
|
-
```
|
|
245
|
-
|
|
246
|
-
### Graph Visualization
|
|
247
|
-
Generate visual representations of your agent workflow:
|
|
248
|
-
```bash
|
|
249
|
-
# Generate architecture diagram (using default config/hardware_store/supervisor_postgres.yaml)
|
|
250
|
-
dao-ai graph -o architecture.png
|
|
251
|
-
|
|
252
|
-
# Generate diagram from specific config
|
|
253
|
-
dao-ai graph -o workflow.png -c config/custom.yaml
|
|
254
|
-
```
|
|
255
|
-
|
|
256
|
-
### Deployment
|
|
257
|
-
Deploy your multi-agent system using Databricks Asset Bundles:
|
|
258
|
-
```bash
|
|
259
|
-
# Deploy the system
|
|
260
|
-
dao-ai bundle --deploy
|
|
261
|
-
|
|
262
|
-
# Run the deployed system
|
|
263
|
-
dao-ai bundle --run
|
|
264
|
-
|
|
265
|
-
# Use specific Databricks profile
|
|
266
|
-
dao-ai bundle --deploy --run --profile my-profile
|
|
267
|
-
```
|
|
268
|
-
|
|
269
|
-
### Verbose Output
|
|
270
|
-
Add `-v`, `-vv`, `-vvv`, or `-vvvv` flags for increasing levels of verbosity (ERROR, WARNING, INFO, DEBUG, TRACE).
|
|
271
|
-
|
|
272
|
-
## Python API
|
|
273
|
-
|
|
274
|
-
The framework provides a comprehensive Python API for programmatic access to all functionality. The main entry point is the `AppConfig` class, which provides methods for agent lifecycle management, vector search operations, and configuration utilities.
|
|
275
|
-
|
|
276
|
-
### Quick Start
|
|
277
|
-
|
|
278
|
-
```python
|
|
279
|
-
from retail_ai.config import AppConfig
|
|
280
|
-
|
|
281
|
-
# Load configuration from file
|
|
282
|
-
config = AppConfig.from_file(path="config/hardware_store/supervisor_postgres.yaml")
|
|
283
|
-
```
|
|
284
|
-
|
|
285
|
-
### Agent Lifecycle Management
|
|
286
|
-
|
|
287
|
-
#### Creating Agents
|
|
288
|
-
Package and register your multi-agent system as an MLflow model:
|
|
289
|
-
|
|
290
|
-
```python
|
|
291
|
-
# Create agent with default settings
|
|
292
|
-
config.create_agent()
|
|
293
|
-
|
|
294
|
-
# Create agent with additional requirements and code paths
|
|
295
|
-
config.create_agent(
|
|
296
|
-
additional_pip_reqs=["custom-package==1.0.0"],
|
|
297
|
-
additional_code_paths=["./custom_modules"]
|
|
298
|
-
)
|
|
299
|
-
```
|
|
300
|
-
|
|
301
|
-
#### Deploying Agents
|
|
302
|
-
Deploy your registered agent to a Databricks serving endpoint:
|
|
303
|
-
|
|
304
|
-
```python
|
|
305
|
-
# Deploy agent to serving endpoint
|
|
306
|
-
config.deploy_agent()
|
|
307
|
-
```
|
|
308
|
-
|
|
309
|
-
The deployment process:
|
|
310
|
-
1. Retrieves the latest model version from MLflow
|
|
311
|
-
2. Creates or updates a Databricks model serving endpoint
|
|
312
|
-
3. Configures scaling, environment variables, and permissions
|
|
313
|
-
4. Sets up proper authentication and resource access
|
|
314
|
-
|
|
315
|
-
### Vector Search Operations
|
|
316
|
-
|
|
317
|
-
#### Creating Vector Search Infrastructure
|
|
318
|
-
Create vector search endpoints and indexes from your configuration:
|
|
319
|
-
|
|
320
|
-
```python
|
|
321
|
-
# Access vector stores from configuration
|
|
322
|
-
vector_stores = config.resources.vector_stores
|
|
323
|
-
|
|
324
|
-
# Create all vector stores
|
|
325
|
-
for name, vector_store in vector_stores.items():
|
|
326
|
-
print(f"Creating vector store: {name}")
|
|
327
|
-
vector_store.create()
|
|
328
|
-
```
|
|
329
|
-
|
|
330
|
-
#### Using Vector Search
|
|
331
|
-
Query your vector search indexes for retrieval-augmented generation:
|
|
332
|
-
|
|
333
|
-
```python
|
|
334
|
-
# Method 1: Direct index access
|
|
335
|
-
from retail_ai.config import RetrieverModel
|
|
336
|
-
|
|
337
|
-
question = "What products do you have in stock?"
|
|
338
|
-
|
|
339
|
-
for name, retriever in config.retrievers.items():
|
|
340
|
-
# Get the vector search index
|
|
341
|
-
index = retriever.vector_store.as_index()
|
|
342
|
-
|
|
343
|
-
# Perform similarity search
|
|
344
|
-
results = index.similarity_search(
|
|
345
|
-
query_text=question,
|
|
346
|
-
columns=retriever.columns,
|
|
347
|
-
**retriever.search_parameters.model_dump()
|
|
348
|
-
)
|
|
349
|
-
|
|
350
|
-
chunks = results.get('result', {}).get('data_array', [])
|
|
351
|
-
print(f"Found {len(chunks)} relevant results")
|
|
352
|
-
```
|
|
353
|
-
|
|
354
|
-
```python
|
|
355
|
-
# Method 2: LangChain integration
|
|
356
|
-
from databricks_langchain import DatabricksVectorSearch
|
|
357
|
-
|
|
358
|
-
for name, retriever in config.retrievers.items():
|
|
359
|
-
# Create LangChain vector store
|
|
360
|
-
vector_search = DatabricksVectorSearch(
|
|
361
|
-
endpoint=retriever.vector_store.endpoint.name,
|
|
362
|
-
index_name=retriever.vector_store.index.full_name,
|
|
363
|
-
columns=retriever.columns,
|
|
364
|
-
)
|
|
365
|
-
|
|
366
|
-
# Search using LangChain interface
|
|
367
|
-
documents = vector_search.similarity_search(
|
|
368
|
-
query=question,
|
|
369
|
-
**retriever.search_parameters.model_dump()
|
|
370
|
-
)
|
|
371
|
-
|
|
372
|
-
print(f"Found {len(documents)} documents")
|
|
373
|
-
```
|
|
374
|
-
|
|
375
|
-
### Configuration Utilities
|
|
376
|
-
|
|
377
|
-
The `AppConfig` class provides helper methods to find and filter configuration components:
|
|
378
|
-
|
|
379
|
-
#### Finding Agents
|
|
380
|
-
```python
|
|
381
|
-
# Get all agents
|
|
382
|
-
all_agents = config.find_agents()
|
|
383
|
-
|
|
384
|
-
# Find agents with specific criteria
|
|
385
|
-
def has_vector_search(agent):
|
|
386
|
-
return any("vector_search" in tool.name.lower() for tool in agent.tools)
|
|
387
|
-
|
|
388
|
-
vector_agents = config.find_agents(predicate=has_vector_search)
|
|
389
|
-
```
|
|
390
|
-
|
|
391
|
-
#### Finding Tools and Guardrails
|
|
392
|
-
```python
|
|
393
|
-
# Get all tools
|
|
394
|
-
all_tools = config.find_tools()
|
|
395
|
-
|
|
396
|
-
# Get all guardrails
|
|
397
|
-
all_guardrails = config.find_guardrails()
|
|
398
|
-
|
|
399
|
-
# Find tools by type
|
|
400
|
-
def is_python_tool(tool):
|
|
401
|
-
return tool.function.type == "python"
|
|
402
|
-
|
|
403
|
-
python_tools = config.find_tools(predicate=is_python_tool)
|
|
404
|
-
```
|
|
405
|
-
|
|
406
|
-
### Visualization
|
|
407
|
-
|
|
408
|
-
Generate and save workflow diagrams:
|
|
409
|
-
|
|
410
|
-
```python
|
|
411
|
-
# Display graph in notebook
|
|
412
|
-
config.display_graph()
|
|
413
|
-
|
|
414
|
-
# Save architecture diagram
|
|
415
|
-
config.save_image("docs/my_architecture.png")
|
|
416
|
-
```
|
|
417
|
-
|
|
418
|
-
### Complete Example
|
|
419
|
-
|
|
420
|
-
See [`notebooks/05_agent_as_code_driver.py`](notebooks/05_agent_as_code_driver.py) for a complete example:
|
|
421
|
-
|
|
422
|
-
```python
|
|
423
|
-
from retail_ai.config import AppConfig
|
|
424
|
-
from pathlib import Path
|
|
425
|
-
|
|
426
|
-
# Load configuration
|
|
427
|
-
config = AppConfig.from_file("config/hardware_store/supervisor_postgres.yaml")
|
|
428
|
-
|
|
429
|
-
# Visualize the workflow
|
|
430
|
-
config.display_graph()
|
|
431
|
-
|
|
432
|
-
# Save architecture diagram
|
|
433
|
-
path = Path("docs") / f"{config.app.name}_architecture.png"
|
|
434
|
-
config.save_image(path)
|
|
435
|
-
|
|
436
|
-
# Create and deploy the agent
|
|
437
|
-
config.create_agent()
|
|
438
|
-
config.deploy_agent()
|
|
439
|
-
```
|
|
440
|
-
|
|
441
|
-
For vector search examples, see [`notebooks/02_provision_vector_search.py`](notebooks/02_provision_vector_search.py).
|
|
442
|
-
|
|
443
|
-
### Available Notebooks
|
|
444
|
-
|
|
445
|
-
The framework includes several example notebooks demonstrating different aspects:
|
|
446
|
-
|
|
447
|
-
| Notebook | Description | Key Methods Demonstrated |
|
|
448
|
-
|----------|-------------|-------------------------|
|
|
449
|
-
| [`01_ingest_and_transform.py`](notebooks/01_ingest_and_transform.py) | Data ingestion and transformation | Dataset creation and SQL execution |
|
|
450
|
-
| [`02_provision_vector_search.py`](notebooks/02_provision_vector_search.py) | Vector search setup and usage | `vector_store.create()`, `as_index()` |
|
|
451
|
-
| [`03_generate_evaluation_data.py`](notebooks/03_generate_evaluation_data.py) | Generate synthetic evaluation datasets | Data generation and evaluation setup |
|
|
452
|
-
| [`04_unity_catalog_tools.py`](notebooks/04_unity_catalog_tools.py) | Unity Catalog function deployment | SQL function creation and testing |
|
|
453
|
-
| [`05_agent_as_code_driver.py`](notebooks/05_agent_as_code_driver.py) | **Complete agent lifecycle** | `create_agent()`, `deploy_agent()` |
|
|
454
|
-
| [`06_run_evaluation.py`](notebooks/06_run_evaluation.py) | Agent evaluation and testing | Evaluation framework usage |
|
|
455
|
-
| [`08_run_examples.py`](notebooks/08_run_examples.py) | End-to-end example queries | Agent interaction and testing |
|
|
456
|
-
|
|
457
|
-
## Configuration
|
|
458
|
-
|
|
459
|
-
Configuration is managed through [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml). This file defines all components of the Retail AI system, including resources, tools, agents, and the overall application setup.
|
|
460
|
-
|
|
461
|
-
**Note**: The configuration file location is configurable throughout the framework. You can specify a different configuration file using the `-c` or `--config` flag in CLI commands, or by setting the appropriate parameters in the Python API.
|
|
462
|
-
|
|
463
|
-
### Basic Structure of [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml)
|
|
464
|
-
|
|
465
|
-
The [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml) is organized into several top-level keys:
|
|
466
|
-
|
|
467
|
-
```yaml
|
|
468
|
-
# filepath: /Users/nate/development/dao-ai/config/hardware_store/supervisor_postgres.yaml
|
|
469
|
-
schemas:
|
|
470
|
-
# ... schema definitions ...
|
|
471
|
-
|
|
472
|
-
resources:
|
|
473
|
-
# ... resource definitions (LLMs, vector stores, etc.) ...
|
|
474
|
-
|
|
475
|
-
tools:
|
|
476
|
-
# ... tool definitions ...
|
|
477
|
-
|
|
478
|
-
agents:
|
|
479
|
-
# ... agent definitions ...
|
|
480
|
-
|
|
481
|
-
app:
|
|
482
|
-
# ... application configuration ...
|
|
483
|
-
|
|
484
|
-
# Other sections like guardrails, retrievers, evaluation, datasets
|
|
485
|
-
```
|
|
486
|
-
|
|
487
|
-
### Loading and Using Configuration
|
|
488
|
-
|
|
489
|
-
The configuration can be loaded and used programmatically through the `AppConfig` class:
|
|
490
|
-
|
|
491
|
-
```python
|
|
492
|
-
from retail_ai.config import AppConfig
|
|
493
|
-
|
|
494
|
-
# Load configuration from file
|
|
495
|
-
config = AppConfig.from_file("config/hardware_store/supervisor_postgres.yaml")
|
|
496
|
-
|
|
497
|
-
# Access different configuration sections
|
|
498
|
-
print(f"Available agents: {list(config.agents.keys())}")
|
|
499
|
-
print(f"Available tools: {list(config.tools.keys())}")
|
|
500
|
-
print(f"Vector stores: {list(config.resources.vector_stores.keys())}")
|
|
501
|
-
|
|
502
|
-
# Use configuration methods for deployment
|
|
503
|
-
config.create_agent() # Package as MLflow model
|
|
504
|
-
config.deploy_agent() # Deploy to serving endpoint
|
|
505
|
-
```
|
|
506
|
-
|
|
507
|
-
The configuration supports both CLI and programmatic workflows, with the Python API providing more flexibility for complex deployment scenarios.
|
|
508
|
-
|
|
509
|
-
### Developing and Configuring Tools
|
|
510
|
-
|
|
511
|
-
Tools are functions that agents can use to interact with external systems or perform specific tasks. They are defined under the `tools` key in [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml). Each tool has a unique name and contains a `function` specification.
|
|
512
|
-
|
|
513
|
-
There are four types of tools supported:
|
|
514
|
-
|
|
515
|
-
#### 1. Python Tools (`type: python`)
|
|
516
|
-
These tools directly map to Python functions. The `name` field should correspond to a function that can be imported and called directly.
|
|
517
|
-
|
|
518
|
-
**Configuration Example:**
|
|
519
|
-
```yaml
|
|
520
|
-
tools:
|
|
521
|
-
my_python_tool:
|
|
522
|
-
name: my_python_tool
|
|
523
|
-
function:
|
|
524
|
-
type: python
|
|
525
|
-
name: retail_ai.tools.my_function_name
|
|
526
|
-
schema: *retail_schema # Optional schema definition
|
|
527
|
-
```
|
|
528
|
-
**Development:**
|
|
529
|
-
Implement the Python function in the specified module (e.g., `retail_ai/tools.py`). The function will be imported and called directly when the tool is invoked.
|
|
530
|
-
|
|
531
|
-
#### 2. Factory Tools (`type: factory`)
|
|
532
|
-
Factory tools use factory functions that return initialized LangChain `BaseTool` instances. This is useful for tools requiring complex initialization or configuration.
|
|
533
|
-
|
|
534
|
-
**Configuration Example:**
|
|
535
|
-
```yaml
|
|
536
|
-
tools:
|
|
537
|
-
vector_search_tool:
|
|
538
|
-
name: vector_search
|
|
539
|
-
function:
|
|
540
|
-
type: factory
|
|
541
|
-
name: retail_ai.tools.create_vector_search_tool
|
|
542
|
-
args:
|
|
543
|
-
retriever: *products_retriever
|
|
544
|
-
name: product_vector_search_tool
|
|
545
|
-
description: "Search for products using vector search"
|
|
546
|
-
```
|
|
547
|
-
**Development:**
|
|
548
|
-
Implement the factory function (e.g., `create_vector_search_tool`) in `retail_ai/tools.py`. This function should accept the specified `args` and return a fully configured `BaseTool` object.
|
|
549
|
-
|
|
550
|
-
#### 3. Unity Catalog Tools (`type: unity_catalog`)
|
|
551
|
-
These tools represent SQL functions registered in Databricks Unity Catalog. They reference functions by their Unity Catalog schema and name.
|
|
552
|
-
|
|
553
|
-
**Configuration Example:**
|
|
554
|
-
```yaml
|
|
555
|
-
tools:
|
|
556
|
-
find_product_by_sku_uc_tool:
|
|
557
|
-
name: find_product_by_sku_uc
|
|
558
|
-
function:
|
|
559
|
-
type: unity_catalog
|
|
560
|
-
name: find_product_by_sku
|
|
561
|
-
schema: *retail_schema
|
|
562
|
-
```
|
|
563
|
-
**Development:**
|
|
564
|
-
Create the corresponding SQL function in your Databricks Unity Catalog using the specified schema and function name. The tool will automatically generate the appropriate function signature and documentation.
|
|
565
|
-
|
|
566
|
-
### Developing Unity Catalog Functions
|
|
567
|
-
|
|
568
|
-
Unity Catalog functions provide the backbone for data access in the multi-agent system. The framework automatically deploys these functions from SQL DDL files during system initialization.
|
|
569
|
-
|
|
570
|
-
#### Function Deployment Configuration
|
|
571
|
-
|
|
572
|
-
Unity Catalog functions are defined in the `unity_catalog_functions` section of [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml). Each function specification includes:
|
|
573
|
-
|
|
574
|
-
- **Function metadata**: Schema and name for Unity Catalog registration
|
|
575
|
-
- **DDL file path**: Location of the SQL file containing the function definition
|
|
576
|
-
- **Test parameters**: Optional test data for function validation
|
|
577
|
-
|
|
578
|
-
**Configuration Example from [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml):**
|
|
579
|
-
```yaml
|
|
580
|
-
unity_catalog_functions:
|
|
581
|
-
- function:
|
|
582
|
-
schema: *retail_schema # Reference to schema configuration
|
|
583
|
-
name: find_product_by_sku # Function name in Unity Catalog
|
|
584
|
-
ddl: ../functions/retail/find_product_by_sku.sql # Path to SQL DDL file
|
|
585
|
-
test: # Optional test configuration
|
|
586
|
-
parameters:
|
|
587
|
-
sku: ["00176279"] # Test parameters for validation
|
|
588
|
-
- function:
|
|
589
|
-
schema: *retail_schema
|
|
590
|
-
name: find_store_inventory_by_sku
|
|
591
|
-
ddl: ../functions/retail/find_store_inventory_by_sku.sql
|
|
592
|
-
test:
|
|
593
|
-
parameters:
|
|
594
|
-
store: "35048" # Multiple parameters for complex functions
|
|
595
|
-
sku: ["00176279"]
|
|
596
|
-
```
|
|
597
|
-
|
|
598
|
-
#### SQL Function Structure
|
|
599
|
-
|
|
600
|
-
SQL files should follow this structure for proper deployment:
|
|
601
|
-
|
|
602
|
-
**File Structure Example** (`functions/retail/find_product_by_sku.sql`):
|
|
603
|
-
```sql
|
|
604
|
-
-- Function to find product details by SKU
|
|
605
|
-
CREATE OR REPLACE FUNCTION {catalog_name}.{schema_name}.find_product_by_sku(
|
|
606
|
-
sku ARRAY<STRING> COMMENT 'One or more unique identifiers for retrieve. SKU values are between 5-8 alpha numeric characters'
|
|
607
|
-
)
|
|
608
|
-
RETURNS TABLE(
|
|
609
|
-
product_id BIGINT COMMENT 'Unique identifier for each product in the catalog',
|
|
610
|
-
sku STRING COMMENT 'Stock Keeping Unit - unique internal product identifier code',
|
|
611
|
-
upc STRING COMMENT 'Universal Product Code - standardized barcode number for product identification',
|
|
612
|
-
brand_name STRING COMMENT 'Name of the manufacturer or brand that produces the product',
|
|
613
|
-
product_name STRING COMMENT 'Display name of the product as shown to customers',
|
|
614
|
-
-- ... additional columns
|
|
615
|
-
)
|
|
616
|
-
READS SQL DATA
|
|
617
|
-
COMMENT 'Retrieves detailed information about a specific product by its SKU. This function is designed for product information retrieval in retail applications.'
|
|
618
|
-
RETURN
|
|
619
|
-
SELECT
|
|
620
|
-
product_id,
|
|
621
|
-
sku,
|
|
622
|
-
upc,
|
|
623
|
-
brand_name,
|
|
624
|
-
product_name
|
|
625
|
-
-- ... additional columns
|
|
626
|
-
FROM products
|
|
627
|
-
WHERE ARRAY_CONTAINS(find_product_by_sku.sku, products.sku);
|
|
628
|
-
```
|
|
629
|
-
|
|
630
|
-
**Key Requirements:**
|
|
631
|
-
- Use `{catalog_name}.{schema_name}` placeholders - these are automatically replaced during deployment
|
|
632
|
-
- Include comprehensive `COMMENT` attributes for all parameters and return columns
|
|
633
|
-
- Provide a clear function-level comment describing purpose and use cases
|
|
634
|
-
- Use `READS SQL DATA` for functions that query data
|
|
635
|
-
- Follow consistent naming conventions for parameters and return values
|
|
636
|
-
|
|
637
|
-
#### Test Configuration
|
|
638
|
-
|
|
639
|
-
The optional `test` section allows you to define test parameters for automatic function validation:
|
|
640
|
-
|
|
641
|
-
```yaml
|
|
642
|
-
test:
|
|
643
|
-
parameters:
|
|
644
|
-
sku: ["00176279"] # Single parameter
|
|
645
|
-
# OR for multi-parameter functions:
|
|
646
|
-
store: "35048" # Multiple parameters
|
|
647
|
-
sku: ["00176279"]
|
|
648
|
-
```
|
|
649
|
-
|
|
650
|
-
**Test Benefits:**
|
|
651
|
-
- **Validation**: Ensures functions work correctly after deployment
|
|
652
|
-
- **Documentation**: Provides example usage for other developers
|
|
653
|
-
- **CI/CD Integration**: Enables automated testing in deployment pipelines
|
|
654
|
-
|
|
655
|
-
**Note**: Test parameters should use realistic data from your datasets to ensure meaningful validation. The framework will execute these tests automatically during deployment to verify function correctness.
|
|
656
|
-
|
|
657
|
-
#### 4. MCP (Model Context Protocol) Tools (`type: mcp`)
|
|
658
|
-
MCP tools allow interaction with external services that implement the Model Context Protocol, supporting both HTTP and stdio transports.
|
|
659
|
-
|
|
660
|
-
**Configuration Example (Direct URL):**
|
|
661
|
-
```yaml
|
|
662
|
-
tools:
|
|
663
|
-
weather_tool_mcp:
|
|
664
|
-
name: weather
|
|
665
|
-
function:
|
|
666
|
-
type: mcp
|
|
667
|
-
name: weather
|
|
668
|
-
transport: streamable_http
|
|
669
|
-
url: http://localhost:8000/mcp
|
|
670
|
-
```
|
|
671
|
-
|
|
672
|
-
**Configuration Example (Unity Catalog Connection):**
|
|
673
|
-
MCP tools can also use Unity Catalog Connections for secure, governed access with on-behalf-of-user capabilities. The connection provides OAuth authentication, while the URL specifies the endpoint:
|
|
674
|
-
```yaml
|
|
675
|
-
resources:
|
|
676
|
-
connections:
|
|
677
|
-
github_connection:
|
|
678
|
-
name: github_u2m_connection # UC Connection name
|
|
679
|
-
|
|
680
|
-
tools:
|
|
681
|
-
github_mcp:
|
|
682
|
-
name: github_mcp
|
|
683
|
-
function:
|
|
684
|
-
type: mcp
|
|
685
|
-
name: github_mcp
|
|
686
|
-
transport: streamable_http
|
|
687
|
-
url: https://workspace.databricks.com/api/2.0/mcp/external/github_u2m_connection # MCP endpoint URL
|
|
688
|
-
connection: *github_connection # UC Connection provides OAuth authentication
|
|
689
|
-
```
|
|
690
|
-
|
|
691
|
-
**Development:**
|
|
692
|
-
- **For direct URL connections**: Ensure the MCP service is running and accessible at the specified URL or command. Provide OAuth credentials (client_id, client_secret) or PAT for authentication.
|
|
693
|
-
- **For UC Connection**: URL is required to specify the endpoint. The connection provides OAuth authentication via the workspace client. Ensure the connection is configured in Unity Catalog with appropriate MCP scopes (`mcp.genie`, `mcp.functions`, `mcp.vectorsearch`, `mcp.external`).
|
|
694
|
-
- The framework will handle the MCP protocol communication automatically, including session management and authentication.
|
|
695
|
-
|
|
696
|
-
### Configuring New Agents
|
|
697
|
-
|
|
698
|
-
Agents are specialized AI assistants defined under the `agents` key in [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml). Each agent has a unique name and specific configuration.
|
|
699
|
-
|
|
700
|
-
**Configuration Example:**
|
|
701
|
-
```yaml
|
|
702
|
-
agents:
|
|
703
|
-
general:
|
|
704
|
-
name: general
|
|
705
|
-
description: "General retail store assistant for home improvement and hardware store inquiries"
|
|
706
|
-
model: *tool_calling_llm
|
|
707
|
-
tools:
|
|
708
|
-
- *find_product_details_by_description_tool
|
|
709
|
-
- *vector_search_tool
|
|
710
|
-
guardrails: []
|
|
711
|
-
checkpointer: *checkpointer
|
|
712
|
-
prompt: |
|
|
713
|
-
You are a helpful retail store assistant for a home improvement and hardware store.
|
|
714
|
-
You have access to search tools to find current information about products, pricing, and store policies.
|
|
715
|
-
|
|
716
|
-
#### CRITICAL INSTRUCTION: ALWAYS USE SEARCH TOOLS FIRST
|
|
717
|
-
Before answering ANY question:
|
|
718
|
-
- ALWAYS use your available search tools to find the most current and accurate information
|
|
719
|
-
- Search for specific details about store policies, product availability, pricing, and services
|
|
720
|
-
```
|
|
721
|
-
|
|
722
|
-
**Agent Configuration Fields:**
|
|
723
|
-
- `name`: Unique identifier for the agent
|
|
724
|
-
- `description`: Human-readable description of the agent's purpose
|
|
725
|
-
- `model`: Reference to an LLM model (using YAML anchors like `*tool_calling_llm`)
|
|
726
|
-
- `tools`: Array of tool references (using YAML anchors like `*search_tool`)
|
|
727
|
-
- `guardrails`: Array of guardrail references (can be empty `[]`)
|
|
728
|
-
- `checkpointer`: Reference to a checkpointer for conversation state (optional)
|
|
729
|
-
- `prompt`: System prompt that defines the agent's behavior and instructions
|
|
730
|
-
|
|
731
|
-
**To configure a new agent:**
|
|
732
|
-
1. Add a new entry under the `agents` section with a unique key
|
|
733
|
-
2. Define the required fields: `name`, `description`, `model`, `tools`, and `prompt`
|
|
734
|
-
3. Optionally configure `guardrails` and `checkpointer`
|
|
735
|
-
4. Reference the agent in the application configuration using YAML anchors
|
|
736
|
-
|
|
737
|
-
### Assigning Tools to Agents
|
|
738
|
-
|
|
739
|
-
Tools are assigned to agents by referencing them using YAML anchors in the agent's `tools` array. Each tool must be defined in the `tools` section with an anchor (using `&tool_name`), then referenced in the agent configuration (using `*tool_name`).
|
|
740
|
-
|
|
741
|
-
**Example:**
|
|
742
|
-
```yaml
|
|
743
|
-
tools:
|
|
744
|
-
search_tool: &search_tool
|
|
745
|
-
name: search
|
|
746
|
-
function:
|
|
747
|
-
type: factory
|
|
748
|
-
name: retail_ai.tools.search_tool
|
|
749
|
-
args: {}
|
|
750
|
-
|
|
751
|
-
genie_tool: &genie_tool
|
|
752
|
-
name: genie
|
|
753
|
-
function:
|
|
754
|
-
type: factory
|
|
755
|
-
name: retail_ai.tools.create_genie_tool
|
|
756
|
-
args:
|
|
757
|
-
genie_room: *retail_genie_room
|
|
758
|
-
|
|
759
|
-
agents:
|
|
760
|
-
general:
|
|
761
|
-
name: general
|
|
762
|
-
description: "General retail store assistant"
|
|
763
|
-
model: *tool_calling_llm
|
|
764
|
-
tools:
|
|
765
|
-
- *search_tool # Reference to the search_tool anchor
|
|
766
|
-
- *genie_tool # Reference to the genie_tool anchor
|
|
767
|
-
# ... other agent configuration
|
|
768
|
-
```
|
|
769
|
-
|
|
770
|
-
This YAML anchor system allows for:
|
|
771
|
-
- **Reusability**: The same tool can be assigned to multiple agents
|
|
772
|
-
- **Maintainability**: Tool configuration is centralized in one place
|
|
773
|
-
- **Consistency**: Tools are guaranteed to have the same configuration across agents
|
|
774
|
-
|
|
775
|
-
### Assigning Agents to the Application and Configuring Orchestration
|
|
776
|
-
|
|
777
|
-
Agents are made available to the application by listing their YAML anchors (defined in the `agents:` section) within the `agents` array under the `app` section. The `app.orchestration` section defines how these agents interact.
|
|
778
|
-
|
|
779
|
-
**Orchestration Configuration:**
|
|
780
|
-
|
|
781
|
-
The `orchestration` block within the `app` section allows you to define the interaction pattern. Your current configuration primarily uses a **Supervisor** pattern.
|
|
782
|
-
|
|
783
|
-
```yaml
|
|
784
|
-
# filepath: /Users/nate/development/dao-ai/config/hardware_store/supervisor_postgres.yaml
|
|
785
|
-
# ...
|
|
786
|
-
# app:
|
|
787
|
-
# ...
|
|
788
|
-
# agents:
|
|
789
|
-
# - *orders
|
|
790
|
-
# - *diy
|
|
791
|
-
# - *product
|
|
792
|
-
# # ... other agents referenced by their anchors
|
|
793
|
-
# - *general
|
|
794
|
-
# orchestration:
|
|
795
|
-
# supervisor:
|
|
796
|
-
# model: *tool_calling_llm # LLM for the supervisor agent
|
|
797
|
-
# default_agent: *general # Agent to handle tasks if no specific agent is chosen
|
|
798
|
-
# # swarm: # Example of how a swarm might be configured if activated
|
|
799
|
-
# # model: *tool_calling_llm
|
|
800
|
-
# ...
|
|
801
|
-
```
|
|
802
|
-
|
|
803
|
-
**Orchestration Patterns:**
|
|
804
|
-
|
|
805
|
-
1. **Supervisor Pattern (Currently Active)**
|
|
806
|
-
* Your configuration defines a `supervisor` block under `app.orchestration`.
|
|
807
|
-
* `model`: Specifies the LLM (e.g., `*tool_calling_llm`) that the supervisor itself will use for its decision-making and routing logic.
|
|
808
|
-
* `default_agent`: Specifies an agent (e.g., `*general`) that the supervisor will delegate to if it cannot determine a more specialized agent from the `app.agents` list or if the query is general.
|
|
809
|
-
* The supervisor is responsible for receiving the initial user query, deciding which specialized agent (from the `app.agents` list) is best suited to handle it, and then passing the query to that agent. If no specific agent is a clear match, or if the query is general, it falls back to the `default_agent`.
|
|
810
|
-
|
|
811
|
-
2. **Swarm Pattern (Commented Out)**
|
|
812
|
-
* Your configuration includes a commented-out `swarm` block. If activated, this would imply a different interaction model.
|
|
813
|
-
* In a swarm, agents might collaborate more directly or work in parallel on different aspects of a query. The `model` under `swarm` would likely define the LLM used by the agents within the swarm or by a coordinating element of the swarm.
|
|
814
|
-
* The specific implementation of how a swarm pattern behaves would be defined in your `retail_ai/graph.py` and `retail_ai/nodes.py`.
|
|
815
|
-
|
|
816
|
-
## Integration Hooks
|
|
817
|
-
|
|
818
|
-
The DAO framework provides several hook integration points that allow you to customize agent behavior and application lifecycle. These hooks enable you to inject custom logic at key points in the system without modifying the core framework code.
|
|
819
|
-
|
|
820
|
-
### Hook Types
|
|
821
|
-
|
|
822
|
-
#### Agent-Level Hooks
|
|
823
|
-
|
|
824
|
-
**Agent hooks** are defined at the individual agent level and allow you to customize specific agent behavior:
|
|
825
|
-
|
|
826
|
-
##### `create_agent_hook`
|
|
827
|
-
Used to provide a completely custom agent implementation. When this is provided all other configuration is ignored. See: **Hook Implementation**
|
|
828
|
-
|
|
829
|
-
```yaml
|
|
830
|
-
agents:
|
|
831
|
-
custom_agent:
|
|
832
|
-
name: custom_agent
|
|
833
|
-
description: "Agent with custom initialization"
|
|
834
|
-
model: *tool_calling_llm
|
|
835
|
-
create_agent_hook: my_package.hooks.initialize_custom_agent
|
|
836
|
-
# ... other agent configuration
|
|
837
|
-
```
|
|
838
|
-
|
|
839
|
-
##### `pre_agent_hook`
|
|
840
|
-
Executed before an agent processes a message. Ideal for request preprocessing, logging, validation, or context injection. See: **Hook Implementation**
|
|
841
|
-
|
|
842
|
-
```yaml
|
|
843
|
-
agents:
|
|
844
|
-
logging_agent:
|
|
845
|
-
name: logging_agent
|
|
846
|
-
description: "Agent with request logging"
|
|
847
|
-
model: *tool_calling_llm
|
|
848
|
-
pre_agent_hook: my_package.hooks.log_incoming_request
|
|
849
|
-
# ... other agent configuration
|
|
850
|
-
```
|
|
851
|
-
|
|
852
|
-
##### `post_agent_hook`
|
|
853
|
-
Executed after an agent completes processing a message. Perfect for response post-processing, logging, metrics collection, or cleanup operations. See: **Hook Implementation**
|
|
854
|
-
|
|
855
|
-
```yaml
|
|
856
|
-
agents:
|
|
857
|
-
analytics_agent:
|
|
858
|
-
name: analytics_agent
|
|
859
|
-
description: "Agent with response analytics"
|
|
860
|
-
model: *tool_calling_llm
|
|
861
|
-
post_agent_hook: my_package.hooks.collect_response_metrics
|
|
862
|
-
# ... other agent configuration
|
|
863
|
-
```
|
|
864
|
-
|
|
865
|
-
#### Application-Level Hooks
|
|
866
|
-
|
|
867
|
-
**Application hooks** operate at the global application level and affect the entire system lifecycle:
|
|
868
|
-
|
|
869
|
-
##### `initialization_hooks`
|
|
870
|
-
Executed when the application starts up via `AppConfig.from_file()`. Use these for system initialization, resource setup, database connections, or external service configuration. See: **Hook Implementation**
|
|
871
|
-
|
|
872
|
-
```yaml
|
|
873
|
-
app:
|
|
874
|
-
name: my_retail_app
|
|
875
|
-
initialization_hooks:
|
|
876
|
-
- my_package.hooks.setup_database_connections
|
|
877
|
-
- my_package.hooks.initialize_external_apis
|
|
878
|
-
- my_package.hooks.setup_monitoring
|
|
879
|
-
# ... other app configuration
|
|
880
|
-
```
|
|
881
|
-
|
|
882
|
-
##### `shutdown_hooks`
|
|
883
|
-
Executed when the application shuts down (registered via `atexit`). Essential for cleanup operations, closing connections, saving state, or performing final logging. See: **Hook Implementation**
|
|
884
|
-
|
|
885
|
-
```yaml
|
|
886
|
-
app:
|
|
887
|
-
name: my_retail_app
|
|
888
|
-
shutdown_hooks:
|
|
889
|
-
- my_package.hooks.cleanup_database_connections
|
|
890
|
-
- my_package.hooks.save_session_data
|
|
891
|
-
- my_package.hooks.send_shutdown_metrics
|
|
892
|
-
# ... other app configuration
|
|
893
|
-
```
|
|
894
|
-
|
|
895
|
-
##### `message_hooks`
|
|
896
|
-
Executed for every message processed by the system. Useful for global logging, authentication, rate limiting, or message transformation. See: **Hook Implementation**
|
|
897
|
-
|
|
898
|
-
```yaml
|
|
899
|
-
app:
|
|
900
|
-
name: my_retail_app
|
|
901
|
-
message_hooks:
|
|
902
|
-
- my_package.hooks.authenticate_user
|
|
903
|
-
- my_package.hooks.apply_rate_limiting
|
|
904
|
-
- my_package.hooks.transform_message_format
|
|
905
|
-
# ... other app configuration
|
|
906
|
-
```
|
|
907
|
-
|
|
908
|
-
### Hook Implementation
|
|
909
|
-
|
|
910
|
-
Hooks can be implemented as either:
|
|
911
|
-
|
|
912
|
-
1. **Python Functions**: Direct function references
|
|
913
|
-
```yaml
|
|
914
|
-
initialization_hooks: my_package.hooks.setup_function
|
|
915
|
-
```
|
|
916
|
-
|
|
917
|
-
2. **Factory Functions**: Functions that return configured tools or handlers
|
|
918
|
-
```yaml
|
|
919
|
-
initialization_hooks:
|
|
920
|
-
type: factory
|
|
921
|
-
name: my_package.hooks.create_setup_handler
|
|
922
|
-
args:
|
|
923
|
-
config_param: "value"
|
|
924
|
-
```
|
|
925
|
-
|
|
926
|
-
3. **Hook Lists**: Multiple hooks executed in sequence
|
|
927
|
-
```yaml
|
|
928
|
-
initialization_hooks:
|
|
929
|
-
- my_package.hooks.setup_database
|
|
930
|
-
- my_package.hooks.setup_cache
|
|
931
|
-
- my_package.hooks.setup_monitoring
|
|
932
|
-
```
|
|
933
|
-
|
|
934
|
-
### Hook Function Signatures
|
|
935
|
-
|
|
936
|
-
Each hook type expects specific function signatures:
|
|
937
|
-
|
|
938
|
-
#### Agent Hooks
|
|
939
|
-
```python
|
|
940
|
-
# create_agent_hook
|
|
941
|
-
def initialize_custom_agent(state: dict, config: dict) -> dict:
|
|
942
|
-
"""Custom agent initialization logic"""
|
|
943
|
-
pass
|
|
944
|
-
|
|
945
|
-
# pre_agent_hook
|
|
946
|
-
def log_incoming_request(state: dict, config: dict) -> dict:
|
|
947
|
-
"""Pre-process incoming request"""
|
|
948
|
-
return state
|
|
949
|
-
|
|
950
|
-
# post_agent_hook
|
|
951
|
-
def collect_response_metrics(state: dict, config: dict) -> dict:
|
|
952
|
-
"""Post-process agent response"""
|
|
953
|
-
return state
|
|
954
|
-
```
|
|
955
|
-
|
|
956
|
-
#### Application Hooks
|
|
957
|
-
```python
|
|
958
|
-
# initialization_hooks
|
|
959
|
-
def setup_database_connections(config: AppConfig) -> None:
|
|
960
|
-
"""Initialize database connections"""
|
|
961
|
-
pass
|
|
962
|
-
|
|
963
|
-
# shutdown_hooks
|
|
964
|
-
def cleanup_resources(config: AppConfig) -> None:
|
|
965
|
-
"""Clean up resources on shutdown"""
|
|
966
|
-
pass
|
|
967
|
-
|
|
968
|
-
# message_hooks
|
|
969
|
-
def authenticate_user(state: dict, config: dict) -> dict:
|
|
970
|
-
"""Authenticate and authorize user requests"""
|
|
971
|
-
return state
|
|
972
|
-
```
|
|
973
|
-
|
|
974
|
-
### Use Cases and Examples
|
|
975
|
-
|
|
976
|
-
#### Common Hook Patterns
|
|
977
|
-
|
|
978
|
-
**Logging and Monitoring**:
|
|
979
|
-
```python
|
|
980
|
-
def log_agent_performance(state: dict, config: AppConfig) -> dict:
|
|
981
|
-
"""Log agent response times and quality metrics"""
|
|
982
|
-
start_time = state.get('start_time')
|
|
983
|
-
if start_time:
|
|
984
|
-
duration = time.time() - start_time
|
|
985
|
-
logger.info(f"Agent response time: {duration:.2f}s")
|
|
986
|
-
return state
|
|
987
|
-
```
|
|
988
|
-
|
|
989
|
-
**Authentication and Authorization**:
|
|
990
|
-
```python
|
|
991
|
-
def validate_user_permissions(state: dict, config: AppConfig) -> dict:
|
|
992
|
-
"""Validate user has permission for requested operation"""
|
|
993
|
-
user_id = state.get('user_id')
|
|
994
|
-
if not has_permission(user_id, state.get('operation')):
|
|
995
|
-
raise UnauthorizedError("Insufficient permissions")
|
|
996
|
-
return state
|
|
997
|
-
```
|
|
998
|
-
|
|
999
|
-
**Resource Management**:
|
|
1000
|
-
```python
|
|
1001
|
-
def initialize_vector_search(config: AppConfig) -> None:
|
|
1002
|
-
"""Initialize vector search connections during startup"""
|
|
1003
|
-
for vs_name, vs_config in config.resources.vector_stores.items():
|
|
1004
|
-
vs_config.create()
|
|
1005
|
-
logger.info(f"Vector store {vs_name} initialized")
|
|
1006
|
-
```
|
|
1007
|
-
|
|
1008
|
-
**State Enrichment**:
|
|
1009
|
-
```python
|
|
1010
|
-
def enrich_user_context(state: dict, config: AppConfig) -> dict:
|
|
1011
|
-
"""Add user profile and preferences to state"""
|
|
1012
|
-
user_id = state.get('user_id')
|
|
1013
|
-
if user_id:
|
|
1014
|
-
user_profile = get_user_profile(user_id)
|
|
1015
|
-
state['user_context'] = user_profile
|
|
1016
|
-
return state
|
|
1017
|
-
```
|
|
1018
|
-
|
|
1019
|
-
### Best Practices
|
|
1020
|
-
|
|
1021
|
-
1. **Keep hooks lightweight**: Avoid heavy computations that could slow down message processing
|
|
1022
|
-
2. **Handle errors gracefully**: Use try-catch blocks to prevent hook failures from breaking the system
|
|
1023
|
-
3. **Use appropriate hook types**: Choose agent-level vs application-level hooks based on scope
|
|
1024
|
-
4. **Maintain state immutability**: Return modified copies of state rather than mutating in-place
|
|
1025
|
-
5. **Log hook execution**: Include logging for troubleshooting and monitoring
|
|
1026
|
-
6. **Test hooks independently**: Write unit tests for hook functions separate from the main application
|
|
1027
|
-
|
|
1028
|
-
|
|
1029
|
-
## Development
|
|
1030
|
-
|
|
1031
|
-
### Project Structure
|
|
1032
|
-
|
|
1033
|
-
- `retail_ai/`: Core package
|
|
1034
|
-
- `config.py`: Pydantic configuration models with full validation
|
|
1035
|
-
- `graph.py`: LangGraph workflow definition
|
|
1036
|
-
- `nodes.py`: Agent node factories and implementations
|
|
1037
|
-
- `tools.py`: Tool creation and factory functions, implementations for Python tools
|
|
1038
|
-
- `vector_search.py`: Vector search utilities
|
|
1039
|
-
- `state.py`: State management for conversations
|
|
1040
|
-
- `tests/`: Test suite with configuration fixtures
|
|
1041
|
-
- `schemas/`: JSON schemas for configuration validation
|
|
1042
|
-
- `notebooks/`: Jupyter notebooks for setup and experimentation
|
|
1043
|
-
- `docs/`: Documentation files, including architecture diagrams.
|
|
1044
|
-
- `config/`: Contains [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml).
|
|
1045
|
-
|
|
1046
|
-
### Building the Package
|
|
1047
|
-
|
|
1048
|
-
```bash
|
|
1049
|
-
# Install development dependencies
|
|
1050
|
-
make depends
|
|
1051
|
-
|
|
1052
|
-
# Build the package
|
|
1053
|
-
make install
|
|
1054
|
-
|
|
1055
|
-
# Run tests
|
|
1056
|
-
make test
|
|
1057
|
-
|
|
1058
|
-
# Format code
|
|
1059
|
-
make format
|
|
1060
|
-
```
|
|
1061
|
-
|
|
1062
|
-
## Deployment with Databricks Bundle CLI
|
|
1063
|
-
|
|
1064
|
-
The agent can be deployed using the existing Databricks Bundle CLI configuration:
|
|
1065
|
-
|
|
1066
|
-
1. Ensure Databricks CLI is installed and configured:
|
|
1067
|
-
```bash
|
|
1068
|
-
pip install databricks-cli
|
|
1069
|
-
databricks configure
|
|
1070
|
-
```
|
|
1071
|
-
|
|
1072
|
-
2. Deploy using the existing `databricks.yml`:
|
|
1073
|
-
```bash
|
|
1074
|
-
databricks bundle deploy
|
|
1075
|
-
```
|
|
1076
|
-
|
|
1077
|
-
3. Check deployment status:
|
|
1078
|
-
```bash
|
|
1079
|
-
databricks bundle status
|
|
1080
|
-
```
|
|
1081
|
-
|
|
1082
|
-
## Usage
|
|
1083
|
-
|
|
1084
|
-
Once deployed, interact with the agent:
|
|
1085
|
-
|
|
1086
|
-
```python
|
|
1087
|
-
from mlflow.deployments import get_deploy_client
|
|
1088
|
-
|
|
1089
|
-
client = get_deploy_client("databricks")
|
|
1090
|
-
response = client.predict(
|
|
1091
|
-
endpoint="retail_ai_agent", # Matches endpoint_name in model_config.yaml
|
|
1092
|
-
inputs={
|
|
1093
|
-
"messages": [
|
|
1094
|
-
{"role": "user", "content": "Can you recommend a lamp for my oak side tables?"}
|
|
1095
|
-
]
|
|
1096
|
-
}
|
|
1097
|
-
)
|
|
1098
|
-
|
|
1099
|
-
print(response["message"]["content"])
|
|
1100
|
-
```
|
|
1101
|
-
|
|
1102
|
-
### Advanced Configuration
|
|
1103
|
-
|
|
1104
|
-
You can also pass additional configuration parameters to customize the agent's behavior:
|
|
1105
|
-
|
|
1106
|
-
```python
|
|
1107
|
-
response = client.predict(
|
|
1108
|
-
endpoint="retail_ai_agent",
|
|
1109
|
-
inputs={
|
|
1110
|
-
"messages": [
|
|
1111
|
-
{"role": "user", "content": "Can you recommend a lamp for my oak side tables?"}
|
|
1112
|
-
],
|
|
1113
|
-
"configurable": {
|
|
1114
|
-
"thread_id": "1",
|
|
1115
|
-
"user_id": "my_user_id",
|
|
1116
|
-
"store_num": 87887
|
|
1117
|
-
}
|
|
1118
|
-
}
|
|
1119
|
-
)
|
|
1120
|
-
```
|
|
1121
|
-
|
|
1122
|
-
The `configurable` section supports:
|
|
1123
|
-
- **`thread_id`**: Unique identifier for conversation threading and state management
|
|
1124
|
-
- **`user_id`**: User identifier for personalization and tracking
|
|
1125
|
-
- **`store_num`**: Store number for location-specific recommendations and inventory
|
|
1126
|
-
|
|
1127
|
-
## Customization
|
|
1128
|
-
|
|
1129
|
-
To customize the agent:
|
|
1130
|
-
|
|
1131
|
-
1. **Update [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml)**:
|
|
1132
|
-
- Add tools in the `tools` section
|
|
1133
|
-
- Create agents in the `agents` section
|
|
1134
|
-
- Configure resources (LLMs, vector stores, etc.)
|
|
1135
|
-
- Adjust orchestration patterns as described above.
|
|
1136
|
-
|
|
1137
|
-
2. **Implement new tools** in `retail_ai/tools.py` (for Python and Factory tools) or in Unity Catalog (for UC tools).
|
|
1138
|
-
|
|
1139
|
-
3. **Extend workflows** in `retail_ai/graph.py` to support the chosen orchestration patterns and agent interactions.
|
|
1140
|
-
|
|
1141
|
-
## Testing
|
|
1142
|
-
|
|
1143
|
-
```bash
|
|
1144
|
-
# Run all tests
|
|
1145
|
-
make test
|
|
1146
|
-
```
|
|
1147
|
-
|
|
1148
|
-
## Logging
|
|
1149
|
-
|
|
1150
|
-
The primary log level for the application is configured in [`model_config.yaml`](config/hardware_store/supervisor_postgres.yaml) under the `app.log_level` field.
|
|
1151
|
-
|
|
1152
|
-
**Configuration Example:**
|
|
1153
|
-
```yaml
|
|
1154
|
-
# filepath: /Users/nate/development/dao-ai/config/hardware_store/supervisor_postgres.yaml
|
|
1155
|
-
app:
|
|
1156
|
-
log_level: INFO # Supported levels: DEBUG, INFO, WARNING, ERROR, CRITICAL
|
|
1157
|
-
# ... other app configurations ...
|
|
1158
|
-
```
|
|
1159
|
-
|
|
1160
|
-
This setting controls the verbosity of logs produced by the `retail_ai` package.
|
|
1161
|
-
|
|
1162
|
-
The system also includes:
|
|
1163
|
-
- **MLflow tracing** for request tracking.
|
|
1164
|
-
- **Structured logging** is used internally.
|
|
1165
|
-
|
|
1166
|
-
## License
|
|
1167
|
-
|
|
1168
|
-
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
|