absolute_zero_reasoner/testtime/task_generator.py

"""
TestTime Task Generator

AZR 추론용 프롬프트 기반 Induction/Deduction/Abduction 태스크 생성
요구사항 3: "AZR처럼 템플릿을 활용하여 induction, deduction, abduction 문제를 생성"
"""

from typing import Dict, List, Any, Optional, Tuple
import random

from .config import TestTimeConfig
from .logger import TestTimeLogger
# AZR 추론용 프롬프트 직접 사용
from ..data_construction.prompts import get_code_problem_predictor_prompt
from .solution_generator import InitialSolutionGenerator


class TestTimeTaskGenerator:
    """IPO 트리플에서 3종 태스크 생성"""
    
    def __init__(self, config: TestTimeConfig, logger: Optional[TestTimeLogger] = None):
        self.config = config
        self.logger = logger or TestTimeLogger()
        
        # AZR 추론용 프롬프트 직접 사용 (get_code_problem_predictor_prompt)
        # 함수 코드 정리용 solution generator 인스턴스 생성
        self.solution_generator = InitialSolutionGenerator(None, None, config, logger)
    
    def generate_tasks(self, ipo_triples: List[Dict[str, Any]], 
                      problem_id: str, round_num: int = 1) -> Dict[str, List[Dict[str, Any]]]:
        """IPO 트리플에서 3종 태스크 생성 (각 트리플마다 3가지 태스크 모두 생성)"""
        
        self.logger.log_info(f"🎯 Generating tasks for {problem_id} from {len(ipo_triples)} triples")
        
        # 🔧 수정: 분배 로직 제거, 각 IPO 트리플에서 3가지 태스크 모두 생성
        induction_tasks = []
        deduction_tasks = []
        abduction_tasks = []
        
        for i, triple in enumerate(ipo_triples):
            # 각 트리플에서 induction 태스크 생성
            induction_task = self._generate_single_induction_task(triple, i, problem_id, round_num)
            if induction_task:
                induction_tasks.append(induction_task)
            
            # 각 트리플에서 deduction 태스크 생성  
            deduction_task = self._generate_single_deduction_task(triple, i, problem_id, round_num)
            if deduction_task:
                deduction_tasks.append(deduction_task)
                
            # 각 트리플에서 abduction 태스크 생성
            abduction_task = self._generate_single_abduction_task(triple, i, problem_id, round_num)
            if abduction_task:
                abduction_tasks.append(abduction_task)
        
        all_tasks = {
            'induction': induction_tasks,
            'deduction': deduction_tasks,
            'abduction': abduction_tasks
        }
        
        # 로깅
        task_counts = {k: len(v) for k, v in all_tasks.items()}
        total_generated = sum(task_counts.values())
        
        self.logger.log_info(f"✅ Generated {len(induction_tasks)} induction, {len(deduction_tasks)} deduction, {len(abduction_tasks)} abduction tasks")
        
        self.logger.log_task_generation(
            problem_id, 
            induction_tasks, 
            deduction_tasks, 
            abduction_tasks
        )
        
        return all_tasks
    
    def _generate_single_induction_task(self, triple: Dict[str, Any], index: int, problem_id: str, round_num: int) -> Optional[Dict[str, Any]]:
        """단일 IPO 트리플에서 induction 태스크 생성"""
        
        try:
            # 입력-출력 쌍 준비
            # 평가를 위해서는 실제 인자(triple['input'])를 사용
            input_output_pairs = [(triple['input'], triple['actual_output'])]
            
            # 표시용으로는 full_input_str 사용
            display_input = triple.get('full_input_str', triple['input'])
            
            # 🔧 수정: clean한 함수 코드만 추출 (test case 제거)
            clean_program = self._extract_clean_function_code(triple['program'])
            
            # 매개변수로 받은 problem_id 사용 (AZR 통합용)
            original_problem_id = triple.get('id', '').split('_triple_')[0]  # 원본 추출 로직 보존
            
            # HumanEval인 경우 특별 처리
            if 'HumanEval' in problem_id:
                # 원본 프로그램에서 함수 설명 추출 (doctest 예시가 있는 원본에서)
                extracted_message = self._extract_function_description(triple['program'])
                if not extracted_message:
                    extracted_message = "Find the function that produces these outputs from these inputs."
            else:
                # MBPP는 기존 방식 유지
                extracted_message = InitialSolutionGenerator.extract_docstring_from_function(clean_program)
            
            # 사용자 정의: input_output_pairs + message → program
            # 프롬프트용으로는 display 입력 사용
            display_pairs = [(display_input, triple['actual_output'])]
            azr_prompt = get_code_problem_predictor_prompt(
                problem_type='code_f',
                snippet=clean_program,  # 🔧 수정: clean한 코드 사용
                input_output_pairs=display_pairs,
                message=extracted_message
            )
            
            # AZR 메타데이터 생성
            source_program_id = triple.get('source_program_id', f'program_{index//3}')
            ipo_index = triple.get('ipo_index', index % 3)
            
            task = {
                'task_id': f'induction_{index}',
                'task_type': 'induction',
                'triple_id': triple['id'],
                'source_program_id': source_program_id,  # 🆕 추가
                'ipo_index': ipo_index,                  # 🆕 추가
                'ipo_triple': {                          # 🆕 추가
                    'input': triple['input'],
                    'output': triple['actual_output'], 
                    'program': triple['program']
                },
                'prompt': azr_prompt,
                'expected_solution': clean_program,  # 🔧 수정: clean한 코드 사용
                'evaluation_data': {
                    'input_output_pairs': input_output_pairs,  # 평가용으로는 실제 인자 사용
                    'original_function': triple['program']
                },
                
                # 🆕 AZR 학습용 메타데이터
                'uid': f"{problem_id}_round_{round_num}_induction_{index}",
                'ipo_group_id': f"{problem_id}_program_{source_program_id}_ipo_{ipo_index}",
                'original_problem_id': problem_id,
                'round': round_num,
                'extra_info': {'metric': 'code_f'},  # induction task는 code_f
                'basic_accuracy': 0.0,  # 초기값, evaluation에서 업데이트됨
                'ground_truth': clean_program  # AZR parquet 형식에서 사용
            }
            
            return task
            
        except Exception as e:
            self.logger.log_error(f"Failed to generate induction task for triple {triple.get('id', 'unknown')}: {e}")
            return None
    
    def _generate_single_deduction_task(self, triple: Dict[str, Any], index: int, problem_id: str, round_num: int) -> Optional[Dict[str, Any]]:
        """단일 IPO 트리플에서 deduction 태스크 생성"""
        
        try:
            # 매개변수로 받은 problem_id 사용 (AZR 통합용)
            original_problem_id = triple.get('id', '').split('_triple_')[0]  # 원본 추출 로직 보존
            
            # HumanEval인 경우 doctest 예시 제거
            if 'HumanEval' in original_problem_id:
                clean_program = self._remove_doctest_examples(triple['program'])
            else:
                # MBPP는 기존 방식 유지
                clean_program = self._extract_clean_function_code(triple['program'])
            
            # 사용자 정의: program + input → output  
            azr_prompt = get_code_problem_predictor_prompt(
                problem_type='code_o',  # 프로그램+입력→출력
                snippet=clean_program,  # 🔧 수정: clean한 코드 사용
                input_args=triple['input']
            )
            
            # AZR 메타데이터 생성
            source_program_id = triple.get('source_program_id', f'program_{index//3}')
            ipo_index = triple.get('ipo_index', index % 3)
            
            task = {
                'task_id': f'deduction_{index}',
                'task_type': 'deduction',
                'triple_id': triple['id'],
                'source_program_id': source_program_id,  # 🆕 추가
                'ipo_index': ipo_index,                  # 🆕 추가
                'ipo_triple': {                          # 🆕 추가
                    'input': triple['input'],
                    'output': triple['actual_output'], 
                    'program': triple['program']
                },
                'prompt': azr_prompt,
                'expected_solution': triple['actual_output'],  # 🔧 수정: expected_solution으로 통일
                'evaluation_data': {
                    'function_code': clean_program,  # 🔧 수정: clean한 코드 사용 (complete_pipeline과 일치)
                    'test_input': triple['input'],  # 🔧 수정: complete_pipeline과 일치
                    'original_function': triple['program']
                },
                
                # 🆕 AZR 학습용 메타데이터
                'uid': f"{problem_id}_round_{round_num}_deduction_{index}",
                'ipo_group_id': f"{problem_id}_program_{source_program_id}_ipo_{ipo_index}",
                'original_problem_id': problem_id,
                'round': round_num,
                'extra_info': {'metric': 'code_o'},  # deduction task는 code_o
                'basic_accuracy': 0.0,  # 초기값, evaluation에서 업데이트됨
                'ground_truth': triple['actual_output']  # AZR parquet 형식에서 사용
            }
            
            return task
            
        except Exception as e:
            self.logger.log_error(f"Failed to generate deduction task for triple {triple.get('id', 'unknown')}: {e}")
            return None
    
    def _generate_single_abduction_task(self, triple: Dict[str, Any], index: int, problem_id: str, round_num: int) -> Optional[Dict[str, Any]]:
        """단일 IPO 트리플에서 abduction 태스크 생성"""
        
        try:
            # 매개변수로 받은 problem_id 사용 (AZR 통합용)
            original_problem_id = triple.get('id', '').split('_triple_')[0]  # 원본 추출 로직 보존
            
            # HumanEval인 경우 doctest 예시 제거
            if 'HumanEval' in original_problem_id:
                clean_program = self._remove_doctest_examples(triple['program'])
            else:
                # MBPP는 기존 방식 유지
                clean_program = self._extract_clean_function_code(triple['program'])
            
            # 사용자 정의: program + output → input
            azr_prompt = get_code_problem_predictor_prompt(
                problem_type='code_i',  # 프로그램+출력→입력
                snippet=clean_program,  # 🔧 수정: clean한 코드 사용
                output=triple['actual_output']  # 🔧 수정: output 파라미터 사용
            )
            
            # AZR 메타데이터 생성
            source_program_id = triple.get('source_program_id', f'program_{index//3}')
            ipo_index = triple.get('ipo_index', index % 3)
            
            task = {
                'task_id': f'abduction_{index}',
                'task_type': 'abduction',
                'triple_id': triple['id'],
                'source_program_id': source_program_id,  # 🆕 추가
                'ipo_index': ipo_index,                  # 🆕 추가
                'ipo_triple': {                          # 🆕 추가
                    'input': triple['input'],
                    'output': triple['actual_output'], 
                    'program': triple['program']
                },
                'prompt': azr_prompt,
                'expected_solution': triple.get('full_input_str', triple['input']),  # 🔧 수정: 전체 함수 호출 사용
                'evaluation_data': {
                    'function_code': clean_program,  # 🔧 수정: clean한 코드 사용 (complete_pipeline과 일치)
                    'expected_output': triple['actual_output'],  # 🔧 수정: complete_pipeline과 일치
                    'original_function': triple['program']
                },
                
                # 🆕 AZR 학습용 메타데이터
                'uid': f"{problem_id}_round_{round_num}_abduction_{index}",
                'ipo_group_id': f"{problem_id}_program_{source_program_id}_ipo_{ipo_index}",
                'original_problem_id': problem_id,
                'round': round_num,
                'extra_info': {'metric': 'code_i'},  # abduction task는 code_i
                'basic_accuracy': 0.0,  # 초기값, evaluation에서 업데이트됨
                'ground_truth': triple.get('full_input_str', triple['input'])  # AZR parquet 형식에서 사용
            }
            
            return task
            
        except Exception as e:
            self.logger.log_error(f"Failed to generate abduction task for triple {triple.get('id', 'unknown')}: {e}")
            return None
    
    def generate_induction_tasks(self, ipo_triples: List[Dict[str, Any]], 
                                count: int) -> List[Dict[str, Any]]:
        """Induction 태스크: 입력-출력 쌍에서 프로그램 추론 (사용자 정의 유지)"""
        
        tasks = []
        selected_triples = random.sample(ipo_triples, min(count, len(ipo_triples)))
        
        for i, triple in enumerate(selected_triples):
            # 입력-출력 쌍 준비
            input_output_pairs = [(triple['input'], triple['actual_output'])]
            
            # 🔧 수정: clean한 함수 코드만 추출 (test case 제거)
            clean_program = self._extract_clean_function_code(triple['program'])
            
            # LLM이 생성한 함수에서 docstring 추출해서 message로 사용
            extracted_message = InitialSolutionGenerator.extract_docstring_from_function(clean_program)
            
            # 사용자 정의: input_output_pairs + message → program
            azr_prompt = get_code_problem_predictor_prompt(
                problem_type='code_f',
                snippet=clean_program,  # 🔧 수정: clean한 코드 사용
                input_output_pairs=input_output_pairs,
                message=extracted_message
            )
            
            task = {
                'task_id': f'induction_{i}',
                'task_type': 'induction',
                'triple_id': triple['id'],
                'prompt': azr_prompt,
                'expected_solution': clean_program,  # 🔧 수정: clean한 코드 사용
                'evaluation_data': {
                    'input_output_pairs': input_output_pairs,
                    'original_function': triple['program']
                }
            }
            
            tasks.append(task)
        
        return tasks
    
    def generate_deduction_tasks(self, ipo_triples: List[Dict[str, Any]], 
                               count: int) -> List[Dict[str, Any]]:
        """Deduction 태스크: 프로그램+입력에서 출력 예측 (사용자 정의에 맞게 수정)"""
        
        tasks = []
        selected_triples = random.sample(ipo_triples, min(count, len(ipo_triples)))
        
        for i, triple in enumerate(selected_triples):
            # 🔧 수정: clean한 함수 코드만 추출 (test case 제거)
            clean_program = self._extract_clean_function_code(triple['program'])
            
            # 사용자 정의: program + input → output
            azr_prompt = get_code_problem_predictor_prompt(
                problem_type='code_o',  # 프로그램+입력→출력
                snippet=clean_program,  # 🔧 수정: clean한 코드 사용
                input_args=triple['input']
            )
            
            task = {
                'task_id': f'deduction_{i}',
                'task_type': 'deduction',
                'triple_id': triple['id'],
                'prompt': azr_prompt,
                'expected_solution': triple['actual_output'],
                'evaluation_data': {
                    'function_code': clean_program,  # 🔧 수정: clean한 코드 사용
                    'test_input': triple['input']
                }
            }
            
            tasks.append(task)
        
        return tasks
    
    def generate_abduction_tasks(self, ipo_triples: List[Dict[str, Any]], 
                               count: int) -> List[Dict[str, Any]]:
        """Abduction 태스크: 프로그램+출력에서 입력 예측 (사용자 정의에 맞게 수정)"""
        
        tasks = []
        selected_triples = random.sample(ipo_triples, min(count, len(ipo_triples)))
        
        for i, triple in enumerate(selected_triples):
            # 🔧 수정: clean한 함수 코드만 추출 (test case 제거)
            clean_program = self._extract_clean_function_code(triple['program'])
            
            # 사용자 정의: program + output → input
            azr_prompt = get_code_problem_predictor_prompt(
                problem_type='code_i',  # 프로그램+출력→입력
                snippet=clean_program,  # 🔧 수정: clean한 코드 사용
                output=triple['actual_output']
            )
            
            task = {
                'task_id': f'abduction_{i}',
                'task_type': 'abduction',
                'triple_id': triple['id'],
                'prompt': azr_prompt,
                'expected_solution': triple.get('full_input_str', triple['input']),  # 🔧 수정: 전체 함수 호출 사용
                'evaluation_data': {
                    'function_code': clean_program,  # 🔧 수정: clean한 코드 사용
                    'expected_output': triple['actual_output']
                }
            }
            
            tasks.append(task)
        
        return tasks
    
    def _extract_clean_function_code(self, program_with_tests: str) -> str:
        """🔧 수정: 프로그램에서 test case와 assert문을 제거하고 순수한 함수 코드만 추출"""
        
        # solution_generator의 _extract_function_code 메서드 사용
        clean_code = self.solution_generator._extract_function_code(program_with_tests)
        
        # 로깅 (디버깅용)
        if "assert" in program_with_tests or "# Test" in program_with_tests:
            self.logger.log_info("🧹 Cleaned function code (removed test cases)")
        
        return clean_code
    
    def get_task_statistics(self, all_tasks: Dict[str, List[Dict[str, Any]]]) -> Dict[str, Any]:
        """태스크 생성 통계"""
        
        stats = {
            'total_tasks': sum(len(tasks) for tasks in all_tasks.values()),
            'tasks_by_type': {task_type: len(tasks) for task_type, tasks in all_tasks.items()},
            'task_types': list(all_tasks.keys())
        }
        
        return stats
    
    def _remove_doctest_examples(self, code: str) -> str:
        """HumanEval 코드에서 doctest 예시 제거"""
        import re
        
        lines = code.split('\n')
        result_lines = []
        in_docstring = False
        docstring_indent = 0
        skip_next = False
        
        for line in lines:
            stripped = line.strip()
            
            # docstring 시작/끝 감지
            if '"""' in line or "'''" in line:
                if not in_docstring:
                    in_docstring = True
                    docstring_indent = len(line) - len(line.lstrip())
                    result_lines.append(line)
                else:
                    in_docstring = False
                    result_lines.append(line)
                continue
            
            # doctest 예시 라인 건너뛰기
            if in_docstring:
                if stripped.startswith('>>>'):
                    skip_next = True  # 다음 라인(결과)도 건너뛰기
                    continue
                elif skip_next and stripped and not stripped.startswith('>>>'):
                    skip_next = False
                    continue
                else:
                    skip_next = False
            
            result_lines.append(line)
        
        return '\n'.join(result_lines)
    
    def _extract_function_description(self, code: str) -> str:
        """docstring에서 함수 설명 추출 (예시 제외)"""
        import re
        
        # 여러 형태의 docstring 매칭
        patterns = [
            r'"""(.*?)"""', # triple double quotes
            r"'''(.*?)'''", # triple single quotes
        ]
        
        for pattern in patterns:
            match = re.search(pattern, code, re.DOTALL)
            if match:
                description = match.group(1).strip()
                # 예시 전까지의 모든 설명 추출
                result_lines = []
                lines = description.split('\n')
                for line in lines:
                    cleaned_line = line.strip()
                    # >>> 예시가 시작되면 중단
                    if cleaned_line.startswith('>>>'):
                        break
                    # 빈 줄이 아니고 예시가 아닌 경우 추가
                    if cleaned_line:
                        result_lines.append(cleaned_line)
                
                # 모든 설명 라인을 공백으로 연결
                if result_lines:
                    return ' '.join(result_lines)
        
        return ""