Artifex/control_plane/authoring/management/commands/fiction_ideas.py

from __future__ import annotations

import json
from pathlib import Path

from django.core.management.base import BaseCommand, CommandError

from control_plane.authoring.models import BookStateVersion, DocumentAuthority, SceneIdeation, Work
from control_plane.authoring.prompts import SCENE_IDEA_TYPES
from control_plane.authoring.standalone_scenes import (
    SceneIdeationService,
    export_scene_ideation_markdown,
)
from model_router.providers import providers_from_resources
from model_router.router import ModelRouter


class Command(BaseCommand):
    help = "Propose cited scene ideas and select one into the standalone scene workflow."

    def add_arguments(self, parser) -> None:
        parser.add_argument("action", choices=["propose", "show", "export", "select"])
        parser.add_argument("--id")
        parser.add_argument("--series-slug")
        parser.add_argument("--work-slug")
        parser.add_argument("--target-book")
        parser.add_argument("--book-state")
        parser.add_argument("--chapter-key")
        parser.add_argument("--focus", default="")
        parser.add_argument("--candidate-count", type=int, default=10)
        parser.add_argument(
            "--scene-type",
            action="append",
            choices=SCENE_IDEA_TYPES,
        )
        parser.add_argument(
            "--include-authority",
            action="append",
            choices=DocumentAuthority.values,
        )
        parser.add_argument("--pin-document", action="append", default=[])
        parser.add_argument("--governing-document", action="append", default=[])
        parser.add_argument("--candidate-id")
        parser.add_argument("--target-words", type=int)
        parser.add_argument("--model")
        parser.add_argument("--output", type=Path)
        parser.add_argument("--compact", action="store_true")

    def handle(self, *args, **options) -> None:
        action = options["action"]
        try:
            if action == "propose":
                service = self._service()
                work = self._work(options)
                idea = service.propose(
                    work=work,
                    target_book=str(options.get("target_book") or ""),
                    focus=options["focus"],
                    candidate_count=options["candidate_count"],
                    scene_types=options["scene_type"],
                    authorities=options["include_authority"],
                    pinned_document_keys=options["pin_document"],
                    governing_document_keys=options["governing_document"],
                    detail_level="compact" if options["compact"] else "full",
                    model_hint=options["model"],
                    book_state=self._book_state(options),
                )
                self._write_idea(idea)
                return
            idea = self._idea(options)
            if action == "export":
                output = options.get("output")
                if output is None:
                    raise CommandError("export requires --output")
                export_scene_ideation_markdown(idea, output, compact=options["compact"])
                self.stdout.write(str(output))
                return
            if action == "select":
                service = self._service()
                candidate_id = str(options.get("candidate_id") or "").strip()
                if not candidate_id:
                    raise CommandError("select requires --candidate-id")
                scene, created = service.select_candidate(
                    idea,
                    candidate_id=candidate_id,
                    target_words=options["target_words"],
                    book_chapter_key=options.get("chapter_key"),
                )
                idea.refresh_from_db()
                self.stdout.write(
                    json.dumps(
                        {
                            "created": created,
                            "scene_id": str(scene.id),
                            "scene_status": scene.status,
                            "scene_title": scene.title,
                            "idea": self._payload(idea),
                        },
                        ensure_ascii=False,
                        indent=2,
                    )
                )
                return
            self._write_idea(idea)
        except (OSError, RuntimeError, TypeError, ValueError) as exc:
            raise CommandError(str(exc)) from exc

    @staticmethod
    def _service() -> SceneIdeationService:
        return SceneIdeationService(ModelRouter(providers_from_resources(), persist_requests=True))

    def _work(self, options: dict) -> Work:
        if not options.get("series_slug") or not options.get("work_slug"):
            raise CommandError("propose requires --series-slug and --work-slug")
        work = Work.objects.filter(
            series__slug=options["series_slug"], slug=options["work_slug"]
        ).first()
        if work is None:
            raise CommandError("work not found; register sources first")
        return work

    def _idea(self, options: dict) -> SceneIdeation:
        if not options.get("id"):
            raise CommandError(f"{options['action']} requires --id")
        idea = SceneIdeation.objects.select_related("work__series").filter(id=options["id"]).first()
        if idea is None:
            raise CommandError("scene ideation not found")
        return idea

    @staticmethod
    def _book_state(options: dict) -> BookStateVersion | None:
        state_id = options.get("book_state")
        if not state_id:
            return None
        state = BookStateVersion.objects.filter(id=state_id).first()
        if state is None:
            raise CommandError("book state not found")
        return state

    def _write_idea(self, idea: SceneIdeation) -> None:
        self.stdout.write(json.dumps(self._payload(idea), ensure_ascii=False, indent=2))

    @staticmethod
    def _payload(idea: SceneIdeation) -> dict:
        return {
            "id": str(idea.id),
            "series": idea.work.series.slug,
            "work": idea.work.slug,
            "book_state_id": str(idea.book_state_id) if idea.book_state_id else None,
            "target_book": idea.target_book,
            "requested_scene_types": idea.requested_scene_types,
            "focus": idea.focus,
            "authorities": idea.authorities,
            "context_pack_sha256": idea.context_pack_sha256,
            "governing_document_keys": (idea.context_pack or {}).get("governing_document_keys")
            or [],
            "citations": (idea.context_pack or {}).get("citations") or [],
            "candidates": idea.candidates,
            "generation_metadata": idea.generation_metadata,
        }