MoFin/prompt_manager/registry.py

"""提示词注册表 — 所有提示词的CRUD和版本管理"""

import json
import shutil
from pathlib import Path
from datetime import datetime
from typing import Optional

from .models import PromptDef, PromptVersion, PROMPT_CATEGORIES

# 数据文件路径
DATA_DIR = Path("/home/hmo/MoFin/data/prompts")
REGISTRY_PATH = DATA_DIR / "registry.json"
VERSIONS_DIR = DATA_DIR / "versions"


def ensure_dirs():
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    VERSIONS_DIR.mkdir(parents=True, exist_ok=True)


def load_registry() -> dict:
    """加载注册表"""
    ensure_dirs()
    if not REGISTRY_PATH.exists():
        return {"prompts": [], "updated_at": datetime.now().isoformat()}
    try:
        return json.loads(REGISTRY_PATH.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, FileNotFoundError):
        return {"prompts": [], "updated_at": datetime.now().isoformat()}


def save_registry(data: dict):
    ensure_dirs()
    data["updated_at"] = datetime.now().isoformat()
    REGISTRY_PATH.write_text(
        json.dumps(data, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )


# ═══════════════════════════════════════════
# CRUD
# ═══════════════════════════════════════════

def list_prompts(category: str = None) -> list:
    """列出所有提示词，可选按分类过滤"""
    reg = load_registry()
    prompts = []
    for p in reg.get("prompts", []):
        if category and p.get("category") != category:
            continue
        prompts.append(dict(PromptDef.from_dict(p).to_dict()))
    return prompts


def get_prompt(prompt_id: str) -> Optional[PromptDef]:
    """获取单个提示词"""
    reg = load_registry()
    for p in reg.get("prompts", []):
        if p["id"] == prompt_id:
            return PromptDef.from_dict(p)
    return None


def add_prompt(prompt_def: PromptDef):
    """添加新提示词"""
    reg = load_registry()
    # 检查是否已存在
    for i, p in enumerate(reg.get("prompts", [])):
        if p["id"] == prompt_def.id:
            reg["prompts"][i] = prompt_def.to_dict()
            break
    else:
        reg.setdefault("prompts", []).append(prompt_def.to_dict())
    save_registry(reg)


def update_prompt(prompt_id: str, updates: dict):
    """更新提示词元数据"""
    reg = load_registry()
    for p in reg.get("prompts", []):
        if p["id"] == prompt_id:
            for k, v in updates.items():
                if k not in ("id", "versions", "created_at"):
                    p[k] = v
            break
    save_registry(reg)


def delete_prompt(prompt_id: str):
    """删除提示词"""
    reg = load_registry()
    reg["prompts"] = [p for p in reg.get("prompts", []) if p["id"] != prompt_id]
    save_registry(reg)


def get_version_content(prompt_id: str, version: str) -> Optional[str]:
    """读取版本内容文件"""
    content_path = VERSIONS_DIR / f"{prompt_id}-{version}.md"
    if content_path.exists():
        return content_path.read_text(encoding="utf-8")
    # 回退：从版本记录的 content 字段读取
    prompt = get_prompt(prompt_id)
    if prompt:
        for v in prompt.versions:
            if v.version == version:
                return v.content
    return None


# ═══════════════════════════════════════════
# 版本管理
# ═══════════════════════════════════════════

def add_version(prompt_id: str, version: PromptVersion):
    """添加新版本到提示词"""
    prompt = get_prompt(prompt_id)
    if not prompt:
        raise ValueError(f"提示词 '{prompt_id}' 不存在")

    # 检查版本号是否已存在
    for i, v in enumerate(prompt.versions):
        if v.version == version.version:
            prompt.versions[i] = version
            break
    else:
        prompt.versions.append(version)

    # 保存版本内容到独立文件
    content_path = VERSIONS_DIR / f"{prompt_id}-{version.version}.md"
    content_path.write_text(version.content, encoding="utf-8")

    # 更新 content_path 只存路径，content 字段只存摘要
    version.content_path = str(content_path)
    version.content = version.content[:200] + "..." if len(version.content) > 200 else version.content

    # 设置当前版本
    prompt.current_version = version.version
    prompt.updated_at = datetime.now().isoformat()

    # 保存
    add_prompt(prompt)


def set_active_version(prompt_id: str, version_str: str):
    """切换当前活跃版本"""
    prompt = get_prompt(prompt_id)
    if not prompt:
        raise ValueError(f"提示词 '{prompt_id}' 不存在")

    # 验证版本存在
    version_exists = any(v.version == version_str for v in prompt.versions)
    if not version_exists:
        raise ValueError(f"版本 '{version_str}' 不存在于 '{prompt_id}'")

    # 将所有版本标记为非活跃，目标版本标记为活跃
    for v in prompt.versions:
        if v.version == version_str:
            v.status = "active"
        elif v.status == "active":
            v.status = "deprecated"

    prompt.current_version = version_str
    prompt.updated_at = datetime.now().isoformat()
    add_prompt(prompt)


# ═══════════════════════════════════════════
# 分类信息
# ═══════════════════════════════════════════

def get_categories() -> dict:
    """获取所有分类"""
    return dict(PROMPT_CATEGORIES)


def get_version_history(prompt_id: str) -> list:
    """获取某个提示词的所有版本历史"""
    prompt = get_prompt(prompt_id)
    if not prompt:
        return []
    history = []
    for v in prompt.versions:
        history.append({
            "version": v.version,
            "label": v.label,
            "status": v.status,
            "created_at": v.created_at,
            "changelog": v.changelog,
            "tags": v.tags,
            "is_current": v.version == prompt.current_version,
        })
    return sorted(history, key=lambda x: x["version"])