dao-ai 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

dao_ai/nodes.py

@@ -25,8 +25,14 @@ 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 create_hitl_middleware_from_tool_models
-from dao_ai.middleware.summarization import create_summarization_middleware
+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
 from dao_ai.state import AgentState, Context
 from dao_ai.tools import create_tools
@@ -50,14 +56,25 @@ def _create_middleware_list(
         List of middleware instances (can include both AgentMiddleware and
         LangChain built-in middleware)
     """
-    logger.debug(f"Building middleware list for agent '{agent.name}'")
+    logger.debug("Building middleware list for agent", agent=agent.name)
     middleware_list: list[Any] = []
 
     # Add configured middleware using factory pattern
     if agent.middleware:
-        logger.debug(f"Processing {len(agent.middleware)} configured middleware")
+        middleware_names: list[str] = [mw.name for mw in agent.middleware]
+        logger.info(
+            "Middleware configuration",
+            agent=agent.name,
+            middleware_count=len(agent.middleware),
+            middleware_names=middleware_names,
+        )
     for middleware_config in agent.middleware:
-        middleware = create_factory_middleware(
+        logger.trace(
+            "Creating middleware for agent",
+            agent=agent.name,
+            middleware_name=middleware_config.name,
+        )
+        middleware: AgentMiddleware[AgentState, Context] = create_factory_middleware(
             function_name=middleware_config.name,
             args=middleware_config.args,
         )
@@ -66,7 +83,13 @@ def _create_middleware_list(
 
     # Add guardrails as middleware
     if agent.guardrails:
-        logger.debug(f"Adding {len(agent.guardrails)} guardrail middleware")
+        guardrail_names: list[str] = [gr.name for gr in agent.guardrails]
+        logger.info(
+            "Guardrails configuration",
+            agent=agent.name,
+            guardrails_count=len(agent.guardrails),
+            guardrail_names=guardrail_names,
+        )
     for guardrail in agent.guardrails:
         # Extract template string from PromptModel if needed
         prompt_str: str
@@ -75,28 +98,54 @@ def _create_middleware_list(
         else:
             prompt_str = guardrail.prompt
 
-        guardrail_middleware = GuardrailMiddleware(
+        guardrail_middleware: GuardrailMiddleware = GuardrailMiddleware(
             name=guardrail.name,
             model=guardrail.model.as_chat_model(),
             prompt=prompt_str,
             num_retries=guardrail.num_retries or 3,
         )
-        logger.debug(f"Created guardrail middleware: {guardrail.name}")
+        logger.trace(
+            "Created guardrail middleware", guardrail=guardrail.name, agent=agent.name
+        )
         middleware_list.append(guardrail_middleware)
 
     # Add summarization middleware if chat_history is configured
     if chat_history is not None:
-        logger.debug("Adding summarization middleware")
-        summarization_middleware = create_summarization_middleware(chat_history)
+        logger.info(
+            "Chat history configuration",
+            agent=agent.name,
+            max_tokens=chat_history.max_tokens,
+            summary_model=chat_history.model.name,
+        )
+        summarization_middleware: LoggingSummarizationMiddleware = (
+            create_summarization_middleware(chat_history)
+        )
         middleware_list.append(summarization_middleware)
 
     # Add human-in-the-loop middleware if any tools require it
-    hitl_middleware = create_hitl_middleware_from_tool_models(tool_models)
+    hitl_middleware: HumanInTheLoopMiddleware | None = (
+        create_hitl_middleware_from_tool_models(tool_models)
+    )
     if hitl_middleware is not None:
-        logger.debug("Added human-in-the-loop middleware")
+        # Log which tools require HITL
+        hitl_tool_names: list[str] = [
+            tool.name
+            for tool in tool_models
+            if hasattr(tool.function, "human_in_the_loop")
+            and tool.function.human_in_the_loop is not None
+        ]
+        logger.info(
+            "Human-in-the-Loop configuration",
+            agent=agent.name,
+            hitl_tools=hitl_tool_names,
+        )
         middleware_list.append(hitl_middleware)
 
-    logger.debug(f"Total middleware count: {len(middleware_list)}")
+    logger.info(
+        "Middleware summary",
+        agent=agent.name,
+        total_middleware_count=len(middleware_list),
+    )
     return middleware_list
 
 
@@ -122,26 +171,71 @@ def create_agent_node(
     Returns:
         RunnableLike: An agent node that processes state and returns responses
     """
-    logger.debug(f"Creating agent node for {agent.name}")
+    logger.info("Creating agent node", agent=agent.name)
+
+    # Log agent configuration details
+    logger.info(
+        "Agent configuration",
+        agent=agent.name,
+        model=agent.model.name,
+        description=agent.description or "No description",
+    )
 
     llm: LanguageModelLike = agent.model.as_chat_model()
 
     tool_models: Sequence[ToolModel] = agent.tools
     if not additional_tools:
         additional_tools = []
+
+    # Log tools being created
+    tool_names: list[str] = [tool.name for tool in tool_models]
+    logger.info(
+        "Tools configuration",
+        agent=agent.name,
+        tools_count=len(tool_models),
+        tool_names=tool_names,
+    )
+
     tools: list[BaseTool] = list(create_tools(tool_models)) + list(additional_tools)
 
+    if additional_tools:
+        logger.debug(
+            "Additional tools added",
+            agent=agent.name,
+            additional_count=len(additional_tools),
+        )
+
     if memory and memory.store:
         namespace: tuple[str, ...] = ("memory",)
         if memory.store.namespace:
             namespace = namespace + (memory.store.namespace,)
-        logger.debug(f"Memory store namespace: {namespace}")
+        logger.info(
+            "Memory configuration",
+            agent=agent.name,
+            has_store=True,
+            has_checkpointer=memory.checkpointer is not None,
+            namespace=namespace,
+        )
+    elif memory:
+        logger.info(
+            "Memory configuration",
+            agent=agent.name,
+            has_store=False,
+            has_checkpointer=memory.checkpointer is not None,
+        )
 
+    # Add memory tools if store is configured
+    if memory and memory.store:
         # Use Databricks-compatible search_memory tool (omits problematic filter field)
         tools += [
             create_manage_memory_tool(namespace=namespace),
             create_search_memory_tool(namespace=namespace),
         ]
+        logger.debug(
+            "Memory tools added",
+            agent=agent.name,
+            tools=["manage_memory", "search_memory"],
+        )
 
     # Create middleware list from configuration
     middleware_list = _create_middleware_list(
@@ -150,7 +244,27 @@ def create_agent_node(
         chat_history=chat_history,
     )
 
-    logger.debug(f"Created {len(middleware_list)} middleware for agent {agent.name}")
+    # Log prompt configuration
+    if agent.prompt:
+        if isinstance(agent.prompt, PromptModel):
+            logger.info(
+                "Prompt configuration",
+                agent=agent.name,
+                prompt_type="PromptModel",
+                prompt_name=agent.prompt.name,
+            )
+        else:
+            prompt_preview: str = (
+                agent.prompt[:100] + "..." if len(agent.prompt) > 100 else agent.prompt
+            )
+            logger.info(
+                "Prompt configuration",
+                agent=agent.name,
+                prompt_type="string",
+                prompt_preview=prompt_preview,
+            )
+    else:
+        logger.debug("No custom prompt configured", agent=agent.name)
 
     checkpointer: bool = memory is not None and memory.checkpointer is not None
 
@@ -167,18 +281,31 @@ def create_agent_node(
         try:
             response_format = agent.response_format.as_strategy()
             if response_format is not None:
-                logger.debug(
-                    f"Agent '{agent.name}' using structured output: {type(response_format).__name__}"
+                logger.info(
+                    "Response format configuration",
+                    agent=agent.name,
+                    format_type=type(response_format).__name__,
+                    structured_output=True,
                 )
         except ValueError as e:
             logger.error(
-                f"Failed to configure structured output for agent {agent.name}: {e}"
+                "Failed to configure structured output for agent",
+                agent=agent.name,
+                error=str(e),
             )
             raise
 
     # Use LangChain v1's create_agent with middleware
     # AgentState extends MessagesState with additional DAO AI fields
     # System prompt is provided via middleware (dynamic_prompt)
+    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(
         name=agent.name,
         model=llm,
@@ -192,4 +319,6 @@ def create_agent_node(
 
     compiled_agent.name = agent.name
 
+    logger.info("Agent node created successfully", agent=agent.name)
+
     return compiled_agent

dao_ai/optimization.py

@@ -245,7 +245,7 @@ class DAOAgentAdapter(GEPAAdapter[_TrainingExample, _Trajectory, str]):
                         )
 
                 except Exception as e:
-                    logger.warning(f"Error evaluating example: {e}")
+                    logger.warning("Error evaluating example", error=str(e))
                     outputs.append("")
                     scores.append(0.0)
 
@@ -362,7 +362,9 @@ def _convert_dataset(
         )
         examples.append(example)
 
-    logger.debug(f"Converted {len(examples)} dataset entries to training examples")
+    logger.debug(
+        "Converted dataset entries to training examples", examples_count=len(examples)
+    )
     return examples
 
 
@@ -400,7 +402,7 @@ def _register_optimized_prompt(
     prompt_name: str = prompt.full_name
     optimization_timestamp: str = datetime.now(timezone.utc).isoformat()
 
-    logger.info(f"Registering optimized prompt: {prompt_name}")
+    logger.info("Registering optimized prompt", prompt_name=prompt_name)
 
     # Build comprehensive tags for the prompt registry
     tags: dict[str, str] = {
@@ -442,7 +444,11 @@ def _register_optimized_prompt(
         tags=tags,
     )
 
-    logger.info(f"Registered as version {version.version}")
+    logger.success(
+        "Registered optimized prompt version",
+        prompt_name=prompt_name,
+        version=version.version,
+    )
 
     # Set 'latest' alias for most recently optimized version
     mlflow.genai.set_prompt_alias(
@@ -450,7 +456,7 @@ def _register_optimized_prompt(
         alias="latest",
         version=version.version,
     )
-    logger.info(f"Set 'latest' alias for version {version.version}")
+    logger.info("Set 'latest' alias", prompt_name=prompt_name, version=version.version)
 
     # Set 'champion' alias if there was actual improvement
     if improvement > 0:
@@ -459,7 +465,9 @@ def _register_optimized_prompt(
             alias="champion",
             version=version.version,
         )
-        logger.info(f"Set 'champion' alias for version {version.version}")
+        logger.success(
+            "Set 'champion' alias", prompt_name=prompt_name, version=version.version
+        )
 
     return version
 
@@ -518,7 +526,7 @@ def optimize_prompt(
         if result.improved:
             print(f"Improved by {result.improvement:.1%}")
     """
-    logger.info(f"Starting GEPA optimization for prompt: {prompt.name}")
+    logger.info("Starting GEPA optimization", prompt_name=prompt.name)
 
     # Get the original template
     original_template = prompt.template
@@ -535,11 +543,11 @@ def optimize_prompt(
     trainset = examples[:split_idx]
     valset = examples[split_idx:] if split_idx < len(examples) else examples
 
-    logger.info(f"Using {len(trainset)} train, {len(valset)} val examples")
+    logger.info("Dataset split", train_size=len(trainset), val_size=len(valset))
 
     # Get reflection model
     reflection_model_name = reflection_model or agent.model.uri
-    logger.info(f"Using reflection model: {reflection_model_name}")
+    logger.info("Using reflection model", model=reflection_model_name)
 
     # Create adapter
     adapter = DAOAgentAdapter(agent_model=agent, metric_fn=metric)
@@ -548,7 +556,7 @@ def optimize_prompt(
     seed_candidate = {"prompt": original_template}
 
     # Run GEPA optimization
-    logger.info(f"Running GEPA optimization (max {num_candidates} evaluations)...")
+    logger.info("Running GEPA optimization", max_evaluations=num_candidates)
 
     try:
         result: GEPAResult = optimize(
@@ -562,7 +570,7 @@ def optimize_prompt(
             skip_perfect_score=True,
         )
     except Exception as e:
-        logger.error(f"GEPA optimization failed: {e}")
+        logger.error("GEPA optimization failed", error=str(e))
         return OptimizationResult(
             optimized_prompt=prompt,
             optimized_template=original_template,
@@ -596,10 +604,12 @@ def optimize_prompt(
         else 0.0
     )
 
-    logger.info("Optimization complete!")
-    logger.info(f"Original score: {original_score:.3f}")
-    logger.info(f"Optimized score: {optimized_score:.3f}")
-    logger.info(f"Improvement: {improvement:.1%}")
+    logger.success(
+        "Optimization complete",
+        original_score=f"{original_score:.3f}",
+        optimized_score=f"{optimized_score:.3f}",
+        improvement=f"{improvement:.1%}",
+    )
 
     # Register if improved
     registered_version: Optional[PromptVersion] = None
@@ -623,7 +633,7 @@ def optimize_prompt(
                 val_size=len(valset),
             )
         except Exception as e:
-            logger.error(f"Failed to register optimized prompt: {e}")
+            logger.error("Failed to register optimized prompt", error=str(e))
 
     # Build optimized prompt model with comprehensive tags
     optimized_tags: dict[str, str] = {

dao_ai/orchestration/core.py

@@ -46,7 +46,7 @@ def create_store(orchestration: OrchestrationModel) -> BaseStore | None:
     """
     if orchestration.memory and orchestration.memory.store:
         store = orchestration.memory.store.as_store()
-        logger.debug(f"Using memory store: {store}")
+        logger.debug("Memory store configured", store_type=type(store).__name__)
         return store
     return None
 
@@ -65,7 +65,9 @@ def create_checkpointer(
     """
     if orchestration.memory and orchestration.memory.checkpointer:
         checkpointer = orchestration.memory.checkpointer.as_checkpointer()
-        logger.debug(f"Using checkpointer: {checkpointer}")
+        logger.debug(
+            "Checkpointer configured", checkpointer_type=type(checkpointer).__name__
+        )
         return checkpointer
     return None
 
@@ -164,9 +166,11 @@ def create_agent_node_handler(
         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)"
+        logger.trace(
+            "Agent receiving filtered messages",
+            agent=agent_name,
+            filtered_count=len(filtered_messages),
+            original_count=len(original_messages),
         )
 
         # Create state with filtered messages for the agent
@@ -183,8 +187,11 @@ def create_agent_node_handler(
         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})"
+            "Agent completed",
+            agent=agent_name,
+            response_count=len(response_messages),
+            total_messages=len(result_messages),
+            output_mode=output_mode,
         )
 
         # Return state update with extracted response
@@ -218,7 +225,7 @@ def create_handoff_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}'")
+        logger.debug("Handoff to agent", target_agent=target_agent_name)
 
         return Command(
             update={

dao_ai/orchestration/supervisor.py

@@ -73,7 +73,7 @@ def _create_handoff_back_to_supervisor_tool() -> BaseTool:
             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]}...")
+        logger.debug("Agent handing back to supervisor", summary_preview=summary[:100])
 
         return Command(
             update={
@@ -163,11 +163,15 @@ def create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
 
     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
 
+    logger.info(
+        "Creating supervisor graph",
+        pattern="handoff",
+        agents_count=len(config.app.agents),
+    )
+
     # Create handoff tools for supervisor to route to agents
     handoff_tools: list[BaseTool] = []
     for registered_agent in config.app.agents:
@@ -177,21 +181,31 @@ def create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
             description=description,
         )
         handoff_tools.append(handoff_tool)
-        logger.debug(f"Created handoff tool for supervisor: {registered_agent.name}")
+        logger.debug("Created handoff tool for supervisor", agent=registered_agent.name)
 
     # Create supervisor's own tools (e.g., memory tools)
+    logger.debug(
+        "Creating tools for supervisor", tools_count=len(supervisor_config.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(
+        logger.trace(
+            "Creating middleware for supervisor",
+            middleware_name=middleware_config.name,
+        )
+        middleware: LangchainAgentMiddleware = 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}")
+            logger.debug(
+                "Created supervisor middleware", middleware=middleware_config.name
+            )
 
     # Set up memory store and checkpointer
     store: BaseStore | None = create_store(orchestration)
@@ -204,7 +218,7 @@ def create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
         and orchestration.memory.store.namespace
     ):
         namespace: tuple[str, ...] = ("memory", orchestration.memory.store.namespace)
-        logger.debug(f"Memory store namespace: {namespace}")
+        logger.debug("Memory store namespace configured", namespace=namespace)
         # Use Databricks-compatible search_memory tool (omits problematic filter field)
         supervisor_tools += [
             create_manage_memory_tool(namespace=namespace),
@@ -235,7 +249,7 @@ def create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
             additional_tools=[supervisor_handoff],
         )
         agent_subgraphs[registered_agent.name] = agent_subgraph
-        logger.debug(f"Created worker agent subgraph: {registered_agent.name}")
+        logger.debug("Created worker agent subgraph", agent=registered_agent.name)
 
     # Build the workflow graph
     # All agents are nodes, handoffs route between them via Command

dao_ai/orchestration/swarm.py

@@ -69,14 +69,14 @@ def _handoffs_for_agent(
             )
 
         if handoff_to_agent is None:
-            logger.warning(
-                f"Handoff agent not found in configuration for agent {agent.name}"
-            )
+            logger.warning("Handoff agent not found in configuration", agent=agent.name)
             continue
         if agent.name == handoff_to_agent.name:
             continue
         logger.debug(
-            f"Creating handoff tool from agent {agent.name} to {handoff_to_agent.name}"
+            "Creating handoff tool",
+            from_agent=agent.name,
+            to_agent=handoff_to_agent.name,
         )
 
         handoff_description: str = get_handoff_description(handoff_to_agent)
@@ -116,19 +116,22 @@ def _create_swarm_router(
 
         # If no active agent set, use default
         if not active_agent:
-            logger.debug(
-                f"No active_agent in state, routing to default: {default_agent}"
+            logger.trace(
+                "No active agent in state, routing to default",
+                default_agent=default_agent,
             )
             return default_agent
 
         # Validate active_agent exists
         if active_agent in agent_names:
-            logger.debug(f"Routing to active_agent: {active_agent}")
+            logger.trace("Routing to active agent", active_agent=active_agent)
             return active_agent
 
         # Fallback to default if active_agent is invalid
         logger.warning(
-            f"Invalid active_agent '{active_agent}', routing to default: {default_agent}"
+            "Invalid active agent, routing to default",
+            active_agent=active_agent,
+            default_agent=default_agent,
         )
         return default_agent
 
@@ -157,8 +160,6 @@ def create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
 
     See: https://github.com/langchain-ai/langgraph-swarm-py
     """
-    logger.debug("Creating swarm graph (handoff pattern)")
-
     orchestration: OrchestrationModel = config.app.orchestration
     swarm: SwarmModel = orchestration.swarm
 
@@ -169,10 +170,27 @@ def create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
     else:
         default_agent = swarm.default_agent
 
+    logger.info(
+        "Creating swarm graph",
+        pattern="handoff",
+        default_agent=default_agent,
+        agents_count=len(config.app.agents),
+    )
+
     # Create agent subgraphs with their specific handoff tools
     # Each agent gets handoff tools only for agents they're allowed to hand off to
     agent_subgraphs: dict[str, CompiledStateGraph] = {}
     memory: MemoryModel | None = orchestration.memory
+
+    # Get swarm-level middleware to apply to all agents
+    swarm_middleware: list = swarm.middleware if swarm.middleware else []
+    if swarm_middleware:
+        logger.info(
+            "Applying swarm-level middleware to all agents",
+            middleware_count=len(swarm_middleware),
+            middleware_names=[mw.name for mw in swarm_middleware],
+        )
+
     for registered_agent in config.app.agents:
         # Get handoff tools for this agent
         handoff_tools: Sequence[BaseTool] = _handoffs_for_agent(
@@ -180,14 +198,41 @@ def create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
             config=config,
         )
 
+        # Merge swarm-level middleware with agent-specific middleware
+        # Swarm middleware is applied first, then agent middleware
+        if swarm_middleware:
+            from copy import deepcopy
+
+            # Create a copy of the agent to avoid modifying the original
+            agent_with_middleware = deepcopy(registered_agent)
+
+            # Combine swarm middleware (first) with agent middleware
+            agent_with_middleware.middleware = (
+                swarm_middleware + agent_with_middleware.middleware
+            )
+
+            logger.debug(
+                "Merged middleware for agent",
+                agent=registered_agent.name,
+                swarm_middleware_count=len(swarm_middleware),
+                agent_middleware_count=len(registered_agent.middleware),
+                total_middleware_count=len(agent_with_middleware.middleware),
+            )
+        else:
+            agent_with_middleware = registered_agent
+
         agent_subgraph: CompiledStateGraph = create_agent_node(
-            agent=registered_agent,
+            agent=agent_with_middleware,
             memory=memory,
             chat_history=config.app.chat_history,
             additional_tools=handoff_tools,
         )
         agent_subgraphs[registered_agent.name] = agent_subgraph
-        logger.debug(f"Created swarm agent subgraph: {registered_agent.name}")
+        logger.debug(
+            "Created swarm agent subgraph",
+            agent=registered_agent.name,
+            handoffs_count=len(handoff_tools),
+        )
 
     # Set up memory store and checkpointer
     store: BaseStore | None = create_store(orchestration)