article-writer/app.py

#!/usr/bin/env python3
"""
文章编写系统 - Article Writer System
功能：主题管理、素材收集、Prompt编辑、大模型生成文章
"""

import os
import json
import uuid
import time
import sqlite3
import requests
from datetime import datetime
from flask import Flask, render_template, request, jsonify, redirect, url_for, send_from_directory
from werkzeug.utils import secure_filename

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
UPLOAD_DIR = os.path.join(BASE_DIR, 'static', 'uploads')
DB_PATH = os.path.join(BASE_DIR, 'articles.db')
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp', 'svg', 'bmp'}

# LLM配置（通过环境变量读取，避免密钥泄露；运行时可从数据库覆盖）
LLM_URL = os.environ.get('LLM_URL', 'https://www.autodl.art/api/v1/chat/completions')
LLM_API_KEY = os.environ.get('LLM_API_KEY', '')
LLM_MODEL = os.environ.get('LLM_MODEL', 'qwen3.6-plus')

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB


def get_db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def init_db():
    conn = get_db()
    conn.executescript('''
        CREATE TABLE IF NOT EXISTS topics (
            id TEXT PRIMARY KEY,
            name TEXT NOT NULL,
            description TEXT DEFAULT '',
            prompt TEXT DEFAULT '',
            created_at TEXT DEFAULT (datetime('now','localtime')),
            updated_at TEXT DEFAULT (datetime('now','localtime'))
        );
        CREATE TABLE IF NOT EXISTS materials (
            id TEXT PRIMARY KEY,
            topic_id TEXT NOT NULL,
            type TEXT NOT NULL DEFAULT 'text',
            content TEXT DEFAULT '',
            file_path TEXT DEFAULT '',
            sort_order INTEGER DEFAULT 0,
            created_at TEXT DEFAULT (datetime('now','localtime')),
            FOREIGN KEY (topic_id) REFERENCES topics(id) ON DELETE CASCADE
        );
        CREATE TABLE IF NOT EXISTS articles (
            id TEXT PRIMARY KEY,
            topic_id TEXT NOT NULL,
            content TEXT DEFAULT '',
            prompt_used TEXT DEFAULT '',
            model TEXT DEFAULT '',
            llm_url TEXT DEFAULT '',
            created_at TEXT DEFAULT (datetime('now','localtime')),
            FOREIGN KEY (topic_id) REFERENCES topics(id) ON DELETE CASCADE
        );
        CREATE TABLE IF NOT EXISTS llm_config (
            id INTEGER PRIMARY KEY CHECK (id = 1),
            url TEXT DEFAULT '',
            api_key TEXT DEFAULT '',
            model TEXT DEFAULT '',
            updated_at TEXT DEFAULT (datetime('now','localtime'))
        );
        CREATE TABLE IF NOT EXISTS prompt_history (
            id TEXT PRIMARY KEY,
            topic_id TEXT NOT NULL,
            prompt TEXT NOT NULL DEFAULT '',
            model TEXT DEFAULT '',
            created_at TEXT DEFAULT (datetime('now','localtime')),
            FOREIGN KEY (topic_id) REFERENCES topics(id) ON DELETE CASCADE
        );
    ''')
    # 初始化默认配置
    conn.execute('INSERT OR IGNORE INTO llm_config (id, url, api_key, model) VALUES (1, ?, ?, ?)',
                 (LLM_URL, LLM_API_KEY, LLM_MODEL))
    conn.commit()
    conn.close()


def get_llm_config():
    """从数据库读取LLM配置，如无则用环境变量默认值"""
    conn = get_db()
    row = conn.execute('SELECT * FROM llm_config WHERE id=1').fetchone()
    conn.close()
    if row and row['url']:
        return {'url': row['url'], 'api_key': row['api_key'], 'model': row['model'] or LLM_MODEL}
    return {'url': LLM_URL, 'api_key': LLM_API_KEY, 'model': LLM_MODEL}


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


# ==================== 页面路由 ====================

@app.route('/')
def index():
    conn = get_db()
    topics = conn.execute(
        'SELECT t.*, (SELECT COUNT(*) FROM materials WHERE topic_id=t.id) as mat_count '
        'FROM topics t ORDER BY updated_at DESC'
    ).fetchall()
    conn.close()
    return render_template('index.html', topics=topics)


@app.route('/topic/<topic_id>')
def topic_detail(topic_id):
    conn = get_db()
    topic = conn.execute('SELECT * FROM topics WHERE id=?', (topic_id,)).fetchone()
    if not topic:
        conn.close()
        return redirect(url_for('index'))
    materials = conn.execute(
        'SELECT * FROM materials WHERE topic_id=? ORDER BY sort_order, created_at',
        (topic_id,)
    ).fetchall()
    articles = conn.execute(
        'SELECT * FROM articles WHERE topic_id=? ORDER BY created_at DESC',
        (topic_id,)
    ).fetchall()
    conn.close()
    return render_template('topic.html', topic=topic, materials=materials, articles=articles)


# ==================== 主题API ====================

@app.route('/api/topics', methods=['POST'])
def create_topic():
    data = request.get_json(force=True)
    name = data.get('name', '').strip()
    if not name:
        return jsonify({'error': '主题名不能为空'}), 400
    conn = get_db()
    tid = str(uuid.uuid4())[:8]
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    conn.execute(
        'INSERT INTO topics (id, name, description, prompt, created_at, updated_at) VALUES (?,?,?,?,?,?)',
        (tid, name, data.get('description', ''), data.get('prompt', ''), now, now)
    )
    conn.commit()
    topic = conn.execute('SELECT * FROM topics WHERE id=?', (tid,)).fetchone()
    conn.close()
    return jsonify(dict(topic))


@app.route('/api/topics/<topic_id>', methods=['PUT'])
def update_topic(topic_id):
    data = request.get_json(force=True)
    conn = get_db()
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    fields = []
    values = []
    for key in ['name', 'description', 'prompt']:
        if key in data:
            fields.append(f'{key}=?')
            values.append(data[key])
    if fields:
        fields.append('updated_at=?')
        values.append(now)
        values.append(topic_id)
        conn.execute(f'UPDATE topics SET {",".join(fields)} WHERE id=?', values)
        conn.commit()
    topic = conn.execute('SELECT * FROM topics WHERE id=?', (topic_id,)).fetchone()
    conn.close()
    return jsonify(dict(topic))


@app.route('/api/topics/<topic_id>', methods=['DELETE'])
def delete_topic(topic_id):
    conn = get_db()
    # 删除关联的素材文件
    materials = conn.execute('SELECT file_path FROM materials WHERE topic_id=? AND type="image"', (topic_id,)).fetchall()
    for m in materials:
        if m['file_path']:
            fp = os.path.join(BASE_DIR, 'static', m['file_path'])
            if os.path.exists(fp):
                os.remove(fp)
    conn.execute('DELETE FROM topics WHERE id=?', (topic_id,))
    conn.commit()
    conn.close()
    return jsonify({'ok': True})


# ==================== 素材API ====================

@app.route('/api/topics/<topic_id>/materials', methods=['POST'])
def add_material(topic_id):
    conn = get_db()
    # 检查主题存在
    if not conn.execute('SELECT id FROM topics WHERE id=?', (topic_id,)).fetchone():
        conn.close()
        return jsonify({'error': '主题不存在'}), 404

    mtype = request.form.get('type', 'text')
    mid = str(uuid.uuid4())[:8]
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    if mtype == 'image':
        file = request.files.get('file')
        if not file or not allowed_file(file.filename):
            conn.close()
            return jsonify({'error': '不支持的图片格式'}), 400
        ext = secure_filename(file.filename).rsplit('.', 1)[-1].lower()
        filename = f"{topic_id}_{mid}_{int(time.time())}.{ext}"
        filepath = os.path.join(UPLOAD_DIR, filename)
        file.save(filepath)
        conn.execute(
            'INSERT INTO materials (id, topic_id, type, content, file_path, created_at) VALUES (?,?,?,?,?,?)',
            (mid, topic_id, 'image', request.form.get('content', ''), f'uploads/{filename}', now)
        )
    else:
        content = request.form.get('content', '').strip()
        if not content:
            conn.close()
            return jsonify({'error': '文本内容不能为空'}), 400
        conn.execute(
            'INSERT INTO materials (id, topic_id, type, content, created_at) VALUES (?,?,?,?,?)',
            (mid, topic_id, 'text', content, now)
        )

    conn.execute('UPDATE topics SET updated_at=? WHERE id=?', (now, topic_id))
    conn.commit()
    material = conn.execute('SELECT * FROM materials WHERE id=?', (mid,)).fetchone()
    conn.close()
    return jsonify(dict(material))


@app.route('/api/materials/<material_id>', methods=['PUT'])
def update_material(material_id):
    data = request.get_json(force=True)
    conn = get_db()
    fields = []
    values = []
    for key in ['content', 'sort_order']:
        if key in data:
            fields.append(f'{key}=?')
            values.append(data[key])
    if fields:
        values.append(material_id)
        conn.execute(f'UPDATE materials SET {",".join(fields)} WHERE id=?', values)
        conn.commit()
    material = conn.execute('SELECT * FROM materials WHERE id=?', (material_id,)).fetchone()
    conn.close()
    return jsonify(dict(material)) if material else jsonify({'error': '未找到'}), 404


@app.route('/api/materials/<material_id>', methods=['DELETE'])
def delete_material(material_id):
    conn = get_db()
    mat = conn.execute('SELECT * FROM materials WHERE id=?', (material_id,)).fetchone()
    if mat and mat['file_path']:
        fp = os.path.join(BASE_DIR, 'static', mat['file_path'])
        if os.path.exists(fp):
            os.remove(fp)
    conn.execute('DELETE FROM materials WHERE id=?', (material_id,))
    conn.commit()
    conn.close()
    return jsonify({'ok': True})


# ==================== LLM配置API ====================

@app.route('/api/llm-config', methods=['GET'])
def get_llm_config_api():
    config = get_llm_config()
    # 隐藏部分API Key
    key = config['api_key']
    masked = key[:8] + '***' + key[-4:] if len(key) > 12 else '***'
    return jsonify({
        'url': config['url'],
        'api_key_masked': masked,
        'api_key_set': bool(key),
        'model': config['model']
    })


@app.route('/api/llm-config', methods=['PUT'])
def update_llm_config():
    data = request.get_json(force=True)
    conn = get_db()
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    # 读取当前配置
    current = conn.execute('SELECT * FROM llm_config WHERE id=1').fetchone()

    url = data.get('url', current['url'] if current else LLM_URL)
    # 如果前端传了空key则保留旧key
    api_key = data.get('api_key', '') or (current['api_key'] if current else LLM_API_KEY)
    model = data.get('model', current['model'] if current else LLM_MODEL)

    conn.execute('''
        INSERT INTO llm_config (id, url, api_key, model, updated_at) VALUES (1, ?, ?, ?, ?)
        ON CONFLICT(id) DO UPDATE SET url=?, api_key=?, model=?, updated_at=?
    ''', (1, url, api_key, model, now, url, api_key, model, now))
    conn.commit()
    conn.close()
    return jsonify({'ok': True})


# ==================== Prompt历史API ====================

@app.route('/api/topics/<topic_id>/prompt-history', methods=['GET'])
def get_prompt_history(topic_id):
    conn = get_db()
    history = conn.execute(
        'SELECT * FROM prompt_history WHERE topic_id=? ORDER BY created_at DESC LIMIT 20',
        (topic_id,)
    ).fetchall()
    conn.close()
    return jsonify([dict(h) for h in history])


# ==================== 文章生成API ====================

@app.route('/api/topics/<topic_id>/generate', methods=['POST'])
def generate_article(topic_id):
    conn = get_db()
    topic = conn.execute('SELECT * FROM topics WHERE id=?', (topic_id,)).fetchone()
    if not topic:
        conn.close()
        return jsonify({'error': '主题不存在'}), 404

    materials = conn.execute(
        'SELECT * FROM materials WHERE topic_id=? ORDER BY sort_order, created_at',
        (topic_id,)
    ).fetchall()
    conn.close()

    # 构建素材文本
    material_parts = []
    for m in materials:
        if m['type'] == 'text':
            material_parts.append(f"【文本素材】\n{m['content']}")
        elif m['type'] == 'image':
            desc = m['content'] or '（图片）'
            material_parts.append(f"【图片素材】{desc}")

    materials_text = '\n\n'.join(material_parts) if material_parts else '（暂无素材）'

    # 构建prompt
    custom_prompt = topic['prompt'] or '请根据以下素材撰写一篇结构清晰、内容丰富的文章：'

    # 获取LLM配置
    llm_config = get_llm_config()
    llm_url = llm_config['url']
    llm_key = llm_config['api_key']
    llm_model = llm_config['model']

    if not llm_key:
        return jsonify({'error': '请先配置LLM API Key（首页右上角设置）'}), 400

    user_message = f"""{custom_prompt}

主题：{topic['name']}
{f'主题说明：{topic["description"]}' if topic['description'] else ''}

=== 素材 ===
{materials_text}

=== 要求 ===
1. 充分利用提供的素材
2. 文章结构完整，层次分明
3. 语言流畅，逻辑清晰
4. 不要输出思考过程，直接给出文章内容"""

    # 调用大模型
    try:
        resp = requests.post(
            llm_url,
            headers={
                'Authorization': f'Bearer {llm_key}',
                'Content-Type': 'application/json'
            },
            json={
                'model': llm_model,
                'messages': [
                    {'role': 'user', 'content': user_message}
                ],
                'temperature': 0.7,
                'max_tokens': 8192
            },
            timeout=120
        )
        resp.raise_for_status()
        result = resp.json()
        article_content = result['choices'][0]['message']['content']
    except Exception as e:
        return jsonify({'error': f'生成失败: {str(e)}'}), 500

    # 保存prompt历史
    phid = str(uuid.uuid4())[:8]
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    conn = get_db()
    conn.execute(
        'INSERT INTO prompt_history (id, topic_id, prompt, model, created_at) VALUES (?,?,?,?,?)',
        (phid, topic_id, custom_prompt, llm_model, now)
    )

    # 保存文章
    aid = str(uuid.uuid4())[:8]
    conn.execute(
        'INSERT INTO articles (id, topic_id, content, prompt_used, model, llm_url, created_at) VALUES (?,?,?,?,?,?,?)',
        (aid, topic_id, article_content, custom_prompt, llm_model, llm_url, now)
    )
    conn.execute('UPDATE topics SET updated_at=? WHERE id=?', (now, topic_id))
    conn.commit()
    article = conn.execute('SELECT * FROM articles WHERE id=?', (aid,)).fetchone()
    conn.close()
    return jsonify(dict(article))


@app.route('/api/articles/<article_id>', methods=['DELETE'])
def delete_article(article_id):
    conn = get_db()
    conn.execute('DELETE FROM articles WHERE id=?', (article_id,))
    conn.commit()
    conn.close()
    return jsonify({'ok': True})


@app.route('/api/articles/<article_id>', methods=['PUT'])
def update_article(article_id):
    data = request.get_json(force=True)
    conn = get_db()
    conn.execute('UPDATE articles SET content=? WHERE id=?', (data.get('content', ''), article_id))
    conn.commit()
    conn.close()
    return jsonify({'ok': True})


# ==================== 启动 ====================

if __name__ == '__main__':
    os.makedirs(UPLOAD_DIR, exist_ok=True)
    init_db()
    print('=' * 50)
    print('  文章编写系统启动')
    print('  访问地址: http://localhost:19200')
    print('=' * 50)
    app.run(host='0.0.0.0', port=19200, debug=False)