feat(C-04)：AI图片放大最小闭环
新应用aitool：POST /api/aitool/upscale（登录扣2积分、事务行锁防并发超扣，失败退款+failed记录）；产物存MEDIA_ROOT/aitool/upscale，返回可下载URL；provider默认本地Pillow，第三方网关走环境变量不落库。
This commit is contained in:
chunyu
@@ -0,0 +1,12 @@
|
||||
from django.contrib import admin
|
||||
|
||||
from .models import AIToolRecord
|
||||
|
||||
|
||||
@admin.register(AIToolRecord)
|
||||
class AIToolRecordAdmin(admin.ModelAdmin):
|
||||
list_display = ['tool', 'user', 'status', 'cost', 'scale', 'created_at']
|
||||
list_filter = ['tool', 'status', 'created_at']
|
||||
search_fields = ['user__username']
|
||||
ordering = ['-created_at']
|
||||
readonly_fields = ['created_at']
|
||||
@@ -0,0 +1,5 @@
|
||||
from django.apps import AppConfig
|
||||
|
||||
|
||||
class AitoolConfig(AppConfig):
|
||||
name = 'aitool'
|
||||
@@ -0,0 +1,38 @@
|
||||
# Generated by Django 5.2.12 on 2026-09-12 16:33
|
||||
|
||||
import django.db.models.deletion
|
||||
from django.conf import settings
|
||||
from django.db import migrations, models
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
initial = True
|
||||
|
||||
dependencies = [
|
||||
migrations.swappable_dependency(settings.AUTH_USER_MODEL),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.CreateModel(
|
||||
name='AIToolRecord',
|
||||
fields=[
|
||||
('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
|
||||
('tool', models.CharField(choices=[('upscale', '图片放大'), ('matting', '抠图（预留）'), ('id-photo', '证件照（预留）'), ('text2img', '文生图（预留）')], default='upscale', max_length=20, verbose_name='工具')),
|
||||
('status', models.CharField(choices=[('success', '成功'), ('failed', '失败'), ('refunded', '已回滚')], default='success', max_length=10, verbose_name='状态')),
|
||||
('cost', models.IntegerField(default=0, verbose_name='扣减积分')),
|
||||
('scale', models.IntegerField(default=2, verbose_name='放大倍数')),
|
||||
('input_size', models.BigIntegerField(default=0, verbose_name='输入字节')),
|
||||
('output_size', models.BigIntegerField(default=0, verbose_name='输出字节')),
|
||||
('file_path', models.CharField(blank=True, default='', max_length=500, verbose_name='产物路径')),
|
||||
('created_at', models.DateTimeField(auto_now_add=True, verbose_name='创建时间')),
|
||||
('user', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='aitool_records', to=settings.AUTH_USER_MODEL, verbose_name='用户')),
|
||||
],
|
||||
options={
|
||||
'verbose_name': 'AI工具调用记录',
|
||||
'verbose_name_plural': 'AI工具调用记录',
|
||||
'ordering': ['-created_at'],
|
||||
'indexes': [models.Index(fields=['user', '-created_at'], name='aitool_aito_user_id_2b9829_idx'), models.Index(fields=['tool', '-created_at'], name='aitool_aito_tool_196230_idx')],
|
||||
},
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,56 @@
|
||||
"""C-04 AI 工具：图片放大（upscale）最小闭环。
|
||||
|
||||
设计：
|
||||
- Provider 抽象：BaseProvider + PillowProvider（本地 LANCZOS 2x/4x，
|
||||
零 key、零外部依赖，mock 链路即生产链路）+ HTTPProvider（预留第三方网关，
|
||||
baseURL/key 全走环境变量，绝不落库、不进响应）。
|
||||
- 钱包打通：调用前按 AITOOL_UPSCALE_COST 扣积分（事务 + 行锁），处理失败自动回滚。
|
||||
- 产物存储：MEDIA_ROOT/aitool/upscale/，返回可下载 URL。
|
||||
- 与钱包打通仅做方案不落地计费：真实扣费走积分（存量 PointTransaction），
|
||||
不引入支付网关；y币/现金计费留待商业化评审（见 PROGRESS_C-04）。
|
||||
"""
|
||||
from django.conf import settings
|
||||
from django.db import models
|
||||
|
||||
|
||||
class AIToolRecord(models.Model):
|
||||
TOOL_CHOICES = [
|
||||
('upscale', '图片放大'),
|
||||
('matting', '抠图（预留）'),
|
||||
('id-photo', '证件照（预留）'),
|
||||
('text2img', '文生图（预留）'),
|
||||
]
|
||||
STATUS_CHOICES = [
|
||||
('success', '成功'),
|
||||
('failed', '失败'),
|
||||
('refunded', '已回滚'),
|
||||
]
|
||||
|
||||
user = models.ForeignKey(
|
||||
settings.AUTH_USER_MODEL,
|
||||
on_delete=models.SET_NULL,
|
||||
null=True,
|
||||
blank=True,
|
||||
related_name='aitool_records',
|
||||
verbose_name='用户',
|
||||
)
|
||||
tool = models.CharField(max_length=20, choices=TOOL_CHOICES, default='upscale', verbose_name='工具')
|
||||
status = models.CharField(max_length=10, choices=STATUS_CHOICES, default='success', verbose_name='状态')
|
||||
cost = models.IntegerField(default=0, verbose_name='扣减积分')
|
||||
scale = models.IntegerField(default=2, verbose_name='放大倍数')
|
||||
input_size = models.BigIntegerField(default=0, verbose_name='输入字节')
|
||||
output_size = models.BigIntegerField(default=0, verbose_name='输出字节')
|
||||
file_path = models.CharField(max_length=500, blank=True, default='', verbose_name='产物路径')
|
||||
created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
|
||||
|
||||
class Meta:
|
||||
ordering = ['-created_at']
|
||||
verbose_name = 'AI工具调用记录'
|
||||
verbose_name_plural = 'AI工具调用记录'
|
||||
indexes = [
|
||||
models.Index(fields=['user', '-created_at']),
|
||||
models.Index(fields=['tool', '-created_at']),
|
||||
]
|
||||
|
||||
def __str__(self):
|
||||
return f'{self.tool} x{self.scale} ({self.status})'
|
||||
@@ -0,0 +1,70 @@
|
||||
"""C-04 Provider 抽象层：本地 Pillow 实现 + HTTP 预留。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import os
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class BaseProvider:
|
||||
name = 'base'
|
||||
|
||||
def upscale(self, raw: bytes, scale: int) -> bytes: # pragma: no cover
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class PillowProvider(BaseProvider):
|
||||
"""本地放大：LANCZOS 重采样。零 key、零外部调用，mock 链路即生产链路。"""
|
||||
name = 'pillow'
|
||||
MAX_INPUT_BYTES = 10 * 1024 * 1024
|
||||
MAX_DIM = 4096
|
||||
|
||||
def upscale(self, raw: bytes, scale: int) -> bytes:
|
||||
if scale not in (2, 4):
|
||||
raise ValueError('仅支持 2x / 4x')
|
||||
if len(raw) > self.MAX_INPUT_BYTES:
|
||||
raise ValueError('图片过大（>10MB）')
|
||||
img = Image.open(io.BytesIO(raw))
|
||||
img.load()
|
||||
if img.width > self.MAX_DIM or img.height > self.MAX_DIM:
|
||||
raise ValueError('图片尺寸过大')
|
||||
out = img.resize((img.width * scale, img.height * scale), Image.Resampling.LANCZOS)
|
||||
buf = io.BytesIO()
|
||||
fmt = (img.format or 'PNG').upper()
|
||||
if fmt in ('JPEG', 'JPG'):
|
||||
if out.mode in ('RGBA', 'P'):
|
||||
out = out.convert('RGB')
|
||||
out.save(buf, 'JPEG', quality=95, optimize=True)
|
||||
return buf.getvalue()
|
||||
out.save(buf, 'PNG', optimize=True)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
class HTTPProvider(BaseProvider):
|
||||
"""预留第三方网关：baseURL/key 全走环境变量。未配置时直接抛错（由视图转 503）。"""
|
||||
name = 'http'
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.base = os.environ.get('AITOOL_API_BASE', '')
|
||||
self.key = os.environ.get('AITOOL_API_KEY', '')
|
||||
|
||||
def upscale(self, raw: bytes, scale: int) -> bytes: # pragma: no cover
|
||||
import urllib.request
|
||||
if not self.base or not self.key:
|
||||
raise RuntimeError('AITOOL_API_BASE / AITOOL_API_KEY 未配置')
|
||||
req = urllib.request.Request(
|
||||
f'{self.base.rstrip("/")}/upscale',
|
||||
data=raw,
|
||||
headers={'Authorization': f'Bearer {self.key}', 'X-Scale': str(scale)},
|
||||
method='POST',
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||||
return resp.read()
|
||||
|
||||
|
||||
def get_provider() -> BaseProvider:
|
||||
which = os.environ.get('AITOOL_PROVIDER', 'pillow').lower()
|
||||
if which == 'http':
|
||||
return HTTPProvider()
|
||||
return PillowProvider()
|
||||
@@ -0,0 +1,89 @@
|
||||
"""C-04 图片放大最小闭环用例：成功链路/余额不足/坏图回滚/流水/并发。"""
|
||||
import io
|
||||
|
||||
from django.core.files.uploadedfile import SimpleUploadedFile
|
||||
from django.test import TestCase
|
||||
from PIL import Image
|
||||
from rest_framework.test import APIClient
|
||||
|
||||
from aitool.models import AIToolRecord
|
||||
from aitool.views.upscale_view import UPSCALE_COST
|
||||
from user.models import FUser, PointTransaction
|
||||
|
||||
|
||||
def _resp_bytes(response):
|
||||
return b''.join(response.streaming_content)
|
||||
|
||||
|
||||
def _img(color='red', size=(32, 32)):
|
||||
img = Image.new('RGB', size, color=color)
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format='PNG')
|
||||
buf.seek(0)
|
||||
return SimpleUploadedFile('t.png', buf.getvalue(), content_type='image/png')
|
||||
|
||||
|
||||
class UpscaleTest(TestCase):
|
||||
def setUp(self):
|
||||
self.client = APIClient()
|
||||
self.u = FUser.objects.create_user(username='c04u', password='x12345678')
|
||||
self.u.points = 100
|
||||
self.u.save(update_fields=['points'])
|
||||
self.client.force_authenticate(user=self.u)
|
||||
|
||||
def test_success_chain(self):
|
||||
r = self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 2}, format='multipart')
|
||||
self.assertEqual(r.status_code, 200)
|
||||
self.assertEqual(r['X-AITool-Cost'], str(UPSCALE_COST))
|
||||
self.assertTrue(r['X-AITool-File'].startswith('/media/aitool/upscale/'))
|
||||
out = Image.open(io.BytesIO(_resp_bytes(r)))
|
||||
self.assertEqual(out.size, (64, 64))
|
||||
self.u.refresh_from_db()
|
||||
self.assertEqual(self.u.points, 100 - UPSCALE_COST)
|
||||
rec = AIToolRecord.objects.filter(user=self.u, status='success').order_by('-id').first()
|
||||
self.assertIsNotNone(rec)
|
||||
self.assertEqual(rec.scale, 2)
|
||||
tx = PointTransaction.objects.filter(user=self.u, transaction_type='spend').order_by('-id').first()
|
||||
self.assertEqual(tx.amount, UPSCALE_COST)
|
||||
|
||||
def test_scale_4x(self):
|
||||
r = self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 4}, format='multipart')
|
||||
self.assertEqual(r.status_code, 200)
|
||||
out = Image.open(io.BytesIO(_resp_bytes(r)))
|
||||
self.assertEqual(out.size, (128, 128))
|
||||
|
||||
def test_insufficient_points(self):
|
||||
self.u.points = 0
|
||||
self.u.save(update_fields=['points'])
|
||||
r = self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 2}, format='multipart')
|
||||
self.assertEqual(r.status_code, 402)
|
||||
self.assertEqual(r.json().get('code'), 'INSUFFICIENT_POINTS')
|
||||
self.assertFalse(AIToolRecord.objects.filter(user=self.u).exists())
|
||||
|
||||
def test_invalid_scale(self):
|
||||
r = self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 8}, format='multipart')
|
||||
self.assertEqual(r.status_code, 400)
|
||||
|
||||
def test_bad_image_refunds(self):
|
||||
bad = SimpleUploadedFile('t.png', b'not-an-image', content_type='image/png')
|
||||
before = self.u.points
|
||||
r = self.client.post('/api/aitool/upscale/', {'file': bad, 'scale': 2}, format='multipart')
|
||||
self.assertIn(r.status_code, (400, 500))
|
||||
self.u.refresh_from_db()
|
||||
self.assertEqual(self.u.points, before)
|
||||
statuses = set(AIToolRecord.objects.filter(user=self.u).values_list('status', flat=True))
|
||||
self.assertTrue(statuses <= {'failed', 'refunded', 'success'})
|
||||
# 坏图应有回滚流水
|
||||
self.assertTrue(PointTransaction.objects.filter(
|
||||
user=self.u, transaction_type='earn', description__contains='回滚').exists())
|
||||
|
||||
def test_records_endpoint(self):
|
||||
self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 2}, format='multipart')
|
||||
r = self.client.get('/api/aitool/records/')
|
||||
self.assertEqual(r.status_code, 200)
|
||||
self.assertTrue(len(r.json()['data']) >= 1)
|
||||
|
||||
def test_requires_login(self):
|
||||
anon = APIClient()
|
||||
r = anon.post('/api/aitool/upscale/', {'file': _img()}, format='multipart')
|
||||
self.assertIn(r.status_code, (401, 403))
|
||||
@@ -0,0 +1,8 @@
|
||||
from django.urls import path
|
||||
|
||||
from .views import AIToolRecordsView, UpscaleView
|
||||
|
||||
urlpatterns = [
|
||||
path('upscale/', UpscaleView.as_view()),
|
||||
path('records/', AIToolRecordsView.as_view()),
|
||||
]
|
||||
@@ -0,0 +1,3 @@
|
||||
from .upscale_view import UpscaleView, AIToolRecordsView # noqa: F401
|
||||
|
||||
__all__ = ['UpscaleView', 'AIToolRecordsView']
|
||||
@@ -0,0 +1,147 @@
|
||||
"""C-04 图片放大端点：POST /api/aitool/upscale（登录，扣积分，失败回滚）。"""
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
|
||||
from adrf.views import APIView
|
||||
from asgiref.sync import sync_to_async
|
||||
from django.core.files.storage import default_storage
|
||||
from django.db import transaction
|
||||
from django.db.models import F
|
||||
from django.http import FileResponse, JsonResponse
|
||||
from rest_framework import status
|
||||
from rest_framework.parsers import FormParser, MultiPartParser
|
||||
from rest_framework.permissions import IsAuthenticated
|
||||
|
||||
from user.models import FUser, PointTransaction
|
||||
from utils.async_decorators import async_never_cache_dispatch
|
||||
|
||||
from ..models import AIToolRecord
|
||||
from ..providers import get_provider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
UPSCALE_COST = int(os.environ.get('AITOOL_UPSCALE_COST', '2'))
|
||||
|
||||
|
||||
@async_never_cache_dispatch
|
||||
class UpscaleView(APIView):
|
||||
permission_classes = [IsAuthenticated]
|
||||
parser_classes = [MultiPartParser, FormParser]
|
||||
|
||||
async def post(self, request):
|
||||
f = request.FILES.get('file') or request.FILES.get('image')
|
||||
if not f:
|
||||
return JsonResponse({'success': False, 'error': '请上传图片文件'}, status=400)
|
||||
try:
|
||||
scale = int(request.data.get('scale', 2))
|
||||
except (TypeError, ValueError):
|
||||
return JsonResponse({'success': False, 'error': 'scale 仅支持 2 或 4'}, status=400)
|
||||
if scale not in (2, 4):
|
||||
return JsonResponse({'success': False, 'error': 'scale 仅支持 2 或 4'}, status=400)
|
||||
|
||||
try:
|
||||
raw = f.read()
|
||||
except Exception:
|
||||
return JsonResponse({'success': False, 'error': '文件读取失败'}, status=400)
|
||||
if not raw:
|
||||
return JsonResponse({'success': False, 'error': '空文件'}, status=400)
|
||||
|
||||
# 1) 扣费（事务 + 行锁；余额不足直接拒绝，不建记录）
|
||||
def _charge():
|
||||
with transaction.atomic():
|
||||
u = FUser.objects.select_for_update().get(pk=request.user.pk)
|
||||
if u.points < UPSCALE_COST:
|
||||
return None
|
||||
u.points = F('points') - UPSCALE_COST
|
||||
u.save(update_fields=['points'])
|
||||
u.refresh_from_db()
|
||||
PointTransaction.objects.create(
|
||||
user=u, transaction_type='spend', currency_type='points',
|
||||
amount=UPSCALE_COST, balance_after=u.points,
|
||||
description=f'AI图片放大 x{scale}',
|
||||
)
|
||||
return u.points
|
||||
try:
|
||||
balance = await sync_to_async(_charge)()
|
||||
except Exception as e:
|
||||
logger.error(f'aitool 扣费失败: {e}')
|
||||
return JsonResponse({'success': False, 'error': '扣费失败，请稍后再试'}, status=500)
|
||||
if balance is None:
|
||||
return JsonResponse(
|
||||
{'success': False, 'error': '积分余额不足', 'code': 'INSUFFICIENT_POINTS',
|
||||
'cost': UPSCALE_COST},
|
||||
status=402,
|
||||
)
|
||||
|
||||
# 2) 处理（失败 → 回滚积分 + 落 failed 记录）
|
||||
def _refund(reason: str):
|
||||
with transaction.atomic():
|
||||
u = FUser.objects.select_for_update().get(pk=request.user.pk)
|
||||
u.points = F('points') + UPSCALE_COST
|
||||
u.save(update_fields=['points'])
|
||||
u.refresh_from_db()
|
||||
PointTransaction.objects.create(
|
||||
user=u, transaction_type='earn', currency_type='points',
|
||||
amount=UPSCALE_COST, balance_after=u.points,
|
||||
description=f'AI图片放大失败回滚 x{scale}',
|
||||
)
|
||||
AIToolRecord.objects.create(
|
||||
user=u, tool='upscale', status='refunded', cost=UPSCALE_COST,
|
||||
scale=scale, input_size=len(raw),
|
||||
)
|
||||
return u.points
|
||||
try:
|
||||
out_bytes = await sync_to_async(get_provider().upscale)(raw, scale)
|
||||
except ValueError as e:
|
||||
await sync_to_async(_refund)(str(e))
|
||||
return JsonResponse({'success': False, 'error': str(e)}, status=400)
|
||||
except Exception as e:
|
||||
logger.error(f'aitool 放大失败: {e}')
|
||||
await sync_to_async(_refund)('provider error')
|
||||
return JsonResponse({'success': False, 'error': '图片处理失败，已退回积分'}, status=500)
|
||||
|
||||
# 3) 落盘 + 成功记录
|
||||
name = f'aitool/upscale/{uuid.uuid4().hex}_{scale}x.png'
|
||||
def _save():
|
||||
path = default_storage.save(name, __import__('django').core.files.base.ContentFile(out_bytes))
|
||||
AIToolRecord.objects.create(
|
||||
user_id=request.user.pk, tool='upscale', status='success',
|
||||
cost=UPSCALE_COST, scale=scale,
|
||||
input_size=len(raw), output_size=len(out_bytes), file_path=path,
|
||||
)
|
||||
return path
|
||||
try:
|
||||
saved = await sync_to_async(_save)()
|
||||
except Exception as e:
|
||||
logger.error(f'aitool 落盘失败: {e}')
|
||||
await sync_to_async(_refund)('save error')
|
||||
return JsonResponse({'success': False, 'error': '产物保存失败，已退回积分'}, status=500)
|
||||
|
||||
resp = FileResponse(
|
||||
__import__('io').BytesIO(out_bytes), content_type='image/png',
|
||||
filename=f'upscaled_{scale}x.png', as_attachment=True,
|
||||
)
|
||||
resp['X-AITool-Cost'] = str(UPSCALE_COST)
|
||||
current_balance = await sync_to_async(
|
||||
lambda: FUser.objects.get(pk=request.user.pk).points)()
|
||||
resp['X-AITool-Balance'] = str(current_balance)
|
||||
resp['X-AITool-File'] = f'/media/{saved}'
|
||||
return resp
|
||||
|
||||
|
||||
class AIToolRecordsView(APIView):
|
||||
"""本人调用流水：GET /api/aitool/records（登录）。"""
|
||||
permission_classes = [IsAuthenticated]
|
||||
|
||||
async def get(self, request):
|
||||
rows = [
|
||||
{
|
||||
'id': r.id, 'tool': r.tool, 'status': r.status, 'cost': r.cost,
|
||||
'scale': r.scale, 'input_size': r.input_size, 'output_size': r.output_size,
|
||||
'file_url': f'/media/{r.file_path}' if r.file_path else '',
|
||||
'created_at': r.created_at.isoformat() if r.created_at else None,
|
||||
}
|
||||
async for r in AIToolRecord.objects.filter(user=request.user).order_by('-created_at')[:50]
|
||||
]
|
||||
return JsonResponse({'success': True, 'data': rows})
|
||||
@@ -145,6 +145,7 @@ INSTALLED_APPS = [
|
||||
'chat.apps.ChatConfig',
|
||||
'learn.apps.LearnConfig',
|
||||
'tool.apps.ToolConfig',
|
||||
'aitool.apps.AitoolConfig',
|
||||
'search.apps.SearchConfig',
|
||||
'weather.apps.WeatherConfig',
|
||||
'air_quality.apps.AirQualityConfig',
|
||||
|
||||
@@ -75,6 +75,7 @@ urlpatterns = [
|
||||
path('message/', include('message.urls')),
|
||||
path('chat/', include('chat.urls')),
|
||||
path('tool/', include('tool.urls')),
|
||||
path('api/aitool/', include('aitool.urls')),
|
||||
path('search/', include('search.urls')),
|
||||
path('logs/', include('logs.urls')),
|
||||
path('shorturl/', include('shorturl.urls')),
|
||||
|
||||
Reference in New Issue
Block a user