narrative_classifier.py

import requests
import json
from typing import List, Dict

from taxonomy import STORY_ARCS, STORY_DEFINITION, get_taxonomy_formatted


class NarrativeClassifier:
    def __init__(self, api_key: str, group_id: str):
        self.api_key = api_key
        self.group_id = group_id
        self.base_url = "https://api.minimaxi.chat/v1"

    def _build_prompt(self, segments: List[Dict]) -> str:
        """Build classification prompt."""

        # Format segments
        segments_text = ""
        for seg in segments:
            segments_text += f"\n[{seg['start']:.1f}s - {seg['end']:.1f}s]"
            segments_text += f"\nVisual: {seg['visual']}"
            if seg['speech']:
                segments_text += f"\nSpeech: \"{seg['speech']}\""
            segments_text += "\n"

        # Format story arcs
        arcs_text = ""
        for arc_name, arc_info in STORY_ARCS.items():
            arcs_text += f"\n- {arc_name}: {' -> '.join(arc_info['sequence'])}"

        prompt = f"""You are an expert in advertising narrative structure analysis.

Analyze this video advertisement segment by segment.

## SEGMENTS TO ANALYZE:
{segments_text}

## FUNCTIONAL ROLE TAXONOMY:
{get_taxonomy_formatted()}

## KNOWN STORY ARCS:
{arcs_text}

## STORY DEFINITION:
{STORY_DEFINITION}

## YOUR TASK:

1. For each segment, determine the PRIMARY functional role from the taxonomy
2. Determine if this ad contains a STORY (YES/NO)
3. Identify which STORY ARC best matches (or "Custom" if none match)
4. List any MISSING elements that could strengthen the ad

## RESPONSE FORMAT (use exactly this JSON structure):

```json
{{
    "segments": [
        {{
            "timestamp": "0.0-2.0s",
            "functional_role": "Hook",
            "role_category": "OPENING",
            "reasoning": "Opens with provocative question to grab attention"
        }}
    ],
    "has_story": true,
    "story_explanation": "Brief explanation of why story is present/absent",
    "story_arc": "Problem-Solution-Outcome",
    "detected_sequence": ["Hook", "Problem Setup", "Solution Reveal", "Call-to-Action"],
    "missing_elements": ["Social Proof", "Outcome"]
}}
```

Respond ONLY with valid JSON, no other text."""

        return prompt

    def classify(self, segments: List[Dict]) -> Dict:
        """
        Classify each segment and detect overall story arc.

        Returns:
            {
                "segments": [...],
                "has_story": True/False,
                "story_arc": "...",
                "detected_sequence": [...],
                "missing_elements": [...],
                "raw_response": "..."
            }
        """
        url = f"{self.base_url}/text/chatcompletion_v2"

        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }

        prompt = self._build_prompt(segments)

        payload = {
            "model": "MiniMax-Text-01",
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.3  # Lower temperature for more consistent classification
        }

        response = requests.post(url, headers=headers, json=payload)

        if response.status_code != 200:
            print(f"Classification API error: {response.text}")
            return self._fallback_result(segments)

        result = response.json()
        raw_response = result['choices'][0]['message']['content']

        # Parse JSON from response
        try:
            # Extract JSON from response (may be wrapped in markdown code block)
            json_str = raw_response
            if "```json" in json_str:
                json_str = json_str.split("```json")[1].split("```")[0]
            elif "```" in json_str:
                json_str = json_str.split("```")[1].split("```")[0]

            parsed = json.loads(json_str.strip())

            # Merge with original segment data
            for i, seg_analysis in enumerate(parsed.get('segments', [])):
                if i < len(segments):
                    segments[i]['functional_role'] = seg_analysis.get('functional_role', 'Unknown')
                    segments[i]['role_category'] = seg_analysis.get('role_category', 'OTHER')
                    segments[i]['reasoning'] = seg_analysis.get('reasoning', '')

            return {
                "segments": segments,
                "has_story": parsed.get('has_story', False),
                "story_explanation": parsed.get('story_explanation', ''),
                "story_arc": parsed.get('story_arc', 'Unknown'),
                "detected_sequence": parsed.get('detected_sequence', []),
                "missing_elements": parsed.get('missing_elements', []),
                "raw_response": raw_response
            }

        except json.JSONDecodeError as e:
            print(f"JSON parse error: {e}")
            print(f"Raw response: {raw_response}")
            return self._fallback_result(segments, raw_response)

    def _fallback_result(self, segments: List[Dict], raw_response: str = "") -> Dict:
        """Return fallback result when parsing fails."""
        for seg in segments:
            seg['functional_role'] = 'Unknown'
            seg['role_category'] = 'OTHER'
            seg['reasoning'] = 'Classification failed'

        return {
            "segments": segments,
            "has_story": False,
            "story_explanation": "Unable to determine",
            "story_arc": "Unknown",
            "detected_sequence": [],
            "missing_elements": [],
            "raw_response": raw_response
        }