chore: Fix unit tests

Author: anubhav756

packages/toolbox-llamaindex/tests/test_async_tools.py

@@ -17,6 +17,7 @@
 
 import pytest
 import pytest_asyncio
+from llama_index.core.tools.types import ToolOutput
 from pydantic import ValidationError
 from toolbox_core.protocol import ParameterSchema as CoreParameterSchema
 from toolbox_core.tool import ToolboxTool as ToolboxCoreTool
@@ -261,7 +262,14 @@ async def test_toolbox_tool_call_requires_auth_strict(self, auth_toolbox_tool):
 
     async def test_toolbox_tool_call(self, toolbox_tool):
         result = await toolbox_tool.acall(param1="test-value", param2=123)
-        assert result == "test-result"
+        assert result == ToolOutput(
+            content="test-result",
+            tool_name="test_tool",
+            raw_input={"param1": "test-value", "param2": 123},
+            raw_output="test-result",
+            is_error=False,
+        )
+
         core_tool = toolbox_tool._AsyncToolboxTool__core_tool
         core_tool._ToolboxTool__session.post.assert_called_once_with(
             "http://test_url/api/tool/test_tool/invoke",
@@ -281,7 +289,13 @@ async def test_toolbox_tool_call_with_bound_params(
     ):
         tool = toolbox_tool.bind_params(bound_param_map)
         result = await tool.acall(param2=123)
-        assert result == "test-result"
+        assert result == ToolOutput(
+            content="test-result",
+            tool_name="test_tool",
+            raw_input={"param2": 123},
+            raw_output="test-result",
+            is_error=False,
+        )
         core_tool = tool._AsyncToolboxTool__core_tool
         core_tool._ToolboxTool__session.post.assert_called_once_with(
             "http://test_url/api/tool/test_tool/invoke",
@@ -294,7 +308,14 @@ async def test_toolbox_tool_call_with_auth_tokens(self, auth_toolbox_tool):
             {"test-auth-source": lambda: "test-token"}
         )
         result = await tool.acall(param2=123)
-        assert result == "test-result"
+        assert result == ToolOutput(
+            content="test-result",
+            tool_name="test_tool",
+            raw_input={"param2": 123},
+            raw_output="test-result",
+            is_error=False,
+        )
+
         core_tool = tool._AsyncToolboxTool__core_tool
         core_tool._ToolboxTool__session.post.assert_called_once_with(
             "https://test-url/api/tool/test_tool/invoke",
@@ -325,7 +346,13 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure(
             {"test-auth-source": lambda: "test-token"}
         )
         result = await tool_with_getter.acall(param2=123)
-        assert result == "test-result"
+        assert result == ToolOutput(
+            content="test-result",
+            tool_name="test_tool",
+            raw_input={"param2": 123},
+            raw_output="test-result",
+            is_error=False,
+        )
 
         modified_core_tool_in_new_tool = tool_with_getter._AsyncToolboxTool__core_tool
         assert (

packages/toolbox-llamaindex/tests/test_tools.py

@@ -16,6 +16,7 @@
 from unittest.mock import AsyncMock, Mock, call, patch
 
 import pytest
+from llama_index.core.tools.types import ToolOutput
 from pydantic import BaseModel
 from toolbox_core.protocol import ParameterSchema as CoreParameterSchema
 from toolbox_core.sync_tool import ToolboxSyncTool as ToolboxCoreSyncTool
@@ -271,8 +272,14 @@ def test_toolbox_tool_add_auth_token_getter(
 
     def test_toolbox_tool_run(self, toolbox_tool, mock_core_tool):
         kwargs_to_run = {"param1": "run_value1", "param2": 100}
-        expected_result = "sync_run_output"
-        mock_core_tool.return_value = expected_result
+        expected_result = ToolOutput(
+            content="sync_run_output",
+            tool_name="test_tool_name_for_llamaindex",
+            raw_input=kwargs_to_run,
+            raw_output="sync_run_output",
+            is_error=False,
+        )
+        mock_core_tool.return_value = "sync_run_output"
 
         result = toolbox_tool.call(**kwargs_to_run)
 
@@ -286,9 +293,15 @@ async def test_toolbox_tool_arun(
         self, mock_to_thread_in_tools, toolbox_tool, mock_core_tool
     ):
         kwargs_to_run = {"param1": "arun_value1", "param2": 200}
-        expected_result = "async_run_output"
+        expected_result = ToolOutput(
+            content="async_run_output",
+            tool_name="test_tool_name_for_llamaindex",
+            raw_input=kwargs_to_run,
+            raw_output="async_run_output",
+            is_error=False,
+        )
 
-        mock_core_tool.return_value = expected_result
+        mock_core_tool.return_value = "async_run_output"
 
         async def to_thread_side_effect(func, *args, **kwargs_for_func):
             return func(**kwargs_for_func)