zwt13703
463 lines
17 KiB
Python
463 lines
17 KiB
Python
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.reference_file import ReferenceFile
|
|
from models.template import Template
|
|
from schemas.schemas import GenerateFullRequest, GenerateTestRequest, ReferenceFileUpdate, Response
|
|
from services.ai_service import call_ai, stream_ai_preview
|
|
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 delete_object, split_bucket_path, upload_bytes, get_presigned_url
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def _serialize_document(document: Document) -> dict:
|
|
request_payload = {}
|
|
if document.request_payload_json:
|
|
try:
|
|
request_payload = json.loads(document.request_payload_json)
|
|
except Exception:
|
|
request_payload = {}
|
|
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,
|
|
"request_payload": request_payload,
|
|
"created_at": document.created_at,
|
|
"updated_at": document.updated_at,
|
|
}
|
|
|
|
|
|
def _serialize_reference_file(file: ReferenceFile) -> dict:
|
|
return {
|
|
"id": file.id,
|
|
"file_name": file.file_name,
|
|
"file_path": file.file_path,
|
|
"file_size": file.file_size,
|
|
"content_type": file.content_type,
|
|
"created_at": file.created_at.isoformat() if file.created_at else None,
|
|
}
|
|
|
|
|
|
async def _build_reference_name_mapping(db: AsyncSession, file_paths: list[str]) -> dict[str, str]:
|
|
if not file_paths:
|
|
return {}
|
|
result = await db.execute(
|
|
select(ReferenceFile.file_path, ReferenceFile.file_name).where(ReferenceFile.file_path.in_(file_paths))
|
|
)
|
|
return {file_path: file_name for file_path, file_name in result.all()}
|
|
|
|
|
|
async def _build_reference_records_mapping(db: AsyncSession, file_paths: list[str]) -> dict[str, ReferenceFile]:
|
|
if not file_paths:
|
|
return {}
|
|
result = await db.execute(select(ReferenceFile).where(ReferenceFile.file_path.in_(file_paths)))
|
|
return {item.file_path: item for item in result.scalars().all()}
|
|
|
|
@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_name_mapping = await _build_reference_name_mapping(db, body.file_paths or [])
|
|
file_summaries = (
|
|
await asyncio.to_thread(summarize_minio_files, body.file_paths or [], file_name_mapping)
|
|
if body.file_paths
|
|
else []
|
|
)
|
|
if model is None or model.status != "enabled":
|
|
content = build_mock_content(paragraph)
|
|
message = "当前未找到可用模型，返回本地模拟生成结果。"
|
|
else:
|
|
setattr(paragraph, "enable_reasoning", bool(model.enable_reasoning))
|
|
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("/test-stream")
|
|
async def generate_test_stream(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)
|
|
|
|
if model is None or model.status != "enabled":
|
|
raise HTTPException(status_code=400, detail="当前段落未配置可用的流式模型")
|
|
if not model.supports_streaming:
|
|
raise HTTPException(status_code=400, detail="当前模型未开启流式传输")
|
|
|
|
file_name_mapping = await _build_reference_name_mapping(db, body.file_paths or [])
|
|
file_summaries = (
|
|
await asyncio.to_thread(summarize_minio_files, body.file_paths or [], file_name_mapping)
|
|
if body.file_paths
|
|
else []
|
|
)
|
|
setattr(paragraph, "enable_reasoning", bool(model.enable_reasoning))
|
|
|
|
async def event_stream():
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps(
|
|
{
|
|
"type": "meta",
|
|
"message": f"正在通过模型 {model.name} 流式生成...",
|
|
"file_summaries": file_summaries,
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
}
|
|
try:
|
|
async for chunk in stream_ai_preview(paragraph, model, file_summaries):
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps({"type": "delta", "content": chunk}, ensure_ascii=False),
|
|
}
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps({"type": "done"}, ensure_ascii=False),
|
|
}
|
|
except Exception as error:
|
|
fallback_message = str(error)
|
|
if "503" in fallback_message or "temporarily unavailable" in fallback_message.lower():
|
|
try:
|
|
result = await call_ai(paragraph, model, file_summaries)
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps(
|
|
{
|
|
"type": "meta",
|
|
"message": "流式通道暂时不可用，已自动回退为普通返回。",
|
|
"file_summaries": file_summaries,
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
}
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps({"type": "delta", "content": result.raw_text}, ensure_ascii=False),
|
|
}
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps({"type": "done"}, ensure_ascii=False),
|
|
}
|
|
return
|
|
except Exception as fallback_error:
|
|
fallback_message = f"{fallback_message}；普通调用回退也失败：{fallback_error}"
|
|
yield {
|
|
"event": "message",
|
|
"data": json.dumps({"type": "error", "message": fallback_message}, ensure_ascii=False),
|
|
}
|
|
|
|
return EventSourceResponse(
|
|
event_stream(),
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
@router.post("/upload")
|
|
async def upload_reference_file(file: UploadFile = File(...), db: AsyncSession = Depends(get_db)):
|
|
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:
|
|
allowed = " / ".join(settings.ALLOWED_EXTENSIONS)
|
|
raise HTTPException(status_code=400, detail=f"文件类型不支持：{ext or '无扩展名'}。当前支持：{allowed}")
|
|
|
|
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",
|
|
)
|
|
record = ReferenceFile(
|
|
file_name=file.filename,
|
|
file_path=f"{settings.MINIO_BUCKET_UPLOADS}/{object_name}",
|
|
file_size=len(content),
|
|
content_type=file.content_type or "application/octet-stream",
|
|
)
|
|
db.add(record)
|
|
await db.commit()
|
|
await db.refresh(record)
|
|
return Response(
|
|
data=_serialize_reference_file(record)
|
|
)
|
|
|
|
|
|
@router.get("/reference-files")
|
|
async def list_reference_files(
|
|
page: int = Query(1, ge=1),
|
|
page_size: int = Query(20, ge=1, le=100),
|
|
keyword: str = Query("", description="按文件名搜索"),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
stmt = select(ReferenceFile)
|
|
count_stmt = select(func.count(ReferenceFile.id))
|
|
if keyword:
|
|
like_keyword = f"%{keyword.strip()}%"
|
|
stmt = stmt.where(ReferenceFile.file_name.like(like_keyword))
|
|
count_stmt = count_stmt.where(ReferenceFile.file_name.like(like_keyword))
|
|
|
|
total = (await db.execute(count_stmt)).scalar_one()
|
|
result = await db.execute(
|
|
stmt.order_by(ReferenceFile.id.desc()).offset((page - 1) * page_size).limit(page_size)
|
|
)
|
|
items = [_serialize_reference_file(item) for item in result.scalars().all()]
|
|
return Response(data={"items": items, "total": total, "page": page, "page_size": page_size})
|
|
|
|
|
|
@router.get("/reference-files/{file_id}/download")
|
|
async def download_reference_file(file_id: int, db: AsyncSession = Depends(get_db)):
|
|
record = await db.get(ReferenceFile, file_id)
|
|
if record is None:
|
|
raise HTTPException(status_code=404, detail="附件记录不存在")
|
|
bucket, object_name = split_bucket_path(record.file_path)
|
|
return Response(data={"url": get_presigned_url(bucket, object_name), "file_name": record.file_name})
|
|
|
|
|
|
@router.put("/reference-files/{file_id}")
|
|
async def update_reference_file(file_id: int, body: ReferenceFileUpdate, db: AsyncSession = Depends(get_db)):
|
|
record = await db.get(ReferenceFile, file_id)
|
|
if record is None:
|
|
raise HTTPException(status_code=404, detail="附件记录不存在")
|
|
record.file_name = body.file_name.strip()
|
|
await db.commit()
|
|
await db.refresh(record)
|
|
return Response(data=_serialize_reference_file(record))
|
|
|
|
|
|
@router.delete("/reference-files/{file_id}")
|
|
async def delete_reference_file(file_id: int, db: AsyncSession = Depends(get_db)):
|
|
record = await db.get(ReferenceFile, file_id)
|
|
if record is None:
|
|
raise HTTPException(status_code=404, detail="附件记录不存在")
|
|
try:
|
|
bucket, object_name = split_bucket_path(record.file_path)
|
|
await asyncio.to_thread(delete_object, bucket, object_name)
|
|
except Exception:
|
|
pass
|
|
await db.delete(record)
|
|
await db.commit()
|
|
return Response(data={"id": file_id})
|
|
|
|
|
|
@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="模板下暂无可生成段落")
|
|
|
|
normalized_file_map: dict[str, list[str]] = {}
|
|
for key, values in (body.file_map or {}).items():
|
|
if not isinstance(values, list):
|
|
continue
|
|
normalized_file_map[str(key)] = [item for item in values if item]
|
|
all_file_paths = [item for values in normalized_file_map.values() for item in values]
|
|
file_records = await _build_reference_records_mapping(db, all_file_paths)
|
|
paragraph_snapshot = []
|
|
for paragraph in paragraphs:
|
|
selected_paths = normalized_file_map.get(str(paragraph.id), [])
|
|
paragraph_snapshot.append(
|
|
{
|
|
"paragraph_id": paragraph.id,
|
|
"title": paragraph.title,
|
|
"sort_index": paragraph.sort_index,
|
|
"need_file": bool(paragraph.need_file),
|
|
"file_note": paragraph.file_note or "",
|
|
"selected_files": [
|
|
_serialize_reference_file(file_records[file_path])
|
|
for file_path in selected_paths
|
|
if file_path in file_records
|
|
],
|
|
}
|
|
)
|
|
request_payload = {
|
|
"template_name": template.name,
|
|
"file_map": normalized_file_map,
|
|
"paragraphs": paragraph_snapshot,
|
|
}
|
|
|
|
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="",
|
|
request_payload_json=json.dumps(request_payload, ensure_ascii=False),
|
|
)
|
|
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})
|