|
1 | 1 | import json |
2 | 2 | from typing import Optional |
3 | 3 | from redis.exceptions import RedisError |
4 | | -from pydantic_core import core_schema |
5 | 4 |
|
6 | 5 | from src.common.connection import RedisConnectionManager |
7 | 6 | from src.common.server import mcp |
8 | 7 |
|
9 | 8 |
|
10 | | -# Custom type that accepts any JSON value but generates a proper schema |
11 | | -class JsonValue: |
12 | | - """Accepts any JSON-serializable value.""" |
13 | | - |
14 | | - @classmethod |
15 | | - def __get_pydantic_core_schema__(cls, _source_type, _handler): |
16 | | - """Define how Pydantic should validate this type.""" |
17 | | - # Accept any value |
18 | | - return core_schema.any_schema() |
19 | | - |
20 | | - @classmethod |
21 | | - def __get_pydantic_json_schema__(cls, _core_schema, _handler): |
22 | | - """Define the JSON schema for this type.""" |
23 | | - # Return a schema that accepts string, number, boolean, object, array, or null |
24 | | - return { |
25 | | - "anyOf": [ |
26 | | - {"type": "string"}, |
27 | | - {"type": "number"}, |
28 | | - {"type": "boolean"}, |
29 | | - {"type": "object"}, |
30 | | - {"type": "array", "items": {"type": "string"}}, |
31 | | - {"type": "null"}, |
32 | | - ] |
33 | | - } |
34 | | - |
35 | | - |
36 | 9 | @mcp.tool() |
37 | 10 | async def json_set( |
38 | 11 | name: str, |
39 | 12 | path: str, |
40 | | - value: JsonValue, |
| 13 | + value: str, |
41 | 14 | expire_seconds: Optional[int] = None, |
42 | 15 | ) -> str: |
43 | 16 | """Set a JSON value in Redis at a given path with an optional expiration time. |
44 | 17 |
|
45 | 18 | Args: |
46 | 19 | name: The Redis key where the JSON document is stored. |
47 | 20 | path: The JSON path where the value should be set. |
48 | | - value: The JSON value to store. |
| 21 | + value: The JSON value to store (as JSON string, or will be auto-converted). |
49 | 22 | expire_seconds: Optional; time in seconds after which the key should expire. |
50 | 23 |
|
51 | 24 | Returns: |
52 | 25 | A success message or an error message. |
53 | 26 | """ |
| 27 | + # Try to parse the value as JSON, if it fails, treat it as a plain string |
| 28 | + try: |
| 29 | + parsed_value = json.loads(value) |
| 30 | + except (json.JSONDecodeError, TypeError): |
| 31 | + parsed_value = value |
| 32 | + |
54 | 33 | try: |
55 | 34 | r = RedisConnectionManager.get_connection() |
56 | | - r.json().set(name, path, value) |
| 35 | + r.json().set(name, path, parsed_value) |
57 | 36 |
|
58 | 37 | if expire_seconds is not None: |
59 | 38 | r.expire(name, expire_seconds) |
|
0 commit comments