feat(langchain-sdk): Correctly parse Manifest API response as JSON (#143)

The Manifest API returns a JSON payload. Previously, it was parsed as
YAML, which worked due to YAML's superset relationship with JSON. This
change explicitly parses the response as JSON for improved robustness
and security by ensuring strict adherence to the expected format.

Author: anubhav756

src/toolbox_langchain_sdk/client.py

@@ -6,7 +6,7 @@
 from langchain_core.tools import StructuredTool
 from pydantic import BaseModel
 
-from .utils import ManifestSchema, _invoke_tool, _load_yaml, _schema_to_model
+from .utils import ManifestSchema, _invoke_tool, _load_manifest, _schema_to_model
 
 
 class ToolboxClient:
@@ -51,8 +51,8 @@ def __del__(self):
 
     async def _load_tool_manifest(self, tool_name: str) -> ManifestSchema:
         """
-        Fetches and parses the YAML manifest for the given tool from the Toolbox
-        service.
+        Fetches and parses the manifest schema for the given tool from the
+        Toolbox service.
 
         Args:
             tool_name: The name of the tool to load.
@@ -61,13 +61,13 @@ async def _load_tool_manifest(self, tool_name: str) -> ManifestSchema:
             The parsed Toolbox manifest.
         """
         url = f"{self._url}/api/tool/{tool_name}"
-        return await _load_yaml(url, self._session)
+        return await _load_manifest(url, self._session)
 
     async def _load_toolset_manifest(
         self, toolset_name: Optional[str] = None
     ) -> ManifestSchema:
         """
-        Fetches and parses the YAML manifest from the Toolbox service.
+        Fetches and parses the manifest schema from the Toolbox service.
 
         Args:
             toolset_name: The name of the toolset to load.
@@ -78,7 +78,7 @@ async def _load_toolset_manifest(
             The parsed Toolbox manifest.
         """
         url = f"{self._url}/api/toolset/{toolset_name or ''}"
-        return await _load_yaml(url, self._session)
+        return await _load_manifest(url, self._session)
 
     def _validate_auth(self, tool_name: str) -> bool:
         """

src/toolbox_langchain_sdk/utils.py

@@ -1,7 +1,7 @@
+import json
 import warnings
 from typing import Any, Callable, Optional, Type, cast
 
-import yaml
 from aiohttp import ClientSession
 from pydantic import BaseModel, Field, create_model
 
@@ -23,12 +23,13 @@ class ManifestSchema(BaseModel):
     tools: dict[str, ToolSchema]
 
 
-async def _load_yaml(url: str, session: ClientSession) -> ManifestSchema:
+async def _load_manifest(url: str, session: ClientSession) -> ManifestSchema:
     """
-    Asynchronously fetches and parses the YAML data from the given URL.
+    Asynchronously fetches and parses the JSON manifest schema from the given
+    URL.
 
     Args:
-        url: The base URL to fetch the YAML from.
+        url: The base URL to fetch the JSON from.
         session: The HTTP client session
 
     Returns:
@@ -37,18 +38,20 @@ async def _load_yaml(url: str, session: ClientSession) -> ManifestSchema:
     async with session.get(url) as response:
         response.raise_for_status()
         try:
-            parsed_yaml = yaml.safe_load(await response.text())
-        except yaml.YAMLError as e:
-            raise yaml.YAMLError(f"Failed to parse YAML from {url}: {e}") from e
+            parsed_json = json.loads(await response.text())
+        except json.JSONDecodeError as e:
+            raise json.JSONDecodeError(
+                f"Failed to parse JSON from {url}: {e}", e.doc, e.pos
+            ) from e
         try:
-            return ManifestSchema(**parsed_yaml)
+            return ManifestSchema(**parsed_json)
         except ValueError as e:
-            raise ValueError(f"Invalid YAML data from {url}: {e}") from e
+            raise ValueError(f"Invalid JSON data from {url}: {e}") from e
 
 
 def _schema_to_model(model_name: str, schema: list[ParameterSchema]) -> Type[BaseModel]:
     """
-    Converts a schema (from the YAML manifest) to a Pydantic BaseModel class.
+    Converts the given manifest schema to a Pydantic BaseModel class.
 
     Args:
         model_name: The name of the model to create.

tests/test_client.py

@@ -71,60 +71,60 @@ async def test_close_not_closing_session():
 
 
 @pytest.mark.asyncio
-@patch("toolbox_langchain_sdk.client._load_yaml")
-async def test_load_tool_manifest_success(mock_load_yaml):
+@patch("toolbox_langchain_sdk.client._load_manifest")
+async def test_load_tool_manifest_success(mock_load_manifest):
     client = ToolboxClient("https://my-toolbox.com", session=aiohttp.ClientSession())
-    mock_load_yaml.return_value = ManifestSchema(**manifest_data)
+    mock_load_manifest.return_value = ManifestSchema(**manifest_data)
 
     result = await client._load_tool_manifest("test_tool")
     assert result == ManifestSchema(**manifest_data)
-    mock_load_yaml.assert_called_once_with(
+    mock_load_manifest.assert_called_once_with(
         "https://my-toolbox.com/api/tool/test_tool", client._session
     )
 
 
 @pytest.mark.asyncio
-@patch("toolbox_langchain_sdk.client._load_yaml")
-async def test_load_tool_manifest_failure(mock_load_yaml):
+@patch("toolbox_langchain_sdk.client._load_manifest")
+async def test_load_tool_manifest_failure(mock_load_manifest):
     client = ToolboxClient("https://my-toolbox.com", session=aiohttp.ClientSession())
-    mock_load_yaml.side_effect = Exception("Failed to load YAML")
+    mock_load_manifest.side_effect = Exception("Failed to load manifest")
 
     with pytest.raises(Exception) as e:
         await client._load_tool_manifest("test_tool")
-    assert str(e.value) == "Failed to load YAML"
+    assert str(e.value) == "Failed to load manifest"
 
 
 @pytest.mark.asyncio
-@patch("toolbox_langchain_sdk.client._load_yaml")
-async def test_load_toolset_manifest_success(mock_load_yaml):
+@patch("toolbox_langchain_sdk.client._load_manifest")
+async def test_load_toolset_manifest_success(mock_load_manifest):
     client = ToolboxClient("https://my-toolbox.com", session=aiohttp.ClientSession())
-    mock_load_yaml.return_value = ManifestSchema(**manifest_data)
+    mock_load_manifest.return_value = ManifestSchema(**manifest_data)
 
     # Test with toolset name
     result = await client._load_toolset_manifest(toolset_name="test_toolset")
     assert result == ManifestSchema(**manifest_data)
-    mock_load_yaml.assert_called_once_with(
+    mock_load_manifest.assert_called_once_with(
         "https://my-toolbox.com/api/toolset/test_toolset", client._session
     )
-    mock_load_yaml.reset_mock()
+    mock_load_manifest.reset_mock()
 
     # Test without toolset name
     result = await client._load_toolset_manifest()
     assert result == ManifestSchema(**manifest_data)
-    mock_load_yaml.assert_called_once_with(
+    mock_load_manifest.assert_called_once_with(
         "https://my-toolbox.com/api/toolset/", client._session
     )
 
 
 @pytest.mark.asyncio
-@patch("toolbox_langchain_sdk.client._load_yaml")
-async def test_load_toolset_manifest_failure(mock_load_yaml):
+@patch("toolbox_langchain_sdk.client._load_manifest")
+async def test_load_toolset_manifest_failure(mock_load_manifest):
     client = ToolboxClient("https://my-toolbox.com", session=aiohttp.ClientSession())
-    mock_load_yaml.side_effect = Exception("Failed to load YAML")
+    mock_load_manifest.side_effect = Exception("Failed to load manifest")
 
     with pytest.raises(Exception) as e:
         await client._load_toolset_manifest(toolset_name="test_toolset")
-    assert str(e.value) == "Failed to load YAML"
+    assert str(e.value) == "Failed to load manifest"
 
 
 @pytest.mark.asyncio
@@ -541,7 +541,7 @@ async def test_process_auth_params(
 
 
 @pytest.mark.asyncio
-@patch("toolbox_langchain_sdk.client._load_yaml")
+@patch("toolbox_langchain_sdk.client._load_manifest")
 @pytest.mark.parametrize(
     "params, auth_headers, expected_tool_param_auth",
     [
@@ -569,13 +569,13 @@ async def test_process_auth_params(
     ],
 )
 async def test_load_tool(
-    mock_load_yaml, params, auth_headers, expected_tool_param_auth
+    mock_load_manifest, params, auth_headers, expected_tool_param_auth
 ):
     """Test load_tool with and without auth headers."""
     client = ToolboxClient("http://test-url")
 
     # Replace with your desired mock manifest data
-    mock_load_yaml.return_value = ManifestSchema(
+    mock_load_manifest.return_value = ManifestSchema(
         serverVersion="1.0",
         tools={
             "tool_name": ToolSchema(
@@ -594,7 +594,7 @@ async def test_load_tool(
 
 
 @pytest.mark.asyncio
-@patch("toolbox_langchain_sdk.client._load_yaml")
+@patch("toolbox_langchain_sdk.client._load_manifest")
 @pytest.mark.parametrize(
     "params, auth_headers, expected_tool_param_auth, expected_num_tools",
     [
@@ -624,7 +624,7 @@ async def test_load_tool(
     ],
 )
 async def test_load_toolset(
-    mock_load_yaml,
+    mock_load_manifest,
     params,
     auth_headers,
     expected_tool_param_auth,
@@ -634,7 +634,7 @@ async def test_load_toolset(
     client = ToolboxClient("http://test-url")
 
     # Replace with your desired mock manifest data
-    mock_load_yaml.return_value = ManifestSchema(
+    mock_load_manifest.return_value = ManifestSchema(
         serverVersion="1.0",
         tools={
             "tool_name": ToolSchema(

tests/test_utils.py

@@ -1,43 +1,52 @@
 import asyncio
+import json
 import warnings
 from typing import Union
 from unittest.mock import AsyncMock, Mock, patch
 
 import aiohttp
 import pytest
-import yaml
 from aiohttp import ClientSession
 from pydantic import BaseModel
 
 from toolbox_langchain_sdk.utils import (
     ParameterSchema,
     _convert_none_to_empty_string,
     _invoke_tool,
-    _load_yaml,
+    _load_manifest,
     _parse_type,
     _schema_to_model,
 )
 
 URL = "https://my-toolbox.com/test"
 MOCK_MANIFEST = """
-serverVersion: 0.0.1
-tools:
-    test_tool:
-        summary: Test Tool
-        description: This is a test tool.
-        parameters:
-          - name: param1
-            type: string
-            description: Parameter 1
-          - name: param2
-            type: integer
-            description: Parameter 2
+{
+  "serverVersion": "0.0.1",
+  "tools": {
+    "test_tool": {
+      "summary": "Test Tool",
+      "description": "This is a test tool.",
+      "parameters": [
+        {
+          "name": "param1",
+          "type": "string",
+          "description": "Parameter 1"
+        },
+        {
+          "name": "param2",
+          "type": "integer",
+          "description": "Parameter 2"
+        }
+      ]
+    }
+  }
+}
 """
 
 
 class TestUtils:
     @pytest.fixture(scope="module")
-    def mock_yaml(self):
+    def mock_manifest(self):
         return aiohttp.ClientResponse(
             method="GET",
             url=aiohttp.client.URL(URL),
@@ -52,13 +61,13 @@ def mock_yaml(self):
 
     @pytest.mark.asyncio
     @patch("aiohttp.ClientSession.get")
-    async def test_load_yaml(self, mock_get, mock_yaml):
-        mock_yaml.raise_for_status = Mock()
-        mock_yaml.text = AsyncMock(return_value=MOCK_MANIFEST)
+    async def test_load_manifest(self, mock_get, mock_manifest):
+        mock_manifest.raise_for_status = Mock()
+        mock_manifest.text = AsyncMock(return_value=MOCK_MANIFEST)
 
-        mock_get.return_value = mock_yaml
+        mock_get.return_value = mock_manifest
         session = aiohttp.ClientSession()
-        manifest = await _load_yaml(URL, session)
+        manifest = await _load_manifest(URL, session)
         await session.close()
         mock_get.assert_called_once_with(URL)
 
@@ -74,55 +83,55 @@ async def test_load_yaml(self, mock_get, mock_yaml):
 
     @pytest.mark.asyncio
     @patch("aiohttp.ClientSession.get")
-    async def test_load_yaml_invalid_yaml(self, mock_get, mock_yaml):
-        mock_yaml.raise_for_status = Mock()
-        mock_yaml.text = AsyncMock(return_value="{ invalid yaml")
-        mock_get.return_value = mock_yaml
+    async def test_load_manifest_invalid_manifest(self, mock_get, mock_manifest):
+        mock_manifest.raise_for_status = Mock()
+        mock_manifest.text = AsyncMock(return_value="{ invalid manifest")
+        mock_get.return_value = mock_manifest
 
         with pytest.raises(Exception) as e:
             session = aiohttp.ClientSession()
-            await _load_yaml(URL, session)
+            await _load_manifest(URL, session)
             await session.close()
             mock_get.assert_called_once_with(URL)
 
         mock_get.assert_called_once_with(URL)
-        assert isinstance(e.value, yaml.YAMLError)
+        assert isinstance(e.value, json.JSONDecodeError)
         assert (
             str(e.value)
-            == "Failed to parse YAML from https://my-toolbox.com/test: while parsing a flow mapping\n  in \"<unicode string>\", line 1, column 1:\n    { invalid yaml\n    ^\nexpected ',' or '}', but got '<stream end>'\n  in \"<unicode string>\", line 1, column 15:\n    { invalid yaml\n                  ^"
+            == "Failed to parse JSON from https://my-toolbox.com/test: Expecting property name enclosed in double quotes: line 1 column 3 (char 2): line 1 column 3 (char 2)"
         )
 
     @pytest.mark.asyncio
     @patch("aiohttp.ClientSession.get")
-    async def test_load_yaml_invalid_manifest(self, mock_get, mock_yaml):
-        mock_yaml.raise_for_status = Mock()
-        mock_yaml.text = AsyncMock(return_value="{ invalid yaml }")
-        mock_get.return_value = mock_yaml
+    async def test_load_manifest_invalid_manifest(self, mock_get, mock_manifest):
+        mock_manifest.raise_for_status = Mock()
+        mock_manifest.text = AsyncMock(return_value='{ "something": "invalid" }')
+        mock_get.return_value = mock_manifest
 
         with pytest.raises(Exception) as e:
             session = aiohttp.ClientSession()
-            await _load_yaml(URL, session)
+            await _load_manifest(URL, session)
             await session.close()
             mock_get.assert_called_once_with(URL)
 
         mock_get.assert_called_once_with(URL)
         assert isinstance(e.value, ValueError)
         assert (
             str(e.value)
-            == "Invalid YAML data from https://my-toolbox.com/test: 2 validation errors for ManifestSchema\nserverVersion\n  Field required [type=missing, input_value={'invalid yaml': None}, input_type=dict]\n    For further information visit https://errors.pydantic.dev/2.10/v/missing\ntools\n  Field required [type=missing, input_value={'invalid yaml': None}, input_type=dict]\n    For further information visit https://errors.pydantic.dev/2.10/v/missing"
+            == "Invalid JSON data from https://my-toolbox.com/test: 2 validation errors for ManifestSchema\nserverVersion\n  Field required [type=missing, input_value={'something': 'invalid'}, input_type=dict]\n    For further information visit https://errors.pydantic.dev/2.10/v/missing\ntools\n  Field required [type=missing, input_value={'something': 'invalid'}, input_type=dict]\n    For further information visit https://errors.pydantic.dev/2.10/v/missing"
         )
 
     @pytest.mark.asyncio
     @patch("aiohttp.ClientSession.get")
-    async def test_load_yaml_api_error(self, mock_get, mock_yaml):
+    async def test_load_manifest_api_error(self, mock_get, mock_manifest):
         error = aiohttp.ClientError("Simulated HTTP Error")
-        mock_yaml.raise_for_status = Mock()
-        mock_yaml.text = AsyncMock(side_effect=error)
-        mock_get.return_value = mock_yaml
+        mock_manifest.raise_for_status = Mock()
+        mock_manifest.text = AsyncMock(side_effect=error)
+        mock_get.return_value = mock_manifest
 
         with pytest.raises(aiohttp.ClientError) as exc_info:
             session = aiohttp.ClientSession()
-            await _load_yaml(URL, session)
+            await _load_manifest(URL, session)
             await session.close()
         mock_get.assert_called_once_with(URL)
         assert exc_info.value == error