mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 00:58:30 +00:00
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 <joao@crewai.com>
This commit is contained in:
@@ -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."""
|
"""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
|
result_as_answer: bool = False
|
||||||
"""Flag to check if the tool should be the final agent answer."""
|
"""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")
|
@field_validator("args_schema", mode="before")
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -251,13 +255,14 @@ def to_langchain(
|
|||||||
return [t.to_structured_tool() if isinstance(t, BaseTool) else t for t in tools]
|
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.
|
Decorator to create a tool from a function.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
*args: Positional arguments, either the function to decorate or the tool name.
|
*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.
|
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:
|
def _make_with_name(tool_name: str) -> Callable:
|
||||||
@@ -284,6 +289,8 @@ def tool(*args, result_as_answer=False):
|
|||||||
func=f,
|
func=f,
|
||||||
args_schema=args_schema,
|
args_schema=args_schema,
|
||||||
result_as_answer=result_as_answer,
|
result_as_answer=result_as_answer,
|
||||||
|
max_usage_count=max_usage_count,
|
||||||
|
current_usage_count=0,
|
||||||
)
|
)
|
||||||
|
|
||||||
return _make_tool
|
return _make_tool
|
||||||
|
|||||||
@@ -200,6 +200,17 @@ class ToolUsage:
|
|||||||
None,
|
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:
|
if result is None:
|
||||||
try:
|
try:
|
||||||
if calling.tool_name in [
|
if calling.tool_name in [
|
||||||
@@ -300,6 +311,9 @@ class ToolUsage:
|
|||||||
if self.agent and hasattr(self.agent, "tools_results"):
|
if self.agent and hasattr(self.agent, "tools_results"):
|
||||||
self.agent.tools_results.append(data)
|
self.agent.tools_results.append(data)
|
||||||
|
|
||||||
|
if available_tool and hasattr(available_tool, 'current_usage_count'):
|
||||||
|
available_tool.current_usage_count += 1
|
||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
def _format_result(self, result: Any) -> str:
|
def _format_result(self, result: Any) -> str:
|
||||||
|
|||||||
68
tests/tools/test_tool_usage_limit.py
Normal file
68
tests/tools/test_tool_usage_limit.py
Normal file
@@ -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
|
||||||
Reference in New Issue
Block a user