hexdag 0.5.0.dev1__py3-none-any.whl

hexdag/builtin/nodes/agent_node.py

@@ -0,0 +1,696 @@
+"""ReActAgentNode - Multi-step reasoning agent."""
+
+import ast
+import json
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, NotRequired, TypedDict
+
+from pydantic import BaseModel, ConfigDict
+
+from hexdag.builtin.adapters.unified_tool_router import UnifiedToolRouter
+from hexdag.core.context import get_port, get_ports
+from hexdag.core.domain.dag import NodeSpec
+from hexdag.core.logging import get_logger
+from hexdag.core.orchestration.prompt import PromptInput
+from hexdag.core.orchestration.prompt.template import PromptTemplate
+from hexdag.core.ports.tool_router import ToolRouter
+from hexdag.core.protocols import to_dict
+
+from .base_node_factory import BaseNodeFactory
+from .llm_node import LLMNode
+from .tool_utils import ToolCallFormat, ToolParser
+
+logger = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from types import MappingProxyType
+
+
+class PhaseContext(TypedDict):
+    """Context structure for phase transitions in agents.
+
+    Attributes
+    ----------
+    previous_phase : str, optional
+        The phase the agent is transitioning from
+    reason : str, optional
+        Explanation for why the phase change is occurring
+    carried_data : dict[str, Any], optional
+        Data to carry forward from the previous phase
+    target_output : str, optional
+        Expected output format or goal for the new phase
+    iteration : int, optional
+        Current iteration number if in a loop or retry scenario
+    metadata : dict[str, Any], optional
+        Additional metadata about the phase transition
+    """
+
+    previous_phase: NotRequired[str]
+    reason: NotRequired[str]
+    carried_data: NotRequired[dict[str, Any]]
+    target_output: NotRequired[str]
+    iteration: NotRequired[int]
+    metadata: NotRequired[dict[str, Any]]
+
+
+class AgentState(BaseModel):
+    """Pydantic model for agent state - provides type safety and validation."""
+
+    # Original input data (preserved)
+    input_data: dict[str, Any] = {}
+
+    # Agent reasoning state
+    reasoning_steps: list[str] = []
+    tool_results: list[str] = []
+    tools_used: list[str] = []
+    current_phase: str = "main"
+    phase_history: list[str] = ["main"]
+    phase_contexts: dict[str, PhaseContext] = {}  # Store typed context for each phase
+    step: int = 0
+    response: str = ""
+
+    # Loop iteration tracking
+    loop_iteration: int = 0
+
+    model_config = ConfigDict(extra="allow")  # Allow additional fields from input mapping
+
+
+@dataclass(frozen=True, slots=True)
+class AgentConfig:
+    """Agent configuration for multi-step reasoning (legacy - kept for backward compatibility)."""
+
+    max_steps: int = 20
+    tool_call_style: ToolCallFormat = ToolCallFormat.MIXED
+
+
+class Agent:
+    """Configuration for Agent Node.
+
+    Attributes
+    ----------
+    max_steps : int
+        Maximum number of reasoning steps (default: 20)
+    tool_call_style : ToolCallFormat
+        Format for tool calls - MIXED, FUNCTION_CALL, or JSON (default: MIXED)
+    """
+
+    max_steps: int = 20
+    tool_call_style: ToolCallFormat = ToolCallFormat.MIXED
+
+
+class ReActAgentNode(BaseNodeFactory):
+    """Multi-step reasoning agent.
+
+    This agent:
+    1. Uses loop control internally for iteration control
+    2. Implements single-step reasoning logic
+    3. Maintains clean agent interface for users
+    4. Leverages proven loop control patterns
+    5. Supports all agent features (tools, phases, events)
+
+    Architecture:
+    ```
+    Agent(input) -> Loop -> SingleStep -> Loop -> SingleStep -> ... -> Output
+    ```
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize with dependencies."""
+        self.llm_node = LLMNode()
+        self.tool_parser = ToolParser()
+
+    def __call__(
+        self,
+        name: str,
+        main_prompt: PromptInput,
+        continuation_prompts: dict[str, PromptInput] | None = None,
+        output_schema: dict[str, type] | type[BaseModel] | None = None,
+        config: AgentConfig | None = None,
+        deps: list[str] | None = None,
+        **kwargs: Any,
+    ) -> NodeSpec:
+        """Create a multi-step reasoning agent with internal loop control.
+
+        Args
+        ----
+            name: Agent name
+            main_prompt: Initial reasoning prompt
+            continuation_prompts: Phase-specific prompts
+            output_schema: Custom output schema for tool_end results
+            config: Agent configuration
+            deps: Dependencies
+            **kwargs: Additional parameters
+
+        Returns
+        -------
+        NodeSpec
+            A configured node specification for the agent
+        """
+        config = config or AgentConfig()
+
+        # Infer input schema from prompt
+        input_schema = self._infer_input_schema(main_prompt)
+
+        input_model = self.create_pydantic_model(f"{name}Input", input_schema)
+        if input_model is None:
+            input_model = type(f"{name}Input", (BaseModel,), {"__annotations__": {"input": str}})
+        output_model = self.create_pydantic_model(f"{name}Output", output_schema) or type(
+            f"{name}Output", (BaseModel,), {"__annotations__": {"output": str}}
+        )
+
+        agent_fn = self._create_agent_with_loop(
+            name, main_prompt, continuation_prompts or {}, output_model, config
+        )
+
+        # Use universal input mapping method
+        return self.create_node_with_mapping(
+            name=name,
+            wrapped_fn=agent_fn,
+            input_schema=input_schema,
+            output_schema=output_model,
+            deps=deps,
+            **kwargs,
+        )
+
+    def _infer_input_schema(self, prompt: PromptInput) -> dict[str, Any]:
+        """Infer input schema from prompt template.
+
+        Returns
+        -------
+        dict[str, Any]
+            Inferred input schema mapping
+        """
+        # Use the shared implementation from BaseNodeFactory
+        # AgentNode doesn't filter special params, so pass None
+        return BaseNodeFactory.infer_input_schema_from_template(prompt, special_params=None)
+
+    def _get_current_prompt(
+        self,
+        main_prompt: PromptInput,
+        continuation_prompts: dict[str, PromptInput],
+        current_phase: str,
+    ) -> PromptInput:
+        """Get the appropriate prompt for the current phase.
+
+        Returns
+        -------
+        PromptInput
+            The prompt to use for the current phase
+        """
+        if current_phase != "main" and current_phase in continuation_prompts:
+            return continuation_prompts[current_phase]
+        return main_prompt
+
+    def _create_agent_with_loop(
+        self,
+        name: str,
+        main_prompt: PromptInput,
+        continuation_prompts: dict[str, PromptInput],
+        output_model: type[BaseModel],
+        config: AgentConfig,
+    ) -> Callable[..., Any]:
+        """Create agent function with internal loop composition for multi-step iteration.
+
+        Returns
+        -------
+        Callable[..., Any]
+            Agent function with internal loop control
+        """
+
+        async def single_step_executor(input_data: Any) -> Any:
+            """Execute single reasoning step."""
+            from hexdag.core.context import get_port
+
+            ports: MappingProxyType[str, Any] | dict[Any, Any] = get_ports() or {}
+
+            state = self._initialize_or_update_state(input_data)
+
+            # Execute single reasoning step
+            updated_state = await self._execute_single_step(
+                state, name, main_prompt, continuation_prompts, config, dict(ports)
+            )
+
+            final_output = await self._check_for_final_output(
+                updated_state, output_model, get_port("event_manager")
+            )
+            if final_output is not None:
+                return final_output
+
+            # Return AgentState directly (Pydantic-first design)
+            return updated_state
+
+        # Define success condition using loop concepts
+        def success_condition(result: Any) -> bool:
+            """Check if agent should stop iterating."""
+            # Stop if we got the final structured output (not AgentState)
+            if not isinstance(result, AgentState):
+                return True
+
+            # Stop if we reached max steps
+            if result.step >= config.max_steps:
+                return True
+
+            # Stop if tool_end was detected
+            return "tool_end" in result.response.lower()
+
+        async def agent_with_internal_loop(input_data: Any) -> Any:
+            """Agent executor with internal loop control."""
+            node_logger = logger.bind(node=name, node_type="agent_node")
+            start_time = time.perf_counter()
+
+            # Log agent start
+            node_logger.info(
+                "Starting agent",
+                max_steps=config.max_steps,
+                tool_call_style=config.tool_call_style.value,
+            )
+
+            # Start with initial input
+            current_result = input_data
+
+            # Run the loop until success condition is met or max iterations reached
+            for step_num in range(config.max_steps):
+                # Log step start
+                node_logger.debug(
+                    "Agent step starting",
+                    step=step_num + 1,
+                    max_steps=config.max_steps,
+                )
+
+                # Execute single step
+                step_result = await single_step_executor(current_result)
+
+                # If not AgentState, it's the final output
+                if not isinstance(step_result, AgentState):
+                    duration_ms = (time.perf_counter() - start_time) * 1000
+                    node_logger.info(
+                        "Agent completed with direct output",
+                        total_steps=step_num + 1,
+                        duration_ms=f"{duration_ms:.2f}",
+                        output_type=type(step_result).__name__,
+                    )
+                    return step_result
+
+                # Log step completion with state info
+                node_logger.debug(
+                    "Agent step completed",
+                    step=step_num + 1,
+                    phase=step_result.current_phase,
+                    tools_used_count=len(step_result.tools_used),
+                )
+
+                # Check success condition
+                if success_condition(step_result):
+                    final_output = await self._check_for_final_output(
+                        step_result,
+                        output_model,
+                        get_port("event_manager"),
+                    )
+                    if final_output is not None:
+                        duration_ms = (time.perf_counter() - start_time) * 1000
+                        node_logger.info(
+                            "Agent completed",
+                            total_steps=step_num + 1,
+                            tools_used=step_result.tools_used,
+                            phases=step_result.phase_history,
+                            duration_ms=f"{duration_ms:.2f}",
+                        )
+                        return final_output
+                    return step_result
+
+                # Continue with next iteration (pass AgentState directly)
+                current_result = step_result
+
+            # If we reach here, max steps reached
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            node_logger.warning(
+                "Agent reached max steps",
+                max_steps=config.max_steps,
+                duration_ms=f"{duration_ms:.2f}",
+            )
+            return current_result
+
+        return agent_with_internal_loop
+
+    def _initialize_or_update_state(self, input_data: Any) -> AgentState:
+        """Initialize new state or update existing state from loop iteration.
+
+        Returns
+        -------
+        AgentState
+            Initialized or updated agent state
+        """
+        # Case 1: Already AgentState (from previous iteration) - return as-is
+        if isinstance(input_data, AgentState):
+            return input_data
+
+        # Case 2: Dict with AgentState fields (legacy/backward compatibility)
+        if isinstance(input_data, dict) and "reasoning_steps" in input_data:
+            state: AgentState = AgentState.model_validate(input_data)
+            return state
+
+        # Case 3: Fresh input (first iteration) - wrap in AgentState
+        try:
+            raw_input = to_dict(input_data)
+        except TypeError:
+            # Fallback for non-dict types
+            raw_input = {"input": str(input_data)}
+
+        return AgentState(input_data=raw_input)
+
+    def _enhance_prompt_with_tools(
+        self, prompt: PromptInput, tool_router: ToolRouter | None, config: AgentConfig
+    ) -> PromptInput:
+        """Add tool instructions to the prompt.
+
+        Returns
+        -------
+        PromptInput
+            Enhanced prompt with tool instructions
+        """
+        if not tool_router:
+            return prompt
+
+        if isinstance(prompt, str):
+            prompt = PromptTemplate(prompt)
+
+        tool_instructions = self._build_tool_instructions(tool_router, config)
+
+        # Use the template's enhance method
+        return prompt + tool_instructions
+
+    def _build_tool_instructions(self, tool_router: ToolRouter, config: AgentConfig) -> str:
+        """Build tool usage instructions based on the configured format.
+
+        Returns
+        -------
+        str
+            Tool usage instructions text
+        """
+        tool_schemas = tool_router.get_all_tool_schemas()
+        if not tool_schemas:
+            return "\n## No tools available"
+
+        tool_list = []
+        for name, schema in tool_schemas.items():
+            params = ", ".join(p["name"] for p in schema.get("parameters", []))
+            tool_list.append(f"- {name}({params}): {schema.get('description', 'No description')}")
+
+        tools_text = "\n".join(tool_list)
+
+        # Generate format-specific usage guidelines
+        usage_guidelines = self._get_format_specific_guidelines(config.tool_call_style)
+
+        return f"""
+## Available Tools
+{tools_text}
+
+## Usage Guidelines
+{usage_guidelines}
+"""
+
+    def _get_format_specific_guidelines(self, format_style: ToolCallFormat) -> str:
+        """Generate format-specific tool calling guidelines.
+
+        Returns
+        -------
+        str
+            Format-specific guidelines text
+        """
+        if format_style == ToolCallFormat.FUNCTION_CALL:
+            return """- Call ONE tool at a time: INVOKE_TOOL: tool_name(param='value')
+- For final answer and structured output: INVOKE_TOOL: tool_end(field1='value1', field2='value2')
+- For phase change: INVOKE_TOOL: change_phase(phase='new_phase', reason='why changing',
+carried_data={'key': 'value'})"""
+
+        if format_style == ToolCallFormat.JSON:
+            return (
+                """- Call ONE tool at a time: INVOKE_TOOL: """
+                """{"tool": "tool_name", "params": {"param": "value"}}\n"""
+                """- For final answer and structured output: INVOKE_TOOL: """
+                """{"tool": "tool_end", "params": {"field1": "value1", "field2": "value2"}}\n"""
+                """- For phase change: INVOKE_TOOL: """
+                """{"tool": "change_phase", "params": {"phase": "new_phase", "reason": "why",
+                "carried_data": {"key": "val"}}}"""
+            )
+
+        # ToolCallFormat.MIXED
+        return """- Call ONE tool at a time using either format:
+  - Function style: INVOKE_TOOL: tool_name(param='value')
+  - JSON style: INVOKE_TOOL: {"tool": "tool_name", "params": {"param": "value"}}
+- For final answer and structured output:
+  - Function: INVOKE_TOOL: tool_end(field1='value1', field2='value2')
+  - JSON: INVOKE_TOOL: {"tool": "tool_end", "params": {"field1": "value1", "field2": "value2"}}
+- For phase change:
+  - Function: INVOKE_TOOL: change_phase(phase='new_phase', reason='why',
+  carried_data={'key': 'val'})
+  - JSON: INVOKE_TOOL: {"tool": "change_phase", "params": {"phase": "new_phase",
+  "reason": "why", "carried_data": {"key": "val"}}}"""
+
+    async def _get_llm_response(
+        self, prompt: PromptInput, llm_input: dict[str, Any], ports: dict[str, Any], node_name: str
+    ) -> str:
+        """Get response from LLM.
+
+        Returns
+        -------
+        str
+            LLM response text
+        """
+        # Ensure we have a proper template (not string)
+        if isinstance(prompt, str):
+            prompt = PromptTemplate(prompt)
+
+        llm_node_spec = self.llm_node.from_template(node_name, template=prompt)
+
+        # Execute LLM with the prepared input (no ports passed - uses ExecutionContext)
+        return await llm_node_spec.fn(llm_input)  # type: ignore[no-any-return]
+
+    async def _execute_single_step(
+        self,
+        state: AgentState,
+        name: str,
+        main_prompt: PromptInput,
+        continuation_prompts: dict[str, PromptInput],
+        config: AgentConfig,
+        ports: dict[str, Any],
+    ) -> AgentState:
+        """Execute a single reasoning step.
+
+        Returns
+        -------
+        AgentState
+            Updated agent state after step execution
+        """
+        event_manager = ports.get("event_manager")
+        tool_router = ports.get("tool_router", UnifiedToolRouter())
+
+        current_step = max(state.loop_iteration, state.step) + 1
+        node_step_name = f"{name}_step_{current_step}"
+
+        current_prompt = self._get_current_prompt(
+            main_prompt, continuation_prompts, state.current_phase
+        )
+
+        # Enhance prompt with tools
+        enhanced_prompt = self._enhance_prompt_with_tools(current_prompt, tool_router, config)
+
+        current_phase_context = state.phase_contexts.get(state.current_phase, {})
+
+        # Build LLM input - only convert to dict when needed for template
+        # Merge state fields with template-specific overrides
+        llm_input = {
+            **state.model_dump(),  # Convert only once, at template boundary
+            **state.input_data,
+            "reasoning_so_far": "\n".join(state.reasoning_steps) or "Starting reasoning...",
+            "phase_context": current_phase_context,
+            "phase_reason": current_phase_context.get("reason", ""),
+            "phase_target": current_phase_context.get("target_output", ""),
+        }
+
+        response = await self._get_llm_response(enhanced_prompt, llm_input, ports, node_step_name)
+
+        # Process tools and phase changes
+        await self._process_tools_and_phases(
+            response, state, tool_router, continuation_prompts, config, event_manager
+        )
+
+        state.reasoning_steps.append(f"Step {current_step}: {response}")
+        state.response = response
+        state.step = current_step
+
+        return state
+
+    def _should_terminate(self, response: str) -> bool:
+        """Check if agent should terminate.
+
+        Returns
+        -------
+        bool
+            True if agent should terminate execution
+        """
+        return "tool_end" in response.lower() or "Tool_END" in response
+
+    async def _process_tools_and_phases(
+        self,
+        response: str,
+        state: AgentState,
+        tool_router: ToolRouter | None,
+        continuation_prompts: dict[str, PromptInput],
+        config: AgentConfig,
+        event_manager: Any,
+    ) -> None:
+        """Process tool calls and phase changes."""
+        if not tool_router:
+            return
+
+        # Parse tool calls
+        tool_calls = self.tool_parser.parse_tool_calls(response, format=config.tool_call_style)
+
+        if tool_calls:
+            logger.debug(
+                "Parsed tool calls",
+                tool_count=len(tool_calls),
+                tools=[tc.name for tc in tool_calls],
+            )
+
+        for tool_call in tool_calls:
+            try:
+                # Log tool execution
+                logger.debug(
+                    "Executing tool",
+                    tool_name=tool_call.name,
+                    params_preview=str(tool_call.params)[:100],
+                )
+
+                # Execute tool
+                result = await tool_router.acall_tool(tool_call.name, tool_call.params)
+
+                state.tool_results.append(f"{tool_call.name}: {result}")
+                state.tools_used.append(tool_call.name)
+
+                if tool_call.name == "change_phase" and isinstance(result, dict):
+                    new_phase = result.get("new_phase")
+                    context = result.get("context", {})
+
+                    if new_phase and new_phase in continuation_prompts:
+                        old_phase = state.current_phase
+
+                        if "previous_phase" not in context:
+                            context["previous_phase"] = state.current_phase
+
+                        state.phase_contexts[new_phase] = context
+
+                        state.current_phase = new_phase
+                        state.phase_history.append(new_phase)
+
+                        # Log phase transition
+                        logger.info(
+                            "Phase transition",
+                            from_phase=old_phase,
+                            to_phase=new_phase,
+                            reason=context.get("reason", ""),
+                        )
+
+                        # If there's carried_data, merge it into state.input_data
+                        if "carried_data" in context and isinstance(context["carried_data"], dict):
+                            state.input_data.update(context["carried_data"])
+
+            except Exception as e:
+                error_msg = f"{tool_call.name}: Error - {e}"
+                state.tool_results.append(error_msg)
+                logger.warning(
+                    "Tool execution failed",
+                    tool_name=tool_call.name,
+                    error=str(e),
+                )
+
+    def _parse_tool_end_result(self, tool_result: str) -> dict[str, Any] | None:
+        """Parse tool_end result string into structured data.
+
+        Args
+        ----
+        tool_result : str
+            Tool result string in format "tool_end: {data}"
+
+        Returns
+        -------
+        dict[str, Any] | None
+            Parsed data dictionary or None if parsing fails
+        """
+        if not tool_result or not tool_result.startswith("tool_end:"):
+            return None
+
+        try:
+            result_str = tool_result.split(":", 1)[1].strip()
+            result_data = ast.literal_eval(result_str)
+
+            if isinstance(result_data, dict):
+                return result_data
+            return None
+        except (json.JSONDecodeError, SyntaxError, ValueError, IndexError):
+            # Failed to parse - return None to skip this result
+            # IndexError: split failed (malformed tool_end output)
+            return None
+
+    async def _emit_agent_metadata(self, state: AgentState, event_manager: Any) -> None:
+        """Emit agent metadata trace event.
+
+        Args
+        ----
+        state : AgentState
+            Current agent state
+        event_manager : Any
+            Event manager instance
+        """
+        if event_manager and hasattr(event_manager, "add_trace"):
+            await event_manager.add_trace(
+                "agent_metadata",
+                {
+                    "reasoning_steps": state.reasoning_steps,
+                    "tools_used": list(set(state.tools_used)),
+                    "reasoning_phases": state.phase_history,
+                    "total_steps": state.step,
+                },
+            )
+
+    async def _check_for_final_output(
+        self,
+        state: AgentState,
+        output_model: type[BaseModel],
+        event_manager: Any,
+    ) -> Any | None:
+        """Check if we have a final output from tool_end calls.
+
+        Returns
+        -------
+        Any | None
+            Final output model instance or None if not found
+        """
+        # Check for tool_end calls with structured output
+        for tool_result in reversed(state.tool_results):
+            parsed_data = self._parse_tool_end_result(tool_result)
+
+            if parsed_data is not None:
+                try:
+                    # Emit metadata before returning final result
+                    await self._emit_agent_metadata(state, event_manager)
+
+                    return output_model.model_validate(parsed_data)
+
+                except (ValueError, TypeError) as e:
+                    # Validation failed - try next tool_end result
+                    logger.debug(
+                        "Failed to validate tool_end result",
+                        output_model=output_model.__name__,
+                        error=str(e),
+                    )
+                    continue  # Skip this tool result and try the next one
+
+        return None
+
+
+# Backward compatibility alias
+ReasoningAgentNode = ReActAgentNode