File size: 1,364 Bytes
6a0af53
 
3e81177
 
 
6a0af53
 
 
 
 
 
 
3e81177
afc8995
 
 
 
 
 
 
 
 
 
 
 
 
6547a33
 
 
98b8e77
afc8995
 
98b8e77
bba8925
98b8e77
afc8995
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
from datetime import datetime, timedelta
from functools import wraps
from io import BytesIO

from fastapi.responses import StreamingResponse

CACHE_SIZE = 50

_cache = {}
_cache_time = {}


def ttl_cache(key_name, media_type=None, ttl_secs=20):
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Assuming the prompt is the key for caching, change as necessary
            key = kwargs.get(key_name)
            ttl = timedelta(seconds=ttl_secs)
            # Check cache
            if key in _cache:
                if datetime.now() - _cache_time[key] > ttl:
                    # Cache has expired
                    del _cache[key]
                    del _cache_time[key]
                else:
                    # if media_type == 'image/png':
                    #     return StreamingResponse(BytesIO(_cache[key]), media_type=media_type)
                    # else:
                    return StreamingResponse(BytesIO(_cache[key]), media_type="image/png")

            # Call the actual function if not in cache or expired
            response, image_data = await func(*args, **kwargs)
            # Cache the content of the response's body.
            _cache[key] = image_data
            _cache_time[key] = datetime.now()

            return response

        return wrapper

    return decorator