mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 09:08:31 +00:00
6b054651a7
* Refactoring task cache to be a tool The previous implementation of the task caching system was early exiting the agent executor due to the fact it was returning an AgentFinish object. This now refactors it to use a cache specific tool that is dynamically added and forced into the agent in case of a task execution that was already executed with the same input.
45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
from typing import Any, Dict
|
|
|
|
from langchain.callbacks.base import BaseCallbackHandler
|
|
|
|
from ..tools.cache_tools import CacheTools
|
|
from .cache_handler import CacheHandler
|
|
|
|
|
|
class ToolsHandler(BaseCallbackHandler):
|
|
"""Callback handler for tool usage."""
|
|
|
|
last_used_tool: Dict[str, Any] = {}
|
|
cache: CacheHandler = None
|
|
|
|
def __init__(self, cache: CacheHandler = None, **kwargs: Any):
|
|
"""Initialize the callback handler."""
|
|
self.cache = cache
|
|
super().__init__(**kwargs)
|
|
|
|
def on_tool_start(
|
|
self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
|
|
) -> Any:
|
|
"""Run when tool starts running."""
|
|
name = serialized.get("name")
|
|
if name not in ["invalid_tool", "_Exception"]:
|
|
tools_usage = {
|
|
"tool": name,
|
|
"input": input_str,
|
|
}
|
|
self.last_used_tool = tools_usage
|
|
|
|
def on_tool_end(self, output: str, **kwargs: Any) -> Any:
|
|
"""Run when tool ends running."""
|
|
if (
|
|
"is not a valid tool" not in output
|
|
and "Invalid or incomplete response" not in output
|
|
and "Invalid Format" not in output
|
|
):
|
|
if self.last_used_tool["tool"] != CacheTools().name:
|
|
self.cache.add(
|
|
tool=self.last_used_tool["tool"],
|
|
input=self.last_used_tool["input"],
|
|
output=output,
|
|
)
|