fix!: Improve response handling from Toolbox server (#69)

anubhav756 · web-flow · commit f4935c09592e · 2025-03-22T01:01:06.000+05:30
* fix: Make all fields required in Tool schema.

Earlier we made all fields as optional since we wanted to keep some fields optional for the LLM. Since Toolbox did not support optional fields, there was no way to know which fields were optional, so as a worst-case, we did a temporary workaround of keeping all fields as optional in the schema generated by Toolbox SDK.

Now, there has been some evidence that the LLMs do not work very well with optional parameters, and so we have decided not to support optional fields for now, neither in Toolbox service nor in the SDK.

This PR removes that temporary fix of making all the fields optional.

This PR also removes an augmentation to the request body where `None` values were converted to empty strings (`''`). This is because now that LLM knows no fields are optional, we can be sure that we would not be getting any `None` values as inputs to the tools. So the function `_convert_none_to_empty_string` is not required anymore.

* chore: Update unit tests.

* fix!: Improve response handling from Toolbox server.

* We remove `response.raise_for_status()` as it masks error reasons thrown by Toolbox server.

* We return the value of `result` key from the response body, so that it can directly be fed to LLMs.
  * This would also prevent situations where the response is `{ "result": '{ "some": "value" }' }` and when we feed that to LLM by stringifying, it becomes something like `"{ "result": '{ \"some\": \"value\" }' }"` which is more cryptic for the LLM because of the extra `\` characters due to double stringification.

* We also check for `error` in the response and throw a `ToolException` with the response if applicable.

* chore: Update test cases.

* fix: Make all fields required in Tool schema.

Earlier we made all fields as optional since we wanted to keep some fields optional for the LLM. Since Toolbox did not support optional fields, there was no way to know which fields were optional, so as a worst-case, we did a temporary workaround of keeping all fields as optional in the schema generated by Toolbox SDK.

Now, there has been some evidence that the LLMs do not work very well with optional parameters, and so we have decided not to support optional fields for now, neither in Toolbox service nor in the SDK.

This PR removes that temporary fix of making all the fields optional.

This PR also removes an augmentation to the request body where `None` values were converted to empty strings (`''`). This is because now that LLM knows no fields are optional, we can be sure that we would not be getting any `None` values as inputs to the tools. So the function `_convert_none_to_empty_string` is not required anymore.

* chore: Update unit tests.

* fix: Make all fields required in Tool schema.

Earlier we made all fields as optional since we wanted to keep some fields optional for the LLM. Since Toolbox did not support optional fields, there was no way to know which fields were optional, so as a worst-case, we did a temporary workaround of keeping all fields as optional in the schema generated by Toolbox SDK.

Now, there has been some evidence that the LLMs do not work very well with optional parameters, and so we have decided not to support optional fields for now, neither in Toolbox service nor in the SDK.

This PR removes that temporary fix of making all the fields optional.

This PR also removes an augmentation to the request body where `None` values were converted to empty strings (`''`). This is because now that LLM knows no fields are optional, we can be sure that we would not be getting any `None` values as inputs to the tools. So the function `_convert_none_to_empty_string` is not required anymore.

* chore: Update unit tests.

* fix: Make all fields required in Tool schema.

Earlier we made all fields as optional since we wanted to keep some fields optional for the LLM. Since Toolbox did not support optional fields, there was no way to know which fields were optional, so as a worst-case, we did a temporary workaround of keeping all fields as optional in the schema generated by Toolbox SDK.

Now, there has been some evidence that the LLMs do not work very well with optional parameters, and so we have decided not to support optional fields for now, neither in Toolbox service nor in the SDK.

This PR removes that temporary fix of making all the fields optional.

This PR also removes an augmentation to the request body where `None` values were converted to empty strings (`''`). This is because now that LLM knows no fields are optional, we can be sure that we would not be getting any `None` values as inputs to the tools. So the function `_convert_none_to_empty_string` is not required anymore.

* chore: Update unit tests.
diff --git a/src/toolbox_langchain/utils.py b/src/toolbox_langchain/utils.py
@@ -18,6 +18,7 @@
 
 from aiohttp import ClientSession
 from deprecated import deprecated
+from langchain_core.tools import ToolException
 from pydantic import BaseModel, Field, create_model
 
 
@@ -187,6 +188,9 @@ async def _invoke_tool(
     Returns:
         A dictionary containing the parsed JSON response from the tool
         invocation.
+
+    Raises:
+        ToolException: If the Toolbox service returns an error.
     """
     url = f"{url}/api/tool/{tool_name}/invoke"
     auth_tokens = _get_auth_tokens(id_token_getters)
@@ -204,9 +208,10 @@ async def _invoke_tool(
         json=data,
         headers=auth_tokens,
     ) as response:
-        # TODO: Remove as it masks error messages.
-        response.raise_for_status()
-        return await response.json()
+        ret = await response.json()
+        if "error" in ret:
+            raise ToolException(ret)
+        return ret.get("result", ret)
 
 
 def _find_auth_params(
diff --git a/tests/test_async_tools.py b/tests/test_async_tools.py
@@ -196,7 +196,7 @@ async def test_toolbox_tool_validate_auth_strict(self, auth_toolbox_tool):
 
     async def test_toolbox_tool_call(self, toolbox_tool):
         result = await toolbox_tool.ainvoke({"param1": "test-value", "param2": 123})
-        assert result == {"result": "test-result"}
+        assert result == "test-result"
         toolbox_tool._AsyncToolboxTool__session.post.assert_called_once_with(
             "http://test_url/api/tool/test_tool/invoke",
             json={"param1": "test-value", "param2": 123},
@@ -215,7 +215,7 @@ async def test_toolbox_tool_call_with_bound_params(
     ):
         tool = toolbox_tool.bind_params(bound_param)
         result = await tool.ainvoke({"param2": 123})
-        assert result == {"result": "test-result"}
+        assert result == "test-result"
         toolbox_tool._AsyncToolboxTool__session.post.assert_called_once_with(
             "http://test_url/api/tool/test_tool/invoke",
             json={"param1": expected_value, "param2": 123},
@@ -227,7 +227,7 @@ async def test_toolbox_tool_call_with_auth_tokens(self, auth_toolbox_tool):
             {"test-auth-source": lambda: "test-token"}
         )
         result = await tool.ainvoke({"param2": 123})
-        assert result == {"result": "test-result"}
+        assert result == "test-result"
         auth_toolbox_tool._AsyncToolboxTool__session.post.assert_called_once_with(
             "https://test-url/api/tool/test_tool/invoke",
             json={"param2": 123},
@@ -244,7 +244,7 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure(self, auth_toolbox_to
                 {"test-auth-source": lambda: "test-token"}
             )
             result = await tool.ainvoke({"param2": 123})
-            assert result == {"result": "test-result"}
+            assert result == "test-result"
             auth_toolbox_tool._AsyncToolboxTool__session.post.assert_called_once_with(
                 "http://test-url/api/tool/test_tool/invoke",
                 json={"param2": 123},
diff --git a/tests/test_e2e.py b/tests/test_e2e.py
@@ -36,7 +36,7 @@
 
 import pytest
 import pytest_asyncio
-from aiohttp import ClientResponseError
+from langchain_core.tools import ToolException
 from pydantic import ValidationError
 
 from toolbox_langchain.client import ToolboxClient
@@ -90,19 +90,17 @@ async def test_aload_toolset_all(self, toolbox):
 
     async def test_run_tool_async(self, get_n_rows_tool):
         response = await get_n_rows_tool.ainvoke({"num_rows": "2"})
-        result = response["result"]
 
-        assert "row1" in result
-        assert "row2" in result
-        assert "row3" not in result
+        assert "row1" in response
+        assert "row2" in response
+        assert "row3" not in response
 
     async def test_run_tool_sync(self, get_n_rows_tool):
         response = get_n_rows_tool.invoke({"num_rows": "2"})
-        result = response["result"]
 
-        assert "row1" in result
-        assert "row2" in result
-        assert "row3" not in result
+        assert "row1" in response
+        assert "row2" in response
+        assert "row3" not in response
 
     async def test_run_tool_missing_params(self, get_n_rows_tool):
         with pytest.raises(ValidationError, match="Field required"):
@@ -120,14 +118,17 @@ async def test_run_tool_unauth_with_auth(self, toolbox, auth_token2):
             "get-row-by-id", auth_tokens={"my-test-auth": lambda: auth_token2}
         )
         response = await tool.ainvoke({"id": "2"})
-        assert "row2" in response["result"]
+        assert "row2" in response
 
     async def test_run_tool_no_auth(self, toolbox):
         """Tests running a tool requiring auth without providing auth."""
         tool = await toolbox.aload_tool(
             "get-row-by-id-auth",
         )
-        with pytest.raises(ClientResponseError, match="401, message='Unauthorized'"):
+        with pytest.raises(
+            ToolException,
+            match="{'status': 'Unauthorized', 'error': 'tool invocation not authorized. Please make sure your specify correct auth headers'}",
+        ):
             await tool.ainvoke({"id": "2"})
 
     async def test_run_tool_wrong_auth(self, toolbox, auth_token2):
@@ -136,7 +137,10 @@ async def test_run_tool_wrong_auth(self, toolbox, auth_token2):
             "get-row-by-id-auth",
         )
         auth_tool = tool.add_auth_token("my-test-auth", lambda: auth_token2)
-        with pytest.raises(ClientResponseError, match="401, message='Unauthorized'"):
+        with pytest.raises(
+            ToolException,
+            match="{'status': 'Unauthorized', 'error': 'tool invocation not authorized. Please make sure your specify correct auth headers'}",
+        ):
             await auth_tool.ainvoke({"id": "2"})
 
     async def test_run_tool_auth(self, toolbox, auth_token1):
@@ -146,7 +150,7 @@ async def test_run_tool_auth(self, toolbox, auth_token1):
         )
         auth_tool = tool.add_auth_token("my-test-auth", lambda: auth_token1)
         response = await auth_tool.ainvoke({"id": "2"})
-        assert "row2" in response["result"]
+        assert "row2" in response
 
     async def test_run_tool_param_auth_no_auth(self, toolbox):
         """Tests running a tool with a param requiring auth, without auth."""
@@ -163,17 +167,19 @@ async def test_run_tool_param_auth(self, toolbox, auth_token1):
             "get-row-by-email-auth", auth_tokens={"my-test-auth": lambda: auth_token1}
         )
         response = await tool.ainvoke({})
-        result = response["result"]
-        assert "row4" in result
-        assert "row5" in result
-        assert "row6" in result
+        assert "row4" in response
+        assert "row5" in response
+        assert "row6" in response
 
     async def test_run_tool_param_auth_no_field(self, toolbox, auth_token1):
         """Tests running a tool with a param requiring auth, with insufficient auth."""
         tool = await toolbox.aload_tool(
             "get-row-by-content-auth", auth_tokens={"my-test-auth": lambda: auth_token1}
         )
-        with pytest.raises(ClientResponseError, match="400, message='Bad Request'"):
+        with pytest.raises(
+            ToolException,
+            match="{'status': 'Bad Request', 'error': 'provided parameters were invalid: error parsing authenticated parameter \"data\": no field named row_data in claims'}",
+        ):
             await tool.ainvoke({})
 
 
@@ -225,19 +231,17 @@ def test_aload_toolset_all(self, toolbox):
     @pytest.mark.asyncio
     async def test_run_tool_async(self, get_n_rows_tool):
         response = await get_n_rows_tool.ainvoke({"num_rows": "2"})
-        result = response["result"]
 
-        assert "row1" in result
-        assert "row2" in result
-        assert "row3" not in result
+        assert "row1" in response
+        assert "row2" in response
+        assert "row3" not in response
 
     def test_run_tool_sync(self, get_n_rows_tool):
         response = get_n_rows_tool.invoke({"num_rows": "2"})
-        result = response["result"]
 
-        assert "row1" in result
-        assert "row2" in result
-        assert "row3" not in result
+        assert "row1" in response
+        assert "row2" in response
+        assert "row3" not in response
 
     def test_run_tool_missing_params(self, get_n_rows_tool):
         with pytest.raises(ValidationError, match="Field required"):
@@ -254,14 +258,17 @@ def test_run_tool_unauth_with_auth(self, toolbox, auth_token2):
             "get-row-by-id", auth_tokens={"my-test-auth": lambda: auth_token2}
         )
         response = tool.invoke({"id": "2"})
-        assert "row2" in response["result"]
+        assert "row2" in response
 
     def test_run_tool_no_auth(self, toolbox):
         """Tests running a tool requiring auth without providing auth."""
         tool = toolbox.load_tool(
             "get-row-by-id-auth",
         )
-        with pytest.raises(ClientResponseError, match="401, message='Unauthorized'"):
+        with pytest.raises(
+            ToolException,
+            match="{'status': 'Unauthorized', 'error': 'tool invocation not authorized. Please make sure your specify correct auth headers'}",
+        ):
             tool.invoke({"id": "2"})
 
     def test_run_tool_wrong_auth(self, toolbox, auth_token2):
@@ -270,7 +277,10 @@ def test_run_tool_wrong_auth(self, toolbox, auth_token2):
             "get-row-by-id-auth",
         )
         auth_tool = tool.add_auth_token("my-test-auth", lambda: auth_token2)
-        with pytest.raises(ClientResponseError, match="401, message='Unauthorized'"):
+        with pytest.raises(
+            ToolException,
+            match="{'status': 'Unauthorized', 'error': 'tool invocation not authorized. Please make sure your specify correct auth headers'}",
+        ):
             auth_tool.invoke({"id": "2"})
 
     def test_run_tool_auth(self, toolbox, auth_token1):
@@ -280,7 +290,7 @@ def test_run_tool_auth(self, toolbox, auth_token1):
         )
         auth_tool = tool.add_auth_token("my-test-auth", lambda: auth_token1)
         response = auth_tool.invoke({"id": "2"})
-        assert "row2" in response["result"]
+        assert "row2" in response
 
     def test_run_tool_param_auth_no_auth(self, toolbox):
         """Tests running a tool with a param requiring auth, without auth."""
@@ -297,15 +307,17 @@ def test_run_tool_param_auth(self, toolbox, auth_token1):
             "get-row-by-email-auth", auth_tokens={"my-test-auth": lambda: auth_token1}
         )
         response = tool.invoke({})
-        result = response["result"]
-        assert "row4" in result
-        assert "row5" in result
-        assert "row6" in result
+        assert "row4" in response
+        assert "row5" in response
+        assert "row6" in response
 
     def test_run_tool_param_auth_no_field(self, toolbox, auth_token1):
         """Tests running a tool with a param requiring auth, with insufficient auth."""
         tool = toolbox.load_tool(
             "get-row-by-content-auth", auth_tokens={"my-test-auth": lambda: auth_token1}
         )
-        with pytest.raises(ClientResponseError, match="400, message='Bad Request'"):
+        with pytest.raises(
+            ToolException,
+            match="{'status': 'Bad Request', 'error': 'provided parameters were invalid: error parsing authenticated parameter \"data\": no field named row_data in claims'}",
+        ):
             tool.invoke({})
