From e2d295ae40c1aefce529333dd00c5799e3dba174 Mon Sep 17 00:00:00 2001 From: shrinitgoyal Date: Tue, 19 Aug 2025 14:54:10 +0530 Subject: [PATCH] handle DB and TTL for redis KV store --- llama_stack/providers/utils/kvstore/config.py | 4 +++- llama_stack/providers/utils/kvstore/redis/redis.py | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py index f00cb1f8b..07349c9cc 100644 --- a/llama_stack/providers/utils/kvstore/config.py +++ b/llama_stack/providers/utils/kvstore/config.py @@ -31,10 +31,12 @@ class RedisKVStoreConfig(CommonConfig): type: Literal[KVStoreType.redis.value] = KVStoreType.redis.value host: str = "localhost" port: int = 6379 + db: int = 0 + ttl: int = None @property def url(self) -> str: - return f"redis://{self.host}:{self.port}" + return f"redis://{self.host}:{self.port}/{self.db}" @classmethod def pip_packages(cls) -> list[str]: diff --git a/llama_stack/providers/utils/kvstore/redis/redis.py b/llama_stack/providers/utils/kvstore/redis/redis.py index 3d2d956c3..c17689272 100644 --- a/llama_stack/providers/utils/kvstore/redis/redis.py +++ b/llama_stack/providers/utils/kvstore/redis/redis.py @@ -26,7 +26,10 @@ class RedisKVStoreImpl(KVStore): async def set(self, key: str, value: str, expiration: datetime | None = None) -> None: key = self._namespaced_key(key) - await self.redis.set(key, value) + if self.config.ttl: + await self.redis.set(key, value, ex=self.config.ttl) + else: + await self.redis.set(key, value) if expiration: await self.redis.expireat(key, expiration)