dao-ai 0.0.28__py3-none-any.whl → 0.1.2__py3-none-any.whl

dao_ai/genie/__init__.py

@@ -0,0 +1,38 @@
+"""
+Genie service implementations and caching layers.
+
+This package provides core Genie functionality that can be used across
+different contexts (tools, direct integration, etc.).
+
+Main exports:
+- GenieService: Core service implementation wrapping Databricks Genie SDK
+- GenieServiceBase: Abstract base class for service implementations
+
+Cache implementations are available in the cache subpackage:
+- dao_ai.genie.cache.lru: LRU (Least Recently Used) cache
+- dao_ai.genie.cache.semantic: Semantic similarity cache using pg_vector
+
+Example usage:
+    from dao_ai.genie import GenieService
+    from dao_ai.genie.cache import LRUCacheService, SemanticCacheService
+"""
+
+from dao_ai.genie.cache import (
+    CacheResult,
+    GenieServiceBase,
+    LRUCacheService,
+    SemanticCacheService,
+    SQLCacheEntry,
+)
+from dao_ai.genie.core import GenieService
+
+__all__ = [
+    # Service classes
+    "GenieService",
+    "GenieServiceBase",
+    # Cache types (from cache subpackage)
+    "CacheResult",
+    "LRUCacheService",
+    "SemanticCacheService",
+    "SQLCacheEntry",
+]

dao_ai/genie/cache/__init__.py

@@ -0,0 +1,43 @@
+"""
+Genie cache implementations.
+
+This package provides caching layers for Genie SQL queries that can be
+chained together using the decorator pattern.
+
+Available cache implementations:
+- LRUCacheService: In-memory LRU cache with O(1) exact match lookup
+- SemanticCacheService: PostgreSQL pg_vector-based semantic similarity cache
+
+Example usage:
+    from dao_ai.genie.cache import LRUCacheService, SemanticCacheService
+
+    # Chain caches: LRU (checked first) -> Semantic (checked second) -> Genie
+    genie_service = SemanticCacheService(
+        impl=GenieService(genie),
+        parameters=semantic_params,
+    )
+    genie_service = LRUCacheService(
+        impl=genie_service,
+        parameters=lru_params,
+    )
+"""
+
+from dao_ai.genie.cache.base import (
+    CacheResult,
+    GenieServiceBase,
+    SQLCacheEntry,
+)
+from dao_ai.genie.cache.core import execute_sql_via_warehouse
+from dao_ai.genie.cache.lru import LRUCacheService
+from dao_ai.genie.cache.semantic import SemanticCacheService
+
+__all__ = [
+    # Base types
+    "CacheResult",
+    "GenieServiceBase",
+    "SQLCacheEntry",
+    "execute_sql_via_warehouse",
+    # Cache implementations
+    "LRUCacheService",
+    "SemanticCacheService",
+]

dao_ai/genie/cache/base.py

@@ -0,0 +1,72 @@
+"""
+Base classes and types for Genie cache implementations.
+
+This module provides the foundational types used across different cache
+implementations (LRU, Semantic, etc.). It contains only abstract base classes
+and data structures - no concrete implementations.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from databricks_ai_bridge.genie import GenieResponse
+
+if TYPE_CHECKING:
+    from dao_ai.genie.cache.base import CacheResult
+
+
+class GenieServiceBase(ABC):
+    """Abstract base class for Genie service implementations."""
+
+    @abstractmethod
+    def ask_question(
+        self, question: str, conversation_id: str | None = None
+    ) -> "CacheResult":
+        """
+        Ask a question to Genie and return the response with cache metadata.
+
+        All implementations return CacheResult to provide consistent cache information,
+        even when caching is disabled (cache_hit=False, served_by=None).
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def space_id(self) -> str:
+        """The space ID for the Genie service."""
+        pass
+
+
+@dataclass
+class SQLCacheEntry:
+    """
+    A cache entry storing the SQL query metadata for re-execution.
+
+    Instead of caching the full result, we cache the SQL query so that
+    on cache hit we can re-execute it to get fresh data.
+    """
+
+    query: str
+    description: str
+    conversation_id: str
+    created_at: datetime
+
+
+@dataclass
+class CacheResult:
+    """
+    Result of a cache-aware query with metadata about cache behavior.
+
+    Attributes:
+        response: The GenieResponse (fresh data, possibly from cached SQL)
+        cache_hit: Whether the SQL query came from cache
+        served_by: Name of the layer that served the cached SQL (None if from origin)
+    """
+
+    response: GenieResponse
+    cache_hit: bool
+    served_by: str | None = None

dao_ai/genie/cache/core.py

@@ -0,0 +1,79 @@
+"""
+Core utilities for Genie cache implementations.
+
+This module provides shared utility functions used by different cache
+implementations (LRU, Semantic, etc.). These are concrete implementations
+of common operations needed across cache types.
+"""
+
+from typing import Any
+
+import pandas as pd
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.sql import StatementResponse, StatementState
+from loguru import logger
+
+from dao_ai.config import WarehouseModel
+
+
+def execute_sql_via_warehouse(
+    warehouse: WarehouseModel,
+    sql: str,
+    layer_name: str = "cache",
+) -> pd.DataFrame | str:
+    """
+    Execute SQL using a Databricks warehouse and return results as DataFrame.
+
+    This is a shared utility for cache implementations that need to re-execute
+    cached SQL queries.
+
+    Args:
+        warehouse: The warehouse configuration for SQL execution
+        sql: The SQL query to execute
+        layer_name: Name of the cache layer (for logging)
+
+    Returns:
+        DataFrame with results, or error message string
+    """
+    w: WorkspaceClient = warehouse.workspace_client
+    warehouse_id: str = str(warehouse.warehouse_id)
+
+    logger.trace("Executing cached SQL", layer=layer_name, sql_prefix=sql[:100])
+
+    statement_response: StatementResponse = w.statement_execution.execute_statement(
+        statement=sql,
+        warehouse_id=warehouse_id,
+        wait_timeout="30s",
+    )
+
+    # Poll for completion if still running
+    while statement_response.status.state in [
+        StatementState.PENDING,
+        StatementState.RUNNING,
+    ]:
+        statement_response = w.statement_execution.get_statement(
+            statement_response.statement_id
+        )
+
+    if statement_response.status.state != StatementState.SUCCEEDED:
+        error_msg: str = f"SQL execution failed: {statement_response.status}"
+        logger.error(
+            "SQL execution failed",
+            layer=layer_name,
+            status=str(statement_response.status),
+        )
+        return error_msg
+
+    # Convert to DataFrame
+    if statement_response.result and statement_response.result.data_array:
+        columns: list[str] = []
+        if statement_response.manifest and statement_response.manifest.schema:
+            columns = [col.name for col in statement_response.manifest.schema.columns]
+
+        data: list[list[Any]] = statement_response.result.data_array
+        if columns:
+            return pd.DataFrame(data, columns=columns)
+        else:
+            return pd.DataFrame(data)
+
+    return pd.DataFrame()

dao_ai/genie/cache/lru.py

@@ -0,0 +1,347 @@
+"""
+LRU (Least Recently Used) cache implementation for Genie SQL queries.
+
+This module provides an in-memory LRU cache that stores SQL queries generated
+by Genie. On cache hit, the cached SQL is re-executed against the warehouse
+to return fresh data while avoiding the Genie NL-to-SQL translation cost.
+"""
+
+from collections import OrderedDict
+from datetime import datetime, timedelta
+from threading import Lock
+from typing import Any
+
+import mlflow
+import pandas as pd
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.sql import StatementResponse, StatementState
+from databricks_ai_bridge.genie import GenieResponse
+from loguru import logger
+
+from dao_ai.config import GenieLRUCacheParametersModel, WarehouseModel
+from dao_ai.genie.cache.base import (
+    CacheResult,
+    GenieServiceBase,
+    SQLCacheEntry,
+)
+
+
+class LRUCacheService(GenieServiceBase):
+    """
+    LRU caching decorator that caches SQL queries and re-executes them.
+
+    This service caches the SQL query generated by Genie (not the result data).
+    On cache hit, it re-executes the cached SQL using the provided warehouse
+    to return fresh data while avoiding the Genie NL-to-SQL translation cost.
+
+    Example:
+        from dao_ai.config import GenieLRUCacheParametersModel, WarehouseModel
+        from dao_ai.genie.cache import LRUCacheService
+
+        cache_params = GenieLRUCacheParametersModel(
+            warehouse=warehouse_model,
+            capacity=100,
+            time_to_live_seconds=86400  # 24 hours
+        )
+        genie = LRUCacheService(
+            impl=GenieService(Genie(space_id="my-space")),
+            parameters=cache_params
+        )
+
+    Thread-safe: Uses a lock to protect cache operations.
+    """
+
+    impl: GenieServiceBase
+    parameters: GenieLRUCacheParametersModel
+    name: str
+    _cache: OrderedDict[str, SQLCacheEntry]
+    _lock: Lock
+
+    def __init__(
+        self,
+        impl: GenieServiceBase,
+        parameters: GenieLRUCacheParametersModel,
+        name: str | None = None,
+    ) -> None:
+        """
+        Initialize the SQL cache service.
+
+        Args:
+            impl: The underlying GenieServiceBase to delegate to on cache miss
+            parameters: Cache configuration including warehouse, capacity, and TTL
+            name: Name for this cache layer (for logging). Defaults to class name.
+        """
+        self.impl = impl
+        self.parameters = parameters
+        self.name = name if name is not None else self.__class__.__name__
+        self._cache = OrderedDict()
+        self._lock = Lock()
+
+    @property
+    def warehouse(self) -> WarehouseModel:
+        """The warehouse used for executing cached SQL queries."""
+        return self.parameters.warehouse
+
+    @property
+    def capacity(self) -> int:
+        """Maximum number of SQL queries to cache."""
+        return self.parameters.capacity
+
+    @property
+    def time_to_live(self) -> timedelta | None:
+        """Duration after which cached queries expire. None means never expires."""
+        ttl = self.parameters.time_to_live_seconds
+        if ttl is None or ttl < 0:
+            return None
+        return timedelta(seconds=ttl)
+
+    @staticmethod
+    def _normalize_key(question: str, conversation_id: str | None = None) -> str:
+        """
+        Normalize the question and conversation_id to create a consistent cache key.
+
+        Args:
+            question: The question text
+            conversation_id: Optional conversation ID to include in the key
+
+        Returns:
+            A normalized cache key combining question and conversation_id
+        """
+        normalized_question = question.strip().lower()
+        if conversation_id:
+            return f"{conversation_id}::{normalized_question}"
+        return normalized_question
+
+    def _is_expired(self, entry: SQLCacheEntry) -> bool:
+        """Check if a cache entry has exceeded its TTL. Returns False if TTL is disabled."""
+        if self.time_to_live is None:
+            return False
+        age: timedelta = datetime.now() - entry.created_at
+        return age > self.time_to_live
+
+    def _evict_oldest(self) -> None:
+        """Remove the oldest (least recently used) entry."""
+        if self._cache:
+            oldest_key: str = next(iter(self._cache))
+            del self._cache[oldest_key]
+            logger.trace(
+                "Evicted cache entry", layer=self.name, key_prefix=oldest_key[:50]
+            )
+
+    def _get(self, key: str) -> SQLCacheEntry | None:
+        """Get from cache, returning None if not found or expired."""
+        if key not in self._cache:
+            return None
+
+        entry: SQLCacheEntry = self._cache[key]
+
+        if self._is_expired(entry):
+            del self._cache[key]
+            logger.trace("Expired cache entry", layer=self.name, key_prefix=key[:50])
+            return None
+
+        self._cache.move_to_end(key)
+        return entry
+
+    def _put(self, key: str, response: GenieResponse) -> None:
+        """Store SQL query in cache, evicting if at capacity."""
+        if key in self._cache:
+            del self._cache[key]
+
+        while len(self._cache) >= self.capacity:
+            self._evict_oldest()
+
+        self._cache[key] = SQLCacheEntry(
+            query=response.query,
+            description=response.description,
+            conversation_id=response.conversation_id,
+            created_at=datetime.now(),
+        )
+        logger.info(
+            "Stored cache entry",
+            layer=self.name,
+            key_prefix=key[:50],
+            sql_prefix=response.query[:50] if response.query else None,
+            cache_size=len(self._cache),
+            capacity=self.capacity,
+        )
+
+    @mlflow.trace(name="execute_cached_sql")
+    def _execute_sql(self, sql: str) -> pd.DataFrame | str:
+        """
+        Execute SQL using the warehouse and return results as DataFrame.
+
+        Args:
+            sql: The SQL query to execute
+
+        Returns:
+            DataFrame with results, or error message string
+        """
+        w: WorkspaceClient = self.warehouse.workspace_client
+        warehouse_id: str = str(self.warehouse.warehouse_id)
+
+        logger.trace("Executing cached SQL", layer=self.name, sql_prefix=sql[:100])
+
+        statement_response: StatementResponse = w.statement_execution.execute_statement(
+            statement=sql,
+            warehouse_id=warehouse_id,
+            wait_timeout="30s",
+        )
+
+        # Poll for completion if still running
+        while statement_response.status.state in [
+            StatementState.PENDING,
+            StatementState.RUNNING,
+        ]:
+            statement_response = w.statement_execution.get_statement(
+                statement_response.statement_id
+            )
+
+        if statement_response.status.state != StatementState.SUCCEEDED:
+            error_msg: str = f"SQL execution failed: {statement_response.status}"
+            logger.error(
+                "SQL execution failed",
+                layer=self.name,
+                status=str(statement_response.status),
+            )
+            return error_msg
+
+        # Convert to DataFrame
+        if statement_response.result and statement_response.result.data_array:
+            columns: list[str] = []
+            if statement_response.manifest and statement_response.manifest.schema:
+                columns = [
+                    col.name for col in statement_response.manifest.schema.columns
+                ]
+
+            data: list[list[Any]] = statement_response.result.data_array
+            if columns:
+                return pd.DataFrame(data, columns=columns)
+            else:
+                return pd.DataFrame(data)
+
+        return pd.DataFrame()
+
+    def ask_question(
+        self, question: str, conversation_id: str | None = None
+    ) -> CacheResult:
+        """
+        Ask a question, using cached SQL query if available.
+
+        On cache hit, re-executes the cached SQL to get fresh data.
+        Returns CacheResult with cache metadata.
+        """
+        return self.ask_question_with_cache_info(question, conversation_id)
+
+    @mlflow.trace(name="genie_lru_cache_lookup")
+    def ask_question_with_cache_info(
+        self,
+        question: str,
+        conversation_id: str | None = None,
+    ) -> CacheResult:
+        """
+        Ask a question with detailed cache hit information.
+
+        On cache hit, the cached SQL is re-executed to return fresh data.
+
+        Args:
+            question: The question to ask
+            conversation_id: Optional conversation ID
+
+        Returns:
+            CacheResult with fresh response and cache metadata
+        """
+        key: str = self._normalize_key(question, conversation_id)
+
+        # Check cache
+        with self._lock:
+            cached: SQLCacheEntry | None = self._get(key)
+
+        if cached is not None:
+            logger.info(
+                "Cache HIT",
+                layer=self.name,
+                question_prefix=question[:50],
+                conversation_id=conversation_id,
+                cache_size=self.size,
+                capacity=self.capacity,
+            )
+
+            # Re-execute the cached SQL to get fresh data
+            result: pd.DataFrame | str = self._execute_sql(cached.query)
+
+            # Use current conversation_id, not the cached one
+            response: GenieResponse = GenieResponse(
+                result=result,
+                query=cached.query,
+                description=cached.description,
+                conversation_id=conversation_id
+                if conversation_id
+                else cached.conversation_id,
+            )
+
+            return CacheResult(response=response, cache_hit=True, served_by=self.name)
+
+        # Cache miss - delegate to wrapped service
+        logger.info(
+            "Cache MISS",
+            layer=self.name,
+            question_prefix=question[:50],
+            conversation_id=conversation_id,
+            cache_size=self.size,
+            capacity=self.capacity,
+            delegating_to=type(self.impl).__name__,
+        )
+
+        result: CacheResult = self.impl.ask_question(question, conversation_id)
+        with self._lock:
+            self._put(key, result.response)
+        return CacheResult(response=result.response, cache_hit=False, served_by=None)
+
+    @property
+    def space_id(self) -> str:
+        return self.impl.space_id
+
+    def invalidate(self, question: str, conversation_id: str | None = None) -> bool:
+        """
+        Remove a specific entry from the cache.
+
+        Args:
+            question: The question text
+            conversation_id: Optional conversation ID to match
+
+        Returns:
+            True if the entry was found and removed, False otherwise
+        """
+        key: str = self._normalize_key(question, conversation_id)
+        with self._lock:
+            if key in self._cache:
+                del self._cache[key]
+                return True
+            return False
+
+    def clear(self) -> int:
+        """Clear all entries from the cache."""
+        with self._lock:
+            count: int = len(self._cache)
+            self._cache.clear()
+            return count
+
+    @property
+    def size(self) -> int:
+        """Current number of entries in the cache."""
+        with self._lock:
+            return len(self._cache)
+
+    def stats(self) -> dict[str, int | float | None]:
+        """Return cache statistics."""
+        with self._lock:
+            expired: int = sum(1 for e in self._cache.values() if self._is_expired(e))
+            ttl = self.time_to_live
+            return {
+                "size": len(self._cache),
+                "capacity": self.capacity,
+                "ttl_seconds": ttl.total_seconds() if ttl else None,
+                "expired_entries": expired,
+                "valid_entries": len(self._cache) - expired,
+            }