38 lines
1.0 KiB
Python
38 lines
1.0 KiB
Python
from dataclasses import dataclass, field
|
|
from threading import Lock
|
|
from uuid import uuid4
|
|
|
|
from app.schemas.changeset import ChangeItem
|
|
from app.schemas.chat import TaskResultType, TaskStatus
|
|
from app.schemas.common import ErrorPayload
|
|
|
|
|
|
@dataclass
|
|
class TaskState:
|
|
task_id: str
|
|
status: TaskStatus = TaskStatus.QUEUED
|
|
result_type: TaskResultType | None = None
|
|
answer: str | None = None
|
|
changeset: list[ChangeItem] = field(default_factory=list)
|
|
error: ErrorPayload | None = None
|
|
|
|
|
|
class TaskStore:
|
|
def __init__(self) -> None:
|
|
self._items: dict[str, TaskState] = {}
|
|
self._lock = Lock()
|
|
|
|
def create(self) -> TaskState:
|
|
task = TaskState(task_id=str(uuid4()))
|
|
with self._lock:
|
|
self._items[task.task_id] = task
|
|
return task
|
|
|
|
def get(self, task_id: str) -> TaskState | None:
|
|
with self._lock:
|
|
return self._items.get(task_id)
|
|
|
|
def save(self, task: TaskState) -> None:
|
|
with self._lock:
|
|
self._items[task.task_id] = task
|