1
+ import json
2
+ from typing import Union
3
+
1
4
from redis .exceptions import RedisError
2
- from redis . typing import EncodableT
5
+ from redis import Redis
3
6
4
7
from src .common .connection import RedisConnectionManager
5
8
from src .common .server import mcp
6
9
7
10
8
11
@mcp .tool ()
9
- async def set (key : str , value : EncodableT , expiration : int = None ) -> str :
12
+ async def set (key : str , value : Union [ str , bytes , int , float , dict ] , expiration : int = None ) -> str :
10
13
"""Set a Redis string value with an optional expiration time.
11
14
12
15
Args:
13
16
key (str): The key to set.
14
- value (str): The value to store.
17
+ value (str, bytes, int, float, dict ): The value to store.
15
18
expiration (int, optional): Expiration time in seconds.
16
19
17
20
Returns:
18
21
str: Confirmation message or an error message.
19
22
"""
23
+ if isinstance (value , bytes ):
24
+ encoded_value = value
25
+ elif isinstance (value , dict ):
26
+ encoded_value = json .dumps (value )
27
+ else :
28
+ encoded_value = str (value )
29
+
30
+ if isinstance (encoded_value , str ):
31
+ encoded_value = encoded_value .encode ("utf-8" )
32
+
20
33
try :
21
- r = RedisConnectionManager .get_connection ()
34
+ r : Redis = RedisConnectionManager .get_connection ()
22
35
if expiration :
23
- r .setex (key , expiration , value )
36
+ r .setex (key , expiration , encoded_value )
24
37
else :
25
- r .set (key , value )
38
+ r .set (key , encoded_value )
39
+
26
40
return f"Successfully set { key } " + (
27
41
f" with expiration { expiration } seconds" if expiration else ""
28
42
)
@@ -31,18 +45,29 @@ async def set(key: str, value: EncodableT, expiration: int = None) -> str:
31
45
32
46
33
47
@mcp .tool ()
34
- async def get (key : str ) -> str :
48
+ async def get (key : str ) -> Union [ str , bytes ] :
35
49
"""Get a Redis string value.
36
50
37
51
Args:
38
52
key (str): The key to retrieve.
39
53
40
54
Returns:
41
- str: The stored value or an error message.
55
+ str, bytes : The stored value or an error message.
42
56
"""
43
57
try :
44
- r = RedisConnectionManager .get_connection ()
58
+ r : Redis = RedisConnectionManager .get_connection ()
45
59
value = r .get (key )
46
- return value if value else f"Key { key } does not exist"
60
+
61
+ if value is None :
62
+ return f"Key { key } does not exist"
63
+
64
+ if isinstance (value , bytes ):
65
+ try :
66
+ text = value .decode ("utf-8" )
67
+ return text
68
+ except UnicodeDecodeError :
69
+ return value
70
+
71
+ return value
47
72
except RedisError as e :
48
73
return f"Error retrieving key { key } : { str (e )} "
0 commit comments