dao-ai 0.1.17__py3-none-any.whl → 0.1.18__py3-none-any.whl

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:
@@ -309,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:
@@ -326,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
@@ -348,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)
@@ -356,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/orchestration/core.py

@@ -9,7 +9,7 @@ This module provides the foundational utilities for multi-agent orchestration:
 - Main orchestration graph factory
 """
 
-from typing import Awaitable, Callable, Literal
+from typing import Any, Awaitable, Callable, Literal
 
 from langchain.tools import ToolRuntime, tool
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
@@ -179,8 +179,16 @@ def create_agent_node_handler(
             "messages": filtered_messages,
         }
 
-        # Invoke the agent
-        result: AgentState = await agent.ainvoke(agent_state, context=runtime.context)
+        # Build config with configurable from context for langmem compatibility
+        # langmem tools expect user_id to be in config.configurable
+        config: dict[str, Any] = {}
+        if runtime.context:
+            config = {"configurable": runtime.context.model_dump()}
+
+        # Invoke the agent with both context and config
+        result: AgentState = await agent.ainvoke(
+            agent_state, context=runtime.context, config=config
+        )
 
         # Extract agent response based on output mode
         result_messages = result.get("messages", [])
@@ -227,15 +235,31 @@ def create_handoff_tool(
         tool_call_id: str = runtime.tool_call_id
         logger.debug("Handoff to agent", target_agent=target_agent_name)
 
+        # Get the AIMessage that triggered this handoff (required for tool_use/tool_result pairing)
+        # LLMs expect tool calls to be paired with their responses, so we must include both
+        # the AIMessage containing the tool call and the ToolMessage acknowledging it.
+        messages: list[BaseMessage] = runtime.state.get("messages", [])
+        last_ai_message: AIMessage | None = None
+        for msg in reversed(messages):
+            if isinstance(msg, AIMessage) and msg.tool_calls:
+                last_ai_message = msg
+                break
+
+        # Build message list with proper pairing
+        update_messages: list[BaseMessage] = []
+        if last_ai_message:
+            update_messages.append(last_ai_message)
+        update_messages.append(
+            ToolMessage(
+                content=f"Transferred to {target_agent_name}",
+                tool_call_id=tool_call_id,
+            )
+        )
+
         return Command(
             update={
                 "active_agent": target_agent_name,
-                "messages": [
-                    ToolMessage(
-                        content=f"Transferred to {target_agent_name}",
-                        tool_call_id=tool_call_id,
-                    )
-                ],
+                "messages": update_messages,
             },
             goto=target_agent_name,
             graph=Command.PARENT,

dao_ai/orchestration/supervisor.py

@@ -13,7 +13,7 @@ from langchain.agents import create_agent
 from langchain.agents.middleware import AgentMiddleware as LangchainAgentMiddleware
 from langchain.tools import ToolRuntime, tool
 from langchain_core.language_models import LanguageModelLike
-from langchain_core.messages import ToolMessage
+from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
 from langchain_core.tools import BaseTool
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph import StateGraph
@@ -75,15 +75,30 @@ def _create_handoff_back_to_supervisor_tool() -> BaseTool:
         tool_call_id: str = runtime.tool_call_id
         logger.debug("Agent handing back to supervisor", summary_preview=summary[:100])
 
+        # Get the AIMessage that triggered this handoff (required for tool_use/tool_result pairing)
+        # LLMs expect tool calls to be paired with their responses, so we must include both
+        # the AIMessage containing the tool call and the ToolMessage acknowledging it.
+        messages: list[BaseMessage] = runtime.state.get("messages", [])
+        last_ai_message: AIMessage | None = None
+        for msg in reversed(messages):
+            if isinstance(msg, AIMessage) and msg.tool_calls:
+                last_ai_message = msg
+                break
+
+        # Build message list with proper pairing
+        update_messages: list[BaseMessage] = []
+        if last_ai_message:
+            update_messages.append(last_ai_message)
+        update_messages.append(
+            ToolMessage(
+                content=f"Task completed: {summary}",
+                tool_call_id=tool_call_id,
+            )
+        )
+
         return Command(
             update={
-                "active_agent": None,
-                "messages": [
-                    ToolMessage(
-                        content=f"Task completed: {summary}",
-                        tool_call_id=tool_call_id,
-                    )
-                ],
+                "messages": update_messages,
             },
             goto=SUPERVISOR_NODE,
             graph=Command.PARENT,

dao_ai/{prompts.py → prompts/__init__.py}

@@ -2,9 +2,11 @@
 Prompt utilities for DAO AI agents.
 
 This module provides utilities for creating dynamic prompts using
-LangChain v1's @dynamic_prompt middleware decorator pattern.
+LangChain v1's @dynamic_prompt middleware decorator pattern, as well as
+paths to prompt template files.
 """
 
+from pathlib import Path
 from typing import Any, Optional
 
 from langchain.agents.middleware import (
@@ -18,6 +20,13 @@ from loguru import logger
 from dao_ai.config import PromptModel
 from dao_ai.state import Context
 
+PROMPTS_DIR = Path(__file__).parent
+
+
+def get_prompt_path(name: str) -> Path:
+    """Get the path to a prompt template file."""
+    return PROMPTS_DIR / name
+
 
 def make_prompt(
     base_system_prompt: Optional[str | PromptModel],

dao_ai/prompts/instructed_retriever_decomposition.yaml

@@ -0,0 +1,58 @@
+name: instructed_retriever_decomposition
+description: Decomposes user queries into multiple search queries with metadata filters
+
+template: |
+  You are a search query decomposition expert. Your task is to break down a user query into one or more focused search queries with appropriate metadata filters. Respond with a JSON object.
+
+  ## Current Time
+  {current_time}
+
+  ## Database Schema
+  {schema_description}
+
+  ## Constraints
+  {constraints}
+
+  ## Few-Shot Examples
+  {examples}
+
+  ## Instructions
+  1. Analyze the user query and identify distinct search intents
+  2. For each intent, create a focused search query text
+  3. Extract metadata filters from the query using the exact filter syntax above
+  4. Resolve relative time references (e.g., "last month", "past year") using the current time
+  5. Generate at most {max_subqueries} search queries
+  6. If no filters apply, set filters to null
+
+  ## User Query
+  {query}
+
+  Generate search queries that together capture all aspects of the user's information need.
+
+variables:
+  - current_time
+  - schema_description
+  - constraints
+  - examples
+  - max_subqueries
+  - query
+
+output_format: |
+  The output must be a JSON object with a "queries" field containing an array of search query objects.
+  Each search query object has:
+  - "text": The search query string
+  - "filters": An array of filter objects, each with "key" (column + optional operator) and "value", or null if no filters
+  
+  Supported filter operators (append to column name):
+  - Equality: {"key": "column", "value": "val"} or {"key": "column", "value": ["val1", "val2"]}
+  - Exclusion: {"key": "column NOT", "value": "val"}
+  - Comparison: {"key": "column <", "value": 100}, also <=, >, >=
+  - Token match: {"key": "column LIKE", "value": "word"}
+  - Exclude token: {"key": "column NOT LIKE", "value": "word"}
+  
+  Examples:
+  - [{"key": "brand_name", "value": "MILWAUKEE"}]
+  - [{"key": "price <", "value": 100}]
+  - [{"key": "brand_name NOT", "value": "DEWALT"}]
+  - [{"key": "brand_name", "value": ["MILWAUKEE", "DEWALT"]}]
+  - [{"key": "description LIKE", "value": "cordless"}]

dao_ai/prompts/instruction_reranker.yaml

@@ -0,0 +1,14 @@
+name: instruction_aware_reranking
+version: "1.1"
+description: Rerank documents based on user instructions and constraints
+
+template: |
+  Rerank these search results for the query "{query}".
+
+  {instructions}
+
+  ## Documents
+
+  {documents}
+
+  Score each document 0.0-1.0 based on relevance to the query and instructions. Return results sorted by score (highest first). Only include documents scoring > 0.1.

dao_ai/prompts/router.yaml

@@ -0,0 +1,37 @@
+name: router_query_classification
+version: "1.0"
+description: Classify query to determine execution mode (standard vs instructed)
+
+template: |
+  You are a query classification system. Your task is to determine the best execution mode for a search query.
+
+  ## Execution Modes
+
+  **standard**: Use for simple keyword or product searches without specific constraints.
+  - General questions about products
+  - Simple keyword searches
+  - Broad category browsing
+
+  **instructed**: Use for queries with explicit constraints that require metadata filtering.
+  - Price constraints ("under $100", "between $50 and $200")
+  - Brand preferences ("Milwaukee", "not DeWalt", "excluding Makita")
+  - Category filters ("power tools", "paint supplies")
+  - Time/recency constraints ("recent", "from last month", "updated this year")
+  - Comparison queries ("compare X and Y")
+  - Multiple combined constraints
+
+  ## Available Schema for Filtering
+
+  {schema_description}
+
+  ## Query to Classify
+
+  "{query}"
+
+  ## Instructions
+
+  Analyze the query and determine:
+  1. Does it contain explicit constraints that can be translated to metadata filters?
+  2. Would the query benefit from being decomposed into subqueries?
+
+  Return your classification as a JSON object with a single field "mode" set to either "standard" or "instructed".

dao_ai/prompts/verifier.yaml

@@ -0,0 +1,46 @@
+name: result_verification
+version: "1.0"
+description: Verify search results satisfy user constraints
+
+template: |
+  You are a result verification system. Your task is to determine whether search results satisfy the user's query constraints.
+
+  ## User Query
+
+  "{query}"
+
+  ## Schema Information
+
+  {schema_description}
+
+  ## Constraints to Verify
+
+  {constraints}
+
+  ## Retrieved Results (Top {num_results})
+
+  {results_summary}
+
+  ## Previous Attempt Feedback (if retry)
+
+  {previous_feedback}
+
+  ## Instructions
+
+  Analyze whether the results satisfy the user's explicit and implicit constraints:
+
+  1. **Intent Match**: Do the results address what the user is looking for?
+  2. **Explicit Constraints**: Are price, brand, category, date constraints met?
+  3. **Relevance**: Are the results actually useful for the user's needs?
+
+  If results do NOT satisfy constraints, suggest specific filter relaxations:
+  - Use "REMOVE" to drop a filter entirely
+  - Use "BROADEN" to widen a range (e.g., price < 100 -> price < 150)
+  - Use specific values to change a filter
+
+  Return a JSON object with:
+  - passed: boolean (true if results are satisfactory)
+  - confidence: float (0.0-1.0, your confidence in the assessment)
+  - feedback: string (brief explanation of issues, if any)
+  - suggested_filter_relaxation: object (filter changes for retry, e.g., {{"brand_name": "REMOVE"}})
+  - unmet_constraints: array of strings (list of constraints not satisfied)

dao_ai/providers/databricks.py

@@ -397,6 +397,8 @@ class DatabricksProvider(ServiceProvider):
 
             pip_requirements += get_installed_packages()
 
+        code_paths = list(dict.fromkeys(code_paths))
+
         logger.trace("Pip requirements prepared", count=len(pip_requirements))
         logger.trace("Code paths prepared", count=len(code_paths))
 
@@ -434,19 +436,38 @@ class DatabricksProvider(ServiceProvider):
             pip_packages_count=len(pip_requirements),
         )
 
-        with mlflow.start_run(run_name=run_name):
-            mlflow.set_tag("type", "agent")
-            mlflow.set_tag("dao_ai", dao_ai_version())
-            logged_agent_info: ModelInfo = mlflow.pyfunc.log_model(
-                python_model=model_path.as_posix(),
-                code_paths=code_paths,
-                model_config=config.model_dump(mode="json", by_alias=True),
-                name="agent",
-                conda_env=conda_env,
-                input_example=input_example,
-                # resources=all_resources,
-                auth_policy=auth_policy,
+        # End any stale runs before starting to ensure clean state on retry
+        if mlflow.active_run():
+            logger.warning(
+                "Ending stale MLflow run before creating new agent",
+                run_id=mlflow.active_run().info.run_id,
+            )
+            mlflow.end_run()
+
+        try:
+            with mlflow.start_run(run_name=run_name):
+                mlflow.set_tag("type", "agent")
+                mlflow.set_tag("dao_ai", dao_ai_version())
+                logged_agent_info: ModelInfo = mlflow.pyfunc.log_model(
+                    python_model=model_path.as_posix(),
+                    code_paths=code_paths,
+                    model_config=config.model_dump(mode="json", by_alias=True),
+                    name="agent",
+                    conda_env=conda_env,
+                    input_example=input_example,
+                    # resources=all_resources,
+                    auth_policy=auth_policy,
+                )
+        except Exception as e:
+            # Ensure run is ended on failure to prevent stale state on retry
+            if mlflow.active_run():
+                mlflow.end_run(status="FAILED")
+            logger.error(
+                "Failed to log model",
+                run_name=run_name,
+                error=str(e),
             )
+            raise
 
         registered_model_name: str = config.app.registered_model.full_name