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

dao_ai/genie/cache/context_aware/in_memory.py

@@ -0,0 +1,609 @@
+"""
+In-memory context-aware Genie cache implementation.
+
+This module provides a context-aware cache that stores embeddings and cache entries
+entirely in memory, without requiring external database dependencies like PostgreSQL
+or Databricks Lakebase. It uses L2 distance for similarity search and supports
+dual embedding matching (question + conversation context).
+
+Use this when:
+- No external database access is available
+- Single-instance deployments (cache not shared across instances)
+- Cache persistence across restarts is not required
+- Cache sizes are moderate (hundreds to low thousands of entries)
+
+For multi-instance deployments or large cache sizes, use PostgresContextAwareGenieService
+with PostgreSQL backend instead.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from threading import Lock
+from typing import Any
+
+import mlflow
+import numpy as np
+from databricks.sdk import WorkspaceClient
+from databricks_ai_bridge.genie import GenieResponse
+from loguru import logger
+
+from dao_ai.config import (
+    GenieInMemorySemanticCacheParametersModel,
+    WarehouseModel,
+)
+from dao_ai.genie.cache.base import (
+    GenieServiceBase,
+    SQLCacheEntry,
+)
+from dao_ai.genie.cache.context_aware.base import ContextAwareGenieService
+
+
+@dataclass
+class InMemoryCacheEntry:
+    """
+    In-memory cache entry storing embeddings and SQL query metadata.
+
+    This dataclass represents a single cache entry stored in memory, including
+    dual embeddings (question + context) for high-precision semantic matching.
+
+    Uses LRU (Least Recently Used) eviction strategy when capacity is reached.
+
+    Attributes:
+        genie_space_id: The Genie space ID this entry belongs to
+        question: The original question text
+        conversation_context: Previous conversation context for embedding
+        question_embedding: Embedding vector for the question
+        context_embedding: Embedding vector for the conversation context
+        sql_query: The SQL query to re-execute on cache hit
+        description: Description of the query
+        conversation_id: The conversation ID where this query originated
+        created_at: When the entry was created
+        last_accessed_at: Last access time for LRU eviction
+        message_id: The original Genie message ID (for feedback on cache hits)
+    """
+
+    genie_space_id: str
+    question: str
+    conversation_context: str
+    question_embedding: list[float]
+    context_embedding: list[float]
+    sql_query: str
+    description: str
+    conversation_id: str
+    created_at: datetime
+    last_accessed_at: datetime  # Track last access time for LRU eviction
+    message_id: str | None = None  # Original Genie message ID for feedback
+
+
+def l2_distance(a: list[float], b: list[float]) -> float:
+    """
+    Calculate L2 (Euclidean) distance between two embedding vectors.
+
+    This uses the same distance metric as PostgreSQL pg_vector to ensure
+    consistent behavior between in-memory and PostgreSQL caches.
+
+    Args:
+        a: First embedding vector
+        b: Second embedding vector
+
+    Returns:
+        L2 distance (0 = identical vectors, larger = more different)
+    """
+    return float(np.linalg.norm(np.array(a) - np.array(b)))
+
+
+def distance_to_similarity(distance: float) -> float:
+    """
+    Convert L2 distance to similarity score in range [0, 1].
+
+    Uses the formula: similarity = 1.0 / (1.0 + distance)
+    This matches the conversion used by PostgreSQL semantic cache.
+
+    Args:
+        distance: L2 distance value
+
+    Returns:
+        Similarity score where 1.0 = perfect match, approaching 0 = very different
+    """
+    return 1.0 / (1.0 + distance)
+
+
+class InMemoryContextAwareGenieService(ContextAwareGenieService):
+    """
+    In-memory context-aware caching decorator using dual embeddings for similarity lookup.
+
+    This service caches the SQL query generated by Genie along with dual embeddings
+    (question + conversation context) for high-precision semantic matching. On
+    subsequent queries, it performs similarity search to find cached queries that
+    match both the question intent AND conversation context.
+
+    Cache entries are partitioned by genie_space_id to ensure queries from different
+    Genie spaces don't return incorrect cache hits.
+
+    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 GenieInMemorySemanticCacheParametersModel
+        from dao_ai.genie.cache.context_aware import InMemoryContextAwareGenieService
+
+        cache_params = GenieInMemorySemanticCacheParametersModel(
+            warehouse=warehouse_model,
+            embedding_model="databricks-gte-large-en",
+            time_to_live_seconds=86400,  # 24 hours
+            similarity_threshold=0.85,
+            capacity=1000,  # Limit to 1000 entries
+        )
+        genie = InMemoryContextAwareGenieService(
+            impl=GenieService(Genie(space_id="my-space")),
+            parameters=cache_params,
+            workspace_client=workspace_client,
+        ).initialize()
+
+    Thread-safe: Uses a lock to protect cache operations.
+    """
+
+    impl: GenieServiceBase
+    parameters: GenieInMemorySemanticCacheParametersModel
+    _workspace_client: WorkspaceClient | None
+    name: str
+    _embeddings: Any  # DatabricksEmbeddings
+    _cache: list[InMemoryCacheEntry]
+    _lock: Lock
+    _embedding_dims: int | None
+    _setup_complete: bool
+
+    def __init__(
+        self,
+        impl: GenieServiceBase,
+        parameters: GenieInMemorySemanticCacheParametersModel,
+        workspace_client: WorkspaceClient | None = None,
+        name: str | None = None,
+    ) -> None:
+        """
+        Initialize the in-memory context-aware cache service.
+
+        Args:
+            impl: The underlying GenieServiceBase to delegate to on cache miss.
+                The space_id will be obtained from impl.space_id.
+            parameters: Cache configuration including warehouse, embedding model, and thresholds
+            workspace_client: Optional WorkspaceClient for retrieving conversation history.
+                If None, conversation context will not be used.
+            name: Name for this cache layer (for logging). Defaults to class name.
+        """
+        self.impl = impl
+        self.parameters = parameters
+        self._workspace_client = workspace_client
+        self.name = name if name is not None else self.__class__.__name__
+        self._embeddings = None
+        self._cache = []
+        self._lock = Lock()
+        self._embedding_dims = None
+        self._setup_complete = False
+
+    def _setup(self) -> None:
+        """Initialize embeddings model lazily."""
+        if self._setup_complete:
+            return
+
+        # Initialize embeddings using base class helper
+        self._initialize_embeddings(
+            self.parameters.embedding_model,
+            self.parameters.embedding_dims,
+        )
+
+        self._setup_complete = True
+        logger.debug(
+            "In-memory context-aware cache initialized",
+            layer=self.name,
+            space_id=self.space_id,
+            dims=self._embedding_dims,
+            capacity=self.parameters.capacity,
+        )
+
+    # Property implementations
+    @property
+    def warehouse(self) -> WarehouseModel:
+        """The warehouse used for executing cached SQL queries."""
+        return self.parameters.warehouse
+
+    @property
+    def time_to_live(self) -> timedelta | None:
+        """Time-to-live for cache entries. None means never expires."""
+        ttl = self.parameters.time_to_live_seconds
+        if ttl is None or ttl < 0:
+            return None
+        return timedelta(seconds=ttl)
+
+    @property
+    def similarity_threshold(self) -> float:
+        """Minimum similarity for cache hit (using L2 distance converted to similarity)."""
+        return self.parameters.similarity_threshold
+
+    @property
+    def context_similarity_threshold(self) -> float:
+        """Minimum similarity for context matching."""
+        return self.parameters.context_similarity_threshold
+
+    @property
+    def question_weight(self) -> float:
+        """Weight for question similarity in combined score."""
+        return self.parameters.question_weight
+
+    @property
+    def context_weight(self) -> float:
+        """Weight for context similarity in combined score."""
+        return self.parameters.context_weight
+
+    def _embed_question(
+        self, question: str, conversation_id: str | None = None
+    ) -> tuple[list[float], list[float], str]:
+        """
+        Generate dual embeddings using Genie API for conversation history.
+
+        Args:
+            question: The question to embed
+            conversation_id: Optional conversation ID for retrieving context
+
+        Returns:
+            Tuple of (question_embedding, context_embedding, conversation_context_string)
+        """
+        return self._embed_question_with_genie_history(
+            question,
+            conversation_id,
+            self.parameters.context_window_size,
+            self.parameters.max_context_tokens,
+        )
+
+    @mlflow.trace(name="semantic_search_in_memory")
+    def _find_similar(
+        self,
+        question: str,
+        conversation_context: str,
+        question_embedding: list[float],
+        context_embedding: list[float],
+        conversation_id: str | None = None,
+    ) -> tuple[SQLCacheEntry, float] | None:
+        """
+        Find a semantically similar cached entry using dual embedding matching.
+
+        Performs linear scan through all cache entries, filtering by space_id and
+        calculating L2 distances for similarity matching.
+
+        Args:
+            question: The original question (for logging)
+            conversation_context: The conversation context string
+            question_embedding: The embedding vector of just the question
+            context_embedding: The embedding vector of the conversation context
+            conversation_id: Optional conversation ID (for logging)
+
+        Returns:
+            Tuple of (SQLCacheEntry, combined_similarity_score) if found, None otherwise
+        """
+        ttl_seconds = self.parameters.time_to_live_seconds
+        ttl_disabled = ttl_seconds is None or ttl_seconds < 0
+
+        question_weight = self.question_weight
+        context_weight = self.context_weight
+
+        best_entry: InMemoryCacheEntry | None = None
+        best_question_sim: float = 0.0
+        best_context_sim: float = 0.0
+        best_combined_sim: float = 0.0
+
+        # Linear scan through all entries
+        with self._lock:
+            entries_to_delete: list[int] = []
+
+            for idx, entry in enumerate(self._cache):
+                # Filter by space_id (partition)
+                if entry.genie_space_id != self.space_id:
+                    continue
+
+                # Check TTL
+                is_valid = True
+                if not ttl_disabled:
+                    age = datetime.now() - entry.created_at
+                    is_valid = age.total_seconds() <= ttl_seconds
+
+                if not is_valid:
+                    entries_to_delete.append(idx)
+                    continue
+
+                # Calculate L2 distances and convert to similarities
+                question_distance = l2_distance(
+                    question_embedding, entry.question_embedding
+                )
+                context_distance = l2_distance(
+                    context_embedding, entry.context_embedding
+                )
+
+                question_sim = distance_to_similarity(question_distance)
+                context_sim = distance_to_similarity(context_distance)
+
+                # Calculate weighted combined similarity
+                combined_sim = (question_weight * question_sim) + (
+                    context_weight * context_sim
+                )
+
+                # Track best match
+                if combined_sim > best_combined_sim:
+                    best_entry = entry
+                    best_question_sim = question_sim
+                    best_context_sim = context_sim
+                    best_combined_sim = combined_sim
+
+            # Delete expired entries
+            for idx in reversed(entries_to_delete):
+                del self._cache[idx]
+                logger.trace("Deleted expired entry", layer=self.name, index=idx)
+
+        # No entries found
+        if best_entry is None:
+            logger.info(
+                "Cache MISS (no entries)",
+                layer=self.name,
+                question=question[:50],
+                space=self.space_id,
+                delegating_to=type(self.impl).__name__,
+            )
+            return None
+
+        # Log best match info
+        logger.debug(
+            "Best match found",
+            layer=self.name,
+            question_sim=f"{best_question_sim:.4f}",
+            context_sim=f"{best_context_sim:.4f}",
+            combined_sim=f"{best_combined_sim:.4f}",
+            cached_question=best_entry.question[:50],
+        )
+
+        # Check BOTH similarity thresholds
+        if best_question_sim < self.similarity_threshold:
+            logger.info(
+                "Cache MISS (question similarity too low)",
+                layer=self.name,
+                question_sim=f"{best_question_sim:.4f}",
+                threshold=self.similarity_threshold,
+            )
+            return None
+
+        if best_context_sim < self.context_similarity_threshold:
+            logger.info(
+                "Cache MISS (context similarity too low)",
+                layer=self.name,
+                context_sim=f"{best_context_sim:.4f}",
+                threshold=self.context_similarity_threshold,
+            )
+            return None
+
+        # Cache HIT - Update last accessed time
+        with self._lock:
+            best_entry.last_accessed_at = datetime.now()
+
+        cache_age_seconds = (datetime.now() - best_entry.created_at).total_seconds()
+        logger.info(
+            "Cache HIT",
+            layer=self.name,
+            question=question[:80],
+            conversation_id=conversation_id,
+            matched_question=best_entry.question[:80],
+            cache_age_seconds=round(cache_age_seconds, 1),
+            question_similarity=f"{best_question_sim:.4f}",
+            context_similarity=f"{best_context_sim:.4f}",
+            combined_similarity=f"{best_combined_sim:.4f}",
+        )
+
+        cache_entry = SQLCacheEntry(
+            query=best_entry.sql_query,
+            description=best_entry.description,
+            conversation_id=best_entry.conversation_id,
+            created_at=best_entry.created_at,
+            message_id=best_entry.message_id,
+            # In-memory caches don't have database row IDs
+            cache_entry_id=None,
+        )
+        return cache_entry, best_combined_sim
+
+    def _store_entry(
+        self,
+        question: str,
+        conversation_context: str,
+        question_embedding: list[float],
+        context_embedding: list[float],
+        response: GenieResponse,
+        message_id: str | None = None,
+    ) -> None:
+        """
+        Store a new cache entry with dual embeddings and message_id.
+
+        If capacity is set and reached, evicts least recently used entry (LRU).
+        """
+        now = datetime.now()
+        new_entry = InMemoryCacheEntry(
+            genie_space_id=self.space_id,
+            question=question,
+            conversation_context=conversation_context,
+            question_embedding=question_embedding,
+            context_embedding=context_embedding,
+            sql_query=response.query,
+            description=response.description,
+            conversation_id=response.conversation_id,
+            created_at=now,
+            last_accessed_at=now,
+            message_id=message_id,
+        )
+
+        with self._lock:
+            # Enforce capacity limit (LRU eviction)
+            if self.parameters.capacity is not None:
+                space_entries = [
+                    e for e in self._cache if e.genie_space_id == self.space_id
+                ]
+
+                while len(space_entries) >= self.parameters.capacity:
+                    # Find and remove least recently used entry
+                    lru_idx = None
+                    lru_time = None
+
+                    for idx, entry in enumerate(self._cache):
+                        if entry.genie_space_id == self.space_id:
+                            if lru_time is None or entry.last_accessed_at < lru_time:
+                                lru_time = entry.last_accessed_at
+                                lru_idx = idx
+
+                    if lru_idx is not None:
+                        evicted = self._cache.pop(lru_idx)
+                        logger.trace(
+                            "Evicted LRU cache entry",
+                            layer=self.name,
+                            question=evicted.question[:50],
+                            capacity=self.parameters.capacity,
+                        )
+                        space_entries = [
+                            e for e in self._cache if e.genie_space_id == self.space_id
+                        ]
+                    else:
+                        break
+
+            self._cache.append(new_entry)
+            logger.debug(
+                "Stored cache entry",
+                layer=self.name,
+                question=question[:50],
+                space=self.space_id,
+                cache_size=len(
+                    [e for e in self._cache if e.genie_space_id == self.space_id]
+                ),
+                capacity=self.parameters.capacity,
+            )
+
+    def _on_stale_cache_entry(self, question: str) -> None:
+        """Remove stale cache entry from memory."""
+        with self._lock:
+            for idx, entry in enumerate(self._cache):
+                if entry.genie_space_id == self.space_id and entry.question == question:
+                    del self._cache[idx]
+                    logger.info(
+                        "Deleted stale cache entry from memory",
+                        layer=self.name,
+                        question=question[:50],
+                    )
+                    break
+
+    def _invalidate_by_question(self, question: str) -> bool:
+        """
+        Invalidate cache entries matching a specific question.
+
+        This method is called when negative feedback is received to remove
+        the corresponding cache entry from the in-memory cache.
+
+        Args:
+            question: The question text to match and invalidate
+
+        Returns:
+            True if an entry was found and invalidated, False otherwise
+        """
+        with self._lock:
+            for idx, entry in enumerate(self._cache):
+                if entry.genie_space_id == self.space_id and entry.question == question:
+                    del self._cache[idx]
+                    logger.info(
+                        "Invalidated cache entry by question",
+                        layer=self.name,
+                        question=question[:50],
+                        space_id=self.space_id,
+                    )
+                    return True
+            return False
+
+    # Note: ask_question_with_cache_info is inherited from ContextAwareGenieService
+    # using the Template Method pattern. InMemoryContextAwareGenieService uses the
+    # default empty hook implementations since it doesn't track prompt history.
+
+    # Template Method implementations for invalidate_expired() and clear()
+
+    def _delete_expired_entries(self, ttl_seconds: int) -> int:
+        """Delete expired entries from the cache."""
+        deleted = 0
+        with self._lock:
+            indices_to_delete: list[int] = []
+            now = datetime.now()
+
+            for idx, entry in enumerate(self._cache):
+                if entry.genie_space_id != self.space_id:
+                    continue
+
+                age = now - entry.created_at
+                if age.total_seconds() > ttl_seconds:
+                    indices_to_delete.append(idx)
+
+            # Delete in reverse order
+            for idx in reversed(indices_to_delete):
+                del self._cache[idx]
+                deleted += 1
+
+            logger.trace(
+                "Deleted expired entries",
+                layer=self.name,
+                deleted_count=deleted,
+            )
+
+        return deleted
+
+    def _delete_all_entries(self) -> int:
+        """Delete all entries for this Genie space."""
+        deleted = 0
+
+        with self._lock:
+            indices_to_delete: list[int] = []
+            for idx, entry in enumerate(self._cache):
+                if entry.genie_space_id == self.space_id:
+                    indices_to_delete.append(idx)
+
+            # Delete in reverse order
+            for idx in reversed(indices_to_delete):
+                del self._cache[idx]
+                deleted += 1
+
+            logger.debug(
+                "Cleared cache entries", layer=self.name, deleted_count=deleted
+            )
+
+        return deleted
+
+    @property
+    def size(self) -> int:
+        """Current number of entries in the cache for this Genie space."""
+        self._setup()
+        with self._lock:
+            return len([e for e in self._cache if e.genie_space_id == self.space_id])
+
+    # Template Method implementations for stats()
+
+    def _count_all_entries(self) -> int:
+        """Count all cache entries for this Genie space."""
+        with self._lock:
+            return len([e for e in self._cache if e.genie_space_id == self.space_id])
+
+    def _count_entries_with_ttl(self, ttl_seconds: int) -> tuple[int, int]:
+        """Count total and expired entries for this Genie space."""
+        now = datetime.now()
+        with self._lock:
+            space_entries = [
+                e for e in self._cache if e.genie_space_id == self.space_id
+            ]
+            total = len(space_entries)
+            expired = 0
+            for entry in space_entries:
+                age = now - entry.created_at
+                if age.total_seconds() > ttl_seconds:
+                    expired += 1
+            return total, expired
+
+    def _get_additional_stats(self) -> dict[str, Any]:
+        """Add capacity info to stats."""
+        return {"capacity": self.parameters.capacity}