diff --git a/db/__init__.py b/db/__init__.py index ae45bae..068bf05 100644 --- a/db/__init__.py +++ b/db/__init__.py @@ -1,11 +1,12 @@ from .init_db import load_sqlite_engine, initialize_sqlite_db from .prompt_store import get_prompt_tinydb, initialize_prompt_store -from .dataset_store import get_all_dataset +from .dataset_store import get_all_dataset, save_dataset __all__ = [ "load_sqlite_engine", "initialize_sqlite_db", "get_prompt_tinydb", "initialize_prompt_store", - "get_all_dataset" + "get_all_dataset", + "save_dataset" ] \ No newline at end of file diff --git a/db/dataset_store.py b/db/dataset_store.py index e73ef33..bafce7e 100644 --- a/db/dataset_store.py +++ b/db/dataset_store.py @@ -39,6 +39,30 @@ def get_all_dataset(workdir: str) -> TinyDB: return db +def save_dataset(db: TinyDB, workdir: str, name: str = None) -> None: + """ + 将TinyDB中的数据集保存为单独的json文件 + + Args: + db (TinyDB): 包含数据集对象的TinyDB实例 + workdir (str): 工作目录路径 + name (str, optional): 要保存的数据集名称,None表示保存所有 + """ + dataset_dir = os.path.join(workdir, "dataset") + os.makedirs(dataset_dir, exist_ok=True) + + datasets = db.all() if name is None else db.search(Query().name == name) + + for dataset in datasets: + try: + filename = f"{dataset.get(dataset['name'])}.json" + filepath = os.path.join(dataset_dir, filename) + + with open(filepath, "w", encoding="utf-8") as f: + json.dump(dataset, f, ensure_ascii=False, indent=2) + except Exception as e: + print(f"Error saving dataset {dataset.get('id', 'unknown')}: {str(e)}") + if __name__ == "__main__": # 定义工作目录路径 workdir = os.path.join(os.path.dirname(__file__), "..", "workdir") @@ -47,4 +71,11 @@ if __name__ == "__main__": # 打印结果 print(f"Found {len(datasets)} datasets:") for ds in datasets.all(): - print(f"- {ds['name']} (ID: {ds['id']})") \ No newline at end of file + print(f"- {ds['name']} (ID: {ds['id']})") + + # 询问要保存的数据集名称 + name = input("输入要保存的数据集名称(直接回车保存所有): ").strip() or None + + # 保存数据集到文件 + save_dataset(datasets, workdir, name) + print(f"Datasets {'all' if name is None else name} saved to json files") \ No newline at end of file