From 7f730fbe02a1bb588af7917a904bd6e101d84207 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 May 2025 06:39:49 +0000 Subject: [PATCH] Add usage limit feature to BaseTool class - Add max_usage_count and current_usage_count attributes to BaseTool - Implement usage limit checking in ToolUsage._use method - Add comprehensive tests for usage limit functionality - Maintain backward compatibility with None default for unlimited usage Co-Authored-By: Joe Moura --- src/crewai/tools/base_tool.py | 9 +++- src/crewai/tools/tool_usage.py | 14 ++++++ tests/tools/test_tool_usage_limit.py | 68 ++++++++++++++++++++++++++++ 3 files changed, 90 insertions(+), 1 deletion(-) create mode 100644 tests/tools/test_tool_usage_limit.py diff --git a/src/crewai/tools/base_tool.py b/src/crewai/tools/base_tool.py index 0e8a7a22b..0e0c35727 100644 --- a/src/crewai/tools/base_tool.py +++ b/src/crewai/tools/base_tool.py @@ -36,6 +36,10 @@ class BaseTool(BaseModel, ABC): """Function that will be used to determine if the tool should be cached, should return a boolean. If None, the tool will be cached.""" result_as_answer: bool = False """Flag to check if the tool should be the final agent answer.""" + max_usage_count: int | None = None + """Maximum number of times this tool can be used. None means unlimited usage.""" + current_usage_count: int = 0 + """Current number of times this tool has been used.""" @field_validator("args_schema", mode="before") @classmethod @@ -251,13 +255,14 @@ def to_langchain( return [t.to_structured_tool() if isinstance(t, BaseTool) else t for t in tools] -def tool(*args, result_as_answer=False): +def tool(*args, result_as_answer=False, max_usage_count=None): """ Decorator to create a tool from a function. Args: *args: Positional arguments, either the function to decorate or the tool name. result_as_answer: Flag to indicate if the tool result should be used as the final agent answer. + max_usage_count: Maximum number of times this tool can be used. None means unlimited usage. """ def _make_with_name(tool_name: str) -> Callable: @@ -284,6 +289,8 @@ def tool(*args, result_as_answer=False): func=f, args_schema=args_schema, result_as_answer=result_as_answer, + max_usage_count=max_usage_count, + current_usage_count=0, ) return _make_tool diff --git a/src/crewai/tools/tool_usage.py b/src/crewai/tools/tool_usage.py index dc5f8f29a..2ac9198c3 100644 --- a/src/crewai/tools/tool_usage.py +++ b/src/crewai/tools/tool_usage.py @@ -200,6 +200,17 @@ class ToolUsage: None, ) + if available_tool and hasattr(available_tool, 'max_usage_count') and available_tool.max_usage_count is not None: + if available_tool.current_usage_count >= available_tool.max_usage_count: + try: + result = f"Tool '{tool.name}' has reached its usage limit of {available_tool.max_usage_count} times and cannot be used anymore." + self._telemetry.tool_usage_error(llm=self.function_calling_llm) + result = self._format_result(result=result) + return result + except Exception: + if self.task: + self.task.increment_tools_errors() + if result is None: try: if calling.tool_name in [ @@ -300,6 +311,9 @@ class ToolUsage: if self.agent and hasattr(self.agent, "tools_results"): self.agent.tools_results.append(data) + if available_tool and hasattr(available_tool, 'current_usage_count'): + available_tool.current_usage_count += 1 + return result def _format_result(self, result: Any) -> str: diff --git a/tests/tools/test_tool_usage_limit.py b/tests/tools/test_tool_usage_limit.py new file mode 100644 index 000000000..132588a26 --- /dev/null +++ b/tests/tools/test_tool_usage_limit.py @@ -0,0 +1,68 @@ +import pytest +from unittest.mock import MagicMock, patch + +from crewai.tools import BaseTool, tool + + +def test_tool_usage_limit(): + """Test that tools respect usage limits.""" + class LimitedTool(BaseTool): + name: str = "Limited Tool" + description: str = "A tool with usage limits for testing" + max_usage_count: int = 2 + + def _run(self, input_text: str) -> str: + return f"Processed {input_text}" + + tool = LimitedTool() + + result1 = tool.run(input_text="test1") + assert result1 == "Processed test1" + assert tool.current_usage_count == 1 + + result2 = tool.run(input_text="test2") + assert result2 == "Processed test2" + assert tool.current_usage_count == 2 + + +def test_unlimited_tool_usage(): + """Test that tools without usage limits work normally.""" + class UnlimitedTool(BaseTool): + name: str = "Unlimited Tool" + description: str = "A tool without usage limits" + + def _run(self, input_text: str) -> str: + return f"Processed {input_text}" + + tool = UnlimitedTool() + + for i in range(5): + result = tool.run(input_text=f"test{i}") + assert result == f"Processed test{i}" + assert tool.current_usage_count == i + 1 + + +def test_tool_decorator_with_usage_limit(): + """Test usage limit with @tool decorator.""" + @tool("Test Tool", max_usage_count=3) + def test_tool(input_text: str) -> str: + """A test tool.""" + return f"Result: {input_text}" + + assert test_tool.max_usage_count == 3 + assert test_tool.current_usage_count == 0 + + result = test_tool.run(input_text="test") + assert result == "Result: test" + assert test_tool.current_usage_count == 1 + + +def test_default_unlimited_usage(): + """Test that tools have unlimited usage by default.""" + @tool("Default Tool") + def default_tool(input_text: str) -> str: + """A default tool.""" + return f"Result: {input_text}" + + assert default_tool.max_usage_count is None + assert default_tool.current_usage_count == 0