Artifex/control_plane/authoring/workflow.py

from __future__ import annotations

from typing import Any

from control_plane.authoring.services import DjangoStoryWorkflowServices
from control_plane.authoring.state import StoryGraphState


def build_story_workflow(services: DjangoStoryWorkflowServices, checkpointer: object):
    try:
        from langgraph.graph import END, StateGraph
        from langgraph.types import interrupt
    except ImportError as exc:
        raise RuntimeError("Story authoring requires LangGraph") from exc

    graph = StateGraph(StoryGraphState)

    def build_context(state: StoryGraphState) -> dict[str, Any]:
        return services.build_context(dict(state))

    def plan_chapter(state: StoryGraphState) -> dict[str, Any]:
        return services.plan_chapter(dict(state))

    def approve_plan(state: StoryGraphState) -> dict[str, Any]:
        approval = services.ensure_approval(
            dict(state),
            "STORY_PLAN_APPROVAL",
            {
                "type": "story_plan",
                "revision_id": state["revision_id"],
                "scene_plan": state.get("scene_plan", {}),
                "allowed_actions": ["approve", "request_revision", "reject"],
            },
        )
        decision = interrupt(approval.payload)
        services.decide_approval(approval.id, decision)
        return {
            "approval_action": str(decision.get("action") or "reject").lower(),
            "human_notes": str(decision.get("notes") or ""),
        }

    def draft_chapter(state: StoryGraphState) -> dict[str, Any]:
        return services.draft_chapter(dict(state))

    def extract_continuity(state: StoryGraphState) -> dict[str, Any]:
        method = getattr(services, "extract_final_state", services.extract_continuity)
        return method(dict(state))

    def quality_review(state: StoryGraphState) -> dict[str, Any]:
        method = getattr(services, "quality_review", None)
        if method is None:
            return {"editorial_finding_ids": []}
        return method(dict(state))

    def judge_state_contract(state: StoryGraphState) -> dict[str, Any]:
        method = getattr(services, "finalize_combined_audit", None)
        if method is None:
            method = getattr(services, "judge_state_contract", None)
        if method is None:
            return {"state_judge_status": "pass"}
        return method(dict(state))

    def decide_patch(state: StoryGraphState) -> dict[str, Any]:
        return services.decide_patch(dict(state))

    def apply_patch(state: StoryGraphState) -> dict[str, Any]:
        return services.apply_automatic_patch(dict(state))

    def verify_patch(state: StoryGraphState) -> dict[str, Any]:
        return services.verify_patch(dict(state))

    def approve_chapter(state: StoryGraphState) -> dict[str, Any]:
        state_payload_method = getattr(services, "state_approval_payload", None)
        state_payload = state_payload_method(dict(state)) if state_payload_method else {}
        approval = services.ensure_approval(
            dict(state),
            "STORY_CHAPTER_APPROVAL",
            {
                "type": "story_chapter",
                "revision_id": state["revision_id"],
                "finding_ids": state.get("editorial_finding_ids", []),
                **state_payload,
                "allowed_actions": ["approve", "request_revision", "reject"],
            },
        )
        decision = interrupt(approval.payload)
        services.decide_approval(approval.id, decision)
        return {
            "approval_action": str(decision.get("action") or "reject").lower(),
            "human_notes": str(decision.get("notes") or ""),
        }

    def commit_chapter(state: StoryGraphState) -> dict[str, Any]:
        return services.commit_chapter(dict(state))

    def publish_story(state: StoryGraphState) -> dict[str, Any]:
        return {"export_uri": services.publish_story(dict(state))}

    graph.add_node("build_context", build_context)
    graph.add_node("plan_chapter", plan_chapter)
    graph.add_node("approve_plan", approve_plan)
    graph.add_node("draft_chapter", draft_chapter)
    graph.add_node("quality_review", quality_review)
    graph.add_node("extract_continuity", extract_continuity)
    graph.add_node("review_draft", judge_state_contract)
    graph.add_node("decide_patch", decide_patch)
    graph.add_node("apply_patch", apply_patch)
    graph.add_node("extract_patched_continuity", extract_continuity)
    graph.add_node("verify_patch", verify_patch)
    graph.add_node("approve_chapter", approve_chapter)
    graph.add_node("commit_chapter", commit_chapter)
    graph.add_node("publish_story", publish_story)
    graph.add_node("manual_revision", lambda state: {})
    graph.add_node("reject", lambda state: {})
    graph.set_entry_point("build_context")
    graph.add_edge("build_context", "plan_chapter")
    graph.add_edge("plan_chapter", "approve_plan")
    graph.add_conditional_edges(
        "approve_plan",
        lambda state: state.get("approval_action", "reject"),
        {"approve": "draft_chapter", "request_revision": "plan_chapter", "reject": "reject"},
    )
    graph.add_edge("draft_chapter", "quality_review")
    graph.add_edge("quality_review", "decide_patch")
    graph.add_conditional_edges(
        "decide_patch",
        lambda state: state.get("patch_decision", "human_review"),
        {"patch": "apply_patch", "human_review": "extract_continuity"},
    )
    graph.add_conditional_edges(
        "apply_patch",
        lambda state: state.get("patch_status", "failed"),
        {"applied": "extract_patched_continuity", "failed": "extract_continuity"},
    )
    graph.add_edge("extract_continuity", "review_draft")
    graph.add_edge("extract_patched_continuity", "verify_patch")
    graph.add_edge("verify_patch", "review_draft")
    graph.add_edge("review_draft", "approve_chapter")
    graph.add_conditional_edges(
        "approve_chapter",
        lambda state: state.get("approval_action", "reject"),
        {"approve": "commit_chapter", "request_revision": "manual_revision", "reject": "reject"},
    )
    graph.add_edge("commit_chapter", "publish_story")
    graph.add_edge("publish_story", END)
    graph.add_edge("manual_revision", END)
    graph.add_edge("reject", END)
    return graph.compile(checkpointer=checkpointer)