dao-ai 0.1.5__py3-none-any.whl → 0.1.20__py3-none-any.whl

dao_ai/tools/router.py

@@ -0,0 +1,89 @@
+"""
+Query router for selecting execution mode based on query characteristics.
+
+Routes to internal execution modes within the same retriever instance:
+- standard: Single similarity_search for simple queries
+- instructed: Decompose -> Parallel Search -> RRF for constrained queries
+"""
+
+from pathlib import Path
+from typing import Any, Literal
+
+import mlflow
+import yaml
+from langchain_core.language_models import BaseChatModel
+from langchain_core.runnables import Runnable
+from loguru import logger
+from mlflow.entities import SpanType
+from pydantic import BaseModel, ConfigDict, Field
+
+# Load prompt template
+_PROMPT_PATH = Path(__file__).parent.parent / "prompts" / "router.yaml"
+
+
+def _load_prompt_template() -> dict[str, Any]:
+    """Load the router prompt template from YAML."""
+    with open(_PROMPT_PATH) as f:
+        return yaml.safe_load(f)
+
+
+class RouterDecision(BaseModel):
+    """Classification of a search query into an execution mode.
+
+    Analyze whether the query contains explicit constraints that map to
+    filterable metadata columns, or is a simple semantic search.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+    mode: Literal["standard", "instructed"] = Field(
+        description=(
+            "The execution mode. "
+            "Use 'standard' for simple semantic searches without constraints. "
+            "Use 'instructed' when the query contains explicit constraints "
+            "that can be translated to metadata filters."
+        )
+    )
+
+
+@mlflow.trace(name="route_query", span_type=SpanType.LLM)
+def route_query(
+    llm: BaseChatModel,
+    query: str,
+    schema_description: str,
+) -> Literal["standard", "instructed"]:
+    """
+    Determine the execution mode for a search query.
+
+    Args:
+        llm: Language model for routing decision
+        query: User's search query
+        schema_description: Column names, types, and filter syntax
+
+    Returns:
+        "standard" for simple queries, "instructed" for constrained queries
+    """
+    prompt_config = _load_prompt_template()
+    prompt_template = prompt_config["template"]
+
+    prompt = prompt_template.format(
+        schema_description=schema_description,
+        query=query,
+    )
+
+    logger.trace("Routing query", query=query[:100])
+
+    # Use LangChain's with_structured_output for automatic strategy selection
+    # (JSON schema vs tool calling based on model capabilities)
+    try:
+        structured_llm: Runnable[str, RouterDecision] = llm.with_structured_output(
+            RouterDecision
+        )
+        decision: RouterDecision = structured_llm.invoke(prompt)
+    except Exception as e:
+        logger.warning("Router failed, defaulting to standard mode", error=str(e))
+        return "standard"
+
+    logger.debug("Router decision", mode=decision.mode, query=query[:50])
+    mlflow.set_tag("router.mode", decision.mode)
+
+    return decision.mode

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/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/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,