feat: add AirShelf core implementation

2026-06-05 10:21:40 +08:00
parent 2ba1058329
commit cfdcd84a30
252 changed files with 70828 additions and 0 deletions
@@ -0,0 +1,102 @@
+from django.db import models
+
+from apps.common.models import TeamOwnedModel, TimeStampedModel
+
+
+class ModelProvider(TimeStampedModel):
+    class Status(models.TextChoices):
+        ACTIVE = "active", "Active"
+        DISABLED = "disabled", "Disabled"
+
+    name = models.CharField(max_length=64, unique=True)
+    display_name = models.CharField(max_length=128)
+    status = models.CharField(max_length=24, choices=Status.choices, default=Status.ACTIVE)
+    base_url = models.URLField(blank=True)
+    metadata = models.JSONField(default=dict, blank=True)
+
+    def __str__(self) -> str:
+        return self.display_name
+
+
+class ModelConfig(TimeStampedModel):
+    class Capability(models.TextChoices):
+        TEXT = "text", "Text"
+        IMAGE = "image", "Image"
+        VIDEO = "video", "Video"
+        VISION = "vision", "Vision"
+        EXPORT = "export", "Export"
+
+    class Status(models.TextChoices):
+        ACTIVE = "active", "Active"
+        DISABLED = "disabled", "Disabled"
+
+    provider = models.ForeignKey(ModelProvider, on_delete=models.CASCADE, related_name="models")
+    name = models.CharField(max_length=128)
+    display_name = models.CharField(max_length=128)
+    capability = models.CharField(max_length=32, choices=Capability.choices)
+    endpoint = models.CharField(max_length=255, blank=True)
+    unit_price = models.DecimalField(max_digits=12, decimal_places=4, default=0)
+    status = models.CharField(max_length=24, choices=Status.choices, default=Status.ACTIVE)
+    rate_limit_per_minute = models.PositiveIntegerField(default=60)
+    metadata = models.JSONField(default=dict, blank=True)
+
+    class Meta:
+        unique_together = [("provider", "name", "capability")]
+
+    def __str__(self) -> str:
+        return f"{self.provider.name}:{self.name}:{self.capability}"
+
+
+class AITask(TeamOwnedModel):
+    class Type(models.TextChoices):
+        SCRIPT_GENERATION = "script_generation", "Script Generation"
+        SCRIPT_OPTIMIZATION = "script_optimization", "Script Optimization"
+        PRODUCT_IMAGE = "product_image", "Product Image"
+        PERSON_IMAGE = "person_image", "Person Image"
+        SCENE_IMAGE = "scene_image", "Scene Image"
+        STORYBOARD = "storyboard", "Storyboard"
+        VIDEO_SEGMENT = "video_segment", "Video Segment"
+        EXPORT = "export", "Export"
+
+    class Status(models.TextChoices):
+        CREATED = "created", "Created"
+        RESERVED = "reserved", "Reserved"
+        SUBMITTED = "submitted", "Submitted"
+        POLLING = "polling", "Polling"
+        POSTPROCESSING = "postprocessing", "Postprocessing"
+        SUCCEEDED = "succeeded", "Succeeded"
+        FAILED = "failed", "Failed"
+        COMPENSATING = "compensating", "Compensating"
+        CANCELLED = "cancelled", "Cancelled"
+
+    project = models.ForeignKey(
+        "projects.Project",
+        on_delete=models.CASCADE,
+        null=True,
+        blank=True,
+        related_name="ai_tasks",
+    )
+    task_type = models.CharField(max_length=48, choices=Type.choices)
+    status = models.CharField(max_length=32, choices=Status.choices, default=Status.CREATED)
+    model_config = models.ForeignKey(ModelConfig, on_delete=models.PROTECT, related_name="tasks")
+    idempotency_key = models.CharField(max_length=128, unique=True)
+    provider_task_id = models.CharField(max_length=255, blank=True)
+    request_payload = models.JSONField(default=dict, blank=True)
+    response_payload = models.JSONField(default=dict, blank=True)
+    estimated_cost = models.DecimalField(max_digits=12, decimal_places=4, default=0)
+    actual_cost = models.DecimalField(max_digits=12, decimal_places=4, default=0)
+    error_code = models.CharField(max_length=64, blank=True)
+    error_message = models.TextField(blank=True)
+    submitted_at = models.DateTimeField(null=True, blank=True)
+    completed_at = models.DateTimeField(null=True, blank=True)
+
+    class Meta:
+        indexes = [
+            models.Index(fields=["team", "status"]),
+            models.Index(fields=["project", "task_type"]),
+            models.Index(fields=["provider_task_id"]),
+        ]
+
+    def __str__(self) -> str:
+        return f"{self.task_type}:{self.status}:{self.id}"
+