File size: 12,942 Bytes
d2b8632
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
"""
Iterative Guideline Refinement Loop for Solo Mode

Orchestrates the automated cycle:
  confusion analysis → guideline suggestions → prompt revision → re-annotation

Monitors agreement rate trends and stops cycling when metrics plateau.
"""

import logging
import threading
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class RefinementCycle:
    """Record of a single refinement cycle."""
    cycle_number: int
    started_at: str
    completed_at: Optional[str] = None
    agreement_rate_before: float = 0.0
    agreement_rate_after: Optional[float] = None
    improvement: Optional[float] = None
    patterns_found: int = 0
    suggestions_generated: int = 0
    rules_applied: int = 0
    reannotation_count: int = 0
    prompt_version_before: int = 0
    prompt_version_after: Optional[int] = None
    status: str = "running"  # running, completed, no_improvement, failed

    def to_dict(self) -> Dict[str, Any]:
        return {
            'cycle_number': self.cycle_number,
            'started_at': self.started_at,
            'completed_at': self.completed_at,
            'agreement_rate_before': self.agreement_rate_before,
            'agreement_rate_after': self.agreement_rate_after,
            'improvement': self.improvement,
            'patterns_found': self.patterns_found,
            'suggestions_generated': self.suggestions_generated,
            'rules_applied': self.rules_applied,
            'reannotation_count': self.reannotation_count,
            'prompt_version_before': self.prompt_version_before,
            'prompt_version_after': self.prompt_version_after,
            'status': self.status,
        }


class RefinementLoop:
    """Orchestrates iterative confusion analysis → prompt revision cycles.

    The loop monitors annotation progress and periodically:
    1. Analyzes confusion patterns from human-LLM disagreements
    2. Generates guideline suggestions for top confusion patterns
    3. Injects suggestions into the prompt (if auto_apply or human-approved)
    4. Triggers re-annotation of low-confidence instances
    5. Measures improvement and decides whether to continue

    The loop stops when:
    - Agreement rate meets the target threshold
    - Improvement plateaus (patience exceeded)
    - Maximum cycles reached
    """

    def __init__(self, solo_config: Any, app_config: Dict[str, Any]):
        self.solo_config = solo_config
        self.app_config = app_config
        self.rl_config = solo_config.refinement_loop

        # Cycle tracking
        self._cycles: List[RefinementCycle] = []
        self._annotations_since_last_check: int = 0
        self._consecutive_no_improvement: int = 0
        self._stopped: bool = False
        self._stop_reason: Optional[str] = None
        self._running: bool = False
        self._lock = threading.Lock()

    @property
    def cycle_count(self) -> int:
        return len(self._cycles)

    @property
    def is_stopped(self) -> bool:
        return self._stopped

    @property
    def stop_reason(self) -> Optional[str]:
        return self._stop_reason

    def record_annotation(self) -> bool:
        """Record that a human annotation was made.

        Returns:
            True if a refinement cycle should be triggered.
        """
        if not self.rl_config.enabled or self._stopped:
            return False

        with self._lock:
            self._annotations_since_last_check += 1
            if self._annotations_since_last_check >= self.rl_config.trigger_interval:
                # Reset counter immediately to prevent multiple triggers
                # before run_cycle starts
                self._annotations_since_last_check = 0
                if self._running:
                    return False  # Cycle already in progress
                return True
            return False

    def should_trigger(self) -> bool:
        """Check if conditions are met to trigger a refinement cycle."""
        if not self.rl_config.enabled or self._stopped or self._running:
            return False

        with self._lock:
            return self._annotations_since_last_check >= self.rl_config.trigger_interval

    def run_cycle(
        self,
        agreement_rate: float,
        prompt_version: int,
        confusion_patterns: List[Any],
        apply_suggestions_fn: Callable[[List[str]], Dict[str, Any]],
        generate_suggestion_fn: Callable[[Any, str], Optional[str]],
        current_prompt: str,
    ) -> RefinementCycle:
        """Execute one refinement cycle.

        Args:
            agreement_rate: Current agreement rate before the cycle.
            prompt_version: Current prompt version number.
            confusion_patterns: List of ConfusionPattern objects from analyzer.
            apply_suggestions_fn: Callable that takes a list of suggestion strings
                and applies them to the prompt. Returns dict with results.
            generate_suggestion_fn: Callable(pattern, current_prompt) -> suggestion.
            current_prompt: The current annotation prompt text.

        Returns:
            RefinementCycle record with results.
        """
        with self._lock:
            if self._running:
                raise RuntimeError("Refinement cycle already running")
            self._running = True
            self._annotations_since_last_check = 0

        cycle = RefinementCycle(
            cycle_number=self.cycle_count + 1,
            started_at=datetime.now().isoformat(),
            agreement_rate_before=agreement_rate,
            prompt_version_before=prompt_version,
            patterns_found=len(confusion_patterns),
        )

        try:
            # Check if max cycles exceeded
            if self.cycle_count >= self.rl_config.max_cycles:
                cycle.status = "max_cycles_reached"
                self._stop("Max refinement cycles reached")
                self._finalize_cycle(cycle)
                return cycle

            # Generate suggestions for top patterns
            suggestions = []
            for pattern in confusion_patterns:
                suggestion = generate_suggestion_fn(pattern, current_prompt)
                if suggestion:
                    suggestions.append(suggestion)

            cycle.suggestions_generated = len(suggestions)

            if not suggestions:
                cycle.status = "no_suggestions"
                self._finalize_cycle(cycle)
                return cycle

            # Apply suggestions
            if self.rl_config.auto_apply_suggestions:
                result = apply_suggestions_fn(suggestions)
                cycle.rules_applied = result.get('categories_incorporated', 0)
                cycle.reannotation_count = result.get('reannotation_count', 0)
                cycle.prompt_version_after = result.get('new_prompt_version')
                cycle.status = "completed"
            else:
                # Suggestions generated but await human approval
                cycle.status = "awaiting_approval"

            self._finalize_cycle(cycle)
            return cycle

        except Exception as e:
            logger.error(f"Refinement cycle {cycle.cycle_number} failed: {e}")
            cycle.status = "failed"
            self._finalize_cycle(cycle)
            return cycle

    def record_post_cycle_metrics(self, agreement_rate_after: float) -> None:
        """Record the agreement rate after a cycle completes and re-annotation settles.

        This should be called after enough new annotations have been collected
        to measure the effect of the refinement.

        Args:
            agreement_rate_after: Agreement rate measured after the cycle.
        """
        with self._lock:
            if not self._cycles:
                return

            last_cycle = self._cycles[-1]
            if last_cycle.agreement_rate_after is not None:
                return  # Already recorded

            last_cycle.agreement_rate_after = agreement_rate_after
            improvement = agreement_rate_after - last_cycle.agreement_rate_before
            last_cycle.improvement = round(improvement, 4)

            if improvement < self.rl_config.min_improvement:
                self._consecutive_no_improvement += 1
                logger.info(
                    f"Refinement cycle {last_cycle.cycle_number}: "
                    f"no significant improvement ({improvement:+.4f}), "
                    f"patience {self._consecutive_no_improvement}/{self.rl_config.patience}"
                )
                if self._consecutive_no_improvement >= self.rl_config.patience:
                    self._stop("Improvement plateaued")
            else:
                self._consecutive_no_improvement = 0
                logger.info(
                    f"Refinement cycle {last_cycle.cycle_number}: "
                    f"improvement {improvement:+.4f}"
                )

    def _finalize_cycle(self, cycle: RefinementCycle) -> None:
        """Finalize a cycle and store it."""
        cycle.completed_at = datetime.now().isoformat()
        with self._lock:
            self._cycles.append(cycle)
            self._running = False

    def _stop(self, reason: str) -> None:
        """Stop the refinement loop."""
        self._stopped = True
        self._stop_reason = reason
        logger.info(f"Refinement loop stopped: {reason}")

    def reset(self) -> None:
        """Reset the refinement loop state, allowing new cycles."""
        with self._lock:
            self._consecutive_no_improvement = 0
            self._stopped = False
            self._stop_reason = None
            self._annotations_since_last_check = 0

    def get_status(self) -> Dict[str, Any]:
        """Get the current refinement loop status."""
        with self._lock:
            last_cycle = self._cycles[-1].to_dict() if self._cycles else None
            last_improvement = None
            if self._cycles and self._cycles[-1].improvement is not None:
                last_improvement = self._cycles[-1].improvement

            return {
                'enabled': self.rl_config.enabled,
                'total_cycles': len(self._cycles),
                'is_running': self._running,
                'is_stopped': self._stopped,
                'stop_reason': self._stop_reason,
                'consecutive_no_improvement': self._consecutive_no_improvement,
                'patience': self.rl_config.patience,
                'max_cycles': self.rl_config.max_cycles,
                'trigger_interval': self.rl_config.trigger_interval,
                'annotations_until_next': max(
                    0,
                    self.rl_config.trigger_interval - self._annotations_since_last_check
                ),
                'last_cycle': last_cycle,
                'last_improvement': last_improvement,
                'cycles': [c.to_dict() for c in self._cycles],
            }

    def to_dict(self) -> Dict[str, Any]:
        """Serialize state for persistence."""
        return {
            'cycles': [c.to_dict() for c in self._cycles],
            'annotations_since_last_check': self._annotations_since_last_check,
            'consecutive_no_improvement': self._consecutive_no_improvement,
            'stopped': self._stopped,
            'stop_reason': self._stop_reason,
        }

    def load_state(self, data: Dict[str, Any]) -> None:
        """Restore state from persistence."""
        with self._lock:
            self._annotations_since_last_check = data.get(
                'annotations_since_last_check', 0
            )
            self._consecutive_no_improvement = data.get(
                'consecutive_no_improvement', 0
            )
            self._stopped = data.get('stopped', False)
            self._stop_reason = data.get('stop_reason')

            self._cycles = []
            for cd in data.get('cycles', []):
                self._cycles.append(RefinementCycle(
                    cycle_number=cd.get('cycle_number', 0),
                    started_at=cd.get('started_at', ''),
                    completed_at=cd.get('completed_at'),
                    agreement_rate_before=cd.get('agreement_rate_before', 0.0),
                    agreement_rate_after=cd.get('agreement_rate_after'),
                    improvement=cd.get('improvement'),
                    patterns_found=cd.get('patterns_found', 0),
                    suggestions_generated=cd.get('suggestions_generated', 0),
                    rules_applied=cd.get('rules_applied', 0),
                    reannotation_count=cd.get('reannotation_count', 0),
                    prompt_version_before=cd.get('prompt_version_before', 0),
                    prompt_version_after=cd.get('prompt_version_after'),
                    status=cd.get('status', 'completed'),
                ))