Agent persistence works

This commit is contained in:
Ashwin Bharambe 2024-09-21 22:43:03 -07:00
parent 4eb0f30891
commit 59f1fe5af8
17 changed files with 136 additions and 90 deletions

View file

@ -12,16 +12,37 @@ def kvstore_dependencies():
return ["aiosqlite", "psycopg2-binary", "redis"]
class InmemoryKVStoreImpl(KVStore):
def __init__(self):
self._store = {}
async def initialize(self) -> None:
pass
async def get(self, key: str) -> Optional[str]:
return self._store.get(key)
async def set(self, key: str, value: str) -> None:
self._store[key] = value
async def range(self, start_key: str, end_key: str) -> List[str]:
return [
self._store[key]
for key in self._store.keys()
if key >= start_key and key < end_key
]
async def kvstore_impl(config: KVStoreConfig) -> KVStore:
if config.type == KVStoreType.redis:
if config.type == KVStoreType.redis.value:
from .redis import RedisKVStoreImpl
impl = RedisKVStoreImpl(config)
elif config.type == KVStoreType.sqlite:
elif config.type == KVStoreType.sqlite.value:
from .sqlite import SqliteKVStoreImpl
impl = SqliteKVStoreImpl(config)
elif config.type == KVStoreType.pgvector:
elif config.type == KVStoreType.postgres.value:
raise NotImplementedError()
else:
raise ValueError(f"Unknown kvstore type {config.type}")