224 lines
8.8 KiB
Python
224 lines
8.8 KiB
Python
"""AI智能排版路由 - 文本排版 + 自动配图"""
|
||
import json
|
||
import uuid
|
||
import httpx
|
||
from datetime import datetime
|
||
from typing import Optional
|
||
from fastapi import APIRouter, Depends, HTTPException
|
||
from pydantic import BaseModel
|
||
from sqlalchemy.orm import Session
|
||
|
||
from database import get_db
|
||
from models.user import User
|
||
from models.ai_model import AIModelConfig
|
||
from routers.auth import get_current_user
|
||
from services.ai_service import ai_service
|
||
|
||
router = APIRouter(prefix="/api/ai", tags=["AI智能排版"])
|
||
|
||
FORMAT_SYSTEM_PROMPT = """你是一位专业的文章排版编辑。你的任务是将用户提供的文章内容重新排版为结构清晰、可读性强的 Markdown 格式。
|
||
|
||
【核心原则】绝对不能修改、删除、改写或替换原文的任何文字内容。原文的每一个字、每一句话都必须原样保留。你可以在合适的位置补充过渡语、小结或说明文字来增强可读性，但必须与原文明确区分，且不能改动原文已有的文字。
|
||
|
||
排版规则：
|
||
1. 分析文章结构，在合适位置添加标题层级（## 和 ###），标题文字从原文中提取
|
||
2. 将原文中的要点整理为列表（有序或无序），但列表内容必须是原文原句
|
||
3. 重要观点用引用块 > 包裹，引用内容必须是原文原句
|
||
4. 关键词和重要内容用 **加粗** 标记
|
||
5. 保留原文中所有图片链接（格式），不要修改或删除
|
||
6. 保留原文中所有URL链接，转为 [链接文字](url) 格式
|
||
7. 适当添加分隔线 --- 划分章节
|
||
8. 长段落拆分为短段落，提高可读性，但不能改变段落中的文字
|
||
9. 如果内容中有流程或步骤，用有序列表清晰展示
|
||
10. 不要添加原文中没有的文字、解释、总结或过渡语
|
||
|
||
同时，你需要分析文章内容，为文章建议 1-2 张配图。对于每张配图，在合适的位置插入占位符，格式为：
|
||
[AI_IMAGE: 图片描述prompt，用英文写，描述要生成的图片内容，风格简洁专业]
|
||
|
||
注意：
|
||
- 占位符要插在文章逻辑合适的位置（如章节开头、流程说明旁边）
|
||
- prompt 用英文描述，风格：flat illustration, modern, professional, tech style
|
||
- 不要超过 2 个图片占位符
|
||
- 如果文章内容不适合配图（如纯代码、纯链接列表），可以不加占位符"""
|
||
|
||
|
||
class FormatRequest(BaseModel):
|
||
model_config = {"protected_namespaces": ()}
|
||
content: str
|
||
generate_images: bool = False # 是否生成配图（默认关闭）
|
||
model_config_id: Optional[int] = None # 指定排版用的文本模型
|
||
|
||
|
||
class FormatResponse(BaseModel):
|
||
formatted_content: str
|
||
images_generated: int = 0
|
||
|
||
|
||
def _get_image_model(db: Session):
|
||
"""从数据库查找已配置的图像生成模型（Seedream endpoint）"""
|
||
# 查找 task_type 包含 image 或 model_name 包含 seedream 的模型
|
||
model = db.query(AIModelConfig).filter(
|
||
AIModelConfig.is_enabled == True,
|
||
AIModelConfig.task_type == "image",
|
||
).first()
|
||
if model:
|
||
return {
|
||
"api_key": model.api_key,
|
||
"base_url": model.base_url or "https://ark.cn-beijing.volces.com/api/v3",
|
||
"model_id": model.model_id,
|
||
}
|
||
return None
|
||
|
||
|
||
async def _generate_image(api_key: str, base_url: str, model_id: str, prompt: str) -> Optional[bytes]:
|
||
"""调用火山方舟 Seedream 生成图片，返回图片二进制数据"""
|
||
url = f"{base_url.rstrip('/')}/images/generations"
|
||
headers = {
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
payload = {
|
||
"model": model_id,
|
||
"prompt": prompt,
|
||
"response_format": "url",
|
||
"size": "1920x1080", # 满足 Seedream 5.0 最低像素要求
|
||
}
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60) as client:
|
||
resp = await client.post(url, json=payload, headers=headers)
|
||
if resp.status_code != 200:
|
||
print(f"图像生成失败: {resp.status_code} - {resp.text}")
|
||
return None
|
||
data = resp.json()
|
||
image_url = data.get("data", [{}])[0].get("url", "")
|
||
if not image_url:
|
||
return None
|
||
# 下载图片（因为 Seedream 返回的是临时链接）
|
||
img_resp = await client.get(image_url)
|
||
if img_resp.status_code == 200:
|
||
return img_resp.content
|
||
return None
|
||
except Exception as e:
|
||
print(f"图像生成异常: {e}")
|
||
return None
|
||
|
||
|
||
def _upload_to_cos(db: Session, image_data: bytes) -> Optional[str]:
|
||
"""将图片上传到 COS，返回永久 URL"""
|
||
from models.system_config import SystemConfig
|
||
keys = ["cos_secret_id", "cos_secret_key", "cos_bucket", "cos_region", "cos_custom_domain"]
|
||
config = {}
|
||
for k in keys:
|
||
row = db.query(SystemConfig).filter(SystemConfig.key == k).first()
|
||
config[k] = row.value if row else ""
|
||
|
||
secret_id = config.get("cos_secret_id", "")
|
||
secret_key = config.get("cos_secret_key", "")
|
||
bucket = config.get("cos_bucket", "")
|
||
region = config.get("cos_region", "")
|
||
|
||
if not all([secret_id, secret_key, bucket, region]):
|
||
return None
|
||
|
||
try:
|
||
from qcloud_cos import CosConfig, CosS3Client
|
||
cos_config = CosConfig(Region=region, SecretId=secret_id, SecretKey=secret_key)
|
||
client = CosS3Client(cos_config)
|
||
|
||
date_prefix = datetime.now().strftime("%Y/%m")
|
||
filename = f"{uuid.uuid4().hex}.png"
|
||
object_key = f"images/{date_prefix}/{filename}"
|
||
|
||
client.put_object(
|
||
Bucket=bucket,
|
||
Body=image_data,
|
||
Key=object_key,
|
||
ContentType="image/png",
|
||
)
|
||
|
||
custom_domain = config.get("cos_custom_domain", "")
|
||
if custom_domain:
|
||
return f"https://{custom_domain}/{object_key}"
|
||
return f"https://{bucket}.cos.{region}.myqcloud.com/{object_key}"
|
||
except Exception as e:
|
||
print(f"COS上传失败: {e}")
|
||
return None
|
||
|
||
|
||
@router.post("/format", response_model=FormatResponse)
|
||
async def format_article(
|
||
req: FormatRequest,
|
||
current_user: User = Depends(get_current_user),
|
||
db: Session = Depends(get_db),
|
||
):
|
||
"""AI智能排版：格式化文章 + 自动生成配图"""
|
||
if not req.content.strip():
|
||
raise HTTPException(status_code=400, detail="文章内容不能为空")
|
||
|
||
# 第1步：AI 排版文本
|
||
messages = [{"role": "user", "content": req.content}]
|
||
try:
|
||
formatted = await ai_service.chat(
|
||
task_type="reasoning",
|
||
messages=messages,
|
||
system_prompt=FORMAT_SYSTEM_PROMPT,
|
||
stream=False,
|
||
model_config_id=req.model_config_id,
|
||
)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"AI排版失败: {str(e)}")
|
||
|
||
if not isinstance(formatted, str):
|
||
raise HTTPException(status_code=500, detail="AI排版返回格式异常")
|
||
|
||
# 过滤掉思考过程（DeepSeek Reasoner 会输出 <think>...</think>）
|
||
import re as _re
|
||
formatted = _re.sub(r'<think>[\s\S]*?</think>\s*', '', formatted).strip()
|
||
# 也过滤 <details> 格式的思考过程
|
||
formatted = _re.sub(r'<details>[\s\S]*?</details>\s*', '', formatted).strip()
|
||
|
||
images_generated = 0
|
||
|
||
# 第2步：生成配图（如果启用）
|
||
if req.generate_images:
|
||
import re
|
||
placeholders = re.findall(r'\[AI_IMAGE:\s*(.+?)\]', formatted)
|
||
|
||
if placeholders:
|
||
image_model = _get_image_model(db)
|
||
if image_model:
|
||
for prompt in placeholders:
|
||
image_data = await _generate_image(
|
||
image_model["api_key"],
|
||
image_model["base_url"],
|
||
image_model["model_id"],
|
||
prompt.strip(),
|
||
)
|
||
if image_data:
|
||
# 上传到 COS
|
||
cos_url = _upload_to_cos(db, image_data)
|
||
if cos_url:
|
||
formatted = formatted.replace(
|
||
f"[AI_IMAGE: {prompt}]",
|
||
f"![{prompt.strip()[:50]}]({cos_url})",
|
||
1,
|
||
)
|
||
images_generated += 1
|
||
continue
|
||
# 生成或上传失败，移除占位符
|
||
formatted = formatted.replace(f"[AI_IMAGE: {prompt}]", "", 1)
|
||
else:
|
||
# 没有配置图像模型，清理所有占位符
|
||
formatted = re.sub(r'\[AI_IMAGE:\s*.+?\]', '', formatted)
|
||
|
||
# 清理可能残留的占位符
|
||
import re
|
||
formatted = re.sub(r'\[AI_IMAGE:\s*.+?\]', '', formatted)
|
||
# 清理多余空行
|
||
formatted = re.sub(r'\n{3,}', '\n\n', formatted).strip()
|
||
|
||
return FormatResponse(
|
||
formatted_content=formatted,
|
||
images_generated=images_generated,
|
||
)
|