dao-ai 0.1.8__py3-none-any.whl → 0.1.9__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/app_server.py +103 -0
- dao_ai/cli.py +15 -1
- dao_ai/config.py +89 -36
- dao_ai/providers/base.py +28 -2
- dao_ai/providers/databricks.py +205 -3
- dao_ai/state.py +1 -0
- {dao_ai-0.1.8.dist-info → dao_ai-0.1.9.dist-info}/METADATA +2 -2
- {dao_ai-0.1.8.dist-info → dao_ai-0.1.9.dist-info}/RECORD +11 -10
- {dao_ai-0.1.8.dist-info → dao_ai-0.1.9.dist-info}/WHEEL +0 -0
- {dao_ai-0.1.8.dist-info → dao_ai-0.1.9.dist-info}/entry_points.txt +0 -0
- {dao_ai-0.1.8.dist-info → dao_ai-0.1.9.dist-info}/licenses/LICENSE +0 -0
dao_ai/app_server.py
ADDED
|
@@ -0,0 +1,103 @@
|
|
|
1
|
+
"""
|
|
2
|
+
App server module for running dao-ai agents as Databricks Apps.
|
|
3
|
+
|
|
4
|
+
This module provides the entry point for deploying dao-ai agents as Databricks Apps
|
|
5
|
+
using MLflow's AgentServer. It follows the same pattern as agent_as_code.py but
|
|
6
|
+
uses the AgentServer for the Databricks Apps runtime.
|
|
7
|
+
|
|
8
|
+
Configuration Loading:
|
|
9
|
+
The config path is specified via the DAO_AI_CONFIG_PATH environment variable,
|
|
10
|
+
or defaults to model_config.yaml in the current directory.
|
|
11
|
+
|
|
12
|
+
Usage:
|
|
13
|
+
# With environment variable
|
|
14
|
+
DAO_AI_CONFIG_PATH=/path/to/config.yaml python -m dao_ai.app_server
|
|
15
|
+
|
|
16
|
+
# With default model_config.yaml in current directory
|
|
17
|
+
python -m dao_ai.app_server
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
import os
|
|
21
|
+
from typing import AsyncGenerator
|
|
22
|
+
|
|
23
|
+
import mlflow
|
|
24
|
+
from dotenv import load_dotenv
|
|
25
|
+
from mlflow.genai.agent_server import AgentServer, invoke, stream
|
|
26
|
+
from mlflow.pyfunc import ResponsesAgent
|
|
27
|
+
from mlflow.types.responses import (
|
|
28
|
+
ResponsesAgentRequest,
|
|
29
|
+
ResponsesAgentResponse,
|
|
30
|
+
ResponsesAgentStreamEvent,
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
from dao_ai.config import AppConfig
|
|
34
|
+
from dao_ai.logging import configure_logging
|
|
35
|
+
|
|
36
|
+
# Load environment variables from .env.local if it exists
|
|
37
|
+
load_dotenv(dotenv_path=".env.local", override=True)
|
|
38
|
+
|
|
39
|
+
# Configure MLflow
|
|
40
|
+
mlflow.set_registry_uri("databricks-uc")
|
|
41
|
+
mlflow.set_tracking_uri("databricks")
|
|
42
|
+
mlflow.langchain.autolog()
|
|
43
|
+
|
|
44
|
+
# Get config path from environment or use default
|
|
45
|
+
config_path: str = os.environ.get("DAO_AI_CONFIG_PATH", "model_config.yaml")
|
|
46
|
+
|
|
47
|
+
# Load configuration using AppConfig.from_file (consistent with CLI, notebook, builder)
|
|
48
|
+
config: AppConfig = AppConfig.from_file(config_path)
|
|
49
|
+
|
|
50
|
+
# Configure logging
|
|
51
|
+
if config.app and config.app.log_level:
|
|
52
|
+
configure_logging(level=config.app.log_level)
|
|
53
|
+
|
|
54
|
+
# Create the ResponsesAgent
|
|
55
|
+
_responses_agent: ResponsesAgent = config.as_responses_agent()
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
@invoke()
|
|
59
|
+
def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
|
|
60
|
+
"""
|
|
61
|
+
Handle non-streaming requests by delegating to the ResponsesAgent.
|
|
62
|
+
|
|
63
|
+
Args:
|
|
64
|
+
request: The incoming ResponsesAgentRequest
|
|
65
|
+
|
|
66
|
+
Returns:
|
|
67
|
+
ResponsesAgentResponse with the complete output
|
|
68
|
+
"""
|
|
69
|
+
return _responses_agent.predict(request)
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
@stream()
|
|
73
|
+
def streaming(
|
|
74
|
+
request: ResponsesAgentRequest,
|
|
75
|
+
) -> AsyncGenerator[ResponsesAgentStreamEvent, None]:
|
|
76
|
+
"""
|
|
77
|
+
Handle streaming requests by delegating to the ResponsesAgent.
|
|
78
|
+
|
|
79
|
+
Args:
|
|
80
|
+
request: The incoming ResponsesAgentRequest
|
|
81
|
+
|
|
82
|
+
Yields:
|
|
83
|
+
ResponsesAgentStreamEvent objects as they are generated
|
|
84
|
+
"""
|
|
85
|
+
# The predict_stream method returns a generator, convert to async generator
|
|
86
|
+
for event in _responses_agent.predict_stream(request):
|
|
87
|
+
yield event
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
# Create the AgentServer instance
|
|
91
|
+
agent_server = AgentServer("ResponsesAgent", enable_chat_proxy=True)
|
|
92
|
+
|
|
93
|
+
# Define the app as a module level variable to enable multiple workers
|
|
94
|
+
app = agent_server.app
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def main() -> None:
|
|
98
|
+
"""Entry point for running the agent server."""
|
|
99
|
+
agent_server.run(app_import_string="dao_ai.app_server:app")
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
if __name__ == "__main__":
|
|
103
|
+
main()
|
dao_ai/cli.py
CHANGED
|
@@ -309,6 +309,14 @@ Examples:
|
|
|
309
309
|
metavar="FILE",
|
|
310
310
|
help="Path to the model configuration file to validate",
|
|
311
311
|
)
|
|
312
|
+
deploy_parser.add_argument(
|
|
313
|
+
"-t",
|
|
314
|
+
"--target",
|
|
315
|
+
type=str,
|
|
316
|
+
choices=["model_serving", "apps"],
|
|
317
|
+
default="model_serving",
|
|
318
|
+
help="Deployment target: 'model_serving' (default) or 'apps'",
|
|
319
|
+
)
|
|
312
320
|
|
|
313
321
|
# List MCP tools command
|
|
314
322
|
list_mcp_parser: ArgumentParser = subparsers.add_parser(
|
|
@@ -729,11 +737,17 @@ def handle_graph_command(options: Namespace) -> None:
|
|
|
729
737
|
|
|
730
738
|
|
|
731
739
|
def handle_deploy_command(options: Namespace) -> None:
|
|
740
|
+
from dao_ai.config import DeploymentTarget
|
|
741
|
+
|
|
732
742
|
logger.debug(f"Validating configuration from {options.config}...")
|
|
733
743
|
try:
|
|
734
744
|
config: AppConfig = AppConfig.from_file(options.config)
|
|
745
|
+
|
|
746
|
+
# Convert target string to enum
|
|
747
|
+
target: DeploymentTarget = DeploymentTarget(options.target)
|
|
748
|
+
|
|
735
749
|
config.create_agent()
|
|
736
|
-
config.deploy_agent()
|
|
750
|
+
config.deploy_agent(target=target)
|
|
737
751
|
sys.exit(0)
|
|
738
752
|
except Exception as e:
|
|
739
753
|
logger.error(f"Deployment failed: {e}")
|
dao_ai/config.py
CHANGED
|
@@ -208,7 +208,9 @@ class IsDatabricksResource(ABC, BaseModel):
|
|
|
208
208
|
Authentication Options:
|
|
209
209
|
----------------------
|
|
210
210
|
1. **On-Behalf-Of User (OBO)**: Set on_behalf_of_user=True to use the
|
|
211
|
-
calling user's identity
|
|
211
|
+
calling user's identity. Implementation varies by deployment:
|
|
212
|
+
- Databricks Apps: Uses X-Forwarded-Access-Token from request headers
|
|
213
|
+
- Model Serving: Uses ModelServingUserCredentials
|
|
212
214
|
|
|
213
215
|
2. **Service Principal (OAuth M2M)**: Provide service_principal or
|
|
214
216
|
(client_id + client_secret + workspace_host) for service principal auth.
|
|
@@ -221,9 +223,17 @@ class IsDatabricksResource(ABC, BaseModel):
|
|
|
221
223
|
|
|
222
224
|
Authentication Priority:
|
|
223
225
|
1. OBO (on_behalf_of_user=True)
|
|
226
|
+
- Checks for forwarded headers (Databricks Apps)
|
|
227
|
+
- Falls back to ModelServingUserCredentials (Model Serving)
|
|
224
228
|
2. Service Principal (client_id + client_secret + workspace_host)
|
|
225
229
|
3. PAT (pat + workspace_host)
|
|
226
230
|
4. Ambient/default authentication
|
|
231
|
+
|
|
232
|
+
Note: When on_behalf_of_user=True, the agent acts as the calling user regardless
|
|
233
|
+
of deployment target. In Databricks Apps, this uses X-Forwarded-Access-Token
|
|
234
|
+
automatically captured by MLflow AgentServer. In Model Serving, this uses
|
|
235
|
+
ModelServingUserCredentials. Forwarded headers are ONLY used when
|
|
236
|
+
on_behalf_of_user=True.
|
|
227
237
|
"""
|
|
228
238
|
|
|
229
239
|
model_config = ConfigDict(use_enum_values=True)
|
|
@@ -235,9 +245,6 @@ class IsDatabricksResource(ABC, BaseModel):
|
|
|
235
245
|
workspace_host: Optional[AnyVariable] = None
|
|
236
246
|
pat: Optional[AnyVariable] = None
|
|
237
247
|
|
|
238
|
-
# Private attribute to cache the workspace client (lazy instantiation)
|
|
239
|
-
_workspace_client: Optional[WorkspaceClient] = PrivateAttr(default=None)
|
|
240
|
-
|
|
241
248
|
@abstractmethod
|
|
242
249
|
def as_resources(self) -> Sequence[DatabricksResource]: ...
|
|
243
250
|
|
|
@@ -273,32 +280,56 @@ class IsDatabricksResource(ABC, BaseModel):
|
|
|
273
280
|
"""
|
|
274
281
|
Get a WorkspaceClient configured with the appropriate authentication.
|
|
275
282
|
|
|
276
|
-
|
|
283
|
+
A new client is created on each access.
|
|
277
284
|
|
|
278
285
|
Authentication priority:
|
|
279
|
-
1.
|
|
280
|
-
|
|
281
|
-
|
|
282
|
-
|
|
283
|
-
|
|
286
|
+
1. On-Behalf-Of User (on_behalf_of_user=True):
|
|
287
|
+
- Forwarded headers (Databricks Apps)
|
|
288
|
+
- ModelServingUserCredentials (Model Serving)
|
|
289
|
+
2. Service Principal (client_id + client_secret + workspace_host)
|
|
290
|
+
3. PAT (pat + workspace_host)
|
|
291
|
+
4. Ambient/default authentication
|
|
284
292
|
"""
|
|
285
|
-
# Return cached client if already instantiated
|
|
286
|
-
if self._workspace_client is not None:
|
|
287
|
-
return self._workspace_client
|
|
288
|
-
|
|
289
293
|
from dao_ai.utils import normalize_host
|
|
290
294
|
|
|
291
295
|
# Check for OBO first (highest priority)
|
|
292
296
|
if self.on_behalf_of_user:
|
|
297
|
+
# NEW: In Databricks Apps, use forwarded headers for per-user auth
|
|
298
|
+
try:
|
|
299
|
+
from mlflow.genai.agent_server import get_request_headers
|
|
300
|
+
|
|
301
|
+
headers = get_request_headers()
|
|
302
|
+
forwarded_token = headers.get("x-forwarded-access-token")
|
|
303
|
+
|
|
304
|
+
if forwarded_token:
|
|
305
|
+
forwarded_user = headers.get("x-forwarded-user", "unknown")
|
|
306
|
+
logger.debug(
|
|
307
|
+
f"Creating WorkspaceClient for {self.__class__.__name__} "
|
|
308
|
+
f"with OBO using forwarded token from Databricks Apps",
|
|
309
|
+
forwarded_user=forwarded_user,
|
|
310
|
+
)
|
|
311
|
+
# Use workspace_host if configured, otherwise SDK will auto-detect
|
|
312
|
+
workspace_host_value: str | None = (
|
|
313
|
+
normalize_host(value_of(self.workspace_host))
|
|
314
|
+
if self.workspace_host
|
|
315
|
+
else None
|
|
316
|
+
)
|
|
317
|
+
return WorkspaceClient(
|
|
318
|
+
host=workspace_host_value,
|
|
319
|
+
token=forwarded_token,
|
|
320
|
+
auth_type="pat",
|
|
321
|
+
)
|
|
322
|
+
except (ImportError, LookupError):
|
|
323
|
+
# mlflow not available or headers not set - fall through to Model Serving
|
|
324
|
+
pass
|
|
325
|
+
|
|
326
|
+
# Fall back to Model Serving OBO (existing behavior)
|
|
293
327
|
credentials_strategy: CredentialsStrategy = ModelServingUserCredentials()
|
|
294
328
|
logger.debug(
|
|
295
329
|
f"Creating WorkspaceClient for {self.__class__.__name__} "
|
|
296
|
-
f"with OBO credentials strategy"
|
|
297
|
-
)
|
|
298
|
-
self._workspace_client = WorkspaceClient(
|
|
299
|
-
credentials_strategy=credentials_strategy
|
|
330
|
+
f"with OBO credentials strategy (Model Serving)"
|
|
300
331
|
)
|
|
301
|
-
return
|
|
332
|
+
return WorkspaceClient(credentials_strategy=credentials_strategy)
|
|
302
333
|
|
|
303
334
|
# Check for service principal credentials
|
|
304
335
|
client_id_value: str | None = (
|
|
@@ -318,13 +349,12 @@ class IsDatabricksResource(ABC, BaseModel):
|
|
|
318
349
|
f"Creating WorkspaceClient for {self.__class__.__name__} with service principal: "
|
|
319
350
|
f"client_id={client_id_value}, host={workspace_host_value}"
|
|
320
351
|
)
|
|
321
|
-
|
|
352
|
+
return WorkspaceClient(
|
|
322
353
|
host=workspace_host_value,
|
|
323
354
|
client_id=client_id_value,
|
|
324
355
|
client_secret=client_secret_value,
|
|
325
356
|
auth_type="oauth-m2m",
|
|
326
357
|
)
|
|
327
|
-
return self._workspace_client
|
|
328
358
|
|
|
329
359
|
# Check for PAT authentication
|
|
330
360
|
pat_value: str | None = value_of(self.pat) if self.pat else None
|
|
@@ -332,20 +362,28 @@ class IsDatabricksResource(ABC, BaseModel):
|
|
|
332
362
|
logger.debug(
|
|
333
363
|
f"Creating WorkspaceClient for {self.__class__.__name__} with PAT"
|
|
334
364
|
)
|
|
335
|
-
|
|
365
|
+
return WorkspaceClient(
|
|
336
366
|
host=workspace_host_value,
|
|
337
367
|
token=pat_value,
|
|
338
368
|
auth_type="pat",
|
|
339
369
|
)
|
|
340
|
-
return self._workspace_client
|
|
341
370
|
|
|
342
371
|
# Default: use ambient authentication
|
|
343
372
|
logger.debug(
|
|
344
373
|
f"Creating WorkspaceClient for {self.__class__.__name__} "
|
|
345
374
|
"with default/ambient authentication"
|
|
346
375
|
)
|
|
347
|
-
|
|
348
|
-
|
|
376
|
+
return WorkspaceClient()
|
|
377
|
+
|
|
378
|
+
|
|
379
|
+
class DeploymentTarget(str, Enum):
|
|
380
|
+
"""Target platform for agent deployment."""
|
|
381
|
+
|
|
382
|
+
MODEL_SERVING = "model_serving"
|
|
383
|
+
"""Deploy to Databricks Model Serving endpoint."""
|
|
384
|
+
|
|
385
|
+
APPS = "apps"
|
|
386
|
+
"""Deploy as a Databricks App."""
|
|
349
387
|
|
|
350
388
|
|
|
351
389
|
class Privilege(str, Enum):
|
|
@@ -865,10 +903,6 @@ class GenieRoomModel(IsDatabricksResource):
|
|
|
865
903
|
pat=self.pat,
|
|
866
904
|
)
|
|
867
905
|
|
|
868
|
-
# Share the cached workspace client if available
|
|
869
|
-
if self._workspace_client is not None:
|
|
870
|
-
warehouse_model._workspace_client = self._workspace_client
|
|
871
|
-
|
|
872
906
|
return warehouse_model
|
|
873
907
|
except Exception as e:
|
|
874
908
|
logger.warning(
|
|
@@ -912,9 +946,6 @@ class GenieRoomModel(IsDatabricksResource):
|
|
|
912
946
|
workspace_host=self.workspace_host,
|
|
913
947
|
pat=self.pat,
|
|
914
948
|
)
|
|
915
|
-
# Share the cached workspace client if available
|
|
916
|
-
if self._workspace_client is not None:
|
|
917
|
-
table_model._workspace_client = self._workspace_client
|
|
918
949
|
|
|
919
950
|
# Verify the table exists before adding
|
|
920
951
|
if not table_model.exists():
|
|
@@ -952,9 +983,6 @@ class GenieRoomModel(IsDatabricksResource):
|
|
|
952
983
|
workspace_host=self.workspace_host,
|
|
953
984
|
pat=self.pat,
|
|
954
985
|
)
|
|
955
|
-
# Share the cached workspace client if available
|
|
956
|
-
if self._workspace_client is not None:
|
|
957
|
-
function_model._workspace_client = self._workspace_client
|
|
958
986
|
|
|
959
987
|
# Verify the function exists before adding
|
|
960
988
|
if not function_model.exists():
|
|
@@ -3255,6 +3283,7 @@ class ResourcesModel(BaseModel):
|
|
|
3255
3283
|
|
|
3256
3284
|
class AppConfig(BaseModel):
|
|
3257
3285
|
model_config = ConfigDict(use_enum_values=True, extra="forbid")
|
|
3286
|
+
version: Optional[str] = None
|
|
3258
3287
|
variables: dict[str, AnyVariable] = Field(default_factory=dict)
|
|
3259
3288
|
service_principals: dict[str, ServicePrincipalModel] = Field(default_factory=dict)
|
|
3260
3289
|
schemas: dict[str, SchemaModel] = Field(default_factory=dict)
|
|
@@ -3275,6 +3304,9 @@ class AppConfig(BaseModel):
|
|
|
3275
3304
|
)
|
|
3276
3305
|
providers: Optional[dict[type | str, Any]] = None
|
|
3277
3306
|
|
|
3307
|
+
# Private attribute to track the source config file path (set by from_file)
|
|
3308
|
+
_source_config_path: str | None = None
|
|
3309
|
+
|
|
3278
3310
|
@classmethod
|
|
3279
3311
|
def from_file(cls, path: PathLike) -> "AppConfig":
|
|
3280
3312
|
path = Path(path).as_posix()
|
|
@@ -3282,12 +3314,20 @@ class AppConfig(BaseModel):
|
|
|
3282
3314
|
model_config: ModelConfig = ModelConfig(development_config=path)
|
|
3283
3315
|
config: AppConfig = AppConfig(**model_config.to_dict())
|
|
3284
3316
|
|
|
3317
|
+
# Store the source config path for later use (e.g., Apps deployment)
|
|
3318
|
+
config._source_config_path = path
|
|
3319
|
+
|
|
3285
3320
|
config.initialize()
|
|
3286
3321
|
|
|
3287
3322
|
atexit.register(config.shutdown)
|
|
3288
3323
|
|
|
3289
3324
|
return config
|
|
3290
3325
|
|
|
3326
|
+
@property
|
|
3327
|
+
def source_config_path(self) -> str | None:
|
|
3328
|
+
"""Get the source config file path if loaded via from_file."""
|
|
3329
|
+
return self._source_config_path
|
|
3330
|
+
|
|
3291
3331
|
def initialize(self) -> None:
|
|
3292
3332
|
from dao_ai.hooks.core import create_hooks
|
|
3293
3333
|
from dao_ai.logging import configure_logging
|
|
@@ -3358,6 +3398,7 @@ class AppConfig(BaseModel):
|
|
|
3358
3398
|
|
|
3359
3399
|
def deploy_agent(
|
|
3360
3400
|
self,
|
|
3401
|
+
target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
|
|
3361
3402
|
w: WorkspaceClient | None = None,
|
|
3362
3403
|
vsc: "VectorSearchClient | None" = None,
|
|
3363
3404
|
pat: str | None = None,
|
|
@@ -3365,6 +3406,18 @@ class AppConfig(BaseModel):
|
|
|
3365
3406
|
client_secret: str | None = None,
|
|
3366
3407
|
workspace_host: str | None = None,
|
|
3367
3408
|
) -> None:
|
|
3409
|
+
"""
|
|
3410
|
+
Deploy the agent to the specified target.
|
|
3411
|
+
|
|
3412
|
+
Args:
|
|
3413
|
+
target: The deployment target (MODEL_SERVING or APPS). Defaults to MODEL_SERVING.
|
|
3414
|
+
w: Optional WorkspaceClient instance
|
|
3415
|
+
vsc: Optional VectorSearchClient instance
|
|
3416
|
+
pat: Optional personal access token for authentication
|
|
3417
|
+
client_id: Optional client ID for service principal authentication
|
|
3418
|
+
client_secret: Optional client secret for service principal authentication
|
|
3419
|
+
workspace_host: Optional workspace host URL
|
|
3420
|
+
"""
|
|
3368
3421
|
from dao_ai.providers.base import ServiceProvider
|
|
3369
3422
|
from dao_ai.providers.databricks import DatabricksProvider
|
|
3370
3423
|
|
|
@@ -3376,7 +3429,7 @@ class AppConfig(BaseModel):
|
|
|
3376
3429
|
client_secret=client_secret,
|
|
3377
3430
|
workspace_host=workspace_host,
|
|
3378
3431
|
)
|
|
3379
|
-
provider.deploy_agent(self)
|
|
3432
|
+
provider.deploy_agent(self, target=target)
|
|
3380
3433
|
|
|
3381
3434
|
def find_agents(
|
|
3382
3435
|
self, predicate: Callable[[AgentModel], bool] | None = None
|
dao_ai/providers/base.py
CHANGED
|
@@ -1,15 +1,19 @@
|
|
|
1
1
|
from abc import ABC, abstractmethod
|
|
2
|
-
from typing import Any, Sequence
|
|
2
|
+
from typing import TYPE_CHECKING, Any, Sequence
|
|
3
3
|
|
|
4
4
|
from dao_ai.config import (
|
|
5
5
|
AppModel,
|
|
6
6
|
DatasetModel,
|
|
7
|
+
DeploymentTarget,
|
|
7
8
|
SchemaModel,
|
|
8
9
|
UnityCatalogFunctionSqlModel,
|
|
9
10
|
VectorStoreModel,
|
|
10
11
|
VolumeModel,
|
|
11
12
|
)
|
|
12
13
|
|
|
14
|
+
if TYPE_CHECKING:
|
|
15
|
+
from dao_ai.config import AppConfig
|
|
16
|
+
|
|
13
17
|
|
|
14
18
|
class ServiceProvider(ABC):
|
|
15
19
|
@abstractmethod
|
|
@@ -52,4 +56,26 @@ class ServiceProvider(ABC):
|
|
|
52
56
|
) -> Any: ...
|
|
53
57
|
|
|
54
58
|
@abstractmethod
|
|
55
|
-
def
|
|
59
|
+
def deploy_model_serving_agent(self, config: "AppConfig") -> Any:
|
|
60
|
+
"""Deploy agent to Databricks Model Serving endpoint."""
|
|
61
|
+
...
|
|
62
|
+
|
|
63
|
+
@abstractmethod
|
|
64
|
+
def deploy_apps_agent(self, config: "AppConfig") -> Any:
|
|
65
|
+
"""Deploy agent as a Databricks App."""
|
|
66
|
+
...
|
|
67
|
+
|
|
68
|
+
@abstractmethod
|
|
69
|
+
def deploy_agent(
|
|
70
|
+
self,
|
|
71
|
+
config: "AppConfig",
|
|
72
|
+
target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
|
|
73
|
+
) -> Any:
|
|
74
|
+
"""
|
|
75
|
+
Deploy agent to the specified target.
|
|
76
|
+
|
|
77
|
+
Args:
|
|
78
|
+
config: The AppConfig containing deployment configuration
|
|
79
|
+
target: The deployment target (MODEL_SERVING or APPS)
|
|
80
|
+
"""
|
|
81
|
+
...
|
dao_ai/providers/databricks.py
CHANGED
|
@@ -23,7 +23,7 @@ from databricks.sdk.service.catalog import (
|
|
|
23
23
|
)
|
|
24
24
|
from databricks.sdk.service.database import DatabaseCredential
|
|
25
25
|
from databricks.sdk.service.iam import User
|
|
26
|
-
from databricks.sdk.service.workspace import GetSecretResponse
|
|
26
|
+
from databricks.sdk.service.workspace import GetSecretResponse, ImportFormat
|
|
27
27
|
from databricks.vector_search.client import VectorSearchClient
|
|
28
28
|
from databricks.vector_search.index import VectorSearchIndex
|
|
29
29
|
from loguru import logger
|
|
@@ -48,6 +48,7 @@ from dao_ai.config import (
|
|
|
48
48
|
DatabaseModel,
|
|
49
49
|
DatabricksAppModel,
|
|
50
50
|
DatasetModel,
|
|
51
|
+
DeploymentTarget,
|
|
51
52
|
FunctionModel,
|
|
52
53
|
GenieRoomModel,
|
|
53
54
|
HasFullName,
|
|
@@ -439,8 +440,19 @@ class DatabricksProvider(ServiceProvider):
|
|
|
439
440
|
version=aliased_model.version,
|
|
440
441
|
)
|
|
441
442
|
|
|
442
|
-
def
|
|
443
|
-
|
|
443
|
+
def deploy_model_serving_agent(self, config: AppConfig) -> None:
|
|
444
|
+
"""
|
|
445
|
+
Deploy agent to Databricks Model Serving endpoint.
|
|
446
|
+
|
|
447
|
+
This is the original deployment method that creates/updates a Model Serving
|
|
448
|
+
endpoint with the registered model.
|
|
449
|
+
|
|
450
|
+
Args:
|
|
451
|
+
config: The AppConfig containing deployment configuration
|
|
452
|
+
"""
|
|
453
|
+
logger.info(
|
|
454
|
+
"Deploying agent to Model Serving", endpoint_name=config.app.endpoint_name
|
|
455
|
+
)
|
|
444
456
|
mlflow.set_registry_uri("databricks-uc")
|
|
445
457
|
|
|
446
458
|
endpoint_name: str = config.app.endpoint_name
|
|
@@ -499,6 +511,196 @@ class DatabricksProvider(ServiceProvider):
|
|
|
499
511
|
permission_level=PermissionLevel[entitlement],
|
|
500
512
|
)
|
|
501
513
|
|
|
514
|
+
def deploy_apps_agent(self, config: AppConfig) -> None:
|
|
515
|
+
"""
|
|
516
|
+
Deploy agent as a Databricks App.
|
|
517
|
+
|
|
518
|
+
This method creates or updates a Databricks App that serves the agent
|
|
519
|
+
using the app_server module.
|
|
520
|
+
|
|
521
|
+
The deployment process:
|
|
522
|
+
1. Determine the workspace source path for the app
|
|
523
|
+
2. Upload the configuration file to the workspace
|
|
524
|
+
3. Create the app if it doesn't exist
|
|
525
|
+
4. Deploy the app
|
|
526
|
+
|
|
527
|
+
Args:
|
|
528
|
+
config: The AppConfig containing deployment configuration
|
|
529
|
+
|
|
530
|
+
Note:
|
|
531
|
+
The config file must be loaded via AppConfig.from_file() so that
|
|
532
|
+
the source_config_path is available for upload.
|
|
533
|
+
"""
|
|
534
|
+
import io
|
|
535
|
+
|
|
536
|
+
from databricks.sdk.service.apps import (
|
|
537
|
+
App,
|
|
538
|
+
AppDeployment,
|
|
539
|
+
AppDeploymentMode,
|
|
540
|
+
AppDeploymentState,
|
|
541
|
+
)
|
|
542
|
+
|
|
543
|
+
# Normalize app name: lowercase, replace underscores with dashes
|
|
544
|
+
raw_name: str = config.app.name
|
|
545
|
+
app_name: str = raw_name.lower().replace("_", "-")
|
|
546
|
+
if app_name != raw_name:
|
|
547
|
+
logger.info(
|
|
548
|
+
"Normalized app name for Databricks Apps",
|
|
549
|
+
original=raw_name,
|
|
550
|
+
normalized=app_name,
|
|
551
|
+
)
|
|
552
|
+
logger.info("Deploying agent to Databricks Apps", app_name=app_name)
|
|
553
|
+
|
|
554
|
+
# Use convention-based workspace path: /Workspace/Users/{user}/apps/{app_name}
|
|
555
|
+
current_user: User = self.w.current_user.me()
|
|
556
|
+
user_name: str = current_user.user_name or "default"
|
|
557
|
+
source_path: str = f"/Workspace/Users/{user_name}/apps/{app_name}"
|
|
558
|
+
|
|
559
|
+
logger.info("Using workspace source path", source_path=source_path)
|
|
560
|
+
|
|
561
|
+
# Upload the configuration file to the workspace
|
|
562
|
+
source_config_path: str | None = config.source_config_path
|
|
563
|
+
if source_config_path:
|
|
564
|
+
# Read the config file and upload to workspace
|
|
565
|
+
config_file_name: str = "model_config.yaml"
|
|
566
|
+
workspace_config_path: str = f"{source_path}/{config_file_name}"
|
|
567
|
+
|
|
568
|
+
logger.info(
|
|
569
|
+
"Uploading config file to workspace",
|
|
570
|
+
source=source_config_path,
|
|
571
|
+
destination=workspace_config_path,
|
|
572
|
+
)
|
|
573
|
+
|
|
574
|
+
# Read the source config file
|
|
575
|
+
with open(source_config_path, "rb") as f:
|
|
576
|
+
config_content: bytes = f.read()
|
|
577
|
+
|
|
578
|
+
# Create the directory if it doesn't exist and upload the file
|
|
579
|
+
try:
|
|
580
|
+
self.w.workspace.mkdirs(source_path)
|
|
581
|
+
except Exception as e:
|
|
582
|
+
logger.debug(f"Directory may already exist: {e}")
|
|
583
|
+
|
|
584
|
+
# Upload the config file
|
|
585
|
+
self.w.workspace.upload(
|
|
586
|
+
path=workspace_config_path,
|
|
587
|
+
content=io.BytesIO(config_content),
|
|
588
|
+
format=ImportFormat.AUTO,
|
|
589
|
+
overwrite=True,
|
|
590
|
+
)
|
|
591
|
+
logger.info("Config file uploaded", path=workspace_config_path)
|
|
592
|
+
else:
|
|
593
|
+
logger.warning(
|
|
594
|
+
"No source config path available. "
|
|
595
|
+
"Ensure DAO_AI_CONFIG_PATH is set in the app environment or "
|
|
596
|
+
"model_config.yaml exists in the app source directory."
|
|
597
|
+
)
|
|
598
|
+
|
|
599
|
+
# Generate and upload app.yaml
|
|
600
|
+
# Use pip to install dao-ai and run the app server
|
|
601
|
+
app_yaml_content: str = """command:
|
|
602
|
+
- /bin/bash
|
|
603
|
+
- -c
|
|
604
|
+
- |
|
|
605
|
+
pip install dao-ai && python -m dao_ai.app_server
|
|
606
|
+
|
|
607
|
+
env:
|
|
608
|
+
- name: MLFLOW_TRACKING_URI
|
|
609
|
+
value: "databricks"
|
|
610
|
+
- name: MLFLOW_REGISTRY_URI
|
|
611
|
+
value: "databricks-uc"
|
|
612
|
+
- name: DAO_AI_CONFIG_PATH
|
|
613
|
+
value: "model_config.yaml"
|
|
614
|
+
"""
|
|
615
|
+
app_yaml_path: str = f"{source_path}/app.yaml"
|
|
616
|
+
self.w.workspace.upload(
|
|
617
|
+
path=app_yaml_path,
|
|
618
|
+
content=io.BytesIO(app_yaml_content.encode("utf-8")),
|
|
619
|
+
format=ImportFormat.AUTO,
|
|
620
|
+
overwrite=True,
|
|
621
|
+
)
|
|
622
|
+
logger.info("app.yaml uploaded", path=app_yaml_path)
|
|
623
|
+
|
|
624
|
+
# Check if app exists
|
|
625
|
+
app_exists: bool = False
|
|
626
|
+
try:
|
|
627
|
+
existing_app: App = self.w.apps.get(name=app_name)
|
|
628
|
+
app_exists = True
|
|
629
|
+
logger.debug("App already exists, updating", app_name=app_name)
|
|
630
|
+
except NotFound:
|
|
631
|
+
logger.debug("Creating new app", app_name=app_name)
|
|
632
|
+
|
|
633
|
+
# Create or get the app
|
|
634
|
+
if not app_exists:
|
|
635
|
+
logger.info("Creating Databricks App", app_name=app_name)
|
|
636
|
+
app_spec = App(
|
|
637
|
+
name=app_name,
|
|
638
|
+
description=config.app.description or f"DAO AI Agent: {app_name}",
|
|
639
|
+
)
|
|
640
|
+
app: App = self.w.apps.create_and_wait(app=app_spec)
|
|
641
|
+
logger.info("App created", app_name=app.name, app_url=app.url)
|
|
642
|
+
else:
|
|
643
|
+
app = existing_app
|
|
644
|
+
|
|
645
|
+
# Deploy the app with source code
|
|
646
|
+
# The app will use the dao_ai.app_server module as the entry point
|
|
647
|
+
logger.info("Deploying app", app_name=app_name)
|
|
648
|
+
|
|
649
|
+
# Create deployment configuration
|
|
650
|
+
app_deployment = AppDeployment(
|
|
651
|
+
mode=AppDeploymentMode.SNAPSHOT,
|
|
652
|
+
source_code_path=source_path,
|
|
653
|
+
)
|
|
654
|
+
|
|
655
|
+
# Deploy the app
|
|
656
|
+
deployment: AppDeployment = self.w.apps.deploy_and_wait(
|
|
657
|
+
app_name=app_name,
|
|
658
|
+
app_deployment=app_deployment,
|
|
659
|
+
)
|
|
660
|
+
|
|
661
|
+
if (
|
|
662
|
+
deployment.status
|
|
663
|
+
and deployment.status.state == AppDeploymentState.SUCCEEDED
|
|
664
|
+
):
|
|
665
|
+
logger.info(
|
|
666
|
+
"App deployed successfully",
|
|
667
|
+
app_name=app_name,
|
|
668
|
+
deployment_id=deployment.deployment_id,
|
|
669
|
+
app_url=app.url if app else None,
|
|
670
|
+
)
|
|
671
|
+
else:
|
|
672
|
+
status_message: str = (
|
|
673
|
+
deployment.status.message if deployment.status else "Unknown error"
|
|
674
|
+
)
|
|
675
|
+
logger.error(
|
|
676
|
+
"App deployment failed",
|
|
677
|
+
app_name=app_name,
|
|
678
|
+
status=status_message,
|
|
679
|
+
)
|
|
680
|
+
raise RuntimeError(f"App deployment failed: {status_message}")
|
|
681
|
+
|
|
682
|
+
def deploy_agent(
|
|
683
|
+
self,
|
|
684
|
+
config: AppConfig,
|
|
685
|
+
target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
|
|
686
|
+
) -> None:
|
|
687
|
+
"""
|
|
688
|
+
Deploy agent to the specified target.
|
|
689
|
+
|
|
690
|
+
This is the main deployment method that routes to the appropriate
|
|
691
|
+
deployment implementation based on the target.
|
|
692
|
+
|
|
693
|
+
Args:
|
|
694
|
+
config: The AppConfig containing deployment configuration
|
|
695
|
+
target: The deployment target (MODEL_SERVING or APPS)
|
|
696
|
+
"""
|
|
697
|
+
if target == DeploymentTarget.MODEL_SERVING:
|
|
698
|
+
self.deploy_model_serving_agent(config)
|
|
699
|
+
elif target == DeploymentTarget.APPS:
|
|
700
|
+
self.deploy_apps_agent(config)
|
|
701
|
+
else:
|
|
702
|
+
raise ValueError(f"Unknown deployment target: {target}")
|
|
703
|
+
|
|
502
704
|
def create_catalog(self, schema: SchemaModel) -> CatalogInfo:
|
|
503
705
|
catalog_info: CatalogInfo
|
|
504
706
|
try:
|
dao_ai/state.py
CHANGED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: dao-ai
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.9
|
|
4
4
|
Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
|
|
5
5
|
Project-URL: Homepage, https://github.com/natefleming/dao-ai
|
|
6
6
|
Project-URL: Documentation, https://natefleming.github.io/dao-ai
|
|
@@ -43,7 +43,7 @@ Requires-Dist: langgraph>=1.0.5
|
|
|
43
43
|
Requires-Dist: langmem>=0.0.30
|
|
44
44
|
Requires-Dist: loguru>=0.7.3
|
|
45
45
|
Requires-Dist: mcp>=1.24.0
|
|
46
|
-
Requires-Dist: mlflow>=3.8.1
|
|
46
|
+
Requires-Dist: mlflow[databricks]>=3.8.1
|
|
47
47
|
Requires-Dist: nest-asyncio>=1.6.0
|
|
48
48
|
Requires-Dist: openevals>=0.1.3
|
|
49
49
|
Requires-Dist: openpyxl>=3.1.5
|
|
@@ -1,8 +1,9 @@
|
|
|
1
1
|
dao_ai/__init__.py,sha256=18P98ExEgUaJ1Byw440Ct1ty59v6nxyWtc5S6Uq2m9Q,1062
|
|
2
2
|
dao_ai/agent_as_code.py,sha256=xIlLDpPVfmDVzLvbdY_V_CrC4Jvj2ItCWJ-NzdrszTo,538
|
|
3
|
+
dao_ai/app_server.py,sha256=QKpl068z-s1gLF67dPW-3fT77i33t_Oab4_ugmxISWs,3010
|
|
3
4
|
dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
|
|
4
|
-
dao_ai/cli.py,sha256=
|
|
5
|
-
dao_ai/config.py,sha256=
|
|
5
|
+
dao_ai/cli.py,sha256=57Kmmi0zgS92ACBTD-gH5hZzW6rPDkbdkRVlFjX4onQ,48604
|
|
6
|
+
dao_ai/config.py,sha256=t7kXU7XjdMaCZ3G9Hn-O9NDOaTS_LaMXX6s5mdyv3dM,127944
|
|
6
7
|
dao_ai/graph.py,sha256=1-uQlo7iXZQTT3uU8aYu0N5rnhw5_g_2YLwVsAs6M-U,1119
|
|
7
8
|
dao_ai/logging.py,sha256=lYy4BmucCHvwW7aI3YQkQXKJtMvtTnPDu9Hnd7_O4oc,1556
|
|
8
9
|
dao_ai/messages.py,sha256=4ZBzO4iFdktGSLrmhHzFjzMIt2tpaL-aQLHOQJysGnY,6959
|
|
@@ -10,7 +11,7 @@ dao_ai/models.py,sha256=AwzwTRTNZF-UOh59HsuXEgFk_YH6q6M-mERNDe64Z8k,81783
|
|
|
10
11
|
dao_ai/nodes.py,sha256=7W6Ek6Uk9-pKa-H06nVCwuDllCrgX02IYy3rHtuL0aM,10777
|
|
11
12
|
dao_ai/optimization.py,sha256=phK6t4wYmWPObCjGUBHdZzsaFXGhQOjhAek2bAEfwXo,22971
|
|
12
13
|
dao_ai/prompts.py,sha256=4cz5bZ7cOzrjyQ8hMp-K4evK6cVYrkGrAGdUl8-KDEM,2784
|
|
13
|
-
dao_ai/state.py,sha256=
|
|
14
|
+
dao_ai/state.py,sha256=ifDTAC7epdowk3Z1CP3Xqw4uH2dIxQEVF3C747dA8yI,6436
|
|
14
15
|
dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
15
16
|
dao_ai/utils.py,sha256=_Urd7Nj2VzrgPKf3NS4E6vt0lWRhEUddBqWN9BksqeE,11543
|
|
16
17
|
dao_ai/vector_search.py,sha256=8d3xROg9zSIYNXjRRl6rSexsJTlufjRl5Fy1ZA8daKA,4019
|
|
@@ -48,8 +49,8 @@ dao_ai/orchestration/core.py,sha256=qoU7uMXBJCth-sqfu0jRE1L0GOn5H4LoZdRUY1Ib3DI,
|
|
|
48
49
|
dao_ai/orchestration/supervisor.py,sha256=alKMEEo9G5LhdpMvTVdAMel234cZj5_MguWl4wFB7XQ,9873
|
|
49
50
|
dao_ai/orchestration/swarm.py,sha256=8tp1eGmsQqqWpaDcjPoJckddPWohZdmmN0RGRJ_xzOA,9198
|
|
50
51
|
dao_ai/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
51
|
-
dao_ai/providers/base.py,sha256
|
|
52
|
-
dao_ai/providers/databricks.py,sha256=
|
|
52
|
+
dao_ai/providers/base.py,sha256=cJGo3UjUTPgS91dv38ePOHwQQtYhIa84ebb167CBXjk,2111
|
|
53
|
+
dao_ai/providers/databricks.py,sha256=X_VjzZogwiSlNpPBWP0iMCCXAvfFDRlbC4AZCHleb2A,68608
|
|
53
54
|
dao_ai/tools/__init__.py,sha256=NfRpAKds_taHbx6gzLPWgtPXve-YpwzkoOAUflwxceM,1734
|
|
54
55
|
dao_ai/tools/agent.py,sha256=plIWALywRjaDSnot13nYehBsrHRpBUpsVZakoGeajOE,1858
|
|
55
56
|
dao_ai/tools/core.py,sha256=bRIN3BZhRQX8-Kpu3HPomliodyskCqjxynQmYbk6Vjs,3783
|
|
@@ -64,8 +65,8 @@ dao_ai/tools/sql.py,sha256=tKd1gjpLuKdQDyfmyYYtMiNRHDW6MGRbdEVaeqyB8Ok,7632
|
|
|
64
65
|
dao_ai/tools/time.py,sha256=tufJniwivq29y0LIffbgeBTIDE6VgrLpmVf8Qr90qjw,9224
|
|
65
66
|
dao_ai/tools/unity_catalog.py,sha256=AjQfW7bvV8NurqDLIyntYRv2eJuTwNdbvex1L5CRjOk,15534
|
|
66
67
|
dao_ai/tools/vector_search.py,sha256=oe2uBwl2TfeJIXPpwiS6Rmz7wcHczSxNyqS9P3hE6co,14542
|
|
67
|
-
dao_ai-0.1.
|
|
68
|
-
dao_ai-0.1.
|
|
69
|
-
dao_ai-0.1.
|
|
70
|
-
dao_ai-0.1.
|
|
71
|
-
dao_ai-0.1.
|
|
68
|
+
dao_ai-0.1.9.dist-info/METADATA,sha256=_cVhNKdywyci5ZyrmMfkF9CiWveo4kdoj9OlaM_w-cg,16697
|
|
69
|
+
dao_ai-0.1.9.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
|
|
70
|
+
dao_ai-0.1.9.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
|
|
71
|
+
dao_ai-0.1.9.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
|
|
72
|
+
dao_ai-0.1.9.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|