diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py index 7b6a79350..e1e3fee95 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["redis"] = 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)