How to return directly ONLY if there was no ToolException?

[Leoccleao]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

My custom tool was created with: '''python
class ConditionalReturnTool(BaseTool):
    name: str = "fetch_data"
    description: str = "Use this tool to test your SQL query. Input should be a SQL query that follows guidelines and examples provided."
    handle_tool_error: bool = _handle_error
    return_direct: bool = True


    def _run(self, query: str) -> str:
        try:
            result = fetch_data(query)
            parsed_result = json.loads(result)
            if parsed_result["result"] == "success":
                return result
            else:
                raise ToolException(result)
        except ToolException as e:
            raise e

    def _arun(self, query: str) -> str:
        raise NotImplementedError("This tool does not support async")
'''

I also edited run method from langchain>tools with '''python
        try:
            print("Running the tool...")
            child_config = patch_config(config, callbacks=run_manager.get_child())
            context = copy_context()
            context.run(_set_config_context, child_config)
            tool_args, tool_kwargs = self._to_args_and_kwargs(tool_input)
            if signature(self._run).parameters.get("run_manager"):
                tool_kwargs["run_manager"] = run_manager

            if config_param := _get_runnable_config_param(self._run):
                tool_kwargs[config_param] = config
            response = context.run(self._run, *tool_args, **tool_kwargs)
            print("Tool ran successfully with response:", response)
            if self.response_format == "content_and_artifact":
                if not isinstance(response, tuple) or len(response) != 2:
                    raise ValueError(
                        "Since response_format='content_and_artifact' "
                        "a two-tuple of the message content and raw tool output is "
                        f"expected. Instead generated response of type: "
                        f"{type(response)}."
                    )
                content, artifact = response
            else:
                content = response
            status = "success"
        except ValidationError as e:
            if not self.handle_validation_error:
                error_to_raise = e
            else:
                content = _handle_validation_error(e, flag=self.handle_validation_error)
            status = "error"
        except ToolException as e:
            print("ToolException encountered:", str(e))
            run_manager.on_tool_error(e)
            # Returning error output immediately, preventing reaching the else block
            return _format_output(str(e), None, tool_call_id, self.name, "error")
        
        except (Exception, KeyboardInterrupt) as e:
            print("Another ToolException encountered:", str(e))
            error_to_raise = e
            status = "error"
        else:
            # This block is only executed if no exceptions were raised
            print("No exceptions encountered. Preparing to return observation.")
            output = _format_output(content, artifact, tool_call_id, self.name, status)
            run_manager.on_tool_end(output, color=color, name=self.name, **kwargs)
            if status == "success" and self.return_direct:
                print("Returning direct as return_direct is set to True.")
                return content
            return output

            # If an exception was raised and caught, the method would return earlier.
        if error_to_raise:
            run_manager.on_tool_error(error_to_raise)
            raise error_to_raise

        return output
'''

Description

To optimize latency and token usage I need to use return direct after a tool is called, BUT it is exiting the agent even if the tool throws an exception.

I need to return directly only if an exception is not thrown.

System Info

System Information

OS: Darwin
OS Version: Darwin Kernel Version 23.2.0: Wed Nov 15 21:53:34 PST 2023; root:xnu-10002.61.3~2/RELEASE_ARM64_T8103
Python Version: 3.9.6 (default, Nov 10 2023, 13:38:27)
[Clang 15.0.0 (clang-1500.1.0.2.5)]

Package Information

langchain_core: 0.2.29
langchain: 0.2.12
langchain_community: 0.2.11
langsmith: 0.1.98
langchain_experimental: 0.0.64
langchain_openai: 0.1.21
langchain_text_splitters: 0.2.0
langchain_together: 0.1.2

[dosubot[bot]]

To ensure that your custom tool returns directly only if no ToolException is thrown, you need to modify both the _run and _arun methods in your ConditionalReturnTool class. Here is how you can achieve this:

Modifying the ConditionalReturnTool Class

class ConditionalReturnTool(BaseTool):
    name: str = "fetch_data"
    description: str = "Use this tool to test your SQL query. Input should be a SQL query that follows guidelines and examples provided."
    handle_tool_error: bool = _handle_error
    return_direct: bool = True

    def _run(self, query: str) -> str:
        try:
            result = fetch_data(query)
            parsed_result = json.loads(result)
            if parsed_result["result"] == "success":
                return result
            else:
                raise ToolException(result)
        except ToolException as e:
            raise e

    def _arun(self, query: str) -> str:
        raise NotImplementedError("This tool does not support async")

Modifying the run Method

You need to ensure that the run method in the LangChain library returns directly only if no ToolException is thrown. Here is the modified run method:

def run(
    self,
    tool_input: Union[str, Dict],
    verbose: Optional[bool] = None,
    start_color: Optional[str] = "green",
    color: Optional[str] = "green",
    callbacks: Callbacks = None,
    *,
    tags: Optional[List[str]] = None,
    metadata: Optional[Dict[str, Any]] = None,
    **kwargs: Any,
) -> Any:
    """Run the tool."""
    run_manager = callback_manager.on_tool_start(
        {"name": self.name, "description": self.description},
        tool_input if isinstance(tool_input, str) else str(tool_input),
        inputs=tool_input if isinstance(tool_input, dict) else None,
        **kwargs
    )
    content = None
    artifact = None
    error_to_raise: Union[Exception, KeyboardInterrupt, None] = None
    try:
        child_config = patch_config(config, callbacks=run_manager.get_child())
        context = copy_context()
        context.run(_set_config_context, child_config)
        tool_args, tool_kwargs = self._to_args_and_kwargs(tool_input)
        if signature(self._run).parameters.get("run_manager"):
            tool_kwargs["run_manager"] = run_manager

        if config_param := _get_runnable_config_param(self._run):
            tool_kwargs[config_param] = config
        response = context.run(self._run, *tool_args, **tool_kwargs)
        if self.response_format == "content_and_artifact":
            if not isinstance(response, tuple) or len(response) != 2:
                raise ValueError(
                    "Since response_format='content_and_artifact' "
                    "a two-tuple of the message content and raw tool output is "
                    f"expected. Instead generated response of type: "
                    f"{type(response)}."
                )
            content, artifact = response
        else:
            content = response
        status = "success"
    except ValidationError as e:
        if not self.handle_validation_error:
            error_to_raise = e
        else:
            content = _handle_validation_error(e, flag=self.handle_validation_error)
        status = "error"
    except ToolException as e:
        run_manager.on_tool_error(e)
        error_to_raise = e
        status = "error"
    except (Exception, KeyboardInterrupt) as e:
        error_to_raise = e
        status = "error"
    else:
        output = _format_output(content, artifact, tool_call_id, self.name, status)
        run_manager.on_tool_end(output, color=color, name=self.name, **kwargs)
        if status == "success" and self.return_direct:
            return content
        return output

    if error_to_raise:
        run_manager.on_tool_error(error_to_raise)
        raise error_to_raise

    return output

In this modified run method, the return_direct logic is placed inside the else block, which is only executed if no exceptions are raised. This ensures that the tool's output is returned directly only if no ToolException is encountered ^[1].

To continue talking to Dosu, mention @dosu.

[Leoccleao]

This way the exception was not passed as return direct, but instead of passing the exception back to agent, it was thrown as error:

Traceback (most recent call last):
File "/Users/leonardo/PycharmProjects/frontend_agent/InternalAgent/utils/sql_tool_test.py", line 65, in fetch_data
raise ToolException({"query": query, "result": "failure", "content": result})
langchain_core.tools.ToolException: {'query': 'SELECT DATE(order_totals.alias_creationDate) AS alias_sale_date, SUM(order_totals.alias_total_items_value) AS alias_total_sales FROM (SELECT "orderId", MAX("creationDate") AS alias_creationDate, SUM("total_Items_value") AS alias_total_items_value FROM sales_history WHERE EXTRACT(MONTH FROM "creationDate") = 4 AND EXTRACT(YEAR FROM "creationDate") = 2024 GROUP BY "orderId") AS alias_order_totals GROUP BY alias_sale_date ORDER BY alias_sale_date', 'result': 'failure', 'content': '{"status":422,"error":"missing FROM-clause entry for table \"order_totals\""}'}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/Users/leonardo/PycharmProjects/frontend_agent/InternalAgent/utils/sql_tool_test.py", line 150, in
answer = sql_agent("Retrieve the daily sales data for April 2024. The data should include the date and the total sales amount for each day.")
File "/Users/leonardo/PycharmProjects/frontend_agent/InternalAgent/utils/sql_tool_test.py", line 146, in sql_agent
output = agent_executor.invoke({"input": input})
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain/chains/base.py", line 166, in invoke
raise e
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain/chains/base.py", line 156, in invoke
self._call(inputs, run_manager=run_manager)
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain/agents/agent.py", line 1612, in _call
next_step_output = self._take_next_step(
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain/agents/agent.py", line 1318, in _take_next_step
[
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain/agents/agent.py", line 1318, in
[
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain/agents/agent.py", line 1403, in _iter_next_step
yield self._perform_agent_action(
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain/agents/agent.py", line 1425, in _perform_agent_action
observation = tool.run(
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain_core/tools.py", line 620, in run
raise error_to_raise
File "/Users/leonardo/PycharmProjects/frontend_agent/venv/lib/python3.9/site-packages/langchain_core/tools.py", line 584, in run
response = context.run(self._run, *tool_args, **tool_kwargs)
File "/Users/leonardo/PycharmProjects/frontend_agent/InternalAgent/utils/sql_tool_test.py", line 99, in _run
raise e
File "/Users/leonardo/PycharmProjects/frontend_agent/InternalAgent/utils/sql_tool_test.py", line 92, in _run
result = fetch_data(query)
File "/Users/leonardo/PycharmProjects/frontend_agent/InternalAgent/utils/sql_tool_test.py", line 70, in fetch_data
raise ToolException(json.dumps({"result": "failure", "content": f"An error occurred: {err}"}))
langchain_core.tools.ToolException: {"result": "failure", "content": "An error occurred: {'query': 'SELECT DATE(order_totals.alias_creationDate) AS alias_sale_date, SUM(order_totals.alias_total_items_value) AS alias_total_sales FROM (SELECT "orderId", MAX("creationDate") AS alias_creationDate, SUM("total_Items_value") AS alias_total_items_value FROM sales_history WHERE EXTRACT(MONTH FROM "creationDate") = 4 AND EXTRACT(YEAR FROM "creationDate") = 2024 GROUP BY "orderId") AS alias_order_totals GROUP BY alias_sale_date ORDER BY alias_sale_date', 'result': 'failure', 'content': '{"status":422,"error":"missing FROM-clause entry for table \\"order_totals\\""}'}"}