andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-11 00:58:30 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Adding custom caching

Browse Source
This commit is contained in:
João Moura
2024-03-22 23:26:53 -03:00
parent d05dcac16f
commit 5977c442b1
10 changed files with 3273 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/crewai/agent.py
View File

@@ -193,13 +193,15 @@ class Agent(BaseModel):
task=task_prompt, context=context
)
tools = self._parse_tools(tools or self.tools)
tools = tools or self.tools
parsed_tools = self._parse_tools(tools)
self.create_agent_executor(tools=tools)
self.agent_executor.tools = tools
self.agent_executor.tools = parsed_tools
self.agent_executor.task = task
self.agent_executor.tools_description = render_text_description(tools)
self.agent_executor.tools_names = self.__tools_names(tools)
self.agent_executor.tools_description = render_text_description(parsed_tools)
self.agent_executor.tools_names = self.__tools_names(parsed_tools)
result = self.agent_executor.invoke(
{
@@ -220,10 +222,11 @@ class Agent(BaseModel):
Args:
cache_handler: An instance of the CacheHandler class.
"""
self.tools_handler = ToolsHandler()
if self.cache:
self.cache_handler = cache_handler
self.tools_handler = ToolsHandler(cache=self.cache_handler)
self.create_agent_executor()
self.tools_handler.cache = cache_handler
self.create_agent_executor()
def set_rpm_controller(self, rpm_controller: RPMController) -> None:
"""Set the rpm controller for the agent.
@@ -257,6 +260,7 @@ class Agent(BaseModel):
"i18n": self.i18n,
"tools": self._parse_tools(tools),
"verbose": self.verbose,
"original_tools": tools,
"handle_parsing_errors": True,
"max_iterations": self.max_iter,
"step_callback": self.step_callback,

2
src/crewai/agents/executor.py
View File

@@ -24,6 +24,7 @@ class CrewAgentExecutor(AgentExecutor):
task: Any = None
tools_description: str = ""
tools_names: str = ""
original_tools: List[Any] = []
function_calling_llm: Any = None
request_within_rpm_limit: Any = None
tools_handler: InstanceOf[ToolsHandler] = None
@@ -202,6 +203,7 @@ class CrewAgentExecutor(AgentExecutor):
tool_usage = ToolUsage(
tools_handler=self.tools_handler,
tools=self.tools,
original_tools=self.original_tools,
tools_description=self.tools_description,
tools_names=self.tools_names,
function_calling_llm=self.function_calling_llm,

15
src/crewai/agents/tools_handler.py
View File

@@ -1,7 +1,7 @@
from typing import Any
from typing import Any, Optional, Union
from ..tools.cache_tools import CacheTools
from ..tools.tool_calling import ToolCalling
from ..tools.tool_calling import InstructorToolCalling, ToolCalling
from .cache.cache_handler import CacheHandler
@@ -11,15 +11,20 @@ class ToolsHandler:
last_used_tool: ToolCalling = {}
cache: CacheHandler
def __init__(self, cache: CacheHandler):
def __init__(self, cache: Optional[CacheHandler] = None):
"""Initialize the callback handler."""
self.cache = cache
self.last_used_tool = {}
def on_tool_use(self, calling: ToolCalling, output: str) -> Any:
def on_tool_use(
self,
calling: Union[ToolCalling, InstructorToolCalling],
output: str,
should_cache: bool = True,
) -> Any:
"""Run when tool ends running."""
self.last_used_tool = calling
if calling.tool_name != CacheTools().name:
if self.cache and should_cache and calling.tool_name != CacheTools().name:
self.cache.add(
tool=calling.tool_name,
input=calling.arguments,

35
src/crewai/tools/tool_usage.py
View File

@@ -29,6 +29,7 @@ class ToolUsage:
task: Task being executed.
tools_handler: Tools handler that will manage the tool usage.
tools: List of tools available for the agent.
original_tools: Original tools available for the agent before being converted to BaseTool.
tools_description: Description of the tools available for the agent.
tools_names: Names of the tools available for the agent.
function_calling_llm: Language model to be used for the tool usage.
@@ -38,6 +39,7 @@ class ToolUsage:
self,
tools_handler: ToolsHandler,
tools: List[BaseTool],
original_tools: List[Any],
tools_description: str,
tools_names: str,
task: Any,
@@ -53,6 +55,7 @@ class ToolUsage:
self.tools_description = tools_description
self.tools_names = tools_names
self.tools_handler = tools_handler
self.original_tools = original_tools
self.tools = tools
self.task = task
self.action = action
@@ -111,7 +114,7 @@ class ToolUsage:
result = None
if self.tools_handler:
if self.tools_handler.cache:
result = self.tools_handler.cache.read(
tool=calling.tool_name, input=calling.arguments
)
@@ -159,7 +162,25 @@ class ToolUsage:
return self.use(calling=calling, tool_string=tool_string)
if self.tools_handler:
self.tools_handler.on_tool_use(calling=calling, output=result)
should_cache = True
print("FORA")
print(tool)
original_tool = next(
(ot for ot in self.original_tools if ot.name == tool.name), None
)
if (
hasattr(original_tool, "cache_function")
and original_tool.cache_function
):
print("CARALHOOOO")
print(original_tool.cache_function)
should_cache = original_tool.cache_function(
calling.arguments, result
)
self.tools_handler.on_tool_use(
calling=calling, output=result, should_cache=should_cache
)
self._printer.print(content=f"\n\n{result}\n", color="yellow")
self._telemetry.tool_usage(
@@ -248,12 +269,12 @@ class ToolUsage:
model=model,
instructions=dedent(
"""\
The schema should have the following structure, only two keys:
- tool_name: str
- arguments: dict (with all arguments being passed)
The schema should have the following structure, only two keys:
- tool_name: str
- arguments: dict (with all arguments being passed)
Example:
{"tool_name": "tool name", "arguments": {"arg_name1": "value", "arg_name2": 2}}""",
Example:
{"tool_name": "tool name", "arguments": {"arg_name1": "value", "arg_name2": 2}}""",
),
max_attemps=1,
)