feat: add AirShelf core implementation

This commit is contained in:
zyc
2026-06-05 10:21:40 +08:00
parent 2ba1058329
commit cfdcd84a30
252 changed files with 70828 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
+21
View File
@@ -0,0 +1,21 @@
from django.contrib import admin
from .models import CreditAccount, CreditLedger, CreditReservation, QuotaPolicy
@admin.register(CreditAccount)
class CreditAccountAdmin(admin.ModelAdmin):
list_display = ("team", "balance", "reserved_balance", "currency", "updated_at")
search_fields = ("team__name",)
@admin.register(CreditLedger)
class CreditLedgerAdmin(admin.ModelAdmin):
list_display = ("team", "user", "project", "task", "ledger_type", "amount", "balance_after", "created_at")
search_fields = ("team__name", "user__username", "project__name", "task__idempotency_key")
list_filter = ("ledger_type",)
admin.site.register(CreditReservation)
admin.site.register(QuotaPolicy)
+7
View File
@@ -0,0 +1,7 @@
from django.apps import AppConfig
class BillingConfig(AppConfig):
default_auto_field = "django.db.models.BigAutoField"
name = "apps.billing"
@@ -0,0 +1,277 @@
# Generated by Django 5.1.15 on 2026-05-29 03:59
import django.db.models.deletion
import uuid
from django.conf import settings
from django.db import migrations, models
class Migration(migrations.Migration):
initial = True
dependencies = [
("accounts", "0001_initial"),
("ai", "0002_initial"),
("projects", "0001_initial"),
migrations.swappable_dependency(settings.AUTH_USER_MODEL),
]
operations = [
migrations.CreateModel(
name="CreditAccount",
fields=[
(
"id",
models.UUIDField(
default=uuid.uuid4,
editable=False,
primary_key=True,
serialize=False,
),
),
("created_at", models.DateTimeField(auto_now_add=True)),
("updated_at", models.DateTimeField(auto_now=True)),
(
"balance",
models.DecimalField(decimal_places=4, default=0, max_digits=14),
),
(
"reserved_balance",
models.DecimalField(decimal_places=4, default=0, max_digits=14),
),
("currency", models.CharField(default="CNY", max_length=16)),
(
"team",
models.OneToOneField(
on_delete=django.db.models.deletion.CASCADE,
related_name="credit_account",
to="accounts.team",
),
),
],
options={
"abstract": False,
},
),
migrations.CreateModel(
name="CreditReservation",
fields=[
(
"id",
models.UUIDField(
default=uuid.uuid4,
editable=False,
primary_key=True,
serialize=False,
),
),
("created_at", models.DateTimeField(auto_now_add=True)),
("updated_at", models.DateTimeField(auto_now=True)),
("amount", models.DecimalField(decimal_places=4, max_digits=14)),
(
"status",
models.CharField(
choices=[
("active", "Active"),
("released", "Released"),
("charged", "Charged"),
("cancelled", "Cancelled"),
],
default="active",
max_length=32,
),
),
("expires_at", models.DateTimeField(blank=True, null=True)),
(
"project",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="credit_reservations",
to="projects.project",
),
),
(
"task",
models.OneToOneField(
on_delete=django.db.models.deletion.CASCADE,
related_name="credit_reservation",
to="ai.aitask",
),
),
(
"team",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="credit_reservations",
to="accounts.team",
),
),
(
"user",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="credit_reservations",
to=settings.AUTH_USER_MODEL,
),
),
],
options={
"abstract": False,
},
),
migrations.CreateModel(
name="QuotaPolicy",
fields=[
(
"id",
models.UUIDField(
default=uuid.uuid4,
editable=False,
primary_key=True,
serialize=False,
),
),
("created_at", models.DateTimeField(auto_now_add=True)),
("updated_at", models.DateTimeField(auto_now=True)),
(
"monthly_limit",
models.DecimalField(
blank=True, decimal_places=4, max_digits=14, null=True
),
),
(
"project_limit",
models.DecimalField(
blank=True, decimal_places=4, max_digits=14, null=True
),
),
(
"per_task_limit",
models.DecimalField(
blank=True, decimal_places=4, max_digits=14, null=True
),
),
("is_active", models.BooleanField(default=True)),
(
"project",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="quota_policies",
to="projects.project",
),
),
(
"team",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="quota_policies",
to="accounts.team",
),
),
(
"user",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="quota_policies",
to=settings.AUTH_USER_MODEL,
),
),
],
options={
"abstract": False,
},
),
migrations.CreateModel(
name="CreditLedger",
fields=[
(
"id",
models.UUIDField(
default=uuid.uuid4,
editable=False,
primary_key=True,
serialize=False,
),
),
("created_at", models.DateTimeField(auto_now_add=True)),
("updated_at", models.DateTimeField(auto_now=True)),
(
"ledger_type",
models.CharField(
choices=[
("recharge", "Recharge"),
("reserve", "Reserve"),
("release", "Release"),
("charge", "Charge"),
("adjustment", "Adjustment"),
("refund", "Refund"),
],
max_length=32,
),
),
("amount", models.DecimalField(decimal_places=4, max_digits=14)),
("balance_after", models.DecimalField(decimal_places=4, max_digits=14)),
("reason", models.CharField(blank=True, max_length=255)),
("metadata", models.JSONField(blank=True, default=dict)),
(
"project",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="credit_ledgers",
to="projects.project",
),
),
(
"task",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="credit_ledgers",
to="ai.aitask",
),
),
(
"team",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="credit_ledgers",
to="accounts.team",
),
),
(
"user",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="credit_ledgers",
to=settings.AUTH_USER_MODEL,
),
),
],
options={
"indexes": [
models.Index(
fields=["team", "ledger_type"],
name="billing_cre_team_id_e0f18f_idx",
),
models.Index(
fields=["project", "task"],
name="billing_cre_project_a79834_idx",
),
],
},
),
]
@@ -0,0 +1 @@
+84
View File
@@ -0,0 +1,84 @@
from django.db import models
from apps.common.models import TimeStampedModel
class CreditAccount(TimeStampedModel):
team = models.OneToOneField("accounts.Team", on_delete=models.CASCADE, related_name="credit_account")
balance = models.DecimalField(max_digits=14, decimal_places=4, default=0)
reserved_balance = models.DecimalField(max_digits=14, decimal_places=4, default=0)
currency = models.CharField(max_length=16, default="CNY")
def __str__(self) -> str:
return f"{self.team} / {self.balance}"
class CreditLedger(TimeStampedModel):
class Type(models.TextChoices):
RECHARGE = "recharge", "Recharge"
RESERVE = "reserve", "Reserve"
RELEASE = "release", "Release"
CHARGE = "charge", "Charge"
ADJUSTMENT = "adjustment", "Adjustment"
REFUND = "refund", "Refund"
team = models.ForeignKey("accounts.Team", on_delete=models.CASCADE, related_name="credit_ledgers")
user = models.ForeignKey("accounts.User", on_delete=models.SET_NULL, null=True, blank=True, related_name="credit_ledgers")
project = models.ForeignKey(
"projects.Project",
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name="credit_ledgers",
)
task = models.ForeignKey("ai.AITask", on_delete=models.SET_NULL, null=True, blank=True, related_name="credit_ledgers")
ledger_type = models.CharField(max_length=32, choices=Type.choices)
amount = models.DecimalField(max_digits=14, decimal_places=4)
balance_after = models.DecimalField(max_digits=14, decimal_places=4)
reason = models.CharField(max_length=255, blank=True)
metadata = models.JSONField(default=dict, blank=True)
class Meta:
indexes = [
models.Index(fields=["team", "ledger_type"]),
models.Index(fields=["project", "task"]),
]
class CreditReservation(TimeStampedModel):
class Status(models.TextChoices):
ACTIVE = "active", "Active"
RELEASED = "released", "Released"
CHARGED = "charged", "Charged"
CANCELLED = "cancelled", "Cancelled"
team = models.ForeignKey("accounts.Team", on_delete=models.CASCADE, related_name="credit_reservations")
user = models.ForeignKey("accounts.User", on_delete=models.SET_NULL, null=True, blank=True, related_name="credit_reservations")
project = models.ForeignKey(
"projects.Project",
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name="credit_reservations",
)
task = models.OneToOneField("ai.AITask", on_delete=models.CASCADE, related_name="credit_reservation")
amount = models.DecimalField(max_digits=14, decimal_places=4)
status = models.CharField(max_length=32, choices=Status.choices, default=Status.ACTIVE)
expires_at = models.DateTimeField(null=True, blank=True)
class QuotaPolicy(TimeStampedModel):
team = models.ForeignKey("accounts.Team", on_delete=models.CASCADE, related_name="quota_policies")
user = models.ForeignKey("accounts.User", on_delete=models.CASCADE, null=True, blank=True, related_name="quota_policies")
project = models.ForeignKey(
"projects.Project",
on_delete=models.CASCADE,
null=True,
blank=True,
related_name="quota_policies",
)
monthly_limit = models.DecimalField(max_digits=14, decimal_places=4, null=True, blank=True)
project_limit = models.DecimalField(max_digits=14, decimal_places=4, null=True, blank=True)
per_task_limit = models.DecimalField(max_digits=14, decimal_places=4, null=True, blank=True)
is_active = models.BooleanField(default=True)
+43
View File
@@ -0,0 +1,43 @@
from rest_framework import serializers
from .models import CreditAccount, CreditLedger, CreditReservation, QuotaPolicy
class CreditAccountSerializer(serializers.ModelSerializer):
class Meta:
model = CreditAccount
fields = ["id", "balance", "reserved_balance", "currency", "updated_at"]
read_only_fields = fields
class CreditLedgerSerializer(serializers.ModelSerializer):
class Meta:
model = CreditLedger
fields = [
"id",
"user",
"project",
"task",
"ledger_type",
"amount",
"balance_after",
"reason",
"metadata",
"created_at",
]
read_only_fields = fields
class CreditReservationSerializer(serializers.ModelSerializer):
class Meta:
model = CreditReservation
fields = ["id", "user", "project", "task", "amount", "status", "expires_at", "created_at"]
read_only_fields = fields
class QuotaPolicySerializer(serializers.ModelSerializer):
class Meta:
model = QuotaPolicy
fields = ["id", "user", "project", "monthly_limit", "project_limit", "per_task_limit", "is_active"]
read_only_fields = ["id"]
@@ -0,0 +1 @@
@@ -0,0 +1,93 @@
from decimal import Decimal
from django.db import transaction
from apps.billing.models import CreditAccount, CreditLedger, CreditReservation
@transaction.atomic
def reserve_credit(*, team, user, task, amount: Decimal) -> CreditReservation:
account, _ = CreditAccount.objects.select_for_update().get_or_create(team=team)
available = account.balance - account.reserved_balance
if available < amount:
raise ValueError("insufficient credit")
account.reserved_balance += amount
account.save(update_fields=["reserved_balance", "updated_at"])
reservation = CreditReservation.objects.create(
team=team,
user=user,
project=task.project,
task=task,
amount=amount,
)
CreditLedger.objects.create(
team=team,
user=user,
project=task.project,
task=task,
ledger_type=CreditLedger.Type.RESERVE,
amount=amount,
balance_after=account.balance,
reason="reserve ai task credit",
)
return reservation
@transaction.atomic
def release_credit(*, reservation: CreditReservation, reason: str = "") -> None:
account = CreditAccount.objects.select_for_update().get(team=reservation.team)
if reservation.status != CreditReservation.Status.ACTIVE:
return
account.reserved_balance -= reservation.amount
account.save(update_fields=["reserved_balance", "updated_at"])
reservation.status = CreditReservation.Status.RELEASED
reservation.save(update_fields=["status", "updated_at"])
CreditLedger.objects.create(
team=reservation.team,
user=reservation.user,
project=reservation.project,
task=reservation.task,
ledger_type=CreditLedger.Type.RELEASE,
amount=reservation.amount,
balance_after=account.balance,
reason=reason or "release reserved credit",
)
@transaction.atomic
def charge_reserved_credit(*, reservation: CreditReservation, actual_amount: Decimal) -> None:
account = CreditAccount.objects.select_for_update().get(team=reservation.team)
if reservation.status != CreditReservation.Status.ACTIVE:
raise ValueError("reservation is not active")
if actual_amount > reservation.amount:
raise ValueError("actual amount exceeds reserved amount")
account.balance -= actual_amount
account.reserved_balance -= reservation.amount
account.save(update_fields=["balance", "reserved_balance", "updated_at"])
reservation.status = CreditReservation.Status.CHARGED
reservation.save(update_fields=["status", "updated_at"])
CreditLedger.objects.create(
team=reservation.team,
user=reservation.user,
project=reservation.project,
task=reservation.task,
ledger_type=CreditLedger.Type.CHARGE,
amount=actual_amount,
balance_after=account.balance,
reason="charge ai task credit",
)
if reservation.amount > actual_amount:
CreditLedger.objects.create(
team=reservation.team,
user=reservation.user,
project=reservation.project,
task=reservation.task,
ledger_type=CreditLedger.Type.RELEASE,
amount=reservation.amount - actual_amount,
balance_after=account.balance,
reason="release unused reserved credit",
)
+60
View File
@@ -0,0 +1,60 @@
from decimal import Decimal
from django.test import TestCase
from apps.accounts.models import Team, TeamMember, User
from apps.ai.models import AITask, ModelConfig, ModelProvider
from apps.billing.models import CreditAccount, CreditLedger, CreditReservation
from apps.billing.services.ledger import charge_reserved_credit, release_credit, reserve_credit
class CreditLedgerTests(TestCase):
def setUp(self):
self.user = User.objects.create_user(username="owner", password="pass")
self.team = Team.objects.create(name="Billing Team", owner=self.user)
TeamMember.objects.create(team=self.team, user=self.user, role=TeamMember.Role.OWNER)
self.account = CreditAccount.objects.create(team=self.team, balance=Decimal("100.0000"))
self.provider = ModelProvider.objects.create(name="volcengine", display_name="Volcano")
self.model = ModelConfig.objects.create(
provider=self.provider,
name="doubao-seed-2-0-pro-260215",
display_name="Doubao",
capability=ModelConfig.Capability.TEXT,
)
self.task = AITask.objects.create(
team=self.team,
created_by=self.user,
task_type=AITask.Type.SCRIPT_GENERATION,
model_config=self.model,
idempotency_key="billing-test-task",
estimated_cost=Decimal("10.0000"),
)
def test_reserve_and_charge_credit(self):
reservation = reserve_credit(team=self.team, user=self.user, task=self.task, amount=Decimal("10.0000"))
self.account.refresh_from_db()
self.assertEqual(reservation.status, CreditReservation.Status.ACTIVE)
self.assertEqual(self.account.balance, Decimal("100.0000"))
self.assertEqual(self.account.reserved_balance, Decimal("10.0000"))
charge_reserved_credit(reservation=reservation, actual_amount=Decimal("8.0000"))
self.account.refresh_from_db()
reservation.refresh_from_db()
self.assertEqual(reservation.status, CreditReservation.Status.CHARGED)
self.assertEqual(self.account.balance, Decimal("92.0000"))
self.assertEqual(self.account.reserved_balance, Decimal("0.0000"))
self.assertEqual(CreditLedger.objects.filter(team=self.team).count(), 3)
def test_release_reserved_credit(self):
reservation = reserve_credit(team=self.team, user=self.user, task=self.task, amount=Decimal("10.0000"))
release_credit(reservation=reservation, reason="model failed")
self.account.refresh_from_db()
reservation.refresh_from_db()
self.assertEqual(reservation.status, CreditReservation.Status.RELEASED)
self.assertEqual(self.account.balance, Decimal("100.0000"))
self.assertEqual(self.account.reserved_balance, Decimal("0.0000"))
self.assertEqual(CreditLedger.objects.filter(ledger_type=CreditLedger.Type.RELEASE).count(), 1)
+9
View File
@@ -0,0 +1,9 @@
from django.urls import path
from .views import ledgers, recharge, summary
urlpatterns = [
path("summary/", summary, name="billing-summary"),
path("ledgers/", ledgers, name="billing-ledgers"),
path("recharge/", recharge, name="billing-recharge"),
]
+80
View File
@@ -0,0 +1,80 @@
from decimal import Decimal, InvalidOperation
from django.db import transaction
from django.db.models import Sum
from rest_framework import status
from rest_framework.decorators import api_view, permission_classes
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response
from apps.common.api import get_current_team
from .models import CreditAccount, CreditLedger
from .serializers import CreditAccountSerializer, CreditLedgerSerializer
@api_view(["GET"])
@permission_classes([IsAuthenticated])
def summary(request):
team = get_current_team(request.user)
account, _ = CreditAccount.objects.get_or_create(team=team)
charged = CreditLedger.objects.filter(team=team, ledger_type=CreditLedger.Type.CHARGE).aggregate(
total=Sum("amount")
)["total"] or 0
return Response(
{
"account": CreditAccountSerializer(account).data,
"charged_total": charged,
}
)
@api_view(["GET"])
@permission_classes([IsAuthenticated])
def ledgers(request):
team = get_current_team(request.user)
queryset = CreditLedger.objects.filter(team=team).select_related("user", "project", "task").order_by("-created_at")
project_id = request.query_params.get("project")
user_id = request.query_params.get("user")
if project_id:
queryset = queryset.filter(project_id=project_id)
if user_id:
queryset = queryset.filter(user_id=user_id)
return Response(CreditLedgerSerializer(queryset[:100], many=True).data)
@api_view(["POST"])
@permission_classes([IsAuthenticated])
def recharge(request):
team = get_current_team(request.user)
try:
amount = Decimal(str(request.data.get("amount", "0")))
bonus = Decimal(str(request.data.get("bonus", "0")))
except (InvalidOperation, TypeError):
return Response({"detail": "invalid amount"}, status=status.HTTP_400_BAD_REQUEST)
if amount <= 0:
return Response({"detail": "amount must be positive"}, status=status.HTTP_400_BAD_REQUEST)
if bonus < 0:
return Response({"detail": "bonus cannot be negative"}, status=status.HTTP_400_BAD_REQUEST)
channel = str(request.data.get("channel") or "manual")[:32]
credited = amount + bonus
with transaction.atomic():
account, _ = CreditAccount.objects.select_for_update().get_or_create(team=team)
account.balance += credited
account.save(update_fields=["balance", "updated_at"])
ledger = CreditLedger.objects.create(
team=team,
user=request.user,
ledger_type=CreditLedger.Type.RECHARGE,
amount=credited,
balance_after=account.balance,
reason="团队充值",
metadata={"channel": channel, "paid_amount": str(amount), "bonus": str(bonus)},
)
return Response(
{
"account": CreditAccountSerializer(account).data,
"ledger": CreditLedgerSerializer(ledger).data,
},
status=status.HTTP_201_CREATED,
)