Spaces:
Sleeping
Sleeping
File size: 1,460 Bytes
025a7d5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
from typing import Any, Dict
from langchain.callbacks.base import BaseCallbackHandler
from ..tools.cache_tools import CacheTools
from .cache.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,
)
|