File size: 457 Bytes
c631b4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from typing import Optional

from pydantic import PrivateAttr


class CacheHandler:
    """Callback handler for tool usage."""

    _cache: PrivateAttr = {}

    def __init__(self):
        self._cache = {}

    def add(self, tool, input, output):
        input = input.strip()
        self._cache[f"{tool}-{input}"] = output

    def read(self, tool, input) -> Optional[str]:
        input = input.strip()
        return self._cache.get(f"{tool}-{input}")