|
4 | 4 |
|
5 | 5 | import pytest |
6 | 6 |
|
7 | | -from haystack.tools import Tool, Toolset, flatten_tools_or_toolsets |
| 7 | +from haystack.tools import Tool, Toolset, flatten_tools_or_toolsets, warm_up_tools |
8 | 8 |
|
9 | 9 |
|
10 | 10 | def add_numbers(a: int, b: int) -> int: |
@@ -171,3 +171,225 @@ def test_flatten_multiple_toolsets(self, add_tool, multiply_tool, subtract_tool) |
171 | 171 | assert result[0].name == "add" |
172 | 172 | assert result[1].name == "multiply" |
173 | 173 | assert result[2].name == "subtract" |
| 174 | + |
| 175 | + |
| 176 | +class WarmupTrackingTool(Tool): |
| 177 | + """A tool that tracks whether warm_up was called.""" |
| 178 | + |
| 179 | + def __init__(self, *args, **kwargs): |
| 180 | + super().__init__(*args, **kwargs) |
| 181 | + self.was_warmed_up = False |
| 182 | + |
| 183 | + def warm_up(self): |
| 184 | + self.was_warmed_up = True |
| 185 | + |
| 186 | + |
| 187 | +class WarmupTrackingToolset(Toolset): |
| 188 | + """A toolset that tracks whether warm_up was called.""" |
| 189 | + |
| 190 | + def __init__(self, tools): |
| 191 | + super().__init__(tools) |
| 192 | + self.was_warmed_up = False |
| 193 | + |
| 194 | + def warm_up(self): |
| 195 | + self.was_warmed_up = True |
| 196 | + # Call parent to warm up individual tools |
| 197 | + super().warm_up() |
| 198 | + |
| 199 | + |
| 200 | +class TestWarmUpTools: |
| 201 | + """Tests for the warm_up_tools() function""" |
| 202 | + |
| 203 | + def test_warm_up_tools_with_none(self): |
| 204 | + """Test that warm_up_tools with None does nothing.""" |
| 205 | + # Should not raise any errors |
| 206 | + warm_up_tools(None) |
| 207 | + |
| 208 | + def test_warm_up_tools_with_single_tool(self): |
| 209 | + """Test that warm_up_tools works with a single tool in a list.""" |
| 210 | + tool = WarmupTrackingTool( |
| 211 | + name="test_tool", |
| 212 | + description="A test tool", |
| 213 | + parameters={"type": "object", "properties": {}}, |
| 214 | + function=lambda: "test", |
| 215 | + ) |
| 216 | + |
| 217 | + assert not tool.was_warmed_up |
| 218 | + warm_up_tools([tool]) |
| 219 | + assert tool.was_warmed_up |
| 220 | + |
| 221 | + def test_warm_up_tools_with_single_toolset(self): |
| 222 | + """ |
| 223 | + Test that when passing a single Toolset, both the Toolset.warm_up() |
| 224 | + and each individual tool's warm_up() are called. |
| 225 | + """ |
| 226 | + tool1 = WarmupTrackingTool( |
| 227 | + name="tool1", |
| 228 | + description="First tool", |
| 229 | + parameters={"type": "object", "properties": {}}, |
| 230 | + function=lambda: "tool1", |
| 231 | + ) |
| 232 | + tool2 = WarmupTrackingTool( |
| 233 | + name="tool2", |
| 234 | + description="Second tool", |
| 235 | + parameters={"type": "object", "properties": {}}, |
| 236 | + function=lambda: "tool2", |
| 237 | + ) |
| 238 | + |
| 239 | + toolset = WarmupTrackingToolset([tool1, tool2]) |
| 240 | + |
| 241 | + assert not toolset.was_warmed_up |
| 242 | + assert not tool1.was_warmed_up |
| 243 | + assert not tool2.was_warmed_up |
| 244 | + |
| 245 | + warm_up_tools(toolset) |
| 246 | + |
| 247 | + # Both the toolset itself and individual tools should be warmed up |
| 248 | + assert toolset.was_warmed_up |
| 249 | + assert tool1.was_warmed_up |
| 250 | + assert tool2.was_warmed_up |
| 251 | + |
| 252 | + def test_warm_up_tools_with_list_containing_toolset(self): |
| 253 | + """Test that when a Toolset is in a list, individual tools inside get warmed up.""" |
| 254 | + tool1 = WarmupTrackingTool( |
| 255 | + name="tool1", |
| 256 | + description="First tool", |
| 257 | + parameters={"type": "object", "properties": {}}, |
| 258 | + function=lambda: "tool1", |
| 259 | + ) |
| 260 | + tool2 = WarmupTrackingTool( |
| 261 | + name="tool2", |
| 262 | + description="Second tool", |
| 263 | + parameters={"type": "object", "properties": {}}, |
| 264 | + function=lambda: "tool2", |
| 265 | + ) |
| 266 | + |
| 267 | + toolset = WarmupTrackingToolset([tool1, tool2]) |
| 268 | + |
| 269 | + assert not toolset.was_warmed_up |
| 270 | + assert not tool1.was_warmed_up |
| 271 | + assert not tool2.was_warmed_up |
| 272 | + |
| 273 | + warm_up_tools([toolset]) |
| 274 | + |
| 275 | + # Both the toolset itself and individual tools should be warmed up |
| 276 | + assert toolset.was_warmed_up |
| 277 | + assert tool1.was_warmed_up |
| 278 | + assert tool2.was_warmed_up |
| 279 | + |
| 280 | + def test_warm_up_tools_with_multiple_toolsets(self): |
| 281 | + """Test multiple Toolsets in a list.""" |
| 282 | + tool1 = WarmupTrackingTool( |
| 283 | + name="tool1", |
| 284 | + description="First tool", |
| 285 | + parameters={"type": "object", "properties": {}}, |
| 286 | + function=lambda: "tool1", |
| 287 | + ) |
| 288 | + tool2 = WarmupTrackingTool( |
| 289 | + name="tool2", |
| 290 | + description="Second tool", |
| 291 | + parameters={"type": "object", "properties": {}}, |
| 292 | + function=lambda: "tool2", |
| 293 | + ) |
| 294 | + tool3 = WarmupTrackingTool( |
| 295 | + name="tool3", |
| 296 | + description="Third tool", |
| 297 | + parameters={"type": "object", "properties": {}}, |
| 298 | + function=lambda: "tool3", |
| 299 | + ) |
| 300 | + |
| 301 | + toolset1 = WarmupTrackingToolset([tool1]) |
| 302 | + toolset2 = WarmupTrackingToolset([tool2, tool3]) |
| 303 | + |
| 304 | + assert not toolset1.was_warmed_up |
| 305 | + assert not toolset2.was_warmed_up |
| 306 | + assert not tool1.was_warmed_up |
| 307 | + assert not tool2.was_warmed_up |
| 308 | + assert not tool3.was_warmed_up |
| 309 | + |
| 310 | + warm_up_tools([toolset1, toolset2]) |
| 311 | + |
| 312 | + # Both toolsets and all individual tools should be warmed up |
| 313 | + assert toolset1.was_warmed_up |
| 314 | + assert toolset2.was_warmed_up |
| 315 | + assert tool1.was_warmed_up |
| 316 | + assert tool2.was_warmed_up |
| 317 | + assert tool3.was_warmed_up |
| 318 | + |
| 319 | + def test_warm_up_tools_with_mixed_tools_and_toolsets(self): |
| 320 | + """Test list with both Tool objects and Toolsets.""" |
| 321 | + standalone_tool = WarmupTrackingTool( |
| 322 | + name="standalone", |
| 323 | + description="Standalone tool", |
| 324 | + parameters={"type": "object", "properties": {}}, |
| 325 | + function=lambda: "standalone", |
| 326 | + ) |
| 327 | + toolset_tool1 = WarmupTrackingTool( |
| 328 | + name="toolset_tool1", |
| 329 | + description="Tool in toolset", |
| 330 | + parameters={"type": "object", "properties": {}}, |
| 331 | + function=lambda: "toolset_tool1", |
| 332 | + ) |
| 333 | + toolset_tool2 = WarmupTrackingTool( |
| 334 | + name="toolset_tool2", |
| 335 | + description="Another tool in toolset", |
| 336 | + parameters={"type": "object", "properties": {}}, |
| 337 | + function=lambda: "toolset_tool2", |
| 338 | + ) |
| 339 | + |
| 340 | + toolset = WarmupTrackingToolset([toolset_tool1, toolset_tool2]) |
| 341 | + |
| 342 | + assert not standalone_tool.was_warmed_up |
| 343 | + assert not toolset.was_warmed_up |
| 344 | + assert not toolset_tool1.was_warmed_up |
| 345 | + assert not toolset_tool2.was_warmed_up |
| 346 | + |
| 347 | + warm_up_tools([standalone_tool, toolset]) |
| 348 | + |
| 349 | + # All tools and the toolset should be warmed up |
| 350 | + assert standalone_tool.was_warmed_up |
| 351 | + assert toolset.was_warmed_up |
| 352 | + assert toolset_tool1.was_warmed_up |
| 353 | + assert toolset_tool2.was_warmed_up |
| 354 | + |
| 355 | + def test_warm_up_tools_idempotency(self): |
| 356 | + """Test that calling warm_up_tools() multiple times is safe.""" |
| 357 | + |
| 358 | + class WarmupCountingTool(Tool): |
| 359 | + """A tool that counts how many times warm_up was called.""" |
| 360 | + |
| 361 | + def __init__(self, *args, **kwargs): |
| 362 | + super().__init__(*args, **kwargs) |
| 363 | + self.warm_up_count = 0 |
| 364 | + |
| 365 | + def warm_up(self): |
| 366 | + self.warm_up_count += 1 |
| 367 | + |
| 368 | + class WarmupCountingToolset(Toolset): |
| 369 | + """A toolset that counts how many times warm_up was called.""" |
| 370 | + |
| 371 | + def __init__(self, tools): |
| 372 | + super().__init__(tools) |
| 373 | + self.warm_up_count = 0 |
| 374 | + |
| 375 | + def warm_up(self): |
| 376 | + self.warm_up_count += 1 |
| 377 | + super().warm_up() # Also warm up individual tools |
| 378 | + |
| 379 | + tool = WarmupCountingTool( |
| 380 | + name="counting_tool", |
| 381 | + description="A counting tool", |
| 382 | + parameters={"type": "object", "properties": {}}, |
| 383 | + function=lambda: "test", |
| 384 | + ) |
| 385 | + toolset = WarmupCountingToolset([tool]) |
| 386 | + |
| 387 | + # Call warm_up_tools multiple times |
| 388 | + warm_up_tools(toolset) |
| 389 | + warm_up_tools(toolset) |
| 390 | + warm_up_tools(toolset) |
| 391 | + |
| 392 | + # warm_up_tools itself doesn't prevent multiple calls, |
| 393 | + # but verify the calls actually happen multiple times |
| 394 | + assert toolset.warm_up_count == 3 |
| 395 | + assert tool.warm_up_count == 3 |
0 commit comments