dao-ai 0.1.7__py3-none-any.whl → 0.1.8__py3-none-any.whl

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/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.dist-info → dao_ai-0.1.8.dist-info}/METADATA

@@ -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.dist-info → dao_ai-0.1.8.dist-info}/RECORD

@@ -2,7 +2,7 @@ dao_ai/__init__.py,sha256=18P98ExEgUaJ1Byw440Ct1ty59v6nxyWtc5S6Uq2m9Q,1062
 dao_ai/agent_as_code.py,sha256=xIlLDpPVfmDVzLvbdY_V_CrC4Jvj2ItCWJ-NzdrszTo,538
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
 dao_ai/cli.py,sha256=7LGrVDRgSBpznr8c8EksAhzPW_8NJ9h4St3DSpx-0z4,48196
-dao_ai/config.py,sha256=GY-n2PtPg4pYtO46KYpHFIGkNkCDVoPA5S7sKLjWpVc,124699
+dao_ai/config.py,sha256=OSAsqb2Rxn3ghnM5Nq7-wh13DizHzWI_6cxuRuT4_j8,125773
 dao_ai/graph.py,sha256=1-uQlo7iXZQTT3uU8aYu0N5rnhw5_g_2YLwVsAs6M-U,1119
 dao_ai/logging.py,sha256=lYy4BmucCHvwW7aI3YQkQXKJtMvtTnPDu9Hnd7_O4oc,1556
 dao_ai/messages.py,sha256=4ZBzO4iFdktGSLrmhHzFjzMIt2tpaL-aQLHOQJysGnY,6959
@@ -55,7 +55,7 @@ dao_ai/tools/agent.py,sha256=plIWALywRjaDSnot13nYehBsrHRpBUpsVZakoGeajOE,1858
 dao_ai/tools/core.py,sha256=bRIN3BZhRQX8-Kpu3HPomliodyskCqjxynQmYbk6Vjs,3783
 dao_ai/tools/email.py,sha256=A3TsCoQgJR7UUWR0g45OPRGDpVoYwctFs1MOZMTt_d4,7389
 dao_ai/tools/genie.py,sha256=4e_5MeAe7kDzHbYeXuNPFbY5z8ci3ouj8l5254CZ2lA,8874
-dao_ai/tools/mcp.py,sha256=0OfP4b4skcjeF2rzkOLYqd65ti1Mj55N_l8VoQlH9qo,17818
+dao_ai/tools/mcp.py,sha256=ZNalYo2atZECatZjMT8w4mHEsaUZJQ_fsCjia7px1nc,18689
 dao_ai/tools/memory.py,sha256=lwObKimAand22Nq3Y63tsv-AXQ5SXUigN9PqRjoWKes,1836
 dao_ai/tools/python.py,sha256=jWFnZPni2sCdtd8D1CqXnZIPHnWkdK27bCJnBXpzhvo,1879
 dao_ai/tools/search.py,sha256=cJ3D9FKr1GAR6xz55dLtRkjtQsI0WRueGt9TPDFpOxc,433
@@ -64,8 +64,8 @@ dao_ai/tools/sql.py,sha256=tKd1gjpLuKdQDyfmyYYtMiNRHDW6MGRbdEVaeqyB8Ok,7632
 dao_ai/tools/time.py,sha256=tufJniwivq29y0LIffbgeBTIDE6VgrLpmVf8Qr90qjw,9224
 dao_ai/tools/unity_catalog.py,sha256=AjQfW7bvV8NurqDLIyntYRv2eJuTwNdbvex1L5CRjOk,15534
 dao_ai/tools/vector_search.py,sha256=oe2uBwl2TfeJIXPpwiS6Rmz7wcHczSxNyqS9P3hE6co,14542
-dao_ai-0.1.7.dist-info/METADATA,sha256=jzaENv6Ic9S-uds3qTC4Pu7chwzYG8y0zvTBni4VCW8,16685
-dao_ai-0.1.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.1.7.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.1.7.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.1.7.dist-info/RECORD,,
+dao_ai-0.1.8.dist-info/METADATA,sha256=kbbCJVZAI-U3czxwWr9z36m14PQk8poQdzOB_6RRLII,16685
+dao_ai-0.1.8.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.1.8.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.1.8.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.1.8.dist-info/RECORD,,