autogen/python/packages/autogen-core/tests/test_workbench.py at efd4d85137024fc6a504c80290c8bda6035de4f5 · microsoft/autogen · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
from typing import Annotated

import pytest
from autogen_core.code_executor import ImportFromModule
from autogen_core.tools import FunctionTool, StaticWorkbench, Workbench


@pytest.mark.asyncio
async def test_static_workbench() -> None:
    def test_tool_func_1(x: Annotated[int, "The number to double."]) -> int:
        return x * 2

    def test_tool_func_2(x: Annotated[int, "The number to add 2."]) -> int:
        raise ValueError("This is a test error")  # Simulate an error

    test_tool_1 = FunctionTool(
        test_tool_func_1,
        name="test_tool_1",
        description="A test tool that doubles a number.",
        global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])],
    )
    test_tool_2 = FunctionTool(
        test_tool_func_2,
        name="test_tool_2",
        description="A test tool that adds 2 to a number.",
        global_imports=[ImportFromModule(module="typing_extensions", imports=["Annotated"])],
    )

    # Create a StaticWorkbench instance with the test tools.
    async with StaticWorkbench(tools=[test_tool_1, test_tool_2]) as workbench:
        # List tools
        tools = await workbench.list_tools()
        assert len(tools) == 2
        assert "description" in tools[0]
        assert "parameters" in tools[0]
        assert tools[0]["name"] == "test_tool_1"
        assert tools[0]["description"] == "A test tool that doubles a number."
        assert tools[0]["parameters"] == {
            "type": "object",
            "properties": {"x": {"type": "integer", "title": "X", "description": "The number to double."}},
            "required": ["x"],
            "additionalProperties": False,
        }
        assert "description" in tools[1]
        assert "parameters" in tools[1]
        assert tools[1]["name"] == "test_tool_2"
        assert tools[1]["description"] == "A test tool that adds 2 to a number."
        assert tools[1]["parameters"] == {
            "type": "object",
            "properties": {"x": {"type": "integer", "title": "X", "description": "The number to add 2."}},
            "required": ["x"],
            "additionalProperties": False,
        }

        # Call tools
        result_1 = await workbench.call_tool("test_tool_1", {"x": 5})
        assert result_1.name == "test_tool_1"
        assert result_1.result[0].type == "TextResultContent"
        assert result_1.result[0].content == "10"
        assert result_1.to_text() == "10"
        assert result_1.is_error is False

        # Call tool with error
        result_2 = await workbench.call_tool("test_tool_2", {"x": 5})
        assert result_2.name == "test_tool_2"
        assert result_2.result[0].type == "TextResultContent"
        assert result_2.result[0].content == "This is a test error"
        assert result_2.to_text() == "This is a test error"
        assert result_2.is_error is True

        # Save state.
        state = await workbench.save_state()
        assert state["type"] == "StaticWorkbenchState"
        assert "tools" in state
        assert len(state["tools"]) == 2

        # Dump config.
        config = workbench.dump_component()

    # Load the workbench from the config.
    async with Workbench.load_component(config) as new_workbench:
        # Load state.
        await new_workbench.load_state(state)

        # Verify that the tools are still available after loading the state.
        tools = await new_workbench.list_tools()
        assert len(tools) == 2
        assert "description" in tools[0]
        assert "parameters" in tools[0]
        assert tools[0]["name"] == "test_tool_1"
        assert tools[0]["description"] == "A test tool that doubles a number."
        assert tools[0]["parameters"] == {
            "type": "object",
            "properties": {"x": {"type": "integer", "title": "X", "description": "The number to double."}},
            "required": ["x"],
            "additionalProperties": False,
        }
        assert "description" in tools[1]
        assert "parameters" in tools[1]
        assert tools[1]["name"] == "test_tool_2"
        assert tools[1]["description"] == "A test tool that adds 2 to a number."
        assert tools[1]["parameters"] == {
            "type": "object",
            "properties": {"x": {"type": "integer", "title": "X", "description": "The number to add 2."}},
            "required": ["x"],
            "additionalProperties": False,
        }

        # Call tools
        result_1 = await new_workbench.call_tool("test_tool_1", {"x": 5})
        assert result_1.name == "test_tool_1"
        assert result_1.result[0].type == "TextResultContent"
        assert result_1.result[0].content == "10"
        assert result_1.to_text() == "10"
        assert result_1.is_error is False

        # Call tool with error
        result_2 = await new_workbench.call_tool("test_tool_2", {"x": 5})
        assert result_2.name == "test_tool_2"
        assert result_2.result[0].type == "TextResultContent"
        assert result_2.result[0].content == "This is a test error"
        assert result_2.to_text() == "This is a test error"
        assert result_2.is_error is True