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

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

@@ -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.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=rUm2wg0TPfj6YwzSoNxy6rgHi6GKWxXIRJ3NgGOjB04,123037
+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.6.dist-info/METADATA,sha256=eux_1l0ANLlbQnRnq0IRgc8Q-ksY6sFmQa1_Vzd5_Kc,16685
-dao_ai-0.1.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.1.6.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.1.6.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.1.6.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,,