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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.7
+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.7 → dao_ai-0.1.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dao-ai"
-version = "0.1.7"
+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.7 → 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.7 → 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.7 → dao_ai-0.1.8}/src/dao_ai/config.py

@@ -418,11 +418,11 @@ 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 
+    The `url` is dynamically retrieved from the workspace client by calling
     `apps.get(name)` and returning the app's URL.
-    
+
     Example:
         ```yaml
         resources:
@@ -431,7 +431,7 @@ class DatabricksAppModel(IsDatabricksResource, HasFullName):
               name: my-databricks-app
         ```
     """
-    
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
     """The unique instance name of the Databricks App in the workspace."""
@@ -440,10 +440,10 @@ class DatabricksAppModel(IsDatabricksResource, HasFullName):
     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.
         """
@@ -455,7 +455,6 @@ class DatabricksAppModel(IsDatabricksResource, HasFullName):
             )
         return app.url
 
-
     @property
     def full_name(self) -> str:
         return self.name
@@ -761,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 [
@@ -786,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
 
@@ -998,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
 
 
@@ -2007,7 +2029,7 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
         # DBSQL MCP server (serverless, workspace-level)
         if self.sql:
             return f"{workspace_host}/api/2.0/mcp/sql"
-        
+
         # Databricks App
         if self.app:
             return self.app.url

{dao_ai-0.1.7 → 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.7 → 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.7 → 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)

{dao_ai-0.1.7 → dao_ai-0.1.8}/tests/dao_ai/test_genie_room_model.py

@@ -400,6 +400,62 @@ class TestGenieRoomModelSerialization:
             # The genie_space_id is stored as the 'name' attribute
             assert resources[0].name == "test-space-123"
 
+    def test_name_populated_from_space(
+        self, mock_workspace_client, mock_genie_space_with_serialized_data
+    ):
+        """Test that name is automatically populated from GenieSpace.title if not provided."""
+        # Set a title on the mock space
+        mock_genie_space_with_serialized_data.title = "My Retail Analytics Space"
+
+        with patch("dao_ai.config.WorkspaceClient", return_value=mock_workspace_client):
+            mock_workspace_client.genie.get_space.return_value = (
+                mock_genie_space_with_serialized_data
+            )
+
+            # Create without name
+            genie_room = GenieRoomModel(space_id="test-space-123")
+
+            # Name should be populated from the space title
+            assert genie_room.name == "My Retail Analytics Space"
+
+    def test_name_not_overridden_if_provided(
+        self, mock_workspace_client, mock_genie_space_with_serialized_data
+    ):
+        """Test that provided name is not overridden by GenieSpace title."""
+        # Set a title on the mock space
+        mock_genie_space_with_serialized_data.title = "Space Title from API"
+
+        with patch("dao_ai.config.WorkspaceClient", return_value=mock_workspace_client):
+            mock_workspace_client.genie.get_space.return_value = (
+                mock_genie_space_with_serialized_data
+            )
+
+            # Create with explicit name
+            genie_room = GenieRoomModel(
+                name="My Custom Name", space_id="test-space-123"
+            )
+
+            # Custom name should be preserved
+            assert genie_room.name == "My Custom Name"
+
+    def test_name_handles_none_from_space(
+        self, mock_workspace_client, mock_genie_space_with_serialized_data
+    ):
+        """Test that None title from space is handled gracefully."""
+        # Set title to None
+        mock_genie_space_with_serialized_data.title = None
+
+        with patch("dao_ai.config.WorkspaceClient", return_value=mock_workspace_client):
+            mock_workspace_client.genie.get_space.return_value = (
+                mock_genie_space_with_serialized_data
+            )
+
+            # Create without name
+            genie_room = GenieRoomModel(space_id="test-space-123")
+
+            # Name should remain None
+            assert genie_room.name is None
+
     def test_description_populated_from_space(
         self, mock_workspace_client, mock_genie_space_with_serialized_data
     ):
@@ -415,13 +471,33 @@ class TestGenieRoomModelSerialization:
             )
 
             # Create without description
-            genie_room = GenieRoomModel(
-                name="test-genie-room", space_id="test-space-123"
-            )
+            genie_room = GenieRoomModel(space_id="test-space-123")
 
             # Description should be populated from the space
             assert genie_room.description == "This is a test Genie space description"
 
+    def test_name_and_description_populated_together(
+        self, mock_workspace_client, mock_genie_space_with_serialized_data
+    ):
+        """Test that both name and description are populated from GenieSpace when not provided."""
+        # Set both title and description on the mock space
+        mock_genie_space_with_serialized_data.title = "Production Analytics"
+        mock_genie_space_with_serialized_data.description = (
+            "Production space for analytics queries"
+        )
+
+        with patch("dao_ai.config.WorkspaceClient", return_value=mock_workspace_client):
+            mock_workspace_client.genie.get_space.return_value = (
+                mock_genie_space_with_serialized_data
+            )
+
+            # Create without name or description
+            genie_room = GenieRoomModel(space_id="test-space-123")
+
+            # Both should be populated from the space
+            assert genie_room.name == "Production Analytics"
+            assert genie_room.description == "Production space for analytics queries"
+
     def test_description_not_overridden_if_provided(
         self, mock_workspace_client, mock_genie_space_with_serialized_data
     ):
@@ -464,6 +540,84 @@ class TestGenieRoomModelSerialization:
             # Description should remain None
             assert genie_room.description is None
 
+    def test_partial_population_name_provided_description_auto(
+        self, mock_workspace_client, mock_genie_space_with_serialized_data
+    ):
+        """Test that description is still populated when name is provided."""
+        # Set both title and description on the mock space
+        mock_genie_space_with_serialized_data.title = "API Space Title"
+        mock_genie_space_with_serialized_data.description = "API Space Description"
+
+        with patch("dao_ai.config.WorkspaceClient", return_value=mock_workspace_client):
+            mock_workspace_client.genie.get_space.return_value = (
+                mock_genie_space_with_serialized_data
+            )
+
+            # Create with name but without description
+            genie_room = GenieRoomModel(name="Custom Name", space_id="test-space-123")
+
+            # Custom name should be preserved, description should be auto-populated
+            assert genie_room.name == "Custom Name"
+            assert genie_room.description == "API Space Description"
+
+    def test_partial_population_description_provided_name_auto(
+        self, mock_workspace_client, mock_genie_space_with_serialized_data
+    ):
+        """Test that name is still populated when description is provided."""
+        # Set both title and description on the mock space
+        mock_genie_space_with_serialized_data.title = "API Space Title"
+        mock_genie_space_with_serialized_data.description = "API Space Description"
+
+        with patch("dao_ai.config.WorkspaceClient", return_value=mock_workspace_client):
+            mock_workspace_client.genie.get_space.return_value = (
+                mock_genie_space_with_serialized_data
+            )
+
+            # Create with description but without name
+            genie_room = GenieRoomModel(
+                description="Custom Description", space_id="test-space-123"
+            )
+
+            # Custom description should be preserved, name should be auto-populated
+            assert genie_room.name == "API Space Title"
+            assert genie_room.description == "Custom Description"
+
+    def test_populate_name_and_description_handles_api_error(
+        self, mock_workspace_client
+    ):
+        """Test that populate_name_and_description validator handles API errors gracefully."""
+        with patch("dao_ai.config.WorkspaceClient", return_value=mock_workspace_client):
+            # Mock API to raise an error
+            mock_workspace_client.genie.get_space.side_effect = Exception(
+                "API Error: Connection timeout"
+            )
+
+            # Create without name or description - should not raise exception
+            genie_room = GenieRoomModel(space_id="test-space-123")
+
+            # Name and description should remain None (not populated due to API error)
+            assert genie_room.name is None
+            assert genie_room.description is None
+
+    def test_populate_name_and_description_not_called_when_both_provided(
+        self, mock_workspace_client
+    ):
+        """Test that API is not called when both name and description are provided."""
+        with patch("dao_ai.config.WorkspaceClient", return_value=mock_workspace_client):
+            # Create with both name and description
+            genie_room = GenieRoomModel(
+                name="Custom Name",
+                description="Custom Description",
+                space_id="test-space-123",
+            )
+
+            # API should not have been called
+            mock_workspace_client.genie.get_space.assert_not_called()
+
+            # Values should be preserved
+            assert genie_room.name == "Custom Name"
+            assert genie_room.description == "Custom Description"
+
     def test_tables_and_functions_inherit_authentication(
         self, mock_workspace_client, mock_genie_space_with_serialized_data
     ):

{dao_ai-0.1.7 → dao_ai-0.1.8}/tests/dao_ai/test_mcp.py

@@ -6,7 +6,13 @@ import pytest
 from langchain_core.tools import BaseTool
 from pydantic import ValidationError
 
-from dao_ai.config import ConnectionModel, DatabricksAppModel, McpFunctionModel, SchemaModel, TransportType
+from dao_ai.config import (
+    ConnectionModel,
+    DatabricksAppModel,
+    McpFunctionModel,
+    SchemaModel,
+    TransportType,
+)
 
 
 @pytest.mark.integration