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

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/memory/postgres.py

@@ -3,6 +3,7 @@ import atexit
 import threading
 from typing import Any, Optional
 
+from databricks_ai_bridge.lakebase import AsyncLakebasePool, LakebasePool
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.postgres import ShallowPostgresSaver
 from langgraph.checkpoint.postgres.aio import AsyncShallowPostgresSaver
@@ -86,13 +87,22 @@ async def _create_async_pool(
 
 
 class AsyncPostgresPoolManager:
+    """
+    Asynchronous PostgreSQL connection pool manager that shares pools
+    based on database configuration.
+
+    For Lakebase connections (when instance_name is provided), uses AsyncLakebasePool
+    from databricks_ai_bridge which handles automatic token rotation and host resolution.
+    For standard PostgreSQL connections, uses psycopg_pool.AsyncConnectionPool.
+    """
+
     _pools: dict[str, AsyncConnectionPool] = {}
+    _lakebase_pools: dict[str, AsyncLakebasePool] = {}
     _lock: asyncio.Lock = asyncio.Lock()
 
     @classmethod
     async def get_pool(cls, database: DatabaseModel) -> AsyncConnectionPool:
         connection_key: str = database.name
-        connection_params: dict[str, Any] = database.connection_params
 
         async with cls._lock:
             if connection_key in cls._pools:
@@ -103,19 +113,43 @@ class AsyncPostgresPoolManager:
 
             logger.debug("Creating new async PostgreSQL pool", database=database.name)
 
-            kwargs: dict[str, Any] = {
-                "row_factory": dict_row,
-                "autocommit": True,
-            } | database.connection_kwargs or {}
-
-            # Create connection pool
-            pool: AsyncConnectionPool = await _create_async_pool(
-                connection_params=connection_params,
-                database_name=database.name,
-                max_pool_size=database.max_pool_size,
-                timeout_seconds=database.timeout_seconds,
-                kwargs=kwargs,
-            )
+            if database.is_lakebase:
+                # Use AsyncLakebasePool for Lakebase connections
+                # AsyncLakebasePool handles automatic token rotation and host resolution
+                lakebase_pool = AsyncLakebasePool(
+                    instance_name=database.instance_name,
+                    workspace_client=database.workspace_client,
+                    min_size=1,
+                    max_size=database.max_pool_size,
+                    timeout=float(database.timeout_seconds),
+                )
+                # Open the async pool
+                await lakebase_pool.open()
+                # Store the AsyncLakebasePool for proper cleanup
+                cls._lakebase_pools[connection_key] = lakebase_pool
+                # Get the underlying AsyncConnectionPool
+                pool = lakebase_pool.pool
+                logger.success(
+                    "Async Lakebase connection pool created",
+                    database=database.name,
+                    instance_name=database.instance_name,
+                    pool_size=database.max_pool_size,
+                )
+            else:
+                # Use standard async PostgreSQL pool for non-Lakebase connections
+                connection_params: dict[str, Any] = database.connection_params
+                kwargs: dict[str, Any] = {
+                    "row_factory": dict_row,
+                    "autocommit": True,
+                } | database.connection_kwargs or {}
+
+                pool = await _create_async_pool(
+                    connection_params=connection_params,
+                    database_name=database.name,
+                    max_pool_size=database.max_pool_size,
+                    timeout_seconds=database.timeout_seconds,
+                    kwargs=kwargs,
+                )
 
             cls._pools[connection_key] = pool
             return pool
@@ -125,7 +159,13 @@ class AsyncPostgresPoolManager:
         connection_key: str = database.name
 
         async with cls._lock:
-            if connection_key in cls._pools:
+            # Close AsyncLakebasePool if it exists (handles underlying pool cleanup)
+            if connection_key in cls._lakebase_pools:
+                lakebase_pool = cls._lakebase_pools.pop(connection_key)
+                await lakebase_pool.close()
+                cls._pools.pop(connection_key, None)
+                logger.debug("Async Lakebase pool closed", database=database.name)
+            elif connection_key in cls._pools:
                 pool = cls._pools.pop(connection_key)
                 await pool.close()
                 logger.debug("Async PostgreSQL pool closed", database=database.name)
@@ -133,9 +173,32 @@ class AsyncPostgresPoolManager:
     @classmethod
     async def close_all_pools(cls):
         async with cls._lock:
+            # Close all AsyncLakebasePool instances first
+            for connection_key, lakebase_pool in cls._lakebase_pools.items():
+                try:
+                    await asyncio.wait_for(lakebase_pool.close(), timeout=2.0)
+                    logger.debug("Async Lakebase pool closed", pool=connection_key)
+                except asyncio.TimeoutError:
+                    logger.warning(
+                        "Timeout closing async Lakebase pool, forcing closure",
+                        pool=connection_key,
+                    )
+                except asyncio.CancelledError:
+                    logger.warning(
+                        "Async Lakebase pool closure cancelled (shutdown in progress)",
+                        pool=connection_key,
+                    )
+                except Exception as e:
+                    logger.error(
+                        "Error closing async Lakebase pool",
+                        pool=connection_key,
+                        error=str(e),
+                    )
+            cls._lakebase_pools.clear()
+
+            # Close any remaining standard async PostgreSQL pools
             for connection_key, pool in cls._pools.items():
                 try:
-                    # Use a short timeout to avoid blocking on pool closure
                     await asyncio.wait_for(pool.close(), timeout=2.0)
                     logger.debug("Async PostgreSQL pool closed", pool=connection_key)
                 except asyncio.TimeoutError:
@@ -178,7 +241,20 @@ class AsyncPostgresStoreManager(StoreManagerBase):
     def _setup(self):
         if self._setup_complete:
             return
-        asyncio.run(self._async_setup())
+        try:
+            # Check if we're already in an async context
+            asyncio.get_running_loop()
+            # If we get here, we're in an async context - raise to caller
+            raise RuntimeError(
+                "Cannot call sync _setup() from async context. "
+                "Use await _async_setup() instead."
+            )
+        except RuntimeError as e:
+            if "no running event loop" in str(e).lower():
+                # No event loop running - safe to use asyncio.run()
+                asyncio.run(self._async_setup())
+            else:
+                raise
 
     async def _async_setup(self):
         if self._setup_complete:
@@ -237,13 +313,25 @@ class AsyncPostgresCheckpointerManager(CheckpointManagerBase):
 
     def _setup(self):
         """
-        Run the async setup. Works in both sync and async contexts when nest_asyncio is applied.
+        Run the async setup. For async contexts, use await _async_setup() directly.
         """
         if self._setup_complete:
             return
 
-        # With nest_asyncio applied in notebooks, asyncio.run() works everywhere
-        asyncio.run(self._async_setup())
+        try:
+            # Check if we're already in an async context
+            asyncio.get_running_loop()
+            # If we get here, we're in an async context - raise to caller
+            raise RuntimeError(
+                "Cannot call sync _setup() from async context. "
+                "Use await _async_setup() instead."
+            )
+        except RuntimeError as e:
+            if "no running event loop" in str(e).lower():
+                # No event loop running - safe to use asyncio.run()
+                asyncio.run(self._async_setup())
+            else:
+                raise
 
     async def _async_setup(self):
         """
@@ -284,15 +372,19 @@ class PostgresPoolManager:
     """
     Synchronous PostgreSQL connection pool manager that shares pools
     based on database configuration.
+
+    For Lakebase connections (when instance_name is provided), uses LakebasePool
+    from databricks_ai_bridge which handles automatic token rotation and host resolution.
+    For standard PostgreSQL connections, uses psycopg_pool.ConnectionPool.
     """
 
     _pools: dict[str, ConnectionPool] = {}
+    _lakebase_pools: dict[str, LakebasePool] = {}
     _lock: threading.Lock = threading.Lock()
 
     @classmethod
     def get_pool(cls, database: DatabaseModel) -> ConnectionPool:
         connection_key: str = str(database.name)
-        connection_params: dict[str, Any] = database.connection_params
 
         with cls._lock:
             if connection_key in cls._pools:
@@ -301,19 +393,41 @@ class PostgresPoolManager:
 
             logger.debug("Creating new PostgreSQL pool", database=database.name)
 
-            kwargs: dict[str, Any] = {
-                "row_factory": dict_row,
-                "autocommit": True,
-            } | database.connection_kwargs or {}
-
-            # Create connection pool
-            pool: ConnectionPool = _create_pool(
-                connection_params=connection_params,
-                database_name=database.name,
-                max_pool_size=database.max_pool_size,
-                timeout_seconds=database.timeout_seconds,
-                kwargs=kwargs,
-            )
+            if database.is_lakebase:
+                # Use LakebasePool for Lakebase connections
+                # LakebasePool handles automatic token rotation and host resolution
+                lakebase_pool = LakebasePool(
+                    instance_name=database.instance_name,
+                    workspace_client=database.workspace_client,
+                    min_size=1,
+                    max_size=database.max_pool_size,
+                    timeout=float(database.timeout_seconds),
+                )
+                # Store the LakebasePool for proper cleanup
+                cls._lakebase_pools[connection_key] = lakebase_pool
+                # Get the underlying ConnectionPool
+                pool = lakebase_pool.pool
+                logger.success(
+                    "Lakebase connection pool created",
+                    database=database.name,
+                    instance_name=database.instance_name,
+                    pool_size=database.max_pool_size,
+                )
+            else:
+                # Use standard PostgreSQL pool for non-Lakebase connections
+                connection_params: dict[str, Any] = database.connection_params
+                kwargs: dict[str, Any] = {
+                    "row_factory": dict_row,
+                    "autocommit": True,
+                } | database.connection_kwargs or {}
+
+                pool = _create_pool(
+                    connection_params=connection_params,
+                    database_name=database.name,
+                    max_pool_size=database.max_pool_size,
+                    timeout_seconds=database.timeout_seconds,
+                    kwargs=kwargs,
+                )
 
             cls._pools[connection_key] = pool
             return pool
@@ -323,7 +437,13 @@ class PostgresPoolManager:
         connection_key: str = database.name
 
         with cls._lock:
-            if connection_key in cls._pools:
+            # Close LakebasePool if it exists (handles underlying pool cleanup)
+            if connection_key in cls._lakebase_pools:
+                lakebase_pool = cls._lakebase_pools.pop(connection_key)
+                lakebase_pool.close()
+                cls._pools.pop(connection_key, None)
+                logger.debug("Lakebase pool closed", database=database.name)
+            elif connection_key in cls._pools:
                 pool = cls._pools.pop(connection_key)
                 pool.close()
                 logger.debug("PostgreSQL pool closed", database=database.name)
@@ -331,16 +451,32 @@ class PostgresPoolManager:
     @classmethod
     def close_all_pools(cls):
         with cls._lock:
-            for connection_key, pool in cls._pools.items():
+            # Close all LakebasePool instances first
+            for connection_key, lakebase_pool in cls._lakebase_pools.items():
                 try:
-                    pool.close()
-                    logger.debug("PostgreSQL pool closed", pool=connection_key)
+                    lakebase_pool.close()
+                    logger.debug("Lakebase pool closed", pool=connection_key)
                 except Exception as e:
                     logger.error(
-                        "Error closing PostgreSQL pool",
+                        "Error closing Lakebase pool",
                         pool=connection_key,
                         error=str(e),
                     )
+            cls._lakebase_pools.clear()
+
+            # Close any remaining standard PostgreSQL pools
+            for connection_key, pool in cls._pools.items():
+                # Skip if already closed via LakebasePool
+                if connection_key not in cls._lakebase_pools:
+                    try:
+                        pool.close()
+                        logger.debug("PostgreSQL pool closed", pool=connection_key)
+                    except Exception as e:
+                        logger.error(
+                            "Error closing PostgreSQL pool",
+                            pool=connection_key,
+                            error=str(e),
+                        )
             cls._pools.clear()
 
 

dao_ai/middleware/__init__.py

@@ -3,8 +3,16 @@
 
 # Re-export LangChain built-in middleware
 from langchain.agents.middleware import (
+    ClearToolUsesEdit,
+    ContextEditingMiddleware,
     HumanInTheLoopMiddleware,
+    LLMToolSelectorMiddleware,
+    ModelCallLimitMiddleware,
+    ModelRetryMiddleware,
+    PIIMiddleware,
     SummarizationMiddleware,
+    ToolCallLimitMiddleware,
+    ToolRetryMiddleware,
     after_agent,
     after_model,
     before_agent,
@@ -37,6 +45,10 @@ from dao_ai.middleware.base import (
     ModelRequest,
     ModelResponse,
 )
+from dao_ai.middleware.context_editing import (
+    create_clear_tool_uses_edit,
+    create_context_editing_middleware,
+)
 from dao_ai.middleware.core import create_factory_middleware
 from dao_ai.middleware.guardrails import (
     ContentFilterMiddleware,
@@ -62,10 +74,16 @@ from dao_ai.middleware.message_validation import (
     create_thread_id_validation_middleware,
     create_user_id_validation_middleware,
 )
+from dao_ai.middleware.model_call_limit import create_model_call_limit_middleware
+from dao_ai.middleware.model_retry import create_model_retry_middleware
+from dao_ai.middleware.pii import create_pii_middleware
 from dao_ai.middleware.summarization import (
     LoggingSummarizationMiddleware,
     create_summarization_middleware,
 )
+from dao_ai.middleware.tool_call_limit import create_tool_call_limit_middleware
+from dao_ai.middleware.tool_retry import create_tool_retry_middleware
+from dao_ai.middleware.tool_selector import create_llm_tool_selector_middleware
 
 __all__ = [
     # Base class (from LangChain)
@@ -85,6 +103,14 @@ __all__ = [
     "SummarizationMiddleware",
     "LoggingSummarizationMiddleware",
     "HumanInTheLoopMiddleware",
+    "ToolCallLimitMiddleware",
+    "ModelCallLimitMiddleware",
+    "ToolRetryMiddleware",
+    "ModelRetryMiddleware",
+    "LLMToolSelectorMiddleware",
+    "ContextEditingMiddleware",
+    "ClearToolUsesEdit",
+    "PIIMiddleware",
     # Core factory function
     "create_factory_middleware",
     # DAO AI middleware implementations
@@ -122,4 +148,16 @@ __all__ = [
     "create_assert_middleware",
     "create_suggest_middleware",
     "create_refine_middleware",
+    # Limit and retry middleware factory functions
+    "create_tool_call_limit_middleware",
+    "create_model_call_limit_middleware",
+    "create_tool_retry_middleware",
+    "create_model_retry_middleware",
+    # Tool selection middleware factory functions
+    "create_llm_tool_selector_middleware",
+    # Context editing middleware factory functions
+    "create_context_editing_middleware",
+    "create_clear_tool_uses_edit",
+    # PII middleware factory functions
+    "create_pii_middleware",
 ]

dao_ai/middleware/assertions.py

@@ -688,7 +688,7 @@ def create_assert_middleware(
         name: Name for function constraints
 
     Returns:
-        AssertMiddleware configured with the constraint
+        List containing AssertMiddleware configured with the constraint
 
     Example:
         # Using a Constraint class
@@ -737,7 +737,7 @@ def create_suggest_middleware(
         name: Name for function constraints
 
     Returns:
-        SuggestMiddleware configured with the constraint
+        List containing SuggestMiddleware configured with the constraint
 
     Example:
         def is_professional(response: str, ctx: dict) -> ConstraintResult:
@@ -783,7 +783,7 @@ def create_refine_middleware(
         select_best: Track and return best response across iterations
 
     Returns:
-        RefineMiddleware configured with the reward function
+        List containing RefineMiddleware configured with the reward function
 
     Example:
         def evaluate_completeness(response: str, ctx: dict) -> float: