134 lines
3.5 KiB
Python
134 lines
3.5 KiB
Python
"""Lightweight generation job/event tracking."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.db.models import GenerationJob, GenerationJobEvent, Story
|
|
|
|
|
|
def _story_snapshot(story: Story | None) -> dict[str, Any]:
|
|
if story is None:
|
|
return {}
|
|
|
|
return {
|
|
"story_id": story.id,
|
|
"mode": story.mode,
|
|
"generation_status": story.generation_status,
|
|
"image_status": story.image_status,
|
|
"audio_status": story.audio_status,
|
|
"retryable_assets": story.retryable_assets,
|
|
"last_error": story.last_error,
|
|
}
|
|
|
|
|
|
def _job_status_from_story(story: Story) -> str:
|
|
if story.generation_status == "failed":
|
|
return "failed"
|
|
if story.generation_status == "degraded_completed":
|
|
return "degraded_completed"
|
|
return "completed"
|
|
|
|
|
|
async def create_generation_job(
|
|
db: AsyncSession,
|
|
*,
|
|
user_id: str,
|
|
output_mode: str,
|
|
input_type: str,
|
|
request_payload: dict[str, Any],
|
|
story_id: int | None = None,
|
|
) -> GenerationJob:
|
|
"""Create a generation job and record its first event."""
|
|
|
|
job = GenerationJob(
|
|
user_id=user_id,
|
|
story_id=story_id,
|
|
output_mode=output_mode,
|
|
input_type=input_type,
|
|
status="running",
|
|
current_step="request_accepted",
|
|
request_payload=request_payload,
|
|
result_snapshot={},
|
|
)
|
|
db.add(job)
|
|
await db.flush()
|
|
await record_generation_event(
|
|
db,
|
|
job=job,
|
|
story_id=story_id,
|
|
event_type="request_accepted",
|
|
status="succeeded",
|
|
message="Generation request accepted.",
|
|
metadata={"output_mode": output_mode, "input_type": input_type},
|
|
commit=False,
|
|
)
|
|
await db.commit()
|
|
await db.refresh(job)
|
|
return job
|
|
|
|
|
|
async def record_generation_event(
|
|
db: AsyncSession,
|
|
*,
|
|
job: GenerationJob,
|
|
event_type: str,
|
|
status: str,
|
|
story_id: int | None = None,
|
|
message: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
commit: bool = True,
|
|
) -> GenerationJobEvent:
|
|
"""Append one event to an existing generation job."""
|
|
|
|
event = GenerationJobEvent(
|
|
job_id=job.id,
|
|
story_id=story_id if story_id is not None else job.story_id,
|
|
event_type=event_type,
|
|
status=status,
|
|
message=message,
|
|
event_metadata=metadata or {},
|
|
)
|
|
db.add(event)
|
|
if commit:
|
|
await db.commit()
|
|
return event
|
|
|
|
|
|
async def finish_generation_job(
|
|
db: AsyncSession,
|
|
*,
|
|
job: GenerationJob,
|
|
story: Story | None,
|
|
status: str | None = None,
|
|
current_step: str,
|
|
error_message: str | None = None,
|
|
message: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> GenerationJob:
|
|
"""Mark a generation job as completed/degraded/failed and append a final event."""
|
|
|
|
job.story_id = story.id if story is not None else job.story_id
|
|
job.status = status or (_job_status_from_story(story) if story is not None else "failed")
|
|
job.current_step = current_step
|
|
job.error_message = error_message
|
|
job.result_snapshot = _story_snapshot(story)
|
|
await record_generation_event(
|
|
db,
|
|
job=job,
|
|
story_id=job.story_id,
|
|
event_type=current_step,
|
|
status=job.status,
|
|
message=message,
|
|
metadata={
|
|
**(metadata or {}),
|
|
"result_snapshot": job.result_snapshot,
|
|
},
|
|
commit=False,
|
|
)
|
|
await db.commit()
|
|
await db.refresh(job)
|
|
return job
|