dao-ai 0.0.18__py3-none-any.whl → 0.0.20__py3-none-any.whl

dao_ai/config.py

@@ -39,6 +39,7 @@ from mlflow.models import ModelConfig
 from mlflow.models.resources import (
     DatabricksFunction,
     DatabricksGenieSpace,
+    DatabricksLakebase,
     DatabricksResource,
     DatabricksServingEndpoint,
     DatabricksSQLWarehouse,
@@ -274,10 +275,8 @@ class TableModel(BaseModel, HasFullName, IsDatabricksResource):
             "_assessment_logs",
             "_request_logs",
         ]
-        
-        excluded_prefixes: Sequence[str] = [
-            "trace_logs_"
-        ]
+
+        excluded_prefixes: Sequence[str] = ["trace_logs_"]
 
         if self.name:
             resources.append(
@@ -299,8 +298,12 @@ class TableModel(BaseModel, HasFullName, IsDatabricksResource):
                         on_behalf_of_user=self.on_behalf_of_user,
                     )
                     for table in tables
-                    if not any(table.name.endswith(suffix) for suffix in excluded_suffixes)
-                    and not any(table.name.startswith(prefix) for prefix in excluded_prefixes)
+                    if not any(
+                        table.name.endswith(suffix) for suffix in excluded_suffixes
+                    )
+                    and not any(
+                        table.name.startswith(prefix) for prefix in excluded_prefixes
+                    )
                 ]
             )
 
@@ -413,7 +416,7 @@ class GenieRoomModel(BaseModel, IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
     description: Optional[str] = None
-    space_id: str
+    space_id: AnyVariable
 
     @property
     def api_scopes(self) -> Sequence[str]:
@@ -428,6 +431,11 @@ class GenieRoomModel(BaseModel, IsDatabricksResource):
             )
         ]
 
+    @model_validator(mode="after")
+    def update_space_id(self):
+        self.space_id = value_of(self.space_id)
+        return self
+
 
 class VolumeModel(BaseModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -648,6 +656,7 @@ class ConnectionModel(BaseModel, HasFullName, IsDatabricksResource):
     def api_scopes(self) -> Sequence[str]:
         return [
             "catalog.connections",
+            "serving.serving-endpoints",
         ]
 
     def as_resources(self) -> Sequence[DatabricksResource]:
@@ -662,7 +671,7 @@ class WarehouseModel(BaseModel, IsDatabricksResource):
     model_config = ConfigDict()
     name: str
     description: Optional[str] = None
-    warehouse_id: str
+    warehouse_id: AnyVariable
 
     @property
     def api_scopes(self) -> Sequence[str]:
@@ -678,8 +687,13 @@ class WarehouseModel(BaseModel, IsDatabricksResource):
             )
         ]
 
+    @model_validator(mode="after")
+    def update_warehouse_id(self):
+        self.warehouse_id = value_of(self.warehouse_id)
+        return self
+
 
-class DatabaseModel(BaseModel):
+class DatabaseModel(BaseModel, IsDatabricksResource):
     model_config = ConfigDict(frozen=True)
     name: str
     description: Optional[str] = None
@@ -695,6 +709,18 @@ class DatabaseModel(BaseModel):
     client_secret: Optional[AnyVariable] = None
     workspace_host: Optional[AnyVariable] = None
 
+    @property
+    def api_scopes(self) -> Sequence[str]:
+        return []
+
+    def as_resources(self) -> Sequence[DatabricksResource]:
+        return [
+            DatabricksLakebase(
+                database_instance_name=self.name,
+                on_behalf_of_user=self.on_behalf_of_user,
+            )
+        ]
+
     @model_validator(mode="after")
     def validate_auth_methods(self):
         oauth_fields: Sequence[Any] = [
@@ -841,7 +867,7 @@ class PythonFunctionModel(BaseFunctionModel, HasFullName):
 
 class FactoryFunctionModel(BaseFunctionModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    args: Optional[dict[str, Any]] = Field(default_factory=dict)
+    args: Optional[dict[str, AnyVariable]] = Field(default_factory=dict)
     type: Literal[FunctionType.FACTORY] = FunctionType.FACTORY
 
     @property
@@ -937,6 +963,7 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
 class UnityCatalogFunctionModel(BaseFunctionModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
+    partial_args: Optional[dict[str, AnyVariable]] = Field(default_factory=dict)
     type: Literal[FunctionType.UNITY_CATALOG] = FunctionType.UNITY_CATALOG
 
     @property
@@ -1190,13 +1217,13 @@ class AppModel(BaseModel):
     endpoint_name: Optional[str] = None
     tags: Optional[dict[str, Any]] = Field(default_factory=dict)
     scale_to_zero: Optional[bool] = True
-    environment_vars: Optional[dict[str, Any]] = Field(default_factory=dict)
+    environment_vars: Optional[dict[str, AnyVariable]] = Field(default_factory=dict)
     budget_policy_id: Optional[str] = None
     workload_size: Optional[WorkloadSize] = "Small"
     permissions: Optional[list[AppPermissionModel]] = Field(default_factory=list)
     agents: list[AgentModel] = Field(default_factory=list)
 
-    orchestration: OrchestrationModel
+    orchestration: Optional[OrchestrationModel] = None
     alias: Optional[str] = None
     initialization_hooks: Optional[FunctionHook | list[FunctionHook]] = Field(
         default_factory=list
@@ -1215,7 +1242,38 @@ class AppModel(BaseModel):
     @model_validator(mode="after")
     def validate_agents_not_empty(self):
         if not self.agents:
-            raise ValueError("agents must contain at least one item")
+            raise ValueError("At least one agent must be specified")
+        return self
+
+    @model_validator(mode="after")
+    def update_environment_vars(self):
+        for key, value in self.environment_vars.items():
+            if isinstance(value, SecretVariableModel):
+                updated_value = str(value)
+            else:
+                updated_value = value_of(value)
+
+            self.environment_vars[key] = updated_value
+        return self
+
+    @model_validator(mode="after")
+    def set_default_orchestration(self):
+        if self.orchestration is None:
+            if len(self.agents) > 1:
+                default_agent: AgentModel = self.agents[0]
+                self.orchestration = OrchestrationModel(
+                    swarm=SupervisorModel(model=default_agent.model)
+                )
+            elif len(self.agents) == 1:
+                default_agent: AgentModel = self.agents[0]
+                self.orchestration = OrchestrationModel(
+                    supervisor=SwarmModel(
+                        model=default_agent.model, default_agent=default_agent
+                    )
+                )
+            else:
+                raise ValueError("At least one agent must be specified")
+
         return self
 
     @model_validator(mode="after")

dao_ai/graph.py

@@ -136,8 +136,8 @@ def _create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
 
     workflow: StateGraph = StateGraph(
         SharedState,
-        input=IncomingState,
-        output=OutgoingState,
+        input_schema=IncomingState,
+        output_schema=OutgoingState,
         context_schema=Context,
     )
 
@@ -200,8 +200,8 @@ def _create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
 
     workflow: StateGraph = StateGraph(
         SharedState,
-        input=IncomingState,
-        output=OutgoingState,
+        input_schema=IncomingState,
+        output_schema=OutgoingState,
         context_schema=Context,
     )
 

dao_ai/guardrails.py

@@ -87,12 +87,12 @@ def judge_node(guardrails: GuardrailModel) -> RunnableLike:
         )
 
         if eval_result["score"]:
-            logger.debug("✅ Response approved by judge")
+            logger.debug("Response approved by judge")
             logger.debug(f"Judge's comment: {eval_result['comment']}")
             return
         else:
             # Otherwise, return the judge's critique as a new user message
-            logger.warning("⚠️ Judge requested improvements")
+            logger.warning("Judge requested improvements")
             comment: str = eval_result["comment"]
             logger.warning(f"Judge's critique: {comment}")
             content: str = "\n".join([human_message.content, comment])

dao_ai/memory/core.py

@@ -99,13 +99,13 @@ class CheckpointManager:
                         checkpointer_manager
                     )
             case StorageType.POSTGRES:
-                from dao_ai.memory.postgres import PostgresCheckpointerManager
+                from dao_ai.memory.postgres import AsyncPostgresCheckpointerManager
 
                 checkpointer_manager = cls.checkpoint_managers.get(
                     checkpointer_model.database.name
                 )
                 if checkpointer_manager is None:
-                    checkpointer_manager = PostgresCheckpointerManager(
+                    checkpointer_manager = AsyncPostgresCheckpointerManager(
                         checkpointer_model
                     )
                     cls.checkpoint_managers[checkpointer_model.database.name] = (

dao_ai/memory/postgres.py

@@ -119,8 +119,9 @@ class AsyncPostgresStoreManager(StoreManagerBase):
                 self.store_model.database
             )
 
-            # Create store with the shared pool
+            # Create store with the shared pool (using patched version)
             self._store = AsyncPostgresStore(conn=self.pool)
+
             await self._store.setup()
 
             self._setup_complete = True

dao_ai/models.py

@@ -227,14 +227,28 @@ class LanggraphResponsesAgent(ResponsesAgent):
         import asyncio
 
         async def _async_invoke():
-            return await self.graph.ainvoke(
-                {"messages": messages}, context=context, config=custom_inputs
-            )
+            try:
+                return await self.graph.ainvoke(
+                    {"messages": messages}, context=context, config=custom_inputs
+                )
+            except Exception as e:
+                logger.error(f"Error in graph.ainvoke: {e}")
+                raise
 
-        loop = asyncio.get_event_loop()
-        response: dict[str, Sequence[BaseMessage]] = loop.run_until_complete(
-            _async_invoke()
-        )
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            # Handle case where no event loop exists (common in some deployment scenarios)
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        try:
+            response: dict[str, Sequence[BaseMessage]] = loop.run_until_complete(
+                _async_invoke()
+            )
+        except Exception as e:
+            logger.error(f"Error in async execution: {e}")
+            raise
 
         # Convert response to ResponsesAgent format
         last_message: BaseMessage = response["messages"][-1]
@@ -243,8 +257,9 @@ class LanggraphResponsesAgent(ResponsesAgent):
             text=last_message.content, id=f"msg_{uuid.uuid4().hex[:8]}"
         )
 
+        custom_outputs = custom_inputs
         return ResponsesAgentResponse(
-            output=[output_item], custom_outputs=request.custom_inputs
+            output=[output_item], custom_outputs=custom_outputs
         )
 
     def predict_stream(
@@ -271,46 +286,59 @@ class LanggraphResponsesAgent(ResponsesAgent):
             item_id = f"msg_{uuid.uuid4().hex[:8]}"
             accumulated_content = ""
 
-            async for nodes, stream_mode, messages_batch in self.graph.astream(
-                {"messages": messages},
-                context=context,
-                config=custom_inputs,
-                stream_mode=["messages", "custom"],
-                subgraphs=True,
-            ):
-                nodes: tuple[str, ...]
-                stream_mode: str
-                messages_batch: Sequence[BaseMessage]
-
-                for message in messages_batch:
-                    if (
-                        isinstance(
-                            message,
-                            (
-                                AIMessageChunk,
-                                AIMessage,
-                            ),
-                        )
-                        and message.content
-                        and "summarization" not in nodes
-                    ):
-                        content = message.content
-                        accumulated_content += content
-
-                        # Yield streaming delta
-                        yield ResponsesAgentStreamEvent(
-                            **self.create_text_delta(delta=content, item_id=item_id)
-                        )
-
-            # Yield final output item
-            yield ResponsesAgentStreamEvent(
-                type="response.output_item.done",
-                item=self.create_text_output_item(text=accumulated_content, id=item_id),
-                custom_outputs=request.custom_inputs,
-            )
+            try:
+                async for nodes, stream_mode, messages_batch in self.graph.astream(
+                    {"messages": messages},
+                    context=context,
+                    config=custom_inputs,
+                    stream_mode=["messages", "custom"],
+                    subgraphs=True,
+                ):
+                    nodes: tuple[str, ...]
+                    stream_mode: str
+                    messages_batch: Sequence[BaseMessage]
+
+                    for message in messages_batch:
+                        if (
+                            isinstance(
+                                message,
+                                (
+                                    AIMessageChunk,
+                                    AIMessage,
+                                ),
+                            )
+                            and message.content
+                            and "summarization" not in nodes
+                        ):
+                            content = message.content
+                            accumulated_content += content
+
+                            # Yield streaming delta
+                            yield ResponsesAgentStreamEvent(
+                                **self.create_text_delta(delta=content, item_id=item_id)
+                            )
+
+                custom_outputs = custom_inputs
+                # Yield final output item
+                yield ResponsesAgentStreamEvent(
+                    type="response.output_item.done",
+                    item=self.create_text_output_item(
+                        text=accumulated_content, id=item_id
+                    ),
+                    custom_outputs=custom_outputs,
+                )
+            except Exception as e:
+                logger.error(f"Error in graph.astream: {e}")
+                raise
 
         # Convert async generator to sync generator
-        loop = asyncio.get_event_loop()
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            # Handle case where no event loop exists (common in some deployment scenarios)
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
         async_gen = _async_stream()
 
         try:
@@ -320,8 +348,14 @@ class LanggraphResponsesAgent(ResponsesAgent):
                     yield item
                 except StopAsyncIteration:
                     break
+                except Exception as e:
+                    logger.error(f"Error in streaming: {e}")
+                    raise
         finally:
-            loop.run_until_complete(async_gen.aclose())
+            try:
+                loop.run_until_complete(async_gen.aclose())
+            except Exception as e:
+                logger.warning(f"Error closing async generator: {e}")
 
     def _extract_text_from_content(
         self,
@@ -522,7 +556,14 @@ def _process_langchain_messages_stream(
                     yield message
 
     # Convert async generator to sync generator
-    loop = asyncio.get_event_loop()
+
+    try:
+        loop = asyncio.get_event_loop()
+    except RuntimeError:
+        # Handle case where no event loop exists (common in some deployment scenarios)
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
     async_gen = _async_stream()
 
     try:

dao_ai/providers/databricks.py

@@ -42,6 +42,7 @@ import dao_ai
 from dao_ai.config import (
     AppConfig,
     ConnectionModel,
+    DatabaseModel,
     DatasetModel,
     FunctionModel,
     GenieRoomModel,
@@ -224,6 +225,7 @@ class DatabricksProvider(ServiceProvider):
         connections: Sequence[ConnectionModel] = list(
             config.resources.connections.values()
         )
+        databases: Sequence[DatabaseModel] = list(config.resources.databases.values())
 
         resources: Sequence[IsDatabricksResource] = (
             llms
@@ -233,6 +235,7 @@ class DatabricksProvider(ServiceProvider):
             + functions
             + tables
             + connections
+            + databases
         )
 
         # Flatten all resources from all models into a single list

dao_ai/tools/human_in_the_loop.py

@@ -87,7 +87,11 @@ def as_human_in_the_loop(
     if isinstance(function, BaseFunctionModel):
         human_in_the_loop: HumanInTheLoopModel | None = function.human_in_the_loop
         if human_in_the_loop:
-            logger.debug(f"Adding human-in-the-loop to tool: {tool.name}")
+            # Get tool name safely - handle RunnableBinding objects
+            tool_name = getattr(tool, "name", None) or getattr(
+                getattr(tool, "bound", None), "name", "unknown_tool"
+            )
+            logger.debug(f"Adding human-in-the-loop to tool: {tool_name}")
             tool = add_human_in_the_loop(
                 tool=tool,
                 interrupt_config=human_in_the_loop.interupt_config,

dao_ai/tools/unity_catalog.py

@@ -1,14 +1,18 @@
-from typing import Sequence
+from typing import Any, Dict, Optional, Sequence, Union
 
-from databricks_langchain import (
-    DatabricksFunctionClient,
-    UCFunctionToolkit,
-)
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.catalog import 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 dao_ai.config import (
+    AnyVariable,
+    CompositeVariableModel,
+    ToolModel,
     UnityCatalogFunctionModel,
+    value_of,
 )
 from dao_ai.tools.human_in_the_loop import as_human_in_the_loop
 
@@ -32,19 +36,330 @@ def create_uc_tools(
 
     logger.debug(f"create_uc_tools: {function}")
 
+    original_function_model = None
     if isinstance(function, UnityCatalogFunctionModel):
-        function = function.full_name
+        original_function_model = function
+        function_name = function.full_name
+    else:
+        function_name = function
 
-    client: DatabricksFunctionClient = DatabricksFunctionClient()
+    # Determine which tools to create
+    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
+        )
+
+        # Use with_partial_args to create the authenticated tool
+        tools = [with_partial_args(tool_model, original_function_model.partial_args)]
+    else:
+        # Fallback to standard UC toolkit approach
+        client: DatabricksFunctionClient = DatabricksFunctionClient()
+
+        toolkit: UCFunctionToolkit = UCFunctionToolkit(
+            function_names=[function_name], client=client
+        )
+
+        tools = toolkit.tools or []
+        logger.debug(f"Retrieved tools: {tools}")
+
+    # Apply human-in-the-loop wrapper to all tools and return
+    return [as_human_in_the_loop(tool=tool, function=function_name) for tool in tools]
+
+
+def _execute_uc_function(
+    client: DatabricksFunctionClient,
+    function_name: str,
+    partial_args: Dict[str, str] = None,
+    **kwargs: Any,
+) -> str:
+    """Execute Unity Catalog function with partial args and provided parameters."""
+
+    # Start with partial args if provided
+    all_params: Dict[str, Any] = dict(partial_args) if partial_args else {}
+
+    # Add any additional kwargs
+    all_params.update(kwargs)
 
-    toolkit: UCFunctionToolkit = UCFunctionToolkit(
-        function_names=[function], client=client
+    logger.debug(
+        f"Calling UC function {function_name} with parameters: {list(all_params.keys())}"
     )
 
-    tools = toolkit.tools or []
+    result = client.execute_function(function_name=function_name, parameters=all_params)
+
+    # Handle errors and extract result
+    if hasattr(result, "error") and result.error:
+        logger.error(f"Unity Catalog function error: {result.error}")
+        raise RuntimeError(f"Function execution failed: {result.error}")
+
+    result_value: str = result.value if hasattr(result, "value") else str(result)
+    logger.debug(f"UC function result: {result_value}")
+    return result_value
+
+
+def _grant_function_permissions(
+    function_name: str,
+    client_id: str,
+    host: Optional[str] = None,
+) -> None:
+    """
+    Grant comprehensive permissions to the service principal for Unity Catalog function execution.
+
+    This includes:
+    - EXECUTE permission on the function itself
+    - USE permission on the containing schema
+    - USE permission on the containing catalog
+    """
+    try:
+        # Initialize workspace client
+        workspace_client = WorkspaceClient(host=host) if host else WorkspaceClient()
+
+        # Parse the function name to get catalog and schema
+        parts = function_name.split(".")
+        if len(parts) != 3:
+            logger.warning(
+                f"Invalid function name format: {function_name}. Expected catalog.schema.function"
+            )
+            return
+
+        catalog_name, schema_name, func_name = parts
+        schema_full_name = f"{catalog_name}.{schema_name}"
+
+        logger.debug(
+            f"Granting comprehensive permissions on function {function_name} to principal {client_id}"
+        )
+
+        # 1. Grant EXECUTE permission on the function
+        try:
+            workspace_client.grants.update(
+                securable_type="function",
+                full_name=function_name,
+                changes=[
+                    PermissionsChange(principal=client_id, add=[Privilege.EXECUTE])
+                ],
+            )
+            logger.debug(f"Granted EXECUTE on function {function_name}")
+        except Exception as e:
+            logger.warning(f"Failed to grant EXECUTE on function {function_name}: {e}")
+
+        # 2. Grant USE_SCHEMA permission on the schema
+        try:
+            workspace_client.grants.update(
+                securable_type="schema",
+                full_name=schema_full_name,
+                changes=[
+                    PermissionsChange(
+                        principal=client_id,
+                        add=[Privilege.USE_SCHEMA],
+                    )
+                ],
+            )
+            logger.debug(f"Granted USE_SCHEMA on schema {schema_full_name}")
+        except Exception as e:
+            logger.warning(
+                f"Failed to grant USE_SCHEMA on schema {schema_full_name}: {e}"
+            )
+
+        # 3. Grant USE_CATALOG and BROWSE permissions on the catalog
+        try:
+            workspace_client.grants.update(
+                securable_type="catalog",
+                full_name=catalog_name,
+                changes=[
+                    PermissionsChange(
+                        principal=client_id,
+                        add=[Privilege.USE_CATALOG, Privilege.BROWSE],
+                    )
+                ],
+            )
+            logger.debug(f"Granted USE_CATALOG and BROWSE on catalog {catalog_name}")
+        except Exception as e:
+            logger.warning(
+                f"Failed to grant catalog permissions on {catalog_name}: {e}"
+            )
+
+        logger.debug(
+            f"Successfully granted comprehensive permissions on {function_name} to {client_id}"
+        )
+
+    except Exception as e:
+        logger.warning(
+            f"Failed to grant permissions on function {function_name} to {client_id}: {e}"
+        )
+        # Don't fail the tool creation if permission granting fails
+        pass
+
+
+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.
+
+    Args:
+        original_schema: The original Pydantic model class
+        exclude_fields: Set of field names to exclude from the schema
+
+    Returns:
+        A new Pydantic model class with the specified fields removed
+    """
+    from pydantic import BaseModel, Field, create_model
+    from pydantic.fields import PydanticUndefined
+
+    try:
+        # Get the original model's fields (Pydantic v2)
+        original_fields = original_schema.model_fields
+        filtered_field_definitions = {}
+
+        for name, field in original_fields.items():
+            if name not in exclude_fields:
+                # Reconstruct the field definition for create_model
+                field_type = field.annotation
+                field_default = (
+                    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)
+
+        # If no fields remain after filtering, return a generic empty schema
+        if not filtered_field_definitions:
+
+            class EmptySchema(BaseModel):
+                """Unity Catalog function with all parameters provided via partial args."""
 
-    logger.debug(f"Retrieved tools: {tools}")
+                pass
 
-    tools = [as_human_in_the_loop(tool=tool, function=function) for tool in tools]
+            return EmptySchema
+
+        # Create the new model dynamically
+        model_name = f"Filtered{original_schema.__name__}"
+        docstring = getattr(
+            original_schema, "__doc__", "Filtered Unity Catalog function parameters."
+        )
+
+        filtered_model = 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}")
+
+        # Fallback to generic schema
+        class GenericFilteredSchema(BaseModel):
+            """Generic filtered schema for Unity Catalog function."""
+
+            pass
+
+        return GenericFilteredSchema
+
+
+def with_partial_args(
+    tool: Union[ToolModel, Dict[str, Any]],
+    partial_args: dict[str, AnyVariable] = {},
+) -> StructuredTool:
+    """
+    Create a Unity Catalog tool with partial arguments pre-filled.
+
+    This function creates a wrapper tool that calls the UC function with partial arguments
+    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
+
+    Returns:
+        StructuredTool: A LangChain tool with partial arguments pre-filled
+    """
+    from unitycatalog.ai.langchain.toolkit import generate_function_input_params_schema
+
+    logger.debug(f"with_partial_args: {tool}")
+
+    # Convert dict-based variables to CompositeVariableModel and resolve their values
+    resolved_args = {}
+    for k, v in partial_args.items():
+        if isinstance(v, dict):
+            resolved_args[k] = value_of(CompositeVariableModel(**v))
+        else:
+            resolved_args[k] = value_of(v)
+
+    logger.debug(f"Resolved partial args: {resolved_args.keys()}")
+
+    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}")
+
+    # Grant permissions if we have credentials
+    if "client_id" in resolved_args:
+        client_id: str = resolved_args["client_id"]
+        host: Optional[str] = resolved_args.get("host")
+        try:
+            _grant_function_permissions(function_name, client_id, host)
+        except Exception as e:
+            logger.warning(f"Failed to grant permissions: {e}")
+
+    # Create the client for function execution
+    client: DatabricksFunctionClient = DatabricksFunctionClient()
+
+    # Try to get the function schema for better tool definition
+    try:
+        function_info = 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.debug(f"Tool description: {tool_description}")
+
+        # Create a modified schema that excludes partial args
+        original_schema = 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())}"
+        )
+
+    except Exception as e:
+        logger.warning(f"Could not introspect function {function_name}: {e}")
+        # Fallback to a generic schema
+        from pydantic import BaseModel
+
+        class GenericUCParams(BaseModel):
+            """Generic parameters for Unity Catalog function."""
+
+            pass
+
+        schema_model = GenericUCParams
+        tool_description = f"Unity Catalog function: {function_name}"
+
+    # Create a wrapper function that calls _execute_uc_function with partial args
+    def uc_function_wrapper(**kwargs) -> str:
+        """Wrapper function that executes Unity Catalog function with partial args."""
+        return _execute_uc_function(
+            client=client,
+            function_name=function_name,
+            partial_args=resolved_args,
+            **kwargs,
+        )
+
+    # Set the function name for the decorator
+    uc_function_wrapper.__name__ = tool.name or function_name.replace(".", "_")
+
+    # Create the tool using LangChain's StructuredTool
+    from langchain_core.tools import StructuredTool
+
+    partial_tool = StructuredTool.from_function(
+        func=uc_function_wrapper,
+        name=tool.name or function_name.replace(".", "_"),
+        description=tool_description,
+        args_schema=schema_model,
+    )
 
-    return tools
+    return partial_tool

{dao_ai-0.0.18.dist-info → dao_ai-0.0.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.18
+Version: 0.0.20
 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
@@ -24,22 +24,22 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: System :: Distributed Computing
 Requires-Python: >=3.12
-Requires-Dist: databricks-agents>=1.2.0
-Requires-Dist: databricks-langchain>=0.4.2
-Requires-Dist: databricks-sdk[openai]>=0.55.0
+Requires-Dist: databricks-agents>=1.6.0
+Requires-Dist: databricks-langchain>=0.8.0
+Requires-Dist: databricks-sdk[openai]>=0.66.0
 Requires-Dist: duckduckgo-search>=8.0.2
 Requires-Dist: grandalf>=0.8
-Requires-Dist: langchain-mcp-adapters>=0.1.9
+Requires-Dist: langchain-mcp-adapters>=0.1.10
 Requires-Dist: langchain-tavily>=0.2.11
 Requires-Dist: langchain>=0.3.27
 Requires-Dist: langgraph-checkpoint-postgres>=2.0.23
 Requires-Dist: langgraph-supervisor>=0.0.29
 Requires-Dist: langgraph-swarm>=0.0.14
-Requires-Dist: langgraph>=0.6.5
+Requires-Dist: langgraph>=0.6.7
 Requires-Dist: langmem>=0.0.29
 Requires-Dist: loguru>=0.7.3
-Requires-Dist: mcp>=1.9.1
-Requires-Dist: mlflow>=3.3.2
+Requires-Dist: mcp>=1.14.1
+Requires-Dist: mlflow>=3.4.0
 Requires-Dist: nest-asyncio>=1.6.0
 Requires-Dist: openevals>=0.0.19
 Requires-Dist: openpyxl>=3.1.5

{dao_ai-0.0.18.dist-info → dao_ai-0.0.20.dist-info}/RECORD

@@ -3,11 +3,11 @@ dao_ai/agent_as_code.py,sha256=kPSeDz2-1jRaed1TMs4LA3VECoyqe9_Ed2beRLB9gXQ,472
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
 dao_ai/chat_models.py,sha256=uhwwOTeLyHWqoTTgHrs4n5iSyTwe4EQcLKnh3jRxPWI,8626
 dao_ai/cli.py,sha256=Aez2TQW3Q8Ho1IaIkRggt0NevDxAAVPjXkePC5GPJF0,20429
-dao_ai/config.py,sha256=JlYC8N_7UL8VVkdSepiCUnR9NA5OsCVAigLjse7dMFM,49922
-dao_ai/graph.py,sha256=kXaGLGFVekDWqm-AHzti6LmrXnyi99VQ-AdCGuNb_xM,7831
-dao_ai/guardrails.py,sha256=-Qh0f_2Db9t4Nbrrx9FM7tnpqShjMoyxepZ0HByItfU,4027
+dao_ai/config.py,sha256=ZO5ei45gnhqg1BtD0R9aekJz4ClmiTw2GHhOk4Idil4,51958
+dao_ai/graph.py,sha256=gmD9mxODfXuvn9xWeBfewm1FiuVAWMLEdnZz7DNmSH0,7859
+dao_ai/guardrails.py,sha256=4TKArDONRy8RwHzOT1plZ1rhy3x9GF_aeGpPCRl6wYA,4016
 dao_ai/messages.py,sha256=xl_3-WcFqZKCFCiov8sZOPljTdM3gX3fCHhxq-xFg2U,7005
-dao_ai/models.py,sha256=h_xFMK5FHQwPApEAYhvrt69y7ZUljmqThHTjp-yde_o,25368
+dao_ai/models.py,sha256=Xb23U-lhDG8KyNRIijcJ4InluadlaGNy4rrYx7Cjgfg,26939
 dao_ai/nodes.py,sha256=SSuFNTXOdFaKg_aX-yUkQO7fM9wvNGu14lPXKDapU1U,8461
 dao_ai/prompts.py,sha256=vpmIbWs_szXUgNNDs5Gh2LcxKZti5pHDKSfoClUcgX0,1289
 dao_ai/state.py,sha256=GwbMbd1TWZx1T5iQrEOX6_rpxOitlmyeJ8dMr2o_pag,1031
@@ -18,23 +18,23 @@ dao_ai/hooks/__init__.py,sha256=LlHGIuiZt6vGW8K5AQo1XJEkBP5vDVtMhq0IdjcLrD4,417
 dao_ai/hooks/core.py,sha256=ZShHctUSoauhBgdf1cecy9-D7J6-sGn-pKjuRMumW5U,6663
 dao_ai/memory/__init__.py,sha256=1kHx_p9abKYFQ6EYD05nuc1GS5HXVEpufmjBGw_7Uho,260
 dao_ai/memory/base.py,sha256=99nfr2UZJ4jmfTL_KrqUlRSCoRxzkZyWyx5WqeUoMdQ,338
-dao_ai/memory/core.py,sha256=K45iCEFbqJCVxMi4m3vmBJi4c6TQ-UtKGzyugDTkPP0,4141
-dao_ai/memory/postgres.py,sha256=YILzA7xtqawPAOLFaGG_i17zW7cQxXTzTD8yd-ipe8k,12480
+dao_ai/memory/core.py,sha256=g7chjBgVgx3iKjR2hghl0QL1j3802uIM_e7mgszur9M,4151
+dao_ai/memory/postgres.py,sha256=pxxMjGotgqjrKhx0lVR3EAjSZTQgBpiPZOB0-cyjprc,12505
 dao_ai/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dao_ai/providers/base.py,sha256=-fjKypCOk28h6vioPfMj9YZSw_3Kcbi2nMuAyY7vX9k,1383
-dao_ai/providers/databricks.py,sha256=KLYrLccOA3Uws9nWJcJUZTbMz-MdR_onhlQeztbplCM,28073
+dao_ai/providers/databricks.py,sha256=fZ8mGotfA3W3t5yUej2xGmGHSybjBFYr895mOctT418,28203
 dao_ai/tools/__init__.py,sha256=ye6MHaJY7tUnJ8336YJiLxuZr55zDPNdOw6gm7j5jlc,1103
 dao_ai/tools/agent.py,sha256=WbQnyziiT12TLMrA7xK0VuOU029tdmUBXbUl-R1VZ0Q,1886
 dao_ai/tools/core.py,sha256=Kei33S8vrmvPOAyrFNekaWmV2jqZ-IPS1QDSvU7RZF0,1984
 dao_ai/tools/genie.py,sha256=GzV5lfDYKmzW_lSLxAsPaTwnzX6GxQOB1UcLaTDqpfY,2787
-dao_ai/tools/human_in_the_loop.py,sha256=IBmQJmpxkdDxnBNyABc_-dZhhsQlTNTkPyUXgkHKIgY,3466
+dao_ai/tools/human_in_the_loop.py,sha256=yk35MO9eNETnYFH-sqlgR-G24TrEgXpJlnZUustsLkI,3681
 dao_ai/tools/mcp.py,sha256=auEt_dwv4J26fr5AgLmwmnAsI894-cyuvkvjItzAUxs,4419
 dao_ai/tools/python.py,sha256=XcQiTMshZyLUTVR5peB3vqsoUoAAy8gol9_pcrhddfI,1831
 dao_ai/tools/time.py,sha256=Y-23qdnNHzwjvnfkWvYsE7PoWS1hfeKy44tA7sCnNac,8759
-dao_ai/tools/unity_catalog.py,sha256=PXfLj2EgyQgaXq4Qq3t25AmTC4KyVCF_-sCtg6enens,1404
+dao_ai/tools/unity_catalog.py,sha256=uX_h52BuBAr4c9UeqSMI7DNz3BPRLeai5tBVW4sJqRI,13113
 dao_ai/tools/vector_search.py,sha256=EDYQs51zIPaAP0ma1D81wJT77GQ-v-cjb2XrFVWfWdg,2621
-dao_ai-0.0.18.dist-info/METADATA,sha256=9lTAXjEqQHxl6dmRMyiqUnYT1Nh_wJpSeJXRG8bGZGg,41378
-dao_ai-0.0.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dao_ai-0.0.18.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.0.18.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.0.18.dist-info/RECORD,,
+dao_ai-0.0.20.dist-info/METADATA,sha256=gWNRLhswz5sCe1vxbBQ6dGlgiObI9nI829Q5DQRqRRY,41380
+dao_ai-0.0.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dao_ai-0.0.20.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.0.20.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.0.20.dist-info/RECORD,,