dao-ai 0.0.34__py3-none-any.whl → 0.0.36__py3-none-any.whl

dao_ai/tools/__init__.py

@@ -1,3 +1,4 @@
+from dao_ai.genie.cache import LRUCacheService, SemanticCacheService
 from dao_ai.hooks.core import create_hooks
 from dao_ai.tools.agent import create_agent_endpoint_tool
 from dao_ai.tools.core import (
@@ -35,7 +36,9 @@ __all__ = [
     "current_time_tool",
     "format_time_tool",
     "is_business_hours_tool",
+    "LRUCacheService",
     "search_tool",
+    "SemanticCacheService",
     "time_difference_tool",
     "time_in_timezone_tool",
     "time_until_tool",

dao_ai/tools/genie/__init__.py

@@ -0,0 +1,236 @@
+"""
+Genie tools for natural language queries to databases.
+
+This package provides tools for interacting with Databricks Genie to translate
+natural language questions into SQL queries.
+
+Main exports:
+- create_genie_tool: Factory function to create a Genie tool with optional caching
+
+Cache implementations are available in the genie cache package:
+- dao_ai.genie.cache.lru: LRU (Least Recently Used) cache
+- dao_ai.genie.cache.semantic: Semantic similarity cache using pg_vector
+"""
+
+import json
+import os
+from textwrap import dedent
+from typing import Annotated, Any, Callable
+
+import pandas as pd
+from databricks_ai_bridge.genie import Genie, GenieResponse
+from langchain.tools import tool
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import InjectedToolCallId
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from loguru import logger
+from pydantic import BaseModel
+
+from dao_ai.config import (
+    AnyVariable,
+    CompositeVariableModel,
+    GenieLRUCacheParametersModel,
+    GenieRoomModel,
+    GenieSemanticCacheParametersModel,
+    value_of,
+)
+from dao_ai.genie import GenieService
+from dao_ai.genie.cache import (
+    CacheResult,
+    GenieServiceBase,
+    LRUCacheService,
+    SemanticCacheService,
+    SQLCacheEntry,
+)
+
+
+class GenieToolInput(BaseModel):
+    """Input schema for Genie tool - only includes user-facing parameters."""
+
+    question: str
+
+
+def _response_to_json(response: GenieResponse) -> str:
+    """Convert GenieResponse to JSON string, handling DataFrame results."""
+    # Convert result to string if it's a DataFrame
+    result: str | pd.DataFrame = response.result
+    if isinstance(result, pd.DataFrame):
+        result = result.to_markdown()
+
+    data: dict[str, Any] = {
+        "result": result,
+        "query": response.query,
+        "description": response.description,
+        "conversation_id": response.conversation_id,
+    }
+    return json.dumps(data)
+
+
+def create_genie_tool(
+    genie_room: GenieRoomModel | dict[str, Any],
+    name: str | None = None,
+    description: str | None = None,
+    persist_conversation: bool = True,
+    truncate_results: bool = False,
+    lru_cache_parameters: GenieLRUCacheParametersModel | dict[str, Any] | None = None,
+    semantic_cache_parameters: GenieSemanticCacheParametersModel
+    | dict[str, Any]
+    | None = None,
+) -> Callable[..., Command]:
+    """
+    Create a tool for interacting with Databricks Genie for natural language queries to databases.
+
+    This factory function generates a tool that leverages Databricks Genie to translate natural
+    language questions into SQL queries and execute them against retail databases. This enables
+    answering questions about inventory, sales, and other structured retail data.
+
+    Args:
+        genie_room: GenieRoomModel or dict containing Genie configuration
+        name: Optional custom name for the tool. If None, uses default "genie_tool"
+        description: Optional custom description for the tool. If None, uses default description
+        persist_conversation: Whether to persist conversation IDs across tool calls for
+            multi-turn conversations within the same Genie space
+        truncate_results: Whether to truncate large query results to fit token limits
+        lru_cache_parameters: Optional LRU cache configuration for SQL query caching
+        semantic_cache_parameters: Optional semantic cache configuration using pg_vector
+            for similarity-based query matching
+
+    Returns:
+        A LangGraph tool that processes natural language queries through Genie
+    """
+    logger.debug("create_genie_tool")
+    logger.debug(f"genie_room type: {type(genie_room)}")
+    logger.debug(f"genie_room: {genie_room}")
+    logger.debug(f"persist_conversation: {persist_conversation}")
+    logger.debug(f"truncate_results: {truncate_results}")
+    logger.debug(f"name: {name}")
+    logger.debug(f"description: {description}")
+    logger.debug(f"genie_room: {genie_room}")
+    logger.debug(f"persist_conversation: {persist_conversation}")
+    logger.debug(f"truncate_results: {truncate_results}")
+    logger.debug(f"lru_cache_parameters: {lru_cache_parameters}")
+    logger.debug(f"semantic_cache_parameters: {semantic_cache_parameters}")
+
+    if isinstance(genie_room, dict):
+        genie_room = GenieRoomModel(**genie_room)
+
+    if isinstance(lru_cache_parameters, dict):
+        lru_cache_parameters = GenieLRUCacheParametersModel(**lru_cache_parameters)
+
+    if isinstance(semantic_cache_parameters, dict):
+        semantic_cache_parameters = GenieSemanticCacheParametersModel(
+            **semantic_cache_parameters
+        )
+
+    space_id: AnyVariable = genie_room.space_id or os.environ.get(
+        "DATABRICKS_GENIE_SPACE_ID"
+    )
+    if isinstance(space_id, dict):
+        space_id = CompositeVariableModel(**space_id)
+    space_id = value_of(space_id)
+
+    default_description: str = dedent("""
+    This tool lets you have a conversation and chat with tabular data about <topic>. You should ask
+    questions about the data and the tool will try to answer them.
+    Please ask simple clear questions that can be answer by sql queries. If you need to do statistics or other forms of testing defer to using another tool.
+    Try to ask for aggregations on the data and ask very simple questions.
+    Prefer to call this tool multiple times rather than asking a complex question.
+    """)
+
+    tool_description: str = (
+        description if description is not None else default_description
+    )
+    tool_name: str = name if name is not None else "genie_tool"
+
+    function_docs = """
+
+Args:
+question (str): The question to ask to ask Genie about your data. Ask simple, clear questions about your tabular data. For complex analysis, ask multiple simple questions rather than one complex question.
+
+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,
+    )
+
+    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,
+            genie_space_id=space_id,
+        )
+
+    # 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,
+        )
+
+    @tool(
+        name_or_callable=tool_name,
+        description=tool_description,
+    )
+    def genie_tool(
+        question: Annotated[str, "The question to ask Genie about your data"],
+        state: Annotated[dict, InjectedState],
+        tool_call_id: Annotated[str, InjectedToolCallId],
+    ) -> Command:
+        """Process a natural language question through Databricks Genie."""
+        # Get existing conversation mapping and retrieve conversation ID for this space
+        conversation_ids: dict[str, str] = state.get("genie_conversation_ids", {})
+        existing_conversation_id: str | None = conversation_ids.get(space_id)
+        logger.debug(
+            f"Existing conversation ID for space {space_id}: {existing_conversation_id}"
+        )
+
+        response: GenieResponse = genie_service.ask_question(
+            question, conversation_id=existing_conversation_id
+        )
+
+        current_conversation_id: str = response.conversation_id
+        logger.debug(
+            f"Current conversation ID for space {space_id}: {current_conversation_id}"
+        )
+
+        # Update the conversation mapping with the new conversation ID for this space
+
+        update: dict[str, Any] = {
+            "messages": [
+                ToolMessage(_response_to_json(response), tool_call_id=tool_call_id)
+            ],
+        }
+
+        if persist_conversation:
+            updated_conversation_ids: dict[str, str] = conversation_ids.copy()
+            updated_conversation_ids[space_id] = current_conversation_id
+            update["genie_conversation_ids"] = updated_conversation_ids
+
+        return Command(update=update)
+
+    return genie_tool
+
+
+# Re-export cache types for convenience
+__all__ = [
+    # Main tool
+    "create_genie_tool",
+    # Input types
+    "GenieToolInput",
+    # Service base classes
+    "GenieService",
+    "GenieServiceBase",
+    # Cache types (from cache subpackage)
+    "CacheResult",
+    "LRUCacheService",
+    "SemanticCacheService",
+    "SQLCacheEntry",
+]

dao_ai/tools/genie.py

@@ -1,3 +1,14 @@
+"""
+Genie tool for natural language queries to databases.
+
+This module provides the tool factory for creating LangGraph tools that
+interact with Databricks Genie.
+
+For the core Genie service and cache implementations, see:
+- dao_ai.genie: GenieService, GenieServiceBase
+- dao_ai.genie.cache: LRUCacheService, SemanticCacheService
+"""
+
 import json
 import os
 from textwrap import dedent
@@ -11,17 +22,24 @@ from langchain_core.tools import InjectedToolCallId
 from langgraph.prebuilt import InjectedState
 from langgraph.types import Command
 from loguru import logger
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
-from dao_ai.config import AnyVariable, CompositeVariableModel, GenieRoomModel, value_of
+from dao_ai.config import (
+    AnyVariable,
+    CompositeVariableModel,
+    GenieLRUCacheParametersModel,
+    GenieRoomModel,
+    GenieSemanticCacheParametersModel,
+    value_of,
+)
+from dao_ai.genie import GenieService, GenieServiceBase
+from dao_ai.genie.cache import LRUCacheService, SemanticCacheService
 
 
 class GenieToolInput(BaseModel):
-    """Input schema for the Genie tool."""
+    """Input schema for Genie tool - only includes user-facing parameters."""
 
-    question: str = Field(
-        description="The question to ask Genie about your data. Ask simple, clear questions about your tabular data. For complex analysis, ask multiple simple questions rather than one complex question."
-    )
+    question: str
 
 
 def _response_to_json(response: GenieResponse) -> str:
@@ -46,6 +64,10 @@ def create_genie_tool(
     description: str | None = None,
     persist_conversation: bool = True,
     truncate_results: bool = False,
+    lru_cache_parameters: GenieLRUCacheParametersModel | dict[str, Any] | None = None,
+    semantic_cache_parameters: GenieSemanticCacheParametersModel
+    | dict[str, Any]
+    | None = None,
 ) -> Callable[..., Command]:
     """
     Create a tool for interacting with Databricks Genie for natural language queries to databases.
@@ -61,6 +83,9 @@ def create_genie_tool(
         persist_conversation: Whether to persist conversation IDs across tool calls for
             multi-turn conversations within the same Genie space
         truncate_results: Whether to truncate large query results to fit token limits
+        lru_cache_parameters: Optional LRU cache configuration for SQL query caching
+        semantic_cache_parameters: Optional semantic cache configuration using pg_vector
+            for similarity-based query matching
 
     Returns:
         A LangGraph tool that processes natural language queries through Genie
@@ -75,10 +100,20 @@ def create_genie_tool(
     logger.debug(f"genie_room: {genie_room}")
     logger.debug(f"persist_conversation: {persist_conversation}")
     logger.debug(f"truncate_results: {truncate_results}")
+    logger.debug(f"lru_cache_parameters: {lru_cache_parameters}")
+    logger.debug(f"semantic_cache_parameters: {semantic_cache_parameters}")
 
     if isinstance(genie_room, dict):
         genie_room = GenieRoomModel(**genie_room)
 
+    if isinstance(lru_cache_parameters, dict):
+        lru_cache_parameters = GenieLRUCacheParametersModel(**lru_cache_parameters)
+
+    if isinstance(semantic_cache_parameters, dict):
+        semantic_cache_parameters = GenieSemanticCacheParametersModel(
+            **semantic_cache_parameters
+        )
+
     space_id: AnyVariable = genie_room.space_id or os.environ.get(
         "DATABRICKS_GENIE_SPACE_ID"
     )
@@ -108,6 +143,29 @@ 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,
+    )
+
+    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,
+            genie_space_id=space_id,
+        ).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,
+        )
+
     @tool(
         name_or_callable=tool_name,
         description=tool_description,
@@ -117,12 +175,6 @@ GenieResponse: A response object containing the conversation ID and result from
         state: Annotated[dict, InjectedState],
         tool_call_id: Annotated[str, InjectedToolCallId],
     ) -> Command:
-        genie: Genie = Genie(
-            space_id=space_id,
-            client=genie_room.workspace_client,
-            truncate_results=truncate_results,
-        )
-
         """Process a natural language question through Databricks Genie."""
         # Get existing conversation mapping and retrieve conversation ID for this space
         conversation_ids: dict[str, str] = state.get("genie_conversation_ids", {})
@@ -131,7 +183,7 @@ GenieResponse: A response object containing the conversation ID and result from
             f"Existing conversation ID for space {space_id}: {existing_conversation_id}"
         )
 
-        response: GenieResponse = genie.ask_question(
+        response: GenieResponse = genie_service.ask_question(
             question, conversation_id=existing_conversation_id
         )
 
@@ -153,8 +205,6 @@ GenieResponse: A response object containing the conversation ID and result from
             updated_conversation_ids[space_id] = current_conversation_id
             update["genie_conversation_ids"] = updated_conversation_ids
 
-        logger.debug(f"State update: {update}")
-
         return Command(update=update)
 
     return genie_tool

dao_ai/tools/unity_catalog.py

@@ -15,6 +15,7 @@ from dao_ai.config import (
     value_of,
 )
 from dao_ai.tools.human_in_the_loop import as_human_in_the_loop
+from dao_ai.utils import normalize_host
 
 
 def create_uc_tools(
@@ -299,9 +300,11 @@ def with_partial_args(
             if "client_secret" not in resolved_args:
                 resolved_args["client_secret"] = value_of(sp.client_secret)
 
-    # Normalize host/workspace_host - accept either key
+    # Normalize host/workspace_host - accept either key, ensure https:// scheme
     if "workspace_host" in resolved_args and "host" not in resolved_args:
-        resolved_args["host"] = resolved_args.pop("workspace_host")
+        resolved_args["host"] = normalize_host(resolved_args.pop("workspace_host"))
+    elif "host" in resolved_args:
+        resolved_args["host"] = normalize_host(resolved_args["host"])
 
     # Default host from WorkspaceClient if not provided
     if "host" not in resolved_args:

dao_ai/tools/vector_search.py

@@ -20,6 +20,7 @@ from dao_ai.config import (
     RetrieverModel,
     VectorStoreModel,
 )
+from dao_ai.utils import normalize_host
 
 
 def create_vector_search_tool(
@@ -108,8 +109,9 @@ def create_vector_search_tool(
     # The workspace_client parameter in DatabricksVectorSearch is only used to detect
     # model serving mode - it doesn't pass credentials to VectorSearchClient.
     client_args: dict[str, Any] = {}
-    if os.environ.get("DATABRICKS_HOST"):
-        client_args["workspace_url"] = os.environ.get("DATABRICKS_HOST")
+    databricks_host = normalize_host(os.environ.get("DATABRICKS_HOST"))
+    if databricks_host:
+        client_args["workspace_url"] = databricks_host
     if os.environ.get("DATABRICKS_TOKEN"):
         client_args["personal_access_token"] = os.environ.get("DATABRICKS_TOKEN")
     if os.environ.get("DATABRICKS_CLIENT_ID"):

dao_ai/utils.py

@@ -38,6 +38,30 @@ def normalize_name(name: str) -> str:
     return normalized.strip("_")
 
 
+def normalize_host(host: str | None) -> str | None:
+    """Ensure host URL has https:// scheme.
+
+    The DATABRICKS_HOST environment variable should always include the https://
+    scheme, but some environments (e.g., Databricks Apps infrastructure) may
+    provide the host without it. This function normalizes the host to ensure
+    it has the proper scheme.
+
+    Args:
+        host: The host URL, with or without scheme
+
+    Returns:
+        The host URL with https:// scheme, or None if host is None/empty
+    """
+    if not host:
+        return None
+    host = host.strip()
+    if not host:
+        return None
+    if not host.startswith("http://") and not host.startswith("https://"):
+        return f"https://{host}"
+    return host
+
+
 def get_default_databricks_host() -> str | None:
     """Get the default Databricks workspace host.
 
@@ -46,19 +70,19 @@ def get_default_databricks_host() -> str | None:
     2. WorkspaceClient ambient authentication (e.g., from ~/.databrickscfg)
 
     Returns:
-        The Databricks workspace host URL, or None if not available.
+        The Databricks workspace host URL (with https:// scheme), or None if not available.
     """
     # Try environment variable first
     host: str | None = os.environ.get("DATABRICKS_HOST")
     if host:
-        return host
+        return normalize_host(host)
 
     # Fall back to WorkspaceClient
     try:
         from databricks.sdk import WorkspaceClient
 
         w: WorkspaceClient = WorkspaceClient()
-        return w.config.host
+        return normalize_host(w.config.host)
     except Exception:
         logger.debug("Could not get default Databricks host from WorkspaceClient")
         return None