feat(db): 添加 prompt 存储功能

- 新增 prompt_store 模块,使用 TinyDB 存储 prompt 模板
- 在全局变量中添加 prompt_store 实例
- 更新 main.py,初始化 prompt 存储
- 新增 prompt 模板的 Pydantic 模型
- 更新 requirements.txt,添加 tinydb 依赖
This commit is contained in:
carry
2025-04-09 09:58:42 +08:00
parent 2d905a0270
commit 67281fe06a
6 changed files with 87 additions and 5 deletions

View File

@@ -1,3 +1,9 @@
from .init_db import get_sqlite_engine, initialize_sqlite_db
from .prompt_store import get_prompt_tinydb, initialize_prompt_store
__all__ = ['get_sqlite_engine', 'initialize_sqlite_db']
__all__ = [
"get_sqlite_engine",
"initialize_sqlite_db",
"get_prompt_tinydb",
"initialize_prompt_store"
]

60
db/prompt_store.py Normal file
View File

@@ -0,0 +1,60 @@
import os
import json
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数据库实例
"""
db.insert(promptTempleta(name="default",
description="默认提示词模板",
content="""项目名为:{ project_name }
请依据以下该项目官方文档的部分内容,创造合适的对话数据集用于微调一个了解该项目的小模型的语料,要求兼顾文档中间尽可能多的信息点,使用中文
文档节选:{ content }
按照如下json格式返回{ json }""").model_dump())
# TinyDB不需要显式创建表结构首次使用时自动创建
if __name__ == "__main__":
# 定义工作目录路径
workdir = os.path.join(os.path.dirname(__file__), "..", "workdir")
# 获取数据库实例
db = get_prompt_tinydb(workdir)
# 初始化prompt存储
initialize_prompt_store(db)