import json import os import uuid from datetime import datetime from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from sse_starlette.sse import EventSourceResponse from config import settings from database import get_db from models.ai_model import AiModel from models.document import Document from models.generation_log import GenerationLog from models.paragraph import Paragraph from models.template import Template from schemas.schemas import GenerateFullRequest, GenerateTestRequest, Response from services.ai_service import call_ai from services.generation_runtime import ( build_mock_content, generation_progress, request_cancel, run_generation, update_progress, ) from services.minio_client import upload_bytes router = APIRouter() def _serialize_document(document: Document) -> dict: return { "id": document.id, "template_id": document.template_id, "name": document.name, "para_count_done": document.para_count_done, "para_count_total": document.para_count_total, "status": document.status, "file_path": document.file_path, "error": document.error, "created_at": document.created_at, "updated_at": document.updated_at, } @router.post("/test") async def generate_test(body: GenerateTestRequest, db: AsyncSession = Depends(get_db)): paragraph = await db.get(Paragraph, body.paragraph_id) if paragraph is None or paragraph.template_id != body.template_id: raise HTTPException(status_code=404, detail="段落不存在") model = None if body.model_id: model = await db.get(AiModel, body.model_id) elif paragraph.model_id: model = await db.get(AiModel, paragraph.model_id) if model is None or model.status != "enabled": content = build_mock_content(paragraph) message = "当前未找到可用模型，返回本地模拟生成结果。" else: result = await call_ai(paragraph, model) content = result.content message = f"已通过模型 {result.used_model} 生成。" return Response( data={ "paragraph_id": paragraph.id, "content": content, "message": message, } ) @router.post("/upload") async def upload_reference_file(file: UploadFile = File(...)): if not file.filename: raise HTTPException(status_code=400, detail="文件名不能为空") ext = os.path.splitext(file.filename)[1].lower() if ext not in settings.ALLOWED_EXTENSIONS: raise HTTPException(status_code=400, detail="文件类型不支持") content = await file.read() if not content: raise HTTPException(status_code=400, detail="上传文件不能为空") if len(content) > settings.MAX_UPLOAD_SIZE: raise HTTPException(status_code=400, detail="文件大小超过限制") object_name = f"{datetime.now().strftime('%Y%m%d')}/{uuid.uuid4().hex}{ext}" await asyncio.to_thread( upload_bytes, settings.MINIO_BUCKET_UPLOADS, object_name, content, file.content_type or "application/octet-stream", ) return Response( data={ "file_name": file.filename, "file_path": f"{settings.MINIO_BUCKET_UPLOADS}/{object_name}", } ) @router.post("/full") async def generate_full(body: GenerateFullRequest, db: AsyncSession = Depends(get_db)): template = await db.get(Template, body.template_id) if template is None: raise HTTPException(status_code=404, detail="模板不存在") result = await db.execute( select(Paragraph) .where(Paragraph.template_id == body.template_id) .order_by(Paragraph.sort_index.asc(), Paragraph.id.asc()) ) paragraphs = result.scalars().all() if not paragraphs: raise HTTPException(status_code=400, detail="模板下暂无可生成段落") document = Document( template_id=template.id, name=f"{template.name}-{datetime.now().strftime('%Y%m%d%H%M%S')}", para_count_done=0, para_count_total=len(paragraphs), status="generating", file_path="", error="", ) db.add(document) await db.flush() await db.commit() await db.refresh(document) update_progress(document.id, status="pending", percent=0, done=0, total=len(paragraphs), message="任务已创建") asyncio.create_task(run_generation(document.id, template.id)) return Response(data=_serialize_document(document)) @router.get("/progress/{document_id}") async def generate_progress(document_id: int): async def event_generator(): while True: state = generation_progress.get( document_id, {"status": "pending", "percent": 0, "message": "等待中", "done": 0, "total": 0}, ) yield { "event": "progress", "data": json.dumps(state, ensure_ascii=False), } if state.get("status") in {"completed", "failed", "cancelled"}: break await asyncio.sleep(1) return EventSourceResponse(event_generator()) @router.get("/documents") async def list_documents( page: int = Query(1, ge=1), page_size: int = Query(20, ge=1, le=100), db: AsyncSession = Depends(get_db), ): total = (await db.execute(select(func.count(Document.id)))).scalar_one() result = await db.execute( select(Document) .order_by(Document.id.desc()) .offset((page - 1) * page_size) .limit(page_size) ) items = [_serialize_document(item) for item in result.scalars().all()] return Response(data={"items": items, "total": total, "page": page, "page_size": page_size}) @router.get("/documents/{document_id}") async def get_document(document_id: int, db: AsyncSession = Depends(get_db)): document = await db.get(Document, document_id) if document is None: raise HTTPException(status_code=404, detail="生成记录不存在") log_result = await db.execute( select(GenerationLog, Paragraph) .join(Paragraph, Paragraph.id == GenerationLog.paragraph_id) .where(GenerationLog.document_id == document_id) .order_by(Paragraph.sort_index.asc(), Paragraph.id.asc()) ) items = [] for log, paragraph in log_result.all(): items.append( { "id": log.id, "paragraph_id": paragraph.id, "title": paragraph.title, "sort_index": paragraph.sort_index, "status": log.status, "content": json.loads(log.content) if log.content else {"content": []}, } ) payload = _serialize_document(document) payload["logs"] = items return Response(data=payload) @router.post("/cancel/{document_id}") async def cancel_document(document_id: int, db: AsyncSession = Depends(get_db)): document = await db.get(Document, document_id) if document is None: raise HTTPException(status_code=404, detail="生成记录不存在") if document.status in {"completed", "failed", "cancelled"}: return Response(data=_serialize_document(document)) request_cancel(document_id) return Response(data=_serialize_document(document)) @router.delete("/documents/{document_id}") async def delete_document(document_id: int, db: AsyncSession = Depends(get_db)): document = await db.get(Document, document_id) if document is None: raise HTTPException(status_code=404, detail="生成记录不存在") result = await db.execute(select(GenerationLog).where(GenerationLog.document_id == document_id)) for log in result.scalars().all(): await db.delete(log) await db.delete(document) await db.commit() return Response(data={"id": document_id})