Work on new save

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)]
 
 

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()")