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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.13
+Version: 0.1.15
 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.15}/config/examples/02_mcp/managed_mcp.yaml

@@ -3,19 +3,19 @@
 variables:
   client_id: &client_id                          # Service principal client ID
     options:
-      - env: RETAIL_AI_DATABRICKS_CLIENT_ID
-      - scope: retail_ai
+      - scope: retail_consumer_goods
         secret: RETAIL_AI_DATABRICKS_CLIENT_ID
+      - env: RETAIL_AI_DATABRICKS_CLIENT_ID
   client_secret: &client_secret                  # Service principal secret
     options:
-      - env: RETAIL_AI_DATABRICKS_CLIENT_SECRET
-      - scope: retail_ai
+      - scope: retail_consumer_goods
         secret: RETAIL_AI_DATABRICKS_CLIENT_SECRET
+      - env: RETAIL_AI_DATABRICKS_CLIENT_SECRET
   workspace_host: &workspace_host                # Databricks workspace URL
     options:
-      - env: RETAIL_AI_DATABRICKS_HOST
-      - scope: retail_ai
+      - scope: retail_consumer_goods
         secret: RETAIL_AI_DATABRICKS_HOST
+      - env: RETAIL_AI_DATABRICKS_HOST
 
 schemas:
   retail_schema: &retail_schema                  # Unity Catalog schema configuration
@@ -112,6 +112,10 @@ app:
   registered_model:                              # MLflow registered model configuration
     schema: *retail_schema                       # Schema where model will be registered
     name: mcp_agent                              # Model name in MLflow registry
+  environment_vars:
+    RETAIL_AI_DATABRICKS_CLIENT_ID: *client_id
+    RETAIL_AI_DATABRICKS_CLIENT_SECRET: *client_secret
+    RETAIL_AI_DATABRICKS_HOST: *workspace_host
   agents:                                        # List of agents included in the system
     - *mcp_agent                                 # MCP agent
   orchestration:                                 # Agent orchestration configuration

{dao_ai-0.1.13 → dao_ai-0.1.15}/config/examples/15_complete_applications/quick_serve_restaurant.yaml

@@ -144,38 +144,7 @@ resources:
       description: "A warehouse for shared endpoints"  # Description
       warehouse_id: 148ccb90800933a1                # Databricks warehouse ID
 
-  # ---------------------------------------------------------------------------
-  # DATABASES
-  # ---------------------------------------------------------------------------
-  # Configure external databases (e.g., PostgreSQL for checkpoints and memory)
-  # Supports both traditional user/password auth and OAuth2 client credentials
-  databases:
-    # PostgreSQL database for agent memory and checkpoints
-    qsr_database: &qsr_database
-      name: "retail-consumer-goods"
-      description: "Database for agent memory and checkpoints" 
-      
-      # Database connection parameters
-      # These can be set via environment variables or Databricks secrets
-      host: *database_host_var
-      port: *port_var
-      database: *database_var
-      client_id: *client_id_var
-      client_secret: *client_secret_var
-      workspace_host: *workspace_host_var
-
-          
-      # Optional: Additional PostgreSQL connection parameters
-      # Uncomment and modify as needed for your specific setup
-      # connection_kwargs:                            # Additional connection parameters
-      #   autocommit: True                            # Enable autocommit mode
-      #   prepare_threshold: 0                        # Disable prepared statements
-      #   sslmode: require                            # Require SSL connection
-      #   connect_timeout: 10                         # Connection timeout in seconds
-      
-      # Optional: Connection pool configuration
-      # max_pool_size: 20                             # Maximum connections in pool
-      # timeout: 5                                    # Connection timeout in seconds
+
 
 # =============================================================================
 # RETRIEVERS CONFIGURATION

{dao_ai-0.1.13 → dao_ai-0.1.15}/databricks.yaml.template

@@ -98,22 +98,22 @@ resources:
               config-path: ${var.config_path}
           environment_key: dao-ai-env
           
-# Cloud-specific targets with appropriate node types
-# The CLI auto-selects the target based on workspace cloud detection
+# App-specific cloud targets - ensures each app has its own deployment identity per cloud
+# The CLI auto-detects cloud provider and selects the appropriate target (e.g., app_name-azure)
 targets:
-  azure:
+  __APP_NAME__-azure:
     mode: development
     variables:
       cloud: azure
       node_type: Standard_D4ads_v5
   
-  aws:
+  __APP_NAME__-aws:
     mode: development
     variables:
       cloud: aws
       node_type: i3.xlarge
   
-  gcp:
+  __APP_NAME__-gcp:
     mode: development
     variables:
       cloud: gcp

{dao_ai-0.1.13 → dao_ai-0.1.15}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dao-ai"
-version = "0.1.13"
+version = "0.1.15"
 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.15}/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.15}/src/dao_ai/cli.py

@@ -63,16 +63,26 @@ def detect_cloud_provider(profile: Optional[str] = None) -> Optional[str]:
         Cloud provider string ('azure', 'aws', 'gcp') or None if detection fails
     """
     try:
+        import os
         from databricks.sdk import WorkspaceClient
 
+        # Check for environment variables that might override profile
+        if profile and os.environ.get("DATABRICKS_HOST"):
+            logger.warning(
+                f"DATABRICKS_HOST environment variable is set, which may override --profile {profile}"
+            )
+
         # Create workspace client with optional profile
         if profile:
+            logger.debug(f"Creating WorkspaceClient with profile: {profile}")
             w = WorkspaceClient(profile=profile)
         else:
+            logger.debug("Creating WorkspaceClient with default/ambient credentials")
             w = WorkspaceClient()
 
         # Get the workspace URL from config
         host = w.config.host
+        logger.debug(f"WorkspaceClient host: {host}, profile used: {profile}")
         if not host:
             logger.warning("Could not determine workspace URL for cloud detection")
             return None
@@ -1143,7 +1153,7 @@ def run_databricks_command(
     app_config: AppConfig = AppConfig.from_file(config_path) if config_path else None
     normalized_name: str = normalize_name(app_config.app.name) if app_config else None
 
-    # Auto-detect cloud provider if not specified
+    # Auto-detect cloud provider if not specified (used for node_type selection)
     if not cloud:
         cloud = detect_cloud_provider(profile)
         if cloud:
@@ -1156,10 +1166,12 @@ def run_databricks_command(
     if config_path and app_config:
         generate_bundle_from_template(config_path, normalized_name)
 
-    # Use cloud as target (azure, aws, gcp) - can be overridden with explicit --target
+    # Use app-specific cloud target: {app_name}-{cloud}
+    # This ensures each app has unique deployment identity while supporting cloud-specific settings
+    # Can be overridden with explicit --target
     if not target:
-        target = cloud
-        logger.debug(f"Using cloud-based target: {target}")
+        target = f"{normalized_name}-{cloud}"
+        logger.info(f"Using app-specific cloud target: {target}")
 
     # Build databricks command
     # --profile is a global flag, --target is a subcommand flag for 'bundle'

{dao_ai-0.1.13 → dao_ai-0.1.15}/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,57 @@ 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
+
+        logger.trace(f"workspace_client_from called", context=context, on_behalf_of_user=self.on_behalf_of_user)
+
+        # 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: str = 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.15}/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.15}/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.15}/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.15}/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 = (