dao-ai 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

dao_ai/tools/unity_catalog.py

@@ -1,17 +1,17 @@
-from typing import Any, Dict, Optional, Sequence, Union
+from typing import Any, Dict, Optional, Sequence, Set
 
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.service.catalog import PermissionsChange, Privilege
+from databricks.sdk.service.catalog import FunctionInfo, PermissionsChange, Privilege
 from databricks_langchain import DatabricksFunctionClient, UCFunctionToolkit
 from langchain_core.runnables.base import RunnableLike
 from langchain_core.tools import StructuredTool
 from loguru import logger
+from pydantic import BaseModel
 from unitycatalog.ai.core.base import FunctionExecutionResult
 
 from dao_ai.config import (
     AnyVariable,
     CompositeVariableModel,
-    ToolModel,
     UnityCatalogFunctionModel,
     value_of,
 )
@@ -34,36 +34,39 @@ def create_uc_tools(
     Returns:
         A sequence of BaseTool objects that wrap the specified UC functions
     """
+    original_function_model: UnityCatalogFunctionModel | None = None
+    workspace_client: WorkspaceClient | None = None
+    function_name: str
 
-    logger.debug(f"create_uc_tools: {function}")
-
-    original_function_model = None
     if isinstance(function, UnityCatalogFunctionModel):
         original_function_model = function
-        function_name = function.full_name
+        function_name = function.resource.full_name
+        workspace_client = function.resource.workspace_client
     else:
         function_name = function
 
+    logger.trace("Creating UC tools", function_name=function_name)
+
     # Determine which tools to create
+    tools: list[RunnableLike]
     if original_function_model and original_function_model.partial_args:
-        logger.debug("Found partial_args, creating custom tool with partial arguments")
-        # Create a ToolModel wrapper for the with_partial_args function
-        tool_model = ToolModel(
-            name=original_function_model.name, function=original_function_model
+        logger.debug(
+            "Creating custom tool with partial arguments", function_name=function_name
         )
-
-        # Use with_partial_args to create the authenticated tool
-        tools = [with_partial_args(tool_model, original_function_model.partial_args)]
+        # Use with_partial_args directly with UnityCatalogFunctionModel
+        tools = [with_partial_args(original_function_model)]
     else:
         # Fallback to standard UC toolkit approach
-        client: DatabricksFunctionClient = DatabricksFunctionClient()
+        client: DatabricksFunctionClient = DatabricksFunctionClient(
+            client=workspace_client
+        )
 
         toolkit: UCFunctionToolkit = UCFunctionToolkit(
             function_names=[function_name], client=client
         )
 
         tools = toolkit.tools or []
-        logger.debug(f"Retrieved tools: {tools}")
+        logger.trace("Retrieved tools", tools_count=len(tools))
 
     # HITL is now handled at middleware level via HumanInTheLoopMiddleware
     return list(tools)
@@ -72,7 +75,7 @@ def create_uc_tools(
 def _execute_uc_function(
     client: DatabricksFunctionClient,
     function_name: str,
-    partial_args: Dict[str, str] = None,
+    partial_args: Optional[Dict[str, str]] = None,
     **kwargs: Any,
 ) -> str:
     """Execute Unity Catalog function with partial args and provided parameters."""
@@ -84,7 +87,9 @@ def _execute_uc_function(
     all_params.update(kwargs)
 
     logger.debug(
-        f"Calling UC function {function_name} with parameters: {list(all_params.keys())}"
+        "Calling UC function",
+        function_name=function_name,
+        parameters=list(all_params.keys()),
     )
 
     result: FunctionExecutionResult = client.execute_function(
@@ -93,11 +98,19 @@ def _execute_uc_function(
 
     # Handle errors and extract result
     if result.error:
-        logger.error(f"Unity Catalog function error: {result.error}")
+        logger.error(
+            "Unity Catalog function error",
+            function_name=function_name,
+            error=result.error,
+        )
         raise RuntimeError(f"Function execution failed: {result.error}")
 
     result_value: str = result.value if result.value is not None else str(result)
-    logger.debug(f"UC function result: {result_value}")
+    logger.trace(
+        "UC function result",
+        function_name=function_name,
+        result_length=len(str(result_value)),
+    )
     return result_value
 
 
@@ -116,21 +129,29 @@ def _grant_function_permissions(
     """
     try:
         # Initialize workspace client
-        workspace_client = WorkspaceClient(host=host) if host else WorkspaceClient()
+        workspace_client: WorkspaceClient = (
+            WorkspaceClient(host=host) if host else WorkspaceClient()
+        )
 
         # Parse the function name to get catalog and schema
-        parts = function_name.split(".")
+        parts: list[str] = function_name.split(".")
         if len(parts) != 3:
             logger.warning(
-                f"Invalid function name format: {function_name}. Expected catalog.schema.function"
+                "Invalid function name format, expected catalog.schema.function",
+                function_name=function_name,
             )
             return
 
+        catalog_name: str
+        schema_name: str
+        func_name: str
         catalog_name, schema_name, func_name = parts
-        schema_full_name = f"{catalog_name}.{schema_name}"
+        schema_full_name: str = f"{catalog_name}.{schema_name}"
 
         logger.debug(
-            f"Granting comprehensive permissions on function {function_name} to principal {client_id}"
+            "Granting comprehensive permissions",
+            function_name=function_name,
+            principal=client_id,
         )
 
         # 1. Grant EXECUTE permission on the function
@@ -142,9 +163,13 @@ def _grant_function_permissions(
                     PermissionsChange(principal=client_id, add=[Privilege.EXECUTE])
                 ],
             )
-            logger.debug(f"Granted EXECUTE on function {function_name}")
+            logger.trace("Granted EXECUTE permission", function_name=function_name)
         except Exception as e:
-            logger.warning(f"Failed to grant EXECUTE on function {function_name}: {e}")
+            logger.warning(
+                "Failed to grant EXECUTE permission",
+                function_name=function_name,
+                error=str(e),
+            )
 
         # 2. Grant USE_SCHEMA permission on the schema
         try:
@@ -158,10 +183,12 @@ def _grant_function_permissions(
                     )
                 ],
             )
-            logger.debug(f"Granted USE_SCHEMA on schema {schema_full_name}")
+            logger.trace("Granted USE_SCHEMA permission", schema=schema_full_name)
         except Exception as e:
             logger.warning(
-                f"Failed to grant USE_SCHEMA on schema {schema_full_name}: {e}"
+                "Failed to grant USE_SCHEMA permission",
+                schema=schema_full_name,
+                error=str(e),
             )
 
         # 3. Grant USE_CATALOG and BROWSE permissions on the catalog
@@ -176,25 +203,34 @@ def _grant_function_permissions(
                     )
                 ],
             )
-            logger.debug(f"Granted USE_CATALOG and BROWSE on catalog {catalog_name}")
+            logger.trace(
+                "Granted USE_CATALOG and BROWSE permissions", catalog=catalog_name
+            )
         except Exception as e:
             logger.warning(
-                f"Failed to grant catalog permissions on {catalog_name}: {e}"
+                "Failed to grant catalog permissions",
+                catalog=catalog_name,
+                error=str(e),
             )
 
         logger.debug(
-            f"Successfully granted comprehensive permissions on {function_name} to {client_id}"
+            "Successfully granted comprehensive permissions",
+            function_name=function_name,
+            principal=client_id,
         )
 
     except Exception as e:
         logger.warning(
-            f"Failed to grant permissions on function {function_name} to {client_id}: {e}"
+            "Failed to grant permissions",
+            function_name=function_name,
+            principal=client_id,
+            error=str(e),
         )
         # Don't fail the tool creation if permission granting fails
         pass
 
 
-def _create_filtered_schema(original_schema: type, exclude_fields: set[str]) -> type:
+def _create_filtered_schema(original_schema: type, exclude_fields: Set[str]) -> type:
     """
     Create a new Pydantic model that excludes specified fields from the original schema.
 
@@ -205,23 +241,27 @@ def _create_filtered_schema(original_schema: type, exclude_fields: set[str]) ->
     Returns:
         A new Pydantic model class with the specified fields removed
     """
-    from pydantic import BaseModel, Field, create_model
-    from pydantic.fields import PydanticUndefined
+    from pydantic import Field, create_model
+    from pydantic.fields import FieldInfo, PydanticUndefined
 
     try:
         # Get the original model's fields (Pydantic v2)
-        original_fields = original_schema.model_fields
-        filtered_field_definitions = {}
+        original_fields: dict[str, FieldInfo] = original_schema.model_fields
+        filtered_field_definitions: dict[str, tuple[type, FieldInfo]] = {}
 
-        for name, field in original_fields.items():
-            if name not in exclude_fields:
+        field_name: str
+        field: FieldInfo
+        for field_name, field in original_fields.items():
+            if field_name not in exclude_fields:
                 # Reconstruct the field definition for create_model
-                field_type = field.annotation
-                field_default = (
+                field_type: type = field.annotation
+                field_default: Any = (
                     field.default if field.default is not PydanticUndefined else ...
                 )
-                field_info = Field(default=field_default, description=field.description)
-                filtered_field_definitions[name] = (field_type, field_info)
+                field_info: FieldInfo = Field(
+                    default=field_default, description=field.description
+                )
+                filtered_field_definitions[field_name] = (field_type, field_info)
 
         # If no fields remain after filtering, return a generic empty schema
         if not filtered_field_definitions:
@@ -234,18 +274,18 @@ def _create_filtered_schema(original_schema: type, exclude_fields: set[str]) ->
             return EmptySchema
 
         # Create the new model dynamically
-        model_name = f"Filtered{original_schema.__name__}"
-        docstring = getattr(
+        model_name: str = f"Filtered{original_schema.__name__}"
+        docstring: str = getattr(
             original_schema, "__doc__", "Filtered Unity Catalog function parameters."
         )
 
-        filtered_model = create_model(
+        filtered_model: type[BaseModel] = create_model(
             model_name, __doc__=docstring, **filtered_field_definitions
         )
         return filtered_model
 
     except Exception as e:
-        logger.warning(f"Failed to create filtered schema: {e}")
+        logger.warning("Failed to create filtered schema", error=str(e))
 
         # Fallback to generic schema
         class GenericFilteredSchema(BaseModel):
@@ -257,8 +297,7 @@ def _create_filtered_schema(original_schema: type, exclude_fields: set[str]) ->
 
 
 def with_partial_args(
-    tool: Union[ToolModel, Dict[str, Any]],
-    partial_args: dict[str, AnyVariable] = {},
+    uc_function: UnityCatalogFunctionModel,
 ) -> StructuredTool:
     """
     Create a Unity Catalog tool with partial arguments pre-filled.
@@ -267,12 +306,8 @@ def with_partial_args(
     already resolved, so the caller only needs to provide the remaining parameters.
 
     Args:
-        tool: ToolModel containing the Unity Catalog function configuration
-        partial_args: Dictionary of arguments to pre-fill in the tool.
-            Supports:
-            - client_id, client_secret: OAuth credentials directly
-            - service_principal: ServicePrincipalModel with client_id and client_secret
-            - host or workspace_host: Databricks workspace host
+        uc_function: UnityCatalogFunctionModel containing the function configuration
+            and partial_args to pre-fill.
 
     Returns:
         StructuredTool: A LangChain tool with partial arguments pre-filled
@@ -281,10 +316,12 @@ def with_partial_args(
 
     from dao_ai.config import ServicePrincipalModel
 
-    logger.debug(f"with_partial_args: {tool}")
+    partial_args: dict[str, AnyVariable] = uc_function.partial_args or {}
 
     # Convert dict-based variables to CompositeVariableModel and resolve their values
     resolved_args: dict[str, Any] = {}
+    k: str
+    v: AnyVariable
     for k, v in partial_args.items():
         if isinstance(v, dict):
             resolved_args[k] = value_of(CompositeVariableModel(**v))
@@ -293,7 +330,7 @@ def with_partial_args(
 
     # Handle service_principal - expand into client_id and client_secret
     if "service_principal" in resolved_args:
-        sp = resolved_args.pop("service_principal")
+        sp: Any = resolved_args.pop("service_principal")
         if isinstance(sp, dict):
             sp = ServicePrincipalModel(**sp)
         if isinstance(sp, ServicePrincipalModel):
@@ -316,17 +353,17 @@ def with_partial_args(
         if host:
             resolved_args["host"] = host
 
-    logger.debug(f"Resolved partial args: {resolved_args.keys()}")
+    # Get function info from the resource
+    function_name: str = uc_function.resource.full_name
+    tool_name: str = uc_function.resource.name or function_name.replace(".", "_")
+    workspace_client: WorkspaceClient = uc_function.resource.workspace_client
 
-    if isinstance(tool, dict):
-        tool = ToolModel(**tool)
-
-    unity_catalog_function = tool.function
-    if isinstance(unity_catalog_function, dict):
-        unity_catalog_function = UnityCatalogFunctionModel(**unity_catalog_function)
-
-    function_name: str = unity_catalog_function.full_name
-    logger.debug(f"Creating UC tool with partial args for: {function_name}")
+    logger.debug(
+        "Creating UC tool with partial args",
+        function_name=function_name,
+        tool_name=tool_name,
+        partial_args=list(resolved_args.keys()),
+    )
 
     # Grant permissions if we have credentials
     if "client_id" in resolved_args:
@@ -335,36 +372,44 @@ def with_partial_args(
         try:
             _grant_function_permissions(function_name, client_id, host)
         except Exception as e:
-            logger.warning(f"Failed to grant permissions: {e}")
+            logger.warning(
+                "Failed to grant permissions", function_name=function_name, error=str(e)
+            )
 
-    # Create the client for function execution
-    client: DatabricksFunctionClient = DatabricksFunctionClient()
+    # Create the client for function execution using the resource's workspace client
+    client: DatabricksFunctionClient = DatabricksFunctionClient(client=workspace_client)
 
     # Try to get the function schema for better tool definition
+    schema_model: type[BaseModel]
+    tool_description: str
     try:
-        function_info = client.get_function(function_name)
+        function_info: FunctionInfo = client.get_function(function_name)
         schema_info = generate_function_input_params_schema(function_info)
         tool_description = (
             function_info.comment or f"Unity Catalog function: {function_name}"
         )
 
-        logger.debug(
-            f"Generated schema for function {function_name}: {schema_info.pydantic_model}"
+        logger.trace(
+            "Generated function schema",
+            function_name=function_name,
+            schema=schema_info.pydantic_model.__name__,
         )
-        logger.debug(f"Tool description: {tool_description}")
 
         # Create a modified schema that excludes partial args
-        original_schema = schema_info.pydantic_model
+        original_schema: type = schema_info.pydantic_model
         schema_model = _create_filtered_schema(original_schema, resolved_args.keys())
-        logger.debug(
-            f"Filtered schema excludes partial args: {list(resolved_args.keys())}"
+        logger.trace(
+            "Filtered schema to exclude partial args",
+            function_name=function_name,
+            excluded_args=list(resolved_args.keys()),
         )
 
     except Exception as e:
-        logger.warning(f"Could not introspect function {function_name}: {e}")
-        # Fallback to a generic schema
-        from pydantic import BaseModel
+        logger.warning(
+            "Could not introspect function", function_name=function_name, error=str(e)
+        )
 
+        # Fallback to a generic schema
         class GenericUCParams(BaseModel):
             """Generic parameters for Unity Catalog function."""
 
@@ -384,14 +429,12 @@ def with_partial_args(
         )
 
     # Set the function name for the decorator
-    uc_function_wrapper.__name__ = tool.name or function_name.replace(".", "_")
+    uc_function_wrapper.__name__ = tool_name
 
     # Create the tool using LangChain's StructuredTool
-    from langchain_core.tools import StructuredTool
-
-    partial_tool = StructuredTool.from_function(
+    partial_tool: StructuredTool = StructuredTool.from_function(
         func=uc_function_wrapper,
-        name=tool.name or function_name.replace(".", "_"),
+        name=tool_name,
         description=tool_description,
         args_schema=schema_model,
     )