agents/comedy_writer.py

"""Comedy Writer agent - Humor enhancement specialist."""
from typing import Any, Dict
from base_agent import BaseAgent


class ComedyWriterAgent(BaseAgent):
    """Comedy Writer agent that enhances script with humor."""

    def __init__(self):
        """Initialize the Comedy Writer agent."""
        system_prompt = (
            "You are the Comedy Writer. Punch up the script — sharpen punchlines, tighten "
            "comic timing, make the first 3 seconds irresistible. Don't break the story.\n\n"
            "CRITICAL:\n"
            "1. You MUST punch up the EXISTING script provided in 'dialogue_polished_script'.\n"
            "2. Do NOT create a new standalone scene or new characters.\n"
            "3. Do NOT introduce characters not in the character_list.\n"
            "4. Maintain all existing character names and story structure.\n"
            "5. If the input script is empty, return an error immediately.\n\n"
            "Return the improved full script with a note on what you changed as JSON with keys: "
            "comedy_sharpened_script, punch_up_notes, hook_rewrite_for_opening."
        )
        super().__init__(
            agent_id="comedy-writer",
            agent_name="Comedy Writer",
            system_prompt=system_prompt,
        )

    def add_humor(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Add humor and punch-ups to the script.

        Args:
            inputs: Dictionary with keys:
                - dialogue_polished_script: From Dialogue Specialist (MUST be non-empty)
                - hook_brief_from_showrunner: Hook context
                - character_list: List of valid character names
                - tone_brief: Tone reference from Showrunner

        Returns:
            Dictionary with comedy-enhanced script and notes
        """
        # Validate input
        polished_script = inputs.get("dialogue_polished_script", "")
        if not polished_script or (isinstance(polished_script, str) and not polished_script.strip()):
            return {
                "comedy_sharpened_script": polished_script,
                "punch_up_notes": "ERROR: Received empty dialogue_polished_script. Cannot punch up empty content.",
                "hook_rewrite_for_opening": "",
            }
        
        outputs = self.process(inputs)
        
        # Save state
        self.save_state(
            {
                "inputs": inputs,
                "outputs": outputs,
            },
            filename="comedy_punch_up.json",
        )
        
        return outputs