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

dao_ai/nodes.py

@@ -26,11 +26,9 @@ from dao_ai.config import (
 from dao_ai.middleware.core import create_factory_middleware
 from dao_ai.middleware.guardrails import GuardrailMiddleware
 from dao_ai.middleware.human_in_the_loop import (
-    HumanInTheLoopMiddleware,
     create_hitl_middleware_from_tool_models,
 )
 from dao_ai.middleware.summarization import (
-    LoggingSummarizationMiddleware,
     create_summarization_middleware,
 )
 from dao_ai.prompts import make_prompt
@@ -78,8 +76,7 @@ def _create_middleware_list(
             function_name=middleware_config.name,
             args=middleware_config.args,
         )
-        if middleware is not None:
-            middleware_list.append(middleware)
+        middleware_list.append(middleware)
 
     # Add guardrails as middleware
     if agent.guardrails:
@@ -117,16 +114,12 @@ def _create_middleware_list(
             max_tokens=chat_history.max_tokens,
             summary_model=chat_history.model.name,
         )
-        summarization_middleware: LoggingSummarizationMiddleware = (
-            create_summarization_middleware(chat_history)
-        )
+        summarization_middleware = create_summarization_middleware(chat_history)
         middleware_list.append(summarization_middleware)
 
     # Add human-in-the-loop middleware if any tools require it
-    hitl_middleware: HumanInTheLoopMiddleware | None = (
-        create_hitl_middleware_from_tool_models(tool_models)
-    )
-    if hitl_middleware is not None:
+    hitl_middlewares = create_hitl_middleware_from_tool_models(tool_models)
+    if hitl_middlewares:
         # Log which tools require HITL
         hitl_tool_names: list[str] = [
             tool.name
@@ -139,7 +132,7 @@ def _create_middleware_list(
             agent=agent.name,
             hitl_tools=hitl_tool_names,
         )
-        middleware_list.append(hitl_middleware)
+        middleware_list.append(hitl_middlewares)
 
     logger.info(
         "Middleware summary",
@@ -266,8 +259,6 @@ def create_agent_node(
     else:
         logger.debug("No custom prompt configured", agent=agent.name)
 
-    checkpointer: bool = memory is not None and memory.checkpointer is not None
-
     # Get the prompt as middleware (always returns AgentMiddleware or None)
     prompt_middleware: AgentMiddleware | None = make_prompt(agent.prompt)
 
@@ -298,12 +289,14 @@ def create_agent_node(
     # Use LangChain v1's create_agent with middleware
     # AgentState extends MessagesState with additional DAO AI fields
     # System prompt is provided via middleware (dynamic_prompt)
+    # NOTE: checkpointer=False because these agents are used as subgraphs
+    # within the parent orchestration graph (swarm/supervisor) which handles
+    # checkpointing at the root level. Subgraphs cannot have checkpointer=True.
     logger.info(
         "Creating LangChain agent",
         agent=agent.name,
         tools_count=len(tools),
         middleware_count=len(middleware_list),
-        has_checkpointer=checkpointer,
     )
 
     compiled_agent: CompiledStateGraph = create_agent(
@@ -311,7 +304,7 @@ def create_agent_node(
         model=llm,
         tools=tools,
         middleware=middleware_list,
-        checkpointer=checkpointer,
+        checkpointer=False,
         state_schema=AgentState,
         context_schema=Context,
         response_format=response_format,  # Add structured output support

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,
@@ -190,6 +205,7 @@ def create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
     supervisor_tools: list[BaseTool] = list(create_tools(supervisor_config.tools))
 
     # Create middleware from configuration
+    # All middleware factories return list[AgentMiddleware] for composability
     middlewares: list[AgentMiddleware] = []
 
     for middleware_config in supervisor_config.middleware:
@@ -201,11 +217,11 @@ def create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
             function_name=middleware_config.name,
             args=middleware_config.args,
         )
-        if middleware is not None:
-            middlewares.append(middleware)
-            logger.debug(
-                "Created supervisor middleware", middleware=middleware_config.name
-            )
+        middlewares.append(middleware)
+        logger.debug(
+            "Created supervisor middleware",
+            middleware=middleware_config.name,
+        )
 
     # Set up memory store and checkpointer
     store: BaseStore | None = create_store(orchestration)

dao_ai/orchestration/swarm.py

@@ -167,8 +167,13 @@ def create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
     default_agent: str
     if isinstance(swarm.default_agent, AgentModel):
         default_agent = swarm.default_agent.name
-    else:
+    elif swarm.default_agent is not None:
         default_agent = swarm.default_agent
+    elif len(config.app.agents) > 0:
+        # Fallback to first agent if no default specified
+        default_agent = config.app.agents[0].name
+    else:
+        raise ValueError("Swarm requires at least one agent and a default_agent")
 
     logger.info(
         "Creating swarm graph",

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],
@@ -61,19 +70,14 @@ def make_prompt(
     @dynamic_prompt
     def dynamic_system_prompt(request: ModelRequest) -> str:
         """Generate dynamic system prompt based on runtime context."""
-        # Get parameters from runtime context
+        # Initialize parameters for template variables
         params: dict[str, Any] = {
             input_variable: "" for input_variable in prompt_template.input_variables
         }
 
-        # Access context from runtime
+        # Apply context fields as template parameters
         context: Context = request.runtime.context
         if context:
-            if context.user_id and "user_id" in params:
-                params["user_id"] = context.user_id
-            if context.thread_id and "thread_id" in params:
-                params["thread_id"] = context.thread_id
-            # Apply all context fields as template parameters
             context_dict = context.model_dump()
             for key, value in context_dict.items():
                 if key in params and value is not None:
@@ -89,56 +93,3 @@ def make_prompt(
         return formatted_prompt
 
     return dynamic_system_prompt
-
-
-def create_prompt_middleware(
-    base_system_prompt: Optional[str | PromptModel],
-) -> AgentMiddleware | None:
-    """
-    Create a dynamic prompt middleware from configuration.
-
-    This always returns an AgentMiddleware suitable for use with
-    LangChain v1's middleware system.
-
-    Args:
-        base_system_prompt: The system prompt string or PromptModel
-
-    Returns:
-        An AgentMiddleware created by @dynamic_prompt, or None if no prompt
-    """
-    if not base_system_prompt:
-        return None
-
-    # Extract template string from PromptModel or use string directly
-    template_str: str
-    if isinstance(base_system_prompt, PromptModel):
-        template_str = base_system_prompt.template
-    else:
-        template_str = base_system_prompt
-
-    prompt_template: PromptTemplate = PromptTemplate.from_template(template_str)
-
-    @dynamic_prompt
-    def prompt_middleware(request: ModelRequest) -> str:
-        """Generate system prompt based on runtime context."""
-        # Get parameters from runtime context
-        params: dict[str, Any] = {
-            input_variable: "" for input_variable in prompt_template.input_variables
-        }
-
-        # Access context from runtime
-        context: Context = request.runtime.context
-        if context:
-            # Apply all context fields as template parameters
-            context_dict = context.model_dump()
-            for key, value in context_dict.items():
-                if key in params and value is not None:
-                    params[key] = value
-
-        # Format the prompt
-        formatted_prompt: str = prompt_template.format(**params)
-        logger.trace("Formatted dynamic prompt with context")
-
-        return formatted_prompt
-
-    return prompt_middleware

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/base.py

@@ -1,15 +1,19 @@
 from abc import ABC, abstractmethod
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from dao_ai.config import (
     AppModel,
     DatasetModel,
+    DeploymentTarget,
     SchemaModel,
     UnityCatalogFunctionSqlModel,
     VectorStoreModel,
     VolumeModel,
 )
 
+if TYPE_CHECKING:
+    from dao_ai.config import AppConfig
+
 
 class ServiceProvider(ABC):
     @abstractmethod
@@ -52,4 +56,26 @@ class ServiceProvider(ABC):
     ) -> Any: ...
 
     @abstractmethod
-    def deploy_agent(self, config: AppModel) -> Any: ...
+    def deploy_model_serving_agent(self, config: "AppConfig") -> Any:
+        """Deploy agent to Databricks Model Serving endpoint."""
+        ...
+
+    @abstractmethod
+    def deploy_apps_agent(self, config: "AppConfig") -> Any:
+        """Deploy agent as a Databricks App."""
+        ...
+
+    @abstractmethod
+    def deploy_agent(
+        self,
+        config: "AppConfig",
+        target: DeploymentTarget = DeploymentTarget.MODEL_SERVING,
+    ) -> Any:
+        """
+        Deploy agent to the specified target.
+
+        Args:
+            config: The AppConfig containing deployment configuration
+            target: The deployment target (MODEL_SERVING or APPS)
+        """
+        ...