dao-ai 0.1.19__py3-none-any.whl → 0.1.21__py3-none-any.whl

dao_ai/genie/cache/lru.py

@@ -6,15 +6,16 @@ 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 __future__ import annotations
+
 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.sdk.service.dashboards import GenieFeedbackRating
 from databricks_ai_bridge.genie import GenieResponse
 from loguru import logger
 
@@ -24,6 +25,7 @@ from dao_ai.genie.cache.base import (
     GenieServiceBase,
     SQLCacheEntry,
 )
+from dao_ai.genie.cache.core import execute_sql_via_warehouse
 
 
 class LRUCacheService(GenieServiceBase):
@@ -141,8 +143,20 @@ class LRUCacheService(GenieServiceBase):
         self._cache.move_to_end(key)
         return entry
 
-    def _put(self, key: str, response: GenieResponse) -> None:
+    def _put(
+        self, key: str, response: GenieResponse, message_id: str | None = None
+    ) -> None:
         """Store SQL query in cache, evicting if at capacity."""
+        # Skip caching if query is empty or whitespace
+        if not response.query or not response.query.strip():
+            logger.warning(
+                "Not caching: response has no SQL query",
+                layer=self.name,
+                key=key[:50],
+                description=response.description[:80] if response.description else None,
+            )
+            return
+
         if key in self._cache:
             del self._cache[key]
 
@@ -154,6 +168,9 @@ class LRUCacheService(GenieServiceBase):
             description=response.description,
             conversation_id=response.conversation_id,
             created_at=datetime.now(),
+            message_id=message_id,
+            # LRU cache is in-memory only, no database row ID
+            cache_entry_id=None,
         )
         logger.debug(
             "Stored cache entry",
@@ -162,6 +179,7 @@ class LRUCacheService(GenieServiceBase):
             sql=response.query[:50] if response.query else None,
             cache_size=len(self._cache),
             capacity=self.capacity,
+            message_id=message_id,
         )
 
     @mlflow.trace(name="execute_cached_sql")
@@ -175,50 +193,22 @@ class LRUCacheService(GenieServiceBase):
         Returns:
             DataFrame with results, or error message string
         """
-        w: WorkspaceClient = self.warehouse.workspace_client
-        warehouse_id: str = str(self.warehouse.warehouse_id)
-
-        logger.trace("Executing cached SQL", layer=self.name, sql=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}"
+        # Validate SQL is not empty
+        if not sql or not sql.strip():
+            error_msg: str = "Cannot execute empty SQL query"
             logger.error(
-                "SQL execution failed",
+                "SQL execution failed: empty query",
                 layer=self.name,
-                status=str(statement_response.status),
+                sql=repr(sql),
             )
             return error_msg
 
-        # Convert to DataFrame
-        if statement_response.result and statement_response.result.data_array:
-            columns: list[str] = []
-            if statement_response.manifest and statement_response.manifest.schema:
-                columns = [
-                    col.name for col in statement_response.manifest.schema.columns
-                ]
-
-            data: list[list[Any]] = statement_response.result.data_array
-            if columns:
-                return pd.DataFrame(data, columns=columns)
-            else:
-                return pd.DataFrame(data)
-
-        return pd.DataFrame()
+        # Use shared utility function for SQL execution
+        return execute_sql_via_warehouse(
+            warehouse=self.warehouse,
+            sql=sql,
+            layer_name=self.name,
+        )
 
     def ask_question(
         self, question: str, conversation_id: str | None = None
@@ -256,33 +246,131 @@ class LRUCacheService(GenieServiceBase):
             cached: SQLCacheEntry | None = self._get(key)
 
         if cached is not None:
-            cache_age_seconds = (datetime.now() - cached.created_at).total_seconds()
-            logger.info(
-                "Cache HIT",
-                layer=self.name,
-                question=question[:80],
-                conversation_id=conversation_id,
-                cached_sql=cached.query[:80] if cached.query else None,
-                cache_age_seconds=round(cache_age_seconds, 1),
-                cache_size=self.size,
-                capacity=self.capacity,
-                ttl_seconds=self.parameters.time_to_live_seconds,
-            )
-
-            # Re-execute the cached SQL to get fresh data
-            result: pd.DataFrame | str = self._execute_sql(cached.query)
-
-            # Use current conversation_id, not the cached one
-            response: GenieResponse = GenieResponse(
-                result=result,
-                query=cached.query,
-                description=cached.description,
-                conversation_id=conversation_id
-                if conversation_id
-                else cached.conversation_id,
-            )
-
-            return CacheResult(response=response, cache_hit=True, served_by=self.name)
+            # Defensive check: if cached query is empty, treat as cache miss
+            if not cached.query or not cached.query.strip():
+                logger.warning(
+                    "Cache HIT but query is empty, treating as MISS",
+                    layer=self.name,
+                    question=question[:80],
+                    conversation_id=conversation_id,
+                    key=key[:50],
+                )
+                # Invalidate this bad cache entry
+                with self._lock:
+                    if key in self._cache:
+                        del self._cache[key]
+                # Fall through to cache miss logic below
+            else:
+                cache_age_seconds = (datetime.now() - cached.created_at).total_seconds()
+                logger.info(
+                    "Cache HIT",
+                    layer=self.name,
+                    question=question[:80],
+                    conversation_id=conversation_id,
+                    cached_sql=cached.query[:80] if cached.query else None,
+                    cache_age_seconds=round(cache_age_seconds, 1),
+                    cache_size=self.size,
+                    capacity=self.capacity,
+                    ttl_seconds=self.parameters.time_to_live_seconds,
+                )
+
+                # Re-execute the cached SQL to get fresh data
+                result: pd.DataFrame | str = self._execute_sql(cached.query)
+
+                # Check if SQL execution failed (returns error string instead of DataFrame)
+                if isinstance(result, str):
+                    logger.warning(
+                        "Cached SQL execution failed, falling back to Genie",
+                        layer=self.name,
+                        question=question[:80],
+                        conversation_id=conversation_id,
+                        cached_sql=cached.query[:80],
+                        error=result[:200],
+                        cache_key=key[:50],
+                    )
+
+                    # Invalidate the bad cache entry
+                    with self._lock:
+                        if key in self._cache:
+                            del self._cache[key]
+                            logger.info(
+                                "Invalidated stale cache entry",
+                                layer=self.name,
+                                cache_key=key[:50],
+                                cache_size=len(self._cache),
+                                capacity=self.capacity,
+                            )
+
+                    # Fall back to Genie to get fresh SQL
+                    logger.info(
+                        "Delegating to Genie for fresh SQL",
+                        layer=self.name,
+                        question=question[:80],
+                        delegating_to=type(self.impl).__name__,
+                    )
+                    fallback_result: CacheResult = self.impl.ask_question(
+                        question, conversation_id
+                    )
+
+                    # Store the fresh SQL in cache (including message_id for feedback)
+                    if fallback_result.response.query:
+                        with self._lock:
+                            self._put(
+                                key,
+                                fallback_result.response,
+                                message_id=fallback_result.message_id,
+                            )
+                        logger.info(
+                            "Stored fresh SQL from fallback",
+                            layer=self.name,
+                            fresh_sql=fallback_result.response.query[:80],
+                            cache_size=len(self._cache),
+                            capacity=self.capacity,
+                            message_id=fallback_result.message_id,
+                        )
+                    else:
+                        logger.warning(
+                            "Fallback response has no SQL query to cache",
+                            layer=self.name,
+                            question=question[:80],
+                        )
+
+                    logger.info(
+                        "Fallback completed successfully",
+                        layer=self.name,
+                        question=question[:80],
+                        fallback_from="stale_cache",
+                        has_result=fallback_result.response.result is not None,
+                    )
+
+                    # Return as cache miss (fallback scenario)
+                    # Propagate message_id from fallback result
+                    return CacheResult(
+                        response=fallback_result.response,
+                        cache_hit=False,
+                        served_by=None,
+                        message_id=fallback_result.message_id,
+                    )
+
+                # Use current conversation_id, not the cached one
+                response: GenieResponse = GenieResponse(
+                    result=result,
+                    query=cached.query,
+                    description=cached.description,
+                    conversation_id=conversation_id
+                    if conversation_id
+                    else cached.conversation_id,
+                )
+
+                # Cache hit - include message_id from original response for feedback support
+                return CacheResult(
+                    response=response,
+                    cache_hit=True,
+                    served_by=self.name,
+                    message_id=cached.message_id,
+                    # LRU cache is in-memory only, no cache_entry_id for traceability
+                    cache_entry_id=None,
+                )
 
         # Cache miss - delegate to wrapped service
         logger.info(
@@ -298,7 +386,7 @@ class LRUCacheService(GenieServiceBase):
 
         result: CacheResult = self.impl.ask_question(question, conversation_id)
         with self._lock:
-            self._put(key, result.response)
+            self._put(key, result.response, message_id=result.message_id)
         # Propagate the inner cache's result - if it was a hit there, preserve that info
         return result
 
@@ -306,6 +394,11 @@ class LRUCacheService(GenieServiceBase):
     def space_id(self) -> str:
         return self.impl.space_id
 
+    @property
+    def workspace_client(self) -> WorkspaceClient | None:
+        """Get workspace client by delegating to impl."""
+        return self.impl.workspace_client
+
     def invalidate(self, question: str, conversation_id: str | None = None) -> bool:
         """
         Remove a specific entry from the cache.
@@ -349,3 +442,88 @@ class LRUCacheService(GenieServiceBase):
                 "expired_entries": expired,
                 "valid_entries": len(self._cache) - expired,
             }
+
+    @mlflow.trace(name="genie_lru_cache_send_feedback")
+    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 with cache invalidation.
+
+        For LRU cache, this method:
+        1. If was_cache_hit is False: forwards feedback to the underlying service
+        2. If rating is NEGATIVE: invalidates any matching cache entries
+
+        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.
+            was_cache_hit: Whether the response being rated was served from cache.
+
+        Note:
+            For cached responses (was_cache_hit=True), only cache invalidation is
+            performed. No feedback is sent to the Genie API because cached responses
+            don't have a corresponding Genie message.
+
+            Future Enhancement: To enable full Genie feedback for cached responses,
+            the cache would need to store the original message_id. See GenieServiceBase
+            docstring for details on required changes.
+        """
+        # Handle cache invalidation on negative feedback
+        invalidated = False
+        if rating == GenieFeedbackRating.NEGATIVE:
+            # For LRU cache, we invalidate by conversation_id since that's part of the key
+            # Iterate through cache and remove entries matching the conversation_id
+            with self._lock:
+                keys_to_remove: list[str] = []
+                for key, entry in self._cache.items():
+                    if entry.conversation_id == conversation_id:
+                        keys_to_remove.append(key)
+
+                for key in keys_to_remove:
+                    del self._cache[key]
+                    invalidated = True
+                    logger.info(
+                        "Invalidated cache entry due to negative feedback",
+                        layer=self.name,
+                        cache_key=key[:50],
+                        conversation_id=conversation_id,
+                    )
+
+            if not keys_to_remove:
+                logger.debug(
+                    "No cache entries found to invalidate for negative feedback",
+                    layer=self.name,
+                    conversation_id=conversation_id,
+                )
+
+        # Forward feedback to underlying service if not a cache hit
+        # For cache hits, there's no Genie message to provide feedback on
+        if was_cache_hit:
+            logger.info(
+                "Skipping Genie API feedback - response was served from cache",
+                layer=self.name,
+                conversation_id=conversation_id,
+                rating=rating.value if rating else None,
+                cache_invalidated=invalidated,
+            )
+            return
+
+        # Forward to underlying service
+        logger.debug(
+            "Forwarding feedback to underlying service",
+            layer=self.name,
+            conversation_id=conversation_id,
+            rating=rating.value if rating else None,
+            delegating_to=type(self.impl).__name__,
+        )
+        self.impl.send_feedback(
+            conversation_id=conversation_id,
+            rating=rating,
+            message_id=message_id,
+            was_cache_hit=False,  # Already handled, so pass False
+        )

dao_ai/genie/core.py

@@ -1,35 +1,259 @@
 """
 Core Genie service implementation.
 
-This module provides the concrete implementation of GenieServiceBase
-that wraps the Databricks Genie SDK.
+This module provides:
+- Extended Genie and GenieResponse classes that capture message_id
+- GenieService: Concrete implementation of GenieServiceBase
+
+The extended classes wrap the databricks_ai_bridge versions to add message_id
+support, which is needed for sending feedback to the Genie API.
 """
 
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Union
+
 import mlflow
-from databricks_ai_bridge.genie import Genie, GenieResponse
+import pandas as pd
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.dashboards import GenieFeedbackRating
+from databricks_ai_bridge.genie import Genie as DatabricksGenie
+from databricks_ai_bridge.genie import GenieResponse as DatabricksGenieResponse
+from loguru import logger
 
 from dao_ai.genie.cache import CacheResult, GenieServiceBase
+from dao_ai.genie.cache.base import get_latest_message_id
+
+if TYPE_CHECKING:
+    from typing import Optional
+
+
+# =============================================================================
+# Extended Genie Classes with message_id Support
+# =============================================================================
+
+
+@dataclass
+class GenieResponse(DatabricksGenieResponse):
+    """
+    Extended GenieResponse that includes message_id.
+
+    This extends the databricks_ai_bridge GenieResponse to capture the message_id
+    from API responses, which is required for sending feedback to the Genie API.
+
+    Attributes:
+        result: The query result as string or DataFrame
+        query: The generated SQL query
+        description: Description of the query
+        conversation_id: The conversation ID
+        message_id: The message ID (NEW - enables feedback without extra API call)
+    """
+
+    result: Union[str, pd.DataFrame] = ""
+    query: Optional[str] = ""
+    description: Optional[str] = ""
+    conversation_id: Optional[str] = None
+    message_id: Optional[str] = None
+
+
+class Genie(DatabricksGenie):
+    """
+    Extended Genie that captures message_id in responses.
+
+    This extends the databricks_ai_bridge Genie to return GenieResponse objects
+    that include the message_id from the API response. This enables sending
+    feedback without requiring an additional API call to look up the message ID.
+
+    Usage:
+        genie = Genie(space_id="my-space")
+        response = genie.ask_question("What are total sales?")
+        print(response.message_id)  # Now available!
+
+    The original databricks_ai_bridge classes are available as:
+        - DatabricksGenie
+        - DatabricksGenieResponse
+    """
+
+    def ask_question(
+        self, question: str, conversation_id: str | None = None
+    ) -> GenieResponse:
+        """
+        Ask a question and return response with message_id.
+
+        This overrides the parent method to capture the message_id from the
+        API response and include it in the returned GenieResponse.
+
+        Args:
+            question: The question to ask
+            conversation_id: Optional conversation ID for follow-up questions
+
+        Returns:
+            GenieResponse with message_id populated
+        """
+        with mlflow.start_span(name="ask_question"):
+            # Start or continue conversation
+            if not conversation_id:
+                resp = self.start_conversation(question)
+            else:
+                resp = self.create_message(conversation_id, question)
+
+            # Capture message_id from the API response
+            message_id = resp.get("message_id")
+
+            # Poll for the result using parent's method
+            genie_response = self.poll_for_result(resp["conversation_id"], message_id)
+
+            # Ensure conversation_id is set
+            if not genie_response.conversation_id:
+                genie_response.conversation_id = resp["conversation_id"]
+
+            # Return our extended response with message_id
+            return GenieResponse(
+                result=genie_response.result,
+                query=genie_response.query,
+                description=genie_response.description,
+                conversation_id=genie_response.conversation_id,
+                message_id=message_id,
+            )
+
+
+# =============================================================================
+# GenieService Implementation
+# =============================================================================
 
 
 class GenieService(GenieServiceBase):
-    """Concrete implementation of GenieServiceBase using the Genie SDK."""
+    """
+    Concrete implementation of GenieServiceBase using the extended Genie.
+
+    This service wraps the extended Genie class and provides the GenieServiceBase
+    interface for use with cache layers.
+    """
 
     genie: Genie
+    _workspace_client: WorkspaceClient | None
+
+    def __init__(
+        self,
+        genie: Genie | DatabricksGenie,
+        workspace_client: WorkspaceClient | None = None,
+    ) -> None:
+        """
+        Initialize the GenieService.
+
+        Args:
+            genie: The Genie instance for asking questions. Can be either our
+                extended Genie or the original DatabricksGenie.
+            workspace_client: Optional WorkspaceClient for feedback API.
+                If not provided, one will be created lazily when needed.
+        """
+        self.genie = genie  # type: ignore[assignment]
+        self._workspace_client = workspace_client
+
+    @property
+    def workspace_client(self) -> WorkspaceClient:
+        """
+        Get or create a WorkspaceClient for API calls.
 
-    def __init__(self, genie: Genie) -> None:
-        self.genie = genie
+        Lazily creates a WorkspaceClient using default credentials if not provided.
+        """
+        if self._workspace_client is None:
+            self._workspace_client = WorkspaceClient()
+        return self._workspace_client
 
     @mlflow.trace(name="genie_ask_question")
     def ask_question(
         self, question: str, conversation_id: str | None = None
     ) -> CacheResult:
-        """Ask question to Genie and return CacheResult (no caching at this level)."""
-        response: GenieResponse = self.genie.ask_question(
-            question, conversation_id=conversation_id
-        )
+        """
+        Ask question to Genie and return CacheResult.
+
+        No caching at this level - returns cache miss with fresh response.
+        If using our extended Genie, the message_id will be captured in the response.
+        """
+        response = self.genie.ask_question(question, conversation_id=conversation_id)
+
+        # Extract message_id if available (from our extended GenieResponse)
+        message_id = getattr(response, "message_id", None)
+
         # No caching at this level - return cache miss
-        return CacheResult(response=response, cache_hit=False, served_by=None)
+        return CacheResult(
+            response=response,
+            cache_hit=False,
+            served_by=None,
+            message_id=message_id,
+        )
 
     @property
     def space_id(self) -> str:
         return self.genie.space_id
+
+    @mlflow.trace(name="genie_send_feedback")
+    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.
+
+        For the core GenieService, this always sends feedback to the Genie API
+        (the was_cache_hit parameter is ignored here - it's used by cache wrappers).
+
+        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.
+            was_cache_hit: Ignored by GenieService. Cache wrappers use this to decide
+                whether to forward feedback to the underlying service.
+        """
+        # Look up message_id if not provided
+        if message_id is None:
+            message_id = get_latest_message_id(
+                workspace_client=self.workspace_client,
+                space_id=self.space_id,
+                conversation_id=conversation_id,
+            )
+            if message_id is None:
+                logger.warning(
+                    "Could not find message_id for feedback, skipping",
+                    space_id=self.space_id,
+                    conversation_id=conversation_id,
+                    rating=rating.value if rating else None,
+                )
+                return
+
+        logger.info(
+            "Sending feedback to Genie",
+            space_id=self.space_id,
+            conversation_id=conversation_id,
+            message_id=message_id,
+            rating=rating.value if rating else None,
+        )
+
+        try:
+            self.workspace_client.genie.send_message_feedback(
+                space_id=self.space_id,
+                conversation_id=conversation_id,
+                message_id=message_id,
+                rating=rating,
+            )
+            logger.debug(
+                "Feedback sent successfully",
+                space_id=self.space_id,
+                conversation_id=conversation_id,
+                message_id=message_id,
+            )
+        except Exception as e:
+            logger.error(
+                "Failed to send feedback to Genie",
+                space_id=self.space_id,
+                conversation_id=conversation_id,
+                message_id=message_id,
+                rating=rating.value if rating else None,
+                error=str(e),
+                exc_info=True,
+            )

dao_ai/middleware/__init__.py

@@ -1,5 +1,5 @@
 # DAO AI Middleware Module
-# This module provides middleware implementations compatible with LangChain v1's create_agent
+# Middleware implementations compatible with LangChain v1's create_agent
 
 # Re-export LangChain built-in middleware
 from langchain.agents.middleware import (
@@ -82,6 +82,10 @@ from dao_ai.middleware.summarization import (
     create_summarization_middleware,
 )
 from dao_ai.middleware.tool_call_limit import create_tool_call_limit_middleware
+from dao_ai.middleware.tool_call_observability import (
+    ToolCallObservabilityMiddleware,
+    create_tool_call_observability_middleware,
+)
 from dao_ai.middleware.tool_retry import create_tool_retry_middleware
 from dao_ai.middleware.tool_selector import create_llm_tool_selector_middleware
 
@@ -160,4 +164,7 @@ __all__ = [
     "create_clear_tool_uses_edit",
     # PII middleware factory functions
     "create_pii_middleware",
+    # Tool call observability middleware
+    "ToolCallObservabilityMiddleware",
+    "create_tool_call_observability_middleware",
 ]