77 lines
2.2 KiB
Python
77 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from typing import Any
|
|
|
|
|
|
class JobStore:
|
|
def __init__(self) -> None:
|
|
self._jobs: dict[str, dict[str, Any]] = {}
|
|
self._lock = threading.RLock()
|
|
|
|
def create(self, session_id: str, user_id: str, message: str) -> dict[str, Any]:
|
|
job_id = uuid.uuid4().hex
|
|
job = {
|
|
"job_id": job_id,
|
|
"session_id": session_id,
|
|
"user_id": user_id,
|
|
"message": message,
|
|
"status": "queued",
|
|
"created_at": time.time(),
|
|
"updated_at": time.time(),
|
|
"events": [],
|
|
"answer": None,
|
|
"usage": None,
|
|
"error": None,
|
|
}
|
|
with self._lock:
|
|
self._jobs[job_id] = job
|
|
return job
|
|
|
|
def get(self, job_id: str) -> dict[str, Any] | None:
|
|
with self._lock:
|
|
job = self._jobs.get(job_id)
|
|
if not job:
|
|
return None
|
|
return {
|
|
key: (value.copy() if isinstance(value, list) else value)
|
|
for key, value in job.items()
|
|
}
|
|
|
|
def append_event(self, job_id: str, event: dict[str, Any]) -> None:
|
|
with self._lock:
|
|
job = self._jobs[job_id]
|
|
seq = len(job["events"]) + 1
|
|
job["events"].append({"seq": seq, **event})
|
|
job["updated_at"] = time.time()
|
|
|
|
def set_status(self, job_id: str, status: str) -> None:
|
|
with self._lock:
|
|
job = self._jobs[job_id]
|
|
job["status"] = status
|
|
job["updated_at"] = time.time()
|
|
|
|
def finish(
|
|
self,
|
|
job_id: str,
|
|
*,
|
|
answer: str,
|
|
usage: dict[str, Any] | None,
|
|
) -> None:
|
|
with self._lock:
|
|
job = self._jobs[job_id]
|
|
job["status"] = "completed"
|
|
job["answer"] = answer
|
|
job["usage"] = usage
|
|
job["updated_at"] = time.time()
|
|
|
|
def fail(self, job_id: str, error_message: str) -> None:
|
|
with self._lock:
|
|
job = self._jobs[job_id]
|
|
job["status"] = "failed"
|
|
job["error"] = error_message
|
|
job["updated_at"] = time.time()
|
|
|