agents/story_editor.py

"""Story Editor agent - Structure specialist for content."""
from typing import Any, Dict
from base_agent import BaseAgent


class StoryEditorAgent(BaseAgent):
    """Story Editor agent that develops structural outlines."""

    def __init__(self):
        """Initialize the Story Editor agent."""
        system_prompt = (
            "You are the Story Editor. Given a showrunner directive, produce a structured "
            "episode outline with clear act breaks, a hook, rising tension, and a satisfying end. "
            "\n\n"
            "CRITICAL CHARACTER CONSISTENCY:\n"
            "You will receive a 'character_list' in the inputs. You MUST use ONLY the characters "
            "from this list. Do NOT invent new character names. If the character list is "
            "['Alex', 'Jordan', 'Casey', 'Morgan'], then use ONLY these names throughout your outline. "
            "No other names like Maya, Derek, Zoe, or any invented characters.\n\n"
            "Flag any continuity risks. Output the outline as a JSON object with keys: "
            "episode_outline, act_structure, story_notes_for_writers."
        )
        super().__init__(
            agent_id="story-editor",
            agent_name="Story Editor",
            system_prompt=system_prompt,
        )

    def generate_outline(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Generate episode outline from showrunner directive.

        Args:
            inputs: Dictionary with keys:
                - episode_directive: From Showrunner
                - series_continuity_log: Continuity reference
                - character_list: List of valid character names (MUST use only these)

        Returns:
            Dictionary with generated outline and related outputs
        """
        outputs = self.process(inputs)
        
        # Save state
        self.save_state(
            {
                "inputs": inputs,
                "outputs": outputs,
            },
            filename="outline.json",
        )
        
        return outputs