This commit is contained in:
Nuno Campos
2024-08-21 09:30:22 -07:00
parent 656f89e16a
commit 8f8f3849fc
+6 -6
View File
@@ -1,5 +1,5 @@
from collections import defaultdict
from typing import List
from typing import List, Optional
from langgraph.kv.base import BaseKV, V
@@ -8,12 +8,12 @@ class MemoryKV(BaseKV):
def __init__(self) -> None:
self.data: dict[str, dict[str, V]] = defaultdict(dict)
def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], V | None]:
def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], Optional[V]]:
return {pair: self.data[pair[0]].get(pair[1]) for pair in pairs}
async def aget(
self, pairs: List[tuple[str, str]]
) -> dict[tuple[str, str], V | None]:
) -> dict[tuple[str, str], Optional[V]]:
return self.get(pairs)
def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
@@ -22,12 +22,12 @@ class MemoryKV(BaseKV):
async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
return self.list(prefixes)
def put(self, writes: List[tuple[str, str, V | None]]) -> None:
def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
for namespace, key, value in writes:
if value is None:
self.data[namespace].pop(key, None)
else:
self.data[namespace][key] = value
async def aput(self, writes: List[tuple[str, str, V | None]]) -> None:
self.put(writes)
async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
return self.put(writes)