extract_questions.py

#!/usr/bin/env python3
"""
Question List Extractor

This script extracts questions from QA JSON files and creates a consolidated question bank
that can be used with the S3 PDF catalog.
"""

import json
import glob
import os
from datetime import datetime
from collections import defaultdict

def extract_questions_from_qa_files():
    """Extract all questions from QA JSON files"""
    print("📋 Extracting questions from QA JSON files...")
    
    qa_files = glob.glob("qa_results/*.json")
    question_bank = {}
    all_questions = []
    
    for json_file in qa_files:
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # Get document info
            doc_info = data.get('document_info', {})
            session_info = data.get('session_info', {})
            
            filename = doc_info.get('filename') or session_info.get('document_filename', '')
            
            # Extract questions and answers
            qa_pairs = data.get('qa_pairs', [])
            
            if qa_pairs and filename:
                questions_for_doc = []
                
                for qa in qa_pairs:
                    question_data = {
                        'id': qa.get('question_id'),
                        'question': qa.get('question'),
                        'answer': qa.get('answer', ''),
                        'source_document': filename
                    }
                    questions_for_doc.append(question_data)
                    all_questions.append(question_data)
                
                question_bank[filename] = {
                    'source_json': json_file,
                    'document_info': doc_info,
                    'total_questions': len(questions_for_doc),
                    'questions': questions_for_doc
                }
                
                print(f"   ✅ {filename}: {len(questions_for_doc)} questions")
            
        except Exception as e:
            print(f"   ❌ Error processing {json_file}: {e}")
    
    return question_bank, all_questions

def create_consolidated_question_list():
    """Create a consolidated question list from all sources"""
    question_bank, all_questions = extract_questions_from_qa_files()
    
    # Group questions by topic/category (basic keyword matching)
    topics = defaultdict(list)
    
    # Define topic keywords
    topic_keywords = {
        'Premium & Payment': ['premium', 'payment', 'grace period', 'installment'],
        'Waiting Period': ['waiting period', 'waiting', 'pre-existing'],
        'Maternity': ['maternity', 'pregnancy', 'delivery', 'prenatal', 'postnatal'],
        'Surgery & Treatment': ['surgery', 'operation', 'treatment', 'cataract'],
        'Coverage & Benefits': ['coverage', 'benefit', 'claim', 'reimbursement'],
        'Hospital & Medical': ['hospital', 'medical', 'doctor', 'nursing'],
        'Policy Terms': ['policy', 'renewal', 'cancellation', 'terms'],
        'Exclusions': ['exclusion', 'not covered', 'excluded'],
        'AYUSH': ['ayush', 'ayurveda', 'homeopathy', 'unani', 'siddha'],
        'Organ Donation': ['organ', 'donor', 'transplant'],
        'Health Checkup': ['health check', 'preventive', 'checkup'],
        'Discounts': ['discount', 'ncd', 'no claim bonus']
    }
    
    # Categorize questions
    for question_data in all_questions:
        question_text = question_data['question'].lower()
        categorized = False
        
        for topic, keywords in topic_keywords.items():
            if any(keyword in question_text for keyword in keywords):
                topics[topic].append(question_data)
                categorized = True
                break
        
        if not categorized:
            topics['General'].append(question_data)
    
    # Create consolidated structure
    consolidated = {
        'extraction_info': {
            'timestamp': datetime.now().isoformat(),
            'total_documents': len(question_bank),
            'total_questions': len(all_questions),
            'source_files': list(question_bank.keys())
        },
        'by_document': question_bank,
        'by_topic': dict(topics),
        'all_questions': all_questions,
        'question_templates': generate_question_templates(all_questions)
    }
    
    return consolidated

def generate_question_templates(all_questions):
    """Generate reusable question templates"""
    templates = {
        'premium_payment': [
            "What is the grace period for premium payment?",
            "How should premiums be paid?",
            "What happens if premium payment is delayed?"
        ],
        'waiting_period': [
            "What is the waiting period for pre-existing diseases?",
            "What is the waiting period for [specific condition]?",
            "Are there any waiting period waivers?"
        ],
        'maternity': [
            "Does this policy cover maternity expenses?",
            "What are the conditions for maternity coverage?",
            "What is the waiting period for maternity benefits?"
        ],
        'coverage': [
            "What is covered under this policy?",
            "What are the exclusions?",
            "What is the sum insured amount?"
        ],
        'claims': [
            "How to file a claim?",
            "What documents are required for claims?",
            "What is the claim settlement process?"
        ]
    }
    
    return templates

def save_question_bank(consolidated_data, filename="question_bank.json"):
    """Save the consolidated question bank"""
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(consolidated_data, f, indent=2, ensure_ascii=False)
    
    print(f"💾 Question bank saved to: {filename}")
    return filename

def merge_with_s3_catalog(s3_catalog_file="s3_pdf_catalog.json", question_bank_file="question_bank.json"):
    """Merge question bank with S3 catalog"""
    try:
        # Load S3 catalog
        with open(s3_catalog_file, 'r', encoding='utf-8') as f:
            s3_catalog = json.load(f)
        
        # Load question bank
        with open(question_bank_file, 'r', encoding='utf-8') as f:
            question_bank = json.load(f)
        
        # Merge data
        for doc_path, doc_data in s3_catalog['documents'].items():
            filename = doc_data['file_info']['filename']
            
            # Try to find matching questions
            if filename in question_bank['by_document']:
                doc_data['questions_data'] = question_bank['by_document'][filename]
                print(f"   ✅ Merged questions for {filename}")
            else:
                # Try partial matching
                for qa_filename, qa_data in question_bank['by_document'].items():
                    if os.path.splitext(filename)[0] in os.path.splitext(qa_filename)[0]:
                        doc_data['questions_data'] = qa_data
                        print(f"   ✅ Merged questions for {filename} (partial match with {qa_filename})")
                        break
        
        # Add question bank to catalog
        s3_catalog['question_bank'] = question_bank
        
        # Save merged catalog
        merged_filename = "s3_pdf_catalog_with_questions.json"
        with open(merged_filename, 'w', encoding='utf-8') as f:
            json.dump(s3_catalog, f, indent=2, ensure_ascii=False)
        
        print(f"🔗 Merged catalog saved to: {merged_filename}")
        return merged_filename
        
    except FileNotFoundError as e:
        print(f"❌ File not found: {e}")
        return None

def print_summary(consolidated_data):
    """Print summary of extracted questions"""
    print("\n" + "="*50)
    print("📊 QUESTION EXTRACTION SUMMARY")
    print("="*50)
    
    info = consolidated_data['extraction_info']
    print(f"📁 Documents processed: {info['total_documents']}")
    print(f"❓ Total questions: {info['total_questions']}")
    
    print(f"\n📂 Questions by document:")
    for doc, data in consolidated_data['by_document'].items():
        print(f"   • {doc}: {data['total_questions']} questions")
    
    print(f"\n🏷️ Questions by topic:")
    for topic, questions in consolidated_data['by_topic'].items():
        print(f"   • {topic}: {len(questions)} questions")
    
    print(f"\n📝 Sample questions:")
    for i, q in enumerate(consolidated_data['all_questions'][:5]):
        print(f"   {i+1}. {q['question'][:80]}...")

def main():
    """Main function"""
    print("🚀 Question List Extractor")
    print("="*40)
    
    # Extract and consolidate questions
    consolidated_data = create_consolidated_question_list()
    
    # Save question bank
    question_bank_file = save_question_bank(consolidated_data)
    
    # Try to merge with S3 catalog if it exists
    if os.path.exists("s3_pdf_catalog.json"):
        print("\n🔗 Merging with S3 catalog...")
        merge_with_s3_catalog()
    else:
        print("\n💡 Run s3_pdf_manager.py first to create S3 catalog for merging")
    
    # Print summary
    print_summary(consolidated_data)
    
    print(f"\n✅ Question extraction completed!")
    print(f"📄 Check {question_bank_file} for complete question bank")

if __name__ == "__main__":
    main()