Cache should raise KeyError when key is missing

This commit is contained in:
Andrei Betlen 2023-05-05 12:21:49 -04:00
parent 24fc38754b
commit 5be0efa5f8

View file

@ -33,12 +33,10 @@ class LlamaCache:
return k
return None
def __getitem__(
self, key: Sequence[llama_cpp.llama_token]
) -> Optional["LlamaState"]:
def __getitem__(self, key: Sequence[llama_cpp.llama_token]) -> "LlamaState":
_key = self._find_key(tuple(key))
if _key is None:
return None
raise KeyError(f"Key not found: {key}")
return self.cache_state[_key]
def __contains__(self, key: Sequence[llama_cpp.llama_token]) -> bool: