agents/dialogue_specialist.py

"""Dialogue Specialist agent - Voice polish expert."""
from typing import Any, Dict
from base_agent import BaseAgent


class DialogueSpecialistAgent(BaseAgent):
    """Dialogue Specialist agent that refines dialogue for authenticity."""

    def __init__(self):
        """Initialize the Dialogue Specialist agent."""
        system_prompt = (
            "You are the Dialogue Specialist. Polish the script's dialogue only — don't touch "
            "action lines or structure. Make each character's speech feel authentic, natural, and "
            "distinct.\n\n"
            "CRITICAL:\n"
            "1. You will receive 'first_draft_script' as a STRING (not empty). If it's empty, "
            "return an error immediately.\n"
            "2. Do NOT modify character names. Use ONLY the character_list provided.\n"
            "3. Return the FULL script with improved dialogue.\n\n"
            "Return the full script with improved dialogue as JSON with keys: "
            "dialogue_polished_script, voice_consistency_notes."
        )
        super().__init__(
            agent_id="dialogue-specialist",
            agent_name="Dialogue Specialist",
            system_prompt=system_prompt,
        )

    def polish_dialogue(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Polish dialogue in the script.

        Args:
            inputs: Dictionary with keys:
                - first_draft_script: From Lead Writer (MUST be non-empty string)
                - character_voice_guide: Character voice definitions
                - character_list: List of valid character names
                - dialect_slang_reference: Reference material

        Returns:
            Dictionary with polished script and voice notes
        """
        # Validate input
        first_draft = inputs.get("first_draft_script", "")
        if not first_draft or (isinstance(first_draft, str) and not first_draft.strip()):
            return {
                "dialogue_polished_script": first_draft,
                "voice_consistency_notes": "ERROR: Received empty first_draft_script. Cannot proceed with dialogue polish.",
            }
        
        outputs = self.process(inputs)
        
        # Save state
        self.save_state(
            {
                "inputs": inputs,
                "outputs": outputs,
            },
            filename="polished_dialogue.json",
        )
        
        return outputs