import asyncio 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.file_summary import summarize_minio_files 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="段落不存在") if body.prompt_text: paragraph.prompt_text = body.prompt_text 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) file_summaries = await asyncio.to_thread(summarize_minio_files, body.file_paths or []) if body.file_paths else [] if model is None or model.status != "enabled": content = build_mock_content(paragraph) message = "当前未找到可用模型，返回本地模拟生成结果。" else: result = await call_ai(paragraph, model, file_summaries) content = result.content message = f"已通过模型 {result.used_model} 生成。" return Response( data={ "paragraph_id": paragraph.id, "content": content, "message": message, "file_summaries": file_summaries, } ) @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})