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:
@@ -17,12 +20,23 @@ async def set(key: str, value: EncodableT, expiration: int = None) -> str:
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,7 +45,7 @@ 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:
@@ -43,6 +57,17 @@ async def get(key: str) -> str:
43
57
try :
44
58
r = 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