-
Notifications
You must be signed in to change notification settings - Fork 33
/
story.py
67 lines (55 loc) · 1.9 KB
/
story.py
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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import json
import os
from ..config import current_dir_path
file_path = os.path.join(current_dir_path, "file")
story_path = ""
class story_json_tool:
@classmethod
def INPUT_TYPES(s):
# 获取file_path文件夹下的所有json文件的文件名
paths = [f for f in os.listdir(file_path) if f.endswith(".json")]
return {
"required": {
"path": (paths, {"default": "story.json"}),
"is_enable": ("BOOLEAN", {"default": True}),
},
"optional": {},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("tools",)
FUNCTION = "file"
# OUTPUT_NODE = False
CATEGORY = "大模型派对(llm_party)/工具(tools)"
def file(self, path, is_enable=True):
if is_enable == False:
return (None,)
path = os.path.join(file_path, path)
global story_path
story_path = path
output = [
{
"type": "function",
"function": {
"name": "read_story_json",
"description": "用于查询小说对应章节的内容",
"parameters": {
"type": "object",
"properties": {
"chapter": {
"type": "string",
"description": "章节的数目,例如:0, 5, 1-1, 2-2-1, 3-23-42-12",
}
},
"required": ["chapter"],
},
},
}
]
out = json.dumps(output, ensure_ascii=False)
return (out,)
def read_story_json(chapter):
global story_path
with open(story_path, "r", encoding="utf-8") as f:
data = json.load(f)
out = json.dumps(data[str(chapter)], ensure_ascii=False, indent=4)
return out