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

dao_ai/genie/cache/context_aware/postgres.py

@@ -0,0 +1,1166 @@
+"""
+PostgreSQL pg_vector-based context-aware Genie cache implementation.
+
+This module provides a context-aware cache that uses PostgreSQL with pg_vector
+for semantic similarity search. It supports both standard PostgreSQL and
+Databricks Lakebase connections via the DatabaseModel abstraction.
+
+Features:
+- Dual embedding matching (question + conversation context)
+- pg_vector similarity search with L2 distance
+- Prompt history tracking for conversation context
+- TTL-based expiration with refresh-on-hit
+- Space-partitioned cache entries
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+from typing import Any, Self
+
+import mlflow
+from databricks.sdk import WorkspaceClient
+from databricks_ai_bridge.genie import GenieResponse
+from loguru import logger
+
+from dao_ai.config import (
+    DatabaseModel,
+    GenieContextAwareCacheParametersModel,
+    WarehouseModel,
+)
+from dao_ai.genie.cache.base import (
+    CacheResult,
+    GenieServiceBase,
+    SQLCacheEntry,
+)
+from dao_ai.genie.cache.context_aware.persistent import (
+    DbRow,
+    PersistentContextAwareGenieCacheService,
+)
+
+
+class PostgresContextAwareGenieService(PersistentContextAwareGenieCacheService):
+    """
+    PostgreSQL pg_vector-based context-aware caching decorator.
+
+    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 using pg_vector to find
+    cached queries that match both the question intent AND conversation context.
+
+    Supports both standard PostgreSQL and Databricks Lakebase via DatabaseModel.
+
+    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 GenieContextAwareCacheParametersModel, DatabaseModel
+        from dao_ai.genie.cache.context_aware import PostgresContextAwareGenieService
+
+        cache_params = GenieContextAwareCacheParametersModel(
+            database=database_model,
+            warehouse=warehouse_model,
+            embedding_model="databricks-gte-large-en",
+            time_to_live_seconds=86400,  # 24 hours
+            similarity_threshold=0.85
+        )
+        genie = PostgresContextAwareGenieService(
+            impl=GenieService(Genie(space_id="my-space")),
+            parameters=cache_params
+        )
+
+    Thread-safe: Uses connection pooling from psycopg_pool.
+    """
+
+    impl: GenieServiceBase
+    parameters: GenieContextAwareCacheParametersModel
+    _workspace_client: WorkspaceClient | None
+    name: str
+    _embeddings: Any  # DatabricksEmbeddings
+    _pool: Any  # ConnectionPool
+    _embedding_dims: int | None
+    _setup_complete: bool
+
+    def __init__(
+        self,
+        impl: GenieServiceBase,
+        parameters: GenieContextAwareCacheParametersModel,
+        workspace_client: WorkspaceClient | None = None,
+        name: str | None = None,
+    ) -> None:
+        """
+        Initialize the PostgreSQL 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 database, warehouse, embedding model
+            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._pool = None
+        self._embedding_dims = None
+        self._setup_complete = False
+        self._prompt_stored_for_current_request = False
+
+    def _setup(self) -> None:
+        """Initialize embeddings and database connection pool lazily."""
+        if self._setup_complete:
+            return
+
+        from dao_ai.memory.postgres import PostgresPoolManager
+
+        # Initialize embeddings using base class helper
+        self._initialize_embeddings(
+            self.parameters.embedding_model,
+            self.parameters.embedding_dims,
+        )
+
+        # Get connection pool
+        self._pool = PostgresPoolManager.get_pool(self.parameters.database)
+
+        # Ensure table exists
+        self._create_table_if_not_exists()
+
+        self._setup_complete = True
+        logger.debug(
+            "PostgreSQL context-aware cache initialized",
+            layer=self.name,
+            space_id=self.space_id,
+            table_name=self.table_name,
+            dims=self._embedding_dims,
+        )
+
+    # Property implementations
+    @property
+    def database(self) -> DatabaseModel:
+        """The database used for storing cache entries."""
+        return self.parameters.database
+
+    @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 time_to_live_seconds(self) -> int | None:
+        """TTL in seconds (None or negative = never expires)."""
+        return self.parameters.time_to_live_seconds
+
+    @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
+
+    @property
+    def table_name(self) -> str:
+        """Name of the cache table."""
+        return self.parameters.table_name
+
+    @property
+    def prompt_history_table(self) -> str:
+        """Name of the prompt history table."""
+        return self.parameters.prompt_history_table
+
+    @property
+    def context_window_size(self) -> int:
+        """Number of previous prompts to include in context."""
+        return self.parameters.context_window_size
+
+    @property
+    def max_context_tokens(self) -> int:
+        """Maximum tokens for context string."""
+        return self.parameters.max_context_tokens
+
+    @property
+    def max_prompt_history_length(self) -> int:
+        """Maximum number of prompts to keep per conversation."""
+        return self.parameters.max_prompt_history_length
+
+    def _create_table_if_not_exists(self) -> None:
+        """Create the cache table and prompt history table with pg_vector extension."""
+        create_extension_sql: str = "CREATE EXTENSION IF NOT EXISTS vector"
+
+        # Check if table exists and get current embedding dimensions
+        check_dims_sql: str = """
+            SELECT atttypmod 
+            FROM pg_attribute 
+            WHERE attrelid = %s::regclass 
+              AND attname = 'question_embedding'
+        """
+
+        create_table_sql: str = f"""
+            CREATE TABLE IF NOT EXISTS {self.table_name} (
+                id SERIAL PRIMARY KEY,
+                genie_space_id TEXT NOT NULL,
+                question TEXT NOT NULL,
+                conversation_context TEXT,
+                context_string TEXT,
+                question_embedding vector({self.embedding_dims}),
+                context_embedding vector({self.embedding_dims}),
+                sql_query TEXT NOT NULL,
+                description TEXT,
+                conversation_id TEXT,
+                message_id TEXT,
+                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
+            )
+        """
+
+        # Migration: Add message_id column if it doesn't exist
+        add_message_id_sql: str = f"""
+            ALTER TABLE {self.table_name} 
+            ADD COLUMN IF NOT EXISTS message_id TEXT
+        """
+        # Index for efficient similarity search partitioned by genie_space_id
+        create_question_embedding_index_sql: str = f"""
+            CREATE INDEX IF NOT EXISTS {self.table_name}_question_embedding_idx 
+            ON {self.table_name} 
+            USING ivfflat (question_embedding vector_l2_ops)
+            WITH (lists = 100)
+        """
+        create_context_embedding_index_sql: str = f"""
+            CREATE INDEX IF NOT EXISTS {self.table_name}_context_embedding_idx 
+            ON {self.table_name} 
+            USING ivfflat (context_embedding vector_l2_ops)
+            WITH (lists = 100)
+        """
+        create_space_index_sql: str = f"""
+            CREATE INDEX IF NOT EXISTS {self.table_name}_space_idx 
+            ON {self.table_name} (genie_space_id)
+        """
+        create_unique_question_index_sql: str = f"""
+            CREATE UNIQUE INDEX IF NOT EXISTS {self.table_name}_unique_question_idx
+            ON {self.table_name} (genie_space_id, question)
+        """
+
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(create_extension_sql)
+
+                # Check if table exists and verify embedding dimensions
+                try:
+                    cur.execute(check_dims_sql, (self.table_name,))
+                    row: DbRow | None = cur.fetchone()
+                    if row is not None:
+                        current_dims = row.get("atttypmod", 0)
+                        if current_dims != self.embedding_dims:
+                            logger.warning(
+                                "Embedding dimension mismatch, dropping and recreating table",
+                                layer=self.name,
+                                table_dims=current_dims,
+                                expected_dims=self.embedding_dims,
+                                table_name=self.table_name,
+                            )
+                            cur.execute(f"DROP TABLE {self.table_name}")
+                except Exception:
+                    pass
+
+                try:
+                    cur.execute(create_table_sql)
+                except Exception as e:
+                    logger.debug(
+                        f"Table creation skipped (may already exist): {e}",
+                        layer=self.name,
+                    )
+
+                # Migration: Add message_id column if it doesn't exist (for existing tables)
+                try:
+                    cur.execute(add_message_id_sql)
+                    logger.debug(
+                        "Added message_id column (or already exists)",
+                        layer=self.name,
+                        table_name=self.table_name,
+                    )
+                except Exception as e:
+                    # Column might already exist or other error
+                    logger.debug(
+                        f"message_id column migration skipped: {e}",
+                        layer=self.name,
+                    )
+
+                # Create indexes
+                for idx_name, idx_sql in [
+                    (f"{self.table_name}_space_idx", create_space_index_sql),
+                    (
+                        f"{self.table_name}_question_embedding_idx",
+                        create_question_embedding_index_sql,
+                    ),
+                    (
+                        f"{self.table_name}_context_embedding_idx",
+                        create_context_embedding_index_sql,
+                    ),
+                    (
+                        f"{self.table_name}_unique_question_idx",
+                        create_unique_question_index_sql,
+                    ),
+                ]:
+                    if self._index_exists(cur, idx_name):
+                        logger.debug(
+                            f"Index {idx_name} already exists", layer=self.name
+                        )
+                        continue
+                    try:
+                        cur.execute(idx_sql)
+                    except Exception as e:
+                        logger.warning(
+                            f"Could not create {idx_name}: {e}", layer=self.name
+                        )
+
+                # Create prompt history table
+                try:
+                    self._create_prompt_history_table(cur)
+                except Exception as e:
+                    logger.error(
+                        f"Failed to create prompt history table: {e}",
+                        layer=self.name,
+                        exc_info=True,
+                    )
+
+    def _create_prompt_history_table(self, cur: Any) -> None:
+        """Create the prompt history table for tracking user prompts."""
+        prompt_table_name = self.prompt_history_table
+
+        create_prompt_table_sql: str = f"""
+            CREATE TABLE IF NOT EXISTS {prompt_table_name} (
+                id SERIAL PRIMARY KEY,
+                genie_space_id TEXT NOT NULL,
+                conversation_id TEXT NOT NULL,
+                prompt TEXT NOT NULL,
+                cache_hit BOOLEAN DEFAULT FALSE,
+                cache_entry_id INTEGER,
+                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
+            )
+        """
+
+        # Migration: Add cache_entry_id column if it doesn't exist
+        add_cache_entry_id_sql: str = f"""
+            ALTER TABLE {prompt_table_name} 
+            ADD COLUMN IF NOT EXISTS cache_entry_id INTEGER
+        """
+
+        create_conversation_index_sql: str = f"""
+            CREATE INDEX IF NOT EXISTS {prompt_table_name}_conversation_idx 
+            ON {prompt_table_name} (genie_space_id, conversation_id, created_at DESC)
+        """
+        create_space_index_sql: str = f"""
+            CREATE INDEX IF NOT EXISTS {prompt_table_name}_space_idx 
+            ON {prompt_table_name} (genie_space_id, created_at DESC)
+        """
+        create_unique_prompt_index_sql: str = f"""
+            CREATE UNIQUE INDEX IF NOT EXISTS {prompt_table_name}_unique_prompt_idx
+            ON {prompt_table_name} (genie_space_id, conversation_id, prompt)
+        """
+        create_cache_entry_index_sql: str = f"""
+            CREATE INDEX IF NOT EXISTS {prompt_table_name}_cache_entry_idx
+            ON {prompt_table_name} (cache_entry_id)
+            WHERE cache_entry_id IS NOT NULL
+        """
+
+        try:
+            cur.execute(create_prompt_table_sql)
+        except Exception as e:
+            if "duplicate key" in str(e) or "already exists" in str(e):
+                logger.debug("Prompt history table already exists", layer=self.name)
+            else:
+                raise
+
+        # Migration: Add cache_entry_id column if it doesn't exist (for existing tables)
+        try:
+            cur.execute(add_cache_entry_id_sql)
+            logger.debug(
+                "Added cache_entry_id column (or already exists)",
+                layer=self.name,
+                table_name=prompt_table_name,
+            )
+        except Exception as e:
+            logger.debug(
+                f"cache_entry_id column migration skipped: {e}",
+                layer=self.name,
+            )
+
+        for idx_name, idx_sql in [
+            (f"{prompt_table_name}_conversation_idx", create_conversation_index_sql),
+            (f"{prompt_table_name}_space_idx", create_space_index_sql),
+            (f"{prompt_table_name}_unique_prompt_idx", create_unique_prompt_index_sql),
+            (f"{prompt_table_name}_cache_entry_idx", create_cache_entry_index_sql),
+        ]:
+            if self._index_exists(cur, idx_name):
+                continue
+            try:
+                cur.execute(idx_sql)
+            except Exception as e:
+                if "duplicate key" not in str(e) and "already exists" not in str(e):
+                    logger.warning(
+                        f"Could not create index {idx_name}: {e}", layer=self.name
+                    )
+
+        logger.info(
+            "Prompt history table ready", layer=self.name, table=prompt_table_name
+        )
+
+    @mlflow.trace(name="semantic_search_postgres")
+    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 pg_vector."""
+        ttl_seconds = self.time_to_live_seconds
+        ttl_disabled = ttl_seconds is None or ttl_seconds < 0
+
+        if ttl_disabled:
+            is_valid_expr = "TRUE"
+        else:
+            is_valid_expr = f"created_at > NOW() - INTERVAL '{ttl_seconds} seconds'"
+
+        question_weight = self.question_weight
+        context_weight = self.context_weight
+
+        search_sql: str = f"""
+            SELECT 
+                id,
+                question,
+                conversation_context,
+                sql_query,
+                description,
+                conversation_id,
+                message_id,
+                created_at,
+                1.0 / (1.0 + (question_embedding <-> %s::vector)) as question_similarity,
+                1.0 / (1.0 + (context_embedding <-> %s::vector)) as context_similarity,
+                ({question_weight} * (1.0 / (1.0 + (question_embedding <-> %s::vector)))) +
+                ({context_weight} * (1.0 / (1.0 + (context_embedding <-> %s::vector)))) as combined_similarity,
+                {is_valid_expr} as is_valid
+            FROM {self.table_name}
+            WHERE genie_space_id = %s
+            ORDER BY combined_similarity DESC
+            LIMIT 1
+        """
+
+        question_emb_str = f"[{','.join(str(x) for x in question_embedding)}]"
+        context_emb_str = f"[{','.join(str(x) for x in context_embedding)}]"
+
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    search_sql,
+                    (
+                        question_emb_str,
+                        context_emb_str,
+                        question_emb_str,
+                        context_emb_str,
+                        self.space_id,
+                    ),
+                )
+                row: DbRow | None = cur.fetchone()
+
+                if row is None:
+                    logger.info(
+                        "Cache MISS (no entries)",
+                        layer=self.name,
+                        question=question[:50],
+                        space=self.space_id,
+                    )
+                    return None
+
+                entry_id = row.get("id")
+                cached_question = row.get("question", "")
+                sql_query = row["sql_query"]
+                description = row.get("description", "")
+                conversation_id_cached = row.get("conversation_id", "")
+                created_at = row["created_at"]
+                question_similarity = row["question_similarity"]
+                context_similarity = row["context_similarity"]
+                combined_similarity = row["combined_similarity"]
+                is_valid = row.get("is_valid", False)
+
+                logger.debug(
+                    "Best match found",
+                    layer=self.name,
+                    question_sim=f"{question_similarity:.4f}",
+                    context_sim=f"{context_similarity:.4f}",
+                    combined_sim=f"{combined_similarity:.4f}",
+                    is_valid=is_valid,
+                )
+
+                if question_similarity < self.similarity_threshold:
+                    logger.info(
+                        "Cache MISS (question similarity too low)",
+                        layer=self.name,
+                        question_sim=f"{question_similarity:.4f}",
+                        threshold=self.similarity_threshold,
+                    )
+                    return None
+
+                if context_similarity < self.context_similarity_threshold:
+                    logger.info(
+                        "Cache MISS (context similarity too low)",
+                        layer=self.name,
+                        context_sim=f"{context_similarity:.4f}",
+                        threshold=self.context_similarity_threshold,
+                    )
+                    return None
+
+                if not is_valid:
+                    cur.execute(
+                        f"DELETE FROM {self.table_name} WHERE id = %s", (entry_id,)
+                    )
+                    logger.info("Cache MISS (expired, deleted)", layer=self.name)
+                    return None
+
+                cache_age_seconds = None
+                if created_at:
+                    cache_age_seconds = (
+                        datetime.now(created_at.tzinfo) - created_at
+                    ).total_seconds()
+
+                logger.info(
+                    "Cache HIT",
+                    layer=self.name,
+                    question=question[:80],
+                    matched_question=cached_question[:80],
+                    cache_age_seconds=round(cache_age_seconds, 1)
+                    if cache_age_seconds
+                    else None,
+                    question_similarity=f"{question_similarity:.4f}",
+                    context_similarity=f"{context_similarity:.4f}",
+                    combined_similarity=f"{combined_similarity:.4f}",
+                )
+
+                message_id_cached = row.get("message_id")
+
+                entry = SQLCacheEntry(
+                    query=sql_query,
+                    description=description,
+                    conversation_id=conversation_id_cached,
+                    created_at=created_at,
+                    message_id=message_id_cached,
+                    cache_entry_id=entry_id,
+                )
+                return entry, combined_similarity
+
+    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."""
+        insert_sql: str = f"""
+            INSERT INTO {self.table_name} 
+            (genie_space_id, question, conversation_context, context_string, 
+             question_embedding, context_embedding, sql_query, description, 
+             conversation_id, message_id)
+            VALUES (%s, %s, %s, %s, %s::vector, %s::vector, %s, %s, %s, %s)
+        """
+        question_emb_str = f"[{','.join(str(x) for x in question_embedding)}]"
+        context_emb_str = f"[{','.join(str(x) for x in context_embedding)}]"
+
+        if conversation_context:
+            full_context_string = f"{conversation_context}\nCurrent: {question}"
+        else:
+            full_context_string = question
+
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    insert_sql,
+                    (
+                        self.space_id,
+                        question,
+                        conversation_context,
+                        full_context_string,
+                        question_emb_str,
+                        context_emb_str,
+                        response.query,
+                        response.description,
+                        response.conversation_id,
+                        message_id,
+                    ),
+                )
+                logger.debug(
+                    "Stored cache entry",
+                    layer=self.name,
+                    question=question[:50],
+                    space=self.space_id,
+                    message_id=message_id,
+                )
+
+    def _on_stale_cache_entry(self, question: str) -> None:
+        """Delete stale cache entry from database."""
+        delete_sql = (
+            f"DELETE FROM {self.table_name} WHERE genie_space_id = %s AND question = %s"
+        )
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(delete_sql, (self.space_id, question))
+                deleted_rows = cur.rowcount
+                logger.info(
+                    "Deleted stale cache entry",
+                    layer=self.name,
+                    deleted_rows=deleted_rows,
+                    space_id=self.space_id,
+                )
+
+    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 PostgreSQL database.
+
+        Args:
+            question: The question text to match and invalidate
+
+        Returns:
+            True if an entry was found and invalidated, False otherwise
+        """
+        delete_sql = (
+            f"DELETE FROM {self.table_name} WHERE genie_space_id = %s AND question = %s"
+        )
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(delete_sql, (self.space_id, question))
+                deleted_rows = cur.rowcount if isinstance(cur.rowcount, int) else 0
+                if deleted_rows > 0:
+                    logger.info(
+                        "Invalidated cache entry by question",
+                        layer=self.name,
+                        question=question[:50],
+                        deleted_rows=deleted_rows,
+                        space_id=self.space_id,
+                    )
+                    return True
+                return False
+
+    # Template Method hook implementations
+
+    def _before_cache_lookup(self, question: str, conversation_id: str | None) -> None:
+        """Store prompt before cache lookup."""
+        if conversation_id:
+            self._store_user_prompt(
+                prompt=question,
+                conversation_id=conversation_id,
+                cache_hit=False,
+            )
+            # Track that we stored the prompt
+            self._prompt_stored_for_current_request = True
+        else:
+            self._prompt_stored_for_current_request = False
+
+    def _after_cache_hit(
+        self,
+        question: str,
+        conversation_id: str | None,
+        result: CacheResult,
+    ) -> None:
+        """Update cache_hit flag and cache_entry_id after a cache hit."""
+        if result.cache_hit and self._prompt_stored_for_current_request:
+            actual_conv_id = result.response.conversation_id or conversation_id
+            if actual_conv_id:
+                self._update_prompt_cache_hit(
+                    conversation_id=actual_conv_id,
+                    prompt=question,
+                    cache_hit=True,
+                    cache_entry_id=result.cache_entry_id,
+                )
+
+    def _after_cache_miss(
+        self,
+        question: str,
+        conversation_id: str | None,
+        result: CacheResult,
+    ) -> None:
+        """Store prompt if not done earlier (when conversation_id comes from response)."""
+        if (
+            not self._prompt_stored_for_current_request
+            and result.response.conversation_id
+        ):
+            self._store_user_prompt(
+                prompt=question,
+                conversation_id=result.response.conversation_id,
+                cache_hit=False,
+            )
+
+    # Template Method implementations for invalidate_expired() and clear()
+
+    def _get_empty_expiration_result(self) -> dict[str, int]:
+        """Return empty dict for PostgresContextAwareGenieService."""
+        return {"cache": 0, "prompt_history": 0}
+
+    def _delete_expired_entries(self, ttl_seconds: int) -> dict[str, int]:
+        """Delete expired entries from cache and prompt history."""
+        prompt_ttl_seconds = self.parameters.prompt_history_ttl_seconds
+        if prompt_ttl_seconds is None:
+            prompt_ttl_seconds = ttl_seconds
+
+        result: dict[str, int] = {"cache": 0, "prompt_history": 0}
+
+        # Delete expired cache entries
+        delete_cache_sql = f"""
+            DELETE FROM {self.table_name}
+            WHERE genie_space_id = %s
+              AND created_at < NOW() - INTERVAL '%s seconds'
+        """
+
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(delete_cache_sql, (self.space_id, ttl_seconds))
+                deleted = cur.rowcount if isinstance(cur.rowcount, int) else 0
+                result["cache"] = deleted
+                logger.debug(
+                    "Deleted expired cache entries",
+                    layer=self.name,
+                    deleted_count=deleted,
+                )
+
+        # Delete expired prompt history
+        if prompt_ttl_seconds is not None and prompt_ttl_seconds >= 0:
+            try:
+                delete_prompt_sql = f"""
+                    DELETE FROM {self.prompt_history_table}
+                    WHERE genie_space_id = %s
+                      AND created_at < NOW() - INTERVAL '%s seconds'
+                """
+                with self._pool.connection() as conn:
+                    with conn.cursor() as cur:
+                        cur.execute(
+                            delete_prompt_sql, (self.space_id, prompt_ttl_seconds)
+                        )
+                        deleted = cur.rowcount if isinstance(cur.rowcount, int) else 0
+                        result["prompt_history"] = deleted
+            except Exception as e:
+                logger.warning(
+                    f"Failed to clean up prompt history: {e}", layer=self.name
+                )
+
+        return result
+
+    def _delete_all_entries(self) -> int:
+        """Delete all cache entries for this Genie space."""
+        delete_sql = f"DELETE FROM {self.table_name} WHERE genie_space_id = %s"
+
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(delete_sql, (self.space_id,))
+                deleted: int = cur.rowcount
+                logger.debug(
+                    "Cleared cache entries", layer=self.name, deleted_count=deleted
+                )
+                return deleted
+
+    # Template Method implementations for stats()
+
+    def _count_all_entries(self) -> int:
+        """Count all cache entries for this Genie space."""
+        count_sql = (
+            f"SELECT COUNT(*) as total FROM {self.table_name} WHERE genie_space_id = %s"
+        )
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(count_sql, (self.space_id,))
+                row = cur.fetchone()
+                return row.get("total", 0) if row else 0
+
+    def _count_entries_with_ttl(self, ttl_seconds: int) -> tuple[int, int]:
+        """Count total and expired entries for this Genie space."""
+        stats_sql = f"""
+            SELECT
+                COUNT(*) as total,
+                COUNT(*) FILTER (WHERE created_at <= NOW() - INTERVAL '%s seconds') as expired
+            FROM {self.table_name}
+            WHERE genie_space_id = %s
+        """
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(stats_sql, (ttl_seconds, self.space_id))
+                row = cur.fetchone()
+                if row:
+                    return row.get("total", 0), row.get("expired", 0)
+                return 0, 0
+
+    def _get_additional_stats(self) -> dict[str, Any]:
+        """Add prompt history stats."""
+        prompt_stats_sql = f"""
+            SELECT
+                COUNT(*) as total_prompts,
+                COUNT(*) FILTER (WHERE cache_hit = true) as cache_hit_prompts,
+                COUNT(*) FILTER (WHERE cache_hit = false) as cache_miss_prompts,
+                COUNT(DISTINCT conversation_id) as total_conversations
+            FROM {self.prompt_history_table}
+            WHERE genie_space_id = %s
+        """
+        with self._pool.connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(prompt_stats_sql, (self.space_id,))
+                row = cur.fetchone()
+                if row:
+                    total_prompts = row.get("total_prompts", 0)
+                    return {
+                        "prompt_history": {
+                            "total_prompts": total_prompts,
+                            "cache_hit_prompts": row.get("cache_hit_prompts", 0),
+                            "cache_miss_prompts": row.get("cache_miss_prompts", 0),
+                            "total_conversations": row.get("total_conversations", 0),
+                            "cache_hit_rate": (
+                                row.get("cache_hit_prompts", 0) / total_prompts
+                                if total_prompts > 0
+                                else 0.0
+                            ),
+                        }
+                    }
+        return {}
+
+    def from_space(
+        self,
+        space_id: str | None = None,
+        *,
+        include_all_messages: bool = True,
+        from_datetime: datetime | None = None,
+        to_datetime: datetime | None = None,
+        max_messages: int | None = None,
+    ) -> Self:
+        """Populate cache from existing Genie space conversations.
+
+        Fetches all conversations from a Genie space and populates:
+        1. Prompt history table - all user messages
+        2. Cache embeddings table - messages with SQL query attachments
+
+        Uses ON CONFLICT DO NOTHING to avoid duplicate entries.
+
+        Args:
+            space_id: Genie space ID to import from (defaults to self.space_id)
+            include_all_messages: If True, fetch all users' conversations
+            from_datetime: Only include messages after this time
+            to_datetime: Only include messages before this time
+            max_messages: Limit to last N messages (most recent first)
+
+        Returns:
+            self for method chaining
+        """
+        if self.workspace_client is None:
+            raise ValueError("workspace_client is required for from_space()")
+
+        self._setup()
+        target_space_id = space_id or self.space_id
+
+        logger.info(
+            "Starting from_space import",
+            layer=self.name,
+            space_id=target_space_id,
+            include_all_messages=include_all_messages,
+        )
+
+        stats = {
+            "conversations_processed": 0,
+            "prompts_imported": 0,
+            "prompts_skipped": 0,
+            "cache_entries_imported": 0,
+            "cache_entries_skipped": 0,
+            "errors": 0,
+        }
+
+        from databricks.sdk.service.dashboards import GenieMessage
+
+        all_messages: list[tuple[str, GenieMessage]] = []
+        page_token: str | None = None
+
+        while True:
+            try:
+                response = self.workspace_client.genie.list_conversations(
+                    space_id=target_space_id,
+                    include_all=include_all_messages,
+                    page_token=page_token,
+                )
+            except Exception as e:
+                logger.error(f"Failed to list conversations: {e}", layer=self.name)
+                stats["errors"] += 1
+                break
+
+            if response.conversations is None:
+                break
+
+            for conversation in response.conversations:
+                if conversation.conversation_id is None:
+                    continue
+
+                stats["conversations_processed"] += 1
+
+                try:
+                    messages_response = (
+                        self.workspace_client.genie.list_conversation_messages(
+                            space_id=target_space_id,
+                            conversation_id=conversation.conversation_id,
+                        )
+                    )
+                except Exception as e:
+                    logger.warning(f"Failed to fetch messages: {e}", layer=self.name)
+                    stats["errors"] += 1
+                    continue
+
+                if messages_response.messages is None:
+                    continue
+
+                for message in messages_response.messages:
+                    all_messages.append((conversation.conversation_id, message))
+
+                if max_messages and len(all_messages) >= max_messages:
+                    break
+
+            if max_messages and len(all_messages) >= max_messages:
+                break
+
+            page_token = response.next_page_token
+            if page_token is None:
+                break
+
+        # Sort and limit
+        all_messages.sort(
+            key=lambda x: x[1].created_timestamp if x[1].created_timestamp else 0,
+            reverse=True,
+        )
+        if max_messages:
+            all_messages = all_messages[:max_messages]
+
+        # Group messages by conversation_id for context building
+        from collections import defaultdict
+
+        messages_by_conversation: dict[str, list[tuple[str, GenieMessage]]] = (
+            defaultdict(list)
+        )
+        for conv_id, msg in all_messages:
+            messages_by_conversation[conv_id].append((conv_id, msg))
+
+        # Sort each conversation's messages by timestamp (oldest first for context building)
+        for conv_id in messages_by_conversation:
+            messages_by_conversation[conv_id].sort(
+                key=lambda x: x[1].created_timestamp if x[1].created_timestamp else 0
+            )
+
+        # Process messages
+        for conversation_id, message in all_messages:
+            if message.content is None:
+                continue
+
+            message_created_at = None
+            if message.created_timestamp:
+                message_created_at = datetime.fromtimestamp(
+                    message.created_timestamp / 1000.0,
+                    tz=from_datetime.tzinfo if from_datetime else None,
+                )
+
+            if message_created_at:
+                if from_datetime and message_created_at < from_datetime:
+                    continue
+                if to_datetime and message_created_at > to_datetime:
+                    continue
+
+            # Store prompt
+            prompt_stored = self._store_prompt_if_not_exists(
+                prompt=message.content,
+                conversation_id=conversation_id,
+                space_id=target_space_id,
+                cache_hit=False,
+                created_at=message_created_at,
+            )
+
+            if prompt_stored:
+                stats["prompts_imported"] += 1
+            else:
+                stats["prompts_skipped"] += 1
+
+            # Check for SQL attachments
+            if message.attachments:
+                for attachment in message.attachments:
+                    if attachment.query and attachment.query.query:
+                        try:
+                            # Build conversation context from prior messages
+                            # Uses same "Previous: {content}" format as normal operations
+                            prior_messages: list[str] = []
+                            conv_messages = messages_by_conversation.get(
+                                conversation_id, []
+                            )
+                            for _, prior_msg in conv_messages:
+                                if (
+                                    prior_msg.created_timestamp
+                                    and message.created_timestamp
+                                ):
+                                    if (
+                                        prior_msg.created_timestamp
+                                        < message.created_timestamp
+                                    ):
+                                        if prior_msg.content:
+                                            content = prior_msg.content
+                                            if len(content) > 500:
+                                                content = content[:500] + "..."
+                                            prior_messages.append(
+                                                f"Previous: {content}"
+                                            )
+
+                            # Limit to context_window_size (most recent N messages)
+                            context_window = self.context_window_size
+                            if len(prior_messages) > context_window:
+                                prior_messages = prior_messages[-context_window:]
+
+                            conversation_context = "\n".join(prior_messages)
+
+                            # Generate embeddings
+                            question_embedding = self._embeddings.embed_query(
+                                message.content
+                            )
+                            if conversation_context:
+                                context_embedding = self._embeddings.embed_query(
+                                    conversation_context
+                                )
+                            else:
+                                # Zero vector when no prior context (first message)
+                                context_embedding = [0.0] * len(question_embedding)
+
+                            cache_stored = self._store_cache_entry_if_not_exists(
+                                question=message.content,
+                                conversation_context=conversation_context,
+                                question_embedding=question_embedding,
+                                context_embedding=context_embedding,
+                                sql_query=attachment.query.query,
+                                description=attachment.query.description or "",
+                                conversation_id=conversation_id,
+                                space_id=target_space_id,
+                            )
+
+                            if cache_stored:
+                                stats["cache_entries_imported"] += 1
+                            else:
+                                stats["cache_entries_skipped"] += 1
+                        except Exception as e:
+                            logger.warning(
+                                f"Failed to generate embeddings: {e}", layer=self.name
+                            )
+                            stats["errors"] += 1
+
+        logger.info("Completed from_space import", layer=self.name, **stats)
+        return self
+
+    def _store_prompt_if_not_exists(
+        self,
+        prompt: str,
+        conversation_id: str,
+        space_id: str | None = None,
+        cache_hit: bool = False,
+        created_at: datetime | None = None,
+    ) -> bool:
+        """Store prompt with ON CONFLICT DO NOTHING."""
+        target_space_id = space_id or self.space_id
+        prompt_table_name = self.prompt_history_table
+
+        if created_at:
+            insert_sql = f"""
+                INSERT INTO {prompt_table_name} 
+                (genie_space_id, conversation_id, prompt, cache_hit, created_at)
+                VALUES (%s, %s, %s, %s, %s)
+                ON CONFLICT (genie_space_id, conversation_id, prompt) DO NOTHING
+            """
+            params = (target_space_id, conversation_id, prompt, cache_hit, created_at)
+        else:
+            insert_sql = f"""
+                INSERT INTO {prompt_table_name} 
+                (genie_space_id, conversation_id, prompt, cache_hit)
+                VALUES (%s, %s, %s, %s)
+                ON CONFLICT (genie_space_id, conversation_id, prompt) DO NOTHING
+            """
+            params = (target_space_id, conversation_id, prompt, cache_hit)
+
+        try:
+            with self._pool.connection() as conn:
+                with conn.cursor() as cur:
+                    cur.execute(insert_sql, params)
+                    return cur.rowcount > 0 if isinstance(cur.rowcount, int) else False
+        except Exception:
+            return False
+
+    def _store_cache_entry_if_not_exists(
+        self,
+        question: str,
+        conversation_context: str,
+        question_embedding: list[float],
+        context_embedding: list[float],
+        sql_query: str,
+        description: str | None = None,
+        conversation_id: str | None = None,
+        space_id: str | None = None,
+    ) -> bool:
+        """Store cache entry with ON CONFLICT DO NOTHING."""
+        target_space_id = space_id or self.space_id
+
+        insert_sql = f"""
+            INSERT INTO {self.table_name} 
+            (genie_space_id, question, conversation_context, context_string, 
+             question_embedding, context_embedding, sql_query, description, conversation_id)
+            VALUES (%s, %s, %s, %s, %s::vector, %s::vector, %s, %s, %s)
+            ON CONFLICT (genie_space_id, question) DO NOTHING
+        """
+        question_emb_str = f"[{','.join(str(x) for x in question_embedding)}]"
+        context_emb_str = f"[{','.join(str(x) for x in context_embedding)}]"
+        full_context = (
+            f"{conversation_context}\nCurrent: {question}"
+            if conversation_context
+            else question
+        )
+
+        try:
+            with self._pool.connection() as conn:
+                with conn.cursor() as cur:
+                    cur.execute(
+                        insert_sql,
+                        (
+                            target_space_id,
+                            question,
+                            conversation_context,
+                            full_context,
+                            question_emb_str,
+                            context_emb_str,
+                            sql_query,
+                            description or "",
+                            conversation_id or "",
+                        ),
+                    )
+                    return cur.rowcount > 0 if isinstance(cur.rowcount, int) else False
+        except Exception:
+            return False