dao-ai 0.0.36__py3-none-any.whl → 0.1.0__py3-none-any.whl

dao_ai/orchestration/__init__.py

@@ -0,0 +1,52 @@
+"""
+Orchestration patterns for DAO AI multi-agent systems.
+
+This package provides factory functions for creating LangGraph workflows
+that orchestrate multiple agents using the supervisor and swarm patterns.
+
+Supervisor Pattern:
+    A central supervisor coordinates specialized worker agents. The supervisor
+    hands off control to agents who then control the conversation. Agents can
+    hand back to the supervisor when done or hand off to other agents.
+
+Swarm Pattern:
+    Agents can directly transfer control to each other using handoff tools.
+    The active agent changes, and the user may continue interacting with
+    the new agent. This provides decentralized, peer-to-peer collaboration.
+
+Both patterns use Command(goto=...) for routing between agent nodes in
+the workflow graph.
+
+See: https://docs.langchain.com/oss/python/langchain/multi-agent
+See: https://github.com/langchain-ai/langgraph-supervisor-py
+See: https://github.com/langchain-ai/langgraph-swarm-py
+"""
+
+from dao_ai.orchestration.core import (
+    SUPERVISOR_NODE,
+    OutputMode,
+    create_agent_node_handler,
+    create_checkpointer,
+    create_handoff_tool,
+    create_orchestration_graph,
+    create_store,
+    extract_agent_response,
+    filter_messages_for_agent,
+    get_handoff_description,
+)
+
+__all__ = [
+    # Constants
+    "SUPERVISOR_NODE",
+    "OutputMode",
+    # Core utilities
+    "create_store",
+    "create_checkpointer",
+    "filter_messages_for_agent",
+    "extract_agent_response",
+    "create_agent_node_handler",
+    "create_handoff_tool",
+    "get_handoff_description",
+    # Main factory
+    "create_orchestration_graph",
+]

dao_ai/orchestration/core.py

@@ -0,0 +1,287 @@
+"""
+Core orchestration utilities and infrastructure.
+
+This module provides the foundational utilities for multi-agent orchestration:
+- Memory and checkpointer creation
+- Message filtering and extraction
+- Agent node handlers
+- Handoff tools
+- Main orchestration graph factory
+"""
+
+from typing import Awaitable, Callable, Literal
+
+from langchain.tools import ToolRuntime, tool
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
+from langchain_core.tools import BaseTool
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.runtime import Runtime
+from langgraph.store.base import BaseStore
+from langgraph.types import Command
+from loguru import logger
+
+from dao_ai.config import AgentModel, AppConfig, OrchestrationModel
+from dao_ai.messages import last_ai_message
+from dao_ai.state import AgentState, Context
+
+# Constant for supervisor node name
+SUPERVISOR_NODE = "supervisor"
+
+# Output mode for agent responses
+# - "full_history": Include all messages from the agent's execution
+# - "last_message": Include only the final AI message from the agent
+OutputMode = Literal["full_history", "last_message"]
+
+
+def create_store(orchestration: OrchestrationModel) -> BaseStore | None:
+    """
+    Create a memory store from orchestration config.
+
+    Args:
+        orchestration: The orchestration configuration
+
+    Returns:
+        The configured store, or None if not configured
+    """
+    if orchestration.memory and orchestration.memory.store:
+        store = orchestration.memory.store.as_store()
+        logger.debug(f"Using memory store: {store}")
+        return store
+    return None
+
+
+def create_checkpointer(
+    orchestration: OrchestrationModel,
+) -> BaseCheckpointSaver | None:
+    """
+    Create a checkpointer from orchestration config.
+
+    Args:
+        orchestration: The orchestration configuration
+
+    Returns:
+        The configured checkpointer, or None if not configured
+    """
+    if orchestration.memory and orchestration.memory.checkpointer:
+        checkpointer = orchestration.memory.checkpointer.as_checkpointer()
+        logger.debug(f"Using checkpointer: {checkpointer}")
+        return checkpointer
+    return None
+
+
+def filter_messages_for_agent(messages: list[BaseMessage]) -> list[BaseMessage]:
+    """
+    Filter messages for a worker agent to avoid tool_use/tool_result pairing errors.
+
+    When the supervisor hands off to an agent, the agent should only see:
+    - HumanMessage (user queries)
+    - AIMessage with content (previous responses, but not tool calls)
+
+    This prevents the agent from seeing orphaned ToolMessages or AIMessages
+    with tool_calls that don't belong to the agent's context.
+
+    Args:
+        messages: The full message history from parent state
+
+    Returns:
+        Filtered messages safe for the agent to process
+    """
+    filtered: list[BaseMessage] = []
+    for msg in messages:
+        if isinstance(msg, HumanMessage):
+            # Always include user messages
+            filtered.append(msg)
+        elif isinstance(msg, AIMessage):
+            # Include AI messages but strip tool_calls to avoid confusion
+            if msg.content and not msg.tool_calls:
+                filtered.append(msg)
+            elif msg.content and msg.tool_calls:
+                # Include content but create clean AIMessage without tool_calls
+                filtered.append(AIMessage(content=msg.content, id=msg.id))
+        # Skip ToolMessages - they belong to the supervisor's context
+    return filtered
+
+
+def extract_agent_response(
+    messages: list[BaseMessage],
+    output_mode: OutputMode = "last_message",
+) -> list[BaseMessage]:
+    """
+    Extract the agent's response based on the output mode.
+
+    Args:
+        messages: The agent's full message history after execution
+        output_mode: How to extract the response
+            - "full_history": Return all messages (may cause issues)
+            - "last_message": Return only the final AI message
+
+    Returns:
+        Messages to include in the parent state update
+    """
+    if output_mode == "full_history":
+        return messages
+
+    # Find the last AI message with content
+    final_response: AIMessage | None = last_ai_message(messages)
+
+    if final_response:
+        # Return clean AIMessage without tool_calls
+        if final_response.tool_calls:
+            return [AIMessage(content=final_response.content, id=final_response.id)]
+        return [final_response]
+
+    return []
+
+
+def create_agent_node_handler(
+    agent_name: str,
+    agent: CompiledStateGraph,
+    output_mode: OutputMode = "last_message",
+) -> Callable[[AgentState, Runtime[Context]], Awaitable[AgentState]]:
+    """
+    Create a handler that wraps an agent subgraph with message filtering.
+
+    This filters messages before passing to the agent and extracts only
+    the relevant response, avoiding tool_use/tool_result pairing errors.
+
+    Used by both supervisor and swarm patterns to ensure consistent
+    message handling when agents are CompiledStateGraphs.
+
+    Based on langgraph-supervisor-py output_mode pattern.
+
+    Args:
+        agent_name: Name of the agent (for logging)
+        agent: The compiled agent subgraph
+        output_mode: How to extract response ("last_message" or "full_history")
+
+    Returns:
+        An async handler function for the workflow node
+    """
+
+    async def handler(state: AgentState, runtime: Runtime[Context]) -> AgentState:
+        # Filter messages to avoid tool_use/tool_result pairing errors
+        original_messages = state.get("messages", [])
+        filtered_messages = filter_messages_for_agent(original_messages)
+
+        logger.debug(
+            f"Agent '{agent_name}' receiving {len(filtered_messages)} filtered "
+            f"messages (from {len(original_messages)} total)"
+        )
+
+        # Create state with filtered messages for the agent
+        agent_state: AgentState = {
+            **state,
+            "messages": filtered_messages,
+        }
+
+        # Invoke the agent
+        result: AgentState = await agent.ainvoke(agent_state, context=runtime.context)
+
+        # Extract agent response based on output mode
+        result_messages = result.get("messages", [])
+        response_messages = extract_agent_response(result_messages, output_mode)
+
+        logger.debug(
+            f"Agent '{agent_name}' completed. Returning {len(response_messages)} "
+            f"messages (from {len(result_messages)} total, mode={output_mode})"
+        )
+
+        # Return state update with extracted response
+        return {
+            **result,
+            "messages": response_messages,
+        }
+
+    return handler
+
+
+def create_handoff_tool(
+    target_agent_name: str,
+    description: str,
+) -> BaseTool:
+    """
+    Create a handoff tool that transfers control to another agent.
+
+    The tool returns a Command object with goto to directly route
+    to the target agent node in the parent graph.
+
+    Args:
+        target_agent_name: The name of the agent to hand off to
+        description: Description of what this agent handles
+
+    Returns:
+        A tool that triggers a handoff to the target agent via Command
+    """
+
+    @tool
+    def handoff_tool(runtime: ToolRuntime[Context, AgentState]) -> Command:
+        """Transfer control to another agent."""
+        tool_call_id: str = runtime.tool_call_id
+        logger.debug(f"Handoff to agent '{target_agent_name}'")
+
+        return Command(
+            update={
+                "active_agent": target_agent_name,
+                "messages": [
+                    ToolMessage(
+                        content=f"Transferred to {target_agent_name}",
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            },
+            goto=target_agent_name,
+            graph=Command.PARENT,
+        )
+
+    # Set the tool name and description
+    handoff_tool.name = f"handoff_to_{target_agent_name}"
+    handoff_tool.__doc__ = f"Transfer to {target_agent_name}: {description}"
+    handoff_tool.description = f"Transfer to {target_agent_name}: {description}"
+
+    return handoff_tool
+
+
+def get_handoff_description(agent: AgentModel) -> str:
+    """
+    Get the handoff description for an agent.
+
+    Priority: handoff_prompt > description > default message
+
+    Args:
+        agent: The agent to get the handoff description for
+
+    Returns:
+        The handoff description string
+    """
+    return (
+        agent.handoff_prompt
+        or agent.description
+        or f"Handles {agent.name} related tasks and inquiries"
+    )
+
+
+def create_orchestration_graph(config: AppConfig) -> CompiledStateGraph:
+    """
+    Create the main orchestration graph based on the configuration.
+
+    This factory function creates either a supervisor or swarm graph
+    depending on the configuration.
+
+    Args:
+        config: The application configuration
+
+    Returns:
+        A compiled LangGraph state machine
+    """
+    from dao_ai.orchestration.supervisor import create_supervisor_graph
+    from dao_ai.orchestration.swarm import create_swarm_graph
+
+    orchestration: OrchestrationModel = config.app.orchestration
+    if orchestration.supervisor:
+        return create_supervisor_graph(config)
+
+    if orchestration.swarm:
+        return create_swarm_graph(config)
+
+    raise ValueError("No valid orchestration model found in the configuration.")

dao_ai/orchestration/supervisor.py

@@ -0,0 +1,264 @@
+"""
+Supervisor pattern for multi-agent orchestration.
+
+The supervisor pattern uses a central supervisor agent that coordinates
+specialized worker agents. The supervisor hands off control to agents
+who then control the conversation. Agents can hand back to the supervisor
+when done.
+
+Based on: https://github.com/langchain-ai/langgraph-supervisor-py
+"""
+
+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.tools import BaseTool
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.store.base import BaseStore
+from langgraph.types import Command
+from langmem import create_manage_memory_tool
+from loguru import logger
+
+from dao_ai.config import (
+    AppConfig,
+    MemoryModel,
+    OrchestrationModel,
+    SupervisorModel,
+)
+from dao_ai.middleware.base import AgentMiddleware
+from dao_ai.middleware.core import create_factory_middleware
+from dao_ai.nodes import create_agent_node
+from dao_ai.orchestration import (
+    SUPERVISOR_NODE,
+    create_agent_node_handler,
+    create_checkpointer,
+    create_handoff_tool,
+    create_store,
+    get_handoff_description,
+)
+from dao_ai.prompts import make_prompt
+from dao_ai.state import AgentState, Context
+from dao_ai.tools import create_tools
+from dao_ai.tools.memory import create_search_memory_tool
+
+
+def _create_handoff_back_to_supervisor_tool() -> BaseTool:
+    """
+    Create a tool for agents to hand control back to the supervisor.
+
+    This is used in the supervisor pattern when an agent has completed
+    its task and wants to return control to the supervisor for further
+    coordination or to complete the conversation.
+
+    Returns:
+        A tool that routes back to the supervisor node
+    """
+
+    @tool
+    def handoff_to_supervisor(
+        summary: str,
+        runtime: ToolRuntime[Context, AgentState],
+    ) -> Command:
+        """
+        Hand control back to the supervisor.
+
+        Use this when you have completed your task and want to return
+        control to the supervisor for further coordination.
+
+        Args:
+            summary: A brief summary of what was accomplished
+        """
+        tool_call_id: str = runtime.tool_call_id
+        logger.debug(f"Handoff back to supervisor with summary: {summary[:100]}...")
+
+        return Command(
+            update={
+                "active_agent": None,
+                "messages": [
+                    ToolMessage(
+                        content=f"Task completed: {summary}",
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            },
+            goto=SUPERVISOR_NODE,
+            graph=Command.PARENT,
+        )
+
+    return handoff_to_supervisor
+
+
+def _create_supervisor_agent(
+    config: AppConfig,
+    tools: list[BaseTool],
+    handoff_tools: list[BaseTool],
+    middlewares: list[AgentMiddleware],
+) -> CompiledStateGraph:
+    """
+    Create a supervisor agent with handoff tools for each worker agent.
+
+    The supervisor coordinates worker agents by handing off control.
+    Worker agents take over the conversation and can hand back to
+    the supervisor when done.
+
+    Args:
+        config: Application configuration
+        tools: Additional tools for the supervisor (e.g., memory tools)
+        handoff_tools: Handoff tools to route to worker agents
+        middlewares: Middleware to apply to the supervisor
+
+    Returns:
+        Compiled supervisor agent
+    """
+    orchestration: OrchestrationModel = config.app.orchestration
+    supervisor: SupervisorModel = orchestration.supervisor
+
+    all_tools: list[BaseTool] = list(tools) + list(handoff_tools)
+
+    model: LanguageModelLike = supervisor.model.as_chat_model()
+
+    # Get the prompt as middleware (always returns AgentMiddleware or None)
+    prompt_middleware: LangchainAgentMiddleware | None = make_prompt(supervisor.prompt)
+
+    # Add prompt middleware at the beginning for priority
+    if prompt_middleware is not None:
+        middlewares.insert(0, prompt_middleware)
+
+    # Create the supervisor agent
+    # Handoff tools route to worker agents in the parent workflow graph
+    supervisor_agent: CompiledStateGraph = create_agent(
+        name=SUPERVISOR_NODE,
+        model=model,
+        tools=all_tools,
+        middleware=middlewares,
+        state_schema=AgentState,
+        context_schema=Context,
+    )
+
+    return supervisor_agent
+
+
+def create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
+    """
+    Create a supervisor-based multi-agent system using handoffs.
+
+    This implements a supervisor pattern where:
+    1. Supervisor receives user input and decides which agent to hand off to
+    2. Agent takes control of the conversation and interacts with user
+    3. Agent can hand back to supervisor or complete the task
+
+    The supervisor and all worker agents are nodes in a workflow graph.
+    Handoff tools use Command(goto=..., graph=Command.PARENT) to route
+    between nodes.
+
+    Args:
+        config: The application configuration
+
+    Returns:
+        A compiled LangGraph state machine
+
+    Based on: https://github.com/langchain-ai/langgraph-supervisor-py
+    """
+    logger.debug("Creating supervisor graph (handoff pattern)")
+
+    orchestration: OrchestrationModel = config.app.orchestration
+    supervisor_config: SupervisorModel = orchestration.supervisor
+
+    # Create handoff tools for supervisor to route to agents
+    handoff_tools: list[BaseTool] = []
+    for registered_agent in config.app.agents:
+        description: str = get_handoff_description(registered_agent)
+        handoff_tool: BaseTool = create_handoff_tool(
+            target_agent_name=registered_agent.name,
+            description=description,
+        )
+        handoff_tools.append(handoff_tool)
+        logger.debug(f"Created handoff tool for supervisor: {registered_agent.name}")
+
+    # Create supervisor's own tools (e.g., memory tools)
+    supervisor_tools: list[BaseTool] = list(create_tools(supervisor_config.tools))
+
+    # Create middleware from configuration
+    middlewares: list[AgentMiddleware] = []
+    for middleware_config in supervisor_config.middleware:
+        middleware = create_factory_middleware(
+            function_name=middleware_config.name,
+            args=middleware_config.args,
+        )
+        if middleware is not None:
+            middlewares.append(middleware)
+            logger.debug(f"Created supervisor middleware: {middleware_config.name}")
+
+    # Set up memory store and checkpointer
+    store: BaseStore | None = create_store(orchestration)
+    checkpointer: BaseCheckpointSaver | None = create_checkpointer(orchestration)
+
+    # Add memory tools if store is configured with namespace
+    if (
+        orchestration.memory
+        and orchestration.memory.store
+        and orchestration.memory.store.namespace
+    ):
+        namespace: tuple[str, ...] = ("memory", orchestration.memory.store.namespace)
+        logger.debug(f"Memory store namespace: {namespace}")
+        # Use Databricks-compatible search_memory tool (omits problematic filter field)
+        supervisor_tools += [
+            create_manage_memory_tool(namespace=namespace),
+            create_search_memory_tool(namespace=namespace),
+        ]
+
+    # Create the supervisor agent
+    supervisor_agent: CompiledStateGraph = _create_supervisor_agent(
+        config=config,
+        tools=supervisor_tools,
+        handoff_tools=handoff_tools,
+        middlewares=middlewares,
+    )
+
+    # Create worker agent subgraphs
+    # Each worker gets a handoff_to_supervisor tool to return control
+    agent_subgraphs: dict[str, CompiledStateGraph] = {}
+    memory: MemoryModel | None = orchestration.memory
+    for registered_agent in config.app.agents:
+        # Create handoff back to supervisor tool
+        supervisor_handoff: BaseTool = _create_handoff_back_to_supervisor_tool()
+
+        # Create the worker agent with handoff back to supervisor
+        agent_subgraph: CompiledStateGraph = create_agent_node(
+            agent=registered_agent,
+            memory=memory,
+            chat_history=config.app.chat_history,
+            additional_tools=[supervisor_handoff],
+        )
+        agent_subgraphs[registered_agent.name] = agent_subgraph
+        logger.debug(f"Created worker agent subgraph: {registered_agent.name}")
+
+    # Build the workflow graph
+    # All agents are nodes, handoffs route between them via Command
+    workflow: StateGraph = StateGraph(
+        AgentState,
+        input=AgentState,
+        output=AgentState,
+        context_schema=Context,
+    )
+
+    # Add supervisor node
+    workflow.add_node(SUPERVISOR_NODE, supervisor_agent)
+
+    # Add worker agent nodes with message filtering handlers
+    for agent_name, agent_subgraph in agent_subgraphs.items():
+        handler = create_agent_node_handler(
+            agent_name=agent_name,
+            agent=agent_subgraph,
+            output_mode="last_message",  # Only return final response to avoid issues
+        )
+        workflow.add_node(agent_name, handler)
+
+    # Supervisor is the entry point
+    workflow.set_entry_point(SUPERVISOR_NODE)
+
+    return workflow.compile(checkpointer=checkpointer, store=store)