Coverage for chatgpt_proxy / cache / cache.py: 69%

1# MIT License 

2# 

3# Copyright (c) 2025 Tuomo Kriikkula 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy 

6# of this software and associated documentation files (the "Software"), to deal 

7# in the Software without restriction, including without limitation the rights 

8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

9# copies of the Software, and to permit persons to whom the Software is 

10# furnished to do so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in all 

13# copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

21# SOFTWARE. 

22 

23import os 

24from enum import StrEnum 

25 

26import aiocache 

27import redis.asyncio as redis 

28from aiocache.serializers import MsgPackSerializer 

29 

30from chatgpt_proxy.log import logger 

31from chatgpt_proxy.utils import is_prod_env 

32 

33 

34# TODO: think about these namespaces and how to split and use them. 

35class CacheNamespace(StrEnum): 

36 Database = "db" 

37 App = "app" 

38 

39 

40_default_cache = "redis" if is_prod_env else "memory" 

41_cache_method = os.getenv("CHATGPT_PROXY_CACHE_METHOD", _default_cache).lower().strip() 

42 

43 

44def setup_memory_cache(namespace: CacheNamespace) -> aiocache.SimpleMemoryCache: 

45 return aiocache.SimpleMemoryCache( 

46 namespace=namespace, 

47 serializer=MsgPackSerializer(), 

48 ) 

49 

50 

51def setup_redis_cache(namespace: CacheNamespace) -> aiocache.RedisCache: 

52 redis_url = os.environ["REDIS_URL"] 

53 redis_client = redis.Redis.from_url(redis_url) 

54 return aiocache.RedisCache( 

55 redis_client, 

56 namespace=namespace, 

57 serializer=MsgPackSerializer(), 

58 ) 

59 

60 

61db_cache: aiocache.BaseCache 

62app_cache: aiocache.BaseCache 

63 

64 

65def setup_cache(namespace: CacheNamespace) -> aiocache.BaseCache: 

66 global _cache_method 

67 

68 if _cache_method == "redis": 

69 if "REDIS_URL" not in os.environ and not is_prod_env: 

70 logger.warning( 

71 f"requested {namespace} cache method is 'redis', but no REDIS_URL is set, " 

72 f"falling back to in-memory cache (is_prod_env={is_prod_env})" 

73 ) 

74 cache = setup_memory_cache(namespace) 

75 else: 

76 cache = setup_redis_cache(namespace) 

77 pass # TODO 

78 elif _cache_method == "memory": 

79 cache = setup_memory_cache(namespace) 

80 pass # TODO 

81 else: 

82 logger.error("invalid cache method: '{}', defaulting to memory") 

83 _cache_method = "memory" 

84 cache = setup_memory_cache(namespace) 

85 return cache 

86 

87 

88db_cache = setup_cache(CacheNamespace.Database) 

89app_cache = setup_cache(CacheNamespace.App) 

90 

91# NOTE: this does not work here, so instead we'll close the cache when 

92# the Sanic application exits. 

93# asyncio_atexit.register(cache.close)