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

dao_ai/genie/cache/base.py

@@ -13,10 +13,92 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import TYPE_CHECKING
 
-from databricks_ai_bridge.genie import GenieResponse
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.dashboards import GenieFeedbackRating
+from loguru import logger
 
 if TYPE_CHECKING:
     from dao_ai.genie.cache.base import CacheResult
+    from dao_ai.genie.core import GenieResponse
+
+
+def get_latest_message_id(
+    workspace_client: WorkspaceClient,
+    space_id: str,
+    conversation_id: str,
+) -> str | None:
+    """
+    Look up the most recent message ID for a conversation.
+
+    This is used when sending feedback without a specific message_id.
+
+    Args:
+        workspace_client: The Databricks workspace client
+        space_id: The Genie space ID
+        conversation_id: The conversation ID to look up
+
+    Returns:
+        The message_id of the most recent message, or None if not found
+    """
+    try:
+        response = workspace_client.genie.list_conversation_messages(
+            space_id=space_id,
+            conversation_id=conversation_id,
+        )
+        if response.messages:
+            # Messages are returned in order; get the last one (most recent)
+            messages = list(response.messages)
+            if messages:
+                # Use message_id if available, fall back to id (legacy)
+                latest = messages[-1]
+                return latest.message_id if latest.message_id else latest.id
+        return None
+    except Exception as e:
+        logger.warning(
+            "Failed to look up latest message_id",
+            space_id=space_id,
+            conversation_id=conversation_id,
+            error=str(e),
+        )
+        return None
+
+
+def get_message_content(
+    workspace_client: WorkspaceClient,
+    space_id: str,
+    conversation_id: str,
+    message_id: str,
+) -> str | None:
+    """
+    Get the content (question text) of a specific message.
+
+    This is used to find matching cache entries when invalidating on negative feedback.
+
+    Args:
+        workspace_client: The Databricks workspace client
+        space_id: The Genie space ID
+        conversation_id: The conversation ID
+        message_id: The message ID to look up
+
+    Returns:
+        The message content (question text), or None if not found
+    """
+    try:
+        message = workspace_client.genie.get_message(
+            space_id=space_id,
+            conversation_id=conversation_id,
+            message_id=message_id,
+        )
+        return message.content if message else None
+    except Exception as e:
+        logger.warning(
+            "Failed to get message content",
+            space_id=space_id,
+            conversation_id=conversation_id,
+            message_id=message_id,
+            error=str(e),
+        )
+        return None
 
 
 class GenieServiceBase(ABC):
@@ -40,6 +122,46 @@ class GenieServiceBase(ABC):
         """The space ID for the Genie service."""
         pass
 
+    @property
+    @abstractmethod
+    def workspace_client(self) -> WorkspaceClient | None:
+        """Optional WorkspaceClient for API operations like from_space() and feedback."""
+        pass
+
+    @abstractmethod
+    def send_feedback(
+        self,
+        conversation_id: str,
+        rating: GenieFeedbackRating,
+        message_id: str | None = None,
+        was_cache_hit: bool = False,
+    ) -> None:
+        """
+        Send feedback for a Genie message.
+
+        This method sends user feedback (positive/negative/none) to the Genie service
+        and handles cache invalidation for negative feedback.
+
+        Args:
+            conversation_id: The conversation containing the message
+            rating: The feedback rating (POSITIVE, NEGATIVE, or NONE)
+            message_id: Optional message ID. If None, looks up the most recent message
+                in the conversation.
+            was_cache_hit: Whether the response being rated was served from cache.
+                If True and rating is NEGATIVE, the cache entry is invalidated but
+                no feedback is sent to the Genie API (since no Genie message exists
+                for cached responses).
+
+        Note:
+            For cached responses (was_cache_hit=True), the message_id from the
+            original Genie response is stored in the cache entry. This enables
+            sending feedback to Genie even for cache hits.
+
+            The cache_entry_id in CacheResult can be used to trace which cache
+            entry was used for a particular prompt in genie_prompt_history.
+        """
+        pass
+
 
 @dataclass
 class SQLCacheEntry:
@@ -48,12 +170,22 @@ class SQLCacheEntry:
 
     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.
+
+    Attributes:
+        query: The SQL query to execute
+        description: Description of the query
+        conversation_id: The conversation ID where this query originated
+        created_at: When the entry was created
+        message_id: The original Genie message ID (for feedback on cache hits)
+        cache_entry_id: The database row ID (for persistent caches)
     """
 
     query: str
     description: str
     conversation_id: str
     created_at: datetime
+    message_id: str | None = None
+    cache_entry_id: int | None = None
 
 
 @dataclass
@@ -62,11 +194,20 @@ class CacheResult:
     Result of a cache-aware query with metadata about cache behavior.
 
     Attributes:
-        response: The GenieResponse (fresh data, possibly from cached SQL)
+        response: The GenieResponse (fresh data, possibly from cached SQL).
+            This is the extended dao_ai.genie.GenieResponse which includes message_id,
+            not the base databricks_ai_bridge.genie.GenieResponse.
         cache_hit: Whether the SQL query came from cache
         served_by: Name of the layer that served the cached SQL (None if from origin)
+        message_id: The Genie message ID (for sending feedback). Available from
+            the extended Genie class on cache miss, or from stored cache entry on hit.
+        cache_entry_id: The database row ID of the cache entry that served this hit.
+            Only populated on cache hits from persistent caches (PostgreSQL).
+            Can be used to trace back to genie_prompt_history entries.
     """
 
     response: GenieResponse
     cache_hit: bool
     served_by: str | None = None
+    message_id: str | None = None
+    cache_entry_id: int | None = None

dao_ai/genie/cache/context_aware/__init__.py

@@ -0,0 +1,31 @@
+"""
+Context-aware cache implementations for Genie SQL queries.
+
+This package provides context-aware caching layers that use semantic similarity
+matching with dual embeddings (question + conversation context) for high-precision
+cache lookups.
+
+Available implementations:
+- InMemoryContextAwareGenieService: In-memory cache with L2 distance matching
+- PostgresContextAwareGenieService: PostgreSQL pg_vector-based persistent cache
+
+Base classes:
+- ContextAwareGenieService: Abstract base for all context-aware cache implementations
+- PersistentContextAwareGenieCacheService: Abstract base for database-backed implementations
+"""
+
+from dao_ai.genie.cache.context_aware.base import ContextAwareGenieService
+from dao_ai.genie.cache.context_aware.in_memory import InMemoryContextAwareGenieService
+from dao_ai.genie.cache.context_aware.persistent import (
+    PersistentContextAwareGenieCacheService,
+)
+from dao_ai.genie.cache.context_aware.postgres import PostgresContextAwareGenieService
+
+__all__ = [
+    # Base classes
+    "ContextAwareGenieService",
+    "PersistentContextAwareGenieCacheService",
+    # Implementations
+    "InMemoryContextAwareGenieService",
+    "PostgresContextAwareGenieService",
+]