97 lines
4 KiB
Python
97 lines
4 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from django.utils import timezone
|
||
|
|
|
||
|
|
from control_plane.authoring.models import ChapterRevision
|
||
|
|
from graph.bootstrap import champion_story_authoring_graph_v2
|
||
|
|
from graph.models import GraphRun, GraphRunStatus
|
||
|
|
|
||
|
|
|
||
|
|
class StoryWorkflowRunner:
|
||
|
|
def __init__(self, workflow: object) -> None:
|
||
|
|
self.workflow = workflow
|
||
|
|
|
||
|
|
def start(self, revision: ChapterRevision, *, max_revisions: int = 2) -> GraphRun:
|
||
|
|
version = champion_story_authoring_graph_v2()
|
||
|
|
graph_run = GraphRun.objects.create(
|
||
|
|
execution_graph_version=version,
|
||
|
|
project=revision.chapter.story.project,
|
||
|
|
status=GraphRunStatus.RUNNING,
|
||
|
|
started_at=timezone.now(),
|
||
|
|
current_node="build_context",
|
||
|
|
metadata={
|
||
|
|
"revision_id": str(revision.id),
|
||
|
|
"initial_revision_id": str(revision.id),
|
||
|
|
"current_revision_id": str(revision.id),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
thread_id = f"story:{revision.chapter.story_id}:chapter:{revision.chapter.number}:revision:{revision.id}"
|
||
|
|
revision.graph_thread_id = thread_id
|
||
|
|
revision.save(update_fields=["graph_thread_id", "updated_at"])
|
||
|
|
initial = {
|
||
|
|
"story_id": str(revision.chapter.story_id),
|
||
|
|
"chapter_id": str(revision.chapter_id),
|
||
|
|
"revision_id": str(revision.id),
|
||
|
|
"graph_run_id": graph_run.id,
|
||
|
|
"thread_id": thread_id,
|
||
|
|
"editorial_finding_ids": [],
|
||
|
|
"patch_finding_ids": [],
|
||
|
|
"patch_attempted": False,
|
||
|
|
"patch_status": "not_needed",
|
||
|
|
"verification_status": "not_needed",
|
||
|
|
}
|
||
|
|
return self._invoke(graph_run, initial)
|
||
|
|
|
||
|
|
def resume(self, graph_run_id: int, decision: dict[str, Any]) -> GraphRun:
|
||
|
|
from langgraph.types import Command
|
||
|
|
|
||
|
|
graph_run = GraphRun.objects.get(id=graph_run_id)
|
||
|
|
if graph_run.status == GraphRunStatus.CANCELLED:
|
||
|
|
raise RuntimeError("cancelled story runs cannot be resumed")
|
||
|
|
graph_run.status = GraphRunStatus.RUNNING
|
||
|
|
graph_run.failure_reason = ""
|
||
|
|
graph_run.save(update_fields=["status", "failure_reason", "updated_at"])
|
||
|
|
value = None if decision.get("action") == "retry" else Command(resume=decision)
|
||
|
|
return self._invoke(graph_run, value)
|
||
|
|
|
||
|
|
def _invoke(self, graph_run: GraphRun, value: object) -> GraphRun:
|
||
|
|
thread_id = ChapterRevision.objects.get(
|
||
|
|
id=graph_run.metadata["revision_id"]
|
||
|
|
).graph_thread_id
|
||
|
|
config = {"configurable": {"thread_id": thread_id}}
|
||
|
|
try:
|
||
|
|
self.workflow.invoke(value, config=config)
|
||
|
|
snapshot = self.workflow.get_state(config)
|
||
|
|
except Exception as exc:
|
||
|
|
graph_run.status = GraphRunStatus.FAILED
|
||
|
|
graph_run.failure_reason = str(exc)[:4000]
|
||
|
|
graph_run.completed_at = timezone.now()
|
||
|
|
graph_run.save(
|
||
|
|
update_fields=["status", "failure_reason", "completed_at", "updated_at"]
|
||
|
|
)
|
||
|
|
raise
|
||
|
|
next_nodes = tuple(snapshot.next or ())
|
||
|
|
current_revision_id = str(snapshot.values.get("revision_id") or graph_run.metadata["revision_id"])
|
||
|
|
graph_run.metadata = {
|
||
|
|
**graph_run.metadata,
|
||
|
|
"current_revision_id": current_revision_id,
|
||
|
|
}
|
||
|
|
if next_nodes:
|
||
|
|
graph_run.status = GraphRunStatus.PAUSED
|
||
|
|
graph_run.current_node = str(next_nodes[0])
|
||
|
|
graph_run.failure_reason = "AWAITING_STORY_APPROVAL"
|
||
|
|
graph_run.save(
|
||
|
|
update_fields=["status", "current_node", "failure_reason", "metadata", "updated_at"]
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
graph_run.status = GraphRunStatus.COMPLETE
|
||
|
|
graph_run.current_node = "complete"
|
||
|
|
graph_run.completed_at = timezone.now()
|
||
|
|
graph_run.metadata = {**graph_run.metadata, "final_state": dict(snapshot.values)}
|
||
|
|
graph_run.save(
|
||
|
|
update_fields=["status", "current_node", "completed_at", "metadata", "updated_at"]
|
||
|
|
)
|
||
|
|
return graph_run
|