Skip to content

Commit

Permalink
bugfix: fix LFI issue
Browse files Browse the repository at this point in the history
  • Loading branch information
GaiZhenbiao committed Sep 18, 2024
1 parent d4451bb commit 22007a7
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 24 deletions.
8 changes: 4 additions & 4 deletions ChuanhuChatbot.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,8 @@ def create_new_model():
historySearchTextbox = gr.Textbox(show_label=False, container=False, placeholder=i18n(
"搜索(支持正则)..."), lines=1, elem_id="history-search-tb")
with gr.Column(min_width=52, scale=1, elem_id="gr-history-header-btns"):
uploadFileBtn = gr.UploadButton(
interactive=True, label="", file_types=[".json"], elem_id="gr-history-upload-btn")
uploadHistoryBtn = gr.UploadButton(
interactive=True, label="", file_types=[".json"], elem_id="gr-history-upload-btn", type="binary")
historyRefreshBtn = gr.Button("", elem_id="gr-history-refresh-btn")


Expand Down Expand Up @@ -708,8 +708,8 @@ def create_greeting(request: gr.Request):
js='(a,b)=>{return clearChatbot(a,b);}',
)
historySelectList.select(**load_history_from_file_args)
uploadFileBtn.upload(upload_chat_history, [current_model, uploadFileBtn], [
saveFileName, systemPromptTxt, chatbot, single_turn_checkbox, temperature_slider, top_p_slider, n_choices_slider, stop_sequence_txt, max_context_length_slider, max_generation_slider, presence_penalty_slider, frequency_penalty_slider, logit_bias_txt, user_identifier_txt, use_streaming_checkbox]).then(**refresh_history_args)
uploadHistoryBtn.upload(upload_chat_history, [current_model, uploadHistoryBtn], [
saveFileName, systemPromptTxt, chatbot, single_turn_checkbox, temperature_slider, top_p_slider, n_choices_slider, stop_sequence_txt, max_context_length_slider, max_generation_slider, presence_penalty_slider, frequency_penalty_slider, logit_bias_txt, user_identifier_txt, use_streaming_checkbox, historySelectList]).then(**refresh_history_args)
historyDownloadBtn.click(None, [
user_name, historySelectList], None, js='(a,b)=>{return downloadHistory(a,b,".json");}')
historyMarkdownDownloadBtn.click(None, [
Expand Down
47 changes: 28 additions & 19 deletions modules/models/base_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -986,28 +986,37 @@ def export_markdown(self, filename, chatbot):
filename += ".md"
save_file(filename, self)

def upload_chat_history(self, new_history_file_content=None):
logging.debug(f"{self.user_name} 加载对话历史中……")
if new_history_file_content is not None:
if isinstance(new_history_file_content, bytes):
try:
# Try to parse the content as JSON
json_content = json.loads(new_history_file_content.decode('utf-8'))

# If successful, save the content to a file
new_history_filename = new_auto_history_filename(self.user_name)
new_history_file_path = os.path.join(HISTORY_DIR, self.user_name, new_history_filename)

# Ensure the directory exists
os.makedirs(os.path.dirname(new_history_file_path), exist_ok=True)

# Write the content to the file
with open(new_history_file_path, 'w', encoding='utf-8') as f:
json.dump(json_content, f, ensure_ascii=False, indent=2)

self.history_file_path = new_history_filename
logging.info(f"History file uploaded and saved as {new_history_filename}")
except json.JSONDecodeError:
logging.error("Uploaded content is not valid JSON. Using default history.")
else:
logging.warning("Unexpected type for new_history_file_content. Using default history.")
return *self.load_chat_history(new_history_file_path), init_history_list(self.user_name)

def load_chat_history(self, new_history_file_path=None):
logging.debug(f"{self.user_name} 加载对话历史中……")
if new_history_file_path is not None:
if type(new_history_file_path) != str:
# copy file from new_history_file_path.name to os.path.join(HISTORY_DIR, self.user_name)
new_history_file_path = new_history_file_path.name
target_path = os.path.join(HISTORY_DIR, self.user_name, new_history_file_path)
# Check if the file is in the history directory
assert os.path.realpath(new_history_file_path).startswith(os.path.realpath(HISTORY_DIR))
assert os.path.realpath(target_path).startswith(os.path.realpath(HISTORY_DIR))
assert self.user_name in [i[0] for i in auth_list]
shutil.copyfile(
new_history_file_path,
os.path.join(
HISTORY_DIR,
self.user_name,
os.path.basename(new_history_file_path),
),
)
self.history_file_path = os.path.basename(new_history_file_path)
else:
self.history_file_path = new_history_file_path
self.history_file_path = new_history_file_path
try:
if self.history_file_path == os.path.basename(self.history_file_path):
history_file_path = os.path.join(
Expand Down
2 changes: 1 addition & 1 deletion modules/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ def export_markdown(current_model, *args):


def upload_chat_history(current_model, *args):
return current_model.load_chat_history(*args)
return current_model.upload_chat_history(*args)


def set_token_upper_limit(current_model, *args):
Expand Down

0 comments on commit 22007a7

Please sign in to comment.