File: test_toolset_functions.py

package info (click to toggle)
python-azure 20251014%2Bgit-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 766,472 kB
  • sloc: python: 6,314,744; ansic: 804; javascript: 287; makefile: 198; sh: 198; xml: 109
file content (186 lines) | stat: -rw-r--r-- 7,326 bytes parent folder | download
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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
import pytest

import azure.ai.agents.models as _models


class TestToolSetFunctions:
    def test_remove_code_interpreter_tool(self):
        """Test removing a CodeInterpreterTool from ToolSet."""
        toolset = _models.ToolSet()
        code_interpreter = _models.CodeInterpreterTool(file_ids=["file1", "file2"])
        toolset.add(code_interpreter)

        # Verify tool was added
        assert len(toolset._tools) == 1
        assert isinstance(toolset._tools[0], _models.CodeInterpreterTool)

        # Remove the tool
        toolset.remove(_models.CodeInterpreterTool)

        # Verify tool was removed
        assert len(toolset._tools) == 0

    def test_remove_file_search_tool(self):
        """Test removing a FileSearchTool from ToolSet."""
        toolset = _models.ToolSet()
        file_search = _models.FileSearchTool(vector_store_ids=["vs1", "vs2"])
        toolset.add(file_search)

        # Verify tool was added
        assert len(toolset._tools) == 1
        assert isinstance(toolset._tools[0], _models.FileSearchTool)

        # Remove the tool
        toolset.remove(_models.FileSearchTool)

        # Verify tool was removed
        assert len(toolset._tools) == 0

    def test_add_and_remove_openapi_tool_by_name(self):
        """Test removing a specific API definition from OpenApiTool by name."""
        from azure.ai.agents.models import OpenApiAuthDetails

        toolset = _models.ToolSet()
        auth = OpenApiAuthDetails(type="api_key")
        openapi_tool = _models.OpenApiTool(name="api1", description="First API", spec={"openapi": "3.0.0"}, auth=auth)

        # Add another definition to the same tool
        openapi_tool2 = _models.OpenApiTool(name="api2", description="Second API", spec={"openapi": "3.0.0"}, auth=auth)

        toolset.add(openapi_tool)
        toolset.add(openapi_tool2)

        # Verify tool was added with 2 definitions
        assert len(toolset._tools) == 1
        assert len(toolset._tools[0].definitions) == 2

        # Remove one definition by name
        toolset.remove(_models.OpenApiTool, name="api1")

        # Verify tool still exists but with only 1 definition
        assert len(toolset._tools) == 1
        assert len(toolset._tools[0].definitions) == 1
        assert toolset._tools[0].definitions[0].openapi.name == "api2"

    def test_remove_openapi_tool_entire_tool(self):
        """Test removing entire OpenApiTool without specifying name."""
        from azure.ai.agents.models import OpenApiAuthDetails

        toolset = _models.ToolSet()
        auth = OpenApiAuthDetails(type="api_key")
        openapi_tool = _models.OpenApiTool(name="api1", description="First API", spec={"openapi": "3.0.0"}, auth=auth)

        # Add another definition
        openapi_tool.add_definition(name="api2", description="Second API", spec={"openapi": "3.0.0"}, auth=auth)

        toolset.add(openapi_tool)

        # Verify tool was added with 2 definitions
        assert len(toolset._tools) == 1
        assert len(toolset._tools[0].definitions) == 2

        # Remove entire OpenApiTool without name parameter
        toolset.remove(_models.OpenApiTool)

        # Verify entire tool was removed
        assert len(toolset._tools) == 0

    def test_remove_mcp_tool_by_server_label(self):
        """Test removing a specific McpTool by server label."""
        toolset = _models.ToolSet()

        mcp_tool1 = _models.McpTool(server_label="server1", server_url="http://server1.com", allowed_tools=["tool1"])

        mcp_tool2 = _models.McpTool(server_label="server2", server_url="http://server2.com", allowed_tools=["tool2"])

        toolset.add(mcp_tool1)
        toolset.add(mcp_tool2)

        # Verify both tools were added
        assert len(toolset._tools) == 2

        # Remove one by server label
        toolset.remove(_models.McpTool, server_label="server1")

        # Verify only one tool remains
        assert len(toolset._tools) == 1
        assert toolset._tools[0].server_label == "server2"

    def test_remove_all_mcp_tools(self):
        """Test removing all McpTool instances without specifying server_label."""
        toolset = _models.ToolSet()

        mcp_tool1 = _models.McpTool(server_label="server1", server_url="http://server1.com", allowed_tools=["tool1"])

        mcp_tool2 = _models.McpTool(server_label="server2", server_url="http://server2.com", allowed_tools=["tool2"])

        toolset.add(mcp_tool1)
        toolset.add(mcp_tool2)

        # Add a non-MCP tool
        def dummy_function():
            pass

        function_tool = _models.FunctionTool({dummy_function})
        toolset.add(function_tool)

        # Verify all tools were added
        assert len(toolset._tools) == 3

        # Remove all MCP tools
        toolset.remove(_models.McpTool)

        # Verify only the function tool remains
        assert len(toolset._tools) == 1
        assert isinstance(toolset._tools[0], _models.FunctionTool)

    def test_remove_nonexistent_tool_type(self):
        """Test error when trying to remove a tool type that doesn't exist."""
        toolset = _models.ToolSet()

        with pytest.raises(ValueError, match="Tool of type FunctionTool not found in the ToolSet"):
            toolset.remove(_models.FunctionTool)

    def test_remove_openapi_tool_nonexistent_name(self):
        """Test error when trying to remove nonexistent API definition by name."""
        from azure.ai.agents.models import OpenApiAuthDetails

        toolset = _models.ToolSet()
        auth = OpenApiAuthDetails(type="api_key")
        openapi_tool = _models.OpenApiTool(name="api1", description="First API", spec={"openapi": "3.0.0"}, auth=auth)

        toolset.add(openapi_tool)

        # Try to remove a definition that doesn't exist
        with pytest.raises(ValueError, match="Definition with the name 'nonexistent' does not exist"):
            toolset.remove(_models.OpenApiTool, name="nonexistent")

    def test_remove_openapi_tool_by_name_no_openapi_tool(self):
        """Test error when trying to remove API definition by name but no OpenApiTool exists."""
        toolset = _models.ToolSet()

        with pytest.raises(ValueError, match="Tool of type OpenApiTool not found in the ToolSet"):
            toolset.remove(_models.OpenApiTool, name="api1")

    def test_remove_mcp_tool_nonexistent_server_label(self):
        """Test error when trying to remove McpTool with nonexistent server label."""
        toolset = _models.ToolSet()

        mcp_tool = _models.McpTool(server_label="server1", server_url="http://server1.com", allowed_tools=["tool1"])

        toolset.add(mcp_tool)

        # Try to remove MCP tool with nonexistent server label
        with pytest.raises(ValueError, match="McpTool with server label 'nonexistent' not found in the ToolSet"):
            toolset.remove(_models.McpTool, server_label="nonexistent")

    def test_remove_mcp_tool_no_mcp_tools(self):
        """Test error when trying to remove McpTool but none exist."""
        toolset = _models.ToolSet()

        with pytest.raises(ValueError, match="No tools of type McpTool found in the ToolSet"):
            toolset.remove(_models.McpTool)