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 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 via ModelServingUserCredentials.
211
+ calling user's identity. Implementation varies by deployment:
212
+ - Databricks Apps: Uses X-Forwarded-Access-Token from request headers
213
+ - Model Serving: Uses ModelServingUserCredentials
212
214
 
213
215
  2. **Service Principal (OAuth M2M)**: Provide service_principal or
214
216
  (client_id + client_secret + workspace_host) for service principal auth.
@@ -221,9 +223,17 @@ class IsDatabricksResource(ABC, BaseModel):
221
223
 
222
224
  Authentication Priority:
223
225
  1. OBO (on_behalf_of_user=True)
226
+ - Checks for forwarded headers (Databricks Apps)
227
+ - Falls back to ModelServingUserCredentials (Model Serving)
224
228
  2. Service Principal (client_id + client_secret + workspace_host)
225
229
  3. PAT (pat + workspace_host)
226
230
  4. Ambient/default authentication
231
+
232
+ Note: When on_behalf_of_user=True, the agent acts as the calling user regardless
233
+ of deployment target. In Databricks Apps, this uses X-Forwarded-Access-Token
234
+ automatically captured by MLflow AgentServer. In Model Serving, this uses
235
+ ModelServingUserCredentials. Forwarded headers are ONLY used when
236
+ on_behalf_of_user=True.
227
237
  """
228
238
 
229
239
  model_config = ConfigDict(use_enum_values=True)
@@ -235,9 +245,6 @@ class IsDatabricksResource(ABC, BaseModel):
235
245
  workspace_host: Optional[AnyVariable] = None
236
246
  pat: Optional[AnyVariable] = None
237
247
 
238
- # Private attribute to cache the workspace client (lazy instantiation)
239
- _workspace_client: Optional[WorkspaceClient] = PrivateAttr(default=None)
240
-
241
248
  @abstractmethod
242
249
  def as_resources(self) -> Sequence[DatabricksResource]: ...
243
250
 
@@ -273,32 +280,56 @@ class IsDatabricksResource(ABC, BaseModel):
273
280
  """
274
281
  Get a WorkspaceClient configured with the appropriate authentication.
275
282
 
276
- The client is lazily instantiated on first access and cached for subsequent calls.
283
+ A new client is created on each access.
277
284
 
278
285
  Authentication priority:
279
- 1. If on_behalf_of_user is True, uses ModelServingUserCredentials (OBO)
280
- 2. If service principal credentials are configured (client_id, client_secret,
281
- workspace_host), uses OAuth M2M
282
- 3. If PAT is configured, uses token authentication
283
- 4. Otherwise, uses default/ambient authentication
286
+ 1. On-Behalf-Of User (on_behalf_of_user=True):
287
+ - Forwarded headers (Databricks Apps)
288
+ - ModelServingUserCredentials (Model Serving)
289
+ 2. Service Principal (client_id + client_secret + workspace_host)
290
+ 3. PAT (pat + workspace_host)
291
+ 4. Ambient/default authentication
284
292
  """
285
- # Return cached client if already instantiated
286
- if self._workspace_client is not None:
287
- return self._workspace_client
288
-
289
293
  from dao_ai.utils import normalize_host
290
294
 
291
295
  # Check for OBO first (highest priority)
292
296
  if self.on_behalf_of_user:
297
+ # NEW: In Databricks Apps, use forwarded headers for per-user auth
298
+ try:
299
+ from mlflow.genai.agent_server import get_request_headers
300
+
301
+ headers = get_request_headers()
302
+ forwarded_token = headers.get("x-forwarded-access-token")
303
+
304
+ if forwarded_token:
305
+ forwarded_user = headers.get("x-forwarded-user", "unknown")
306
+ logger.debug(
307
+ f"Creating WorkspaceClient for {self.__class__.__name__} "
308
+ f"with OBO using forwarded token from Databricks Apps",
309
+ forwarded_user=forwarded_user,
310
+ )
311
+ # Use workspace_host if configured, otherwise SDK will auto-detect
312
+ workspace_host_value: str | None = (
313
+ normalize_host(value_of(self.workspace_host))
314
+ if self.workspace_host
315
+ else None
316
+ )
317
+ return WorkspaceClient(
318
+ host=workspace_host_value,
319
+ token=forwarded_token,
320
+ auth_type="pat",
321
+ )
322
+ except (ImportError, LookupError):
323
+ # mlflow not available or headers not set - fall through to Model Serving
324
+ pass
325
+
326
+ # Fall back to Model Serving OBO (existing behavior)
293
327
  credentials_strategy: CredentialsStrategy = ModelServingUserCredentials()
294
328
  logger.debug(
295
329
  f"Creating WorkspaceClient for {self.__class__.__name__} "
296
- f"with OBO credentials strategy"
297
- )
298
- self._workspace_client = WorkspaceClient(
299
- credentials_strategy=credentials_strategy
330
+ f"with OBO credentials strategy (Model Serving)"
300
331
  )
301
- return self._workspace_client
332
+ return WorkspaceClient(credentials_strategy=credentials_strategy)
302
333
 
303
334
  # Check for service principal credentials
304
335
  client_id_value: str | None = (
@@ -318,13 +349,12 @@ class IsDatabricksResource(ABC, BaseModel):
318
349
  f"Creating WorkspaceClient for {self.__class__.__name__} with service principal: "
319
350
  f"client_id={client_id_value}, host={workspace_host_value}"
320
351
  )
321
- self._workspace_client = WorkspaceClient(
352
+ return WorkspaceClient(
322
353
  host=workspace_host_value,
323
354
  client_id=client_id_value,
324
355
  client_secret=client_secret_value,
325
356
  auth_type="oauth-m2m",
326
357
  )
327
- return self._workspace_client
328
358
 
329
359
  # Check for PAT authentication
330
360
  pat_value: str | None = value_of(self.pat) if self.pat else None
@@ -332,20 +362,28 @@ class IsDatabricksResource(ABC, BaseModel):
332
362
  logger.debug(
333
363
  f"Creating WorkspaceClient for {self.__class__.__name__} with PAT"
334
364
  )
335
- self._workspace_client = WorkspaceClient(
365
+ return WorkspaceClient(
336
366
  host=workspace_host_value,
337
367
  token=pat_value,
338
368
  auth_type="pat",
339
369
  )
340
- return self._workspace_client
341
370
 
342
371
  # Default: use ambient authentication
343
372
  logger.debug(
344
373
  f"Creating WorkspaceClient for {self.__class__.__name__} "
345
374
  "with default/ambient authentication"
346
375
  )
347
- self._workspace_client = WorkspaceClient()
348
- return self._workspace_client
376
+ return WorkspaceClient()
377
+
378
+
379
+ class DeploymentTarget(str, Enum):
380
+ """Target platform for agent deployment."""
381
+
382
+ MODEL_SERVING = "model_serving"
383
+ """Deploy to Databricks Model Serving endpoint."""
384
+
385
+ APPS = "apps"
386
+ """Deploy as a Databricks App."""
349
387
 
350
388
 
351
389
  class Privilege(str, Enum):
@@ -865,10 +903,6 @@ class GenieRoomModel(IsDatabricksResource):
865
903
  pat=self.pat,
866
904
  )
867
905
 
868
- # Share the cached workspace client if available
869
- if self._workspace_client is not None:
870
- warehouse_model._workspace_client = self._workspace_client
871
-
872
906
  return warehouse_model
873
907
  except Exception as e:
874
908
  logger.warning(
@@ -912,9 +946,6 @@ class GenieRoomModel(IsDatabricksResource):
912
946
  workspace_host=self.workspace_host,
913
947
  pat=self.pat,
914
948
  )
915
- # Share the cached workspace client if available
916
- if self._workspace_client is not None:
917
- table_model._workspace_client = self._workspace_client
918
949
 
919
950
  # Verify the table exists before adding
920
951
  if not table_model.exists():
@@ -952,9 +983,6 @@ class GenieRoomModel(IsDatabricksResource):
952
983
  workspace_host=self.workspace_host,
953
984
  pat=self.pat,
954
985
  )
955
- # Share the cached workspace client if available
956
- if self._workspace_client is not None:
957
- function_model._workspace_client = self._workspace_client
958
986
 
959
987
  # Verify the function exists before adding
960
988
  if not function_model.exists():
@@ -3255,6 +3283,7 @@ class ResourcesModel(BaseModel):
3255
3283
 
3256
3284
  class AppConfig(BaseModel):
3257
3285
  model_config = ConfigDict(use_enum_values=True, extra="forbid")
3286
+ version: Optional[str] = None
3258
3287
  variables: dict[str, AnyVariable] = Field(default_factory=dict)
3259
3288
  service_principals: dict[str, ServicePrincipalModel] = Field(default_factory=dict)
3260
3289
  schemas: dict[str, SchemaModel] = Field(default_factory=dict)
@@ -3275,6 +3304,9 @@ class AppConfig(BaseModel):
3275
3304
  )
3276
3305
  providers: Optional[dict[type | str, Any]] = None
3277
3306
 
3307
+ # Private attribute to track the source config file path (set by from_file)
3308
+ _source_config_path: str | None = None
3309
+
3278
3310
  @classmethod
3279
3311
  def from_file(cls, path: PathLike) -> "AppConfig":
3280
3312
  path = Path(path).as_posix()
@@ -3282,12 +3314,20 @@ class AppConfig(BaseModel):
3282
3314
  model_config: ModelConfig = ModelConfig(development_config=path)
3283
3315
  config: AppConfig = AppConfig(**model_config.to_dict())
3284
3316
 
3317
+ # Store the source config path for later use (e.g., Apps deployment)
3318
+ config._source_config_path = path
3319
+
3285
3320
  config.initialize()
3286
3321
 
3287
3322
  atexit.register(config.shutdown)
3288
3323
 
3289
3324
  return config
3290
3325
 
3326
+ @property
3327
+ def source_config_path(self) -> str | None:
3328
+ """Get the source config file path if loaded via from_file."""
3329
+ return self._source_config_path
3330
+
3291
3331
  def initialize(self) -> None:
3292
3332
  from dao_ai.hooks.core import create_hooks
3293
3333
  from dao_ai.logging import configure_logging
@@ -3358,6 +3398,7 @@ class AppConfig(BaseModel):
3358
3398
 
3359
3399
  def deploy_agent(
3360
3400
  self,
3401
+ target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
3361
3402
  w: WorkspaceClient | None = None,
3362
3403
  vsc: "VectorSearchClient | None" = None,
3363
3404
  pat: str | None = None,
@@ -3365,6 +3406,18 @@ class AppConfig(BaseModel):
3365
3406
  client_secret: str | None = None,
3366
3407
  workspace_host: str | None = None,
3367
3408
  ) -> None:
3409
+ """
3410
+ Deploy the agent to the specified target.
3411
+
3412
+ Args:
3413
+ target: The deployment target (MODEL_SERVING or APPS). Defaults to MODEL_SERVING.
3414
+ w: Optional WorkspaceClient instance
3415
+ vsc: Optional VectorSearchClient instance
3416
+ pat: Optional personal access token for authentication
3417
+ client_id: Optional client ID for service principal authentication
3418
+ client_secret: Optional client secret for service principal authentication
3419
+ workspace_host: Optional workspace host URL
3420
+ """
3368
3421
  from dao_ai.providers.base import ServiceProvider
3369
3422
  from dao_ai.providers.databricks import DatabricksProvider
3370
3423
 
@@ -3376,7 +3429,7 @@ class AppConfig(BaseModel):
3376
3429
  client_secret=client_secret,
3377
3430
  workspace_host=workspace_host,
3378
3431
  )
3379
- provider.deploy_agent(self)
3432
+ provider.deploy_agent(self, target=target)
3380
3433
 
3381
3434
  def find_agents(
3382
3435
  self, predicate: Callable[[AgentModel], bool] | None = None
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 deploy_agent(self, config: AppModel) -> Any: ...
59
+ def deploy_model_serving_agent(self, config: "AppConfig") -> Any:
60
+ """Deploy agent to Databricks Model Serving endpoint."""
61
+ ...
62
+
63
+ @abstractmethod
64
+ def deploy_apps_agent(self, config: "AppConfig") -> Any:
65
+ """Deploy agent as a Databricks App."""
66
+ ...
67
+
68
+ @abstractmethod
69
+ def deploy_agent(
70
+ self,
71
+ config: "AppConfig",
72
+ target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
73
+ ) -> Any:
74
+ """
75
+ Deploy agent to the specified target.
76
+
77
+ Args:
78
+ config: The AppConfig containing deployment configuration
79
+ target: The deployment target (MODEL_SERVING or APPS)
80
+ """
81
+ ...
@@ -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 deploy_agent(self, config: AppConfig) -> None:
443
- logger.info("Deploying agent", endpoint_name=config.app.endpoint_name)
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
@@ -164,6 +164,7 @@ class Context(BaseModel):
164
164
 
165
165
  user_id: str | None = None
166
166
  thread_id: str | None = None
167
+ headers: dict[str, Any] | None = None
167
168
 
168
169
  @classmethod
169
170
  def from_runnable_config(cls, config: dict[str, Any]) -> "Context":
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dao-ai
3
- Version: 0.1.8
3
+ Version: 0.1.9
4
4
  Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
5
5
  Project-URL: Homepage, https://github.com/natefleming/dao-ai
6
6
  Project-URL: Documentation, https://natefleming.github.io/dao-ai
@@ -43,7 +43,7 @@ Requires-Dist: langgraph>=1.0.5
43
43
  Requires-Dist: langmem>=0.0.30
44
44
  Requires-Dist: loguru>=0.7.3
45
45
  Requires-Dist: mcp>=1.24.0
46
- Requires-Dist: mlflow>=3.8.1
46
+ Requires-Dist: mlflow[databricks]>=3.8.1
47
47
  Requires-Dist: nest-asyncio>=1.6.0
48
48
  Requires-Dist: openevals>=0.1.3
49
49
  Requires-Dist: openpyxl>=3.1.5
@@ -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=7LGrVDRgSBpznr8c8EksAhzPW_8NJ9h4St3DSpx-0z4,48196
5
- dao_ai/config.py,sha256=OSAsqb2Rxn3ghnM5Nq7-wh13DizHzWI_6cxuRuT4_j8,125773
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=0wbbzfQmldkCu26gdTE5j0Rl-_pfilza-YIHPbSWlvI,6394
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=-fjKypCOk28h6vioPfMj9YZSw_3Kcbi2nMuAyY7vX9k,1383
52
- dao_ai/providers/databricks.py,sha256=XxYkyoDYkwGV_Xg1IJBpGOl4d7U5HiFP4RtjjSLgenI,61437
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.8.dist-info/METADATA,sha256=kbbCJVZAI-U3czxwWr9z36m14PQk8poQdzOB_6RRLII,16685
68
- dao_ai-0.1.8.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
69
- dao_ai-0.1.8.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
70
- dao_ai-0.1.8.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
71
- dao_ai-0.1.8.dist-info/RECORD,,
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