agents/lead_writer.py

"""Lead Writer agent - Script creation specialist."""
from typing import Any, Dict
from base_agent import BaseAgent


class LeadWriterAgent(BaseAgent):
    """Lead Writer agent that crafts the first draft script."""

    def __init__(self):
        """Initialize the Lead Writer agent."""
        system_prompt = (
            "You are the Lead Writer. Using the outline and cultural notes, write a complete "
            "episode script. Each character must speak in their own voice. Include scene headings, "
            "action lines, and dialogue. Format it like a proper script.\n\n"
            "CRITICAL CHARACTER CONSISTENCY:\n"
            "You will receive a 'character_list' in the inputs. You MUST use ONLY these characters "
            "in your script. Do NOT invent new characters or use different names. If the list is "
            "['Alex', 'Jordan', 'Casey', 'Morgan'], use ONLY these names. Every character name in "
            "your script MUST match exactly one from the provided list.\n\n"
            "POP CULTURE REFERENCE LIMIT:\n"
            "Limit pop culture references to 1-2 per scene maximum. Avoid reference dumps that "
            "make the dialogue feel forced or unnatural.\n\n"
            "Output as JSON with keys: full_episode_first_draft, scene_descriptions, dialogue, stage_directions."
        )
        super().__init__(
            agent_id="lead-writer",
            agent_name="Lead Writer",
            system_prompt=system_prompt,
        )

    def write_script(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Write episode script from outline and cultural notes.

        Args:
            inputs: Dictionary with keys:
                - approved_outline: From Story Editor
                - cultural_consultant_notes: From Cultural Consultant
                - character_voice_guide: Character definitions
                - character_list: List of valid character names (MUST use only these)
                - story_premise: Story premise from Showrunner

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