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

dao_ai/nodes.py

@@ -1,229 +1,324 @@
-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,
-    AppModel,
     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 (
+    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 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(app_model: AppModel) -> RunnableLike:
-    chat_history: ChatHistoryModel | None = app_model.chat_history
-    if chat_history is None:
-        raise ValueError(
-            "AppModel must have chat_history configured to use summarization"
-        )
+def _create_middleware_list(
+    agent: AgentModel,
+    tool_models: Sequence[ToolModel],
+    chat_history: Optional[ChatHistoryModel] = None,
+) -> list[Any]:
+    """
+    Create a list of middleware instances from agent configuration.
 
-    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
-    )
+    Args:
+        agent: AgentModel configuration
+        tool_models: Tool model configurations (for HITL settings)
+        chat_history: Optional chat history configuration for summarization
 
-    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}"
-    )
+    Returns:
+        List of middleware instances (can include both AgentMiddleware and
+        LangChain built-in middleware)
+    """
+    logger.debug("Building middleware list for agent", agent=agent.name)
+    middleware_list: list[Any] = []
+
+    # Add configured middleware using factory pattern
+    if agent.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:
+        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,
+        )
+        if middleware is not None:
+            middleware_list.append(middleware)
+
+    # Add guardrails as middleware
+    if agent.guardrails:
+        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
+        if isinstance(guardrail.prompt, PromptModel):
+            prompt_str = guardrail.prompt.template
+        else:
+            prompt_str = guardrail.prompt
+
+        guardrail_middleware: GuardrailMiddleware = GuardrailMiddleware(
+            name=guardrail.name,
+            model=guardrail.model.as_chat_model(),
+            prompt=prompt_str,
+            num_retries=guardrail.num_retries or 3,
+        )
+        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.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)
 
-    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",
+    # Add human-in-the-loop middleware if any tools require it
+    hitl_middleware: HumanInTheLoopMiddleware | None = (
+        create_hitl_middleware_from_tool_models(tool_models)
     )
-    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
+    if hitl_middleware is not None:
+        # 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,
         )
-        response_messages = response.get("messages", [])
-        logger.debug(f"Agent returned {len(response_messages)} messages")
-
-        return {"messages": response_messages}
+        middleware_list.append(hitl_middleware)
 
-    return call_agent
+    logger.info(
+        "Middleware summary",
+        agent=agent.name,
+        total_middleware_count=len(middleware_list),
+    )
+    return middleware_list
 
 
 def create_agent_node(
-    app: AppModel,
     agent: AgentModel,
+    memory: Optional[MemoryModel] = None,
+    chat_history: Optional[ChatHistoryModel] = None,
     additional_tools: Optional[Sequence[BaseTool]] = None,
 ) -> RunnableLike:
     """
     Factory function that creates a LangGraph node for a specialized agent.
 
-    This creates a node function that handles user requests using a specialized agent
-    based on the provided agent_type. The function configures the agent with the
-    appropriate model, prompt, tools, and guardrails from the model_config.
+    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:
-        model_config: Configuration containing models, prompts, tools, and guardrails
-        agent_type: Type of agent to create (e.g., "general", "product", "inventory")
+        agent: AgentModel configuration for the agent
+        memory: Optional MemoryModel for memory store configuration
+        chat_history: Optional ChatHistoryModel for chat history summarization
+        additional_tools: Optional sequence of additional tools to add to the agent
 
     Returns:
-        An agent callable function that processes state and returns responses
+        RunnableLike: An agent node that processes state and returns responses
     """
-    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
+    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 = []
-    tools: Sequence[BaseTool] = create_tools(tool_models) + additional_tools
 
-    if app.orchestration.memory and app.orchestration.memory.store:
+    # 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 app.orchestration.memory.store.namespace:
-            namespace = namespace + (app.orchestration.memory.store.namespace,)
-        logger.debug(f"Memory store namespace: {namespace}")
+        if memory.store.namespace:
+            namespace = namespace + (memory.store.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"],
+        )
 
-    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
+    # 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
+
+    # 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.info(
+                    "Response format configuration",
+                    agent=agent.name,
+                    format_type=type(response_format).__name__,
+                    structured_output=True,
+                )
+        except ValueError as e:
+            logger.error(
+                "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,
     )
-    logger.debug(f"post_agent_hook: {post_agent_hook}")
 
-    compiled_agent: CompiledStateGraph = create_react_agent(
+    compiled_agent: CompiledStateGraph = create_agent(
         name=agent.name,
         model=llm,
-        prompt=make_prompt(agent.prompt),
         tools=tools,
-        store=True,
-        checkpointer=True,
-        state_schema=SharedState,
+        middleware=middleware_list,
+        checkpointer=checkpointer,
+        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
-
-    chat_history: ChatHistoryModel = app.chat_history
+    logger.info("Agent node created successfully", agent=agent.name)
 
-    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(app))
-        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