opus-orchestrator-ai/opus_orchestrator/schemas/book.py

"""Pydantic schemas for Opus Orchestrator."""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Optional

from pydantic import BaseModel, Field


class BookType(str, Enum):
    """Type of book being generated."""

    FICTION = "fiction"
    NONFICTION = "nonfiction"


class Genre(str, Enum):
    """Fiction genres."""

    MYSTERY = "mystery"
    ROMANCE = "romance"
    THRILLER = "thriller"
    FANTASY = "fantasy"
    SCIENCE_FICTION = "science-fiction"
    HORROR = "horror"
    LITERARY = "literary"
    OTHER = "other"


class NonfictionGenre(str, Enum):
    """Nonfiction genres."""

    BUSINESS = "business"
    SELF_HELP = "self-help"
    HISTORY = "history"
    BIOGRAPHY = "biography"
    SCIENCE = "science"
    PHILOSOPHY = "philosophy"
    MEMOIR = "memoir"
    ACADEMIC = "academic"
    OTHER = "other"


# --- Input Schemas ---


class BookIntent(BaseModel):
    """User's intent for the book project."""

    book_type: BookType
    genre: Optional[str] = None
    working_title: Optional[str] = None
    target_audience: str = Field(description="Who is this book for?")
    intended_outcome: str = Field(description="What should the final product achieve?")
    tone: Optional[str] = None
    target_word_count: int = Field(default=80000, ge=1000, le=500000)
    special_instructions: Optional[str] = None


class RawContent(BaseModel):
    """Raw content extracted from source."""

    content_type: str = Field(description="Type: outline, notes, stream-of-consciousness, etc.")
    text: str = Field(description="The raw content itself")
    metadata: dict[str, Any] = Field(default_factory=dict)


# --- Blueprint Schemas ---


class ChapterBlueprint(BaseModel):
    """Blueprint for a single chapter."""

    chapter_number: int
    title: str
    summary: str = Field(description="Brief summary of chapter content")
    word_count_target: int
    pov_character: Optional[str] = None
    key_events: list[str] = Field(default_factory=list)
    themes: list[str] = Field(default_factory=list)
    scenes: list[dict[str, Any]] = Field(default_factory=list)


class BookBlueprint(BaseModel):
    """Complete blueprint for the book."""

    title: str
    subtitle: Optional[str] = None
    genre: str
    target_audience: str
    target_word_count: int
    structure: str = Field(description="Overall structural approach")
    themes: list[str] = Field(default_factory=list)
    tone: str
    chapters: list[ChapterBlueprint] = Field(default_factory=list)
    characters: list[dict[str, Any]] = Field(default_factory=list, description="For fiction")
    world_elements: dict[str, Any] = Field(default_factory=dict, description="For fiction")
    key_arguments: list[str] = Field(default_factory=list, description="For nonfiction")


# --- Agent Output Schemas ---


class ChapterDraft(BaseModel):
    """A chapter draft from the writer agent."""

    chapter_number: int
    title: str
    content: str
    word_count: int
    notes: Optional[str] = None


class CritiqueScore(BaseModel):
    """Score from a critic agent."""

    criterion: str
    score: float = Field(ge=0.0, le=1.0)
    strengths: list[str] = Field(default_factory=list)
    weaknesses: list[str] = Field(default_factory=list)
    suggestions: list[str] = Field(default_factory=list)


class ChapterCritique(BaseModel):
    """Complete critique of a chapter."""

    chapter_number: int
    overall_score: float = Field(ge=0.0, le=1.0)
    criteria_scores: list[CritiqueScore]
    consensus_strengths: list[str] = Field(default_factory=list)
    consensus_weaknesses: list[str] = Field(default_factory=list)
    revision_priority: str = Field(description="major_revisions, minor_revisions, approved")
    notes: Optional[str] = None


class Revision(BaseModel):
    """Revision instructions for a chapter."""

    chapter_number: int
    priority: str = Field(description="major or minor")
    changes_required: list[str] = Field(default_factory=list)
    preserve_elements: list[str] = Field(default_factory=list)
    notes: Optional[str] = None


# --- Final Output Schemas ---


class Chapter(BaseModel):
    """Final chapter in the manuscript."""

    chapter_number: int
    title: str
    content: str
    word_count: int


class Manuscript(BaseModel):
    """Complete manuscript."""

    title: str
    subtitle: Optional[str] = None
    author: str = "Opus Orchestrator AI"
    book_type: BookType
    genre: str
    chapters: list[Chapter]
    total_word_count: int
    frontmatter: dict[str, Any] = Field(default_factory=dict)
    backmatter: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)
    generated_at: datetime = Field(default_factory=datetime.utcnow)

    def to_markdown(self) -> str:
        """Convert to markdown format."""
        lines = []

        # Frontmatter
        if self.frontmatter.get("dedication"):
            lines.append(f"*{self.frontmatter['dedication']}*\n")
        if self.frontmatter.get("epigraph"):
            lines.append(f"> {self.frontmatter['epigraph']}\n")

        # Title
        lines.append(f"# {self.title}")
        if self.subtitle:
            lines.append(f"## {self.subtitle}")
        lines.append(f"\n*by {self.author}*\n")

        # TOC
        if self.frontmatter.get("include_toc", True):
            lines.append("## Table of Contents\n")
            for ch in self.chapters:
                lines.append(f"{ch.chapter_number}. [{ch.title}](#chapter-{ch.chapter_number})")
            lines.append("")

        # Chapters
        for ch in self.chapters:
            lines.append(f"## Chapter {ch.chapter_number}: {ch.title}\n")
            lines.append(ch.content)
            lines.append("")

        return "\n".join(lines)


# --- State Schemas ---


class IterationState(BaseModel):
    """State for iteration tracking."""

    round: int = 0
    chapter_number: int
    critiques: list[ChapterCritique] = Field(default_factory=list)
    current_score: float = 0.0
    approved: bool = False
    revisions_requested: list[Revision] = Field(default_factory=list)