146 lines
4.0 KiB
Python
146 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from datetime import datetime
|
|
import threading
|
|
from typing import Any, Optional
|
|
|
|
|
|
_RUNTIME_LOCK = threading.RLock()
|
|
_JOB_STATES: dict[str, dict[str, Any]] = {}
|
|
|
|
|
|
def _now_str() -> str:
|
|
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
|
|
def _chapter_payload(
|
|
*,
|
|
section_key: str,
|
|
section_title: str,
|
|
section_order: int,
|
|
status: str = "pending",
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"sectionKey": section_key,
|
|
"sectionTitle": section_title,
|
|
"sectionOrder": section_order,
|
|
"status": status,
|
|
"content": None,
|
|
"errorMessage": None,
|
|
"updatedAt": _now_str(),
|
|
"promptText": None,
|
|
"evidencePayload": None,
|
|
"validationPayload": None,
|
|
}
|
|
|
|
|
|
def init_job_state(
|
|
*,
|
|
job_id: str,
|
|
project_id: str,
|
|
template_id: Optional[str],
|
|
chapters: list[dict[str, Any]],
|
|
) -> None:
|
|
with _RUNTIME_LOCK:
|
|
_JOB_STATES[job_id] = {
|
|
"jobId": job_id,
|
|
"projectId": project_id,
|
|
"templateId": template_id,
|
|
"status": "pending",
|
|
"progress": 0,
|
|
"currentSectionKey": None,
|
|
"errorMessage": None,
|
|
"createdAt": _now_str(),
|
|
"updatedAt": _now_str(),
|
|
"completedAt": None,
|
|
"chapters": {
|
|
str(item["sectionKey"]): _chapter_payload(
|
|
section_key=str(item["sectionKey"]),
|
|
section_title=str(item["sectionTitle"]),
|
|
section_order=int(item["sectionOrder"]),
|
|
status=str(item.get("status") or "pending"),
|
|
)
|
|
for item in (chapters or [])
|
|
},
|
|
}
|
|
|
|
|
|
def get_job_state(job_id: str) -> Optional[dict[str, Any]]:
|
|
with _RUNTIME_LOCK:
|
|
state = _JOB_STATES.get(job_id)
|
|
return deepcopy(state) if state else None
|
|
|
|
|
|
def update_job_state(job_id: str, **fields: Any) -> None:
|
|
with _RUNTIME_LOCK:
|
|
state = _JOB_STATES.get(job_id)
|
|
if not state:
|
|
return
|
|
state.update(fields)
|
|
state["updatedAt"] = _now_str()
|
|
|
|
|
|
def update_chapter_state(
|
|
job_id: str,
|
|
section_key: str,
|
|
**fields: Any,
|
|
) -> None:
|
|
with _RUNTIME_LOCK:
|
|
state = _JOB_STATES.get(job_id)
|
|
if not state:
|
|
return
|
|
chapter = state.get("chapters", {}).get(section_key)
|
|
if not chapter:
|
|
return
|
|
chapter.update(fields)
|
|
chapter["updatedAt"] = _now_str()
|
|
state["updatedAt"] = _now_str()
|
|
|
|
|
|
def append_chapter_content(
|
|
job_id: str,
|
|
section_key: str,
|
|
delta_text: str,
|
|
*,
|
|
stream_phase: str,
|
|
) -> None:
|
|
if not delta_text:
|
|
return
|
|
with _RUNTIME_LOCK:
|
|
state = _JOB_STATES.get(job_id)
|
|
if not state:
|
|
return
|
|
chapter = state.get("chapters", {}).get(section_key)
|
|
if not chapter:
|
|
return
|
|
current = str(chapter.get("content") or "")
|
|
validation_payload = dict(chapter.get("validationPayload") or {})
|
|
validation_payload["streamPhase"] = stream_phase
|
|
chapter["content"] = current + delta_text
|
|
chapter["validationPayload"] = validation_payload
|
|
chapter["updatedAt"] = _now_str()
|
|
state["currentSectionKey"] = section_key
|
|
state["updatedAt"] = _now_str()
|
|
|
|
|
|
def set_chapter_stream_phase(job_id: str, section_key: str, stream_phase: str) -> None:
|
|
with _RUNTIME_LOCK:
|
|
state = _JOB_STATES.get(job_id)
|
|
if not state:
|
|
return
|
|
chapter = state.get("chapters", {}).get(section_key)
|
|
if not chapter:
|
|
return
|
|
validation_payload = dict(chapter.get("validationPayload") or {})
|
|
validation_payload["streamPhase"] = stream_phase
|
|
chapter["validationPayload"] = validation_payload
|
|
chapter["updatedAt"] = _now_str()
|
|
state["currentSectionKey"] = section_key
|
|
state["updatedAt"] = _now_str()
|
|
|
|
|
|
def remove_job_state(job_id: str) -> None:
|
|
with _RUNTIME_LOCK:
|
|
_JOB_STATES.pop(job_id, None)
|