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

dao_ai/config.py

@@ -28,8 +28,10 @@ from databricks.sdk.service.database import DatabaseInstance
 from databricks.vector_search.client import VectorSearchClient
 from databricks.vector_search.index import VectorSearchIndex
 from databricks_langchain import (
+    DatabricksEmbeddings,
     DatabricksFunctionClient,
 )
+from langchain_core.embeddings import Embeddings
 from langchain_core.language_models import LanguageModelLike
 from langchain_core.messages import BaseMessage, messages_from_dict
 from langchain_core.runnables.base import RunnableLike
@@ -408,6 +410,9 @@ class LLMModel(BaseModel, IsDatabricksResource):
 
         return chat_client
 
+    def as_embeddings_model(self) -> Embeddings:
+        return DatabricksEmbeddings(endpoint=self.name)
+
 
 class VectorSearchEndpointType(str, Enum):
     STANDARD = "STANDARD"
@@ -977,6 +982,30 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
         provider.create_lakebase_instance_role(self)
 
 
+class GenieLRUCacheParametersModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    capacity: int = 1000
+    time_to_live_seconds: int | None = (
+        60 * 60 * 24
+    )  # 1 day default, None or negative = never expires
+    warehouse: WarehouseModel
+
+
+class GenieSemanticCacheParametersModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    time_to_live_seconds: int | None = (
+        60 * 60 * 24
+    )  # 1 day default, None or negative = never expires
+    similarity_threshold: float = (
+        0.85  # Minimum similarity for cache hit (L2 distance converted to 0-1 scale)
+    )
+    embedding_model: str | LLMModel = "databricks-gte-large-en"
+    embedding_dims: int | None = None  # Auto-detected if None
+    database: DatabaseModel
+    warehouse: WarehouseModel
+    table_name: str = "genie_semantic_cache"
+
+
 class SearchParametersModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     num_results: Optional[int] = 10

dao_ai/genie/__init__.py

@@ -0,0 +1,59 @@
+"""
+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
+"""
+
+import mlflow
+from databricks_ai_bridge.genie import Genie, GenieResponse
+
+from dao_ai.genie.cache import (
+    CacheResult,
+    GenieServiceBase,
+    LRUCacheService,
+    SemanticCacheService,
+    SQLCacheEntry,
+)
+
+
+class GenieService(GenieServiceBase):
+    """Concrete implementation of GenieServiceBase using the Genie SDK."""
+
+    genie: Genie
+
+    def __init__(self, genie: Genie) -> None:
+        self.genie = genie
+
+    @mlflow.trace(name="genie_ask_question")
+    def ask_question(
+        self, question: str, conversation_id: str | None = None
+    ) -> GenieResponse:
+        response: GenieResponse = self.genie.ask_question(
+            question, conversation_id=conversation_id
+        )
+        return response
+
+
+__all__ = [
+    # Service classes
+    "GenieService",
+    "GenieServiceBase",
+    # Cache types (from cache subpackage)
+    "CacheResult",
+    "LRUCacheService",
+    "SemanticCacheService",
+    "SQLCacheEntry",
+]

dao_ai/genie/cache/__init__.py

@@ -0,0 +1,44 @@
+"""
+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_space_id=space_id,
+    )
+    genie_service = LRUCacheService(
+        impl=genie_service,
+        parameters=lru_params,
+    )
+"""
+
+from dao_ai.genie.cache.base import (
+    CacheResult,
+    GenieServiceBase,
+    SQLCacheEntry,
+    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,122 @@
+"""
+Base classes and types for Genie cache implementations.
+
+This module provides the foundational types used across different cache
+implementations (LRU, Semantic, etc.).
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+
+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 WarehouseModel
+
+
+class GenieServiceBase(ABC):
+    """Abstract base class for Genie service implementations."""
+
+    @abstractmethod
+    def ask_question(
+        self, question: str, conversation_id: str | None = None
+    ) -> GenieResponse:
+        """Ask a question to Genie and return the response."""
+        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
+
+
+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.debug(f"[{layer_name}] Executing cached SQL: {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(f"[{layer_name}] {error_msg}")
+        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]
+        elif hasattr(statement_response.result, "schema"):
+            columns = [col.name for col in statement_response.result.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,306 @@
+"""
+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) -> str:
+        """Normalize the question to create a consistent cache key."""
+        return question.strip().lower()
+
+    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.debug(f"[{self.name}] Evicted: {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.debug(f"[{self.name}] Expired: {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(
+            f"[{self.name}] Stored cache entry: key='{key[:50]}...' "
+            f"sql='{response.query[:50] if response.query else 'None'}...' "
+            f"(cache_size={len(self._cache)}/{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.debug(f"[{self.name}] Executing cached SQL: {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(f"[{self.name}] {error_msg}")
+            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
+                ]
+            elif hasattr(statement_response.result, "schema"):
+                columns = [col.name for col in statement_response.result.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
+    ) -> GenieResponse:
+        """
+        Ask a question, using cached SQL query if available.
+
+        On cache hit, re-executes the cached SQL to get fresh data.
+        Implements GenieServiceBase for seamless chaining.
+        """
+        result: CacheResult = self.ask_question_with_cache_info(
+            question, conversation_id
+        )
+        return result.response
+
+    @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)
+
+        # Check cache
+        with self._lock:
+            cached: SQLCacheEntry | None = self._get(key)
+
+        if cached is not None:
+            logger.info(
+                f"[{self.name}] Cache HIT: '{question[:50]}...' "
+                f"(cache_size={self.size}/{self.capacity})"
+            )
+
+            # Re-execute the cached SQL to get fresh data
+            result: pd.DataFrame | str = self._execute_sql(cached.query)
+
+            response: GenieResponse = GenieResponse(
+                result=result,
+                query=cached.query,
+                description=cached.description,
+                conversation_id=cached.conversation_id,
+            )
+
+            return CacheResult(response=response, cache_hit=True, served_by=self.name)
+
+        # Cache miss - delegate to wrapped service
+        logger.info(
+            f"[{self.name}] Cache MISS: '{question[:50]}...' "
+            f"(cache_size={self.size}/{self.capacity}, delegating to {type(self.impl).__name__})"
+        )
+
+        response = self.impl.ask_question(question, conversation_id)
+        with self._lock:
+            self._put(key, response)
+        return CacheResult(response=response, cache_hit=False, served_by=None)
+
+    def invalidate(self, question: str) -> bool:
+        """Remove a specific entry from the cache."""
+        key: str = self._normalize_key(question)
+        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,
+            }