86 lines
3.4 KiB
Python
86 lines
3.4 KiB
Python
|
"""Persistent projections of native jobs; this is not another scheduler."""
|
||
|
|
|
||
|
|
from dataclasses import asdict, dataclass, field
|
||
|
|
from datetime import UTC, datetime
|
||
|
|
from uuid import uuid4
|
||
|
|
|
||
|
|
from sqlalchemy import select
|
||
|
|
from pydantic.alias_generators import to_camel
|
||
|
|
from db.models import BackgroundTask, BackgroundTaskEvent
|
||
|
|
from db.session import SessionLocal
|
||
|
|
|
||
|
|
TERMINAL = {"succeeded", "failed", "cancelled"}
|
||
|
|
ACTIVE = {"queued", "submitted", "running", "waiting_input", "cancel_requested", "unknown"}
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class TaskRecord:
|
||
|
|
id: str
|
||
|
|
assistant_id: str
|
||
|
|
conversation_id: str
|
||
|
|
tool_id: str
|
||
|
|
instruction: str
|
||
|
|
origin_key: str
|
||
|
|
connection: dict
|
||
|
|
secrets: dict
|
||
|
|
status: str = "queued"
|
||
|
|
summary: str = "任务已排队"
|
||
|
|
result: str = ""
|
||
|
|
remote_handle: dict = field(default_factory=dict)
|
||
|
|
input_request: dict = field(default_factory=dict)
|
||
|
|
job_id: str = ""
|
||
|
|
sequence: int = 0
|
||
|
|
created_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
|
||
|
|
updated_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
|
||
|
|
|
||
|
|
def public(self):
|
||
|
|
hidden = {"connection", "secrets", "origin_key", "remote_handle"}
|
||
|
|
return {to_camel(key): value for key, value in asdict(self).items() if key not in hidden}
|
||
|
|
|
||
|
|
|
||
|
|
def from_row(row) -> TaskRecord:
|
||
|
|
return TaskRecord(**{key: getattr(row, key) for key in TaskRecord.__dataclass_fields__})
|
||
|
|
|
||
|
|
|
||
|
|
class TaskRepository:
|
||
|
|
async def create(self, record: TaskRecord) -> TaskRecord:
|
||
|
|
async with SessionLocal() as db:
|
||
|
|
# All calls are serialized by the application entry point; unique
|
||
|
|
# origin_key also protects against duplicate tool delivery.
|
||
|
|
existing = (await db.execute(select(BackgroundTask).where(BackgroundTask.origin_key == record.origin_key))).scalar_one_or_none()
|
||
|
|
if existing:
|
||
|
|
return from_row(existing)
|
||
|
|
db.add(BackgroundTask(**asdict(record)))
|
||
|
|
await db.commit()
|
||
|
|
return record
|
||
|
|
|
||
|
|
async def get(self, task_id: str) -> TaskRecord | None:
|
||
|
|
async with SessionLocal() as db:
|
||
|
|
row = await db.get(BackgroundTask, task_id)
|
||
|
|
return from_row(row) if row else None
|
||
|
|
|
||
|
|
async def list(self, assistant_id: str | None = None, *, active=False) -> list[TaskRecord]:
|
||
|
|
statement = select(BackgroundTask).order_by(BackgroundTask.created_at.desc())
|
||
|
|
if not active:
|
||
|
|
statement = statement.limit(100)
|
||
|
|
if assistant_id:
|
||
|
|
statement = statement.where(BackgroundTask.assistant_id == assistant_id)
|
||
|
|
if active:
|
||
|
|
statement = statement.where(BackgroundTask.status.in_(ACTIVE))
|
||
|
|
async with SessionLocal() as db:
|
||
|
|
return [from_row(row) for row in (await db.execute(statement)).scalars().all()]
|
||
|
|
|
||
|
|
async def patch(self, task_id: str, **changes) -> TaskRecord:
|
||
|
|
async with SessionLocal() as db:
|
||
|
|
row = (await db.execute(select(BackgroundTask).where(BackgroundTask.id == task_id).with_for_update())).scalar_one()
|
||
|
|
if row.status in TERMINAL:
|
||
|
|
return from_row(row)
|
||
|
|
for key, value in changes.items():
|
||
|
|
setattr(row, key, value)
|
||
|
|
row.sequence += 1
|
||
|
|
row.updated_at = datetime.now(UTC).isoformat()
|
||
|
|
record = from_row(row)
|
||
|
|
db.add(BackgroundTaskEvent(id="te_" + uuid4().hex, task_id=task_id, sequence=row.sequence, data=record.public()))
|
||
|
|
await db.commit()
|
||
|
|
return record
|