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

dao_ai/nodes.py

@@ -1,101 +1,103 @@
-from typing import Any, Callable, Optional, Sequence
+"""
+Node creation utilities for DAO AI agents.
 
-import mlflow
+This module provides factory functions for creating LangGraph nodes
+that implement agent logic using LangChain v1's create_agent pattern.
+"""
+
+from typing import Any, Optional, Sequence
+
+from langchain.agents import create_agent
+from langchain.agents.middleware import AgentMiddleware
 from langchain_core.language_models import LanguageModelLike
-from langchain_core.messages import AIMessage, AnyMessage, BaseMessage
-from langchain_core.messages.utils import count_tokens_approximately
-from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.base import RunnableLike
 from langchain_core.tools import BaseTool
-from langgraph.graph import StateGraph
 from langgraph.graph.state import CompiledStateGraph
-from langgraph.prebuilt import create_react_agent
-from langgraph.runtime import Runtime
-from langmem import create_manage_memory_tool, create_search_memory_tool
-from langmem.short_term import SummarizationNode
-from langmem.short_term.summarization import TokenCounter
+from langmem import create_manage_memory_tool
 from loguru import logger
 
 from dao_ai.config import (
     AgentModel,
-    AppConfig,
     ChatHistoryModel,
-    FunctionHook,
     MemoryModel,
+    PromptModel,
     ToolModel,
 )
-from dao_ai.guardrails import reflection_guardrail, with_guardrails
-from dao_ai.hooks.core import create_hooks
+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.prompts import make_prompt
-from dao_ai.state import Context, IncomingState, SharedState
+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 summarization_node(chat_history: ChatHistoryModel) -> RunnableLike:
+def _create_middleware_list(
+    agent: AgentModel,
+    tool_models: Sequence[ToolModel],
+    chat_history: Optional[ChatHistoryModel] = None,
+) -> list[Any]:
     """
-    Create a summarization node for managing chat history.
+    Create a list of middleware instances from agent configuration.
 
     Args:
-        chat_history: ChatHistoryModel configuration for summarization
+        agent: AgentModel configuration
+        tool_models: Tool model configurations (for HITL settings)
+        chat_history: Optional chat history configuration for summarization
 
     Returns:
-        RunnableLike: A summarization node that processes messages
+        List of middleware instances (can include both AgentMiddleware and
+        LangChain built-in middleware)
     """
-    if chat_history is None:
-        raise ValueError("chat_history must be provided to use summarization")
-
-    max_tokens: int = chat_history.max_tokens
-    max_tokens_before_summary: int | None = chat_history.max_tokens_before_summary
-    max_messages_before_summary: int | None = chat_history.max_messages_before_summary
-    max_summary_tokens: int | None = chat_history.max_summary_tokens
-    token_counter: TokenCounter = (
-        count_tokens_approximately if max_tokens_before_summary else len
-    )
-
-    logger.debug(
-        f"Creating summarization node with max_tokens: {max_tokens}, "
-        f"max_tokens_before_summary: {max_tokens_before_summary}, "
-        f"max_messages_before_summary: {max_messages_before_summary}, "
-        f"max_summary_tokens: {max_summary_tokens}"
-    )
-
-    summarization_model: LanguageModelLike = chat_history.model.as_chat_model()
-
-    node: RunnableLike = SummarizationNode(
-        model=summarization_model,
-        max_tokens=max_tokens,
-        max_tokens_before_summary=max_tokens_before_summary
-        or max_messages_before_summary,
-        max_summary_tokens=max_summary_tokens,
-        token_counter=token_counter,
-        input_messages_key="messages",
-        output_messages_key="summarized_messages",
-    )
-    return node
-
-
-def call_agent_with_summarized_messages(agent: CompiledStateGraph) -> RunnableLike:
-    async def call_agent(state: SharedState, runtime: Runtime[Context]) -> SharedState:
-        logger.debug(f"Calling agent {agent.name} with summarized messages")
-
-        # Get the summarized messages from the summarization node
-        messages: Sequence[AnyMessage] = state.get("summarized_messages", [])
-        logger.debug(f"Found {len(messages)} summarized messages")
-        logger.trace(f"Summarized messages: {[m.model_dump() for m in messages]}")
-
-        input: dict[str, Any] = {
-            "messages": messages,
-        }
-
-        response: dict[str, Any] = await agent.ainvoke(
-            input=input, context=runtime.context
+    logger.debug(f"Building middleware list for 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")
+    for middleware_config in agent.middleware:
+        middleware = create_factory_middleware(
+            function_name=middleware_config.name,
+            args=middleware_config.args,
+        )
+        if middleware is not None:
+            middleware_list.append(middleware)
+
+    # Add guardrails as middleware
+    if agent.guardrails:
+        logger.debug(f"Adding {len(agent.guardrails)} guardrail middleware")
+    for guardrail in agent.guardrails:
+        # Extract template string from PromptModel if needed
+        prompt_str: str
+        if isinstance(guardrail.prompt, PromptModel):
+            prompt_str = guardrail.prompt.template
+        else:
+            prompt_str = guardrail.prompt
+
+        guardrail_middleware = GuardrailMiddleware(
+            name=guardrail.name,
+            model=guardrail.model.as_chat_model(),
+            prompt=prompt_str,
+            num_retries=guardrail.num_retries or 3,
         )
-        response_messages = response.get("messages", [])
-        logger.debug(f"Agent returned {len(response_messages)} messages")
+        logger.debug(f"Created guardrail middleware: {guardrail.name}")
+        middleware_list.append(guardrail_middleware)
 
-        return {"messages": response_messages}
+    # 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)
+        middleware_list.append(summarization_middleware)
 
-    return call_agent
+    # Add human-in-the-loop middleware if any tools require it
+    hitl_middleware = create_hitl_middleware_from_tool_models(tool_models)
+    if hitl_middleware is not None:
+        logger.debug("Added human-in-the-loop middleware")
+        middleware_list.append(hitl_middleware)
+
+    logger.debug(f"Total middleware count: {len(middleware_list)}")
+    return middleware_list
 
 
 def create_agent_node(
@@ -107,9 +109,9 @@ def create_agent_node(
     """
     Factory function that creates a LangGraph node for a specialized agent.
 
-    This creates a node function that handles user requests using a specialized agent.
-    The function configures the agent with the appropriate model, prompt, tools, and guardrails.
-    If chat_history is provided, it creates a workflow with summarization node.
+    This creates an agent using LangChain v1's create_agent function with
+    middleware for customization. The function configures the agent with
+    the appropriate model, prompt, tools, and middleware.
 
     Args:
         agent: AgentModel configuration for the agent
@@ -122,16 +124,12 @@ def create_agent_node(
     """
     logger.debug(f"Creating agent node for {agent.name}")
 
-    if agent.create_agent_hook:
-        agent_hook = next(iter(create_hooks(agent.create_agent_hook)), None)
-        return agent_hook
-
     llm: LanguageModelLike = agent.model.as_chat_model()
 
     tool_models: Sequence[ToolModel] = agent.tools
     if not additional_tools:
         additional_tools = []
-    tools: Sequence[BaseTool] = create_tools(tool_models) + additional_tools
+    tools: list[BaseTool] = list(create_tools(tool_models)) + list(additional_tools)
 
     if memory and memory.store:
         namespace: tuple[str, ...] = ("memory",)
@@ -139,100 +137,59 @@ def create_agent_node(
             namespace = namespace + (memory.store.namespace,)
         logger.debug(f"Memory store namespace: {namespace}")
 
+        # Use Databricks-compatible search_memory tool (omits problematic filter field)
         tools += [
             create_manage_memory_tool(namespace=namespace),
             create_search_memory_tool(namespace=namespace),
         ]
 
-    pre_agent_hook: Callable[..., Any] = next(
-        iter(create_hooks(agent.pre_agent_hook)), None
+    # Create middleware list from configuration
+    middleware_list = _create_middleware_list(
+        agent=agent,
+        tool_models=tool_models,
+        chat_history=chat_history,
     )
-    logger.debug(f"pre_agent_hook: {pre_agent_hook}")
 
-    post_agent_hook: Callable[..., Any] = next(
-        iter(create_hooks(agent.post_agent_hook)), None
-    )
-    logger.debug(f"post_agent_hook: {post_agent_hook}")
-
-    checkpointer: bool = memory and memory.checkpointer is not None
-
-    compiled_agent: CompiledStateGraph = create_react_agent(
+    logger.debug(f"Created {len(middleware_list)} middleware for 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)
+
+    # Add prompt middleware at the beginning for priority
+    if prompt_middleware is not None:
+        middleware_list.insert(0, prompt_middleware)
+
+    # Configure structured output if response_format is specified
+    response_format: Any = None
+    if agent.response_format is not None:
+        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__}"
+                )
+        except ValueError as e:
+            logger.error(
+                f"Failed to configure structured output for agent {agent.name}: {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)
+    compiled_agent: CompiledStateGraph = create_agent(
         name=agent.name,
         model=llm,
-        prompt=make_prompt(agent.prompt),
         tools=tools,
-        store=True,
+        middleware=middleware_list,
         checkpointer=checkpointer,
-        state_schema=SharedState,
+        state_schema=AgentState,
         context_schema=Context,
-        pre_model_hook=pre_agent_hook,
-        post_model_hook=post_agent_hook,
+        response_format=response_format,  # Add structured output support
     )
 
-    for guardrail_definition in agent.guardrails:
-        guardrail: CompiledStateGraph = reflection_guardrail(guardrail_definition)
-        compiled_agent = with_guardrails(compiled_agent, guardrail)
-
     compiled_agent.name = agent.name
 
-    agent_node: CompiledStateGraph
-
-    if chat_history is None:
-        logger.debug("No chat history configured, using compiled agent directly")
-        agent_node = compiled_agent
-    else:
-        logger.debug("Creating agent node with chat history summarization")
-        workflow: StateGraph = StateGraph(
-            SharedState,
-            config_schema=RunnableConfig,
-            input=SharedState,
-            output=SharedState,
-        )
-        workflow.add_node("summarization", summarization_node(chat_history))
-        workflow.add_node(
-            "agent",
-            call_agent_with_summarized_messages(agent=compiled_agent),
-        )
-        workflow.add_edge("summarization", "agent")
-        workflow.set_entry_point("summarization")
-        agent_node = workflow.compile(name=agent.name)
-
-    return agent_node
-
-
-def message_hook_node(config: AppConfig) -> RunnableLike:
-    message_hooks: Sequence[Callable[..., Any]] = create_hooks(config.app.message_hooks)
-
-    @mlflow.trace()
-    async def message_hook(
-        state: IncomingState, runtime: Runtime[Context]
-    ) -> SharedState:
-        logger.debug("Running message validation")
-        response: dict[str, Any] = {"is_valid": True, "message_error": None}
-
-        for message_hook in message_hooks:
-            message_hook: FunctionHook
-            if message_hook:
-                try:
-                    hook_response: dict[str, Any] = message_hook(
-                        state=state,
-                        runtime=runtime,
-                    )
-                    response.update(hook_response)
-                    logger.debug(f"Hook response: {hook_response}")
-                    if not response.get("is_valid", True):
-                        break
-                except Exception as e:
-                    logger.error(f"Message validation failed: {e}")
-                    response_messages: Sequence[BaseMessage] = [
-                        AIMessage(content=str(e))
-                    ]
-                    return {
-                        "is_valid": False,
-                        "message_error": str(e),
-                        "messages": response_messages,
-                    }
-
-        return response
-
-    return message_hook
+    return compiled_agent