mz/video-gen
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

用户生成资源容量管控

Browse Source
This commit is contained in:
GinHa
2026-06-29 13:50:39 +08:00
parent b21f35582e
commit 25fa41f894
21 changed files with 1314 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
video-gen-api/alembic/versions/78fb32c26a6e_add_user_resource_capacity_config.py
+43
View File
@@ -0,0 +1,43 @@
"""add user resource capacity config
Revision ID: 78fb32c26a6e
Revises: n123456789ab
Create Date: 2026-06-29 13:05:45.888522
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision: str = '78fb32c26a6e'
down_revision: Union[str, None] = 'n123456789ab'
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('user_resource_capacity_configs',
sa.Column('id', sa.String(length=32), nullable=False),
sa.Column('user_id', sa.String(length=32), nullable=False, comment='用户ID'),
sa.Column('enabled', sa.Boolean(), server_default='false', nullable=False, comment='是否启用该用户个人容量限制'),
sa.Column('limit_value', sa.Numeric(precision=18, scale=3), server_default='1', nullable=False, comment='容量数值，最小1，最多3位小数'),
sa.Column('limit_unit', sa.String(length=8), server_default='GB', nullable=False, comment='容量单位：MB/GB/TB'),
sa.Column('limit_bytes', sa.BigInteger(), server_default='1073741824', nullable=False, comment='换算后的容量字节数'),
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(['user_id'], ['users.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('id'),
sa.UniqueConstraint('user_id', name='uq_user_resource_capacity_configs_user')
)
op.create_index(op.f('ix_user_resource_capacity_configs_user_id'), 'user_resource_capacity_configs', ['user_id'], unique=False)
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_index(op.f('ix_user_resource_capacity_configs_user_id'), table_name='user_resource_capacity_configs')
op.drop_table('user_resource_capacity_configs')
# ### end Alembic commands ###
video-gen-api/app/api/admin/__init__.py
+2
View File
@@ -1,6 +1,8 @@
from fastapi import APIRouter
from app.api.admin.video_prompt_schema_config import router as video_prompt_schema_config_router
from app.api.admin.resource_capacity import router as resource_capacity_router
router = APIRouter()
router.include_router(video_prompt_schema_config_router)
router.include_router(resource_capacity_router)
video-gen-api/app/api/admin/resource_capacity.py
+175
View File
@@ -0,0 +1,175 @@
from __future__ import annotations
from fastapi import APIRouter, Depends, Path
from sqlalchemy.ext.asyncio import AsyncSession
from app.dependencies import get_admin_user, get_db
from app.enums.resource_capacity import ResourceCapacityOperationEnum
from app.models.user import User
from app.schemas.resource_capacity import (
AdminUserResourceCapacityOut,
ResourceCapacityConfigOut,
ResourceCapacityConfigUpdate,
)
from app.services.operation_log import log_operation
from app.services.resource_capacity_service import (
build_global_resource_capacity_operation_detail,
build_user_resource_capacity_operation_detail,
delete_user_resource_capacity_config,
get_admin_user_resource_capacity,
get_global_resource_capacity_config,
save_global_resource_capacity_config,
save_user_resource_capacity_config,
)
router = APIRouter(prefix="/admin", tags=["admin-resource-capacity"])
@router.get(
"/resource-capacity/global",
response_model=ResourceCapacityConfigOut,
summary="获取全局资源空间容量配置",
description=(
"获取管理后台全局生成资源空间容量管控配置。"
"配置最终存储在 system_configs 表，key=resource_capacity_limit_config。"
"enabled=false 表示全局容量管控关闭；enabled=true 表示按 limit_value + limit_unit 换算出的 limit_bytes 进行限制。"
"单位枚举：MB=1048576字节，GB=1073741824字节，TB=1099511627776字节。"
),
)
async def get_global_resource_capacity(
admin: User = Depends(get_admin_user),
db: AsyncSession = Depends(get_db),
):
_ = admin
return await get_global_resource_capacity_config(db)
@router.put(
"/resource-capacity/global",
response_model=ResourceCapacityConfigOut,
summary="保存全局资源空间容量配置",
description=(
"保存管理后台全局生成资源空间容量管控配置。"
"enabled=true 时，limit_value 和 limit_unit 必填。"
"limit_value 最小为1，不能为负数，最多支持3位小数。"
"limit_unit 枚举明细：MB=1048576字节，GB=1073741824字节，TB=1099511627776字节。"
"limit_bytes 不允许前端传入，由后端统一换算后保存。"
"本接口会写入后台操作日志，action=更新全局资源容量配置，detail记录修改前后配置快照。"
),
)
async def update_global_resource_capacity(
req: ResourceCapacityConfigUpdate,
admin: User = Depends(get_admin_user),
db: AsyncSession = Depends(get_db),
):
before = await get_global_resource_capacity_config(db)
result = await save_global_resource_capacity_config(db, req)
await log_operation(
db,
admin.id,
admin.username,
ResourceCapacityOperationEnum.UPDATE_GLOBAL_CONFIG.value,
"PUT",
"/admin/resource-capacity/global",
detail=build_global_resource_capacity_operation_detail(before, result),
)
await db.commit()
return result
@router.get(
"/users/{user_id}/resource-capacity",
response_model=AdminUserResourceCapacityOut,
summary="获取指定用户资源空间容量配置",
description=(
"获取指定用户的个人容量配置、全局容量配置以及最终生效容量数据。"
"优先级：用户个人配置存在时优先使用；用户个人配置不存在时使用全局配置。"
"注意：用户个人配置存在且 enabled=false，表示该用户个人明确关闭容量限制，不再回落到全局配置。"
),
)
async def get_user_resource_capacity(
user_id: str = Path(..., description="用户ID，用于查询该用户个人容量配置和最终生效容量数据"),
admin: User = Depends(get_admin_user),
db: AsyncSession = Depends(get_db),
):
_ = admin
return await get_admin_user_resource_capacity(db, user_id)
@router.put(
"/users/{user_id}/resource-capacity",
response_model=AdminUserResourceCapacityOut,
summary="保存指定用户个人资源空间容量配置",
description=(
"新增或更新指定用户个人资源空间容量配置。"
"保存后该用户配置优先级高于全局配置。"
"enabled=true 表示启用该用户个人容量限制；enabled=false 表示该用户个人明确关闭限制，不再回落到全局。"
"limit_value 最小为1，不能为负数，最多3位小数。"
"limit_unit 枚举明细：MB=1048576字节，GB=1073741824字节，TB=1099511627776字节。"
"本接口会写入后台操作日志：首次创建 action=新增用户资源容量配置；已有配置更新 action=更新用户资源容量配置；detail记录修改前后配置快照。"
),
)
async def update_user_resource_capacity(
req: ResourceCapacityConfigUpdate,
user_id: str = Path(..., description="用户ID，用于新增或更新该用户个人容量配置"),
admin: User = Depends(get_admin_user),
db: AsyncSession = Depends(get_db),
):
before = await get_admin_user_resource_capacity(db, user_id)
result = await save_user_resource_capacity_config(db, user_id, req)
is_create = not before.has_user_config
await log_operation(
db,
admin.id,
admin.username,
(
ResourceCapacityOperationEnum.CREATE_USER_CONFIG.value
if is_create
else ResourceCapacityOperationEnum.UPDATE_USER_CONFIG.value
),
"PUT",
f"/admin/users/{user_id}/resource-capacity",
detail=build_user_resource_capacity_operation_detail(
target_user_id=user_id,
operation="create" if is_create else "update",
before=before,
after=result,
),
)
await db.commit()
return result
@router.delete(
"/users/{user_id}/resource-capacity",
response_model=AdminUserResourceCapacityOut,
summary="删除指定用户个人资源空间容量配置",
description=(
"删除指定用户个人容量配置。删除后该用户不再有个人覆盖配置，后续容量限制重新回落到全局配置。"
"本接口会写入后台操作日志，action=删除用户资源容量配置，detail记录删除前个人配置、删除后最终生效配置。"
),
)
async def remove_user_resource_capacity(
user_id: str = Path(..., description="用户ID，用于删除该用户个人容量配置并恢复走全局配置"),
admin: User = Depends(get_admin_user),
db: AsyncSession = Depends(get_db),
):
before = await get_admin_user_resource_capacity(db, user_id)
result = await delete_user_resource_capacity_config(db, user_id)
await log_operation(
db,
admin.id,
admin.username,
ResourceCapacityOperationEnum.DELETE_USER_CONFIG.value,
"DELETE",
f"/admin/users/{user_id}/resource-capacity",
detail=build_user_resource_capacity_operation_detail(
target_user_id=user_id,
operation="delete",
before=before,
after=result,
remark="删除用户个人容量配置，用户恢复使用全局资源容量配置。",
),
)
await db.commit()
return result
video-gen-api/app/api/v1/admin.py
+15 -2
View File
@@ -49,6 +49,7 @@ from app.services.auth import hash_password, verify_password
from app.services.operation_log import log_operation
from app.services.resource_signed_url_service import build_resource_signed_url
from app.services.payment import sync_pending_orders, process_refund
from app.services.resource_capacity_service import batch_get_user_resource_capacity_usage, get_user_resource_capacity_usage
from app.services.generation_billing_service import (
OWNER_GENERATION_RECORD,
@@ -107,7 +108,16 @@ async def list_users(
total = (await db.execute(count_query)).scalar() or 0
result = await db.execute(query.offset((page - 1) * page_size).limit(page_size))
items = result.scalars().all()
return {"items": [AdminUserOut.model_validate(u) for u in items], "total": total}
capacity_map = await batch_get_user_resource_capacity_usage(db, [u.id for u in items])
return {
"items": [
AdminUserOut.model_validate(u)
.model_copy(update={"resource_capacity": capacity_map.get(u.id)})
.model_dump(mode="json")
for u in items
],
"total": total,
}
@router.post("/users", response_model=AdminUserOut)
@@ -183,7 +193,10 @@ async def get_user(
if not user:
raise HTTPException(status_code=404, detail="用户不存在")
user.credits = round(user.credits, 2)
return user
resource_capacity = await get_user_resource_capacity_usage(db, user.id)
return AdminUserOut.model_validate(user).model_copy(
update={"resource_capacity": resource_capacity}
)
@router.post("/users/{user_id}/credits")
video-gen-api/app/api/v1/auth.py
+9 -2
View File
@@ -30,6 +30,7 @@ from app.services.auth import (
verify_password,
)
from app.services.sms import verify_sms_code
from app.services.resource_capacity_service import get_user_resource_capacity_usage
from app.utils.id_gen import generate_id
router = APIRouter(prefix="/auth", tags=["auth"])
@@ -248,9 +249,15 @@ async def logout(current_user: User = Depends(get_current_user_allow_password_pe
@router.get("/me", response_model=UserOut)
async def get_me(current_user: User = Depends(get_current_user_allow_password_pending)):
async def get_me(
current_user: User = Depends(get_current_user_allow_password_pending),
db: AsyncSession = Depends(get_db),
):
current_user.credits = round(current_user.credits, 2)
return current_user
resource_capacity = await get_user_resource_capacity_usage(db, current_user.id)
return UserOut.model_validate(current_user).model_copy(
update={"resource_capacity": resource_capacity}
)
@router.post(
video-gen-api/app/api/v1/generation.py
+5
View File
@@ -34,6 +34,7 @@ from app.services.resource_accounting_service import (
safe_file_size,
)
from app.services.resource_signed_url_service import build_resource_signed_url
from app.services.resource_capacity_service import assert_user_resource_capacity_available
from app.services.generation_billing_service import (
CHARGE_TEXT_PROMPT,
OWNER_GENERATION_RECORD,
@@ -398,6 +399,8 @@ async def generate(
if record.status not in ("prompt_optimized", "failed"):
raise InvalidStatusError("当前状态不允许生成")
await assert_user_resource_capacity_available(db, current_user.id)
attempt_no = await get_next_credit_attempt_no(
db,
owner_type=OWNER_GENERATION_RECORD,
@@ -522,6 +525,8 @@ async def retry_generation(
if record.status != "failed":
raise InvalidStatusError("只有失败的记录可以重试")
await assert_user_resource_capacity_available(db, current_user.id)
attempt_no = await get_next_credit_attempt_no(
db,
owner_type=OWNER_GENERATION_RECORD,
video-gen-api/app/api/v1/generation_ai.py
+3
View File
@@ -33,6 +33,7 @@ from app.services.generation_billing_service import (
)
from app.services.generation_log_service import log_task_event
from app.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
from app.services.resource_capacity_service import assert_user_resource_capacity_available
from app.tasks.celery_app import celery_app
router = APIRouter(
@@ -566,6 +567,8 @@ async def retry_task(
if task.status != "failed":
raise HTTPException(status_code=400, detail="只有失败任务可以重试")
await assert_user_resource_capacity_available(db, current_user.id)
attempt_no = await get_next_credit_attempt_no(
db,
owner_type=OWNER_CHAT_GENERATION_TASK,
video-gen-api/app/enums/__init__.py
+1
View File
@@ -10,3 +10,4 @@ from app.enums.generation_task import *
from app.enums.generation_status import *
from app.enums.sms import *
from app.enums.notification import *
from app.enums.resource_capacity import *
video-gen-api/app/enums/resource_capacity.py
+68
View File
@@ -0,0 +1,68 @@
from __future__ import annotations
from enum import Enum
class ResourceCapacityUnitEnum(str, Enum):
"""生成资源容量单位。"""
MB = "MB"
GB = "GB"
TB = "TB"
@property
def bytes_multiplier(self) -> int:
return RESOURCE_CAPACITY_UNIT_BYTES[self]
@property
def label(self) -> str:
return RESOURCE_CAPACITY_UNIT_LABELS[self]
MB_BYTES = 1048576
GB_BYTES = 1073741824
TB_BYTES = 1099511627776
RESOURCE_CAPACITY_UNIT_BYTES: dict[ResourceCapacityUnitEnum, int] = {
ResourceCapacityUnitEnum.MB: MB_BYTES,
ResourceCapacityUnitEnum.GB: GB_BYTES,
ResourceCapacityUnitEnum.TB: TB_BYTES,
}
RESOURCE_CAPACITY_UNIT_LABELS: dict[ResourceCapacityUnitEnum, str] = {
ResourceCapacityUnitEnum.MB: "MB（1048576 字节）",
ResourceCapacityUnitEnum.GB: "GB（1073741824 字节）",
ResourceCapacityUnitEnum.TB: "TB（1099511627776 字节）",
}
class ResourceCapacitySourceEnum(str, Enum):
"""最终容量配置来源。"""
USER = "user"
GLOBAL = "global"
DISABLED = "disabled"
class ResourceCapacityConfigKeyEnum(str, Enum):
"""SystemConfig 中使用的固定配置 key。"""
RESOURCE_CAPACITY_LIMIT_CONFIG = "resource_capacity_limit_config"
class ResourceCapacityErrorCodeEnum(str, Enum):
"""资源容量管控错误码。"""
RESOURCE_CAPACITY_EXCEEDED = "RESOURCE_CAPACITY_EXCEEDED"
class ResourceCapacityOperationEnum(str, Enum):
"""管理后台资源容量配置操作日志动作。"""
UPDATE_GLOBAL_CONFIG = "更新全局资源容量配置"
CREATE_USER_CONFIG = "新增用户资源容量配置"
UPDATE_USER_CONFIG = "更新用户资源容量配置"
DELETE_USER_CONFIG = "删除用户资源容量配置"
RESOURCE_CAPACITY_EXCEEDED_MESSAGE = "您个人的空间容量已超额，请删除素材资源释放容量或购买更高额度的容量套餐"
video-gen-api/app/models/__init__.py
+2
View File
@@ -21,6 +21,7 @@ from app.models.chat_provider_call_log import ChatProviderCallLog
from app.models.generated_resource import GeneratedResource
from app.models.user_resource_month_stat import UserResourceMonthStat
from app.models.user_resource_total_stat import UserResourceTotalStat
from app.models.user_resource_capacity_config import UserResourceCapacityConfig
from app.models.module_generation_project import ModuleGenerationProject
from app.models.module_generation_step import ModuleGenerationStep
from app.models.shot_replicate_task_set import ShotReplicateTaskSet
@@ -38,6 +39,7 @@ __all__ = [
"MenuConfig", "RechargePackage", "OperationLog",
"ChatGenerationTask", "ChatGenerationTaskEvent", "ChatProviderCallLog",
"GeneratedResource", "UserResourceMonthStat", "UserResourceTotalStat",
"UserResourceCapacityConfig",
"ModuleGenerationProject", "ModuleGenerationStep",
"ShotReplicateTaskSet", "ShotReplicateSegment",
"UserOAuth", "UserOAuthAccount", "UserOAuthApp",
video-gen-api/app/models/user_resource_capacity_config.py
+53
View File
@@ -0,0 +1,53 @@
from __future__ import annotations
from sqlalchemy import BigInteger, Boolean, ForeignKey, Numeric, String, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column
from app.enums.resource_capacity import ResourceCapacityUnitEnum
from app.models.base import Base, TimestampMixin
class UserResourceCapacityConfig(Base, TimestampMixin):
"""用户个人生成资源容量配置。存在记录即代表用户配置已单独设置。"""
__tablename__ = "user_resource_capacity_configs"
__table_args__ = (
UniqueConstraint("user_id", name="uq_user_resource_capacity_configs_user"),
)
id: Mapped[str] = mapped_column(String(32), primary_key=True)
user_id: Mapped[str] = mapped_column(
String(32),
ForeignKey("users.id", ondelete="CASCADE"),
index=True,
nullable=False,
comment="用户ID",
)
enabled: Mapped[bool] = mapped_column(
Boolean,
default=False,
server_default="false",
nullable=False,
comment="是否启用该用户个人容量限制",
)
limit_value: Mapped[float] = mapped_column(
Numeric(18, 3),
default=1,
server_default="1",
nullable=False,
comment="容量数值，最小1，最多3位小数",
)
limit_unit: Mapped[str] = mapped_column(
String(8),
default=ResourceCapacityUnitEnum.GB.value,
server_default=ResourceCapacityUnitEnum.GB.value,
nullable=False,
comment="容量单位：MB/GB/TB",
)
limit_bytes: Mapped[int] = mapped_column(
BigInteger,
default=1024 * 1024 * 1024,
server_default=str(1024 * 1024 * 1024),
nullable=False,
comment="换算后的容量字节数",
)
video-gen-api/app/schemas/admin.py
+2
View File
@@ -1,6 +1,7 @@
from pydantic import BaseModel, Field
from app.schemas.common import NaiveDatetime, NaiveDatetimeOptional
from app.schemas.resource_capacity import ResourceCapacityUsageOut
class CreditAdjustRequest(BaseModel):
@@ -54,6 +55,7 @@ class AdminUserOut(BaseModel):
created_at: NaiveDatetime
last_login_at: NaiveDatetimeOptional = None
allowed_menus: list | None = None
resource_capacity: ResourceCapacityUsageOut | None = None
model_config = {"from_attributes": True}
video-gen-api/app/schemas/resource_capacity.py
+72
View File
@@ -0,0 +1,72 @@
from __future__ import annotations
from decimal import Decimal
from pydantic import BaseModel, Field, field_validator, model_validator
from app.enums.resource_capacity import ResourceCapacitySourceEnum, ResourceCapacityUnitEnum
class ResourceCapacityConfigUpdate(BaseModel):
enabled: bool = Field(
...,
description="是否开启容量管控。true=开启，false=关闭。全局配置关闭表示全局不限制；用户配置关闭表示该用户个人明确关闭限制，优先级高于全局。",
examples=[True],
)
limit_value: Decimal | None = Field(
None,
description="容量数值，最小1，不能为负数，最多3位小数。例如：1、10、10.5、10.500。enabled=true 时必填。",
examples=["10.500"],
)
limit_unit: ResourceCapacityUnitEnum | None = Field(
None,
description="容量单位枚举：MB=1024*1024字节，GB=1024*1024*1024字节，TB=1024*1024*1024*1024字节。enabled=true 时必填。",
examples=[ResourceCapacityUnitEnum.GB.value],
)
@field_validator("limit_value")
@classmethod
def validate_limit_value(cls, value: Decimal | None) -> Decimal | None:
if value is None:
return value
if value < Decimal("1"):
raise ValueError("容量数值不能小于1")
if value.as_tuple().exponent < -3:
raise ValueError("容量数值最多支持3位小数")
return value
@model_validator(mode="after")
def validate_enabled_payload(self) -> "ResourceCapacityConfigUpdate":
if self.enabled:
if self.limit_value is None:
raise ValueError("开启容量管控时必须填写容量数值")
if self.limit_unit is None:
raise ValueError("开启容量管控时必须选择容量单位：MB、GB、TB")
return self
class ResourceCapacityConfigOut(BaseModel):
enabled: bool = Field(..., description="是否开启容量管控")
limit_value: str = Field(..., description="容量数值字符串，最多3位小数")
limit_unit: ResourceCapacityUnitEnum = Field(..., description="容量单位枚举：MB、GB、TB")
limit_bytes: int = Field(..., ge=0, description="按单位换算后的字节数")
class ResourceCapacityUsageOut(BaseModel):
enabled: bool = Field(..., description="当前用户最终是否启用容量管控")
source: ResourceCapacitySourceEnum = Field(..., description="最终配置来源：user=用户个人配置，global=全局配置，disabled=容量管控未开启")
has_user_config: bool = Field(..., description="该用户是否存在个人容量配置记录。注意：存在且 enabled=false 表示用户个人明确关闭限制")
used_bytes: int = Field(..., ge=0, description="当前用户已用有效资源容量，来自 UserResourceTotalStat.active_size_bytes")
available_bytes: int | None = Field(None, description="当前用户可用容量字节数。未开启容量管控时为 null")
total_bytes: int | None = Field(None, description="当前用户总容量字节数。未开启容量管控时为 null")
usage_percent: float | None = Field(None, description="当前用户容量使用百分比，未开启容量管控时为 null")
exceeded: bool = Field(..., description="是否已超额。判断规则：enabled=true 且 used_bytes >= total_bytes")
limit_value: str | None = Field(None, description="最终生效容量数值。未开启容量管控时为 null")
limit_unit: ResourceCapacityUnitEnum | None = Field(None, description="最终生效容量单位。未开启容量管控时为 null")
class AdminUserResourceCapacityOut(BaseModel):
has_user_config: bool = Field(..., description="该用户是否已单独设置容量配置")
user_config: ResourceCapacityConfigOut | None = Field(None, description="用户个人容量配置。未单独设置时为 null")
global_config: ResourceCapacityConfigOut = Field(..., description="当前全局容量配置")
effective: ResourceCapacityUsageOut = Field(..., description="该用户最终生效的容量使用数据")
video-gen-api/app/schemas/user.py
+3
View File
@@ -1,5 +1,7 @@
from pydantic import BaseModel
from app.schemas.resource_capacity import ResourceCapacityUsageOut
class UserOut(BaseModel):
id: str
@@ -12,5 +14,6 @@ class UserOut(BaseModel):
user_type: str = "frontend"
allowed_menus: list | None = None
must_set_password: bool = False
resource_capacity: ResourceCapacityUsageOut | None = None
model_config = {"from_attributes": True}
video-gen-api/app/services/generation_ai_service.py
+3
View File
@@ -36,6 +36,7 @@ from app.services.resource_accounting_service import (
soft_delete_chat_task_resources,
)
from app.services.resource_signed_url_service import build_resource_signed_url
from app.services.resource_capacity_service import assert_user_resource_capacity_available
from app.utils.id_gen import generate_id
IMAGE_DEFAULT_SIZE = "2K"
@@ -225,6 +226,8 @@ async def create_async_generation_task(db: AsyncSession, current_user: User, req
now = datetime.now(timezone.utc)
task_id = generate_id()
await assert_user_resource_capacity_available(db, current_user.id)
if gen_type == "image":
engine = await _get_image_engine(db, req.engine_id)
sizes = _image_supported_sizes(engine)
video-gen-api/app/services/generation_task_factory_service.py
+3
View File
@@ -26,6 +26,7 @@ from app.services.generation_ai_service import (
normalize_px,
)
from app.services.generation_billing_service import OWNER_CHAT_GENERATION_TASK, charge_generation_media_by_params
from app.services.resource_capacity_service import assert_user_resource_capacity_available
from app.utils.id_gen import generate_id
@@ -87,6 +88,8 @@ async def create_chat_generation_task_for_module(
task_id=task_id,
)
await assert_user_resource_capacity_available(db, current_user.id)
if gen_type == "image":
engine = await _get_image_engine(db, engine_id)
sizes = _image_supported_sizes(engine)
video-gen-api/app/services/resource_capacity_service.py
+440
View File
@@ -0,0 +1,440 @@
from __future__ import annotations
import json
from decimal import Decimal, ROUND_HALF_UP
from typing import Any, Iterable
from fastapi import HTTPException
from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import AsyncSession
from app.enums.resource_capacity import (
RESOURCE_CAPACITY_EXCEEDED_MESSAGE,
ResourceCapacityConfigKeyEnum,
ResourceCapacityErrorCodeEnum,
ResourceCapacitySourceEnum,
ResourceCapacityUnitEnum,
)
from app.models.system_config import SystemConfig
from app.models.user import User
from app.models.user_resource_capacity_config import UserResourceCapacityConfig
from app.models.user_resource_total_stat import UserResourceTotalStat
from app.schemas.resource_capacity import (
AdminUserResourceCapacityOut,
ResourceCapacityConfigOut,
ResourceCapacityConfigUpdate,
ResourceCapacityUsageOut,
)
from app.utils.id_gen import generate_id
DEFAULT_LIMIT_VALUE = Decimal("1.000")
DEFAULT_LIMIT_UNIT = ResourceCapacityUnitEnum.GB
DEFAULT_LIMIT_BYTES = DEFAULT_LIMIT_UNIT.bytes_multiplier
def _normalize_limit_value(value: Decimal | int | float | str | None) -> Decimal:
if value is None:
return DEFAULT_LIMIT_VALUE
decimal_value = Decimal(str(value))
return decimal_value.quantize(Decimal("0.001"), rounding=ROUND_HALF_UP)
def _limit_value_to_str(value: Decimal | int | float | str | None) -> str:
return format(_normalize_limit_value(value), "f")
def calculate_limit_bytes(
limit_value: Decimal | int | float | str | None,
limit_unit: ResourceCapacityUnitEnum | str | None,
) -> int:
unit = ResourceCapacityUnitEnum(limit_unit or DEFAULT_LIMIT_UNIT.value)
value = _normalize_limit_value(limit_value)
return int((value * Decimal(unit.bytes_multiplier)).to_integral_value(rounding=ROUND_HALF_UP))
def _default_config_out() -> ResourceCapacityConfigOut:
return ResourceCapacityConfigOut(
enabled=False,
limit_value=_limit_value_to_str(DEFAULT_LIMIT_VALUE),
limit_unit=DEFAULT_LIMIT_UNIT,
limit_bytes=DEFAULT_LIMIT_BYTES,
)
def _config_out(
*,
enabled: bool,
limit_value: Decimal | int | float | str | None,
limit_unit: ResourceCapacityUnitEnum | str | None,
limit_bytes: int | None = None,
) -> ResourceCapacityConfigOut:
unit = ResourceCapacityUnitEnum(limit_unit or DEFAULT_LIMIT_UNIT.value)
value = _normalize_limit_value(limit_value)
return ResourceCapacityConfigOut(
enabled=bool(enabled),
limit_value=_limit_value_to_str(value),
limit_unit=unit,
limit_bytes=int(limit_bytes if limit_bytes is not None else calculate_limit_bytes(value, unit)),
)
def _config_model_to_out(config: UserResourceCapacityConfig | None) -> ResourceCapacityConfigOut | None:
if config is None:
return None
return _config_out(
enabled=config.enabled,
limit_value=config.limit_value,
limit_unit=config.limit_unit,
limit_bytes=config.limit_bytes,
)
def _global_value_to_out(value: str | None) -> ResourceCapacityConfigOut:
if not value:
return _default_config_out()
try:
data = json.loads(value)
if not isinstance(data, dict):
return _default_config_out()
enabled = bool(data.get("enabled", False))
limit_unit = data.get("limit_unit") or DEFAULT_LIMIT_UNIT.value
limit_value = data.get("limit_value") or DEFAULT_LIMIT_VALUE
limit_bytes = data.get("limit_bytes")
return _config_out(
enabled=enabled,
limit_value=limit_value,
limit_unit=limit_unit,
limit_bytes=int(limit_bytes) if limit_bytes is not None else None,
)
except Exception:
return _default_config_out()
def _config_to_json(config: ResourceCapacityConfigOut) -> str:
return json.dumps(
{
"enabled": config.enabled,
"limit_value": config.limit_value,
"limit_unit": config.limit_unit.value,
"limit_bytes": config.limit_bytes,
},
ensure_ascii=False,
)
def _config_snapshot(config: ResourceCapacityConfigOut | None) -> dict[str, Any] | None:
if config is None:
return None
return {
"enabled": config.enabled,
"limit_value": config.limit_value,
"limit_unit": config.limit_unit.value,
"limit_bytes": config.limit_bytes,
}
def _usage_snapshot(usage: ResourceCapacityUsageOut | None) -> dict[str, Any] | None:
if usage is None:
return None
return {
"enabled": usage.enabled,
"source": usage.source.value,
"has_user_config": usage.has_user_config,
"used_bytes": usage.used_bytes,
"available_bytes": usage.available_bytes,
"total_bytes": usage.total_bytes,
"usage_percent": usage.usage_percent,
"exceeded": usage.exceeded,
"limit_value": usage.limit_value,
"limit_unit": usage.limit_unit.value if usage.limit_unit else None,
}
def build_global_resource_capacity_operation_detail(
before: ResourceCapacityConfigOut | None,
after: ResourceCapacityConfigOut | None,
) -> str:
"""构造全局容量配置操作日志详情。"""
return json.dumps(
{
"target": "global_resource_capacity",
"before": _config_snapshot(before),
"after": _config_snapshot(after),
},
ensure_ascii=False,
default=str,
)
def build_user_resource_capacity_operation_detail(
*,
target_user_id: str,
operation: str,
before: AdminUserResourceCapacityOut | None,
after: AdminUserResourceCapacityOut | None,
remark: str | None = None,
) -> str:
"""构造用户个人容量配置操作日志详情。"""
payload: dict[str, Any] = {
"target": "user_resource_capacity",
"target_user_id": target_user_id,
"operation": operation,
"before": {
"has_user_config": before.has_user_config if before else False,
"user_config": _config_snapshot(before.user_config) if before else None,
"effective": _usage_snapshot(before.effective) if before else None,
},
"after": {
"has_user_config": after.has_user_config if after else False,
"user_config": _config_snapshot(after.user_config) if after else None,
"effective": _usage_snapshot(after.effective) if after else None,
},
}
if remark:
payload["remark"] = remark
return json.dumps(payload, ensure_ascii=False, default=str)
def _build_usage_out(
*,
used_bytes: int,
has_user_config: bool,
source: ResourceCapacitySourceEnum,
config: ResourceCapacityConfigOut | None,
) -> ResourceCapacityUsageOut:
used = max(int(used_bytes or 0), 0)
if not config or not config.enabled:
return ResourceCapacityUsageOut(
enabled=False,
source=ResourceCapacitySourceEnum.DISABLED,
has_user_config=has_user_config,
used_bytes=used,
available_bytes=None,
total_bytes=None,
usage_percent=None,
exceeded=False,
limit_value=None,
limit_unit=None,
)
total = max(int(config.limit_bytes or 0), 0)
available = max(total - used, 0)
usage_percent = round((used / total) * 100, 2) if total > 0 else None
return ResourceCapacityUsageOut(
enabled=True,
source=source,
has_user_config=has_user_config,
used_bytes=used,
available_bytes=available,
total_bytes=total,
usage_percent=usage_percent,
exceeded=bool(total > 0 and used >= total),
limit_value=config.limit_value,
limit_unit=config.limit_unit,
)
async def get_global_resource_capacity_config(db: AsyncSession) -> ResourceCapacityConfigOut:
result = await db.execute(
select(SystemConfig.value)
.where(SystemConfig.key == ResourceCapacityConfigKeyEnum.RESOURCE_CAPACITY_LIMIT_CONFIG.value)
.limit(1)
)
return _global_value_to_out(result.scalar_one_or_none())
async def save_global_resource_capacity_config(
db: AsyncSession,
req: ResourceCapacityConfigUpdate,
) -> ResourceCapacityConfigOut:
limit_value = req.limit_value if req.limit_value is not None else DEFAULT_LIMIT_VALUE
limit_unit = req.limit_unit if req.limit_unit is not None else DEFAULT_LIMIT_UNIT
config_out = _config_out(
enabled=req.enabled,
limit_value=limit_value,
limit_unit=limit_unit,
)
result = await db.execute(
select(SystemConfig)
.where(SystemConfig.key == ResourceCapacityConfigKeyEnum.RESOURCE_CAPACITY_LIMIT_CONFIG.value)
.limit(1)
)
config = result.scalar_one_or_none()
if config:
config.value = _config_to_json(config_out)
config.description = "全局生成资源空间容量管控配置"
else:
db.add(
SystemConfig(
id=generate_id(),
key=ResourceCapacityConfigKeyEnum.RESOURCE_CAPACITY_LIMIT_CONFIG.value,
value=_config_to_json(config_out),
description="全局生成资源空间容量管控配置",
)
)
await db.flush()
return config_out
async def _get_user_config(db: AsyncSession, user_id: str) -> UserResourceCapacityConfig | None:
result = await db.execute(
select(UserResourceCapacityConfig)
.where(UserResourceCapacityConfig.user_id == user_id)
.limit(1)
)
return result.scalar_one_or_none()
async def _get_used_bytes(db: AsyncSession, user_id: str) -> int:
result = await db.execute(
select(UserResourceTotalStat.active_size_bytes)
.where(UserResourceTotalStat.user_id == user_id)
.limit(1)
)
return int(result.scalar_one_or_none() or 0)
async def get_user_resource_capacity_usage(
db: AsyncSession,
user_id: str,
) -> ResourceCapacityUsageOut:
global_config = await get_global_resource_capacity_config(db)
user_config = await _get_user_config(db, user_id)
used_bytes = await _get_used_bytes(db, user_id)
user_config_out = _config_model_to_out(user_config)
if user_config_out is not None:
return _build_usage_out(
used_bytes=used_bytes,
has_user_config=True,
source=ResourceCapacitySourceEnum.USER if user_config_out.enabled else ResourceCapacitySourceEnum.DISABLED,
config=user_config_out,
)
return _build_usage_out(
used_bytes=used_bytes,
has_user_config=False,
source=ResourceCapacitySourceEnum.GLOBAL if global_config.enabled else ResourceCapacitySourceEnum.DISABLED,
config=global_config,
)
async def batch_get_user_resource_capacity_usage(
db: AsyncSession,
user_ids: Iterable[str],
) -> dict[str, ResourceCapacityUsageOut]:
ids = [user_id for user_id in dict.fromkeys(user_ids) if user_id]
if not ids:
return {}
global_config = await get_global_resource_capacity_config(db)
config_result = await db.execute(
select(UserResourceCapacityConfig)
.where(UserResourceCapacityConfig.user_id.in_(ids))
)
user_config_map = {item.user_id: item for item in config_result.scalars().all()}
stat_result = await db.execute(
select(UserResourceTotalStat.user_id, UserResourceTotalStat.active_size_bytes)
.where(UserResourceTotalStat.user_id.in_(ids))
)
used_map = {row.user_id: int(row.active_size_bytes or 0) for row in stat_result.all()}
usage_map: dict[str, ResourceCapacityUsageOut] = {}
for user_id in ids:
user_config_out = _config_model_to_out(user_config_map.get(user_id))
if user_config_out is not None:
usage_map[user_id] = _build_usage_out(
used_bytes=used_map.get(user_id, 0),
has_user_config=True,
source=ResourceCapacitySourceEnum.USER if user_config_out.enabled else ResourceCapacitySourceEnum.DISABLED,
config=user_config_out,
)
else:
usage_map[user_id] = _build_usage_out(
used_bytes=used_map.get(user_id, 0),
has_user_config=False,
source=ResourceCapacitySourceEnum.GLOBAL if global_config.enabled else ResourceCapacitySourceEnum.DISABLED,
config=global_config,
)
return usage_map
async def assert_user_resource_capacity_available(db: AsyncSession, user_id: str) -> None:
usage = await get_user_resource_capacity_usage(db, user_id)
if usage.enabled and usage.exceeded:
raise HTTPException(
status_code=400,
detail=RESOURCE_CAPACITY_EXCEEDED_MESSAGE,
headers={"X-Error-Code": ResourceCapacityErrorCodeEnum.RESOURCE_CAPACITY_EXCEEDED.value},
)
async def ensure_user_exists(db: AsyncSession, user_id: str) -> User:
result = await db.execute(select(User).where(User.id == user_id).limit(1))
user = result.scalar_one_or_none()
if not user:
raise HTTPException(status_code=404, detail="用户不存在")
return user
async def get_admin_user_resource_capacity(
db: AsyncSession,
user_id: str,
) -> AdminUserResourceCapacityOut:
await ensure_user_exists(db, user_id)
global_config = await get_global_resource_capacity_config(db)
user_config = await _get_user_config(db, user_id)
effective = await get_user_resource_capacity_usage(db, user_id)
return AdminUserResourceCapacityOut(
has_user_config=user_config is not None,
user_config=_config_model_to_out(user_config),
global_config=global_config,
effective=effective,
)
async def save_user_resource_capacity_config(
db: AsyncSession,
user_id: str,
req: ResourceCapacityConfigUpdate,
) -> AdminUserResourceCapacityOut:
await ensure_user_exists(db, user_id)
limit_value = req.limit_value if req.limit_value is not None else DEFAULT_LIMIT_VALUE
limit_unit = req.limit_unit if req.limit_unit is not None else DEFAULT_LIMIT_UNIT
limit_bytes = calculate_limit_bytes(limit_value, limit_unit)
config = await _get_user_config(db, user_id)
if config:
config.enabled = req.enabled
config.limit_value = _normalize_limit_value(limit_value)
config.limit_unit = ResourceCapacityUnitEnum(limit_unit).value
config.limit_bytes = limit_bytes
else:
db.add(
UserResourceCapacityConfig(
id=generate_id(),
user_id=user_id,
enabled=req.enabled,
limit_value=_normalize_limit_value(limit_value),
limit_unit=ResourceCapacityUnitEnum(limit_unit).value,
limit_bytes=limit_bytes,
)
)
await db.flush()
return await get_admin_user_resource_capacity(db, user_id)
async def delete_user_resource_capacity_config(
db: AsyncSession,
user_id: str,
) -> AdminUserResourceCapacityOut:
await ensure_user_exists(db, user_id)
await db.execute(
delete(UserResourceCapacityConfig).where(UserResourceCapacityConfig.user_id == user_id)
)
await db.flush()
return await get_admin_user_resource_capacity(db, user_id)