data_generator/theme_generator.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import re
import json
import requests
import threading
from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict, Tuple

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import api_key, base_url

# API configuration
API_KEY = api_key
API_PROVIDER = base_url

# Thread-safe print function
print_lock = threading.Lock()
def thread_safe_print(*args, **kwargs):
    with print_lock:
        print(*args, **kwargs)

def query_llm(prompt: str) -> str:
    """
    Query LLM API with a prompt
    Args:
        prompt: The prompt to send to LLM
    Returns:
        str: The response from LLM
    """
    headers = {
        'Authorization': f'Bearer {API_KEY}',
        'Content-Type': 'application/json'
    }
    
    data = {
        'model': 'gpt-5-mini',
        'messages': [
            {
                'role': 'system', 
                'content': 'You are a senior data journalist and infographic designer specialized in creating compelling data stories. Always return valid JSON format only, without any markdown formatting or extra text.'
            },
            {
                'role': 'user', 
                'content': prompt
            }
        ],
        'temperature': 0.7
    }
    
    try:
        response = requests.post(
            f'{API_PROVIDER}/chat/completions',
            headers=headers,
            json=data,
            timeout=120
        )
        response.raise_for_status()
        
        result = response.json()
        return result['choices'][0]['message']['content'].strip()
        
    except requests.exceptions.Timeout:
        thread_safe_print("❌ LLM API 超时（60秒）")
        return None
    except requests.exceptions.HTTPError as e:
        thread_safe_print(f"❌ LLM API HTTP 错误: {e}")
        if hasattr(e.response, 'text'):
            thread_safe_print(f"   响应: {e.response.text[:200]}")
        return None
    except requests.exceptions.RequestException as e:
        thread_safe_print(f"❌ LLM API 请求错误: {e}")
        return None
    except KeyError as e:
        thread_safe_print(f"❌ LLM API 响应格式错误: {e}")
        return None
    except Exception as e:
        thread_safe_print(f"❌ 查询 LLM 时出错: {e}")
        return None

def read_theme_file(file_path: str) -> Dict[str, List[Dict[str, str]]]:
    """
    Read the theme file and parse it into a dictionary with detailed themes
    Args:
        file_path: Path to the theme file
    Returns:
        Dict[str, List[Dict]]: Dictionary where keys are main theme names and values are lists of 
                              dictionaries containing specific themes with their number and text
    """
    themes = {}
    current_main_theme = None
    
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
                
            if line.startswith('#'):
                current_main_theme = line[1:].strip()
                themes[current_main_theme] = []
            elif current_main_theme and re.match(r'^\d+\.', line):
                # Extract the number and the specific theme content
                match = re.match(r'^(\d+)\.\s*(.*)', line)
                if match:
                    number = int(match.group(1))
                    specific_theme = match.group(2)
                    themes[current_main_theme].append({
                        "number": number,
                        "theme": specific_theme
                    })
    
    return themes

def generate_similar_themes(main_theme: str, specific_theme: str, count: int = 15) -> List[Dict]:
    """
    Generate similar themes to a specific theme and return in JSON format
    Args:
        main_theme: The main theme category
        specific_theme: The specific theme to generate similar themes for
        count: Number of similar themes to generate
    Returns:
        List[Dict]: List of new theme dictionaries with id, theme, and description
    """
    prompt = f"""
You are a senior data journalist and infographic designer. Generate {count} compelling, diverse data story themes inspired by this reference:

Main Category: {main_theme}
Reference Theme: {specific_theme}

Create themes that would make excellent real-world infographics with these characteristics:

**DIVERSITY REQUIREMENTS:**
- Mix different angles: comparisons, trends over time, geographic distributions, rankings, cause-effect relationships, surprising statistics, myth-busting facts
- Vary the scope: global, regional, national, city-level, industry-specific, demographic-specific
- Include different time frames: historical analysis, current snapshots, future projections
- Cover various data types: percentages, absolute numbers, ratios, growth rates, correlations

**QUALITY CRITERIA:**
1. Each theme should tell a compelling data story that surprises, educates, or reveals hidden patterns
2. Must be based on realistic, obtainable data (surveys, government statistics, research studies, industry reports)
3. Should have a clear "hook" - why would someone stop scrolling to look at this infographic?
4. Include specific, concrete angles (e.g., "How coffee consumption varies by profession" instead of generic "Coffee consumption trends")
5. Themes should evoke curiosity or challenge common assumptions
6. Consider timely topics, emerging trends, or evergreen insights

**THEME STYLES TO INCLUDE:**
- "Did you know..." style surprising statistics
- "The real cost of..." economic breakdowns
- "A day/year in the life of..." behavioral patterns
- "X vs Y: The ultimate comparison" head-to-head analysis
- "The rise and fall of..." historical trends
- "What [demographic] really thinks about..." opinion data
- "Behind the numbers of..." deep-dive analysis
- "The geography of..." spatial distributions
- "Before and after..." transformation stories

Return ONLY valid JSON in this exact format:
[
  {{
    "id": 1,
    "theme": "[Specific, engaging theme title that could be an infographic headline]",
    "description": "[One-sentence description of the data story and why it's interesting]"
  }},
  ...
]

Generate {count} DIVERSE themes - avoid repetitive patterns or similar angles. Each theme should feel fresh and distinct.
    """
    
    response = query_llm(prompt)
    if not response:
        return []
    
    # Parse the JSON response with robust cleaning
    try:
        # 清理可能的 markdown 代码块
        cleaned_response = response.strip()
        
        # 如果响应被包裹在代码块中
        if cleaned_response.startswith('```'):
            lines = cleaned_response.split('\n')
            # 移除第一行和最后一行的```
            if lines[-1].strip() == '```' or lines[-1].strip().startswith('```'):
                cleaned_response = '\n'.join(lines[1:-1])
            else:
                cleaned_response = '\n'.join(lines[1:])
            # 进一步清理
            cleaned_response = cleaned_response.replace('```json', '').replace('```', '').strip()
        
        # 尝试提取JSON数组
        json_match = re.search(r'(\[[\s\S]*\])', cleaned_response)
        if json_match:
            json_content = json_match.group(1)
        else:
            json_content = cleaned_response
        
        # 解析JSON
        themes_data = json.loads(json_content)
        
        # 验证返回的数据结构
        if isinstance(themes_data, list) and len(themes_data) > 0:
            # 验证每个主题是否有必需的字段
            valid_themes = []
            for theme in themes_data:
                if isinstance(theme, dict) and 'theme' in theme and 'description' in theme:
                    valid_themes.append(theme)
            
            if valid_themes:
                return valid_themes
            else:
                thread_safe_print(f"⚠️  主题 '{specific_theme}' 的响应缺少必需字段")
                return []
        else:
            thread_safe_print(f"⚠️  主题 '{specific_theme}' 的响应不是有效的列表")
            return []
            
    except json.JSONDecodeError as e:
        thread_safe_print(f"❌ 解析 JSON 失败，主题 '{specific_theme}': {e}")
        thread_safe_print(f"   响应内容（前500字符）: {response[:500]}")
        return []
    except Exception as e:
        thread_safe_print(f"❌ 处理响应时出错，主题 '{specific_theme}': {e}")
        return []

def process_specific_theme(main_theme: str, specific_theme_data: Dict, all_results: List[Dict]) -> None:
    """
    Process a specific theme and add generated similar themes to the results
    Args:
        main_theme: The main theme category
        specific_theme_data: Dictionary with number and theme content
        all_results: List to store all results
    """
    specific_theme = specific_theme_data["theme"]
    original_number = specific_theme_data["number"]
    
    thread_safe_print(f"\n{'='*80}")
    thread_safe_print(f"🔄 正在处理主题")
    thread_safe_print(f"   分类: {main_theme}")
    thread_safe_print(f"   主题: {specific_theme}")
    thread_safe_print(f"   编号: {original_number}")
    
    # First add the original theme as the first entry
    with print_lock:
        original_theme_entry = {
            "id": len(all_results) + 1,
            "theme": specific_theme,
            "description": f"Original theme {original_number} from {main_theme} category",
            "main_category": main_theme,
            "is_original": True,
            "original_number": original_number
        }
        all_results.append(original_theme_entry)
    
    thread_safe_print(f"🤖 调用 LLM 生成相似主题...")
    
    # Generate similar themes
    similar_themes = generate_similar_themes(main_theme, specific_theme)
    
    if similar_themes:
        # Add main category and reference to original theme
        for theme in similar_themes:
            theme["main_category"] = main_theme
            theme["is_original"] = False
            theme["related_to_original"] = original_number
        
        with print_lock:
            all_results.extend(similar_themes)
        
        thread_safe_print(f"✅ 成功生成 {len(similar_themes)} 个相似主题")
        thread_safe_print(f"   总进度: {len(all_results)} 个主题已生成")
    else:
        thread_safe_print(f"❌ 生成主题失败: '{specific_theme}'")

def main():
    import time
    
    theme_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "theme.txt")
    output_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "theme_new.json")
    
    print("="*80)
    print("🚀 开始主题生成流程")
    print("="*80)
    print(f"📖 读取主题文件: {theme_file}")
    print(f"💾 输出文件: {output_file}")
    
    # Read the original theme file
    themes = read_theme_file(theme_file)
    
    total_specific_themes = sum(len(specific_themes) for specific_themes in themes.values())
    print(f"✅ 成功读取 {len(themes)} 个主分类，共 {total_specific_themes} 个具体主题")
    print(f"🔧 并行线程数: 4")
    print("="*80)
    
    # Initialize results list for all themes
    all_results = []
    
    start_time = time.time()
    
    # Use a thread pool to process specific themes in parallel
    with ThreadPoolExecutor(max_workers=8) as executor:
        futures = []
        
        for main_theme, specific_themes in themes.items():
            for specific_theme_data in specific_themes:
                future = executor.submit(
                    process_specific_theme, 
                    main_theme, 
                    specific_theme_data, 
                    all_results
                )
                futures.append(future)
        
        # Wait for all tasks to complete
        for future in futures:
            future.result()
    
    elapsed_time = time.time() - start_time
    
    print("\n" + "="*80)
    print("📊 重新分配主题 ID...")
    
    # Reassign IDs to ensure they are sequential across all themes
    for i, theme in enumerate(all_results, 1):
        theme["id"] = i
    
    print(f"💾 保存主题到文件: {output_file}")
    
    # Save all themes to the output file
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(all_results, f, indent=2, ensure_ascii=False)
    
    print("\n" + "="*80)
    print("✅ 所有主题处理完成！")
    print("="*80)
    print(f"📈 统计信息:")
    print(f"   总主题数: {len(all_results)}")
    print(f"   原始主题: {total_specific_themes}")
    print(f"   生成主题: {len(all_results) - total_specific_themes}")
    print(f"   扩展比例: {len(all_results) / total_specific_themes:.1f}x")
    print(f"⏱️  总耗时: {elapsed_time:.1f} 秒")
    print(f"💾 保存路径: {output_file}")
    print("="*80)

if __name__ == "__main__":
    main()