nvidia-nat 1.3.dev0__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/agent/rewoo_agent/agent.py

@@ -13,20 +13,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import json
-# pylint: disable=R0917
 import logging
+import re
 from json import JSONDecodeError
+from typing import Any
 
 from langchain_core.callbacks.base import AsyncCallbackHandler
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages.ai import AIMessage
+from langchain_core.messages.base import BaseMessage
 from langchain_core.messages.human import HumanMessage
 from langchain_core.messages.tool import ToolMessage
 from langchain_core.prompts.chat import ChatPromptTemplate
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
 from pydantic import BaseModel
 from pydantic import Field
 
@@ -41,22 +45,40 @@ from nat.agent.base import BaseAgent
 logger = logging.getLogger(__name__)
 
 
+class ReWOOEvidence(BaseModel):
+    placeholder: str
+    tool: str
+    tool_input: Any
+
+
+class ReWOOPlanStep(BaseModel):
+    plan: str
+    evidence: ReWOOEvidence
+
+
 class ReWOOGraphState(BaseModel):
     """State schema for the ReWOO Agent Graph"""
+    messages: list[BaseMessage] = Field(default_factory=list)  # input and output of the ReWOO Agent
     task: HumanMessage = Field(default_factory=lambda: HumanMessage(content=""))  # the task provided by user
     plan: AIMessage = Field(
         default_factory=lambda: AIMessage(content=""))  # the plan generated by the planner to solve the task
     steps: AIMessage = Field(
         default_factory=lambda: AIMessage(content=""))  # the steps to solve the task, parsed from the plan
+    # New fields for parallel execution support
+    evidence_map: dict[str, ReWOOPlanStep] = Field(default_factory=dict)  # mapping from placeholders to step info
+    execution_levels: list[list[str]] = Field(default_factory=list)  # levels for parallel execution
+    current_level: int = Field(default=0)  # current execution level
     intermediate_results: dict[str, ToolMessage] = Field(default_factory=dict)  # the intermediate results of each step
     result: AIMessage = Field(
         default_factory=lambda: AIMessage(content=""))  # the final result of the task, generated by the solver
 
 
 class ReWOOAgentGraph(BaseAgent):
-    """Configurable LangGraph ReWOO Agent. A ReWOO Agent performs reasoning by interacting with other objects or tools
-    and utilizes their outputs to make decisions. Supports retrying on output parsing errors. Argument
-    "detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
+    """Configurable ReWOO Agent.
+
+    Args:
+        detailed_logs: Toggles logging of inputs, outputs, and intermediate steps.
+    """
 
     def __init__(self,
                  llm: BaseChatModel,
@@ -65,28 +87,34 @@ class ReWOOAgentGraph(BaseAgent):
                  tools: list[BaseTool],
                  use_tool_schema: bool = True,
                  callbacks: list[AsyncCallbackHandler] | None = None,
-                 detailed_logs: bool = False):
-        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+                 detailed_logs: bool = False,
+                 log_response_max_chars: int = 1000,
+                 tool_call_max_retries: int = 3,
+                 raise_tool_call_error: bool = True):
+        super().__init__(llm=llm,
+                         tools=tools,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
 
         logger.debug(
             "%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
             AGENT_LOG_PREFIX)
-        tool_names = ",".join([tool.name for tool in tools[:-1]]) + ',' + tools[-1].name  # prevent trailing ","
-        if not use_tool_schema:
-            tool_names_and_descriptions = "\n".join(
-                [f"{tool.name}: {tool.description}"
-                 for tool in tools[:-1]]) + "\n" + f"{tools[-1].name}: {tools[-1].description}"  # prevent trailing "\n"
-        else:
-            logger.debug("%s Adding the tools' input schema to the tools' description", AGENT_LOG_PREFIX)
-            tool_names_and_descriptions = "\n".join([
-                f"{tool.name}: {tool.description}. {INPUT_SCHEMA_MESSAGE.format(schema=tool.input_schema.model_fields)}"
-                for tool in tools[:-1]
-            ]) + "\n" + (f"{tools[-1].name}: {tools[-1].description}. "
-                         f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
+
+        def describe_tool(tool: BaseTool) -> str:
+            description = f"{tool.name}: {tool.description}"
+            if use_tool_schema:
+                description += f". {INPUT_SCHEMA_MESSAGE.format(schema=tool.input_schema.model_fields)}"
+            return description
+
+        tool_names = ",".join(tool.name for tool in tools)
+        tool_names_and_descriptions = "\n".join(describe_tool(tool) for tool in tools)
 
         self.planner_prompt = planner_prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
         self.solver_prompt = solver_prompt
         self.tools_dict = {tool.name: tool for tool in tools}
+        self.tool_call_max_retries = tool_call_max_retries
+        self.raise_tool_call_error = raise_tool_call_error
 
         logger.debug("%s Initialized ReWOO Agent Graph", AGENT_LOG_PREFIX)
 
@@ -94,34 +122,91 @@ class ReWOOAgentGraph(BaseAgent):
         try:
             return self.tools_dict.get(tool_name)
         except Exception as ex:
-            logger.exception("%s Unable to find tool with the name %s\n%s",
-                             AGENT_LOG_PREFIX,
-                             tool_name,
-                             ex,
-                             exc_info=True)
-            raise ex
+            logger.error("%s Unable to find tool with the name %s\n%s", AGENT_LOG_PREFIX, tool_name, ex)
+            raise
 
     @staticmethod
-    def _get_current_step(state: ReWOOGraphState) -> int:
-        steps = state.steps.content
-        if len(steps) == 0:
-            raise RuntimeError('No steps received in ReWOOGraphState')
+    def _get_current_level_status(state: ReWOOGraphState) -> tuple[int, bool]:
+        """
+        Get the current execution level and whether it's complete.
+
+        Args:
+            state: The ReWOO graph state.
 
-        if len(state.intermediate_results) == len(steps):
-            # all steps are done
-            return -1
+        Returns:
+            tuple of (current_level, is_complete). Level -1 means all execution is complete.
+        """
+        if not state.execution_levels:
+            return -1, True
 
-        return len(state.intermediate_results)
+        current_level = state.current_level
+
+        # Check if we've completed all levels
+        if current_level >= len(state.execution_levels):
+            return -1, True
+
+        # Check if current level is complete
+        current_level_placeholders = state.execution_levels[current_level]
+        level_complete = all(placeholder in state.intermediate_results for placeholder in current_level_placeholders)
+
+        return current_level, level_complete
 
     @staticmethod
-    def _parse_planner_output(planner_output: str) -> AIMessage:
+    def _parse_planner_output(planner_output: str) -> list[ReWOOPlanStep]:
 
         try:
-            steps = json.loads(planner_output)
-        except json.JSONDecodeError as ex:
+            return [ReWOOPlanStep(**step) for step in json.loads(planner_output)]
+        except Exception as ex:
             raise ValueError(f"The output of planner is invalid JSON format: {planner_output}") from ex
 
-        return AIMessage(content=steps)
+    @staticmethod
+    def _parse_planner_dependencies(steps: list[ReWOOPlanStep]) -> tuple[dict[str, ReWOOPlanStep], list[list[str]]]:
+        """
+        Parse planner steps to identify dependencies and create execution levels for parallel processing.
+        This creates a dependency map and identifies which evidence placeholders can be executed in parallel.
+
+        Args:
+            steps: list of plan steps from the planner.
+
+        Returns:
+            A mapping from evidence placeholders to step info and execution levels for parallel processing.
+        """
+        # First pass: collect all evidence placeholders and their info
+        evidences: dict[str, ReWOOPlanStep] = {
+            step.evidence.placeholder: step
+            for step in steps if step.evidence and step.evidence.placeholder
+        }
+
+        # Second pass: find dependencies now that we have all placeholders
+        dependencies = {
+            step.evidence.placeholder: [
+                var for var in re.findall(r"#E\d+", str(step.evidence.tool_input))
+                if var in evidences and var != step.evidence.placeholder
+            ]
+            for step in steps if step.evidence and step.evidence.placeholder
+        }
+
+        # Create execution levels using topological sort
+        levels: list[list[str]] = []
+        remaining = dict(dependencies)
+
+        while remaining:
+            # Find items with no dependencies (can be executed in parallel)
+            ready = [placeholder for placeholder, deps in remaining.items() if not deps]
+
+            if not ready:
+                raise ValueError("Circular dependency detected in planner output")
+
+            levels.append(ready)
+
+            # Remove completed items from remaining
+            for placeholder in ready:
+                remaining.pop(placeholder)
+
+            # Remove completed items from other dependencies
+            for ph, deps in list(remaining.items()):
+                remaining[ph] = list(set(deps) - set(ready))
+        return evidences, levels
 
     @staticmethod
     def _replace_placeholder(placeholder: str, tool_input: str | dict, tool_output: str | dict) -> str | dict:
@@ -141,6 +226,7 @@ class ReWOOAgentGraph(BaseAgent):
 
         else:
             assert False, f"Unexpected type for tool_input: {type(tool_input)}"
+
         return tool_input
 
     @staticmethod
@@ -183,132 +269,217 @@ class ReWOOAgentGraph(BaseAgent):
             if not task:
                 logger.error("%s No task provided to the ReWOO Agent. Please provide a valid task.", AGENT_LOG_PREFIX)
                 return {"result": NO_INPUT_ERROR_MESSAGE}
-
+            chat_history = self._get_chat_history(state.messages)
             plan = await self._stream_llm(
                 planner,
-                {"task": task},
+                {
+                    "task": task, "chat_history": chat_history
+                },
                 RunnableConfig(callbacks=self.callbacks)  # type: ignore
             )
 
             steps = self._parse_planner_output(str(plan.content))
 
+            # Parse dependencies and create execution levels for parallel processing
+            evidence_map, execution_levels = self._parse_planner_dependencies(steps)
+
             if self.detailed_logs:
                 agent_response_log_message = AGENT_CALL_LOG_MESSAGE % (task, str(plan.content))
                 logger.info("ReWOO agent planner output: %s", agent_response_log_message)
+                logger.info("ReWOO agent execution levels: %s", execution_levels)
 
-            return {"plan": plan, "steps": steps}
+            return {
+                "plan": plan,
+                "evidence_map": evidence_map,
+                "execution_levels": execution_levels,
+                "current_level": 0,
+            }
 
         except Exception as ex:
-            logger.exception("%s Failed to call planner_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
-            raise ex
+            logger.error("%s Failed to call planner_node: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
     async def executor_node(self, state: ReWOOGraphState):
+        """
+        Execute tools in parallel for the current dependency level.
+
+        This replaces the sequential execution with parallel execution of tools
+        that have no dependencies between them.
+        """
         try:
             logger.debug("%s Starting the ReWOO Executor Node", AGENT_LOG_PREFIX)
 
-            current_step = self._get_current_step(state)
-            # The executor node should not be invoked after all steps are finished
-            if current_step < 0:
-                logger.error("%s ReWOO Executor is invoked with an invalid step number: %s",
-                             AGENT_LOG_PREFIX,
-                             current_step)
-                raise RuntimeError(f"ReWOO Executor is invoked with an invalid step number: {current_step}")
-
-            steps_content = state.steps.content
-            if isinstance(steps_content, list) and current_step < len(steps_content):
-                step = steps_content[current_step]
-                if isinstance(step, dict) and "evidence" in step:
-                    step_info = step["evidence"]
-                    placeholder = step_info.get("placeholder", "")
-                    tool = step_info.get("tool", "")
-                    tool_input = step_info.get("tool_input", "")
+            current_level, level_complete = self._get_current_level_status(state)
+
+            # Should not be invoked if all levels are complete
+            if current_level < 0:
+                logger.error("%s ReWOO Executor invoked after all levels complete", AGENT_LOG_PREFIX)
+                raise RuntimeError("ReWOO Executor invoked after all levels complete")
+
+            # If current level is already complete, move to next level
+            if level_complete:
+                new_level = current_level + 1
+                logger.debug("%s Level %s complete, moving to level %s", AGENT_LOG_PREFIX, current_level, new_level)
+                return {"current_level": new_level}
+
+            # Get placeholders for current level
+            current_level_placeholders = state.execution_levels[current_level]
+
+            # Filter to only placeholders not yet completed
+            pending_placeholders = list(set(current_level_placeholders) - set(state.intermediate_results.keys()))
+
+            if not pending_placeholders:
+                # All placeholders in this level are done, move to next level
+                new_level = current_level + 1
+                return {"current_level": new_level}
+
+            logger.debug("%s Executing level %s with %s tools in parallel: %s",
+                         AGENT_LOG_PREFIX,
+                         current_level,
+                         len(pending_placeholders),
+                         pending_placeholders)
+
+            # Execute all tools in current level in parallel
+            tasks = []
+            for placeholder in pending_placeholders:
+                step_info = state.evidence_map[placeholder]
+                task = self._execute_single_tool(placeholder, step_info, state.intermediate_results)
+                tasks.append(task)
+
+            # Wait for all tasks in current level to complete
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+
+            # Process results and update intermediate_results
+            updated_intermediate_results = dict(state.intermediate_results)
+
+            for placeholder, result in zip(pending_placeholders, results):
+                if isinstance(result, BaseException):
+                    logger.error("%s Tool execution failed for %s: %s", AGENT_LOG_PREFIX, placeholder, result)
+                    # Create error tool message
+                    error_message = f"Tool execution failed: {str(result)}"
+                    updated_intermediate_results[placeholder] = ToolMessage(content=error_message,
+                                                                            tool_call_id=placeholder)
+                    if self.raise_tool_call_error:
+                        raise result
                 else:
-                    logger.error("%s Invalid step format at index %s", AGENT_LOG_PREFIX, current_step)
-                    return {"intermediate_results": state.intermediate_results}
-            else:
-                logger.error("%s Invalid steps content or index %s", AGENT_LOG_PREFIX, current_step)
-                return {"intermediate_results": state.intermediate_results}
-
-            intermediate_results = state.intermediate_results
-
-            # Replace the placeholder in the tool input with the previous tool output
-            for _placeholder, _tool_output in intermediate_results.items():
-                _tool_output = _tool_output.content
-                # If the content is a list, get the first element which should be a dict
-                if isinstance(_tool_output, list):
-                    _tool_output = _tool_output[0]
-                    assert isinstance(_tool_output, dict)
-
-                tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
-
-            requested_tool = self._get_tool(tool)
-            if not requested_tool:
-                configured_tool_names = list(self.tools_dict.keys())
-                logger.warning(
-                    "%s ReWOO Agent wants to call tool %s. In the ReWOO Agent's configuration within the config file,"
-                    "there is no tool with that name: %s",
-                    AGENT_LOG_PREFIX,
-                    tool,
-                    configured_tool_names)
-
-                intermediate_results[placeholder] = ToolMessage(content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(
-                    tool_name=tool, tools=configured_tool_names),
-                                                                tool_call_id=tool)
-                return {"intermediate_results": intermediate_results}
+                    updated_intermediate_results[placeholder] = result
+                    # Check if the ToolMessage has error status and raise_tool_call_error is True
+                    if (isinstance(result, ToolMessage) and hasattr(result, 'status') and result.status == "error"
+                            and self.raise_tool_call_error):
+                        logger.error("%s Tool call failed for %s: %s", AGENT_LOG_PREFIX, placeholder, result.content)
+                        raise RuntimeError(f"Tool call failed: {result.content}")
 
             if self.detailed_logs:
-                logger.debug("%s Calling tool %s with input: %s", AGENT_LOG_PREFIX, requested_tool.name, tool_input)
+                logger.info("%s Completed level %s with %s tools",
+                            AGENT_LOG_PREFIX,
+                            current_level,
+                            len(pending_placeholders))
 
-            # Run the tool. Try to use structured input, if possible
-            tool_input_parsed = self._parse_tool_input(tool_input)
-            tool_response = await self._call_tool(requested_tool,
-                                                  tool_input_parsed,
-                                                  RunnableConfig(callbacks=self.callbacks),
-                                                  max_retries=3)
-
-            # ToolMessage only accepts str or list[str | dict] as content.
-            # Convert into list if the response is a dict.
-            if isinstance(tool_response, dict):
-                tool_response = [tool_response]
-
-            tool_response_message = ToolMessage(name=tool, tool_call_id=tool, content=tool_response)
-
-            if self.detailed_logs:
-                self._log_tool_response(requested_tool.name, tool_input_parsed, str(tool_response))
-
-            intermediate_results[placeholder] = tool_response_message
-            return {"intermediate_results": intermediate_results}
+            return {"intermediate_results": updated_intermediate_results}
 
         except Exception as ex:
-            logger.exception("%s Failed to call executor_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
-            raise ex
+            logger.error("%s Failed to call executor_node: %s", AGENT_LOG_PREFIX, ex)
+            raise
+
+    async def _execute_single_tool(self,
+                                   placeholder: str,
+                                   step_info: ReWOOPlanStep,
+                                   intermediate_results: dict[str, ToolMessage]) -> ToolMessage:
+        """
+        Execute a single tool with proper placeholder replacement.
+
+        Args:
+            placeholder: The evidence placeholder (e.g., "#E1").
+            step_info: Step information containing tool and tool_input.
+            intermediate_results: Current intermediate results for placeholder replacement.
+
+        Returns:
+            ToolMessage with the tool execution result.
+        """
+        evidence_info = step_info.evidence
+        tool_name = evidence_info.tool
+        tool_input = evidence_info.tool_input
+
+        # Replace placeholders in tool input with previous results
+        for ph_key, tool_output in intermediate_results.items():
+            tool_output_content = tool_output.content
+            # If the content is a list, get the first element which should be a dict
+            if isinstance(tool_output_content, list):
+                tool_output_content = tool_output_content[0]
+                assert isinstance(tool_output_content, dict)
+
+            tool_input = self._replace_placeholder(ph_key, tool_input, tool_output_content)
+
+        # Get the requested tool
+        requested_tool = self._get_tool(tool_name)
+        if not requested_tool:
+            configured_tool_names = list(self.tools_dict.keys())
+            logger.warning(
+                "%s ReWOO Agent wants to call tool %s. In the ReWOO Agent's configuration within the config file,"
+                "there is no tool with that name: %s",
+                AGENT_LOG_PREFIX,
+                tool_name,
+                configured_tool_names)
+
+            return ToolMessage(content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(tool_name=tool_name,
+                                                                           tools=configured_tool_names),
+                               tool_call_id=placeholder)
+
+        if self.detailed_logs:
+            logger.debug("%s Calling tool %s with input: %s", AGENT_LOG_PREFIX, requested_tool.name, tool_input)
+
+        # Parse and execute the tool
+        tool_input_parsed = self._parse_tool_input(tool_input)
+        tool_response = await self._call_tool(
+            requested_tool,
+            tool_input_parsed,
+            RunnableConfig(callbacks=self.callbacks),  # type: ignore
+            max_retries=self.tool_call_max_retries)
+
+        if self.detailed_logs:
+            self._log_tool_response(requested_tool.name, tool_input_parsed, str(tool_response))
+
+        return tool_response
 
     async def solver_node(self, state: ReWOOGraphState):
         try:
             logger.debug("%s Starting the ReWOO Solver Node", AGENT_LOG_PREFIX)
 
             plan = ""
-            # Add the tool outputs of each step to the plan
-            for step in state.steps.content:
-                step_info = step["evidence"]
-                placeholder = step_info.get("placeholder", "")
-                tool_input = step_info.get("tool_input", "")
-
-                intermediate_results = state.intermediate_results
-                for _placeholder, _tool_output in intermediate_results.items():
-                    _tool_output = _tool_output.content
+            # Add the tool outputs of each step to the plan using evidence_map
+            for placeholder, step_info in state.evidence_map.items():
+                evidence_info = step_info.evidence
+                original_tool_input = evidence_info.tool_input
+                tool_name = evidence_info.tool
+
+                # Replace placeholders in tool input with actual results
+                final_tool_input = original_tool_input
+                for ph_key, tool_output in state.intermediate_results.items():
+                    tool_output_content = tool_output.content
                     # If the content is a list, get the first element which should be a dict
-                    if isinstance(_tool_output, list):
-                        _tool_output = _tool_output[0]
-                        assert isinstance(_tool_output, dict)
-
-                    tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
-
-                    placeholder = placeholder.replace(_placeholder, str(_tool_output))
-
-                _plan = step.get("plan")
-                tool = step_info.get("tool")
-                plan += f"Plan: {_plan}\n{placeholder} = {tool}[{tool_input}]"
+                    if isinstance(tool_output_content, list):
+                        tool_output_content = tool_output_content[0]
+                        assert isinstance(tool_output_content, dict)
+
+                    final_tool_input = self._replace_placeholder(ph_key, final_tool_input, tool_output_content)
+
+                # Get the final result for this placeholder
+                final_result = ""
+                if placeholder in state.intermediate_results:
+                    result_content = state.intermediate_results[placeholder].content
+                    if isinstance(result_content, list):
+                        result_content = result_content[0]
+                        if isinstance(result_content, dict):
+                            final_result = str(result_content)
+                    else:
+                        final_result = str(result_content)
+
+                step_plan = step_info.plan
+                plan += '\n'.join([
+                    f"Plan: {step_plan}",
+                    f"{placeholder} = {tool_name}[{final_tool_input}",
+                    f"Result: {final_result}\n\n"
+                ])
 
             task = str(state.task.content)
             solver_prompt = self.solver_prompt.partial(plan=plan)
@@ -324,30 +495,39 @@ class ReWOOAgentGraph(BaseAgent):
             return {"result": output_message}
 
         except Exception as ex:
-            logger.exception("%s Failed to call solver_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
-            raise ex
+            logger.error("%s Failed to call solver_node: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
     async def conditional_edge(self, state: ReWOOGraphState):
         try:
             logger.debug("%s Starting the ReWOO Conditional Edge", AGENT_LOG_PREFIX)
 
-            current_step = self._get_current_step(state)
-            if current_step == -1:
-                logger.debug("%s The ReWOO Executor has finished its task", AGENT_LOG_PREFIX)
+            current_level, level_complete = self._get_current_level_status(state)
+
+            # If all levels are complete, move to solver
+            if current_level == -1:
+                logger.debug("%s All execution levels complete, moving to solver", AGENT_LOG_PREFIX)
                 return AgentDecision.END
 
-            logger.debug("%s The ReWOO Executor is still working on the task", AGENT_LOG_PREFIX)
+            # If current level is complete, check if there are more levels
+            if level_complete:
+                next_level = current_level + 1
+                if next_level >= len(state.execution_levels):
+                    logger.debug("%s All execution levels complete, moving to solver", AGENT_LOG_PREFIX)
+                    return AgentDecision.END
+
+            logger.debug("%s Continuing with executor (level %s, complete: %s)",
+                         AGENT_LOG_PREFIX,
+                         current_level,
+                         level_complete)
             return AgentDecision.TOOL
 
         except Exception as ex:
-            logger.exception("%s Failed to determine whether agent is calling a tool: %s",
-                             AGENT_LOG_PREFIX,
-                             ex,
-                             exc_info=True)
+            logger.exception("%s Failed to determine whether agent is calling a tool: %s", AGENT_LOG_PREFIX, ex)
             logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
             return AgentDecision.END
 
-    async def _build_graph(self, state_schema):
+    async def _build_graph(self, state_schema: type) -> CompiledStateGraph:
         try:
             logger.debug("%s Building and compiling the ReWOO Graph", AGENT_LOG_PREFIX)
 
@@ -357,8 +537,10 @@ class ReWOOAgentGraph(BaseAgent):
             graph.add_node("solver", self.solver_node)
 
             graph.add_edge("planner", "executor")
-            conditional_edge_possible_outputs = {AgentDecision.TOOL: "executor", AgentDecision.END: "solver"}
-            graph.add_conditional_edges("executor", self.conditional_edge, conditional_edge_possible_outputs)
+            graph.add_conditional_edges("executor",
+                                        self.conditional_edge, {
+                                            AgentDecision.TOOL: "executor", AgentDecision.END: "solver"
+                                        })
 
             graph.set_entry_point("planner")
             graph.set_finish_point("solver")
@@ -369,8 +551,8 @@ class ReWOOAgentGraph(BaseAgent):
             return self.graph
 
         except Exception as ex:
-            logger.exception("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            raise ex
+            logger.error("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
     async def build_graph(self):
         try:
@@ -378,8 +560,8 @@ class ReWOOAgentGraph(BaseAgent):
             logger.debug("%s ReWOO Graph built and compiled successfully", AGENT_LOG_PREFIX)
             return self.graph
         except Exception as ex:
-            logger.exception("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            raise ex
+            logger.error("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
     @staticmethod
     def validate_planner_prompt(planner_prompt: str) -> bool:
@@ -395,7 +577,7 @@ class ReWOOAgentGraph(BaseAgent):
                 errors.append(error_message)
         if errors:
             error_text = "\n".join(errors)
-            logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
+            logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
             raise ValueError(error_text)
         return True
 
@@ -406,6 +588,6 @@ class ReWOOAgentGraph(BaseAgent):
             errors.append("The solver prompt cannot be empty.")
         if errors:
             error_text = "\n".join(errors)
-            logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
+            logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
             raise ValueError(error_text)
         return True