init: 初始化项目
This commit is contained in:
zwt13703
@@ -0,0 +1,9 @@
|
||||
venv/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.env
|
||||
.git
|
||||
storage/
|
||||
alembic/versions/
|
||||
*.egg-info/
|
||||
.pytest_cache/
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"hash": "1e9afdbb",
|
||||
"configHash": "ac7453e4",
|
||||
"lockfileHash": "e3b0c442",
|
||||
"browserHash": "29840d3b",
|
||||
"optimized": {},
|
||||
"chunks": {}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"type": "module"
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
FROM python:3.12-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
gcc libpq-dev && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN mkdir -p /app/storage/templates /app/storage/ref_files /app/storage/results
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
@@ -0,0 +1,19 @@
|
||||
FROM python:3.12-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
gcc libpq-dev && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN mkdir -p /app/storage/templates /app/storage/ref_files /app/storage/results
|
||||
|
||||
CMD ["celery", "-A", "app.tasks.celery_app", "worker", "--loglevel=info", "--concurrency=4"]
|
||||
@@ -0,0 +1,37 @@
|
||||
[alembic]
|
||||
script_location = alembic
|
||||
prepend_sys_path = .
|
||||
sqlalchemy.url = postgresql://docforge:docforge@localhost:5432/docforge
|
||||
|
||||
[loggers]
|
||||
keys = root,sqlalchemy,alembic
|
||||
|
||||
[handlers]
|
||||
keys = console
|
||||
|
||||
[formatters]
|
||||
keys = generic
|
||||
|
||||
[logger_root]
|
||||
level = WARN
|
||||
handlers = console
|
||||
|
||||
[logger_sqlalchemy]
|
||||
level = WARN
|
||||
handlers =
|
||||
qualname = sqlalchemy.engine
|
||||
|
||||
[logger_alembic]
|
||||
level = INFO
|
||||
handlers =
|
||||
qualname = alembic
|
||||
|
||||
[handler_console]
|
||||
class = StreamHandler
|
||||
args = (sys.stderr,)
|
||||
level = NOTSET
|
||||
formatter = generic
|
||||
|
||||
[formatter_generic]
|
||||
format = %(levelname)-5.5s [%(name)s] %(message)s
|
||||
datefmt = %H:%M:%S
|
||||
@@ -0,0 +1,51 @@
|
||||
from logging.config import fileConfig
|
||||
from sqlalchemy import engine_from_config, pool
|
||||
from alembic import context
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
|
||||
from app.core.config import get_settings
|
||||
from app.models import Base
|
||||
|
||||
config = context.config
|
||||
|
||||
|
||||
settings = get_settings()
|
||||
config.set_main_option("sqlalchemy.url", settings.DATABASE_URL_SYNC)
|
||||
|
||||
if config.config_file_name is not None:
|
||||
fileConfig(config.config_file_name)
|
||||
|
||||
target_metadata = Base.metadata
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
url = config.get_main_option("sqlalchemy.url")
|
||||
context.configure(
|
||||
url=url,
|
||||
target_metadata=target_metadata,
|
||||
literal_binds=True,
|
||||
dialect_opts={"paramstyle": "named"},
|
||||
)
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
def run_migrations_online() -> None:
|
||||
connectable = engine_from_config(
|
||||
config.get_section(config.config_ini_section, {}),
|
||||
prefix="sqlalchemy.",
|
||||
poolclass=pool.NullPool,
|
||||
)
|
||||
with connectable.connect() as connection:
|
||||
context.configure(connection=connection, target_metadata=target_metadata)
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
if context.is_offline_mode():
|
||||
run_migrations_offline()
|
||||
else:
|
||||
run_migrations_online()
|
||||
@@ -0,0 +1,23 @@
|
||||
"""${message}
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
${imports if imports else ""}
|
||||
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: Union[str, None] = ${repr(down_revision)}
|
||||
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
${downgrades if downgrades else "pass"}
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Add selected_text to generation_points
|
||||
|
||||
Revision ID: 011bde04d871
|
||||
Revises: 857f5a3874dd
|
||||
Create Date: 2026-07-06 18:19:15.771438
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
revision: str = '011bde04d871'
|
||||
down_revision: Union[str, None] = '857f5a3874dd'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.add_column('generation_points', sa.Column('selected_text', sa.Text(), nullable=True))
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_column('generation_points', 'selected_text')
|
||||
# ### end Alembic commands ###
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Initial schema
|
||||
|
||||
Revision ID: 857f5a3874dd
|
||||
Revises:
|
||||
Create Date: 2026-07-06 16:42:37.845634
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
revision: str = '857f5a3874dd'
|
||||
down_revision: Union[str, None] = None
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table('models',
|
||||
sa.Column('name', sa.String(length=100), nullable=False),
|
||||
sa.Column('provider', sa.String(length=50), nullable=False),
|
||||
sa.Column('endpoint', sa.String(length=255), nullable=False),
|
||||
sa.Column('api_key', sa.Text(), nullable=False),
|
||||
sa.Column('extra_params', postgresql.JSONB(astext_type=sa.Text()), nullable=False),
|
||||
sa.Column('is_enabled', sa.Boolean(), nullable=False),
|
||||
sa.Column('remark', sa.Text(), nullable=True),
|
||||
sa.Column('id', sa.UUID(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_table('system_config',
|
||||
sa.Column('key', sa.String(length=50), nullable=False),
|
||||
sa.Column('value', sa.Text(), nullable=True),
|
||||
sa.Column('description', sa.String(length=200), nullable=True),
|
||||
sa.PrimaryKeyConstraint('key')
|
||||
)
|
||||
op.create_table('templates',
|
||||
sa.Column('name', sa.String(length=200), nullable=False),
|
||||
sa.Column('file_path', sa.String(length=500), nullable=False),
|
||||
sa.Column('html_content', sa.Text(), nullable=True),
|
||||
sa.Column('id', sa.UUID(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_table('generation_points',
|
||||
sa.Column('template_id', sa.UUID(), nullable=False),
|
||||
sa.Column('position', postgresql.JSONB(astext_type=sa.Text()), nullable=False),
|
||||
sa.Column('prompt', sa.Text(), nullable=False),
|
||||
sa.Column('model_id', sa.UUID(), nullable=True),
|
||||
sa.Column('ref_file_path', sa.String(length=500), nullable=True),
|
||||
sa.Column('order', sa.Integer(), nullable=False),
|
||||
sa.Column('id', sa.UUID(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
|
||||
sa.ForeignKeyConstraint(['model_id'], ['models.id'], ondelete='SET NULL'),
|
||||
sa.ForeignKeyConstraint(['template_id'], ['templates.id'], ondelete='CASCADE'),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_generation_points_template_id'), 'generation_points', ['template_id'], unique=False)
|
||||
op.create_table('generation_tasks',
|
||||
sa.Column('template_id', sa.UUID(), nullable=False),
|
||||
sa.Column('status', sa.String(length=20), nullable=False),
|
||||
sa.Column('result_file_path', sa.String(length=500), nullable=True),
|
||||
sa.Column('error_msg', sa.Text(), nullable=True),
|
||||
sa.Column('celery_task_id', sa.String(length=100), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column('finished_at', sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column('id', sa.UUID(), nullable=False),
|
||||
sa.ForeignKeyConstraint(['template_id'], ['templates.id'], ondelete='CASCADE'),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_generation_tasks_template_id'), 'generation_tasks', ['template_id'], unique=False)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_index(op.f('ix_generation_tasks_template_id'), table_name='generation_tasks')
|
||||
op.drop_table('generation_tasks')
|
||||
op.drop_index(op.f('ix_generation_points_template_id'), table_name='generation_points')
|
||||
op.drop_table('generation_points')
|
||||
op.drop_table('templates')
|
||||
op.drop_table('system_config')
|
||||
op.drop_table('models')
|
||||
# ### end Alembic commands ###
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Change ref_file_path to Text for multi-file support
|
||||
|
||||
Revision ID: 9066ecf60786
|
||||
Revises: b4b91d4466e4
|
||||
Create Date: 2026-07-06 21:01:41.552375
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
revision: str = '9066ecf60786'
|
||||
down_revision: Union[str, None] = 'b4b91d4466e4'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.alter_column('generation_points', 'ref_file_path',
|
||||
existing_type=sa.VARCHAR(length=500),
|
||||
type_=sa.Text(),
|
||||
existing_nullable=True)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.alter_column('generation_points', 'ref_file_path',
|
||||
existing_type=sa.Text(),
|
||||
type_=sa.VARCHAR(length=500),
|
||||
existing_nullable=True)
|
||||
# ### end Alembic commands ###
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Add need_ref_file and remark to generation_points
|
||||
|
||||
Revision ID: b4b91d4466e4
|
||||
Revises: 011bde04d871
|
||||
Create Date: 2026-07-06 18:27:03.054325
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
revision: str = 'b4b91d4466e4'
|
||||
down_revision: Union[str, None] = '011bde04d871'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.add_column('generation_points', sa.Column('need_ref_file', sa.Boolean(), nullable=False))
|
||||
op.add_column('generation_points', sa.Column('remark', sa.Text(), nullable=True))
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_column('generation_points', 'remark')
|
||||
op.drop_column('generation_points', 'need_ref_file')
|
||||
# ### end Alembic commands ###
|
||||
@@ -0,0 +1,128 @@
|
||||
import uuid
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form, Query
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import select
|
||||
from app.core.database import get_db
|
||||
from app.core.security_middleware import validate_file_extension
|
||||
from app.models.generation_point import GenerationPoint
|
||||
from app.models.template import Template
|
||||
from app.models.ai_model import AIModel
|
||||
from app.schemas.generation_point import (
|
||||
GenerationPointCreate,
|
||||
GenerationPointUpdate,
|
||||
GenerationPointResponse,
|
||||
BatchOrderUpdate,
|
||||
)
|
||||
from app.services.file_storage import save_upload, REF_FILES_DIR
|
||||
|
||||
router = APIRouter(prefix="/generation-points", tags=["生成点管理"])
|
||||
|
||||
|
||||
@router.post("", response_model=GenerationPointResponse)
|
||||
async def create_generation_point(
|
||||
template_id: uuid.UUID = Form(...),
|
||||
position: str = Form(...),
|
||||
prompt: str = Form(...),
|
||||
model_id: uuid.UUID | None = Form(None),
|
||||
order: int = Form(0),
|
||||
selected_text: str | None = Form(None),
|
||||
need_ref_file: bool = Form(False),
|
||||
remark: str | None = Form(None),
|
||||
ref_file: UploadFile | None = File(None),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
import json
|
||||
|
||||
template_result = await db.execute(select(Template).where(Template.id == template_id))
|
||||
if not template_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=404, detail="模板不存在")
|
||||
|
||||
if model_id:
|
||||
model_result = await db.execute(select(AIModel).where(AIModel.id == model_id))
|
||||
if not model_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
|
||||
ref_file_path = None
|
||||
if ref_file and ref_file.filename:
|
||||
validate_file_extension(ref_file.filename)
|
||||
ref_file_path = await save_upload(ref_file, REF_FILES_DIR)
|
||||
|
||||
point = GenerationPoint(
|
||||
template_id=template_id,
|
||||
position=json.loads(position),
|
||||
prompt=prompt,
|
||||
model_id=model_id,
|
||||
order=order,
|
||||
ref_file_path=ref_file_path,
|
||||
selected_text=selected_text,
|
||||
need_ref_file=need_ref_file,
|
||||
remark=remark,
|
||||
)
|
||||
db.add(point)
|
||||
await db.flush()
|
||||
await db.refresh(point)
|
||||
return point
|
||||
|
||||
|
||||
@router.get("", response_model=list[GenerationPointResponse])
|
||||
async def list_generation_points(
|
||||
template_id: uuid.UUID = Query(...),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
query = (
|
||||
select(GenerationPoint)
|
||||
.where(GenerationPoint.template_id == template_id)
|
||||
.order_by(GenerationPoint.order.asc(), GenerationPoint.created_at.asc())
|
||||
)
|
||||
result = await db.execute(query)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
@router.get("/{point_id}", response_model=GenerationPointResponse)
|
||||
async def get_generation_point(point_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(GenerationPoint).where(GenerationPoint.id == point_id))
|
||||
point = result.scalar_one_or_none()
|
||||
if not point:
|
||||
raise HTTPException(status_code=404, detail="生成点不存在")
|
||||
return point
|
||||
|
||||
|
||||
@router.put("/{point_id}", response_model=GenerationPointResponse)
|
||||
async def update_generation_point(
|
||||
point_id: str,
|
||||
data: GenerationPointUpdate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
result = await db.execute(select(GenerationPoint).where(GenerationPoint.id == point_id))
|
||||
point = result.scalar_one_or_none()
|
||||
if not point:
|
||||
raise HTTPException(status_code=404, detail="生成点不存在")
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
for key, value in update_data.items():
|
||||
setattr(point, key, value)
|
||||
|
||||
await db.flush()
|
||||
await db.refresh(point)
|
||||
return point
|
||||
|
||||
|
||||
@router.delete("/{point_id}")
|
||||
async def delete_generation_point(point_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(GenerationPoint).where(GenerationPoint.id == point_id))
|
||||
point = result.scalar_one_or_none()
|
||||
if not point:
|
||||
raise HTTPException(status_code=404, detail="生成点不存在")
|
||||
await db.delete(point)
|
||||
return {"detail": "删除成功"}
|
||||
|
||||
|
||||
@router.post("/batch-order")
|
||||
async def batch_update_order(data: BatchOrderUpdate, db: AsyncSession = Depends(get_db)):
|
||||
for item in data.points:
|
||||
result = await db.execute(select(GenerationPoint).where(GenerationPoint.id == item["id"]))
|
||||
point = result.scalar_one_or_none()
|
||||
if point:
|
||||
point.order = item["order"]
|
||||
await db.flush()
|
||||
return {"detail": "排序更新成功"}
|
||||
@@ -0,0 +1,115 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import select
|
||||
from app.core.database import get_db
|
||||
from app.core.security import encrypt_api_key
|
||||
from app.models.ai_model import AIModel
|
||||
from app.schemas.ai_model import AIModelCreate, AIModelUpdate, AIModelToggle, AIModelResponse
|
||||
from app.services.ai_adapter import call_ai_model
|
||||
|
||||
router = APIRouter(prefix="/models", tags=["AI模型管理"])
|
||||
|
||||
|
||||
@router.post("", response_model=AIModelResponse)
|
||||
async def create_model(data: AIModelCreate, db: AsyncSession = Depends(get_db)):
|
||||
encrypted_key = encrypt_api_key(data.api_key)
|
||||
model = AIModel(
|
||||
name=data.name,
|
||||
provider=data.provider,
|
||||
endpoint=data.endpoint,
|
||||
api_key=encrypted_key,
|
||||
extra_params=data.extra_params,
|
||||
is_enabled=data.is_enabled,
|
||||
remark=data.remark,
|
||||
)
|
||||
db.add(model)
|
||||
await db.flush()
|
||||
await db.refresh(model)
|
||||
return model
|
||||
|
||||
|
||||
@router.get("", response_model=list[AIModelResponse])
|
||||
async def list_models(
|
||||
enabled: bool | None = Query(None, description="过滤启用/禁用"),
|
||||
skip: int = Query(0, ge=0),
|
||||
limit: int = Query(20, ge=1, le=100),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
query = select(AIModel)
|
||||
if enabled is not None:
|
||||
query = query.where(AIModel.is_enabled == enabled)
|
||||
query = query.offset(skip).limit(limit).order_by(AIModel.created_at.desc())
|
||||
result = await db.execute(query)
|
||||
models = result.scalars().all()
|
||||
return models
|
||||
|
||||
|
||||
@router.get("/{model_id}", response_model=AIModelResponse)
|
||||
async def get_model(model_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(AIModel).where(AIModel.id == model_id))
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
return model
|
||||
|
||||
|
||||
@router.put("/{model_id}", response_model=AIModelResponse)
|
||||
async def update_model(model_id: str, data: AIModelUpdate, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(AIModel).where(AIModel.id == model_id))
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
if "api_key" in update_data and update_data["api_key"] is not None:
|
||||
update_data["api_key"] = encrypt_api_key(update_data["api_key"])
|
||||
|
||||
for key, value in update_data.items():
|
||||
setattr(model, key, value)
|
||||
|
||||
await db.flush()
|
||||
await db.refresh(model)
|
||||
return model
|
||||
|
||||
|
||||
@router.delete("/{model_id}")
|
||||
async def delete_model(model_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(AIModel).where(AIModel.id == model_id))
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
await db.delete(model)
|
||||
return {"detail": "删除成功"}
|
||||
|
||||
|
||||
@router.patch("/{model_id}/toggle", response_model=AIModelResponse)
|
||||
async def toggle_model(model_id: str, data: AIModelToggle, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(AIModel).where(AIModel.id == model_id))
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
model.is_enabled = data.is_enabled
|
||||
await db.flush()
|
||||
await db.refresh(model)
|
||||
return model
|
||||
|
||||
|
||||
@router.post("/{model_id}/test")
|
||||
async def test_model(model_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(AIModel).where(AIModel.id == model_id))
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
|
||||
model_config = {
|
||||
"provider": model.provider,
|
||||
"endpoint": model.endpoint,
|
||||
"api_key": model.api_key,
|
||||
"extra_params": model.extra_params,
|
||||
}
|
||||
|
||||
try:
|
||||
response = await call_ai_model(model_config, "请用一句话介绍你自己。")
|
||||
return {"success": True, "result": response}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
@@ -0,0 +1,16 @@
|
||||
from fastapi import APIRouter
|
||||
from app.api.models import router as models_router
|
||||
from app.api.templates import router as templates_router
|
||||
from app.api.generation_points import router as generation_points_router
|
||||
from app.api.tasks import router as tasks_router
|
||||
from app.api.settings import router as settings_router
|
||||
from app.core.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
api_router = APIRouter(prefix=settings.API_V1_PREFIX)
|
||||
api_router.include_router(models_router)
|
||||
api_router.include_router(templates_router)
|
||||
api_router.include_router(generation_points_router)
|
||||
api_router.include_router(tasks_router, prefix="")
|
||||
api_router.include_router(settings_router)
|
||||
@@ -0,0 +1,34 @@
|
||||
from fastapi import APIRouter, Depends
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import select
|
||||
from app.core.database import get_db
|
||||
from app.models.system_config import SystemConfig
|
||||
|
||||
router = APIRouter(prefix="/settings", tags=["系统设置"])
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def get_all_settings(db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(SystemConfig))
|
||||
configs = result.scalars().all()
|
||||
return {c.key: c.value for c in configs}
|
||||
|
||||
|
||||
@router.put("/{key}")
|
||||
async def update_setting(key: str, body: dict, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(SystemConfig).where(SystemConfig.key == key))
|
||||
config = result.scalar_one_or_none()
|
||||
|
||||
value = body.get("value", "")
|
||||
description = body.get("description", "")
|
||||
|
||||
if config:
|
||||
config.value = value
|
||||
if description:
|
||||
config.description = description
|
||||
else:
|
||||
config = SystemConfig(key=key, value=value, description=description)
|
||||
db.add(config)
|
||||
|
||||
await db.flush()
|
||||
return {"key": key, "value": value}
|
||||
@@ -0,0 +1,233 @@
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile, File
|
||||
from fastapi.responses import Response
|
||||
from typing import List
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import select
|
||||
from app.core.database import get_db
|
||||
from app.core.security_middleware import validate_file_extension
|
||||
from app.models.generation_task import GenerationTask
|
||||
from app.models.generation_point import GenerationPoint
|
||||
from app.models.template import Template
|
||||
from app.models.ai_model import AIModel
|
||||
from app.schemas.generation_task import TaskResponse, GenerateResponse, SingleTestRequest, SingleTestResponse, TaskDetailResponse
|
||||
from app.services.document_processor import docx_to_pdf_bytes
|
||||
from app.services.file_storage import get_file_content, save_upload, REF_FILES_DIR
|
||||
from app.services.ai_adapter import call_ai_model
|
||||
from app.services.ref_parser import parse_reference_file, parse_reference_files
|
||||
from app.tasks.generate import generate_document
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/templates/{template_id}/generate", response_model=GenerateResponse)
|
||||
async def trigger_generation(template_id: str, db: AsyncSession = Depends(get_db)):
|
||||
template_result = await db.execute(select(Template).where(Template.id == template_id))
|
||||
if not template_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=404, detail="模板不存在")
|
||||
|
||||
points_result = await db.execute(
|
||||
select(GenerationPoint).where(GenerationPoint.template_id == template_id)
|
||||
)
|
||||
points = points_result.scalars().all()
|
||||
if not points:
|
||||
raise HTTPException(status_code=400, detail="模板没有生成点")
|
||||
|
||||
task = GenerationTask(
|
||||
template_id=uuid.UUID(template_id),
|
||||
status="pending",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
db.add(task)
|
||||
await db.flush()
|
||||
|
||||
celery_task = generate_document.delay(str(task.id))
|
||||
task.celery_task_id = celery_task.id
|
||||
await db.commit()
|
||||
|
||||
return GenerateResponse(task_id=task.id, status="pending")
|
||||
|
||||
|
||||
@router.get("/tasks/{task_id}")
|
||||
async def get_task_status(task_id: str, detail: bool = Query(False), db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(GenerationTask).where(GenerationTask.id == task_id))
|
||||
task = result.scalar_one_or_none()
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
|
||||
if detail:
|
||||
tpl_result = await db.execute(select(Template).where(Template.id == task.template_id))
|
||||
tpl = tpl_result.scalar_one_or_none()
|
||||
points_result = await db.execute(
|
||||
select(GenerationPoint).where(GenerationPoint.template_id == task.template_id).order_by(GenerationPoint.order.asc())
|
||||
)
|
||||
points = points_result.scalars().all()
|
||||
return TaskDetailResponse(
|
||||
id=task.id,
|
||||
template_id=task.template_id,
|
||||
template_name=tpl.name if tpl else "",
|
||||
status=task.status,
|
||||
result_file_path=task.result_file_path,
|
||||
error_msg=task.error_msg,
|
||||
created_at=task.created_at,
|
||||
finished_at=task.finished_at,
|
||||
points=[{
|
||||
"id": str(p.id),
|
||||
"prompt": p.prompt,
|
||||
"position": p.position,
|
||||
"model_id": str(p.model_id) if p.model_id else None,
|
||||
"ref_file_path": p.ref_file_path,
|
||||
"selected_text": p.selected_text,
|
||||
"need_ref_file": p.need_ref_file,
|
||||
"remark": p.remark,
|
||||
"order": p.order,
|
||||
} for p in points],
|
||||
)
|
||||
return task
|
||||
|
||||
|
||||
@router.get("/tasks/{task_id}/download")
|
||||
async def download_task_result(
|
||||
task_id: str,
|
||||
format: str = Query("docx", pattern="^(docx|pdf)$"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
result = await db.execute(select(GenerationTask).where(GenerationTask.id == task_id))
|
||||
task = result.scalar_one_or_none()
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
|
||||
if task.status != "done" or not task.result_file_path:
|
||||
raise HTTPException(status_code=400, detail="任务未完成或无结果文件")
|
||||
|
||||
file_content = await get_file_content(task.result_file_path)
|
||||
|
||||
if format == "pdf":
|
||||
file_content = docx_to_pdf_bytes(file_content)
|
||||
media_type = "application/pdf"
|
||||
filename = f"result_{task_id}.pdf"
|
||||
else:
|
||||
media_type = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
filename = f"result_{task_id}.docx"
|
||||
|
||||
return Response(
|
||||
content=file_content,
|
||||
media_type=media_type,
|
||||
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/tasks")
|
||||
async def list_tasks(
|
||||
skip: int = Query(0, ge=0),
|
||||
limit: int = Query(20, ge=1, le=100),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
query = select(GenerationTask).offset(skip).limit(limit).order_by(GenerationTask.created_at.desc())
|
||||
result = await db.execute(query)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
@router.post("/tasks/{task_id}/cancel")
|
||||
async def cancel_task(task_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(GenerationTask).where(GenerationTask.id == task_id))
|
||||
task = result.scalar_one_or_none()
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
|
||||
if task.status not in ("pending", "processing"):
|
||||
raise HTTPException(status_code=400, detail="任务无法取消")
|
||||
|
||||
if task.celery_task_id:
|
||||
from app.tasks.celery_app import celery_app
|
||||
celery_app.control.revoke(task.celery_task_id, terminate=True)
|
||||
|
||||
task.status = "failed"
|
||||
task.error_msg = "用户取消"
|
||||
task.finished_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
return {"detail": "任务已取消"}
|
||||
|
||||
|
||||
@router.post("/generation-points/{point_id}/test", response_model=SingleTestResponse)
|
||||
async def test_single_point(
|
||||
point_id: str,
|
||||
ref_files: List[UploadFile] = File(default=[]),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
point_result = await db.execute(select(GenerationPoint).where(GenerationPoint.id == point_id))
|
||||
point = point_result.scalar_one_or_none()
|
||||
if not point:
|
||||
raise HTTPException(status_code=404, detail="生成点不存在")
|
||||
|
||||
tmp_paths = []
|
||||
if point.need_ref_file:
|
||||
if not ref_files and not point.ref_file_path:
|
||||
raise HTTPException(status_code=400, detail="此生成点需要上传参考文件")
|
||||
|
||||
for f in ref_files:
|
||||
if f.filename:
|
||||
validate_file_extension(f.filename)
|
||||
path = await save_upload(f, REF_FILES_DIR)
|
||||
tmp_paths.append(path)
|
||||
|
||||
model_config = {"provider": "custom", "endpoint": "", "api_key": "", "extra_params": {}}
|
||||
if point.model_id:
|
||||
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,
|
||||
}
|
||||
|
||||
ref_content = ""
|
||||
if tmp_paths:
|
||||
ref_content = await parse_reference_files(json.dumps(tmp_paths))
|
||||
elif point.ref_file_path:
|
||||
ref_content = await parse_reference_files(point.ref_file_path)
|
||||
|
||||
try:
|
||||
ai_result = await call_ai_model(model_config, point.prompt, ref_content)
|
||||
return SingleTestResponse(result=ai_result)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"AI 调用失败: {str(e)}")
|
||||
|
||||
|
||||
@router.post("/generation-points/{point_id}/upload-ref")
|
||||
async def upload_ref_files(
|
||||
point_id: str,
|
||||
ref_files: List[UploadFile] = File(...),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
point_result = await db.execute(select(GenerationPoint).where(GenerationPoint.id == point_id))
|
||||
point = point_result.scalar_one_or_none()
|
||||
if not point:
|
||||
raise HTTPException(status_code=404, detail="生成点不存在")
|
||||
|
||||
paths = []
|
||||
for f in ref_files:
|
||||
if f.filename:
|
||||
validate_file_extension(f.filename)
|
||||
path = await save_upload(f, REF_FILES_DIR)
|
||||
paths.append(path)
|
||||
|
||||
existing = []
|
||||
if point.ref_file_path:
|
||||
try:
|
||||
existing = json.loads(point.ref_file_path)
|
||||
if not isinstance(existing, list):
|
||||
existing = [point.ref_file_path]
|
||||
except json.JSONDecodeError:
|
||||
existing = [point.ref_file_path] if point.ref_file_path else []
|
||||
|
||||
all_paths = existing + paths
|
||||
point.ref_file_path = json.dumps(all_paths)
|
||||
await db.commit()
|
||||
|
||||
return {"success": True, "count": len(all_paths)}
|
||||
@@ -0,0 +1,150 @@
|
||||
import uuid
|
||||
import bleach
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form, Query
|
||||
from fastapi.responses import Response
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import select
|
||||
from app.core.database import get_db
|
||||
from app.core.security_middleware import validate_file_extension
|
||||
from app.models.template import Template
|
||||
from app.schemas.template import TemplateResponse, TemplateListItem, HTMLContentResponse, HTMLUpdateRequest
|
||||
from app.services.file_storage import save_upload, get_file_content, delete_file, TEMPLATES_DIR
|
||||
from app.services.document_processor import docx_to_html, html_to_docx_bytes, docx_to_pdf_bytes
|
||||
|
||||
ALLOWED_TAGS = [
|
||||
"p", "div", "span", "br", "hr",
|
||||
"h1", "h2", "h3", "h4", "h5", "h6",
|
||||
"ul", "ol", "li",
|
||||
"a", "img", "table", "thead", "tbody", "tr", "td", "th",
|
||||
"b", "i", "u", "strong", "em", "del", "sub", "sup",
|
||||
"pre", "code", "blockquote",
|
||||
]
|
||||
ALLOWED_ATTRS = {
|
||||
"a": ["href", "title", "target"],
|
||||
"img": ["src", "alt", "width", "height"],
|
||||
"td": ["colspan", "rowspan"],
|
||||
"th": ["colspan", "rowspan"],
|
||||
"p": ["style"],
|
||||
"span": ["style"],
|
||||
"div": ["style"],
|
||||
"table": ["style"],
|
||||
}
|
||||
|
||||
router = APIRouter(prefix="/templates", tags=["模板管理"])
|
||||
|
||||
|
||||
@router.post("", response_model=TemplateResponse)
|
||||
async def upload_template(
|
||||
file: UploadFile = File(...),
|
||||
name: str | None = Form(None),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
if not file.filename or not file.filename.endswith(".docx"):
|
||||
raise HTTPException(status_code=400, detail="仅支持 .docx 文件")
|
||||
|
||||
validate_file_extension(file.filename)
|
||||
file_path = await save_upload(file, TEMPLATES_DIR)
|
||||
file_content = await get_file_content(file_path)
|
||||
html_content = await docx_to_html(file_content)
|
||||
|
||||
template = Template(
|
||||
name=name or file.filename.replace(".docx", ""),
|
||||
file_path=file_path,
|
||||
html_content=html_content,
|
||||
)
|
||||
db.add(template)
|
||||
await db.flush()
|
||||
await db.refresh(template)
|
||||
return template
|
||||
|
||||
|
||||
@router.get("", response_model=list[TemplateListItem])
|
||||
async def list_templates(
|
||||
skip: int = Query(0, ge=0),
|
||||
limit: int = Query(20, ge=1, le=100),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
query = select(Template).offset(skip).limit(limit).order_by(Template.created_at.desc())
|
||||
result = await db.execute(query)
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
@router.get("/{template_id}", response_model=TemplateResponse)
|
||||
async def get_template(template_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(Template).where(Template.id == template_id))
|
||||
template = result.scalar_one_or_none()
|
||||
if not template:
|
||||
raise HTTPException(status_code=404, detail="模板不存在")
|
||||
return template
|
||||
|
||||
|
||||
@router.get("/{template_id}/html", response_model=HTMLContentResponse)
|
||||
async def get_template_html(template_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(Template).where(Template.id == template_id))
|
||||
template = result.scalar_one_or_none()
|
||||
if not template:
|
||||
raise HTTPException(status_code=404, detail="模板不存在")
|
||||
return HTMLContentResponse(html_content=template.html_content or "")
|
||||
|
||||
|
||||
@router.put("/{template_id}/html")
|
||||
async def update_template_html(
|
||||
template_id: str,
|
||||
data: HTMLUpdateRequest,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
result = await db.execute(select(Template).where(Template.id == template_id))
|
||||
template = result.scalar_one_or_none()
|
||||
if not template:
|
||||
raise HTTPException(status_code=404, detail="模板不存在")
|
||||
|
||||
template.html_content = bleach.clean(
|
||||
data.html_content, tags=ALLOWED_TAGS, attributes=ALLOWED_ATTRS, strip=True
|
||||
)
|
||||
|
||||
docx_bytes = html_to_docx_bytes(template.html_content)
|
||||
with open(template.file_path, "wb") as f:
|
||||
f.write(docx_bytes)
|
||||
|
||||
await db.flush()
|
||||
return {"detail": "保存成功"}
|
||||
|
||||
|
||||
@router.get("/{template_id}/download")
|
||||
async def download_template(
|
||||
template_id: str,
|
||||
format: str = Query("docx", pattern="^(docx|pdf)$"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
result = await db.execute(select(Template).where(Template.id == template_id))
|
||||
template = result.scalar_one_or_none()
|
||||
if not template:
|
||||
raise HTTPException(status_code=404, detail="模板不存在")
|
||||
|
||||
file_content = await get_file_content(template.file_path)
|
||||
|
||||
if format == "pdf":
|
||||
file_content = docx_to_pdf_bytes(file_content)
|
||||
media_type = "application/pdf"
|
||||
filename = f"{template.name}.pdf"
|
||||
else:
|
||||
media_type = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
filename = f"{template.name}.docx"
|
||||
|
||||
return Response(
|
||||
content=file_content,
|
||||
media_type=media_type,
|
||||
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
||||
)
|
||||
|
||||
|
||||
@router.delete("/{template_id}")
|
||||
async def delete_template(template_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(Template).where(Template.id == template_id))
|
||||
template = result.scalar_one_or_none()
|
||||
if not template:
|
||||
raise HTTPException(status_code=404, detail="模板不存在")
|
||||
|
||||
delete_file(template.file_path)
|
||||
await db.delete(template)
|
||||
return {"detail": "删除成功"}
|
||||
@@ -0,0 +1,47 @@
|
||||
from pydantic_settings import BaseSettings
|
||||
from functools import lru_cache
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
PROJECT_NAME: str = "Doc Forge Reborn"
|
||||
API_V1_PREFIX: str = "/api/v1"
|
||||
DEBUG: bool = False
|
||||
|
||||
POSTGRES_USER: str = "docforge"
|
||||
POSTGRES_PASSWORD: str = "docforge"
|
||||
POSTGRES_DB: str = "docforge"
|
||||
POSTGRES_HOST: str = "localhost"
|
||||
POSTGRES_PORT: int = 5432
|
||||
|
||||
REDIS_URL: str = "redis://localhost:6379/0"
|
||||
CELERY_BROKER_URL: str = "redis://localhost:6379/0"
|
||||
CELERY_RESULT_BACKEND: str = "redis://localhost:6379/1"
|
||||
|
||||
SECRET_KEY: str = "change-me-in-production"
|
||||
DEFAULT_MODEL_ID: str = ""
|
||||
|
||||
STORAGE_ROOT: str = "./storage"
|
||||
MAX_UPLOAD_SIZE: int = 50 * 1024 * 1024
|
||||
|
||||
CORS_ORIGINS: list[str] = ["http://localhost:3000", "http://localhost:5173"]
|
||||
|
||||
@property
|
||||
def DATABASE_URL(self) -> str:
|
||||
return (
|
||||
f"postgresql+asyncpg://{self.POSTGRES_USER}:{self.POSTGRES_PASSWORD}"
|
||||
f"@{self.POSTGRES_HOST}:{self.POSTGRES_PORT}/{self.POSTGRES_DB}"
|
||||
)
|
||||
|
||||
@property
|
||||
def DATABASE_URL_SYNC(self) -> str:
|
||||
return (
|
||||
f"postgresql://{self.POSTGRES_USER}:{self.POSTGRES_PASSWORD}"
|
||||
f"@{self.POSTGRES_HOST}:{self.POSTGRES_PORT}/{self.POSTGRES_DB}"
|
||||
)
|
||||
|
||||
model_config = {"env_file": ".env", "env_file_encoding": "utf-8", "case_sensitive": True}
|
||||
|
||||
|
||||
@lru_cache
|
||||
def get_settings() -> Settings:
|
||||
return Settings()
|
||||
@@ -0,0 +1,35 @@
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
|
||||
from sqlalchemy.orm import DeclarativeBase
|
||||
from app.core.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
engine = create_async_engine(
|
||||
settings.DATABASE_URL,
|
||||
echo=settings.DEBUG,
|
||||
pool_size=20,
|
||||
max_overflow=10,
|
||||
pool_pre_ping=True,
|
||||
)
|
||||
|
||||
async_session_factory = async_sessionmaker(
|
||||
engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False,
|
||||
)
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
pass
|
||||
|
||||
|
||||
async def get_db() -> AsyncSession:
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
yield session
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
finally:
|
||||
await session.close()
|
||||
@@ -0,0 +1,25 @@
|
||||
import base64
|
||||
import hashlib
|
||||
from cryptography.fernet import Fernet
|
||||
from app.core.config import get_settings
|
||||
|
||||
|
||||
def _derive_fernet_key(secret: str) -> bytes:
|
||||
digest = hashlib.sha256(secret.encode("utf-8")).digest()
|
||||
return base64.urlsafe_b64encode(digest)
|
||||
|
||||
|
||||
def _get_fernet() -> Fernet:
|
||||
settings = get_settings()
|
||||
key = _derive_fernet_key(settings.SECRET_KEY)
|
||||
return Fernet(key)
|
||||
|
||||
|
||||
def encrypt_api_key(api_key: str) -> str:
|
||||
f = _get_fernet()
|
||||
return f.encrypt(api_key.encode("utf-8")).decode("utf-8")
|
||||
|
||||
|
||||
def decrypt_api_key(encrypted_key: str) -> str:
|
||||
f = _get_fernet()
|
||||
return f.decrypt(encrypted_key.encode("utf-8")).decode("utf-8")
|
||||
@@ -0,0 +1,32 @@
|
||||
from fastapi import Request, HTTPException
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from app.core.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class SecurityMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
if request.method in ("POST", "PUT", "PATCH"):
|
||||
content_length = request.headers.get("content-length")
|
||||
if content_length and int(content_length) > settings.MAX_UPLOAD_SIZE:
|
||||
raise HTTPException(
|
||||
status_code=413,
|
||||
detail=f"文件大小不能超过 {settings.MAX_UPLOAD_SIZE // 1024 // 1024}MB",
|
||||
)
|
||||
response = await call_next(request)
|
||||
return response
|
||||
|
||||
|
||||
ALLOWED_EXTENSIONS = {".docx", ".txt", ".pdf"}
|
||||
|
||||
|
||||
def validate_file_extension(filename: str) -> str:
|
||||
import os
|
||||
ext = os.path.splitext(filename)[1].lower()
|
||||
if ext not in ALLOWED_EXTENSIONS:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"不支持的文件类型: {ext}，允许的类型: {', '.join(ALLOWED_EXTENSIONS)}",
|
||||
)
|
||||
return ext
|
||||
@@ -0,0 +1,29 @@
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from app.core.config import get_settings
|
||||
from app.core.security_middleware import SecurityMiddleware
|
||||
from app.api.router import api_router
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
app = FastAPI(
|
||||
title=settings.PROJECT_NAME,
|
||||
openapi_url=f"{settings.API_V1_PREFIX}/openapi.json",
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.CORS_ORIGINS,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
app.add_middleware(SecurityMiddleware)
|
||||
|
||||
app.include_router(api_router)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
return {"status": "ok", "version": "0.1.0"}
|
||||
@@ -0,0 +1,17 @@
|
||||
from app.models.base import Base, TimestampMixin, UUIDMixin
|
||||
from app.models.ai_model import AIModel
|
||||
from app.models.template import Template
|
||||
from app.models.generation_point import GenerationPoint
|
||||
from app.models.generation_task import GenerationTask
|
||||
from app.models.system_config import SystemConfig
|
||||
|
||||
__all__ = [
|
||||
"Base",
|
||||
"TimestampMixin",
|
||||
"UUIDMixin",
|
||||
"AIModel",
|
||||
"Template",
|
||||
"GenerationPoint",
|
||||
"GenerationTask",
|
||||
"SystemConfig",
|
||||
]
|
||||
@@ -0,0 +1,18 @@
|
||||
import uuid
|
||||
from sqlalchemy import Boolean, String, Text
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
from app.core.database import Base
|
||||
from app.models.base import TimestampMixin, UUIDMixin
|
||||
|
||||
|
||||
class AIModel(Base, UUIDMixin, TimestampMixin):
|
||||
__tablename__ = "models"
|
||||
|
||||
name: Mapped[str] = mapped_column(String(100), nullable=False)
|
||||
provider: Mapped[str] = mapped_column(String(50), nullable=False)
|
||||
endpoint: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||
api_key: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
extra_params: Mapped[dict] = mapped_column(JSONB, default=dict, nullable=False)
|
||||
is_enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
||||
remark: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
@@ -0,0 +1,17 @@
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from sqlalchemy import DateTime, func
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class TimestampMixin:
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
|
||||
)
|
||||
|
||||
|
||||
class UUIDMixin:
|
||||
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
@@ -0,0 +1,27 @@
|
||||
import uuid
|
||||
from sqlalchemy import Integer, String, Text, ForeignKey
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
from app.core.database import Base
|
||||
from app.models.base import TimestampMixin, UUIDMixin
|
||||
|
||||
|
||||
class GenerationPoint(Base, UUIDMixin, TimestampMixin):
|
||||
__tablename__ = "generation_points"
|
||||
|
||||
template_id: Mapped[uuid.UUID] = mapped_column(
|
||||
UUID(as_uuid=True), ForeignKey("templates.id", ondelete="CASCADE"), nullable=False, index=True
|
||||
)
|
||||
position: Mapped[dict] = mapped_column(JSONB, nullable=False)
|
||||
prompt: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
model_id: Mapped[uuid.UUID | None] = mapped_column(
|
||||
UUID(as_uuid=True), ForeignKey("models.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
ref_file_path: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
need_ref_file: Mapped[bool] = mapped_column(default=False, nullable=False)
|
||||
remark: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
order: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
selected_text: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
|
||||
template = relationship("Template")
|
||||
model = relationship("AIModel")
|
||||
@@ -0,0 +1,23 @@
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from sqlalchemy import String, Text, DateTime, ForeignKey
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
from app.core.database import Base
|
||||
from app.models.base import UUIDMixin
|
||||
|
||||
|
||||
class GenerationTask(Base, UUIDMixin):
|
||||
__tablename__ = "generation_tasks"
|
||||
|
||||
template_id: Mapped[uuid.UUID] = mapped_column(
|
||||
UUID(as_uuid=True), ForeignKey("templates.id", ondelete="CASCADE"), nullable=False, index=True
|
||||
)
|
||||
status: Mapped[str] = mapped_column(String(20), nullable=False, default="pending")
|
||||
result_file_path: Mapped[str | None] = mapped_column(String(500), nullable=True)
|
||||
error_msg: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
celery_task_id: Mapped[str | None] = mapped_column(String(100), nullable=True)
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
|
||||
finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
|
||||
template = relationship("Template")
|
||||
@@ -0,0 +1,11 @@
|
||||
from sqlalchemy import String, Text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class SystemConfig(Base):
|
||||
__tablename__ = "system_config"
|
||||
|
||||
key: Mapped[str] = mapped_column(String(50), primary_key=True)
|
||||
value: Mapped[str] = mapped_column(Text, nullable=True)
|
||||
description: Mapped[str | None] = mapped_column(String(200), nullable=True)
|
||||
@@ -0,0 +1,14 @@
|
||||
import uuid
|
||||
from sqlalchemy import String, Text
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
from app.core.database import Base
|
||||
from app.models.base import TimestampMixin, UUIDMixin
|
||||
|
||||
|
||||
class Template(Base, UUIDMixin, TimestampMixin):
|
||||
__tablename__ = "templates"
|
||||
|
||||
name: Mapped[str] = mapped_column(String(200), nullable=False)
|
||||
file_path: Mapped[str] = mapped_column(String(500), nullable=False)
|
||||
html_content: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
@@ -0,0 +1,46 @@
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from pydantic import BaseModel, Field, field_serializer
|
||||
|
||||
|
||||
class AIModelCreate(BaseModel):
|
||||
name: str = Field(..., max_length=100)
|
||||
provider: str = Field(..., max_length=50)
|
||||
endpoint: str = Field(..., max_length=255)
|
||||
api_key: str
|
||||
extra_params: dict = Field(default_factory=dict)
|
||||
is_enabled: bool = True
|
||||
remark: str | None = None
|
||||
|
||||
|
||||
class AIModelUpdate(BaseModel):
|
||||
name: str | None = Field(None, max_length=100)
|
||||
provider: str | None = Field(None, max_length=50)
|
||||
endpoint: str | None = Field(None, max_length=255)
|
||||
api_key: str | None = None
|
||||
extra_params: dict | None = None
|
||||
is_enabled: bool | None = None
|
||||
remark: str | None = None
|
||||
|
||||
|
||||
class AIModelToggle(BaseModel):
|
||||
is_enabled: bool
|
||||
|
||||
|
||||
class AIModelResponse(BaseModel):
|
||||
id: uuid.UUID
|
||||
name: str
|
||||
provider: str
|
||||
endpoint: str
|
||||
api_key: str
|
||||
extra_params: dict
|
||||
is_enabled: bool
|
||||
remark: str | None = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
@field_serializer("api_key")
|
||||
def mask_api_key(self, value: str) -> str:
|
||||
return "***"
|
||||
@@ -0,0 +1,45 @@
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class GenerationPointCreate(BaseModel):
|
||||
template_id: uuid.UUID
|
||||
position: dict = Field(..., description="选区位置信息")
|
||||
prompt: str
|
||||
model_id: uuid.UUID | None = None
|
||||
order: int = 0
|
||||
selected_text: str | None = None
|
||||
need_ref_file: bool = False
|
||||
remark: str | None = None
|
||||
|
||||
|
||||
class GenerationPointUpdate(BaseModel):
|
||||
position: dict | None = None
|
||||
prompt: str | None = None
|
||||
model_id: uuid.UUID | None = None
|
||||
order: int | None = None
|
||||
selected_text: str | None = None
|
||||
need_ref_file: bool | None = None
|
||||
remark: str | None = None
|
||||
|
||||
|
||||
class GenerationPointResponse(BaseModel):
|
||||
id: uuid.UUID
|
||||
template_id: uuid.UUID
|
||||
position: dict
|
||||
prompt: str
|
||||
model_id: uuid.UUID | None = None
|
||||
ref_file_path: str | None = None
|
||||
need_ref_file: bool = False
|
||||
remark: str | None = None
|
||||
order: int
|
||||
selected_text: str | None = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class BatchOrderUpdate(BaseModel):
|
||||
points: list[dict] = Field(..., description="[{id: uuid, order: int}, ...]")
|
||||
@@ -0,0 +1,44 @@
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class TaskResponse(BaseModel):
|
||||
id: uuid.UUID
|
||||
template_id: uuid.UUID
|
||||
status: str
|
||||
result_file_path: str | None = None
|
||||
error_msg: str | None = None
|
||||
created_at: datetime
|
||||
finished_at: datetime | None = None
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class TaskDetailResponse(BaseModel):
|
||||
id: uuid.UUID
|
||||
template_id: uuid.UUID
|
||||
template_name: str = ""
|
||||
status: str
|
||||
result_file_path: str | None = None
|
||||
error_msg: str | None = None
|
||||
created_at: datetime
|
||||
finished_at: datetime | None = None
|
||||
points: list[dict] = []
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class GenerateResponse(BaseModel):
|
||||
task_id: uuid.UUID
|
||||
status: str
|
||||
|
||||
|
||||
class SingleTestRequest(BaseModel):
|
||||
prompt: str
|
||||
model_id: uuid.UUID | None = None
|
||||
ref_file_path: str | None = None
|
||||
|
||||
|
||||
class SingleTestResponse(BaseModel):
|
||||
result: str
|
||||
@@ -0,0 +1,31 @@
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class TemplateResponse(BaseModel):
|
||||
id: uuid.UUID
|
||||
name: str
|
||||
file_path: str
|
||||
html_content: str | None = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class TemplateListItem(BaseModel):
|
||||
id: uuid.UUID
|
||||
name: str
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class HTMLContentResponse(BaseModel):
|
||||
html_content: str
|
||||
|
||||
|
||||
class HTMLUpdateRequest(BaseModel):
|
||||
html_content: str
|
||||
@@ -0,0 +1,181 @@
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
import httpx
|
||||
from app.core.security import decrypt_api_key
|
||||
|
||||
PROVIDER_OPENAI = "openai"
|
||||
PROVIDER_AZURE = "azure"
|
||||
PROVIDER_CUSTOM = "custom"
|
||||
|
||||
|
||||
class AIAdapter(ABC):
|
||||
@abstractmethod
|
||||
def build_request(self, model_config: dict, prompt: str, ref_content: str | None = None) -> dict:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def parse_response(self, response_data: dict) -> str:
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def provider(self) -> str:
|
||||
pass
|
||||
|
||||
|
||||
class OpenAIAdapter(AIAdapter):
|
||||
@property
|
||||
def provider(self) -> str:
|
||||
return PROVIDER_OPENAI
|
||||
|
||||
def build_request(self, model_config: dict, prompt: str, ref_content: str | None = None) -> dict:
|
||||
extra = model_config.get("extra_params", {})
|
||||
temperature = extra.get("temperature", 0.7)
|
||||
max_tokens = extra.get("max_tokens", 2000)
|
||||
|
||||
messages = [{"role": "system", "content": "你是一个专业的文档内容生成助手。"}]
|
||||
user_content = prompt
|
||||
if ref_content:
|
||||
user_content = f"参考以下内容：\n{ref_content}\n\n任务：{prompt}"
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
return {
|
||||
"url": model_config["endpoint"],
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {decrypt_api_key(model_config['api_key'])}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
"json": {
|
||||
"model": extra.get("model", "gpt-4"),
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
},
|
||||
}
|
||||
|
||||
def parse_response(self, response_data: dict) -> str:
|
||||
return response_data["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
class AzureAdapter(AIAdapter):
|
||||
@property
|
||||
def provider(self) -> str:
|
||||
return PROVIDER_AZURE
|
||||
|
||||
def build_request(self, model_config: dict, prompt: str, ref_content: str | None = None) -> dict:
|
||||
extra = model_config.get("extra_params", {})
|
||||
temperature = extra.get("temperature", 0.7)
|
||||
max_tokens = extra.get("max_tokens", 2000)
|
||||
|
||||
messages = [{"role": "system", "content": "你是一个专业的文档内容生成助手。"}]
|
||||
user_content = prompt
|
||||
if ref_content:
|
||||
user_content = f"参考以下内容：\n{ref_content}\n\n任务：{prompt}"
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
api_version = extra.get("api_version", "2024-02-15-preview")
|
||||
endpoint = model_config["endpoint"]
|
||||
url = f"{endpoint}?api-version={api_version}"
|
||||
|
||||
return {
|
||||
"url": url,
|
||||
"headers": {
|
||||
"api-key": decrypt_api_key(model_config["api_key"]),
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
"json": {
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
},
|
||||
}
|
||||
|
||||
def parse_response(self, response_data: dict) -> str:
|
||||
return response_data["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
class CustomAdapter(AIAdapter):
|
||||
@property
|
||||
def provider(self) -> str:
|
||||
return PROVIDER_CUSTOM
|
||||
|
||||
def build_request(self, model_config: dict, prompt: str, ref_content: str | None = None) -> dict:
|
||||
extra = model_config.get("extra_params", {})
|
||||
user_content = prompt
|
||||
if ref_content:
|
||||
user_content = f"参考以下内容：\n{ref_content}\n\n任务：{prompt}"
|
||||
|
||||
messages = [{"role": "system", "content": "你是一个专业的文档内容生成助手。"}]
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
body = {
|
||||
"model": extra.get("model", "gpt-3.5-turbo"),
|
||||
"messages": messages,
|
||||
"max_tokens": extra.get("max_tokens", 2000),
|
||||
"temperature": extra.get("temperature", 0.7),
|
||||
}
|
||||
body.update({k: v for k, v in extra.items() if k not in ("model", "messages", "max_tokens", "temperature")})
|
||||
|
||||
return {
|
||||
"url": model_config["endpoint"],
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {decrypt_api_key(model_config['api_key'])}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
"json": body,
|
||||
}
|
||||
|
||||
def parse_response(self, response_data: dict) -> str:
|
||||
if "choices" in response_data:
|
||||
return response_data["choices"][0]["message"]["content"]
|
||||
if "response" in response_data:
|
||||
return response_data["response"]
|
||||
if "content" in response_data:
|
||||
return response_data["content"]
|
||||
if "text" in response_data:
|
||||
return response_data["text"]
|
||||
return json.dumps(response_data)
|
||||
|
||||
|
||||
_adapters: dict[str, AIAdapter] = {
|
||||
PROVIDER_OPENAI: OpenAIAdapter(),
|
||||
PROVIDER_AZURE: AzureAdapter(),
|
||||
PROVIDER_CUSTOM: CustomAdapter(),
|
||||
}
|
||||
|
||||
|
||||
def get_adapter(provider: str) -> AIAdapter:
|
||||
adapter = _adapters.get(provider)
|
||||
if not adapter:
|
||||
raise ValueError(f"不支持的供应商: {provider}")
|
||||
return adapter
|
||||
|
||||
|
||||
async def call_ai_model(model_config: dict, prompt: str, ref_content: str | None = None) -> str:
|
||||
adapter = get_adapter(model_config["provider"])
|
||||
request = adapter.build_request(model_config, prompt, ref_content)
|
||||
|
||||
timeout = model_config.get("extra_params", {}).get("timeout", 120)
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=timeout,
|
||||
proxy=None,
|
||||
trust_env=False,
|
||||
) as client:
|
||||
response = await client.post(
|
||||
request["url"],
|
||||
headers=request["headers"],
|
||||
json=request["json"],
|
||||
)
|
||||
response.raise_for_status()
|
||||
return adapter.parse_response(response.json())
|
||||
except httpx.HTTPStatusError as e:
|
||||
detail = e.response.text[:500] if e.response else str(e)
|
||||
raise RuntimeError(f"AI 服务返回错误 ({e.response.status_code}): {detail}")
|
||||
except httpx.TimeoutException:
|
||||
raise RuntimeError("AI 调用超时")
|
||||
except httpx.ConnectError as e:
|
||||
raise RuntimeError(f"无法连接 AI 服务: {e}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"AI 调用异常: {str(e)}")
|
||||
@@ -0,0 +1,168 @@
|
||||
import mammoth
|
||||
from io import BytesIO
|
||||
from bs4 import BeautifulSoup
|
||||
from docx import Document
|
||||
from docx.shared import Pt, Inches
|
||||
from docx.enum.text import WD_ALIGN_PARAGRAPH
|
||||
from app.services.file_storage import get_file_content
|
||||
|
||||
ALIGN_MAP = {
|
||||
WD_ALIGN_PARAGRAPH.CENTER: "center",
|
||||
WD_ALIGN_PARAGRAPH.RIGHT: "right",
|
||||
WD_ALIGN_PARAGRAPH.JUSTIFY: "justify",
|
||||
}
|
||||
|
||||
|
||||
async def docx_to_html(file_content: bytes) -> str:
|
||||
result = mammoth.convert_to_html(BytesIO(file_content))
|
||||
html = result.value
|
||||
|
||||
try:
|
||||
doc = Document(BytesIO(file_content))
|
||||
soup = BeautifulSoup(html, "html.parser")
|
||||
paragraphs = doc.paragraphs
|
||||
html_blocks = soup.find_all(["p", "h1", "h2", "h3", "h4", "h5", "h6", "li"])
|
||||
|
||||
for i, para in enumerate(paragraphs):
|
||||
if i >= len(html_blocks):
|
||||
break
|
||||
if para.alignment and para.alignment in ALIGN_MAP:
|
||||
css = ALIGN_MAP[para.alignment]
|
||||
existing = html_blocks[i].get("style", "")
|
||||
styles = f"text-align:{css}"
|
||||
if existing:
|
||||
styles = existing.rstrip(";") + ";" + styles
|
||||
html_blocks[i]["style"] = styles
|
||||
|
||||
html = str(soup)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return html
|
||||
|
||||
|
||||
def html_to_docx_bytes(html_content: str) -> bytes:
|
||||
from html.parser import HTMLParser
|
||||
|
||||
doc = Document()
|
||||
|
||||
style = doc.styles["Normal"]
|
||||
font = style.font
|
||||
font.name = "Arial"
|
||||
font.size = Pt(11)
|
||||
|
||||
class RichHTMLParser(HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.paragraphs: list[dict] = []
|
||||
self.current = {"runs": [], "align": None}
|
||||
self.in_paragraph = False
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
self.tag_stack: list[str] = []
|
||||
self.heading_level = 0
|
||||
|
||||
def handle_starttag(self, tag, attrs):
|
||||
tag_lower = tag.lower()
|
||||
attrs_dict = dict(attrs)
|
||||
if tag_lower in ("p", "div", "li"):
|
||||
self.in_paragraph = True
|
||||
self.current = {"runs": [], "align": None}
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
style = attrs_dict.get("style", "")
|
||||
if "text-align:center" in style:
|
||||
self.current["align"] = "center"
|
||||
elif "text-align:right" in style:
|
||||
self.current["align"] = "right"
|
||||
elif "text-align:justify" in style:
|
||||
self.current["align"] = "justify"
|
||||
elif tag_lower in ("h1", "h2", "h3", "h4", "h5", "h6"):
|
||||
self.in_paragraph = True
|
||||
self.heading_level = int(tag_lower[1])
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
elif tag_lower in ("strong", "b"):
|
||||
self.current_run["bold"] = True
|
||||
elif tag_lower in ("em", "i"):
|
||||
self.current_run["italic"] = True
|
||||
elif tag_lower == "u":
|
||||
self.current_run["underline"] = True
|
||||
elif tag_lower in ("br",):
|
||||
if self.in_paragraph:
|
||||
self.current["runs"].append(dict(self.current_run))
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
self.tag_stack.append(tag_lower)
|
||||
|
||||
def handle_endtag(self, tag):
|
||||
tag_lower = tag.lower()
|
||||
if tag_lower in ("p", "div", "li", "h1", "h2", "h3", "h4", "h5", "h6"):
|
||||
if self.current_run["text"].strip():
|
||||
self.current["runs"].append(dict(self.current_run))
|
||||
if self.current["runs"]:
|
||||
p = dict(self.current)
|
||||
p["heading"] = self.heading_level
|
||||
self.paragraphs.append(p)
|
||||
self.current = {"runs": []}
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
self.in_paragraph = False
|
||||
self.heading_level = 0
|
||||
if self.tag_stack:
|
||||
self.tag_stack.pop()
|
||||
|
||||
def handle_data(self, data):
|
||||
if self.in_paragraph:
|
||||
self.current_run["text"] += data
|
||||
|
||||
parser = RichHTMLParser()
|
||||
parser.feed(html_content)
|
||||
|
||||
for para_data in parser.paragraphs:
|
||||
heading = para_data.get("heading", 0)
|
||||
if heading > 0:
|
||||
p = doc.add_heading(level=min(heading, 9))
|
||||
else:
|
||||
p = doc.add_paragraph()
|
||||
|
||||
align_val = para_data.get("align")
|
||||
if align_val == "center":
|
||||
p.alignment = WD_ALIGN_PARAGRAPH.CENTER
|
||||
elif align_val == "right":
|
||||
p.alignment = WD_ALIGN_PARAGRAPH.RIGHT
|
||||
elif align_val == "justify":
|
||||
p.alignment = WD_ALIGN_PARAGRAPH.JUSTIFY
|
||||
|
||||
for run_data in para_data.get("runs", []):
|
||||
run = p.add_run(run_data["text"])
|
||||
if run_data.get("bold"):
|
||||
run.bold = True
|
||||
if run_data.get("italic"):
|
||||
run.italic = True
|
||||
if run_data.get("underline"):
|
||||
run.underline = True
|
||||
|
||||
output = BytesIO()
|
||||
doc.save(output)
|
||||
return output.getvalue()
|
||||
|
||||
|
||||
def docx_to_pdf_bytes(file_content: bytes) -> bytes:
|
||||
doc = Document(BytesIO(file_content))
|
||||
|
||||
from io import BytesIO as Bio
|
||||
from reportlab.lib.pagesizes import A4
|
||||
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
|
||||
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
|
||||
from reportlab.lib.enums import TA_LEFT
|
||||
|
||||
buffer = Bio()
|
||||
pdf_doc = SimpleDocTemplate(buffer, pagesize=A4)
|
||||
|
||||
styles = getSampleStyleSheet()
|
||||
story = []
|
||||
|
||||
for para in doc.paragraphs:
|
||||
if para.text.strip():
|
||||
p = Paragraph(para.text, styles["Normal"])
|
||||
story.append(p)
|
||||
story.append(Spacer(1, 6))
|
||||
|
||||
pdf_doc.build(story)
|
||||
return buffer.getvalue()
|
||||
@@ -0,0 +1,43 @@
|
||||
import os
|
||||
import aiofiles
|
||||
from pathlib import Path
|
||||
from fastapi import UploadFile
|
||||
from app.core.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
TEMPLATES_DIR = "templates"
|
||||
REF_FILES_DIR = "ref_files"
|
||||
RESULTS_DIR = "results"
|
||||
|
||||
|
||||
def _ensure_dir(path: str) -> str:
|
||||
os.makedirs(path, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def get_storage_dir(subdir: str) -> str:
|
||||
return _ensure_dir(os.path.join(settings.STORAGE_ROOT, subdir))
|
||||
|
||||
|
||||
async def save_upload(upload_file: UploadFile, subdir: str) -> str:
|
||||
dir_path = get_storage_dir(subdir)
|
||||
file_path = os.path.join(dir_path, upload_file.filename or "unnamed")
|
||||
async with aiofiles.open(file_path, "wb") as f:
|
||||
content = await upload_file.read()
|
||||
await f.write(content)
|
||||
return file_path
|
||||
|
||||
|
||||
async def get_file_content(file_path: str) -> bytes:
|
||||
async with aiofiles.open(file_path, "rb") as f:
|
||||
return await f.read()
|
||||
|
||||
|
||||
def delete_file(file_path: str) -> None:
|
||||
if os.path.exists(file_path):
|
||||
os.remove(file_path)
|
||||
|
||||
|
||||
def get_absolute_path(rel_path: str) -> str:
|
||||
return os.path.abspath(rel_path)
|
||||
@@ -0,0 +1,46 @@
|
||||
import os
|
||||
import json
|
||||
from docx import Document
|
||||
from PyPDF2 import PdfReader
|
||||
from app.services.file_storage import get_file_content
|
||||
|
||||
|
||||
async def parse_reference_file(file_path: str) -> str:
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
content = await get_file_content(file_path)
|
||||
|
||||
if ext == ".txt":
|
||||
return content.decode("utf-8", errors="ignore")
|
||||
if ext == ".docx":
|
||||
from io import BytesIO
|
||||
doc = Document(BytesIO(content))
|
||||
return "\n".join(p.text for p in doc.paragraphs if p.text.strip())
|
||||
if ext == ".pdf":
|
||||
from io import BytesIO
|
||||
reader = PdfReader(BytesIO(content))
|
||||
return "\n".join(page.extract_text() or "" for page in reader.pages)
|
||||
|
||||
raise ValueError(f"不支持的文件格式: {ext}")
|
||||
|
||||
|
||||
async def parse_reference_files(ref_file_path: str | None) -> str:
|
||||
if not ref_file_path:
|
||||
return ""
|
||||
|
||||
try:
|
||||
paths = json.loads(ref_file_path)
|
||||
if isinstance(paths, list):
|
||||
texts = []
|
||||
for path in paths:
|
||||
try:
|
||||
texts.append(await parse_reference_file(path))
|
||||
except Exception:
|
||||
texts.append(f"[无法解析文件: {path}]")
|
||||
return "\n\n".join(texts)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
return await parse_reference_file(ref_file_path)
|
||||
except Exception:
|
||||
return f"[无法解析文件: {ref_file_path}]"
|
||||
@@ -0,0 +1,25 @@
|
||||
from celery import Celery
|
||||
from app.core.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
celery_app = Celery(
|
||||
"doc_forge_reds",
|
||||
broker=settings.CELERY_BROKER_URL,
|
||||
backend=settings.CELERY_RESULT_BACKEND,
|
||||
)
|
||||
|
||||
celery_app.conf.update(
|
||||
task_serializer="json",
|
||||
accept_content=["json"],
|
||||
result_serializer="json",
|
||||
timezone="Asia/Shanghai",
|
||||
enable_utc=True,
|
||||
task_track_started=True,
|
||||
task_acks_late=True,
|
||||
worker_prefetch_multiplier=1,
|
||||
task_soft_time_limit=600,
|
||||
task_time_limit=900,
|
||||
)
|
||||
|
||||
celery_app.autodiscover_tasks(["app.tasks.generate"])
|
||||
@@ -0,0 +1,225 @@
|
||||
import asyncio
|
||||
import os
|
||||
from io import BytesIO
|
||||
from datetime import datetime, timezone
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
|
||||
from docx import Document
|
||||
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.ai_adapter import call_ai_model
|
||||
from app.services.ref_parser import parse_reference_files
|
||||
from app.services.file_storage import get_storage_dir, get_file_content, 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)
|
||||
|
||||
|
||||
def _strip_html(html: str) -> str:
|
||||
from html.parser import HTMLParser
|
||||
|
||||
class Stripper(HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.text = ""
|
||||
def handle_data(self, data):
|
||||
self.text += data
|
||||
|
||||
s = Stripper()
|
||||
s.feed(html)
|
||||
return s.text
|
||||
|
||||
|
||||
def _get_selected_text(html_content: str, position: dict) -> str:
|
||||
start = position.get("start", 0)
|
||||
end = position.get("end", 0)
|
||||
|
||||
plain_text = _strip_html(html_content)
|
||||
if 0 <= start < end <= len(plain_text):
|
||||
return plain_text[start:end].strip()
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def _replace_text_in_docx(doc: Document, old_text: str, new_text: str) -> bool:
|
||||
if not old_text:
|
||||
return False
|
||||
|
||||
for paragraph in doc.paragraphs:
|
||||
if old_text in paragraph.text:
|
||||
inline = paragraph.runs
|
||||
for run in inline:
|
||||
if old_text in run.text:
|
||||
run.text = run.text.replace(old_text, new_text)
|
||||
return True
|
||||
|
||||
full_text = "".join(r.text for r in inline)
|
||||
if old_text in full_text:
|
||||
remaining = old_text
|
||||
for run in inline:
|
||||
if not remaining:
|
||||
break
|
||||
if remaining.startswith(run.text):
|
||||
remaining = remaining[len(run.text):]
|
||||
elif run.text in remaining:
|
||||
idx = remaining.find(run.text)
|
||||
if idx >= 0:
|
||||
remaining = remaining[:idx] + remaining[idx + len(run.text):]
|
||||
if remaining.startswith(run.text):
|
||||
remaining = remaining[len(run.text):]
|
||||
|
||||
if not remaining:
|
||||
chunk_parts = new_text
|
||||
for run in inline:
|
||||
if chunk_parts:
|
||||
chunk_parts = chunk_parts[len(run.text):]
|
||||
|
||||
first_run = inline[0]
|
||||
first_run.text = new_text
|
||||
for run in inline[1:]:
|
||||
run.text = ""
|
||||
return True
|
||||
|
||||
for table in doc.tables:
|
||||
for row in table.rows:
|
||||
for cell in row.cells:
|
||||
for paragraph in cell.paragraphs:
|
||||
if old_text in paragraph.text:
|
||||
for run in paragraph.runs:
|
||||
if old_text in run.text:
|
||||
run.text = run.text.replace(old_text, new_text)
|
||||
return True
|
||||
|
||||
return 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
|
||||
|
||||
if not os.path.exists(template.file_path):
|
||||
task.status = "failed"
|
||||
task.error_msg = f"原始文件不存在: {template.file_path}"
|
||||
task.finished_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
return
|
||||
|
||||
docx_content = await get_file_content(template.file_path)
|
||||
doc = Document(BytesIO(docx_content))
|
||||
|
||||
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:
|
||||
# 先收集所有 AI 调用参数
|
||||
point_data = []
|
||||
for point in points:
|
||||
selected_text = point.selected_text or _get_selected_text(html_content, point.position)
|
||||
if not selected_text:
|
||||
continue
|
||||
|
||||
ref_content = ""
|
||||
if point.ref_file_path:
|
||||
ref_content = await parse_reference_files(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,
|
||||
}
|
||||
|
||||
point_data.append({
|
||||
"point": point,
|
||||
"selected_text": selected_text,
|
||||
"model_config": model_config,
|
||||
"ref_content": ref_content,
|
||||
})
|
||||
|
||||
# 并发调用所有 AI 模型
|
||||
async def _call_one(pd):
|
||||
try:
|
||||
return await call_ai_model(pd["model_config"], pd["point"].prompt, pd["ref_content"])
|
||||
except Exception as e:
|
||||
return f"[生成失败: {e}]"
|
||||
|
||||
coros = [_call_one(pd) for pd in point_data]
|
||||
ai_results = await asyncio.gather(*coros)
|
||||
|
||||
# 按顺序替换文本
|
||||
for pd, ai_result in zip(point_data, ai_results):
|
||||
_replace_text_in_docx(doc, pd["selected_text"], str(ai_result))
|
||||
|
||||
await db.commit()
|
||||
|
||||
result_dir = get_storage_dir(RESULTS_DIR)
|
||||
result_path = os.path.join(result_dir, f"{task_id}.docx")
|
||||
doc.save(result_path)
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,16 @@
|
||||
fastapi==0.111.0
|
||||
uvicorn[standard]==0.30.1
|
||||
sqlalchemy[asyncio]==2.0.30
|
||||
asyncpg==0.29.0
|
||||
alembic==1.13.1
|
||||
pydantic-settings==2.3.4
|
||||
cryptography==42.0.8
|
||||
python-multipart==0.0.9
|
||||
python-docx==1.1.2
|
||||
mammoth==1.7.2
|
||||
PyPDF2==3.0.1
|
||||
httpx==0.27.0
|
||||
celery[redis]==5.4.0
|
||||
redis==5.0.7
|
||||
aiofiles==24.1.0
|
||||
reportlab==4.2.2
|
||||
Reference in New Issue
Block a user