File size: 5,493 Bytes
cfe83fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
from __future__ import annotations
from typing import Any, List, Optional
from pydantic import BaseModel

from viral_script_engine.environment.observations import DebateRound, RewardComponents

_SEVERITY_ORDER = {"high": 3, "medium": 2, "low": 1}

_CRITIQUE_TO_ACTION = {
    "hook_weakness": "hook_rewrite",
    "pacing_issue": "section_reorder",
    "cultural_mismatch": "cultural_ref_sub",
    "cta_buried": "cta_placement",
    "coherence_break": "section_reorder",
    "retention_risk": "hook_rewrite",
}

_ACTION_TO_TARGET = {
    "hook_rewrite": "hook",
    "section_reorder": "body",
    "cultural_ref_sub": "body",
    "cta_placement": "cta",
}


class TrajectoryType:
    CRITIC_FIRST = "critic_first"      # Trajectory A: act on Critic's top claim first
    DEFENDER_FIRST = "defender_first"  # Trajectory B: act on Defender's concern first


class Trajectory(BaseModel):
    trajectory_id: str
    trajectory_type: str
    initial_script: str
    current_script: str
    steps: List[Any] = []
    cumulative_reward: float = 0.0
    final_reward_components: Optional[Any] = None
    terminated: bool = False
    step_count: int = 0

    def get_forced_first_action(
        self,
        critic_claims: List[Any],
        defender_output: Any,
    ) -> dict:
        """
        Returns the forced first action based on trajectory type.

        CRITIC_FIRST: pick the action that addresses the highest-severity CritiqueClaim.
        DEFENDER_FIRST: pick the action that preserves the core_strength_quote.
            If core_strength is in hook → hook_rewrite is risky → pick cta_placement first.
        """
        if self.trajectory_type == TrajectoryType.CRITIC_FIRST:
            return self._critic_first_action(critic_claims)
        return self._defender_first_action(critic_claims, defender_output)

    def _critic_first_action(self, critic_claims: List[Any]) -> dict:
        if not critic_claims:
            return _fallback_action("C1")
        sorted_claims = sorted(
            critic_claims,
            key=lambda c: _SEVERITY_ORDER.get(getattr(c, "severity", "low"), 0),
            reverse=True,
        )
        top = sorted_claims[0]
        action_type = _CRITIQUE_TO_ACTION.get(
            getattr(top, "critique_class", ""), "hook_rewrite"
        )
        return {
            "action_type": action_type,
            "target_section": _ACTION_TO_TARGET.get(action_type, "hook"),
            "instruction": (
                f"Address the top critic concern: "
                f"{getattr(top, 'claim_text', '')[:100]}"
            ),
            "critique_claim_id": getattr(top, "claim_id", "C1"),
            "reasoning": (
                f"CRITIC_FIRST: targeting highest-severity "
                f"{getattr(top, 'critique_class', '')} claim ({getattr(top, 'severity', '')})."
            ),
        }

    def _defender_first_action(self, critic_claims: List[Any], defender_output: Any) -> dict:
        core_quote = ""
        flagged: set = set()

        if defender_output is not None:
            if hasattr(defender_output, "core_strength_quote"):
                core_quote = defender_output.core_strength_quote or ""
                flagged = set(getattr(defender_output, "flagged_critic_claims", []))
            elif isinstance(defender_output, dict):
                core_quote = defender_output.get("core_strength_quote", "")
                flagged = set(defender_output.get("flagged_critic_claims", []))

        # Core strength is "in the hook" if its first 20 chars appear in the leading 100 chars
        hook_portion = self.current_script[:100].lower()
        core_in_hook = bool(core_quote) and core_quote.lower()[:20] in hook_portion

        if core_in_hook:
            # Hook is precious — choose a safe non-hook action first
            action_type = "cta_placement"
            target = "cta"
            instruction = (
                "Improve CTA positioning to boost completion rate "
                "without altering the hook."
            )
            claim_id = (
                getattr(critic_claims[0], "claim_id", "C1")
                if critic_claims else "C1"
            )
        else:
            # Core is in body — safe to improve the hook
            action_type = "hook_rewrite"
            target = "hook"
            instruction = (
                "Rewrite the hook for stronger attention capture "
                "while preserving the core body voice."
            )
            unflagged = [
                c for c in critic_claims
                if getattr(c, "claim_id", "") not in flagged
            ]
            claim = unflagged[0] if unflagged else (critic_claims[0] if critic_claims else None)
            claim_id = getattr(claim, "claim_id", "C1") if claim else "C1"

        return {
            "action_type": action_type,
            "target_section": target,
            "instruction": instruction,
            "critique_claim_id": claim_id,
            "reasoning": (
                "DEFENDER_FIRST: preserving Defender's core strength "
                "and regional voice before addressing critic claims."
            ),
        }


def _fallback_action(claim_id: str = "C1") -> dict:
    return {
        "action_type": "hook_rewrite",
        "target_section": "hook",
        "instruction": "Rewrite the hook to open with a strong immediate claim.",
        "critique_claim_id": claim_id,
        "reasoning": "Fallback: no critic claims available.",
    }