From 4ed1955aff343ed7ee2852e78fbb9a72cf85d30b Mon Sep 17 00:00:00 2001 From: somebody Date: Mon, 28 Nov 2022 16:11:09 -0600 Subject: [PATCH] Work on new save --- aiserver.py | 55 +++++++++++++++++++++++++++++++++----------- koboldai_settings.py | 39 +++++++++++++++++++++++++++++-- 2 files changed, 79 insertions(+), 15 deletions(-) diff --git a/aiserver.py b/aiserver.py index 7a2a8776..33e7ce0e 100644 --- a/aiserver.py +++ b/aiserver.py @@ -8142,10 +8142,12 @@ def get_files_folders(starting_folder): folders = [] files = [] base_path = os.path.abspath(starting_folder).replace("\\", "/") + if advanced_sort is not None: files_to_check = advanced_sort(base_path, desc=desc) else: files_to_check = get_files_sorted(base_path, sort, desc=desc) + for item in files_to_check: item_full_path = os.path.join(base_path, item).replace("\\", "/") if hasattr(os.stat(item_full_path), "st_file_attributes"): @@ -8571,24 +8573,51 @@ def valid_story(file): pass return valid +@logger.catch +def valid_v3_story(path: str) -> bool: + if not os.path.exists(path): return False + if not os.path.isdir(path): return False + if not os.path.exists(os.path.join(path, "story.json")): return False + return True + @logger.catch def story_sort(base_path, desc=False): files = {} for file in os.scandir(path=base_path): - if file.name.endswith(".json"): - filename = os.path.join(base_path, file.name).replace("\\", "/") - if os.path.getsize(filename) < 2*1024*1024: #2MB - with open(filename, "r") as f: - try: - js = json.load(f) - if 'story_name' in js and js['story_name'] in koboldai_vars.story_loads: - files[file.name] = datetime.datetime.strptime(koboldai_vars.story_loads[js['story_name']], "%m/%d/%Y, %H:%M:%S") - else: - files[file.name] = datetime.datetime.fromtimestamp(file.stat().st_mtime) - except: - pass + if file.is_dir(): + if not valid_v3_story(file.path): + continue + + story_path = os.path.join(file.path, "story.json") + story_stat = os.stat(story_path) + + if os.path.getsize(story_path) < 2*1024*1024: #2MB + with open(story_path, "r") as f: + j = json.load(f) + if j.get("story_name") in koboldai_vars.story_loads: + files[file.name] = datetime.datetime.strptime(koboldai_vars.story_loads[j["story_name"]], "%m/%d/%Y, %H:%M:%S") + else: + files[file.name] = datetime.datetime.fromtimestamp(story_stat.st_mtime) else: - files[file.name] = datetime.datetime.fromtimestamp(file.stat().st_mtime) + files[file.name] = datetime.datetime.fromtimestamp(story_stat.st_mtime) + continue + + if not file.name.endswith(".json"): + continue + + filename = os.path.join(base_path, file.name).replace("\\", "/") + if os.path.getsize(filename) < 2*1024*1024: #2MB + with open(filename, "r") as f: + try: + js = json.load(f) + if 'story_name' in js and js['story_name'] in koboldai_vars.story_loads: + files[file.name] = datetime.datetime.strptime(koboldai_vars.story_loads[js['story_name']], "%m/%d/%Y, %H:%M:%S") + else: + files[file.name] = datetime.datetime.fromtimestamp(file.stat().st_mtime) + except: + pass + else: + files[file.name] = datetime.datetime.fromtimestamp(file.stat().st_mtime) return [key[0] for key in sorted(files.items(), key=lambda kv: (kv[1], kv[0]), reverse=desc)] diff --git a/koboldai_settings.py b/koboldai_settings.py index a654f7c8..6e9b9f9d 100644 --- a/koboldai_settings.py +++ b/koboldai_settings.py @@ -871,10 +871,45 @@ class story_settings(settings): ################### must be at bottom ######################### self.no_save = False #Temporary disable save (doesn't save with the file) + + def save_story(self) -> None: + if self.no_save: + return + + if not any([self.prompt, self.memory, self.authornote, len(self.actions), len(self.worldinfo_v2)]): + return + logger.info("Saving") + + save_name = self.story_name or "Untitled" + + # Disambiguate stories by adding (n) if needed + disambiguator = 0 + save_path = os.path.join("stories", save_name) + while os.path.exists(save_path): + try: + # If the stories share a story id, overwrite the existing one. + with open(os.path.join(save_path, "story.json"), "r") as file: + j = json.load(file) + if self.story_id == j["story_id"]: + break + except FileNotFoundError: + raise FileNotFoundError("Malformed save file: Missing story.json") + + disambiguator += 1 + save_path = os.path.join("stories", save_name + (f" ({disambiguator})" if disambiguator else "")) - - def save_story(self): + if not os.path.exists(save_path): + # We are making the story for the first time. Setup the directory structure. + os.mkdir(save_path) + os.mkdir(os.path.join(save_path, "generated_audio")) + os.mkdir(os.path.join(save_path, "generated_images")) + + with open(os.path.join(save_path, "story.json"), "w") as file: + file.write(self.to_json()) + self.gamesaved = True + + def old_save_story(self): if not self.no_save: if self.prompt != "" or self.memory != "" or self.authornote != "" or len(self.actions) > 0 or len(self.worldinfo_v2) > 0: logger.debug("Saving story from koboldai_vars.story_settings.save_story()")