119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
import os
|
|
from datetime import datetime, timezone
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
|
|
from app.tasks.celery_app import celery_app
|
|
from app.core.config import get_settings
|
|
from app.models.generation_task import GenerationTask
|
|
from app.models.generation_point import GenerationPoint
|
|
from app.models.template import Template
|
|
from app.services.document_processor import html_to_docx_bytes
|
|
from app.services.ai_adapter import call_ai_model
|
|
from app.services.ref_parser import parse_reference_file
|
|
from app.services.file_storage import get_storage_dir, RESULTS_DIR
|
|
|
|
|
|
def _create_db_session() -> async_sessionmaker[AsyncSession]:
|
|
settings = get_settings()
|
|
engine = create_async_engine(
|
|
settings.DATABASE_URL,
|
|
echo=False,
|
|
pool_size=5,
|
|
max_overflow=5,
|
|
pool_pre_ping=True,
|
|
)
|
|
return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
|
|
async def _generate_document(task_id: str) -> None:
|
|
session_factory = _create_db_session()
|
|
|
|
async with session_factory() as db:
|
|
task_result = await db.execute(select(GenerationTask).where(GenerationTask.id == task_id))
|
|
task = task_result.scalar_one_or_none()
|
|
if not task:
|
|
return
|
|
|
|
task.status = "processing"
|
|
await db.commit()
|
|
|
|
template_result = await db.execute(select(Template).where(Template.id == task.template_id))
|
|
template = template_result.scalar_one_or_none()
|
|
if not template:
|
|
task.status = "failed"
|
|
task.error_msg = "模板不存在"
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
await db.commit()
|
|
return
|
|
|
|
html_content = template.html_content or ""
|
|
|
|
points_result = await db.execute(
|
|
select(GenerationPoint)
|
|
.where(GenerationPoint.template_id == task.template_id)
|
|
.order_by(GenerationPoint.order.asc())
|
|
)
|
|
points = points_result.scalars().all()
|
|
|
|
try:
|
|
for point in points:
|
|
ref_content = None
|
|
if point.ref_file_path:
|
|
ref_content = await parse_reference_file(point.ref_file_path)
|
|
|
|
model_config = {"provider": "custom", "endpoint": "", "api_key": "", "extra_params": {}}
|
|
if point.model_id:
|
|
from app.models.ai_model import AIModel
|
|
model_result = await db.execute(select(AIModel).where(AIModel.id == point.model_id))
|
|
model = model_result.scalar_one_or_none()
|
|
if model:
|
|
model_config = {
|
|
"provider": model.provider,
|
|
"endpoint": model.endpoint,
|
|
"api_key": model.api_key,
|
|
"extra_params": model.extra_params,
|
|
}
|
|
|
|
ai_result = await call_ai_model(model_config, point.prompt, ref_content)
|
|
|
|
position = point.position
|
|
start = position.get("start", 0)
|
|
end = position.get("end", 0)
|
|
if 0 <= start < end <= len(html_content):
|
|
html_content = html_content[:start] + ai_result + html_content[end:]
|
|
|
|
await db.commit()
|
|
|
|
result_dir = get_storage_dir(RESULTS_DIR)
|
|
result_path = os.path.join(result_dir, f"{task_id}.docx")
|
|
docx_bytes = html_to_docx_bytes(html_content)
|
|
with open(result_path, "wb") as f:
|
|
f.write(docx_bytes)
|
|
|
|
task.status = "done"
|
|
task.result_file_path = result_path
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
await db.commit()
|
|
|
|
except Exception as e:
|
|
task.status = "failed"
|
|
task.error_msg = str(e)
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
await db.commit()
|
|
|
|
await session_factory.engine.dispose()
|
|
|
|
|
|
@celery_app.task(bind=True, name="generate_document")
|
|
def generate_document(self, task_id: str) -> dict:
|
|
import asyncio
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(_generate_document(task_id))
|
|
return {"status": "done", "task_id": task_id}
|
|
except Exception as e:
|
|
return {"status": "failed", "task_id": task_id, "error": str(e)}
|
|
finally:
|
|
loop.close()
|