dao-ai 0.1.13__tar.gz → 0.1.14__tar.gz

{dao_ai-0.1.13 → dao_ai-0.1.14}/PKG-INFO

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dao-ai"
-version = "0.1.13"
+version = "0.1.14"
 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.13 → dao_ai-0.1.14}/src/dao_ai/apps/handlers.py

@@ -14,7 +14,7 @@ from typing import AsyncGenerator
 
 import mlflow
 from dotenv import load_dotenv
-from mlflow.genai.agent_server import invoke, stream
+from mlflow.genai.agent_server import get_request_headers, invoke, stream
 from mlflow.types.responses import (
     ResponsesAgentRequest,
     ResponsesAgentResponse,
@@ -25,6 +25,23 @@ from dao_ai.config import AppConfig
 from dao_ai.logging import configure_logging
 from dao_ai.models import LanggraphResponsesAgent
 
+
+def _inject_headers_into_request(request: ResponsesAgentRequest) -> None:
+    """Inject request headers into custom_inputs for Context propagation.
+
+    Captures headers from the MLflow AgentServer context (where they're available)
+    and injects them into request.custom_inputs.configurable.headers so they
+    flow through to Context and can be used for OBO authentication.
+    """
+    headers: dict[str, str] = get_request_headers()
+    if headers:
+        if request.custom_inputs is None:
+            request.custom_inputs = {}
+        if "configurable" not in request.custom_inputs:
+            request.custom_inputs["configurable"] = {}
+        request.custom_inputs["configurable"]["headers"] = headers
+
+
 # Load environment variables from .env.local if it exists
 load_dotenv(dotenv_path=".env.local", override=True)
 
@@ -61,6 +78,8 @@ async def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentRespons
     Returns:
         ResponsesAgentResponse with the complete output
     """
+    # Capture headers while in the AgentServer async context (before they're lost)
+    _inject_headers_into_request(request)
     return await _responses_agent.apredict(request)
 
 
@@ -80,5 +99,7 @@ async def streaming(
     Yields:
         ResponsesAgentStreamEvent objects as they are generated
     """
+    # Capture headers while in the AgentServer async context (before they're lost)
+    _inject_headers_into_request(request)
     async for event in _responses_agent.apredict_stream(request):
         yield event

{dao_ai-0.1.13 → dao_ai-0.1.14}/src/dao_ai/config.py

@@ -7,6 +7,7 @@ from enum import Enum
 from os import PathLike
 from pathlib import Path
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Iterator,
@@ -18,6 +19,9 @@ from typing import (
     Union,
 )
 
+if TYPE_CHECKING:
+    from dao_ai.state import Context
+
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.credentials_provider import (
     CredentialsStrategy,
@@ -284,8 +288,8 @@ class IsDatabricksResource(ABC, BaseModel):
 
         Authentication priority:
         1. On-Behalf-Of User (on_behalf_of_user=True):
-           - Forwarded headers (Databricks Apps)
-           - ModelServingUserCredentials (Model Serving)
+           - Uses ModelServingUserCredentials (Model Serving)
+           - For Databricks Apps with headers, use workspace_client_from(context)
         2. Service Principal (client_id + client_secret + workspace_host)
         3. PAT (pat + workspace_host)
         4. Ambient/default authentication
@@ -294,37 +298,6 @@ class IsDatabricksResource(ABC, BaseModel):
 
         # Check for OBO first (highest priority)
         if self.on_behalf_of_user:
-            # In Databricks Apps, use forwarded headers for per-user auth
-            from mlflow.genai.agent_server import get_request_headers
-
-            headers = get_request_headers()
-            logger.debug(f"Headers received: {list(headers.keys())}")
-            # Try both lowercase and title-case header names (HTTP headers are case-insensitive)
-            forwarded_token = headers.get("x-forwarded-access-token") or headers.get(
-                "X-Forwarded-Access-Token"
-            )
-
-            if forwarded_token:
-                forwarded_user = headers.get("x-forwarded-user") or headers.get(
-                    "X-Forwarded-User", "unknown"
-                )
-                logger.debug(
-                    f"Creating WorkspaceClient for {self.__class__.__name__} "
-                    f"with OBO using forwarded token from Databricks Apps",
-                    forwarded_user=forwarded_user,
-                )
-                # Use workspace_host if configured, otherwise SDK will auto-detect
-                workspace_host_value: str | None = (
-                    normalize_host(value_of(self.workspace_host))
-                    if self.workspace_host
-                    else None
-                )
-                return WorkspaceClient(
-                    host=workspace_host_value,
-                    token=forwarded_token,
-                    auth_type="pat",
-                )
-
             credentials_strategy: CredentialsStrategy = ModelServingUserCredentials()
             logger.debug(
                 f"Creating WorkspaceClient for {self.__class__.__name__} "
@@ -383,6 +356,55 @@ class IsDatabricksResource(ABC, BaseModel):
         )
         return WorkspaceClient()
 
+    def workspace_client_from(self, context: "Context | None") -> WorkspaceClient:
+        """
+        Get a WorkspaceClient using headers from the provided Context.
+
+        Use this method from tools that have access to ToolRuntime[Context].
+        This allows OBO authentication to work in Databricks Apps where headers
+        are captured at request entry and passed through the Context.
+
+        Args:
+            context: Runtime context containing headers for OBO auth.
+                     If None or no headers, falls back to workspace_client property.
+
+        Returns:
+            WorkspaceClient configured with appropriate authentication.
+        """
+        from dao_ai.utils import normalize_host
+
+        # Check if we have headers in context for OBO
+        if context and context.headers and self.on_behalf_of_user:
+            headers = context.headers
+            # Try both lowercase and title-case header names (HTTP headers are case-insensitive)
+            forwarded_token = headers.get("x-forwarded-access-token") or headers.get(
+                "X-Forwarded-Access-Token"
+            )
+
+            if forwarded_token:
+                forwarded_user = headers.get("x-forwarded-user") or headers.get(
+                    "X-Forwarded-User", "unknown"
+                )
+                logger.debug(
+                    f"Creating WorkspaceClient for {self.__class__.__name__} "
+                    f"with OBO using forwarded token from Context",
+                    forwarded_user=forwarded_user,
+                )
+                # Use workspace_host if configured, otherwise SDK will auto-detect
+                workspace_host_value: str | None = (
+                    normalize_host(value_of(self.workspace_host))
+                    if self.workspace_host
+                    else None
+                )
+                return WorkspaceClient(
+                    host=workspace_host_value,
+                    token=forwarded_token,
+                    auth_type="pat",
+                )
+
+        # Fall back to existing workspace_client property
+        return self.workspace_client
+
 
 class DeploymentTarget(str, Enum):
     """Target platform for agent deployment."""

{dao_ai-0.1.13 → dao_ai-0.1.14}/src/dao_ai/tools/genie.py

@@ -139,29 +139,53 @@ Returns:
 GenieResponse: A response object containing the conversation ID and result from Genie."""
     tool_description = tool_description + function_docs
 
-    genie: Genie = Genie(
-        space_id=space_id,
-        client=genie_room.workspace_client,
-        truncate_results=truncate_results,
-    )
+    # Cache for genie service - created lazily on first call
+    # This allows us to use workspace_client_from with runtime context for OBO
+    _cached_genie_service: GenieServiceBase | None = None
+
+    def _get_genie_service(context: Context | None) -> GenieServiceBase:
+        """Get or create the Genie service, using context for OBO auth if available."""
+        nonlocal _cached_genie_service
+
+        # Use cached service if available (for non-OBO or after first call)
+        # For OBO, we need fresh workspace client each time to use the user's token
+        if _cached_genie_service is not None and not genie_room.on_behalf_of_user:
+            return _cached_genie_service
+
+        # Get workspace client using context for OBO support
+        from databricks.sdk import WorkspaceClient
 
-    genie_service: GenieServiceBase = GenieService(genie)
-
-    # Wrap with semantic cache first (checked second due to decorator pattern)
-    if semantic_cache_parameters is not None:
-        genie_service = SemanticCacheService(
-            impl=genie_service,
-            parameters=semantic_cache_parameters,
-            workspace_client=genie_room.workspace_client,  # Pass workspace client for conversation history
-        ).initialize()  # Eagerly initialize to fail fast and create table
-
-    # Wrap with LRU cache last (checked first - fast O(1) exact match)
-    if lru_cache_parameters is not None:
-        genie_service = LRUCacheService(
-            impl=genie_service,
-            parameters=lru_cache_parameters,
+        workspace_client: WorkspaceClient = genie_room.workspace_client_from(context)
+
+        genie: Genie = Genie(
+            space_id=space_id,
+            client=workspace_client,
+            truncate_results=truncate_results,
         )
 
+        genie_service: GenieServiceBase = GenieService(genie)
+
+        # Wrap with semantic cache first (checked second due to decorator pattern)
+        if semantic_cache_parameters is not None:
+            genie_service = SemanticCacheService(
+                impl=genie_service,
+                parameters=semantic_cache_parameters,
+                workspace_client=workspace_client,
+            ).initialize()
+
+        # Wrap with LRU cache last (checked first - fast O(1) exact match)
+        if lru_cache_parameters is not None:
+            genie_service = LRUCacheService(
+                impl=genie_service,
+                parameters=lru_cache_parameters,
+            )
+
+        # Cache for non-OBO scenarios
+        if not genie_room.on_behalf_of_user:
+            _cached_genie_service = genie_service
+
+        return genie_service
+
     @tool(
         name_or_callable=tool_name,
         description=tool_description,
@@ -177,6 +201,10 @@ GenieResponse: A response object containing the conversation ID and result from
         # Access state through runtime
         state: AgentState = runtime.state
         tool_call_id: str = runtime.tool_call_id
+        context: Context | None = runtime.context
+
+        # Get genie service with OBO support via context
+        genie_service: GenieServiceBase = _get_genie_service(context)
 
         # Ensure space_id is a string for state keys
         space_id_str: str = str(space_id)
@@ -194,6 +222,14 @@ GenieResponse: A response object containing the conversation ID and result from
             conversation_id=existing_conversation_id,
         )
 
+        # Log the prompt being sent to Genie
+        logger.trace(
+            "Sending prompt to Genie",
+            space_id=space_id_str,
+            conversation_id=existing_conversation_id,
+            prompt=question[:500] + "..." if len(question) > 500 else question,
+        )
+
         # Call ask_question which always returns CacheResult with cache metadata
         cache_result: CacheResult = genie_service.ask_question(
             question, conversation_id=existing_conversation_id
@@ -211,6 +247,22 @@ GenieResponse: A response object containing the conversation ID and result from
             cache_key=cache_key,
         )
 
+        # Log truncated response for debugging
+        result_preview: str = str(genie_response.result)
+        if len(result_preview) > 500:
+            result_preview = result_preview[:500] + "..."
+        logger.trace(
+            "Genie response content",
+            question=question[:100] + "..." if len(question) > 100 else question,
+            query=genie_response.query,
+            description=(
+                genie_response.description[:200] + "..."
+                if genie_response.description and len(genie_response.description) > 200
+                else genie_response.description
+            ),
+            result_preview=result_preview,
+        )
+
         # Update session state with cache information
         if persist_conversation:
             session.genie.update_space(

{dao_ai-0.1.13 → dao_ai-0.1.14}/src/dao_ai/tools/mcp.py

@@ -30,6 +30,7 @@ from dao_ai.config import (
     McpFunctionModel,
     TransportType,
 )
+from dao_ai.state import Context
 
 
 @dataclass
@@ -173,6 +174,7 @@ def _get_auth_resource(function: McpFunctionModel) -> IsDatabricksResource:
 
 def _build_connection_config(
     function: McpFunctionModel,
+    context: Context | None = None,
 ) -> dict[str, Any]:
     """
     Build the connection configuration dictionary for MultiServerMCPClient.
@@ -193,6 +195,7 @@ def _build_connection_config(
 
     Args:
         function: The MCP function model configuration.
+        context: Optional runtime context with headers for OBO auth.
 
     Returns:
         A dictionary containing the transport-specific connection settings.
@@ -205,14 +208,17 @@ def _build_connection_config(
         }
 
     # For HTTP transport, use DatabricksOAuthClientProvider with unified auth
+    from databricks.sdk import WorkspaceClient
     from databricks_mcp import DatabricksOAuthClientProvider
 
     # Get the resource to use for authentication
-    auth_resource = _get_auth_resource(function)
+    auth_resource: IsDatabricksResource = _get_auth_resource(function)
 
-    # Get workspace client from the auth resource
-    workspace_client = auth_resource.workspace_client
-    auth_provider = DatabricksOAuthClientProvider(workspace_client)
+    # Get workspace client from the auth resource with OBO support via context
+    workspace_client: WorkspaceClient = auth_resource.workspace_client_from(context)
+    auth_provider: DatabricksOAuthClientProvider = DatabricksOAuthClientProvider(
+        workspace_client
+    )
 
     # Log which resource is providing auth
     resource_name = (
@@ -509,19 +515,28 @@ async def acreate_mcp_tools(
     def _create_tool_wrapper(mcp_tool: Tool) -> RunnableLike:
         """
         Create a LangChain tool wrapper for an MCP tool.
+
+        Supports OBO authentication via context headers.
         """
+        from langchain.tools import ToolRuntime
 
         @create_tool(
             mcp_tool.name,
             description=mcp_tool.description or f"MCP tool: {mcp_tool.name}",
             args_schema=mcp_tool.inputSchema,
         )
-        async def tool_wrapper(**kwargs: Any) -> str:
+        async def tool_wrapper(
+            runtime: ToolRuntime[Context] = None,
+            **kwargs: Any,
+        ) -> str:
             """Execute MCP tool with fresh session."""
             logger.trace("Invoking MCP tool", tool_name=mcp_tool.name, args=kwargs)
 
-            invocation_client = MultiServerMCPClient(
-                {"mcp_function": _build_connection_config(function)}
+            # Get context for OBO support
+            context: Context | None = runtime.context if runtime else None
+
+            invocation_client: MultiServerMCPClient = MultiServerMCPClient(
+                {"mcp_function": _build_connection_config(function, context)}
             )
 
             try:
@@ -530,7 +545,7 @@ async def acreate_mcp_tools(
                         mcp_tool.name, kwargs
                     )
 
-                    text_result = _extract_text_content(result)
+                    text_result: str = _extract_text_content(result)
 
                     logger.trace(
                         "MCP tool completed",
@@ -625,20 +640,28 @@ def create_mcp_tools(
         This wrapper handles:
         - Fresh session creation per invocation (stateless)
         - Content extraction to plain text (avoiding extra fields)
+        - OBO authentication via context headers
         """
+        from langchain.tools import ToolRuntime
 
         @create_tool(
             mcp_tool.name,
             description=mcp_tool.description or f"MCP tool: {mcp_tool.name}",
             args_schema=mcp_tool.inputSchema,
         )
-        async def tool_wrapper(**kwargs: Any) -> str:
+        async def tool_wrapper(
+            runtime: ToolRuntime[Context] = None,
+            **kwargs: Any,
+        ) -> str:
             """Execute MCP tool with fresh session."""
             logger.trace("Invoking MCP tool", tool_name=mcp_tool.name, args=kwargs)
 
-            # Create a fresh client/session for each invocation
-            invocation_client = MultiServerMCPClient(
-                {"mcp_function": _build_connection_config(function)}
+            # Get context for OBO support
+            context: Context | None = runtime.context if runtime else None
+
+            # Create a fresh client/session for each invocation with OBO support
+            invocation_client: MultiServerMCPClient = MultiServerMCPClient(
+                {"mcp_function": _build_connection_config(function, context)}
             )
 
             try:
@@ -648,7 +671,7 @@ def create_mcp_tools(
                     )
 
                     # Extract text content, avoiding extra fields
-                    text_result = _extract_text_content(result)
+                    text_result: str = _extract_text_content(result)
 
                     logger.trace(
                         "MCP tool completed",

{dao_ai-0.1.13 → dao_ai-0.1.14}/src/dao_ai/tools/slack.py

@@ -1,11 +1,13 @@
 from typing import Any, Callable, Optional
 
 from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
+from langchain.tools import ToolRuntime
 from langchain_core.tools import tool
 from loguru import logger
 from requests import Response
 
 from dao_ai.config import ConnectionModel
+from dao_ai.state import Context
 
 
 def _find_channel_id_by_name(
@@ -129,8 +131,17 @@ def create_send_slack_message_tool(
         name_or_callable=name,
         description=description,
     )
-    def send_slack_message(text: str) -> str:
-        response: Response = connection.workspace_client.serving_endpoints.http_request(
+    def send_slack_message(
+        text: str,
+        runtime: ToolRuntime[Context] = None,
+    ) -> str:
+        from databricks.sdk import WorkspaceClient
+
+        # Get workspace client with OBO support via context
+        context: Context | None = runtime.context if runtime else None
+        workspace_client: WorkspaceClient = connection.workspace_client_from(context)
+
+        response: Response = workspace_client.serving_endpoints.http_request(
             conn=connection.name,
             method=ExternalFunctionRequestHttpMethod.POST,
             path="/api/chat.postMessage",

{dao_ai-0.1.13 → dao_ai-0.1.14}/src/dao_ai/tools/sql.py

@@ -7,10 +7,11 @@ pre-configured SQL statements against a Databricks SQL warehouse.
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.sql import StatementResponse, StatementState
-from langchain.tools import tool
+from langchain.tools import ToolRuntime, tool
 from loguru import logger
 
 from dao_ai.config import WarehouseModel, value_of
+from dao_ai.state import Context
 
 
 def create_execute_statement_tool(
@@ -63,7 +64,6 @@ def create_execute_statement_tool(
         description = f"Execute a pre-configured SQL query against the {warehouse.name} warehouse and return the results."
 
     warehouse_id: str = value_of(warehouse.warehouse_id)
-    workspace_client: WorkspaceClient = warehouse.workspace_client
 
     logger.debug(
         "Creating SQL execution tool",
@@ -74,7 +74,7 @@ def create_execute_statement_tool(
     )
 
     @tool(name_or_callable=name, description=description)
-    def execute_statement_tool() -> str:
+    def execute_statement_tool(runtime: ToolRuntime[Context] = None) -> str:
         """
         Execute the pre-configured SQL statement against the Databricks SQL warehouse.
 
@@ -88,6 +88,10 @@ def create_execute_statement_tool(
             sql_preview=statement[:100] + "..." if len(statement) > 100 else statement,
         )
 
+        # Get workspace client with OBO support via context
+        context: Context | None = runtime.context if runtime else None
+        workspace_client: WorkspaceClient = warehouse.workspace_client_from(context)
+
         try:
             # Execute the SQL statement
             statement_response: StatementResponse = (

{dao_ai-0.1.13 → dao_ai-0.1.14}/src/dao_ai/tools/unity_catalog.py

@@ -1,10 +1,11 @@
-from typing import Any, Dict, Optional, Sequence, Set
+from typing import Annotated, Any, Dict, Optional, Sequence, Set
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.catalog import FunctionInfo, PermissionsChange, Privilege
 from databricks_langchain import DatabricksFunctionClient, UCFunctionToolkit
+from langchain.tools import ToolRuntime
 from langchain_core.runnables.base import RunnableLike
-from langchain_core.tools import StructuredTool
+from langchain_core.tools import InjectedToolArg, StructuredTool
 from loguru import logger
 from pydantic import BaseModel
 from unitycatalog.ai.core.base import FunctionExecutionResult
@@ -15,6 +16,7 @@ from dao_ai.config import (
     UnityCatalogFunctionModel,
     value_of,
 )
+from dao_ai.state import Context
 from dao_ai.utils import normalize_host
 
 
@@ -35,13 +37,11 @@ def create_uc_tools(
         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
 
     if isinstance(function, UnityCatalogFunctionModel):
         original_function_model = function
         function_name = function.resource.full_name
-        workspace_client = function.resource.workspace_client
     else:
         function_name = function
 
@@ -56,6 +56,12 @@ def create_uc_tools(
         # Use with_partial_args directly with UnityCatalogFunctionModel
         tools = [with_partial_args(original_function_model)]
     else:
+        # For standard UC toolkit, we need workspace_client at creation time
+        # Use the resource's workspace_client (will use ambient auth if no OBO)
+        workspace_client: WorkspaceClient | None = None
+        if original_function_model:
+            workspace_client = original_function_model.resource.workspace_client
+
         # Fallback to standard UC toolkit approach
         client: DatabricksFunctionClient = DatabricksFunctionClient(
             client=workspace_client
@@ -356,7 +362,6 @@ def with_partial_args(
     # 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
 
     logger.debug(
         "Creating UC tool with partial args",
@@ -365,7 +370,7 @@ def with_partial_args(
         partial_args=list(resolved_args.keys()),
     )
 
-    # Grant permissions if we have credentials
+    # Grant permissions if we have credentials (using ambient auth for setup)
     if "client_id" in resolved_args:
         client_id: str = resolved_args["client_id"]
         host: Optional[str] = resolved_args.get("host")
@@ -376,14 +381,18 @@ def with_partial_args(
                 "Failed to grant permissions", function_name=function_name, error=str(e)
             )
 
-    # Create the client for function execution using the resource's workspace client
-    client: DatabricksFunctionClient = DatabricksFunctionClient(client=workspace_client)
+    # Get workspace client for schema introspection (uses ambient auth at definition time)
+    # Actual execution will use OBO via context
+    setup_workspace_client: WorkspaceClient = uc_function.resource.workspace_client
+    setup_client: DatabricksFunctionClient = DatabricksFunctionClient(
+        client=setup_workspace_client
+    )
 
     # Try to get the function schema for better tool definition
     schema_model: type[BaseModel]
     tool_description: str
     try:
-        function_info: FunctionInfo = client.get_function(function_name)
+        function_info: FunctionInfo = setup_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}"
@@ -419,8 +428,21 @@ def with_partial_args(
         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:
+    # Uses InjectedToolArg to ensure runtime is injected but hidden from the LLM
+    def uc_function_wrapper(
+        runtime: Annotated[ToolRuntime[Context], InjectedToolArg] = None,
+        **kwargs: Any,
+    ) -> str:
         """Wrapper function that executes Unity Catalog function with partial args."""
+        # Get workspace client with OBO support via context
+        context: Context | None = runtime.context if runtime else None
+        workspace_client: WorkspaceClient = uc_function.resource.workspace_client_from(
+            context
+        )
+        client: DatabricksFunctionClient = DatabricksFunctionClient(
+            client=workspace_client
+        )
+
         return _execute_uc_function(
             client=client,
             function_name=function_name,