feat: add AirShelf core implementation
This commit is contained in:
zyc
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
from django.contrib import admin
|
||||
|
||||
from .models import AITask, ModelConfig, ModelProvider
|
||||
|
||||
|
||||
@admin.register(ModelProvider)
|
||||
class ModelProviderAdmin(admin.ModelAdmin):
|
||||
list_display = ("name", "display_name", "status", "updated_at")
|
||||
search_fields = ("name", "display_name")
|
||||
list_filter = ("status",)
|
||||
|
||||
|
||||
@admin.register(ModelConfig)
|
||||
class ModelConfigAdmin(admin.ModelAdmin):
|
||||
list_display = ("provider", "name", "capability", "unit_price", "status", "rate_limit_per_minute")
|
||||
search_fields = ("provider__name", "name", "display_name")
|
||||
list_filter = ("capability", "status")
|
||||
|
||||
|
||||
@admin.register(AITask)
|
||||
class AITaskAdmin(admin.ModelAdmin):
|
||||
list_display = ("id", "team", "project", "task_type", "status", "model_config", "actual_cost", "updated_at")
|
||||
search_fields = ("idempotency_key", "provider_task_id", "project__name", "team__name")
|
||||
list_filter = ("task_type", "status", "model_config__capability")
|
||||
readonly_fields = ("request_payload", "response_payload")
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
from django.apps import AppConfig
|
||||
|
||||
|
||||
class AiConfig(AppConfig):
|
||||
default_auto_field = "django.db.models.BigAutoField"
|
||||
name = "apps.ai"
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
VOLCANO_PROVIDER = {
|
||||
"name": "volcengine",
|
||||
"display_name": "火山引擎(豆包)",
|
||||
"base_url": "https://ark.cn-beijing.volces.com/api/v3",
|
||||
}
|
||||
|
||||
VOLCANO_MODELS = [
|
||||
{
|
||||
"display_name": "Doubao-Seed-2.0-Pro",
|
||||
"name": "doubao-seed-2-0-pro-260215",
|
||||
"capability": "text",
|
||||
"endpoint": "chat/completions",
|
||||
"metadata": {"think": True, "source": "video-flow/data/vendor/volcengine.ts"},
|
||||
},
|
||||
{
|
||||
"display_name": "Doubao-Seed-2.0-Lite",
|
||||
"name": "doubao-seed-2-0-lite-260215",
|
||||
"capability": "text",
|
||||
"endpoint": "chat/completions",
|
||||
"metadata": {"think": True, "source": "video-flow/data/vendor/volcengine.ts"},
|
||||
},
|
||||
{
|
||||
"display_name": "Seedream-5.0",
|
||||
"name": "doubao-seedream-5-0-260128",
|
||||
"capability": "image",
|
||||
"endpoint": "images/generations",
|
||||
"metadata": {
|
||||
"modes": ["text", "singleImage", "multiReference"],
|
||||
"watermark": False,
|
||||
"source": "video-flow/data/vendor/volcengine.ts",
|
||||
},
|
||||
},
|
||||
{
|
||||
"display_name": "Seedream-4.5",
|
||||
"name": "doubao-seedream-4-5-251128",
|
||||
"capability": "image",
|
||||
"endpoint": "images/generations",
|
||||
"metadata": {
|
||||
"modes": ["text", "singleImage", "multiReference"],
|
||||
"watermark": False,
|
||||
"source": "video-flow/data/vendor/volcengine.ts",
|
||||
},
|
||||
},
|
||||
{
|
||||
"display_name": "Seedance-2.0",
|
||||
"name": "doubao-seedance-2-0-260128",
|
||||
"capability": "video",
|
||||
"endpoint": "contents/generations/tasks",
|
||||
"metadata": {
|
||||
"audio": "optional",
|
||||
"modes": ["text", "startFrameOptional", "imageReference:9", "videoReference:3", "audioReference:3"],
|
||||
"durations": list(range(4, 16)),
|
||||
"resolutions": ["480p", "720p"],
|
||||
"watermark": False,
|
||||
"source": "video-flow/data/vendor/volcengine.ts",
|
||||
},
|
||||
},
|
||||
{
|
||||
"display_name": "Seedance-1.5-Pro",
|
||||
"name": "doubao-seedance-1-5-pro-251215",
|
||||
"capability": "video",
|
||||
"endpoint": "contents/generations/tasks",
|
||||
"metadata": {
|
||||
"audio": "optional",
|
||||
"modes": ["text", "startFrameOptional"],
|
||||
"durations": list(range(4, 13)),
|
||||
"resolutions": ["480p", "720p", "1080p"],
|
||||
"watermark": False,
|
||||
"source": "video-flow/data/vendor/volcengine.ts",
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
# 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 = [
|
||||
migrations.swappable_dependency(settings.AUTH_USER_MODEL),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.CreateModel(
|
||||
name="ModelConfig",
|
||||
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)),
|
||||
("name", models.CharField(max_length=128)),
|
||||
("display_name", models.CharField(max_length=128)),
|
||||
(
|
||||
"capability",
|
||||
models.CharField(
|
||||
choices=[
|
||||
("text", "Text"),
|
||||
("image", "Image"),
|
||||
("video", "Video"),
|
||||
("vision", "Vision"),
|
||||
("export", "Export"),
|
||||
],
|
||||
max_length=32,
|
||||
),
|
||||
),
|
||||
("endpoint", models.CharField(blank=True, max_length=255)),
|
||||
(
|
||||
"unit_price",
|
||||
models.DecimalField(decimal_places=4, default=0, max_digits=12),
|
||||
),
|
||||
(
|
||||
"status",
|
||||
models.CharField(
|
||||
choices=[("active", "Active"), ("disabled", "Disabled")],
|
||||
default="active",
|
||||
max_length=24,
|
||||
),
|
||||
),
|
||||
("rate_limit_per_minute", models.PositiveIntegerField(default=60)),
|
||||
("metadata", models.JSONField(blank=True, default=dict)),
|
||||
],
|
||||
),
|
||||
migrations.CreateModel(
|
||||
name="ModelProvider",
|
||||
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)),
|
||||
("name", models.CharField(max_length=64, unique=True)),
|
||||
("display_name", models.CharField(max_length=128)),
|
||||
(
|
||||
"status",
|
||||
models.CharField(
|
||||
choices=[("active", "Active"), ("disabled", "Disabled")],
|
||||
default="active",
|
||||
max_length=24,
|
||||
),
|
||||
),
|
||||
("base_url", models.URLField(blank=True)),
|
||||
("metadata", models.JSONField(blank=True, default=dict)),
|
||||
],
|
||||
options={
|
||||
"abstract": False,
|
||||
},
|
||||
),
|
||||
migrations.CreateModel(
|
||||
name="AITask",
|
||||
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)),
|
||||
(
|
||||
"task_type",
|
||||
models.CharField(
|
||||
choices=[
|
||||
("script_generation", "Script Generation"),
|
||||
("script_optimization", "Script Optimization"),
|
||||
("product_image", "Product Image"),
|
||||
("person_image", "Person Image"),
|
||||
("scene_image", "Scene Image"),
|
||||
("storyboard", "Storyboard"),
|
||||
("video_segment", "Video Segment"),
|
||||
("export", "Export"),
|
||||
],
|
||||
max_length=48,
|
||||
),
|
||||
),
|
||||
(
|
||||
"status",
|
||||
models.CharField(
|
||||
choices=[
|
||||
("created", "Created"),
|
||||
("reserved", "Reserved"),
|
||||
("submitted", "Submitted"),
|
||||
("polling", "Polling"),
|
||||
("postprocessing", "Postprocessing"),
|
||||
("succeeded", "Succeeded"),
|
||||
("failed", "Failed"),
|
||||
("compensating", "Compensating"),
|
||||
("cancelled", "Cancelled"),
|
||||
],
|
||||
default="created",
|
||||
max_length=32,
|
||||
),
|
||||
),
|
||||
("idempotency_key", models.CharField(max_length=128, unique=True)),
|
||||
("provider_task_id", models.CharField(blank=True, max_length=255)),
|
||||
("request_payload", models.JSONField(blank=True, default=dict)),
|
||||
("response_payload", models.JSONField(blank=True, default=dict)),
|
||||
(
|
||||
"estimated_cost",
|
||||
models.DecimalField(decimal_places=4, default=0, max_digits=12),
|
||||
),
|
||||
(
|
||||
"actual_cost",
|
||||
models.DecimalField(decimal_places=4, default=0, max_digits=12),
|
||||
),
|
||||
("error_code", models.CharField(blank=True, max_length=64)),
|
||||
("error_message", models.TextField(blank=True)),
|
||||
("submitted_at", models.DateTimeField(blank=True, null=True)),
|
||||
("completed_at", models.DateTimeField(blank=True, null=True)),
|
||||
(
|
||||
"created_by",
|
||||
models.ForeignKey(
|
||||
blank=True,
|
||||
null=True,
|
||||
on_delete=django.db.models.deletion.SET_NULL,
|
||||
related_name="created_%(class)s_set",
|
||||
to=settings.AUTH_USER_MODEL,
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,78 @@
|
||||
# Generated by Django 5.1.15 on 2026-05-29 03:59
|
||||
|
||||
import django.db.models.deletion
|
||||
from django.db import migrations, models
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
initial = True
|
||||
|
||||
dependencies = [
|
||||
("accounts", "0001_initial"),
|
||||
("ai", "0001_initial"),
|
||||
("projects", "0001_initial"),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.AddField(
|
||||
model_name="aitask",
|
||||
name="project",
|
||||
field=models.ForeignKey(
|
||||
blank=True,
|
||||
null=True,
|
||||
on_delete=django.db.models.deletion.CASCADE,
|
||||
related_name="ai_tasks",
|
||||
to="projects.project",
|
||||
),
|
||||
),
|
||||
migrations.AddField(
|
||||
model_name="aitask",
|
||||
name="team",
|
||||
field=models.ForeignKey(
|
||||
on_delete=django.db.models.deletion.CASCADE,
|
||||
related_name="%(class)s_set",
|
||||
to="accounts.team",
|
||||
),
|
||||
),
|
||||
migrations.AddField(
|
||||
model_name="aitask",
|
||||
name="model_config",
|
||||
field=models.ForeignKey(
|
||||
on_delete=django.db.models.deletion.PROTECT,
|
||||
related_name="tasks",
|
||||
to="ai.modelconfig",
|
||||
),
|
||||
),
|
||||
migrations.AddField(
|
||||
model_name="modelconfig",
|
||||
name="provider",
|
||||
field=models.ForeignKey(
|
||||
on_delete=django.db.models.deletion.CASCADE,
|
||||
related_name="models",
|
||||
to="ai.modelprovider",
|
||||
),
|
||||
),
|
||||
migrations.AddIndex(
|
||||
model_name="aitask",
|
||||
index=models.Index(
|
||||
fields=["team", "status"], name="ai_aitask_team_id_710ece_idx"
|
||||
),
|
||||
),
|
||||
migrations.AddIndex(
|
||||
model_name="aitask",
|
||||
index=models.Index(
|
||||
fields=["project", "task_type"], name="ai_aitask_project_f2850d_idx"
|
||||
),
|
||||
),
|
||||
migrations.AddIndex(
|
||||
model_name="aitask",
|
||||
index=models.Index(
|
||||
fields=["provider_task_id"], name="ai_aitask_provide_67beef_idx"
|
||||
),
|
||||
),
|
||||
migrations.AlterUniqueTogether(
|
||||
name="modelconfig",
|
||||
unique_together={("provider", "name", "capability")},
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -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}"
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
from .base import AIProvider, AIProviderResult
|
||||
from .volcano import VolcanoArkProvider
|
||||
|
||||
|
||||
__all__ = ["AIProvider", "AIProviderResult", "VolcanoArkProvider"]
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AIProviderResult:
|
||||
provider_task_id: str = ""
|
||||
status: str = "succeeded"
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
asset_urls: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
class AIProvider(Protocol):
|
||||
def submit(self, payload: dict[str, Any]) -> AIProviderResult:
|
||||
...
|
||||
|
||||
def poll(self, provider_task_id: str) -> AIProviderResult:
|
||||
...
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
from dataclasses import dataclass
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from django.conf import settings
|
||||
|
||||
from .base import AIProviderResult
|
||||
|
||||
|
||||
@dataclass
|
||||
class VolcanoArkProvider:
|
||||
api_key: str | None = None
|
||||
base_url: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.api_key = self.api_key or settings.VOLCANO.get("ark_api_key")
|
||||
self.base_url = self.base_url or settings.VOLCANO.get("ark_base_url")
|
||||
|
||||
def submit(self, payload: dict[str, Any]) -> AIProviderResult:
|
||||
# The exact endpoint is resolved by ModelConfig; this adapter keeps IO centralized.
|
||||
endpoint = payload.get("endpoint")
|
||||
if not endpoint:
|
||||
raise ValueError("Volcano request payload requires endpoint")
|
||||
|
||||
response = requests.post(
|
||||
f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
json=payload.get("body", {}),
|
||||
timeout=60,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return AIProviderResult(
|
||||
provider_task_id=str(data.get("id") or data.get("task_id") or ""),
|
||||
status=str(data.get("status") or "submitted"),
|
||||
payload=data,
|
||||
)
|
||||
|
||||
def chat_completion(self, *, model: str, messages: list[dict[str, str]], endpoint: str = "chat/completions") -> dict[str, Any]:
|
||||
if not self.api_key:
|
||||
raise ValueError("VOLCANO_ARK_API_KEY is not configured")
|
||||
|
||||
response = requests.post(
|
||||
f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
|
||||
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
|
||||
json={"model": model, "messages": messages},
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
@staticmethod
|
||||
def extract_text(data: dict[str, Any]) -> str:
|
||||
choices = data.get("choices") or []
|
||||
if choices:
|
||||
message = choices[0].get("message") or {}
|
||||
content = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
return "\n".join(str(item.get("text", "")) for item in content if isinstance(item, dict))
|
||||
output = data.get("output")
|
||||
if isinstance(output, str):
|
||||
return output
|
||||
raise ValueError("Volcano response does not contain text content")
|
||||
|
||||
def poll(self, provider_task_id: str) -> AIProviderResult:
|
||||
if not provider_task_id:
|
||||
raise ValueError("provider_task_id is required")
|
||||
|
||||
return AIProviderResult(provider_task_id=provider_task_id, status="polling", payload={})
|
||||
|
||||
def image_generation(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
endpoint: str = "images/generations",
|
||||
image: str | list[str] | None = None,
|
||||
size: str = "2K",
|
||||
) -> dict[str, Any]:
|
||||
if not self.api_key:
|
||||
raise ValueError("VOLCANO_ARK_API_KEY is not configured")
|
||||
body: dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"response_format": "url",
|
||||
"watermark": False,
|
||||
"size": size,
|
||||
"sequential_image_generation": "disabled",
|
||||
}
|
||||
if image:
|
||||
body["image"] = image
|
||||
response = requests.post(
|
||||
f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
|
||||
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
|
||||
json=body,
|
||||
timeout=180,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
def create_video_task(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
endpoint: str,
|
||||
prompt: str,
|
||||
ratio: str = "9:16",
|
||||
duration: int = 15,
|
||||
resolution: str = "720p",
|
||||
reference_images: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if not self.api_key:
|
||||
raise ValueError("VOLCANO_ARK_API_KEY is not configured")
|
||||
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
||||
for image_url in reference_images or []:
|
||||
content.append({"type": "image_url", "image_url": {"url": image_url}, "role": "reference_image"})
|
||||
body = {
|
||||
"model": model,
|
||||
"content": content,
|
||||
"ratio": ratio,
|
||||
"duration": duration,
|
||||
"resolution": resolution,
|
||||
"watermark": False,
|
||||
"generate_audio": False,
|
||||
}
|
||||
response = requests.post(
|
||||
f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
|
||||
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
|
||||
json=body,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
def poll_video_task(self, *, endpoint: str, provider_task_id: str) -> dict[str, Any]:
|
||||
if not self.api_key:
|
||||
raise ValueError("VOLCANO_ARK_API_KEY is not configured")
|
||||
response = requests.get(
|
||||
f"{self.base_url.rstrip('/')}/{endpoint.rstrip('/')}/{provider_task_id}",
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
timeout=60,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
@staticmethod
|
||||
def extract_first_media_url(data: dict[str, Any]) -> str:
|
||||
items = data.get("data") or []
|
||||
for item in items:
|
||||
if item.get("url"):
|
||||
return item["url"]
|
||||
if item.get("b64_json"):
|
||||
return item["b64_json"]
|
||||
content = data.get("content") or {}
|
||||
if content.get("video_url"):
|
||||
return content["video_url"]
|
||||
raise ValueError("Volcano response does not contain media url")
|
||||
|
||||
@staticmethod
|
||||
def media_to_bytes(media: str) -> tuple[BytesIO, str]:
|
||||
if media.startswith("http://") or media.startswith("https://"):
|
||||
response = requests.get(media, timeout=180)
|
||||
response.raise_for_status()
|
||||
return BytesIO(response.content), response.headers.get("content-type", "application/octet-stream")
|
||||
if "," in media and media.startswith("data:"):
|
||||
header, raw = media.split(",", 1)
|
||||
content_type = header.split(";")[0].replace("data:", "") or "application/octet-stream"
|
||||
return BytesIO(base64.b64decode(raw)), content_type
|
||||
return BytesIO(base64.b64decode(media)), "image/png"
|
||||
@@ -0,0 +1,44 @@
|
||||
from rest_framework import serializers
|
||||
|
||||
from .models import AITask, ModelConfig, ModelProvider
|
||||
|
||||
|
||||
class ModelProviderSerializer(serializers.ModelSerializer):
|
||||
class Meta:
|
||||
model = ModelProvider
|
||||
fields = ["id", "name", "display_name", "status", "base_url", "metadata"]
|
||||
read_only_fields = fields
|
||||
|
||||
|
||||
class ModelConfigSerializer(serializers.ModelSerializer):
|
||||
provider = ModelProviderSerializer(read_only=True)
|
||||
|
||||
class Meta:
|
||||
model = ModelConfig
|
||||
fields = ["id", "provider", "name", "display_name", "capability", "endpoint", "unit_price", "status", "metadata"]
|
||||
read_only_fields = fields
|
||||
|
||||
|
||||
class AITaskSerializer(serializers.ModelSerializer):
|
||||
model_config = ModelConfigSerializer(read_only=True)
|
||||
|
||||
class Meta:
|
||||
model = AITask
|
||||
fields = [
|
||||
"id",
|
||||
"project",
|
||||
"task_type",
|
||||
"status",
|
||||
"model_config",
|
||||
"provider_task_id",
|
||||
"estimated_cost",
|
||||
"actual_cost",
|
||||
"error_code",
|
||||
"error_message",
|
||||
"submitted_at",
|
||||
"completed_at",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
]
|
||||
read_only_fields = fields
|
||||
|
||||
@@ -0,0 +1,421 @@
|
||||
import uuid
|
||||
from decimal import Decimal
|
||||
from django.db import transaction
|
||||
from django.utils import timezone
|
||||
|
||||
from apps.ai.models import AITask, ModelConfig
|
||||
from apps.ai.providers import VolcanoArkProvider
|
||||
from apps.assets.models import Asset, AssetFile
|
||||
from apps.assets.storage import TosStorage
|
||||
from apps.billing.services.ledger import charge_reserved_credit, release_credit, reserve_credit
|
||||
from apps.projects.models import (
|
||||
BaseAssetGroup,
|
||||
ExportJob,
|
||||
ProjectStage,
|
||||
ScriptSegment,
|
||||
ScriptVersion,
|
||||
StoryboardFrame,
|
||||
StoryboardVersion,
|
||||
VideoSegment,
|
||||
VideoSegmentVersion,
|
||||
)
|
||||
|
||||
|
||||
def get_default_model(capability: str) -> ModelConfig:
|
||||
return (
|
||||
ModelConfig.objects.select_related("provider")
|
||||
.filter(capability=capability, status=ModelConfig.Status.ACTIVE, provider__status="active")
|
||||
.order_by("created_at")
|
||||
.first()
|
||||
)
|
||||
|
||||
|
||||
def estimate_cost(model_config: ModelConfig) -> Decimal:
|
||||
return model_config.unit_price if model_config.unit_price > 0 else Decimal("1.0000")
|
||||
|
||||
|
||||
def build_script_prompt(*, project, user_prompt: str, selling_point_ids: list[str] | None = None) -> list[dict[str, str]]:
|
||||
product = project.product
|
||||
selling_points = product.selling_points.all()
|
||||
if selling_point_ids:
|
||||
selling_points = selling_points.filter(id__in=selling_point_ids)
|
||||
selling_text = "\n".join(f"- {item.title}: {item.detail}" for item in selling_points)
|
||||
system = (
|
||||
"你是电商短视频脚本导演。请为 9:16 竖屏带货短视频生成 60 秒脚本，"
|
||||
"拆成 4 个 15 秒段落。每段包含旁白、画面描述、商品露出方式和转场建议。"
|
||||
)
|
||||
user = f"""
|
||||
商品标题：{product.title}
|
||||
品牌：{product.brand or "未填写"}
|
||||
类目：{product.category or "未填写"}
|
||||
目标人群：{product.target_audience or "未填写"}
|
||||
商品描述：{product.description or "未填写"}
|
||||
卖点：
|
||||
{selling_text or "未选择卖点，请根据商品信息自行提炼。"}
|
||||
|
||||
用户补充需求：
|
||||
{user_prompt or "生成一条结构完整、节奏清晰、适合投放的带货短视频脚本。"}
|
||||
""".strip()
|
||||
return [{"role": "system", "content": system}, {"role": "user", "content": user}]
|
||||
|
||||
|
||||
def split_script_into_segments(content: str) -> list[str]:
|
||||
blocks = [line.strip() for line in content.splitlines() if line.strip()]
|
||||
if len(blocks) >= 4:
|
||||
return blocks[:4]
|
||||
if not content.strip():
|
||||
return [""] * 4
|
||||
return [content.strip()] + [""] * (4 - len(blocks or [content]))
|
||||
|
||||
|
||||
@transaction.atomic
|
||||
def create_ai_task(*, project, user, task_type: str, model_config: ModelConfig, request_payload: dict) -> AITask:
|
||||
cost = estimate_cost(model_config)
|
||||
task = AITask.objects.create(
|
||||
team=project.team,
|
||||
created_by=user,
|
||||
project=project,
|
||||
task_type=task_type,
|
||||
status=AITask.Status.CREATED,
|
||||
model_config=model_config,
|
||||
idempotency_key=f"{task_type}:{project.id}:{uuid.uuid4()}",
|
||||
request_payload=request_payload,
|
||||
estimated_cost=cost,
|
||||
)
|
||||
reserve_credit(team=project.team, user=user, task=task, amount=cost)
|
||||
task.status = AITask.Status.RESERVED
|
||||
task.save(update_fields=["status", "updated_at"])
|
||||
return task
|
||||
|
||||
|
||||
def generate_project_script(*, project, user, user_prompt: str, selling_point_ids: list[str] | None = None) -> ScriptVersion:
|
||||
model_config = get_default_model(ModelConfig.Capability.TEXT)
|
||||
if model_config is None:
|
||||
raise ValueError("no active text model configured")
|
||||
|
||||
messages = build_script_prompt(project=project, user_prompt=user_prompt, selling_point_ids=selling_point_ids)
|
||||
payload = {"model": model_config.name, "endpoint": model_config.endpoint, "messages": messages}
|
||||
task = create_ai_task(
|
||||
project=project,
|
||||
user=user,
|
||||
task_type=AITask.Type.SCRIPT_GENERATION,
|
||||
model_config=model_config,
|
||||
request_payload=payload,
|
||||
)
|
||||
reservation = task.credit_reservation
|
||||
|
||||
try:
|
||||
task.status = AITask.Status.SUBMITTED
|
||||
task.submitted_at = timezone.now()
|
||||
task.save(update_fields=["status", "submitted_at", "updated_at"])
|
||||
|
||||
provider = VolcanoArkProvider(base_url=model_config.provider.base_url or None)
|
||||
response = provider.chat_completion(model=model_config.name, endpoint=model_config.endpoint, messages=messages)
|
||||
content = provider.extract_text(response)
|
||||
|
||||
with transaction.atomic():
|
||||
task.status = AITask.Status.SUCCEEDED
|
||||
task.response_payload = response
|
||||
task.actual_cost = task.estimated_cost
|
||||
task.completed_at = timezone.now()
|
||||
task.save(update_fields=["status", "response_payload", "actual_cost", "completed_at", "updated_at"])
|
||||
charge_reserved_credit(reservation=reservation, actual_amount=task.actual_cost)
|
||||
|
||||
script = ScriptVersion.objects.create(
|
||||
project=project,
|
||||
task=task,
|
||||
title="AI 脚本",
|
||||
content=content,
|
||||
source="ai",
|
||||
is_adopted=False,
|
||||
)
|
||||
for index, segment_text in enumerate(split_script_into_segments(content)):
|
||||
ScriptSegment.objects.create(
|
||||
script_version=script,
|
||||
sort_order=index,
|
||||
duration_seconds=15,
|
||||
narration=segment_text,
|
||||
visual_prompt=segment_text,
|
||||
)
|
||||
|
||||
stage, _ = ProjectStage.objects.get_or_create(project=project, stage=ProjectStage.Stage.SCRIPT)
|
||||
stage.status = ProjectStage.Status.NEEDS_REVIEW
|
||||
stage.save(update_fields=["status", "updated_at"])
|
||||
return script
|
||||
except Exception as exc:
|
||||
with transaction.atomic():
|
||||
task.status = AITask.Status.FAILED
|
||||
task.error_message = str(exc)
|
||||
task.completed_at = timezone.now()
|
||||
task.save(update_fields=["status", "error_message", "completed_at", "updated_at"])
|
||||
release_credit(reservation=reservation, reason=str(exc))
|
||||
raise
|
||||
|
||||
|
||||
def _store_generated_media(*, team, user, project, task, media: str, name: str, category: str, asset_type: str) -> Asset:
|
||||
fileobj, content_type = VolcanoArkProvider.media_to_bytes(media)
|
||||
suffix = ".png"
|
||||
if "video" in content_type:
|
||||
suffix = ".mp4"
|
||||
elif "jpeg" in content_type:
|
||||
suffix = ".jpg"
|
||||
elif "webp" in content_type:
|
||||
suffix = ".webp"
|
||||
asset_id = uuid.uuid4()
|
||||
object_key = f"teams/{team.id}/projects/{project.id}/generated/{asset_id}{suffix}"
|
||||
stored = TosStorage().upload_fileobj(fileobj=fileobj, object_key=object_key, content_type=content_type)
|
||||
asset = Asset.objects.create(
|
||||
id=asset_id,
|
||||
team=team,
|
||||
created_by=user,
|
||||
name=name,
|
||||
asset_type=asset_type,
|
||||
source=Asset.Source.AI_GENERATED,
|
||||
category=category,
|
||||
origin_task=task,
|
||||
)
|
||||
AssetFile.objects.create(
|
||||
asset=asset,
|
||||
object_key=stored.object_key,
|
||||
bucket=stored.bucket,
|
||||
content_type=stored.content_type,
|
||||
size_bytes=stored.size_bytes,
|
||||
is_primary=True,
|
||||
)
|
||||
return asset
|
||||
|
||||
|
||||
def generate_base_asset(*, project, user, kind: str, prompt: str) -> BaseAssetGroup:
|
||||
model_config = get_default_model(ModelConfig.Capability.IMAGE)
|
||||
if model_config is None:
|
||||
raise ValueError("no active image model configured")
|
||||
payload = {"model": model_config.name, "endpoint": model_config.endpoint, "prompt": prompt, "kind": kind}
|
||||
task = create_ai_task(
|
||||
project=project,
|
||||
user=user,
|
||||
task_type={
|
||||
BaseAssetGroup.Kind.PRODUCT: AITask.Type.PRODUCT_IMAGE,
|
||||
BaseAssetGroup.Kind.PERSON: AITask.Type.PERSON_IMAGE,
|
||||
BaseAssetGroup.Kind.SCENE: AITask.Type.SCENE_IMAGE,
|
||||
}[kind],
|
||||
model_config=model_config,
|
||||
request_payload=payload,
|
||||
)
|
||||
reservation = task.credit_reservation
|
||||
try:
|
||||
provider = VolcanoArkProvider(base_url=model_config.provider.base_url or None)
|
||||
response = provider.image_generation(model=model_config.name, endpoint=model_config.endpoint, prompt=prompt)
|
||||
media = provider.extract_first_media_url(response)
|
||||
with transaction.atomic():
|
||||
task.status = AITask.Status.SUCCEEDED
|
||||
task.response_payload = response
|
||||
task.actual_cost = task.estimated_cost
|
||||
task.completed_at = timezone.now()
|
||||
task.save(update_fields=["status", "response_payload", "actual_cost", "completed_at", "updated_at"])
|
||||
charge_reserved_credit(reservation=reservation, actual_amount=task.actual_cost)
|
||||
category = {
|
||||
BaseAssetGroup.Kind.PRODUCT: Asset.Category.PRODUCT_IMAGE,
|
||||
BaseAssetGroup.Kind.PERSON: Asset.Category.PERSON,
|
||||
BaseAssetGroup.Kind.SCENE: Asset.Category.SCENE,
|
||||
}[kind]
|
||||
asset = _store_generated_media(
|
||||
team=project.team,
|
||||
user=user,
|
||||
project=project,
|
||||
task=task,
|
||||
media=media,
|
||||
name=f"{project.name}-{kind}",
|
||||
category=category,
|
||||
asset_type=Asset.Type.IMAGE,
|
||||
)
|
||||
group = BaseAssetGroup.objects.create(project=project, kind=kind, task=task, prompt=prompt)
|
||||
group.candidate_assets.add(asset)
|
||||
group.adopted_asset = asset
|
||||
group.save(update_fields=["adopted_asset", "updated_at"])
|
||||
return group
|
||||
except Exception as exc:
|
||||
task.status = AITask.Status.FAILED
|
||||
task.error_message = str(exc)
|
||||
task.completed_at = timezone.now()
|
||||
task.save(update_fields=["status", "error_message", "completed_at", "updated_at"])
|
||||
release_credit(reservation=reservation, reason=str(exc))
|
||||
raise
|
||||
|
||||
|
||||
def generate_storyboard(*, project, user, prompt: str = "") -> StoryboardVersion:
|
||||
adopted_script = project.script_versions.filter(is_adopted=True).prefetch_related("segments").first()
|
||||
if adopted_script is None:
|
||||
raise ValueError("script must be adopted before generating storyboard")
|
||||
model_config = get_default_model(ModelConfig.Capability.IMAGE)
|
||||
if model_config is None:
|
||||
raise ValueError("no active image model configured")
|
||||
|
||||
storyboard = StoryboardVersion.objects.create(project=project, prompt=prompt)
|
||||
provider = VolcanoArkProvider(base_url=model_config.provider.base_url or None)
|
||||
for segment in adopted_script.segments.all():
|
||||
task = create_ai_task(
|
||||
project=project,
|
||||
user=user,
|
||||
task_type=AITask.Type.STORYBOARD,
|
||||
model_config=model_config,
|
||||
request_payload={"model": model_config.name, "endpoint": model_config.endpoint, "prompt": segment.visual_prompt},
|
||||
)
|
||||
reservation = task.credit_reservation
|
||||
try:
|
||||
response = provider.image_generation(
|
||||
model=model_config.name,
|
||||
endpoint=model_config.endpoint,
|
||||
prompt=f"{prompt}\n{segment.visual_prompt}".strip(),
|
||||
)
|
||||
media = provider.extract_first_media_url(response)
|
||||
task.status = AITask.Status.SUCCEEDED
|
||||
task.response_payload = response
|
||||
task.actual_cost = task.estimated_cost
|
||||
task.completed_at = timezone.now()
|
||||
task.save(update_fields=["status", "response_payload", "actual_cost", "completed_at", "updated_at"])
|
||||
charge_reserved_credit(reservation=reservation, actual_amount=task.actual_cost)
|
||||
asset = _store_generated_media(
|
||||
team=project.team,
|
||||
user=user,
|
||||
project=project,
|
||||
task=task,
|
||||
media=media,
|
||||
name=f"{project.name}-storyboard-{segment.sort_order + 1}",
|
||||
category=Asset.Category.SCENE,
|
||||
asset_type=Asset.Type.IMAGE,
|
||||
)
|
||||
StoryboardFrame.objects.create(
|
||||
storyboard=storyboard,
|
||||
script_segment=segment,
|
||||
asset=asset,
|
||||
sort_order=segment.sort_order,
|
||||
prompt=segment.visual_prompt,
|
||||
)
|
||||
except Exception as exc:
|
||||
task.status = AITask.Status.FAILED
|
||||
task.error_message = str(exc)
|
||||
task.completed_at = timezone.now()
|
||||
task.save(update_fields=["status", "error_message", "completed_at", "updated_at"])
|
||||
release_credit(reservation=reservation, reason=str(exc))
|
||||
raise
|
||||
storyboard.is_adopted = True
|
||||
storyboard.save(update_fields=["is_adopted", "updated_at"])
|
||||
return storyboard
|
||||
|
||||
|
||||
def submit_video_segment(*, video_segment: VideoSegment, user, prompt: str) -> VideoSegmentVersion | None:
|
||||
model_config = get_default_model(ModelConfig.Capability.VIDEO)
|
||||
if model_config is None:
|
||||
raise ValueError("no active video model configured")
|
||||
project = video_segment.project
|
||||
task = create_ai_task(
|
||||
project=project,
|
||||
user=user,
|
||||
task_type=AITask.Type.VIDEO_SEGMENT,
|
||||
model_config=model_config,
|
||||
request_payload={
|
||||
"model": model_config.name,
|
||||
"endpoint": model_config.endpoint,
|
||||
"prompt": prompt,
|
||||
"duration": video_segment.target_duration_seconds,
|
||||
"ratio": "9:16",
|
||||
"video_segment_id": str(video_segment.id),
|
||||
},
|
||||
)
|
||||
try:
|
||||
provider = VolcanoArkProvider(base_url=model_config.provider.base_url or None)
|
||||
response = provider.create_video_task(
|
||||
model=model_config.name,
|
||||
endpoint=model_config.endpoint,
|
||||
prompt=prompt,
|
||||
duration=video_segment.target_duration_seconds,
|
||||
ratio="9:16",
|
||||
resolution="720p",
|
||||
)
|
||||
task.provider_task_id = str(response.get("id") or response.get("task_id") or "")
|
||||
task.response_payload = response
|
||||
task.status = AITask.Status.SUBMITTED
|
||||
task.submitted_at = timezone.now()
|
||||
task.save(update_fields=["provider_task_id", "response_payload", "status", "submitted_at", "updated_at"])
|
||||
video_segment.status = VideoSegment.Status.RUNNING
|
||||
video_segment.save(update_fields=["status", "updated_at"])
|
||||
return None
|
||||
except Exception as exc:
|
||||
task.status = AITask.Status.FAILED
|
||||
task.error_message = str(exc)
|
||||
task.completed_at = timezone.now()
|
||||
task.save(update_fields=["status", "error_message", "completed_at", "updated_at"])
|
||||
release_credit(reservation=task.credit_reservation, reason=str(exc))
|
||||
video_segment.status = VideoSegment.Status.FAILED
|
||||
video_segment.error_message = str(exc)
|
||||
video_segment.save(update_fields=["status", "error_message", "updated_at"])
|
||||
raise
|
||||
|
||||
|
||||
def poll_video_segment(*, video_segment: VideoSegment, user) -> VideoSegmentVersion | None:
|
||||
task = video_segment.versions.order_by("-created_at").first()
|
||||
ai_task = None
|
||||
if task:
|
||||
ai_task = task.task
|
||||
if ai_task is None:
|
||||
ai_task = video_segment.project.ai_tasks.filter(
|
||||
task_type=AITask.Type.VIDEO_SEGMENT,
|
||||
request_payload__video_segment_id=str(video_segment.id),
|
||||
status__in=[AITask.Status.SUBMITTED, AITask.Status.POLLING],
|
||||
).order_by("-created_at").first()
|
||||
if ai_task is None:
|
||||
raise ValueError("no active video generation task")
|
||||
|
||||
provider = VolcanoArkProvider(base_url=ai_task.model_config.provider.base_url or None)
|
||||
response = provider.poll_video_task(endpoint=ai_task.model_config.endpoint, provider_task_id=ai_task.provider_task_id)
|
||||
remote_status = response.get("status")
|
||||
if remote_status in {"queued", "running", "processing"}:
|
||||
ai_task.status = AITask.Status.POLLING
|
||||
ai_task.response_payload = response
|
||||
ai_task.save(update_fields=["status", "response_payload", "updated_at"])
|
||||
return None
|
||||
if remote_status in {"failed", "expired", "cancelled"}:
|
||||
ai_task.status = AITask.Status.FAILED
|
||||
ai_task.response_payload = response
|
||||
ai_task.error_message = response.get("error", {}).get("message", "video generation failed")
|
||||
ai_task.completed_at = timezone.now()
|
||||
ai_task.save(update_fields=["status", "response_payload", "error_message", "completed_at", "updated_at"])
|
||||
release_credit(reservation=ai_task.credit_reservation, reason=ai_task.error_message)
|
||||
video_segment.status = VideoSegment.Status.FAILED
|
||||
video_segment.error_message = ai_task.error_message
|
||||
video_segment.save(update_fields=["status", "error_message", "updated_at"])
|
||||
return None
|
||||
|
||||
media = provider.extract_first_media_url(response)
|
||||
asset = _store_generated_media(
|
||||
team=video_segment.project.team,
|
||||
user=user,
|
||||
project=video_segment.project,
|
||||
task=ai_task,
|
||||
media=media,
|
||||
name=f"{video_segment.project.name}-segment-{video_segment.sort_order + 1}",
|
||||
category=Asset.Category.VIDEO_CLIP,
|
||||
asset_type=Asset.Type.VIDEO,
|
||||
)
|
||||
ai_task.status = AITask.Status.SUCCEEDED
|
||||
ai_task.response_payload = response
|
||||
ai_task.actual_cost = ai_task.estimated_cost
|
||||
ai_task.completed_at = timezone.now()
|
||||
ai_task.save(update_fields=["status", "response_payload", "actual_cost", "completed_at", "updated_at"])
|
||||
charge_reserved_credit(reservation=ai_task.credit_reservation, actual_amount=ai_task.actual_cost)
|
||||
version = VideoSegmentVersion.objects.create(
|
||||
video_segment=video_segment,
|
||||
task=ai_task,
|
||||
asset=asset,
|
||||
prompt=ai_task.request_payload.get("prompt", ""),
|
||||
is_adopted=True,
|
||||
)
|
||||
video_segment.adopted_version = version
|
||||
video_segment.status = VideoSegment.Status.SUCCEEDED
|
||||
video_segment.error_message = ""
|
||||
video_segment.save(update_fields=["adopted_version", "status", "error_message", "updated_at"])
|
||||
return version
|
||||
|
||||
|
||||
def create_export_job(*, timeline, user) -> ExportJob:
|
||||
return ExportJob.objects.create(timeline=timeline, status=ExportJob.Status.QUEUED)
|
||||
@@ -0,0 +1,12 @@
|
||||
from airshelf.celery import app
|
||||
|
||||
|
||||
@app.task(bind=True, max_retries=3)
|
||||
def submit_ai_task(self, task_id: str) -> str:
|
||||
return task_id
|
||||
|
||||
|
||||
@app.task(bind=True, max_retries=5)
|
||||
def poll_ai_task(self, task_id: str) -> str:
|
||||
return task_id
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
from rest_framework.routers import DefaultRouter
|
||||
|
||||
from .views import AITaskViewSet, ModelConfigViewSet
|
||||
|
||||
router = DefaultRouter()
|
||||
router.register("tasks", AITaskViewSet, basename="ai-task")
|
||||
router.register("models", ModelConfigViewSet, basename="model-config")
|
||||
|
||||
urlpatterns = router.urls
|
||||
@@ -0,0 +1,21 @@
|
||||
from rest_framework.viewsets import ReadOnlyModelViewSet
|
||||
|
||||
from apps.common.api import TeamScopedViewSetMixin
|
||||
|
||||
from .models import AITask, ModelConfig
|
||||
from .serializers import AITaskSerializer, ModelConfigSerializer
|
||||
|
||||
|
||||
class AITaskViewSet(TeamScopedViewSetMixin, ReadOnlyModelViewSet):
|
||||
queryset = AITask.objects.select_related("team", "project", "model_config", "model_config__provider").all()
|
||||
serializer_class = AITaskSerializer
|
||||
search_fields = ["idempotency_key", "provider_task_id", "project__name"]
|
||||
ordering_fields = ["created_at", "updated_at", "completed_at"]
|
||||
|
||||
|
||||
class ModelConfigViewSet(ReadOnlyModelViewSet):
|
||||
queryset = ModelConfig.objects.select_related("provider").filter(status=ModelConfig.Status.ACTIVE)
|
||||
serializer_class = ModelConfigSerializer
|
||||
search_fields = ["name", "display_name", "capability"]
|
||||
ordering_fields = ["created_at", "display_name"]
|
||||
|
||||
Reference in New Issue
Block a user