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

{dao_ai-0.1.6 → dao_ai-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.6
+Version: 0.1.8
 Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
 Project-URL: Homepage, https://github.com/natefleming/dao-ai
 Project-URL: Documentation, https://natefleming.github.io/dao-ai

{dao_ai-0.1.6 → dao_ai-0.1.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dao-ai"
-version = "0.1.6"
+version = "0.1.8"
 description = "DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML."
 readme = "README.md"
 license = { text = "MIT" }

{dao_ai-0.1.6 → dao_ai-0.1.8}/requirements.txt

@@ -63,7 +63,6 @@ grandalf==0.8
 graphene==3.4.3
 graphql-core==3.2.7
 graphql-relay==3.2.0
-greenlet==3.3.0
 grpcio==1.76.0
 grpcio-status==1.76.0
 gunicorn==23.0.0

{dao_ai-0.1.6 → dao_ai-0.1.8}/schemas/model_config_schema.json

@@ -1196,6 +1196,7 @@
     },
     "DatabricksAppModel": {
       "additionalProperties": false,
+      "description": "Configuration for a Databricks App resource.\n\nThe `name` is the unique instance name of the Databricks App within the workspace.\nThe `url` is dynamically retrieved from the workspace client by calling\n`apps.get(name)` and returning the app's URL.\n\nExample:\n    ```yaml\n    resources:\n      apps:\n        my_app:\n          name: my-databricks-app\n    ```",
       "properties": {
         "on_behalf_of_user": {
           "anyOf": [
@@ -1355,40 +1356,10 @@
         "name": {
           "title": "Name",
           "type": "string"
-        },
-        "url": {
-          "anyOf": [
-            {
-              "$ref": "#/$defs/CompositeVariableModel"
-            },
-            {
-              "$ref": "#/$defs/EnvironmentVariableModel"
-            },
-            {
-              "$ref": "#/$defs/SecretVariableModel"
-            },
-            {
-              "$ref": "#/$defs/PrimitiveVariableModel"
-            },
-            {
-              "type": "string"
-            },
-            {
-              "type": "integer"
-            },
-            {
-              "type": "number"
-            },
-            {
-              "type": "boolean"
-            }
-          ],
-          "title": "Url"
         }
       },
       "required": [
-        "name",
-        "url"
+        "name"
       ],
       "title": "DatabricksAppModel",
       "type": "object"
@@ -2082,8 +2053,16 @@
           "title": "Pat"
         },
         "name": {
-          "title": "Name",
-          "type": "string"
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Name"
         },
         "description": {
           "anyOf": [
@@ -2128,7 +2107,6 @@
         }
       },
       "required": [
-        "name",
         "space_id"
       ],
       "title": "GenieRoomModel",
@@ -2935,6 +2913,17 @@
           "title": "Args",
           "type": "array"
         },
+        "app": {
+          "anyOf": [
+            {
+              "$ref": "#/$defs/DatabricksAppModel"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null
+        },
         "connection": {
           "anyOf": [
             {
@@ -4932,6 +4921,7 @@
       "type": "object"
     },
     "WarehouseModel": {
+      "additionalProperties": false,
       "properties": {
         "on_behalf_of_user": {
           "anyOf": [
@@ -5089,8 +5079,16 @@
           "title": "Pat"
         },
         "name": {
-          "title": "Name",
-          "type": "string"
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Name"
         },
         "description": {
           "anyOf": [
@@ -5135,7 +5133,6 @@
         }
       },
       "required": [
-        "name",
         "warehouse_id"
       ],
       "title": "WarehouseModel",

{dao_ai-0.1.6 → dao_ai-0.1.8}/src/dao_ai/config.py

@@ -24,6 +24,7 @@ from databricks.sdk.credentials_provider import (
     ModelServingUserCredentials,
 )
 from databricks.sdk.errors.platform import NotFound
+from databricks.sdk.service.apps import App
 from databricks.sdk.service.catalog import FunctionInfo, TableInfo
 from databricks.sdk.service.dashboards import GenieSpace
 from databricks.sdk.service.database import DatabaseInstance
@@ -147,7 +148,7 @@ class PrimitiveVariableModel(BaseModel, HasValue):
         return str(value)
 
     @model_validator(mode="after")
-    def validate_value(self) -> "PrimitiveVariableModel":
+    def validate_value(self) -> Self:
         if not isinstance(self.as_value(), (str, int, float, bool)):
             raise ValueError("Value must be a primitive type (str, int, float, bool)")
         return self
@@ -415,15 +416,44 @@ class SchemaModel(BaseModel, HasFullName):
 
 
 class DatabricksAppModel(IsDatabricksResource, HasFullName):
+    """
+    Configuration for a Databricks App resource.
+
+    The `name` is the unique instance name of the Databricks App within the workspace.
+    The `url` is dynamically retrieved from the workspace client by calling
+    `apps.get(name)` and returning the app's URL.
+
+    Example:
+        ```yaml
+        resources:
+          apps:
+            my_app:
+              name: my-databricks-app
+        ```
+    """
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
-    url: AnyVariable
+    """The unique instance name of the Databricks App in the workspace."""
 
-    @model_validator(mode="after")
-    def resolve_variables(self) -> Self:
-        """Resolve AnyVariable fields to their actual string values."""
-        self.url = value_of(self.url)
-        return self
+    @property
+    def url(self) -> str:
+        """
+        Retrieve the URL of the Databricks App from the workspace.
+
+        Returns:
+            The URL of the deployed Databricks App.
+
+        Raises:
+            RuntimeError: If the app is not found or URL is not available.
+        """
+        app: App = self.workspace_client.apps.get(self.name)
+        if app.url is None:
+            raise RuntimeError(
+                f"Databricks App '{self.name}' does not have a URL. "
+                "The app may not be deployed yet."
+            )
+        return app.url
 
     @property
     def full_name(self) -> str:
@@ -445,7 +475,7 @@ class TableModel(IsDatabricksResource, HasFullName):
     name: Optional[str] = None
 
     @model_validator(mode="after")
-    def validate_name_or_schema_required(self) -> "TableModel":
+    def validate_name_or_schema_required(self) -> Self:
         if not self.name and not self.schema_model:
             raise ValueError(
                 "Either 'name' or 'schema_model' must be provided for TableModel"
@@ -730,11 +760,20 @@ class FunctionModel(IsDatabricksResource, HasFullName):
 
 
 class WarehouseModel(IsDatabricksResource):
-    model_config = ConfigDict()
-    name: str
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    name: Optional[str] = None
     description: Optional[str] = None
     warehouse_id: AnyVariable
 
+    _warehouse_details: Optional[GetWarehouseResponse] = PrivateAttr(default=None)
+
+    def _get_warehouse_details(self) -> GetWarehouseResponse:
+        if self._warehouse_details is None:
+            self._warehouse_details = self.workspace_client.warehouses.get(
+                id=value_of(self.warehouse_id)
+            )
+        return self._warehouse_details
+
     @property
     def api_scopes(self) -> Sequence[str]:
         return [
@@ -755,10 +794,22 @@ class WarehouseModel(IsDatabricksResource):
         self.warehouse_id = value_of(self.warehouse_id)
         return self
 
+    @model_validator(mode="after")
+    def populate_name(self) -> Self:
+        """Populate name from warehouse details if not provided."""
+        if self.warehouse_id and not self.name:
+            try:
+                warehouse_details = self._get_warehouse_details()
+                if warehouse_details.name:
+                    self.name = warehouse_details.name
+            except Exception as e:
+                logger.debug(f"Could not fetch details from warehouse: {e}")
+        return self
+
 
 class GenieRoomModel(IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    name: str
+    name: Optional[str] = None
     description: Optional[str] = None
     space_id: AnyVariable
 
@@ -967,15 +1018,17 @@ class GenieRoomModel(IsDatabricksResource):
         return self
 
     @model_validator(mode="after")
-    def update_description_from_space(self) -> Self:
-        """Populate description from GenieSpace if not provided."""
-        if not self.description:
+    def populate_name_and_description(self) -> Self:
+        """Populate name and description from GenieSpace if not provided."""
+        if self.space_id and (not self.name or not self.description):
             try:
                 space_details = self._get_space_details()
-                if space_details.description:
+                if not self.name and space_details.title:
+                    self.name = space_details.title
+                if not self.description and space_details.description:
                     self.description = space_details.description
             except Exception as e:
-                logger.debug(f"Could not fetch description from Genie space: {e}")
+                logger.debug(f"Could not fetch details from Genie space: {e}")
         return self
 
 
@@ -1011,7 +1064,7 @@ class VolumePathModel(BaseModel, HasFullName):
     path: Optional[str] = None
 
     @model_validator(mode="after")
-    def validate_path_or_volume(self) -> "VolumePathModel":
+    def validate_path_or_volume(self) -> Self:
         if not self.volume and not self.path:
             raise ValueError("Either 'volume' or 'path' must be provided")
         return self
@@ -1853,6 +1906,7 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
     headers: dict[str, AnyVariable] = Field(default_factory=dict)
     args: list[str] = Field(default_factory=list)
     # MCP-specific fields
+    app: Optional[DatabricksAppModel] = None
     connection: Optional[ConnectionModel] = None
     functions: Optional[SchemaModel] = None
     genie_room: Optional[GenieRoomModel] = None
@@ -1940,6 +1994,7 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
 
         Returns the URL based on the configured source:
         - If url is set, returns it directly
+        - If app is set, retrieves URL from Databricks App via workspace client
         - If connection is set, constructs URL from connection
         - If genie_room is set, constructs Genie MCP URL
         - If sql is set, constructs DBSQL MCP URL (serverless)
@@ -1952,6 +2007,7 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
         - Vector Search: https://{host}/api/2.0/mcp/vector-search/{catalog}/{schema}
         - UC Functions: https://{host}/api/2.0/mcp/functions/{catalog}/{schema}
         - Connection: https://{host}/api/2.0/mcp/external/{connection_name}
+        - Databricks App: Retrieved dynamically from workspace
         """
         # Direct URL provided
         if self.url:
@@ -1974,6 +2030,10 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
         if self.sql:
             return f"{workspace_host}/api/2.0/mcp/sql"
 
+        # Databricks App
+        if self.app:
+            return self.app.url
+
         # Vector Search
         if self.vector_search:
             if (
@@ -1983,33 +2043,35 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
                 raise ValueError(
                     "vector_search must have an index with a schema (catalog/schema) configured"
                 )
-            catalog: str = self.vector_search.index.schema_model.catalog_name
-            schema: str = self.vector_search.index.schema_model.schema_name
+            catalog: str = value_of(self.vector_search.index.schema_model.catalog_name)
+            schema: str = value_of(self.vector_search.index.schema_model.schema_name)
             return f"{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}"
 
         # UC Functions MCP server
         if self.functions:
-            catalog: str = self.functions.catalog_name
-            schema: str = self.functions.schema_name
+            catalog: str = value_of(self.functions.catalog_name)
+            schema: str = value_of(self.functions.schema_name)
             return f"{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}"
 
         raise ValueError(
-            "No URL source configured. Provide one of: url, connection, genie_room, "
+            "No URL source configured. Provide one of: url, app, connection, genie_room, "
             "sql, vector_search, or functions"
         )
 
     @field_serializer("transport")
-    def serialize_transport(self, value) -> str:
+    def serialize_transport(self, value: TransportType) -> str:
+        """Serialize transport enum to string."""
         if isinstance(value, TransportType):
             return value.value
         return str(value)
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self) -> "McpFunctionModel":
+    def validate_mutually_exclusive(self) -> Self:
         """Validate that exactly one URL source is provided."""
         # Count how many URL sources are provided
         url_sources: list[tuple[str, Any]] = [
             ("url", self.url),
+            ("app", self.app),
             ("connection", self.connection),
             ("genie_room", self.genie_room),
             ("sql", self.sql),
@@ -2025,13 +2087,13 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
             if len(provided_sources) == 0:
                 raise ValueError(
                     "For STREAMABLE_HTTP transport, exactly one of the following must be provided: "
-                    "url, connection, genie_room, sql, vector_search, or functions"
+                    "url, app, connection, genie_room, sql, vector_search, or functions"
                 )
             if len(provided_sources) > 1:
                 raise ValueError(
                     f"For STREAMABLE_HTTP transport, only one URL source can be provided. "
                     f"Found: {', '.join(provided_sources)}. "
-                    f"Please provide only one of: url, connection, genie_room, sql, vector_search, or functions"
+                    f"Please provide only one of: url, app, connection, genie_room, sql, vector_search, or functions"
                 )
 
         if self.transport == TransportType.STDIO:
@@ -2043,18 +2105,25 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
         return self
 
     @model_validator(mode="after")
-    def update_url(self) -> "McpFunctionModel":
-        self.url = value_of(self.url)
+    def update_url(self) -> Self:
+        """Resolve AnyVariable to concrete value for URL."""
+        if self.url is not None:
+            resolved_value: Any = value_of(self.url)
+            # Cast to string since URL must be a string
+            self.url = str(resolved_value) if resolved_value else None
         return self
 
     @model_validator(mode="after")
-    def update_headers(self) -> "McpFunctionModel":
+    def update_headers(self) -> Self:
+        """Resolve AnyVariable to concrete values for headers."""
         for key, value in self.headers.items():
-            self.headers[key] = value_of(value)
+            resolved_value: Any = value_of(value)
+            # Headers must be strings
+            self.headers[key] = str(resolved_value) if resolved_value else ""
         return self
 
     @model_validator(mode="after")
-    def validate_tool_filters(self) -> "McpFunctionModel":
+    def validate_tool_filters(self) -> Self:
         """Validate tool filter configuration."""
         from loguru import logger
 

{dao_ai-0.1.6 → dao_ai-0.1.8}/src/dao_ai/tools/mcp.py

@@ -26,9 +26,9 @@ from loguru import logger
 from mcp.types import CallToolResult, TextContent, Tool
 
 from dao_ai.config import (
+    IsDatabricksResource,
     McpFunctionModel,
     TransportType,
-    value_of,
 )
 
 
@@ -143,12 +143,54 @@ def _should_include_tool(
     return True
 
 
+def _get_auth_resource(function: McpFunctionModel) -> IsDatabricksResource:
+    """
+    Get the IsDatabricksResource to use for authentication.
+
+    Follows a priority hierarchy:
+    1. Explicit resource with auth (app, connection, genie_room, vector_search, functions)
+    2. McpFunctionModel itself (which also inherits from IsDatabricksResource)
+
+    Returns the resource whose workspace_client should be used for authentication.
+    """
+    # Check each possible resource source in priority order
+    # These resources may have their own auth configured
+    if function.app:
+        return function.app
+    if function.connection:
+        return function.connection
+    if function.genie_room:
+        return function.genie_room
+    if function.vector_search:
+        return function.vector_search
+    if function.functions:
+        # SchemaModel doesn't have auth - fall through to McpFunctionModel
+        pass
+
+    # Fall back to McpFunctionModel itself (it inherits from IsDatabricksResource)
+    return function
+
+
 def _build_connection_config(
     function: McpFunctionModel,
 ) -> dict[str, Any]:
     """
     Build the connection configuration dictionary for MultiServerMCPClient.
 
+    Authentication Strategy:
+    -----------------------
+    For HTTP transport, authentication is handled consistently using
+    DatabricksOAuthClientProvider with the workspace_client from the appropriate
+    IsDatabricksResource. The auth resource is selected in this priority:
+
+    1. Nested resource (app, connection, genie_room, vector_search) if it has auth
+    2. McpFunctionModel itself (inherits from IsDatabricksResource)
+
+    This approach ensures:
+    - Consistent auth handling across all MCP sources
+    - Automatic token refresh for long-running connections
+    - Support for OBO, service principal, PAT, and ambient auth
+
     Args:
         function: The MCP function model configuration.
 
@@ -162,52 +204,30 @@ def _build_connection_config(
             "transport": function.transport.value,
         }
 
-    # For HTTP transport with UC Connection, use DatabricksOAuthClientProvider
-    if function.connection:
-        from databricks_mcp import DatabricksOAuthClientProvider
+    # For HTTP transport, use DatabricksOAuthClientProvider with unified auth
+    from databricks_mcp import DatabricksOAuthClientProvider
 
-        workspace_client = function.connection.workspace_client
-        auth_provider = DatabricksOAuthClientProvider(workspace_client)
+    # Get the resource to use for authentication
+    auth_resource = _get_auth_resource(function)
 
-        logger.trace(
-            "Using DatabricksOAuthClientProvider for authentication",
-            connection_name=function.connection.name,
-        )
-
-        return {
-            "url": function.mcp_url,
-            "transport": "http",
-            "auth": auth_provider,
-        }
-
-    # For HTTP transport with headers-based authentication
-    headers: dict[str, str] = {
-        key: str(value_of(val)) for key, val in function.headers.items()
-    }
+    # Get workspace client from the auth resource
+    workspace_client = auth_resource.workspace_client
+    auth_provider = DatabricksOAuthClientProvider(workspace_client)
 
-    if "Authorization" not in headers:
-        logger.trace("Generating fresh authentication token")
-
-        from dao_ai.providers.databricks import DatabricksProvider
-
-        try:
-            provider = DatabricksProvider(
-                workspace_host=value_of(function.workspace_host),
-                client_id=value_of(function.client_id),
-                client_secret=value_of(function.client_secret),
-                pat=value_of(function.pat),
-            )
-            headers["Authorization"] = f"Bearer {provider.create_token()}"
-            logger.trace("Generated fresh authentication token")
-        except Exception as e:
-            logger.error("Failed to create fresh token", error=str(e))
-    else:
-        logger.trace("Using existing authentication token")
+    # Log which resource is providing auth
+    resource_name = (
+        getattr(auth_resource, "name", None) or auth_resource.__class__.__name__
+    )
+    logger.trace(
+        "Using DatabricksOAuthClientProvider for authentication",
+        auth_resource=resource_name,
+        resource_type=auth_resource.__class__.__name__,
+    )
 
     return {
         "url": function.mcp_url,
         "transport": "http",
-        "headers": headers,
+        "auth": auth_provider,
     }
 
 

{dao_ai-0.1.6 → dao_ai-0.1.8}/tests/dao_ai/test_config.py

@@ -1,5 +1,4 @@
 import sys
-from unittest.mock import patch
 
 import pytest
 import yaml
@@ -282,21 +281,28 @@ def test_mcp_function_model_partial_oauth_credentials() -> None:
 
 @pytest.mark.unit
 def test_mcp_function_model_existing_authorization_header() -> None:
-    """Test that existing Authorization header is preserved and authentication is skipped."""
-    with patch("dao_ai.providers.databricks.DatabricksProvider") as mock_provider_class:
-        mcp_function = McpFunctionModel(
-            transport=TransportType.STREAMABLE_HTTP,
-            url="https://example.com",
-            headers={"Authorization": "Bearer existing_token"},
-            pat="test_pat",
-            workspace_host="https://test-workspace.cloud.databricks.com",
-        )
+    """Test that Authorization header can be stored in headers dict.
+
+    Note: With unified auth (via DatabricksOAuthClientProvider), the Authorization
+    header in the headers dict is not used for authentication. Instead, the
+    workspace_client from IsDatabricksResource provides authentication.
+
+    This test verifies that the model can store headers, but authentication
+    happens through the OAuth provider at invocation time.
+    """
+    mcp_function = McpFunctionModel(
+        transport=TransportType.STREAMABLE_HTTP,
+        url="https://example.com",
+        headers={"Authorization": "Bearer existing_token"},
+        pat="test_pat",
+        workspace_host="https://test-workspace.cloud.databricks.com",
+    )
 
-        # DatabricksProvider should not be called since Authorization header exists
-        mock_provider_class.assert_not_called()
+    # Header is stored in the model
+    assert mcp_function.headers["Authorization"] == "Bearer existing_token"
 
-        # Existing header should be preserved
-        assert mcp_function.headers["Authorization"] == "Bearer existing_token"
+    # But authentication will happen via workspace_client at invocation time
+    # The stored header won't be used by _build_connection_config
 
 
 @pytest.mark.unit

{dao_ai-0.1.6 → dao_ai-0.1.8}/tests/dao_ai/test_function_parsing.py

@@ -182,7 +182,7 @@ class TestFunctionModelParsing:
 
         with pytest.raises(
             ValidationError,
-            match="exactly one of the following must be provided: url, connection, genie_room, sql, vector_search, or functions",
+            match="url, app, connection, genie_room, sql, vector_search, or functions",
         ):
             ToolModel(**yaml_data)
 

{dao_ai-0.1.6 → dao_ai-0.1.8}/tests/dao_ai/test_genie.py

@@ -1770,6 +1770,7 @@ def test_create_genie_tool_with_cache_parameters() -> None:
 
         # Create mock warehouse with IsDatabricksResource attrs
         mock_warehouse = Mock(spec=WarehouseModel)
+        mock_warehouse.name = "Test Warehouse"
         mock_warehouse.warehouse_id = "test-warehouse"
         mock_warehouse.workspace_client = Mock()
         add_databricks_resource_attrs(mock_warehouse)