-
Notifications
You must be signed in to change notification settings - Fork 20
/
load_persona.py
41 lines (32 loc) · 1.26 KB
/
load_persona.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
import os
from ..config import current_dir_path
persona_path = os.path.join(current_dir_path, "persona")
txt_filenames = [os.path.splitext(file)[0] for file in os.listdir(persona_path) if file.endswith(".txt")]
class load_persona:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"persona_name": (txt_filenames, {"default": None}),
"is_enable": ("BOOLEAN", {"default": True}),
},
"optional": {
"file_content": ("STRING", {"forceInput": True}),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("system_prompt",)
FUNCTION = "persona"
# OUTPUT_NODE = False
CATEGORY = "大模型派对(llm_party)/面具(persona)"
def persona(self, persona_name, is_enable=True, file_content=None):
if is_enable == False:
return (None,)
per_path = os.path.join(persona_path, persona_name + ".txt")
text = ""
if file_content is not None:
text += "##背景知识:\n" + file_content + "\n\n"
if os.path.exists(per_path):
with open(per_path, "r", encoding="utf-8") as f:
text += f.read()
return (text,)