run_detection.py

"""
Batch Detection Script - Multi-Agent Debate Detection

Independently runs multi-agent debate detection workflow.
Saves after processing each item to improve fault tolerance.

Usage:
    python run_detection.py --input_file results/generated.json --output_file results/detected.json
"""

import json
import os
import argparse
import time
from typing import List, Dict
from agents import DebateDetector
from config import DEFAULT_LLM_MODEL


def load_data(input_file: str, resume_index: int = 0) -> List[Dict]:
    """Load input data"""
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    if resume_index > 0:
        print(f"Resuming from index {resume_index}")
        data = data[resume_index:]
    
    return data


def load_existing_results(output_file: str) -> List[Dict]:
    """Load existing results"""
    if os.path.exists(output_file) and os.path.getsize(output_file) > 0:
        try:
            with open(output_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError:
            print(f"Warning: Cannot parse {output_file}, will create new file")
            return []
    return []


def save_result(output_file: str, results: List[Dict]):
    """Save results to file"""
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)


def run_detection_with_retry(news_content: str, llm_model: str, max_retries: int = 3):
    """
    Detection with retry mechanism
    
    Args:
        news_content: News content
        llm_model: LLM model to use
        max_retries: Maximum retry attempts
    
    Returns:
        (debate_record, judgement) or (None, None) if failed
    """
    for attempt in range(max_retries):
        try:
            print(f"  [Attempt {attempt + 1}/{max_retries}] Running debate detection...")
            
            detector = DebateDetector(
                news_topic=news_content,
                llm_model_names=[llm_model]
            )
            
            debate_record, judgement = detector.run()
            
            if judgement:  # If successfully obtained judgement
                print(f"  [Success] Judgement: {judgement}")
                return debate_record, judgement
            else:
                print(f"  [Failed] Could not obtain valid judgement")
                if attempt < max_retries - 1:
                    print(f"  [Waiting] Retrying in 5 seconds...")
                    time.sleep(5)
                    
        except Exception as e:
            print(f"  [Error] Attempt {attempt + 1} failed: {str(e)}")
            if attempt < max_retries - 1:
                print(f"  [Waiting] Retrying in 5 seconds...")
                time.sleep(5)
    
    return None, None


def run_detection(
    input_file: str,
    output_file: str,
    llm_model: str,
    use_optimized: bool,
    max_retries: int,
    resume_index: int = 0,
    max_samples: int = None
):
    """
    Run batch detection workflow
    
    Args:
        input_file: Input file path (can be original data or generated data)
        output_file: Output file path
        llm_model: Model for detector
        use_optimized: Whether to use optimized content (if input is generated data)
        max_retries: Maximum retry attempts per sample
        resume_index: Resume index
        max_samples: Maximum number of samples to process
    """
    print("="*60)
    print("SALF Batch Detection Script")
    print("="*60)
    
    # Load data
    print(f"\nLoading data: {input_file}")
    input_data = load_data(input_file, resume_index)
    
    if max_samples:
        input_data = input_data[:max_samples]
        print(f"Limiting samples to: {max_samples}")
    
    print(f"Samples to process: {len(input_data)}")
    print(f"Using optimized content: {use_optimized}")
    
    # Load existing results
    results = load_existing_results(output_file)
    print(f"Existing results: {len(results)}")
    
    # Process each item
    for idx, item in enumerate(input_data):
        actual_idx = idx + resume_index
        
        print(f"\n{'='*60}")
        print(f"Processing sample {actual_idx + 1}/{resume_index + len(input_data)}")
        print(f"{'='*60}")
        
        # Choose original or optimized content
        if use_optimized and 'optimized_content' in item:
            news_content = item['optimized_content']
            print(f"[Using] Optimized content")
        else:
            news_content = item.get('content', '')
            print(f"[Using] Original content")
        
        label = item.get('label', 1)
        
        # If real news, can choose to skip or still detect
        if label == 0:
            print(f"[Info] Sample {actual_idx} is real news")
        
        try:
            # Run detection (with retry)
            debate_record, judgement = run_detection_with_retry(
                news_content,
                llm_model,
                max_retries
            )
            
            if debate_record is None:
                print(f"[Failed] All retries failed, skipping this sample")
                result = {
                    'content': item.get('content', ''),
                    'label': label,
                    'debate_record': {},
                    'judgement': '',
                    'detection_success': False,
                    'error': 'All retries failed'
                }
            else:
                result = {
                    'content': item.get('content', ''),
                    'label': label,
                    'debate_record': debate_record,
                    'judgement': judgement,
                    'detection_success': True
                }
            
            # Preserve other fields from original data
            result.update({k: v for k, v in item.items() if k not in result})
            
            results.append(result)
            
            # Save after each item
            save_result(output_file, results)
            print(f"\n[Saved] Results saved to {output_file}")
            
        except Exception as e:
            print(f"\n[Error] Error processing sample {actual_idx}: {str(e)}")
            import traceback
            traceback.print_exc()
            
            # Save progress even on error
            result = {
                'content': item.get('content', ''),
                'label': label,
                'debate_record': {},
                'judgement': '',
                'detection_success': False,
                'error': str(e)
            }
            result.update({k: v for k, v in item.items() if k not in result})
            results.append(result)
            save_result(output_file, results)
            
            print(f"[Saved] Error record saved, continuing to next item")
            continue
    
    print(f"\n{'='*60}")
    print(f"Batch detection complete!")
    print(f"Total processed: {len(results)} items")
    print(f"Results saved to: {output_file}")
    print(f"{'='*60}")


def main():
    parser = argparse.ArgumentParser(description='SALF Batch Detection Script')
    
    parser.add_argument('--input_file', type=str, required=True,
                       help='Input data file (JSON format)')
    parser.add_argument('--output_file', type=str, required=True,
                       help='Output result file')
    parser.add_argument('--llm_model', type=str, default=DEFAULT_LLM_MODEL,
                       help='LLM model for detector')
    parser.add_argument('--use_optimized', action='store_true',
                       help='Use optimized content (if input file contains optimized_content field)')
    parser.add_argument('--max_retries', type=int, default=3,
                       help='Maximum retry attempts per sample')
    parser.add_argument('--resume_index', type=int, default=0,
                       help='Resume from specified index')
    parser.add_argument('--max_samples', type=int, default=None,
                       help='Maximum number of samples to process (for testing)')
    
    args = parser.parse_args()
    
    run_detection(
        input_file=args.input_file,
        output_file=args.output_file,
        llm_model=args.llm_model,
        use_optimized=args.use_optimized,
        max_retries=args.max_retries,
        resume_index=args.resume_index,
        max_samples=args.max_samples
    )


if __name__ == "__main__":
    main()