Repository: emcie-co/parlant
File: tests/core/services/tools/test_plugin_client.py (Python)

Code before:
    container: Container,
    agent: Agent,
) -> None:
    class ProductCategory(enum.Enum):
        CATEGORY_A = "category_a"
        CATEGORY_B = "category_b"

    @tool
    async def my_enum_tool(context: ToolContext, category: ProductCategory) -> ToolResult:
        return ToolResult(category.value)

    async with run_service_server([my_enum_tool]) as server:
        async with create_client(server, container[EventBufferFactory]) as client:
            tools = await client.list_tools()

            assert tools
            result = await client.call_tool(
                my_enum_tool.tool.name,
                context,
                arguments={"category": "category_a"},
            )

            assert result.data == "category_a"


async def test_that_a_plugin_tool_that_return_huge_payload_raise_an_error(
    context: ToolContext,
    container: Container,
    agent: Agent,
) -> None:
    @tool
    def huge_payload_tool(context: ToolContext) -> ToolResult:
        huge_payload = {f"key_{i}": "value" for i in range(10000)}
        return ToolResult({"size": len(huge_payload), "payload": huge_payload})

    async with run_service_server([huge_payload_tool]) as server:
        async with create_client(server, container[EventBufferFactory]) as client:
            tools = await client.list_tools()

            assert tools
            with raises(ToolResultError) as exc:
                await client.call_tool(huge_payload_tool.tool.name, context, arguments={})

            assert "Response exceeds 15KB limit" in str(exc.value)


Review comment (suggestion):
`test_that_a_plugin_tool_that_returns_a_huge_payload_raises_an_error`

Code after:
    container: Container,
    agent: Agent,
) -> None:
    class ProductCategory(enum.Enum):
        CATEGORY_A = "category_a"
        CATEGORY_B = "category_b"

    @tool
    async def my_enum_tool(context: ToolContext, category: ProductCategory) -> ToolResult:
        return ToolResult(category.value)

    async with run_service_server([my_enum_tool]) as server:
        async with create_client(server, container[EventBufferFactory]) as client:
            tools = await client.list_tools()

            assert tools
            result = await client.call_tool(
                my_enum_tool.tool.name,
                context,
                arguments={"category": "category_a"},
            )

            assert result.data == "category_a"


async def test_that_a_plugin_tool_that_return_huge_payload_raises_an_error(
    context: ToolContext,
    container: Container,
    agent: Agent,
) -> None:
    @tool
    def huge_payload_tool(context: ToolContext) -> ToolResult:
        huge_payload = {f"key_{i}": "value" for i in range(10000)}
        return ToolResult({"size": len(huge_payload), "payload": huge_payload})

    async with run_service_server([huge_payload_tool]) as server:
        async with create_client(server, container[EventBufferFactory]) as client:
            tools = await client.list_tools()

            assert tools
            with raises(ToolResultError) as exc:
                await client.call_tool(huge_payload_tool.tool.name, context, arguments={})

            assert "Response exceeds 15KB limit" in str(exc.value)