Repository: run-llama/llama_index
Pull request: Openai agent json malformed 11921
# Description

The OpenAIAgent no longer crashes when the function argument is not valid json. Furthermore, a tool_call_parser can be passed when creating the OpenAIAgent that allows the user to handle/ recover from malformed tool calls. For an example, see docs/examples/agent/openai_agent_tool_call_parser.ipynb

Fixes  #11921

## Version Bump?
- [x] Yes
- [ ] No

## Type of Change

Please delete options that are not relevant.

- [x] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)

## How Has This Been Tested?

- [x] Added unit/integration tests and executed them locally

## Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] I have added Google Colab support for the newly added notebooks.
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Bug]: Agent errors when json function argument is malformed
### Bug Description

We use the OpenAIAgent together with some custom tools.
We noticed, that even gpt4-turbo has some issues with creating a valid json in the function calls when the prompt gets big.
For example, the tool has an input named "query" which should be a string (in our concrete case, python code to be executed). So the function arguments string should look like '{"query": "some long string"}', but often is malformed - e.g. 'query = """some long multi-line string"""'.
Currently, this results in an error in the agent, as in agent.openai.step.call_function, the argument_str is parsed via json.loads, without error handling - which fails in those circumstances and a JSONDecodeError is thrown.

### Version

0.10.19

### Steps to Reproduce
```
from llama_index.agent.openai import OpenAIAgent
from dotenv import load_dotenv
from llama_index.llms.openai import OpenAI
from llama_index.core.tools.function_tool import FunctionTool
from pydantic import BaseModel, Field

load_dotenv()


class PythonInput(BaseModel):
    query: str = Field(
        description="python script or command WIHTOUT COMMENTS which will be evaluated by the eval command.\nNever start variable names with numbers!'"
    )


def python_tool(query: str) -> str:
    # acutally, a fake python tool
    print("Python tool called with query:")
    print(query)
    print("______")
    return query


tool = FunctionTool.from_defaults(
    name="python-tool",
    fn=python_tool,
    description="A Python shell. Use this to execute python commands.\nNever start variable names with numbers!\n",
    fn_schema=PythonInput,
)

long_system_message = """You are working with a python shell that has a pandas DataFrame.
The name of the dataframe is `df`.
You have pandas and numpy available as pd and np.This is the description of df:

This DataFrame contains balance sheet data of the company in question for several years.
Each row contains the balance sheet data of a year.
The column 'Year' is very important. It is sorted DESC and contains the year for which the report is valid.
The column 'Currency' is also important, it specifies the currency in which the numbers are reported in.

Some rules to follow:

1. Group calculations per Year and include the Year column in your answer.
2. If you are asked how a metric developed, respond with the absolute values of the last couple of years, not the percentage changes.
3. Return information from the context in the form of a markdown table.
4. Include the Currency column if possible.
5. Don't use dropna on df, you could lose important information.
6. If calculating diffs: use .diff(-1) to calculate differences to previous year. And don't include the Currency column for the diff, only in the assign part. Good example:
df[['Year', 'Total_Assets', 'Cash_Bank_Deposits', 'AS30', 'Goodwill', 'AS32', 'AS40']].diff(-1).assign(Year=df['Year'], Currency=df['Currency'])

This is the description of the relevant columns:
Column 'Trade_Receivables': Trade Receivables 
Column 'Intragroup_Receivables': Intragroup Receivables
Column 'Other_Receivables': Other Receivables
Column 'Subtotal_Inventory': Subtotal Inventory
Column 'Total_Assets': Total Assets 
Column 'Currency': Currency in which all figures for this balance sheet are reported

This is the result of `print(df.head())`:
   Other_Receivables Currency  Year  Total_Assets  Intragroup_Receivables  Subtotal_Inventory  Trade_Receivables
0       2.025012e+09      EUR  2022  1.551012e+10                       0        7.689101e+08       6.813689e+08
1       3.801111e+08      EUR  2019  1.647862e+10                       0        6.358830e+08       7.145821e+08

Use the tool to answer the questions posed to you."""

# Here is a question about the dataframe, that needs multiple lines of code to answer:
question = "For all of the following lines, calculate its share of total assets: Subtotal Liquid Assets, Trade Receivables, Intragroup Receivables, Other Receivables, Subtotal Inventory, Subtotal Tangible Fixed Assets, Subtotal Intangible Fixed Assets, Subtotal Financial Fixed Assets"
# in our production use case, we tried our best to not confuse the llm with the python_tool docstring or the system message, but the llm still got confused.
for i in range(10):
    llm = OpenAI(
        model="gpt-4-1106-preview", temperature=0.0000000001, additional_kwargs={"seed": i, "top_p": 0.0000000001}
    )
    agent = OpenAIAgent.from_tools(tools=[tool], system_prompt=long_system_message, llm=llm, verbose=True)
    print(f"Seed: {i}")
    response = agent.chat(question)
    print(response)

```

### Relevant Logs/Tracbacks

ERROR:root:Error while running agent.chat("For all of the following lines, calculate its share of total assets: Subtotal Liquid Assets, Trade Receivables, Intragroup Receivables, Other Receivables, Subtotal Inventory, Subtotal Tangible Fixed Assets, Subtotal Intangible Fixed Assets, Subtotal Financial Fixed Assets")
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/core/callbacks/utils.py", line 41, in wrapper
    return func(self, *args, **kwargs)
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/core/agent/runner/base.py", line 575, in chat
    chat_response = self._chat(
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/core/agent/runner/base.py", line 520, in _chat
    cur_step_output = self._run_step(
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/core/agent/runner/base.py", line 372, in _run_step
    cur_step_output = self.agent_worker.run_step(step, task, **kwargs)
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/core/callbacks/utils.py", line 41, in wrapper
    return func(self, *args, **kwargs)
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/agent/openai/step.py", line 572, in run_step
    return self._run_step(
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/agent/openai/step.py", line 469, in _run_step
    self._call_function(
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/agent/openai/step.py", line 361, in _call_function
    function_message, tool_output = call_function(
  File "/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/llama_index/agent/openai/step.py", line 99, in call_function
    argument_dict = json.loads(arguments_str)
  File "/usr/lib/python3.10/json/__init__.py", line 346, in loads
    return _default_decoder.decode(s)
  File "/usr/lib/python3.10/json/decoder.py", line 337, in decode
    obj, end = self.raw_decode(s, idx=_w(s, 0).end())
  File "/usr/lib/python3.10/json/decoder.py", line 355, in raw_decode
    raise JSONDecodeError("Expecting value", s, err.value) from None
json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)

```shell
Added user message to memory: For all of the following lines, calculate its share of total assets: Subtotal Liquid Assets, Trade Receivables, Intragroup Receivables, Other Receivables, Subtotal Inventory, Subtotal Tangible Fixed Assets, Subtotal Intangible Fixed Assets, Subtotal Financial Fixed Assets
=== Calling Function ===
Calling function: python-tool with args: query = """
df.assign(
    Subtotal_Liquid_Assets_Share=df['Cash_Bank_Deposits'] / df['Total_Assets'],
    Trade_Receivables_Share=df['Trade_Receivables'] / df['Total_Assets'],
    Intragroup_Receivables_Share=df['Intragroup_Receivables'] / df['Total_Assets'],
    Other_Receivables_Share=df['Other_Receivables'] / df['Total_Assets'],
    Subtotal_Inventory_Share=df['Subtotal_Inventory'] / df['Total_Assets'],
    Subtotal_Tangible_Fixed_Assets_Share=df['AS40'] / df['Total_Assets'],
    Subtotal_Intangible_Fixed_Assets_Share=df['Goodwill'] + df['AS32'] / df['Total_Assets'],
    Subtotal_Financial_Fixed_Assets_Share=df['AS30'] / df['Total_Assets']
)[['Year', 'Currency', 'Subtotal_Liquid_Assets_Share', 'Trade_Receivables_Share', 'Intragroup_Receivables_Share', 'Other_Receivables_Share', 'Subtotal_Inventory_Share', 'Subtotal_Tangible_Fixed_Assets_Share', 'Subtotal_Intangible_Fixed_Assets_Share', 'Subtotal_Financial_Fixed_Assets_Share']]
"""
```

Hints from issue discussion:
<!-- Answer -->
Thank you for providing a detailed description of the issue you're facing. It seems like the problem arises when the `OpenAIAgent` tries to parse a large or complex JSON string into a Python dictionary using `json.loads(arguments_str)`. If the JSON string is malformed, this can lead to a `JSONDecodeError`.

To handle this, you could add error handling around the JSON parsing to catch `JSONDecodeError` and provide a meaningful error message or fallback mechanism. Here's an example of how you could modify the `call_function` and `acall_function` methods to include error handling for JSON parsing:

```python
import json
from json.decoder import JSONDecodeError

def call_function(
    tools: List[BaseTool],
    tool_call: OpenAIToolCall,
    verbose: bool = False
) -> Tuple[ChatMessage, ToolOutput]:
    # Existing code...
    try:
        argument_dict = json.loads(arguments_str)
    except JSONDecodeError as e:
        logger.error(f"Failed to parse JSON arguments: {arguments_str}. Error: {e}")
        # Handle the error appropriately, e.g., return an error message
        return (
            ChatMessage(
                content="Error: Invalid input format.",
                role=MessageRole.TOOL,
                additional_kwargs={
                    "name": name,
                    "tool_call_id": id_,
                },
            ),
            ToolOutput(
                content="Error: Invalid input format.",
                tool_name=name,
                raw_input={"kwargs": arguments_str},
                raw_output="Invalid JSON format.",
            ),
        )
    # Existing code continues...

async def acall_function(
    tools: List[BaseTool], tool_call: OpenAIToolCall, verbose: bool = False
) -> Tuple[ChatMessage, ToolOutput]:
    # Existing code...
    try:
        argument_dict = json.loads(arguments_str)
    except JSONDecodeError as e:
        logger.error(f"Failed to parse JSON arguments: {arguments_str}. Error: {e}")
        # Handle the error appropriately
        return (
            ChatMessage(
                content="Error: Invalid input format.",
                role=MessageRole.TOOL,
                additional_kwargs={
                    "name": name,
                    "tool_call_id": id_,
                },
            ),
            ToolOutput(
                content="Error: Invalid input format.",
                tool_name=name,
                raw_input={"kwargs": arguments_str},
                raw_output="Invalid JSON format.",
            ),
        )
    # Existing code continues...
```

These modifications add a try-except block around the `json.loads(arguments_str)` call to catch `JSONDecodeError` and log an error message. Additionally, it returns an error message as part of the `ChatMessage` and `ToolOutput` indicating that the input format is invalid. This approach allows for graceful handling of JSON parsing errors and provides feedback that can be used to correct the input format.

Please note that you'll need to replace `logger.error` with your preferred logging method, and adjust the error handling to suit your needs.

I hope this helps! If you have any further questions or need more assistance, feel free to ask.

## Sources
- [RouterQueryEngine sometimes raises JSONDecodeError during inference](https://github.com/run-llama/llama_index/issues/2808)
- [llama-index-legacy/llama_index/legacy/agent/openai/step.py](https://github.com/run-llama/llama_index/blob/main/llama-index-legacy/llama_index/legacy/agent/openai/step.py)
- [llama-index-integrations/llms/llama-index-llms-openai/llama_index/llms/openai/utils.py](https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/llms/llama-index-llms-openai/llama_index/llms/openai/utils.py)
- [llama-index-legacy/llama_index/legacy/agent/openai_assistant_agent.py](https://github.com/run-llama/llama_index/blob/main/llama-index-legacy/llama_index/legacy/agent/openai_assistant_agent.py)

<!-- Dosu Comment Footer -->
<details open>
<summary>

#### About Dosu
</summary>

This response is meant to be useful and save you time. It is **not meant to be a precise solution**, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.
</details>
So from my point of view, there are two separate issues:
- The JSONDecodeError is not handled at all
- It could be attempted to parse the input in a smarter way than just using json.loads

I implemented already something for both issues, however, I think it would be even nicer if the user could optionally pass their own tool-call-parser instead. That could be a function that takes the argument_str and returns the argument_dict.

The tool-call-parser could throw ValueError if the input is still malformed and that error is caught and a meaningful ChatMessage and ToolOutput could be returned. In my case I did:

```
    error_message: Optional[str] = None
    try:
        argument_dict = parse_tool_call(arguments_str=arguments_str, tool=tool)
    except ValueError as e:
        error_message = (
            f"Error in calling function {name}: The input json block is malformed:\n```json\n{arguments_str}\n```"
        )
        return (
            ChatMessage(
                content=error_message,
                role=MessageRole.TOOL,
                additional_kwargs={
                    "name": name,
                    "tool_call_id": id_,
                },
            ),
            ToolOutput(
                content=error_message,
                tool_name=name,
                raw_input={"args": arguments_str},
                raw_output=error_message,
            ),
        )
    # Call tool
    # Use default error message except if json parsing fails
    output = call_tool_with_error_handling(tool, argument_dict, error_message=error_message)
```

And here is my implementation for the parser:

```
def parse_tool_call(arguments_str: str, tool: BaseTool) -> Dict:
    """Parse tool call."""
    if len(arguments_str.strip()) == 0:
        # OpenAI returns an empty string for functions containing no args
        return {}
    try:
        tool_call = json.loads(arguments_str)
        if not isinstance(tool_call, dict):
            raise ValueError("Tool call must be a dictionary")
        return tool_call
    except json.JSONDecodeError as e:
        # if there is only one required argument, we can try to parse it
        if tool.metadata.fn_schema and len(tool.metadata.fn_schema.schema()["required"]) == 1:
            arg_name = list(tool.metadata.fn_schema.schema()["required"])[0]
            # check if arguments_str is of format '<arg_name> = """value"""' or '<arg_name> = "value"'
            if arguments_str.startswith(f'{arg_name} = """') and '"""' in arguments_str.split(f'{arg_name} = """')[1]:
                # parse until the first occurence of '"""' after <arg_name> = '""""':
                return {arg_name: arguments_str.split(f'{arg_name} = """')[1].split('"""')[0]}
            elif arguments_str.startswith(f'{arg_name} = "') and '"' in arguments_str.split(f'{arg_name} = "')[1]:
                # parse until the first occurence of '"' after <arg_name> = '"':
                return {arg_name: arguments_str.split(f'{arg_name} = "')[1].split('"')[0]}
            else:
                return {arg_name: arguments_str}
        raise ValueError(f"Invalid tool call: {e!s}")
```

I could open up a PR with the proposed changes, including some tests.

Commit message: adding advanced_tool_call_parser with tests

Patch to review:
diff --git a/docs/examples/agent/openai_agent_tool_call_parser.ipynb b/docs/examples/agent/openai_agent_tool_call_parser.ipynb
new file mode 100644
index 0000000000..326200ceec
--- /dev/null
+++ b/docs/examples/agent/openai_agent_tool_call_parser.ipynb
@@ -0,0 +1,346 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "24103c51",
+   "metadata": {},
+   "source": [
+    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/openai_agent.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "99cea58c-48bc-4af6-8358-df9695659983",
+   "metadata": {},
+   "source": [
+    "# OpenAI Agent with Tool Call Parser"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "673df1fe-eb6c-46ea-9a73-a96e7ae7942e",
+   "metadata": {},
+   "source": [
+    "Unfortunately, the tool calls by OpenAI are not always valid json, especially from older versions of the API. Up to and including the OpenAI API version 1106, this issue is relatively frequent if the argument is a long string (e.g. a python script), see for example [here](https://community.openai.com/t/malformed-json-in-gpt4-1106-function-arguments/685884).\n",
+    "\n",
+    "With the default tool call parser, the OpenAI Agent will fail to parse these tool calls and tries to fix the tool call in the next step. This needs another llm call, which is slow and expensive.\n",
+    "\n",
+    "This notebook demonstrates how to define a custom tool call parser that can handle certain kinds of malformed function calls.\n",
+    "The following steps are copied from the OpenAI Agent notebook, with the addition of a custom tool call parser."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "54b7bc2e-606f-411a-9490-fcfab9236dfc",
+   "metadata": {},
+   "source": [
+    "## Initial Setup "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "23e80e5b-aaee-4f23-b338-7ae62b08141f",
+   "metadata": {},
+   "source": [
+    "Let's start by importing some simple building blocks.  \n",
+    "\n",
+    "The main thing we need is:\n",
+    "1. the OpenAI API (using our own `llama_index` LLM class)\n",
+    "2. a place to keep conversation history \n",
+    "3. a definition for tools that our agent can use."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "41101795",
+   "metadata": {},
+   "source": [
+    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4985c578",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install llama-index-agent-openai\n",
+    "%pip install llama-index-llms-openai"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c61c873d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install llama-index"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9d47283b-025e-4874-88ed-76245b22f82e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from llama_index.core.tools import FunctionTool\n",
+    "\n",
+    "import nest_asyncio\n",
+    "\n",
+    "nest_asyncio.apply()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6fe08eb1-e638-4c00-9103-5c305bfacccf",
+   "metadata": {},
+   "source": [
+    "Let's define some very simple calculator tools for our agent."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3dd3c4a6-f3e0-46f9-ad3b-7ba57d1bc992",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def multiply(a: int, b: int) -> int:\n",
+    "    \"\"\"Multiple two integers and returns the result integer\"\"\"\n",
+    "    return a * b\n",
+    "\n",
+    "\n",
+    "multiply_tool = FunctionTool.from_defaults(fn=multiply)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bfcfb78b-7d4f-48d9-8d4c-ffcded23e7ac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def add(a: int, b: int) -> int:\n",
+    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
+    "    return a + b\n",
+    "\n",
+    "\n",
+    "add_tool = FunctionTool.from_defaults(fn=add)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "30016529",
+   "metadata": {},
+   "source": [
+    "## Definition of the Tool Call Parser\n",
+    "\n",
+    "Sometimes, OpenAI tool calls are not valid json\n",
+    "\n",
+    "When defining your own Tool Call Parser, you need to define a function that takes a OpenAIToolCall and returns a dictionary. The dictionary will be passed as **kwargs to the tool function.\n",
+    "\n",
+    "The Parser should throw a ValueError if the tool call can't be parsed. This will be returned to the agent and it will try to fix the call on the next step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ffe547b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Dict\n",
+    "from llama_index.llms.openai.utils import OpenAIToolCall\n",
+    "import re\n",
+    "\n",
+    "# The same parser is available as\n",
+    "# from llama_index.agent.openai import advanced_tool_call_parser\n",
+    "\n",
+    "\n",
+    "def custom_tool_call_parser(tool_call: OpenAIToolCall) -> Dict:\n",
+    "    r\"\"\"Parse tool calls that are not standard json.\n",
+    "    Also parses tool calls of the following forms:\n",
+    "    variable = \\\"\\\"\\\"Some long text\\\"\\\"\\\"\n",
+    "    variable = \"Some long text\"'\n",
+    "    variable = '''Some long text'''\n",
+    "    variable = 'Some long text'\n",
+    "    \"\"\"\n",
+    "    arguments_str = tool_call.function.arguments\n",
+    "    if len(arguments_str.strip()) == 0:\n",
+    "        # OpenAI returns an empty string for functions containing no args\n",
+    "        return {}\n",
+    "    try:\n",
+    "        tool_call = json.loads(arguments_str)\n",
+    "        if not isinstance(tool_call, dict):\n",
+    "            raise ValueError(\"Tool call must be a dictionary\")\n",
+    "        return tool_call\n",
+    "    except json.JSONDecodeError as e:\n",
+    "        # pattern to match variable names and content within quotes\n",
+    "        pattern = r'([a-zA-Z_][a-zA-Z_0-9]*)\\s*=\\s*[\"\\']+(.*?)[\"\\']+'\n",
+    "        match = re.search(pattern, arguments_str)\n",
+    "\n",
+    "        if match:\n",
+    "            variable_name = match.group(1)  # This is the variable name\n",
+    "            content = match.group(2)  # This is the content within the quotes\n",
+    "            return {variable_name: content}\n",
+    "        raise ValueError(f\"Invalid tool call: {e!s}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "707d30b8-6405-4187-a9ed-6146dcc42167",
+   "metadata": {},
+   "source": [
+    "## Defining the OpenAI Agent with Tool Call Parser"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "38ab3938-1138-43ea-b085-f430b42f5377",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llama_index.agent.openai import OpenAIAgent\n",
+    "from llama_index.llms.openai import OpenAI"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5b762593",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from dotenv import load_dotenv\n",
+    "\n",
+    "load_dotenv()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d852ece7-e5a1-4368-9d59-c7014e0b5b4d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = OpenAI(model=\"gpt-3.5-turbo-0613\")\n",
+    "agent = OpenAIAgent.from_tools(\n",
+    "    [multiply_tool, add_tool],\n",
+    "    llm=llm,\n",
+    "    verbose=True,\n",
+    "    tool_call_parser=custom_tool_call_parser,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "500cbee4",
+   "metadata": {},
+   "source": [
+    "### Chat"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9fd1cad5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Added user message to memory: What is (121 * 3) + 42?\n",
+      "NOT_GIVEN\n",
+      "[{'role': 'user', 'content': 'What is (121 * 3) + 42?'}]\n",
+      "ChatCompletion(id='chatcmpl-94ID6Oa58vLKGFk3enr06hi4PFuhQ', choices=[Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content=None, role='assistant', function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_U2TN1oM7iJBuZ0L78RVG3Z1K', function=Function(arguments='{\\n  \"a\": 121,\\n  \"b\": 3\\n}', name='multiply'), type='function')]))], created=1710810340, model='gpt-3.5-turbo-0613', object='chat.completion', system_fingerprint=None, usage=CompletionUsage(completion_tokens=21, prompt_tokens=117, total_tokens=138))\n",
+      "=== Calling Function ===\n",
+      "Calling function: multiply with args: {\n",
+      "  \"a\": 121,\n",
+      "  \"b\": 3\n",
+      "}\n",
+      "Got output: 363\n",
+      "========================\n",
+      "\n",
+      "NOT_GIVEN\n",
+      "[{'role': 'user', 'content': 'What is (121 * 3) + 42?'}, {'role': 'assistant', 'content': None, 'tool_calls': [ChatCompletionMessageToolCall(id='call_U2TN1oM7iJBuZ0L78RVG3Z1K', function=Function(arguments='{\\n  \"a\": 121,\\n  \"b\": 3\\n}', name='multiply'), type='function')]}, {'role': 'tool', 'content': '363', 'name': 'multiply', 'tool_call_id': 'call_U2TN1oM7iJBuZ0L78RVG3Z1K'}]\n",
+      "ChatCompletion(id='chatcmpl-94ID7RlMkMPCTGJe0jdgllMXfX73M', choices=[Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content=None, role='assistant', function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_68PCmKAEpIdEayjDQNOLWyLv', function=Function(arguments='{\\n  \"a\": 363,\\n  \"b\": 42\\n}', name='add'), type='function')]))], created=1710810341, model='gpt-3.5-turbo-0613', object='chat.completion', system_fingerprint=None, usage=CompletionUsage(completion_tokens=21, prompt_tokens=146, total_tokens=167))\n",
+      "=== Calling Function ===\n",
+      "Calling function: add with args: {\n",
+      "  \"a\": 363,\n",
+      "  \"b\": 42\n",
+      "}\n",
+      "Got output: 405\n",
+      "========================\n",
+      "\n",
+      "NOT_GIVEN\n",
+      "[{'role': 'user', 'content': 'What is (121 * 3) + 42?'}, {'role': 'assistant', 'content': None, 'tool_calls': [ChatCompletionMessageToolCall(id='call_U2TN1oM7iJBuZ0L78RVG3Z1K', function=Function(arguments='{\\n  \"a\": 121,\\n  \"b\": 3\\n}', name='multiply'), type='function')]}, {'role': 'tool', 'content': '363', 'name': 'multiply', 'tool_call_id': 'call_U2TN1oM7iJBuZ0L78RVG3Z1K'}, {'role': 'assistant', 'content': None, 'tool_calls': [ChatCompletionMessageToolCall(id='call_68PCmKAEpIdEayjDQNOLWyLv', function=Function(arguments='{\\n  \"a\": 363,\\n  \"b\": 42\\n}', name='add'), type='function')]}, {'role': 'tool', 'content': '405', 'name': 'add', 'tool_call_id': 'call_68PCmKAEpIdEayjDQNOLWyLv'}]\n",
+      "ChatCompletion(id='chatcmpl-94ID83hvSDLGF2GKbUa537AtiGZbn', choices=[Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content='The result of (121 * 3) + 42 is 405.', role='assistant', function_call=None, tool_calls=None))], created=1710810342, model='gpt-3.5-turbo-0613', object='chat.completion', system_fingerprint=None, usage=CompletionUsage(completion_tokens=17, prompt_tokens=175, total_tokens=192))\n",
+      "The result of (121 * 3) + 42 is 405.\n"
+     ]
+    }
+   ],
+   "source": [
+    "response = agent.chat(\"What is (121 * 3) + 42?\")\n",
+    "print(str(response))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "538bf32f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[ToolOutput(content='363', tool_name='multiply', raw_input={'args': (), 'kwargs': {'a': 121, 'b': 3}}, raw_output=363), ToolOutput(content='405', tool_name='add', raw_input={'args': (), 'kwargs': {'a': 363, 'b': 42}}, raw_output=405)]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# inspect sources\n",
+    "print(response.sources)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/__init__.py b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/__init__.py
index 8526cfa294..d6d6e65c55 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/__init__.py
+++ b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/__init__.py
@@ -1,5 +1,10 @@
 from llama_index.agent.openai.base import OpenAIAgent
 from llama_index.agent.openai.openai_assistant_agent import OpenAIAssistantAgent
-from llama_index.agent.openai.step import OpenAIAgentWorker
+from llama_index.agent.openai.step import OpenAIAgentWorker, advanced_tool_call_parser
 
-__all__ = ["OpenAIAgent", "OpenAIAgentWorker", "OpenAIAssistantAgent"]
+__all__ = [
+    "OpenAIAgent",
+    "OpenAIAgentWorker",
+    "OpenAIAssistantAgent",
+    "advanced_tool_call_parser",
+]
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/base.py b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/base.py
index b604ecf13e..2bc3a048ee 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/base.py
+++ b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/base.py
@@ -10,7 +10,9 @@
 
 from typing import (
     Any,
+    Dict,
     List,
+    Callable,
     Optional,
     Type,
 )
@@ -26,6 +28,7 @@
 from llama_index.core.settings import Settings
 from llama_index.core.tools import BaseTool
 from llama_index.llms.openai import OpenAI
+from llama_index.llms.openai.utils import OpenAIToolCall
 
 DEFAULT_MAX_FUNCTION_CALLS = 5
 
@@ -53,6 +56,7 @@ def __init__(
         default_tool_choice: str = "auto",
         callback_manager: Optional[CallbackManager] = None,
         tool_retriever: Optional[ObjectRetriever[BaseTool]] = None,
+        tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None,
     ) -> None:
         """Init params."""
         callback_manager = callback_manager or llm.callback_manager
@@ -64,6 +68,7 @@ def __init__(
             max_function_calls=max_function_calls,
             callback_manager=callback_manager,
             prefix_messages=prefix_messages,
+            tool_call_parser=tool_call_parser,
         )
         super().__init__(
             step_engine,
@@ -88,6 +93,7 @@ def from_tools(
         callback_manager: Optional[CallbackManager] = None,
         system_prompt: Optional[str] = None,
         prefix_messages: Optional[List[ChatMessage]] = None,
+        tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None,
         **kwargs: Any,
     ) -> "OpenAIAgent":
         """Create an OpenAIAgent from a list of tools.
@@ -133,4 +139,5 @@ def from_tools(
             max_function_calls=max_function_calls,
             callback_manager=callback_manager,
             default_tool_choice=default_tool_choice,
+            tool_call_parser=tool_call_parser,
         )
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
index 3a3f7b22ea..ec9da75b5e 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
+++ b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
@@ -6,7 +6,8 @@
 import uuid
 from functools import partial
 from threading import Thread
-from typing import Any, Dict, List, Optional, Tuple, Union, cast, get_args
+from typing import Any, Dict, List, Callable, Optional, Tuple, Union, cast, get_args
+import re
 
 from llama_index.agent.openai.utils import resolve_tool_choice
 from llama_index.core.agent.types import (
@@ -77,10 +78,54 @@ def call_tool_with_error_handling(
         )
 
 
+def default_tool_call_parser(tool_call: OpenAIToolCall):
+    try:
+        json.loads(tool_call.function.arguments)
+    except json.JSONDecodeError as e:
+        raise ValueError(
+            f"Error in calling tool {tool_call.function.name}: The input json block is malformed:\n```json\n{tool_call.function.arguments}\n```"
+        )
+
+
+def advanced_tool_call_parser(tool_call: OpenAIToolCall) -> Dict:
+    r"""Parse tool calls that are not standard json.
+
+    Also parses tool calls of the following forms:
+    variable = \"\"\"Some long text\"\"\"
+    variable = "Some long text"'
+    variable = '''Some long text'''
+    variable = 'Some long text'
+    """
+    arguments_str = tool_call.function.arguments
+    if len(arguments_str.strip()) == 0:
+        # OpenAI returns an empty string for functions containing no args
+        return {}
+    try:
+        tool_call = json.loads(arguments_str)
+        if not isinstance(tool_call, dict):
+            raise ValueError(
+                f"Error in calling tool {tool_call.function.name}: The input json block is malformed:\n```json\n{tool_call.function.arguments}\n```"
+            )
+        return tool_call
+    except json.JSONDecodeError as e:
+        # pattern to match variable names and content within quotes
+        pattern = r'([a-zA-Z_][a-zA-Z_0-9]*)\s*=\s*["\']+(.*?)["\']+'
+        match = re.search(pattern, arguments_str)
+
+        if match:
+            variable_name = match.group(1)  # This is the variable name
+            content = match.group(2)  # This is the content within the quotes
+            return {variable_name: content}
+        raise ValueError(
+            f"Error in calling tool {tool_call.function.name}: The input json block is malformed:\n```json\n{tool_call.function.arguments}\n```"
+        )
+
+
 def call_function(
     tools: List[BaseTool],
     tool_call: OpenAIToolCall,
     verbose: bool = False,
+    tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None,
 ) -> Tuple[ChatMessage, ToolOutput]:
     """Call a function and return the output as a string."""
     # validations to get passed mypy
@@ -88,6 +133,7 @@ def call_function(
     assert tool_call.function is not None
     assert tool_call.function.name is not None
     assert tool_call.function.arguments is not None
+    tool_call_parser = tool_call_parser or default_tool_call_parser
 
     id_ = tool_call.id
     _function_call = tool_call.function
@@ -97,11 +143,33 @@ def call_function(
         print("=== Calling Function ===")
         print(f"Calling function: {name} with args: {arguments_str}")
     tool = get_function_by_name(tools, name)
-    argument_dict = json.loads(arguments_str)
+    error_message: Optional[str] = None
+    try:
+        argument_dict = tool_call_parser(tool_call)
+    except ValueError as e:
+        error_message = str(e)
+        return (
+            ChatMessage(
+                content=error_message,
+                role=MessageRole.TOOL,
+                additional_kwargs={
+                    "name": name,
+                    "tool_call_id": id_,
+                },
+            ),
+            ToolOutput(
+                content=error_message,
+                tool_name=name,
+                raw_input={"args": arguments_str},
+                raw_output=error_message,
+            ),
+        )
 
     # Call tool
-    # Use default error message
-    output = call_tool_with_error_handling(tool, argument_dict, error_message=None)
+    # Use default error message except if json parsing fails
+    output = call_tool_with_error_handling(
+        tool, argument_dict, error_message=error_message
+    )
     if verbose:
         print(f"Got output: {output!s}")
         print("========================\n")
@@ -119,7 +187,10 @@ def call_function(
 
 
 async def acall_function(
-    tools: List[BaseTool], tool_call: OpenAIToolCall, verbose: bool = False
+    tools: List[BaseTool],
+    tool_call: OpenAIToolCall,
+    verbose: bool = False,
+    tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None,
 ) -> Tuple[ChatMessage, ToolOutput]:
     """Call a function and return the output as a string."""
     # validations to get passed mypy
@@ -127,6 +198,7 @@ async def acall_function(
     assert tool_call.function is not None
     assert tool_call.function.name is not None
     assert tool_call.function.arguments is not None
+    tool_call_parser = tool_call_parser or default_tool_call_parser
 
     id_ = tool_call.id
     _function_call = tool_call.function
@@ -137,7 +209,28 @@ async def acall_function(
         print(f"Calling function: {name} with args: {arguments_str}")
     tool = get_function_by_name(tools, name)
     async_tool = adapt_to_async_tool(tool)
-    argument_dict = json.loads(arguments_str)
+    error_message: Optional[str] = None
+    try:
+        argument_dict = tool_call_parser(tool_call)
+    except ValueError as e:
+        error_message = str(e)
+        return (
+            ChatMessage(
+                content=error_message,
+                role=MessageRole.TOOL,
+                additional_kwargs={
+                    "name": name,
+                    "tool_call_id": id_,
+                },
+            ),
+            ToolOutput(
+                content=error_message,
+                tool_name=name,
+                raw_input={"args": arguments_str},
+                raw_output=error_message,
+            ),
+        )
+
     output = await async_tool.acall(**argument_dict)
     if verbose:
         print(f"Got output: {output!s}")
@@ -167,12 +260,14 @@ def __init__(
         max_function_calls: int = DEFAULT_MAX_FUNCTION_CALLS,
         callback_manager: Optional[CallbackManager] = None,
         tool_retriever: Optional[ObjectRetriever[BaseTool]] = None,
+        tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None,
     ):
         self._llm = llm
         self._verbose = verbose
         self._max_function_calls = max_function_calls
         self.prefix_messages = prefix_messages
         self.callback_manager = callback_manager or self._llm.callback_manager
+        self.tool_call_parser = tool_call_parser or default_tool_call_parser
 
         if len(tools) > 0 and tool_retriever is not None:
             raise ValueError("Cannot specify both tools and tool_retriever")
@@ -196,6 +291,7 @@ def from_tools(
         callback_manager: Optional[CallbackManager] = None,
         system_prompt: Optional[str] = None,
         prefix_messages: Optional[List[ChatMessage]] = None,
+        tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None,
         **kwargs: Any,
     ) -> "OpenAIAgentWorker":
         """Create an OpenAIAgent from a list of tools.
@@ -236,6 +332,7 @@ def from_tools(
             verbose=verbose,
             max_function_calls=max_function_calls,
             callback_manager=callback_manager,
+            tool_call_parser=tool_call_parser,
         )
 
     def get_all_messages(self, task: Task) -> List[ChatMessage]:
@@ -362,7 +459,10 @@ def _call_function(
             },
         ) as event:
             function_message, tool_output = call_function(
-                tools, tool_call, verbose=self._verbose
+                tools,
+                tool_call,
+                verbose=self._verbose,
+                tool_call_parser=self.tool_call_parser,
             )
             event.on_end(payload={EventPayload.FUNCTION_OUTPUT: str(tool_output)})
         sources.append(tool_output)
@@ -391,7 +491,10 @@ async def _acall_function(
             },
         ) as event:
             function_message, tool_output = await acall_function(
-                tools, tool_call, verbose=self._verbose
+                tools,
+                tool_call,
+                verbose=self._verbose,
+                tool_call_parser=self.tool_call_parser,
             )
             event.on_end(payload={EventPayload.FUNCTION_OUTPUT: str(tool_output)})
         sources.append(tool_output)
@@ -446,6 +549,12 @@ def _run_step(
         openai_tools = [tool.metadata.to_openai_tool() for tool in tools]
 
         llm_chat_kwargs = self._get_llm_chat_kwargs(task, openai_tools, tool_choice)
+        for tool in tools:  # self.agent_worker._get_tools("test"):
+            print("Tool metadata for ", tool.metadata.name)
+            print(tool.metadata)
+            print(tool.metadata.fn_schema.schema())
+        for message in llm_chat_kwargs["messages"]:
+            print(message)
 
         agent_chat_response = self._get_agent_response(
             task, mode=mode, **llm_chat_kwargs
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/pyproject.toml b/llama-index-integrations/agent/llama-index-agent-openai/pyproject.toml
index 24b1ee7799..a5d9502c85 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai/pyproject.toml
+++ b/llama-index-integrations/agent/llama-index-agent-openai/pyproject.toml
@@ -28,7 +28,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-agent-openai"
 readme = "README.md"
-version = "0.1.5"
+version = "0.2.0"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/tests/test_openai_agent.py b/llama-index-integrations/agent/llama-index-agent-openai/tests/test_openai_agent.py
new file mode 100644
index 0000000000..02970fa44b
--- /dev/null
+++ b/llama-index-integrations/agent/llama-index-agent-openai/tests/test_openai_agent.py
@@ -0,0 +1,448 @@
+from typing import Any, AsyncGenerator, Generator, List, Sequence
+from unittest.mock import MagicMock, patch
+
+import pytest
+from llama_index.agent.openai.base import OpenAIAgent
+from llama_index.agent.openai.step import (
+    call_tool_with_error_handling,
+    advanced_tool_call_parser,
+)
+from llama_index.core.base.llms.types import ChatMessage, ChatResponse
+from llama_index.core.chat_engine.types import (
+    AgentChatResponse,
+    StreamingAgentChatResponse,
+)
+from llama_index.core.llms.mock import MockLLM
+from llama_index.core.tools.function_tool import FunctionTool
+from llama_index.llms.openai import OpenAI
+
+from openai.types.chat.chat_completion import ChatCompletion, Choice
+from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDelta
+from openai.types.chat.chat_completion_message import ChatCompletionMessage
+from openai.types.chat.chat_completion_message_tool_call import (
+    ChatCompletionMessageToolCall,
+    Function,
+)
+
+
+def mock_chat_completion(*args: Any, **kwargs: Any) -> ChatCompletion:
+    if "functions" in kwargs:
+        if not kwargs["functions"]:
+            raise ValueError("functions must not be empty")
+
+    # Example taken from https://platform.openai.com/docs/api-reference/chat/create
+    return ChatCompletion(
+        id="chatcmpl-abc123",
+        object="chat.completion",
+        created=1677858242,
+        model="gpt-3.5-turbo-0301",
+        usage={"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
+        choices=[
+            Choice(
+                message=ChatCompletionMessage(
+                    role="assistant", content="\n\nThis is a test!"
+                ),
+                finish_reason="stop",
+                index=0,
+                logprobs=None,
+            )
+        ],
+    )
+
+
+def mock_chat_completion_tool_call(
+    function: Function, *args: Any, **kwargs: Any
+) -> ChatCompletion:
+    # Example taken from https://platform.openai.com/docs/api-reference/chat/create
+    return ChatCompletion(
+        id="chatcmpl-abc123",
+        object="chat.completion",
+        created=1677858242,
+        model="gpt-3.5-turbo-0301",
+        usage={"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
+        choices=[
+            Choice(
+                message=ChatCompletionMessage(
+                    role="assistant",
+                    content="\n\nThis is a test!",
+                    tool_calls=[
+                        ChatCompletionMessageToolCall(
+                            id="toolcall-abc123",
+                            function=function,
+                            type="function",
+                        )
+                    ],
+                ),
+                finish_reason="stop",
+                index=0,
+                logprobs=None,
+            )
+        ],
+    )
+
+
+def mock_chat_stream(
+    *args: Any, **kwargs: Any
+) -> Generator[ChatCompletionChunk, None, None]:
+    if "functions" in kwargs:
+        if not kwargs["functions"]:
+            raise ValueError("functions must not be empty")
+
+    yield ChatCompletionChunk(
+        id="chatcmpl-abc123",
+        object="chat.completion.chunk",
+        created=1677858242,
+        model="gpt-3.5-turbo-0301",
+        choices=[
+            ChunkChoice(
+                delta=ChoiceDelta(role="assistant", content="\n\nThis is a test!"),
+                finish_reason="stop",
+                index=0,
+                logprobs=None,
+            )
+        ],
+    )
+
+
+async def mock_achat_completion(*args: Any, **kwargs: Any) -> ChatCompletion:
+    return mock_chat_completion(*args, **kwargs)
+
+
+async def mock_achat_stream(
+    *args: Any, **kwargs: Any
+) -> AsyncGenerator[ChatCompletionChunk, None]:
+    async def _mock_achat_stream(
+        *args: Any, **kwargs: Any
+    ) -> AsyncGenerator[ChatCompletionChunk, None]:
+        if "functions" in kwargs:
+            if not kwargs["functions"]:
+                raise ValueError("functions must not be empty")
+
+        yield ChatCompletionChunk(
+            id="chatcmpl-abc123",
+            object="chat.completion.chunk",
+            created=1677858242,
+            model="gpt-3.5-turbo-0301",
+            choices=[
+                ChunkChoice(
+                    delta=ChoiceDelta(role="assistant", content="\n\nThis is a test!"),
+                    finish_reason="stop",
+                    index=0,
+                    logprobs=None,
+                )
+            ],
+        )
+
+    return _mock_achat_stream(*args, **kwargs)
+
+
+@pytest.fixture()
+def add_tool() -> FunctionTool:
+    def add(a: int, b: int) -> int:
+        """Add two integers and returns the result integer."""
+        return a + b
+
+    return FunctionTool.from_defaults(fn=add)
+
+
+@pytest.fixture()
+def echo_tool() -> FunctionTool:
+    def echo(query: str) -> str:
+        """Echos input."""
+        return query
+
+    return FunctionTool.from_defaults(fn=echo)
+
+
+class MockChatLLM(MockLLM):
+    def __init__(self, responses: List[ChatMessage]) -> None:
+        self._i = 0  # call counter, determines which response to return
+        self._responses = responses  # list of responses to return
+
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        del messages  # unused
+        response = ChatResponse(
+            message=self._responses[self._i],
+        )
+        self._i += 1
+        return response
+
+
+MOCK_ACTION_RESPONSE = """\
+Thought: I need to use a tool to help me answer the question.
+Action: add
+Action Input: {"a": 1, "b": 1}
+"""
+
+MOCK_FINAL_RESPONSE = """\
+Thought: I have enough information to answer the question without using any more tools.
+Answer: 2
+"""
+
+
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+def test_chat_basic(MockSyncOpenAI: MagicMock, add_tool: FunctionTool) -> None:
+    mock_instance = MockSyncOpenAI.return_value
+    mock_instance.chat.completions.create.return_value = mock_chat_completion()
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+
+    agent = OpenAIAgent.from_tools(
+        tools=[add_tool],
+        llm=llm,
+    )
+    response = agent.chat("What is 1 + 1?")
+    assert isinstance(response, AgentChatResponse)
+    assert response.response == "\n\nThis is a test!"
+    assert len(agent.chat_history) == 2
+    assert agent.chat_history[0].content == "What is 1 + 1?"
+    assert agent.chat_history[1].content == "\n\nThis is a test!"
+
+
+@patch("llama_index.llms.openai.base.AsyncOpenAI")
+@pytest.mark.asyncio()
+async def test_achat_basic(MockAsyncOpenAI: MagicMock, add_tool: FunctionTool) -> None:
+    mock_instance = MockAsyncOpenAI.return_value
+    mock_instance.chat.completions.create.return_value = mock_achat_completion()
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+
+    agent = OpenAIAgent.from_tools(
+        tools=[add_tool],
+        llm=llm,
+    )
+    response = await agent.achat("What is 1 + 1?")
+    assert isinstance(response, AgentChatResponse)
+    assert response.response == "\n\nThis is a test!"
+    assert len(agent.chat_history) == 2
+    assert agent.chat_history[0].content == "What is 1 + 1?"
+    assert agent.chat_history[1].content == "\n\nThis is a test!"
+
+
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+def test_stream_chat_basic(MockSyncOpenAI: MagicMock, add_tool: FunctionTool) -> None:
+    mock_instance = MockSyncOpenAI.return_value
+    mock_instance.chat.completions.create.side_effect = mock_chat_stream
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+
+    agent = OpenAIAgent.from_tools(
+        tools=[add_tool],
+        llm=llm,
+    )
+    response = agent.stream_chat("What is 1 + 1?")
+    assert isinstance(response, StreamingAgentChatResponse)
+    # str() strips newline values
+    assert str(response) == "This is a test!"
+    assert len(agent.chat_history) == 2
+    assert agent.chat_history[0].content == "What is 1 + 1?"
+    assert agent.chat_history[1].content == "This is a test!"
+
+
+@patch("llama_index.llms.openai.base.AsyncOpenAI")
+@pytest.mark.asyncio()
+async def test_astream_chat_basic(
+    MockAsyncOpenAI: MagicMock, add_tool: FunctionTool
+) -> None:
+    mock_instance = MockAsyncOpenAI.return_value
+    mock_instance.chat.completions.create.side_effect = mock_achat_stream
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+
+    agent = OpenAIAgent.from_tools(
+        tools=[add_tool],
+        llm=llm,
+    )
+    response_stream = await agent.astream_chat("What is 1 + 1?")
+    async for response in response_stream.async_response_gen():
+        pass
+    assert isinstance(response_stream, StreamingAgentChatResponse)
+    # str() strips newline values
+    assert response == "\n\nThis is a test!"
+    assert len(agent.chat_history) == 2
+    assert agent.chat_history[0].content == "What is 1 + 1?"
+    assert agent.chat_history[1].content == "This is a test!"
+
+
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+def test_chat_no_functions(MockSyncOpenAI: MagicMock) -> None:
+    mock_instance = MockSyncOpenAI.return_value
+    mock_instance.chat.completions.create.return_value = mock_chat_completion()
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+
+    agent = OpenAIAgent.from_tools(
+        llm=llm,
+    )
+    response = agent.chat("What is 1 + 1?")
+    assert isinstance(response, AgentChatResponse)
+    assert response.response == "\n\nThis is a test!"
+
+
+def test_call_tool_with_error_handling() -> None:
+    """Test call tool with error handling."""
+
+    def _add(a: int, b: int) -> int:
+        return a + b
+
+    tool = FunctionTool.from_defaults(fn=_add)
+
+    output = call_tool_with_error_handling(
+        tool, {"a": 1, "b": 1}, error_message="Error!"
+    )
+    assert output.content == "2"
+
+    # try error
+    output = call_tool_with_error_handling(
+        tool, {"a": "1", "b": 1}, error_message="Error!"
+    )
+    assert output.content == "Error!"
+
+
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+def test_call_tool_with_malformed_function_call(
+    MockSyncOpenAI: MagicMock,
+    echo_tool: FunctionTool,
+) -> None:
+    """Test add step."""
+    mock_instance = MockSyncOpenAI.return_value
+    test_result: str = "This is a test"
+    function = Function(name="echo", arguments=f'query = "{test_result}"')
+    mock_instance.chat.completions.create.return_value = mock_chat_completion_tool_call(
+        function=function
+    )
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+    # sync
+    agent = OpenAIAgent.from_tools(
+        tools=[echo_tool],
+        llm=llm,
+    )
+    ## NOTE: can only take a single step before finishing,
+    # since mocked chat output does not call any tools
+    task = agent.create_task(
+        f"This happens if tool call is malformed like:\n{function.arguments}"
+    )
+    step_output = agent.run_step(task.task_id)
+    assert (
+        str(step_output.output.sources[0])
+        == 'Error in calling tool echo: The input json block is malformed:\n```json\nquery = "This is a test"\n```'
+    )
+
+
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+def test_call_tool_with_malformed_function_call_and_parser(
+    MockSyncOpenAI: MagicMock,
+    echo_tool: FunctionTool,
+) -> None:
+    """Test add step."""
+    mock_instance = MockSyncOpenAI.return_value
+    test_result: str = "This is a test"
+    function = Function(name="echo", arguments=f'query = "{test_result}"')
+    mock_instance.chat.completions.create.return_value = mock_chat_completion_tool_call(
+        function=function
+    )
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+    # sync
+    agent = OpenAIAgent.from_tools(
+        tools=[echo_tool],
+        llm=llm,
+        tool_call_parser=advanced_tool_call_parser,
+    )
+    ## NOTE: can only take a single step before finishing,
+    # since mocked chat output does not call any tools
+    task = agent.create_task(
+        f"This happens if tool call is malformed like:\n{function.arguments}"
+    )
+    step_output = agent.run_step(task.task_id)
+    assert str(step_output.output.sources[0]) == test_result
+
+
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+def test_add_step(
+    MockSyncOpenAI: MagicMock,
+    add_tool: FunctionTool,
+) -> None:
+    """Test add step."""
+    mock_instance = MockSyncOpenAI.return_value
+    mock_instance.chat.completions.create.return_value = mock_chat_completion()
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+    # sync
+    agent = OpenAIAgent.from_tools(
+        tools=[add_tool],
+        llm=llm,
+    )
+    ## NOTE: can only take a single step before finishing,
+    # since mocked chat output does not call any tools
+    task = agent.create_task("What is 1 + 1?")
+    step_output = agent.run_step(task.task_id)
+    assert str(step_output) == "\n\nThis is a test!"
+
+    # add human input (not used but should be in memory)
+    task = agent.create_task("What is 1 + 1?")
+    step_output = agent.run_step(task.task_id, input="tmp")
+    chat_history: List[ChatMessage] = task.extra_state["new_memory"].get_all()
+    assert "tmp" in [m.content for m in chat_history]
+
+    # # stream_step
+    # agent = OpenAIAgent.from_tools(
+    #     tools=[add_tool],
+    #     llm=llm,
+    # )
+    # task = agent.create_task("What is 1 + 1?")
+    # # first step
+    # step_output = agent.stream_step(task.task_id)
+    # # add human input (not used but should be in memory)
+    # step_output = agent.stream_step(task.task_id, input="tmp")
+    # chat_history: List[ChatMessage] = task.extra_state["new_memory"].get_all()
+    # assert "tmp" in [m.content for m in chat_history]
+
+
+@patch("llama_index.llms.openai.base.AsyncOpenAI")
+@pytest.mark.asyncio()
+async def test_async_add_step(
+    MockAsyncOpenAI: MagicMock,
+    add_tool: FunctionTool,
+) -> None:
+    mock_instance = MockAsyncOpenAI.return_value
+
+    llm = OpenAI(model="gpt-3.5-turbo")
+    # async
+    agent = OpenAIAgent.from_tools(
+        tools=[add_tool],
+        llm=llm,
+    )
+    task = agent.create_task("What is 1 + 1?")
+    # first step
+    mock_instance.chat.completions.create.return_value = mock_achat_completion()
+    step_output = await agent.arun_step(task.task_id)
+    # add human input (not used but should be in memory)
+    task = agent.create_task("What is 1 + 1?")
+    mock_instance.chat.completions.create.return_value = mock_achat_completion()
+    step_output = await agent.arun_step(task.task_id, input="tmp")
+    chat_history: List[ChatMessage] = task.extra_state["new_memory"].get_all()
+    assert "tmp" in [m.content for m in chat_history]
+
+    # async stream step
+    agent = OpenAIAgent.from_tools(
+        tools=[add_tool],
+        llm=llm,
+    )
+    task = agent.create_task("What is 1 + 1?")
+    # first step
+    mock_instance.chat.completions.create.side_effect = mock_achat_stream
+    step_output = await agent.astream_step(task.task_id)
+    # add human input (not used but should be in memory)
+    task = agent.create_task("What is 1 + 1?")
+    mock_instance.chat.completions.create.side_effect = mock_achat_stream
+
+    # stream the output to ensure it gets written to memory
+    step_output = await agent.astream_step(task.task_id, input="tmp")
+    async for _ in step_output.output.async_response_gen():
+        pass
+
+    chat_history = task.memory.get_all()
+    assert "tmp" in [m.content for m in chat_history]


Code review comments:
path: docs/examples/agent/openai_agent_tool_call_parser.ipynb
text: @user1:
Do you mind rerunning this? I think there are some debug prints (i.e. the `NOT GIVEN` stuff, and others), makes it a little confusing to read the outputs

@author:
No problem at all

diff_hunk: @@ -446,6 +549,12 @@ def _run_step(
         openai_tools = [tool.metadata.to_openai_tool() for tool in tools]
 
         llm_chat_kwargs = self._get_llm_chat_kwargs(task, openai_tools, tool_choice)
+        for tool in tools:  # self.agent_worker._get_tools("test"):
path: llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
text: @user1:
Are these debug prints? Can we remove them?

@author:
Indeed, removing them in the next commit