Nephalem/llama-cpp-telegram_bot
0
1import json
2from os.path import exists
3
4import yaml
5from pathlib import Path
6
7
8class TelegramBotUser:
9 """
10 Class stored individual tg user info (history, message sequence, etc...) and provide some actions
11 """
12
13 default_messages_template = { # dict of messages templates for various situations. Use _VAR_ replacement
14 "mem_lost": "<b>MEMORY LOST!</b>\nSend /start or any text for new session.", # refers to non-existing
15 "retyping": "<i>_NAME2_ retyping...</i>", # added when "regenerate button" working
16 "typing": "<i>_NAME2_ typing...</i>", # added when generating working
17 "char_loaded": "_NAME2_ LOADED!\n_OPEN_TAG__GREETING__CLOSE_TAG_ ", # When new char loaded
18 "preset_loaded": "LOADED PRESET: _OPEN_TAG__CUSTOM_STRING__CLOSE_TAG_", # When new char loaded
19 "model_loaded": "LOADED MODEL: _OPEN_TAG__CUSTOM_STRING__CLOSE_TAG_", # When new char loaded
20 "mem_reset": "MEMORY RESET!\n_OPEN_TAG__GREETING__CLOSE_TAG_", # When history cleared
21 "hist_to_chat": "To load history - forward message to this chat", # download history
22 "hist_loaded": "_NAME2_ LOADED!\n_OPEN_TAG__GREETING__CLOSE_TAG_"
23 "\n\nLAST MESSAGE:\n_OPEN_TAG__CUSTOM_STRING__CLOSE_TAG_", # load history
24 }
25
26 def __init__(self,
27 char_file="",
28 name1="You",
29 name2="Bot",
30 context="",
31 example="",
32 language="en",
33 turn_template="",
34 greeting="Hi!"):
35 """
36 Init User class with default attribute
37 :param name1: username
38 :param name2: current character name
39 :param context: context of conversation, example: "Conversation between Bot and You"
40 :param greeting: just greeting message from bot
41 :return: None
42 """
43 self.char_file: str = char_file
44 self.name1: str = name1
45 self.name2: str = name2
46 self.context: str = context
47 self.example: str = example
48 self.language: str = language
49 self.turn_template: str = turn_template
50 self.user_in: list = [] # "user input history": [["Hi!","Who are you?"]], need for regenerate option
51 self.history: list = [] # "history": [["Hi!", "Hi there!","Who are you?", "I am you assistant."]],
52 self.msg_id: list = [] # "msg_id": [143, 144, 145, 146],
53 self.greeting: str = greeting
54
55 def pop(self):
56 # Converts all data to json string
57 user_in = self.user_in.pop()
58 msg_id = self.msg_id.pop()
59 self.history = self.history[:-2]
60 return user_in, msg_id
61
62 def reset_history(self):
63 # clear all user history
64 self.user_in = []
65 self.history = []
66 self.msg_id = []
67
68 def to_json(self):
69 # Converts all data to json string
70 return json.dumps({
71 "char_file": self.char_file,
72 "name1": self.name1,
73 "name2": self.name2,
74 "context": self.context,
75 "example": self.example,
76 "language": self.language,
77 "turn_template": self.turn_template,
78 "user_in": self.user_in,
79 "history": self.history,
80 "msg_id": self.msg_id,
81 "greeting": self.greeting,
82 })
83
84 def from_json(self, s: str):
85 # Converts json string to internal values
86 data = json.loads(s)
87 try:
88 self.char_file = data["char_file"] if "char_file" in data else ""
89 self.name1 = data["name1"] if "name1" in data else "You"
90 self.name2 = data["name2"] if "name2" in data else "Bot"
91 self.context = data["context"] if "context" in data else ""
92 self.example = data["example"] if "example" in data else ""
93 self.language = data["language"] if "language" in data else "en"
94 self.turn_template = data["turn_template"] if "turn_template" in data else ""
95 self.user_in = data["user_in"]
96 self.history = data["history"]
97 self.msg_id = data["msg_id"]
98 self.greeting = data["greeting"] if "greeting" in data else "Hi!"
99 return True
100 except Exception as exception:
101 print("from_json", exception)
102 return False
103
104 def load_character_file(self, characters_dir_path: str, char_file: str):
105 # Copy default user data. If reading will fail - return default user data
106 try:
107 # Try to read char file.
108 char_file_path = Path(f'{characters_dir_path}/{char_file}')
109 with open(char_file_path, 'r', encoding='utf-8') as user_file:
110 if char_file.split(".")[-1] == "json":
111 data = json.loads(user_file.read())
112 else:
113 data = yaml.safe_load(user_file.read())
114 # load persona and scenario
115 self.char_file = char_file
116 if 'you_name' in data:
117 self.name1 = data['you_name']
118 if 'char_name' in data:
119 self.name2 = data['char_name']
120 if 'name' in data:
121 self.name2 = data['name']
122 if 'turn_template' in data:
123 self.turn_template = data['turn_template']
124 self.context = ''
125 if 'char_persona' in data:
126 self.context += f"{self.name2}'s Persona: {data['char_persona'].strip()}\n"
127 if 'context' in data:
128 self.context += f"{data['context'].strip()}\n"
129 if 'world_scenario' in data:
130 self.context += f"Scenario: {data['world_scenario'].strip()}\n"
131 # add dialogue examples
132 if 'example_dialogue' in data:
133 self.example = f"\n{data['example_dialogue'].strip()}\n"
134 # add char greeting
135 if 'char_greeting' in data:
136 self.greeting = data['char_greeting'].strip()
137 if 'greeting' in data:
138 self.greeting = data['greeting'].strip()
139 self.context = self.replace_context_templates(self.context)
140 self.greeting = self.replace_context_templates(self.greeting)
141 self.example = self.replace_context_templates(self.example)
142 self.msg_id = []
143 self.user_in = []
144 self.history = []
145 except Exception as exception:
146 print("load_char_json_file", exception)
147 finally:
148 return self
149
150 def replace_context_templates(self, s: str) -> str:
151 s = s.replace('{{char}}', self.name2)
152 s = s.replace('{{user}}', self.name1)
153 s = s.replace('<BOT>', self.name2)
154 s = s.replace('<USER>', self.name1)
155 return s
156
157 def find_and_load_user_char_history(self, chat_id, history_dir_path: str):
158 chat_id = str(chat_id)
159 user_char_history_path = f'{history_dir_path}/{str(chat_id)}{self.char_file}.json'
160 user_char_history_old_path = f'{history_dir_path}/{str(chat_id)}{self.name2}.json'
161 if exists(user_char_history_path):
162 self.load_user_history(user_char_history_path)
163 elif exists(user_char_history_old_path):
164 self.load_user_history(user_char_history_old_path)
165
166 def load_user_history(self, file_path):
167 try:
168 if exists(file_path):
169 with open(file_path, 'r', encoding='utf-8') as user_file:
170 data = user_file.read()
171 self.from_json(data)
172 if self.char_file == "":
173 self.char_file = self.name2
174 except Exception as exception:
175 print(f"load_user_history: {exception}")
176
177 def save_user_history(self, chat_id, history_dir_path="history"):178 """179 Save two history file -user+char and default user history files and return their path180 :param chat_id: user chat_id
181 :param history_dir_path: history dir path
182 :return: user_char_file_path, default_user_file_path
183 """
184 if self.char_file == "":185 self.char_file = self.name2186 user_data = self.to_json()187 Path(history_dir_path).mkdir(parents=True, exist_ok=True)188 user_char_file_path = Path(f"{history_dir_path}/{chat_id}{self.char_file}.json")189 with user_char_file_path.open("w", encoding="utf-8") as user_file:
190 user_file.write(user_data)
191
192 default_user_file_path = Path(f"{history_dir_path}/{chat_id}.json")
193 with default_user_file_path.open("w", encoding="utf-8") as user_file:
194 user_file.write(user_data)
195
196 return str(user_char_file_path), str(default_user_file_path)
197