MetaGPT

1
# !/usr/bin/python3
2
# -*- coding: utf-8 -*-
3
"""
4
@Time    : 2023/12/27
5
@Author  : mashenquan
6
@File    : redis.py
7
"""
8
from __future__ import annotations
9

10
import traceback
11
from datetime import timedelta
12

13
import aioredis  # https://aioredis.readthedocs.io/en/latest/getting-started/
14

15
from metagpt.configs.redis_config import RedisConfig
16
from metagpt.logs import logger
17

18

19
class Redis:
20
    def __init__(self, config: RedisConfig = None):
21
        self.config = config
22
        self._client = None
23

24
    async def _connect(self, force=False):
25
        if self._client and not force:
26
            return True
27

28
        try:
29
            self._client = await aioredis.from_url(
30
                self.config.to_url(),
31
                username=self.config.username,
32
                password=self.config.password,
33
                db=self.config.db,
34
            )
35
            return True
36
        except Exception as e:
37
            logger.warning(f"Redis initialization has failed:{e}")
38
        return False
39

40
    async def get(self, key: str) -> bytes | None:
41
        if not await self._connect() or not key:
42
            return None
43
        try:
44
            v = await self._client.get(key)
45
            return v
46
        except Exception as e:
47
            logger.exception(f"{e}, stack:{traceback.format_exc()}")
48
            return None
49

50
    async def set(self, key: str, data: str, timeout_sec: int = None):
51
        if not await self._connect() or not key:
52
            return
53
        try:
54
            ex = None if not timeout_sec else timedelta(seconds=timeout_sec)
55
            await self._client.set(key, data, ex=ex)
56
        except Exception as e:
57
            logger.exception(f"{e}, stack:{traceback.format_exc()}")
58

59
    async def close(self):
60
        if not self._client:
61
            return
62
        await self._client.close()
63
        self._client = None
64