62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
import os
|
||
import sys
|
||
from typing import Optional
|
||
from pathlib import Path
|
||
from datetime import datetime, timezone
|
||
from tinydb import TinyDB, Query
|
||
from tinydb.storages import JSONStorage
|
||
|
||
# 将项目根目录添加到系统路径中，以便能够导入项目中的其他模块
|
||
sys.path.append(str(Path(__file__).resolve().parent.parent))
|
||
from schema.prompt import promptTempleta
|
||
|
||
# 全局变量，用于存储TinyDB实例
|
||
_db_instance: Optional[TinyDB] = None
|
||
# 自定义存储类，用于格式化JSON数据
|
||
def get_prompt_tinydb(workdir: str) -> TinyDB:
|
||
"""
|
||
获取TinyDB实例。如果实例尚未创建，则创建一个新的并返回。
|
||
|
||
Args:
|
||
workdir (str): 工作目录路径，用于确定数据库文件的存储位置。
|
||
|
||
Returns:
|
||
TinyDB: TinyDB数据库实例
|
||
"""
|
||
global _db_instance
|
||
if not _db_instance:
|
||
# 创建数据库目录（如果不存在）
|
||
db_dir = os.path.join(workdir, "db")
|
||
os.makedirs(db_dir, exist_ok=True)
|
||
# 定义数据库文件路径
|
||
db_path = os.path.join(db_dir, "prompts.json")
|
||
# 创建TinyDB实例
|
||
_db_instance = TinyDB(db_path)
|
||
return _db_instance
|
||
|
||
def initialize_prompt_store(db: TinyDB) -> None:
|
||
"""
|
||
初始化prompt模板存储
|
||
|
||
Args:
|
||
db (TinyDB): TinyDB数据库实例
|
||
"""
|
||
# 检查数据库是否为空
|
||
if not db.all(): # 如果数据库中没有数据
|
||
db.insert(promptTempleta(
|
||
id=1,
|
||
name="default",
|
||
description="默认提示词模板",
|
||
content="""项目名为:{project_name}
|
||
请依据以下该项目官方文档的部分内容，创造合适的对话数据集用于微调一个了解该项目的小模型的语料，要求兼顾文档中间尽可能多的信息点，使用中文
|
||
文档节选：{document_slice}""").model_dump())
|
||
# 如果数据库中已有数据，则跳过插入
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 定义工作目录路径
|
||
workdir = os.path.join(os.path.dirname(__file__), "..", "workdir")
|
||
# 获取数据库实例
|
||
db = get_prompt_tinydb(workdir)
|
||
# 初始化prompt存储
|
||
initialize_prompt_store(db) |