joke/api/app/routers/generate.py

"""AI 笑话生成器 API — 提示词从数据库读取"""

import json
from datetime import datetime
from fastapi import APIRouter, Depends, HTTPException
from openai import OpenAI
from sqlalchemy.orm import Session

from app.database import get_db
from app.models.setting import AiSetting
from app.schemas.joke import GenerateRequest, GenerateResponse


router = APIRouter(prefix="/generate", tags=["生成器"])

STYLE_MAP = {
    "cold": "冷幽默 / 无厘头",
    "warm": "温馨幽默 / 暖心搞笑",
    "twist": "反转 / 神转折",
    "pun": "谐音梗 / 文字游戏",
    "sketch": "段子 / 吐槽调侃",
    "irony": "讽刺幽默 / 黑色幽默",
}

LENGTH_MAP = {
    "short": "30-80字，非常简短",
    "medium": "80-150字，正常长度",
    "long": "150-300字，可以描述一个小场景",
}


def _load_prompt(setting: AiSetting) -> str:
    """从数据库读取生成提示词，没有则返回默认值"""
    prompt = setting.generate_prompt
    if not prompt or not prompt.strip():
        prompt = GENERATE_PROMPT_DEFAULT
    return prompt


GENERATE_PROMPT_DEFAULT = """你是一位幽默大师，专门创作轻松搞笑的短笑话。

{context}

要求：
1. 根据场景和关键词创作一条原创笑话
2. 笑话要有反转或意外结局
3. 语言风格：{style}
4. 字数要求：{length_requirement}
5. 直接输出笑话内容，不需要解释

请严格按照以下 JSON 格式输出，不要加任何额外说明：
{{
  "title": "笑话标题",
  "content": "笑话正文",
  "score": 8,
  "reason": "这个笑话巧妙结合了场景和关键词，结尾有反转"
}}

score 是 1-10 的整数评分，reason 是用一句话说明亮点。"""


def _build_prompt(
    scenarios: list[str],
    keywords: list[str],
    style: str = "twist",
    length: str = "medium",
    setting: AiSetting | None = None,
) -> str:
    """构建 AI prompt，支持风格和长度控制"""
    parts = []
    if scenarios:
        parts.append(f"场景：{', '.join(scenarios)}")
    if keywords:
        parts.append(f"关键词：{', '.join(keywords)}")
    if not parts:
        parts.append("场景：日常生活的各种趣事（不指定具体场景）")

    prompt_template = _load_prompt(setting) if setting else GENERATE_PROMPT_DEFAULT
    return prompt_template.format(
        context="\n".join(parts),
        style=STYLE_MAP.get(style, "反转 / 神转折"),
        length_requirement=LENGTH_MAP.get(length, "80-150字，正常长度"),
    )


@router.post("", response_model=GenerateResponse)
def generate_joke(
    req: GenerateRequest,
    db: Session = Depends(get_db),
):
    """调用 AI 生成笑话，提示词从数据库读取"""
    # 获取激活的 AI 配置
    setting = db.query(AiSetting).filter(AiSetting.is_active == True).first()
    if not setting:
        setting = db.query(AiSetting).first()

    if not setting or not setting.api_key:
        raise HTTPException(status_code=503, detail="AI 服务未配置，请联系管理员")

    try:
        client = OpenAI(base_url=setting.api_base, api_key=setting.api_key)
        prompt = _build_prompt(req.scenarios, req.keywords, req.style, req.length, setting)

        response = client.chat.completions.create(
            model=setting.model_name,
            messages=[{"role": "user", "content": prompt}],
            temperature=req.temperature,
            max_tokens=setting.max_tokens,
        )

        raw = response.choices[0].message.content
        return _parse_json_response(raw)

    except json.JSONDecodeError:
        raise HTTPException(status_code=500, detail="AI 返回格式异常，请重新生成")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"生成失败: {str(e)}")


def _parse_json_response(raw: str | None) -> GenerateResponse:
    """解析 AI 返回的 JSON 格式"""
    if not raw or not raw.strip():
        return GenerateResponse(title="生成的笑话", content="（内容生成失败，请重新生成）", score=0)

    raw = raw.strip()
    if raw.startswith("```"):
        lines = raw.split("\n")
        if len(lines) >= 3:
            raw = "\n".join(lines[1:-1]).strip()

    if raw.endswith(","):
        raw = raw[:-1]

    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        try:
            start = raw.index("{")
            end = raw.rindex("}") + 1
            raw = raw[start:end]
            data = json.loads(raw)
        except (ValueError, json.JSONDecodeError):
            return _parse_response(raw)

    return GenerateResponse(
        title=data.get("title", "生成的笑话"),
        content=data.get("content", ""),
        score=data.get("score", 0),
        reason=data.get("reason", ""),
        created_at=datetime.now(),
    )


def _parse_response(raw: str | None) -> GenerateResponse:
    """解析 AI 返回内容（非 JSON 回退）"""
    if not raw or not raw.strip():
        raise ValueError("AI 返回内容为空")

    title = ""
    content = raw

    for line in raw.split("\n"):
        line = line.strip()
        if line.startswith("标题：") or line.startswith("标题:"):
            title = line.split("：", 1)[-1].split(":", 1)[-1].strip()
            lines = content.split("\n")
            for i, l in enumerate(lines):
                if l.strip() == line:
                    lines[i] = ""
                    break
            content = "\n".join(lines).strip()
            break

    if not title:
        first_line = raw.split("\n")[0].strip()
        if first_line.startswith("标题"):
            first_line = first_line.split("：", 1)[-1].split(":", 1)[-1].strip()
        title = first_line[:30] if len(first_line) > 30 else first_line

    if not content.strip():
        content = "（内容生成失败，请重新生成）"

    return GenerateResponse(
        title=title or "生成的笑话",
        content=content.strip(),
        created_at=datetime.now(),
    )