adding cache_function to base_tool

This commit is contained in:
João Moura
2024-03-19 13:33:20 -03:00
parent cf4f49c6e9
commit b8e86b0fb4
2 changed files with 33 additions and 8 deletions

View File

@@ -1,5 +1,5 @@
from abc import ABC, abstractmethod
from typing import Any, Callable, cast, Optional, Type
from typing import Any, Callable, Optional, Type
from pydantic import BaseModel, model_validator
from pydantic.v1 import BaseModel as V1BaseModel
@@ -14,6 +14,9 @@ class BaseTool(BaseModel, ABC):
args_schema: Optional[Type[V1BaseModel]] = None
"""The schema for the arguments that the tool accepts."""
description_updated: bool = False
"""Flag to check if the description has been updated."""
cache_function: Optional[Callable] = lambda: True
"""Function that will be used to determine if the tool should be cached, should return a boolean. If None, the tool will be cached."""
@model_validator(mode="after")
def _check_args_schema(self):

View File

@@ -1,6 +1,4 @@
import json
import pydantic_core
import pytest
from typing import Callable
from crewai_tools import BaseTool, tool
def test_creating_a_tool_using_annotation():
@@ -11,14 +9,14 @@ def test_creating_a_tool_using_annotation():
# Assert all the right attributes were defined
assert my_tool.name == "Name of my tool"
assert my_tool.description == "Clear description for what this tool is useful for, you agent will need this information to use it."
assert my_tool.description == "Name of my tool(question: 'string') - Clear description for what this tool is useful for, you agent will need this information to use it."
assert my_tool.args_schema.schema()["properties"] == {'question': {'title': 'Question', 'type': 'string'}}
assert my_tool.func("What is the meaning of life?") == "What is the meaning of life?"
# Assert the langchain tool conversion worked as expected
converted_tool = my_tool.to_langchain()
assert converted_tool.name == "Name of my tool"
assert converted_tool.description == "Clear description for what this tool is useful for, you agent will need this information to use it."
assert converted_tool.description == "Name of my tool(question: 'string') - Clear description for what this tool is useful for, you agent will need this information to use it."
assert converted_tool.args_schema.schema()["properties"] == {'question': {'title': 'Question', 'type': 'string'}}
assert converted_tool.func("What is the meaning of life?") == "What is the meaning of life?"
@@ -33,14 +31,38 @@ def test_creating_a_tool_using_baseclass():
my_tool = MyCustomTool()
# Assert all the right attributes were defined
assert my_tool.name == "Name of my tool"
assert my_tool.description == "Clear description for what this tool is useful for, you agent will need this information to use it."
assert my_tool.description == "Name of my tool(question: 'string') - Clear description for what this tool is useful for, you agent will need this information to use it."
assert my_tool.args_schema.schema()["properties"] == {'question': {'title': 'Question', 'type': 'string'}}
assert my_tool.run("What is the meaning of life?") == "What is the meaning of life?"
# Assert the langchain tool conversion worked as expected
converted_tool = my_tool.to_langchain()
assert converted_tool.name == "Name of my tool"
assert converted_tool.description == "Clear description for what this tool is useful for, you agent will need this information to use it."
assert converted_tool.description == "Name of my tool(question: 'string') - Clear description for what this tool is useful for, you agent will need this information to use it."
assert converted_tool.args_schema.schema()["properties"] == {'question': {'title': 'Question', 'type': 'string'}}
assert converted_tool.run("What is the meaning of life?") == "What is the meaning of life?"
def test_setting_cache_function():
class MyCustomTool(BaseTool):
name: str = "Name of my tool"
description: str = "Clear description for what this tool is useful for, you agent will need this information to use it."
cache_function: Callable = lambda: False
def _run(self, question: str) -> str:
return question
my_tool = MyCustomTool()
# Assert all the right attributes were defined
assert my_tool.cache_function() == False
def test_default_cache_function_is_true():
class MyCustomTool(BaseTool):
name: str = "Name of my tool"
description: str = "Clear description for what this tool is useful for, you agent will need this information to use it."
def _run(self, question: str) -> str:
return question
my_tool = MyCustomTool()
# Assert all the right attributes were defined
assert my_tool.cache_function() == True