diff --git a/llama_stack/providers/impls/meta_reference/agents/persistence.py b/llama_stack/providers/impls/meta_reference/agents/persistence.py index 37ac75d6a..bbd63e893 100644 --- a/llama_stack/providers/impls/meta_reference/agents/persistence.py +++ b/llama_stack/providers/impls/meta_reference/agents/persistence.py @@ -68,10 +68,7 @@ class AgentPersistence: ) async def get_session_turns(self, session_id: str) -> List[Turn]: - values = await self.kvstore.range( - start_key=f"session:{self.agent_id}:{session_id}:", - end_key=f"session:{self.agent_id}:{session_id}:\xff\xff\xff\xff", - ) + values = await self.kvstore.get_match(key_to_match=f"session:{self.agent_id}:{session_id}:") turns = [] for value in values: try: diff --git a/llama_stack/providers/utils/kvstore/api.py b/llama_stack/providers/utils/kvstore/api.py index ba5b206c0..0570b199c 100644 --- a/llama_stack/providers/utils/kvstore/api.py +++ b/llama_stack/providers/utils/kvstore/api.py @@ -18,4 +18,4 @@ class KVStore(Protocol): async def delete(self, key: str) -> None: ... - async def range(self, start_key: str, end_key: str) -> List[str]: ... + async def get_match(self, key_to_match: str) -> List[str]: ... diff --git a/llama_stack/providers/utils/kvstore/kvstore.py b/llama_stack/providers/utils/kvstore/kvstore.py index a3cabc206..5165342e7 100644 --- a/llama_stack/providers/utils/kvstore/kvstore.py +++ b/llama_stack/providers/utils/kvstore/kvstore.py @@ -25,11 +25,11 @@ class InmemoryKVStoreImpl(KVStore): 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]: + async def get_match(self, key_to_match: str) -> List[str]: return [ - self._store[key] + self._store.get[key] for key in self._store.keys() - if key >= start_key and key < end_key + if key.startswith(key_to_match) ] diff --git a/llama_stack/providers/utils/kvstore/redis/redis.py b/llama_stack/providers/utils/kvstore/redis/redis.py index fb264b15c..0d961f0bd 100644 --- a/llama_stack/providers/utils/kvstore/redis/redis.py +++ b/llama_stack/providers/utils/kvstore/redis/redis.py @@ -45,8 +45,22 @@ class RedisKVStoreImpl(KVStore): key = self._namespaced_key(key) await self.redis.delete(key) - async def range(self, start_key: str, end_key: str) -> List[str]: - start_key = self._namespaced_key(start_key) - end_key = self._namespaced_key(end_key) + async def get_match(self, key_to_match: str) -> List[str]: + key_to_match = self._namespaced_key(key_to_match) - return await self.redis.zrangebylex(start_key, end_key) + cursor = 0 + keys = set() + + while True: + cursor, keys_chunk = await self.redis.scan(cursor=cursor, match=f"{key_to_match}*", count=100) + keys.update(key.decode() for key in keys_chunk) + if cursor == 0: + break + + if not keys: + return [] + + values = await self.redis.mget(*keys) + values = [value.decode() for value in values if value is not None] + + return sorted(values) diff --git a/llama_stack/providers/utils/kvstore/sqlite/sqlite.py b/llama_stack/providers/utils/kvstore/sqlite/sqlite.py index 1c5311d10..eeeaf38ae 100644 --- a/llama_stack/providers/utils/kvstore/sqlite/sqlite.py +++ b/llama_stack/providers/utils/kvstore/sqlite/sqlite.py @@ -60,11 +60,11 @@ class SqliteKVStoreImpl(KVStore): await db.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,)) await db.commit() - async def range(self, start_key: str, end_key: str) -> List[str]: + async def get_match(self, key_to_match: str) -> List[str]: async with aiosqlite.connect(self.db_path) as db: async with db.execute( - f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?", - (start_key, end_key), + f"SELECT key, value, expiration FROM {self.table_name} WHERE key LIKE ?", + (f"{key_to_match}%",), ) as cursor: result = [] async for row in cursor: