batch_parser_stream.py

"""
试卷批量解析工具 - 流式版 + 实时监控
- 流式传输，实时显示AI输出
- WebSocket服务器，前端可实时监控
- 支持多并发流式处理
"""
import os
import sys
import re
import json
import time
import asyncio
import argparse
import threading
from pathlib import Path
from typing import Optional, Dict, List, Set
from docx import Document
from openai import OpenAI
import http.server
import socketserver
import webbrowser

# ============ 配置区域 ============
BASE_DIR = Path(__file__).parent


def _load_env_file(path: Path) -> None:
    """Minimal .env loader to avoid hardcoding secrets in code."""
    if not path.exists():
        return

    try:
        for raw in path.read_text(encoding="utf-8", errors="ignore").splitlines():
            line = raw.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, value = line.split("=", 1)
            key = key.strip()
            value = value.strip().strip('"').strip("'")
            if key and key not in os.environ:
                os.environ[key] = value
    except Exception:
        # Best-effort only.
        return


_load_env_file(BASE_DIR / ".env")
_load_env_file(BASE_DIR / "backend" / ".env")

API_BASE = os.getenv("API_BASE", "https://api.siliconflow.cn/v1")
API_KEY = os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY") or ""
MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/DeepSeek-V3.2")
INPUT_DIR = BASE_DIR / "exam_papers_to_parse"
OUTPUT_DIR = BASE_DIR / "parsed_results_stream"

# 全局状态（供前端监控）
STREAM_STATUS = {}  # {file_id: {"status": "streaming/done/error", "content": "", "tokens": 0}}
STATUS_LOCK = threading.Lock()

# ============ Prompt模板 ============
SYSTEM_PROMPT = """你是一个专业的上海高考英语试卷分片解析专家。将试卷（含答案和解析）按题型分类提取为结构化JSON。

## 题型分类 (Standard Shanghai Gaokao Format)
1. **grammar** - 语法填空 (Section A, 10题, 题号21-30). 一篇短文，10个挖空.
2. **vocabulary** - 词汇选择 (Section B, 11选10, 题号31-40). 一篇短文，文后给出一个11个选项的Word Bank (A-K).
3. **cloze** - 完形填空 (Cloze Test, 15题, 题号41-55). 一篇短文，15个挖空，每个空有4个选项.
4. **reading_a** - 阅读理解A篇 (3-4题, 题号56-59).
5. **reading_b** - 阅读理解B篇 (3-4题, 题号60-62).
6. **reading_c** - 阅读理解C篇 (3-4题, 题号63-66).
7. **reading_d** - 六选四 (Section C, 4题, 题号67-70). 一篇短文，4个挖空，6个句子选项.
8. **summary** - 概要写作 (Section D). 一篇短文.

## 输出JSON格式
```json
{
  "exam_meta": {
    "year": 2024,
    "region": "宝山区",
    "exam_type": "一模",
    "term": "上学期"
  },
  "sections": [
    {
      "section_type": "grammar",
      "section_name": "Section A Grammar",
      "start_question": 21,
      "end_question": 30,
      "material": "完整填空文章内容...",
      "items": [
        {"id": "21", "blank_context": "上下文", "answer": "taking", "explanation": "解析...", "grammar_point": "Non-finite verb"}
      ]
    },
    {
      "section_type": "vocabulary",
      "section_name": "Section B Vocabulary",
      "start_question": 31,
      "end_question": 40,
      "material": "完整文章内容...",
      "word_bank": ["A. abundance", "B. broadcast", "C. ..."], 
      "items": [
        {"id": "31", "answer": "A", "explanation": "解析..."}
      ]
    },
    {
      "section_type": "cloze",
      "section_name": "Cloze Test",
      "start_question": 41,
      "end_question": 55,
      "material": "完整文章内容...",
      "items": [
        {"id": "41", "blank_context": "句内上下文", "options": ["A. good", "B. bad", "C. ..."], "answer": "B", "explanation": "..."}
      ]
    },
    {
      "section_type": "reading_a",
      "section_name": "Reading Comprehension A",
      "material": "文章全文...",
      "items": [{"id": "56", "question_stem": "题目...", "options": {"A": "...", "B": "..."}, "answer": "C", "explanation": "..."}]
    },
    {
      "section_type": "reading_d",
      "section_name": "Reading Comprehension D (六选四)",
      "material": "原文...",
      "word_bank": ["A. Sentence 1", "B. Sentence 2", ...],
      "items": [{"id": "67", "answer": "A", "explanation": "..."}]
    },
    {
      "section_type": "summary",
      "section_name": "Summary Writing",
      "material": "原文",
      "sample_answer": "参考范文",
      "explanation": "写作思路"
    }
  ]
}
```

## 重要规则
1. **Word Bank处理**: Vocabulary和Reading D必须提取`word_bank`字段。
2. **完整性**: 必须包含`answer`和`explanation`。
3. **内容清洗**: `material`字段中不要包含由于OCR产生的页眉页脚干扰。
4. **格式**: 严格遵守JSON格式，不要输出Markdown格式。
"""


USER_PROMPT_TEMPLATE = """请完整解析以下上海高考英语试卷（含答案和解析），提取所有8个题型。

注意：使用试卷中的原始题号（有些试卷没有听力，题号可能从1开始）。

## 试卷内容
{content}

直接输出JSON格式结果。"""


def extract_exam_info(filename: str, folder_path: str) -> Dict:
    info = {"year": None, "region": None, "exam_type": None, "term": None}
    full_path = f"{folder_path}/{filename}"
    
    year_match = re.search(r'(20\d{2})', full_path)
    if year_match:
        info["year"] = int(year_match.group(1))
    
    regions = ["黄浦区", "徐汇区", "长宁区", "静安区", "普陀区", "虹口区", "杨浦区",
               "闵行区", "宝山区", "嘉定区", "浦东新区", "金山区", "松江区", "青浦区", "奉贤区", "崇明区"]
    for region in regions:
        if region in full_path:
            info["region"] = region
            break
    
    if "一模" in full_path or "上学期" in full_path:
        info["exam_type"] = "一模"
        info["term"] = "上学期"
    elif "二模" in full_path or "下学期" in full_path:
        info["exam_type"] = "二模"
        info["term"] = "下学期"
    
    return info


DOC_TEXT_CACHE = {}

def preload_doc_files(files: List[Path]):
    doc_files = [f for f in files if f.suffix.lower() == '.doc']
    if not doc_files:
        return
    
    print(f"\n正在预处理 {len(doc_files)} 个.doc文件...")
    try:
        import win32com.client
        import pythoncom
        pythoncom.CoInitialize()
        word = win32com.client.Dispatch("Word.Application")
        word.Visible = False
        
        for i, f in enumerate(doc_files):
            try:
                doc = word.Documents.Open(str(f.resolve()))
                text = doc.Content.Text
                doc.Close(False)
                DOC_TEXT_CACHE[str(f)] = text
                print(f"  [{i+1}/{len(doc_files)}] ✓ {f.name[:40]}...")
            except Exception as e:
                DOC_TEXT_CACHE[str(f)] = ""
        
        word.Quit()
        pythoncom.CoUninitialize()
    except Exception as e:
        print(f"预处理出错: {e}")


def extract_text(file_path: Path) -> str:
    if file_path.suffix.lower() == '.doc':
        return DOC_TEXT_CACHE.get(str(file_path), "")
    elif file_path.suffix.lower() == '.docx':
        try:
            doc = Document(str(file_path))
            return "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
        except:
            return ""
    return ""


class StreamingExamParser:
    """流式试卷解析器"""
    
    def __init__(self):
        self.client = OpenAI(api_key=API_KEY, base_url=API_BASE)
        self.output_dir = OUTPUT_DIR
        self.section_dirs = {
            "grammar": self.output_dir / "Grammar",
            "vocabulary": self.output_dir / "Vocabulary",
            "cloze": self.output_dir / "Cloze",
            "reading_a": self.output_dir / "Reading_A",
            "reading_b": self.output_dir / "Reading_B",
            "reading_c": self.output_dir / "Reading_C",
            "reading_d": self.output_dir / "Reading_D",
            "summary": self.output_dir / "Summary",
            "full": self.output_dir / "Full_Papers"
        }
        self.debug_dir = self.output_dir / "debug"
        for d in list(self.section_dirs.values()) + [self.debug_dir]:
            d.mkdir(parents=True, exist_ok=True)
        
        self.success_count = 0
        self.fail_count = 0
        self.lock = threading.Lock()
    
    def call_api_streaming(self, content: str, file_id: str) -> Optional[Dict]:
        """流式调用API"""
        user_prompt = USER_PROMPT_TEMPLATE.format(content=content[:60000])
        
        with STATUS_LOCK:
            STREAM_STATUS[file_id] = {"status": "streaming", "content": "", "tokens": 0, "start_time": time.time()}
        
        for attempt in range(3):
            try:
                full_response = ""
                token_count = 0
                
                stream = self.client.chat.completions.create(
                    model=MODEL_NAME,
                    messages=[
                        {"role": "system", "content": SYSTEM_PROMPT},
                        {"role": "user", "content": user_prompt}
                    ],
                    temperature=0.1,
                    max_tokens=24000,
                    stream=True
                )
                
                for chunk in stream:
                    if chunk.choices[0].delta.content:
                        text = chunk.choices[0].delta.content
                        full_response += text
                        token_count += 1
                        
                        # 更新全局状态（每2个token更新一次，更流畅）
                        if token_count % 2 == 0:
                            with STATUS_LOCK:
                                STREAM_STATUS[file_id]["content"] = full_response[-800:]  # 保留更多字符
                                STREAM_STATUS[file_id]["tokens"] = token_count
                        
                        # 终端实时显示
                            print(text, end="", flush=True)
                
                print()  # 换行
                
                # 保存原始响应
                debug_file = self.debug_dir / f"{file_id}.txt"
                with open(debug_file, "w", encoding="utf-8") as f:
                    f.write(full_response)
                
                # 解析JSON
                json_str = self._extract_json(full_response)
                if json_str:
                    json_str = self._clean_json(json_str)
                    data = json.loads(json_str)
                    if "sections" in data:
                        with STATUS_LOCK:
                            STREAM_STATUS[file_id]["status"] = "done"
                            STREAM_STATUS[file_id]["sections"] = len(data.get("sections", []))
                            STREAM_STATUS[file_id]["end_time"] = time.time()
                        return data
                
                with STATUS_LOCK:
                    STREAM_STATUS[file_id]["status"] = "error"
                    STREAM_STATUS[file_id]["error"] = "无法解析JSON"
                    STREAM_STATUS[file_id]["end_time"] = time.time()
                    
            except Exception as e:
                print(f"\n⚠️ 错误: {e}")
                with STATUS_LOCK:
                    STREAM_STATUS[file_id]["status"] = "error"
                    STREAM_STATUS[file_id]["error"] = str(e)[:50]
                    STREAM_STATUS[file_id]["end_time"] = time.time()
                
                if attempt < 2:
                    time.sleep(3)
        
        return None
    
    def _clean_json(self, json_str: str) -> str:
        result = []
        in_string = False
        escape_next = False
        
        for char in json_str:
            if escape_next:
                result.append(char)
                escape_next = False
                continue
            if char == '\\':
                escape_next = True
                result.append(char)
                continue
            if char == '"':
                in_string = not in_string
                result.append(char)
                continue
            if in_string:
                if char == '\n':
                    result.append('\\n')
                elif char == '\r':
                    result.append('\\r')
                elif char == '\t':
                    result.append('\\t')
                elif ord(char) < 32:
                    result.append(f'\\u{ord(char):04x}')
                else:
                    result.append(char)
            else:
                result.append(char)
        
        return ''.join(result)
    
    def _extract_json(self, content: str) -> Optional[str]:
        content = content.strip()
        if content.startswith("```"):
            content = re.sub(r'^```\w*\n?', '', content)
            content = re.sub(r'\n?```$', '', content)
            content = content.strip()
        
        if content.startswith('{') and content.endswith('}'):
            return content
        
        start = content.find('{')
        end = content.rfind('}')
        if start != -1 and end > start:
            return content[start:end+1]
        return None
    
    def save_result(self, data: Dict, file_path: Path):
        meta = data.get("exam_meta", {})
        file_info = extract_exam_info(file_path.name, str(file_path.parent))
        for k, v in file_info.items():
            if v and not meta.get(k):
                meta[k] = v
        data["exam_meta"] = meta
        
        year = meta.get("year", "unknown")
        region = meta.get("region", "unknown")
        exam_type = meta.get("exam_type", "unknown")
        prefix = f"{year}_{region}_{exam_type}".replace(" ", "_")
        
        full_path = self.section_dirs["full"] / f"{prefix}_full.json"
        counter = 1
        while full_path.exists():
            full_path = self.section_dirs["full"] / f"{prefix}_full_{counter}.json"
            counter += 1
        
        with open(full_path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        sections = data.get("sections", [])
        saved_types = []
        for section in sections:
            section_type = section.get("section_type", "").lower()
            if section_type in self.section_dirs:
                section_data = {"exam_meta": meta, "section": section}
                section_path = self.section_dirs[section_type] / f"{prefix}_{section_type}.json"
                counter = 1
                while section_path.exists():
                    section_path = self.section_dirs[section_type] / f"{prefix}_{section_type}_{counter}.json"
                    counter += 1
                with open(section_path, "w", encoding="utf-8") as f:
                    json.dump(section_data, f, ensure_ascii=False, indent=2)
                saved_types.append(section_type)
        
        return saved_types
    
    def process_file(self, file_path: Path, index: int, total: int) -> bool:
        # 使用文件名+哈希确保唯一性，避免相似文件名冲突
        import hashlib
        name_hash = hashlib.md5(str(file_path).encode()).hexdigest()[:6]
        file_id = f"{file_path.stem[:24]}_{name_hash}"
        print(f"\n{'='*60}")
        print(f"[{index}/{total}] 📄 {file_path.name}")
        print(f"{'='*60}")
        
        content = extract_text(file_path)
        if len(content) < 500:
            print(f"⚠️ 内容太短，跳过")
            return False
        
        print(f"📝 内容: {len(content)} 字符\n")
        print("🤖 AI输出:")
        print("-" * 40)
        
        result = self.call_api_streaming(content, file_id)
        
        print("-" * 40)
        
        if result:
            saved_types = self.save_result(result, file_path)
            print(f"✅ 成功! {len(saved_types)} 个题型: {', '.join(saved_types)}")
            with self.lock:
                self.success_count += 1
            return True
        else:
            print(f"❌ 解析失败")
            with self.lock:
                self.fail_count += 1
            return False
    
    def process_all(self, files: List[Path], parallel: int = 1):
        print(f"\n{'='*60}")
        print(f"试卷解析工具 - 流式版")
        print(f"{'='*60}")
        print(f"模型: {MODEL_NAME}")
        print(f"待处理: {len(files)} 个文件")
        print(f"并发数: {parallel}")
        print(f"{'='*60}")
        
        preload_doc_files(files)
        start_time = time.time()
        
        if parallel > 1:
            global PARALLEL_MODE
            PARALLEL_MODE = True  # 并行模式，禁用终端输出
            from concurrent.futures import ThreadPoolExecutor, as_completed
            with ThreadPoolExecutor(max_workers=parallel) as executor:
                futures = {executor.submit(self.process_file, f, i, len(files)): f for i, f in enumerate(files, 1)}
                for future in as_completed(futures):
                    try:
                        future.result()
                    except Exception as e:
                        print(f"⚠️ 错误: {e}")
        else:
            for i, f in enumerate(files, 1):
                self.process_file(f, i, len(files))
        
        elapsed = time.time() - start_time
        print(f"\n{'='*60}")
        print(f"完成! 成功: {self.success_count} | 失败: {self.fail_count}")
        print(f"耗时: {elapsed/60:.1f}分钟")
        print(f"{'='*60}")


# ============ 监控API服务器 ============
def create_monitor_html():
    """创建监控页面"""
    html = '''<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <title>试卷解析监控</title>
    <style>
        * { margin: 0; padding: 0; box-sizing: border-box; }
        body { font-family: 'Microsoft YaHei', sans-serif; background: #1a1a2e; color: #eee; padding: 20px; }
        h1 { text-align: center; margin-bottom: 20px; color: #4fc3f7; }
        .grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(400px, 1fr)); gap: 15px; }
        .card { background: #16213e; border-radius: 10px; padding: 15px; border: 1px solid #0f3460; }
        .card.streaming { border-color: #4caf50; box-shadow: 0 0 10px rgba(76,175,80,0.3); }
        .card.done { border-color: #4fc3f7; }
        .card.error { border-color: #f44336; }
        .card-header { display: flex; justify-content: space-between; margin-bottom: 10px; }
        .card-title { font-weight: bold; color: #4fc3f7; font-size: 14px; }
        .card-status { font-size: 12px; padding: 2px 8px; border-radius: 10px; }
        .streaming .card-status { background: #4caf50; }
        .done .card-status { background: #2196f3; }
        .error .card-status { background: #f44336; }
        .card-content { background: #0f3460; padding: 10px; border-radius: 5px; font-family: monospace; font-size: 12px; 
                        max-height: 200px; overflow-y: auto; white-space: pre-wrap; word-break: break-all; }
        .card-footer { margin-top: 10px; font-size: 12px; color: #888; display: flex; justify-content: space-between; }
        .stats { text-align: center; margin-bottom: 20px; font-size: 18px; }
        .stats span { margin: 0 15px; }
        .streaming-count { color: #4caf50; }
        .done-count { color: #4fc3f7; }
        .error-count { color: #f44336; }
        .total-stats { text-align: center; margin-bottom: 25px; padding: 15px; background: linear-gradient(135deg, #1f4037, #99f2c8); 
                       border-radius: 15px; font-size: 24px; }
        .total-stats .big-num { font-size: 48px; font-weight: bold; color: #fff; text-shadow: 2px 2px 4px rgba(0,0,0,0.3); }
        .total-stats .label { font-size: 14px; color: rgba(255,255,255,0.8); }
        .total-stats .speed { color: #ffeb3b; font-weight: bold; }
    </style>
</head>
<body>
    <h1>🎯 试卷解析实时监控</h1>
    <div class="total-stats">
        <span class="big-num" id="total-tokens">0</span> <span class="label">Total Tokens</span>
        &nbsp;&nbsp;|&nbsp;&nbsp;
        <span class="speed" id="total-speed">0</span> <span class="label">tokens/s (当前)</span>
    </div>
    <div class="stats">
        <span class="streaming-count">⏳ 进行中: <b id="streaming-num">0</b></span>
        <span class="done-count">✅ 完成: <b id="done-num">0</b></span>
        <span class="error-count">❌ 失败: <b id="error-num">0</b></span>
    </div>
    <div class="grid" id="cards"></div>
    
    <script>
        let lastTotalTokens = 0;
        let lastStreamingTokens = 0;
        let lastTime = Date.now();
        
        async function fetchStatus() {
            try {
                const res = await fetch('/api/status');
                const data = await res.json();
                updateUI(data);
            } catch(e) {}
            setTimeout(fetchStatus, 100);  // 100ms轮询，更流畅
        }
        
        function updateUI(data) {
            const cards = document.getElementById('cards');
            let streaming = 0, done = 0, error = 0;
            let totalTokens = 0;
            let streamingTokens = 0;  // 只统计正在流式传输的tokens
            let html = '';
            
            for (const [id, info] of Object.entries(data)) {
                const status = info.status || 'unknown';
                if (status === 'streaming') {
                    streaming++;
                    streamingTokens += (info.tokens || 0);
                }
                else if (status === 'done') done++;
                else if (status === 'error') error++;
                
                totalTokens += (info.tokens || 0);
                
                // 完成后使用end_time，否则使用当前时间
                let elapsed;
                if (status === 'done' || status === 'error') {
                    elapsed = info.end_time && info.start_time ? ((info.end_time - info.start_time) | 0) : 0;
                } else {
                    elapsed = info.start_time ? ((Date.now()/1000 - info.start_time) | 0) : 0;
                }
                const speed = elapsed > 0 ? ((info.tokens || 0) / elapsed).toFixed(1) : '0';
                
                html += `
                <div class="card ${status}">
                    <div class="card-header">
                        <span class="card-title">${id}</span>
                        <span class="card-status">${status === 'streaming' ? '🔄 生成中' : status === 'done' ? '✅ 完成' : '❌ 失败'}</span>
                    </div>
                    <div class="card-content">${info.content || info.error || '等待...'}</div>
                    <div class="card-footer">
                        <span>Tokens: ${info.tokens || 0}</span>
                        <span>${speed} tokens/s</span>
                        <span>${elapsed}s</span>
                        ${info.sections ? '<span>题型: ' + info.sections + '</span>' : ''}
                    </div>
                </div>`;
            }
            
            // 计算实时速度 (只统计正在streaming的tokens增量，避免负数)
            const now = Date.now();
            const deltaTime = (now - lastTime) / 1000;
            const deltaTokens = streamingTokens - lastStreamingTokens;
            // 速度至少为0
            const currentSpeed = deltaTime > 0 ? Math.max(0, deltaTokens / deltaTime).toFixed(1) : '0';
            lastStreamingTokens = streamingTokens;
            lastTotalTokens = totalTokens;
            lastTime = now;
            lastTime = now;
            
            cards.innerHTML = html;
            document.getElementById('streaming-num').textContent = streaming;
            document.getElementById('done-num').textContent = done;
            document.getElementById('error-num').textContent = error;
            document.getElementById('total-tokens').textContent = totalTokens.toLocaleString();
            document.getElementById('total-speed').textContent = currentSpeed;
        }
        
        fetchStatus();
    </script>
</body>
</html>'''
    return html


class MonitorHandler(http.server.SimpleHTTPRequestHandler):
    def do_GET(self):
        if self.path == '/':
            self.send_response(200)
            self.send_header('Content-type', 'text/html; charset=utf-8')
            self.end_headers()
            self.wfile.write(create_monitor_html().encode('utf-8'))
        elif self.path == '/api/status':
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.send_header('Access-Control-Allow-Origin', '*')
            self.end_headers()
            with STATUS_LOCK:
                self.wfile.write(json.dumps(STREAM_STATUS).encode('utf-8'))
        else:
            self.send_error(404)
    
    def log_message(self, format, *args):
        pass  # 禁止日志输出


def start_monitor_server(port=8765):
    """启动监控服务器"""
    with socketserver.TCPServer(("", port), MonitorHandler) as httpd:
        print(f"📊 监控页面: http://localhost:{port}")
        httpd.serve_forever()


def collect_files(source_dir: Path) -> List[Path]:
    files = []
    for pattern in ["**/*解析*.docx", "**/*解析*.doc"]:
        for f in source_dir.glob(pattern):
            if "原卷" in f.name or f.name.startswith("~$"):
                continue
            files.append(f)
    return files


def main():
    parser = argparse.ArgumentParser(description="试卷解析工具 - 流式版")
    parser.add_argument("--input", "-i", action="store_true", help="从exam_papers_to_parse目录读取")
    parser.add_argument("--test", "-t", type=int, default=0, help="测试模式")
    parser.add_argument("--parallel", "-p", type=int, default=1, help="并发数")
    parser.add_argument("--monitor", "-m", action="store_true", help="启动监控页面")
    args = parser.parse_args()
    
    try:
        if args.input and INPUT_DIR.exists():
            # 只获取根目录下的文件，忽略子文件夹
            files = [f for f in INPUT_DIR.iterdir() 
                     if f.is_file() and f.suffix.lower() in ['.doc', '.docx'] and not f.name.startswith('~$')]
        else:
            files = collect_files(BASE_DIR / "17-24mock")
        
        if not files:
            print("未找到文件!")
            return
        
        if args.test > 0:
            files = files[:args.test]
        
        # 启动监控服务器
        if args.monitor:
            monitor_thread = threading.Thread(target=start_monitor_server, daemon=True)
            monitor_thread.start()
            webbrowser.open("http://localhost:8765")
            time.sleep(1)
        
        exam_parser = StreamingExamParser()
        exam_parser.process_all(files, parallel=args.parallel)
        
    except KeyboardInterrupt:
        print("\n用户中断")
        sys.exit(0)


if __name__ == "__main__":
    import signal
    signal.signal(signal.SIGINT, lambda s, f: sys.exit(0))
    main()