vent_agent_2.0/extensions/main_fan/workflow_documents.py

"""逐页提取：Python 文本层优先；GLM OCR 失败的扫描/混合页由 Qwen 补识别。"""
from __future__ import annotations

import asyncio
import base64
from dataclasses import dataclass
import hashlib
import json
from pathlib import Path
import re
import threading
import time
from urllib.parse import urlsplit

import httpx
import pypdfium2 as pdfium
import pypdfium2.raw as pdfium_raw

import config
from agents.model_factory import _is_qwen35_model, _qwen_thinking_payload, normalize_thinking
from extensions.main_fan.exploration import _write
from obs import get_logger
from runtime.python_sandbox import digest, workspace_file

logger = get_logger(__name__)
VERSION = 'main-fan-workflow-pdf-v3'
_PDFIUM_LOCK = threading.RLock()
QWEN_OCR_PROMPT = (
    '请对这张报告页面进行完整OCR。按原页顺序逐字转录标题、报告编号、表头、单位和全部数据行；'
    '表格分别输出为Markdown表格，保留测点序号及工况行。数字保持原值，不计算、不纠错、不解释；'
    '图中只转录实际可读的标签，不从曲线猜数；看不清的单元格写[无法辨认]。只输出识别结果。')


def service_url(value: str) -> str:
    parts = urlsplit(value)
    if parts.scheme not in {'http', 'https'} or not parts.netloc or parts.username or parts.password or parts.query or parts.fragment:
        raise ValueError('模型地址须为不含凭证、查询参数的 HTTP(S) API 基址')
    return value.rstrip('/')


@dataclass(frozen=True)
class OCRSettings:
    base_url: str = 'http://10.244.3.155:10803/v1'
    model: str = 'glm-ocr'
    api_key: str = ''
    timeout: int = 120
    max_tokens: int = 4096
    concurrency: int = 2
    dpi: int = 150
    max_side: int = 2200

    @classmethod
    def from_config(cls):
        value = cls(
            base_url=service_url(config.get_str('MAIN_FAN_OCR_BASE_URL', cls.base_url)),
            model=config.get_str('MAIN_FAN_OCR_MODEL', cls.model),
            api_key=config.get_str('MAIN_FAN_OCR_API_KEY'),
            timeout=config.get_int('MAIN_FAN_OCR_TIMEOUT_SECONDS', cls.timeout),
            max_tokens=config.get_int('MAIN_FAN_OCR_MAX_TOKENS', cls.max_tokens),
            concurrency=config.get_int('MAIN_FAN_OCR_CONCURRENCY', cls.concurrency),
            dpi=config.get_int('MAIN_FAN_PDF_DPI', cls.dpi),
            max_side=config.get_int('MAIN_FAN_PDF_MAX_SIDE', cls.max_side),
        )
        if (not value.model.strip() or value.timeout <= 0 or not 256 <= value.max_tokens <= 8192
                or not 1 <= value.concurrency <= 4 or not 100 <= value.dpi <= 300
                or not 1000 <= value.max_side <= 3000):
            raise ValueError('PDF/OCR 配置无效：检查模型名、超时、输出长度、并发及图片尺寸')
        return value

    def cache_signature(self):
        # 只缓存原件的原始提取文字，不缓存任何矿井归属或研判结论。
        return [VERSION, self.base_url, self.model, self.max_tokens, self.dpi, self.max_side,
                config.OPENAI_BASE_URL, config.MODEL, normalize_thinking(None),
                config.get_int('MAIN_FAN_QWEN_OCR_MAX_TOKENS', 8192)]


def repeated_ocr(text):
    """仅拦截长串循环，避免把表内少量合法重复值误判为失败。"""
    previous, consecutive = '', 0
    for line in text.splitlines():
        line = re.sub(r'\s+', ' ', line).strip()
        if not line:
            continue
        consecutive = consecutive + 1 if line == previous else 1
        previous = line
        if consecutive >= 12 and sum(char.isalnum() for char in line) >= 4:
            return True
    tail = text[-2000:]
    for width in range(8, 81):
        unit = tail[-width:]
        if sum(char.isalnum() for char in unit) >= 4 and tail.endswith(unit * 12):
            return True
    return False


def usable_text(text: str) -> bool:
    cleaned = re.sub(r'契约锁|扫描全能王|CamScanner', '', text, flags=re.I)
    meaningful = sum(char.isalnum() for char in cleaned)
    return meaningful >= 24 and text.count('\ufffd') <= max(1, len(text) // 100)


def large_raster(page) -> bool:
    """有文字的页也可能嵌着整张扫描表；不能把页眉/水印当作完整内容。"""
    width, height = page.get_size()
    area = width * height
    image_area = 0
    for obj in page.get_objects(filter=[pdfium_raw.FPDF_PAGEOBJ_IMAGE]):
        left, bottom, right, top = obj.get_bounds()
        covered = max(0, min(width, right) - max(0, left)) * max(0, min(height, top) - max(0, bottom))
        image_area += covered
        if area and image_area / area >= .12:
            return True
    return False


def prepare_pages(source: Path, directory: Path, settings: OCRSettings):
    """PDFium 不跨线程同时使用；渲染后立即释放页面和位图。"""
    pages = []
    directory.mkdir(parents=True, exist_ok=True)
    with _PDFIUM_LOCK, pdfium.PdfDocument(source) as pdf:
        for index in range(len(pdf)):
            page = pdf[index]
            stem = directory / f'page-{index + 1:04d}'
            row = {'page': index + 1, 'status': 'pending', 'method': 'python', 'text': ''}
            step = 'python_text'
            try:
                # PdfTextPage 只有 close()，不支持上下文管理器协议。
                textpage = page.get_textpage()
                try:
                    text = textpage.get_text_range().replace('\r\n', '\n').strip()
                finally:
                    textpage.close()
                step = 'classify_page'
                raster = large_raster(page)
                if usable_text(text) and not raster:
                    row.update(status='completed', text=text)
                elif not text and not list(page.get_objects()):
                    row.update(status='completed', blank=True)
                else:
                    row['method'] = 'glm_ocr'
                    step = 'render_for_ocr'
                    # 保存原始文本层，扫描页 OCR 失败时也不能冒充完整提取。
                    stem.with_suffix('.native.txt').write_text(text, encoding='utf-8')
                    width, height = page.get_size()
                    scale = min(settings.dpi / 72, settings.max_side / max(width, height))
                    bitmap = page.render(scale=scale)
                    try:
                        image = bitmap.to_pil().convert('RGB')
                        try:
                            image.save(stem.with_suffix('.png'))
                        finally:
                            image.close()
                    finally:
                        bitmap.close()
                    row['image_name'] = stem.with_suffix('.png').name
                pages.append(row)
            except Exception as exc:
                pages.append({**row, 'status': 'failed', 'error_stage': step,
                              'error': type(exc).__name__ + ': ' + str(exc)[:240]})
            finally:
                page.close()
    return pages


def _pages_digest(pages):
    return hashlib.sha256(json.dumps(pages, ensure_ascii=False, sort_keys=True).encode()).hexdigest()


def load_cache(path, source_sha256, signature, page_count):
    try:
        value = json.loads(path.read_text(encoding='utf-8'))
        pages = value['pages']
        if (value['source_sha256'] == source_sha256 and value['signature'] == signature
                and value['pages_sha256'] == _pages_digest(pages)
                and [row['page'] for row in pages] == list(range(1, page_count + 1))
                and all(row['status'] == 'completed' and row.get('method') in {'python', 'glm_ocr', 'qwen_ocr'} and isinstance(row['text'], str)
                        and (row['text'].strip() or row.get('blank')) for row in pages)):
            return pages
    except (OSError, ValueError, KeyError, TypeError):
        pass
    return None


class PDFExtractor:
    def __init__(self, user, session, prefix, cache_root, *, settings=None, client=None,
                 qwen_client=None, progress=None, on_usage=None, on_qwen_usage=None,
                 on_qwen_call=None, cache_policy='auto'):
        self.user, self.session, self.prefix = user, session, prefix
        self.cache_root = Path(cache_root) / 'workflow_pdf'
        self.settings = settings or OCRSettings.from_config()
        self.client = client
        self.qwen_client = qwen_client
        self.progress = progress
        self.on_usage = on_usage
        self.on_qwen_usage = on_qwen_usage
        self.on_qwen_call = on_qwen_call
        if cache_policy not in {'auto', 'refresh', 'off'}:
            raise ValueError('PDF 缓存策略须为 auto / refresh / off')
        self.cache_policy = cache_policy
        self.files = {}

    def remember(self, path):
        virtual = self.prefix + '/' + path.name
        self.files[virtual] = digest(path)
        return virtual

    async def _ocr_attempt(self, row, directory, client, encoded, *, fallback=False):
        started = time.monotonic()
        image = directory / row['image_name']
        stem = directory / (f"page-{row['page']:04d}" + ('.qwen' if fallback else ''))
        result = {'method': 'qwen_ocr' if fallback else 'glm_ocr', 'status': 'failed'}
        try:
            if fallback:
                if not _is_qwen35_model(config.MODEL):
                    raise ValueError('OCR 补识别需要配置 Qwen3.5 视觉模型')
                endpoint, model, api_key = service_url(config.OPENAI_BASE_URL), config.MODEL, config.OPENAI_API_KEY
                maximum = config.get_int('MAIN_FAN_QWEN_OCR_MAX_TOKENS', 8192)
                timeout = config.get_int('MAIN_FAN_QWEN_OCR_TIMEOUT_SECONDS', 240)
                if maximum <= 0 or timeout <= 0:
                    raise ValueError('Qwen OCR 输出长度和超时必须大于 0')
                prompt = QWEN_OCR_PROMPT
                thinking = _qwen_thinking_payload(normalize_thinking(None))
            else:
                endpoint, model, api_key = self.settings.base_url, self.settings.model, self.settings.api_key
                maximum, timeout = self.settings.max_tokens, self.settings.timeout
                prompt, thinking = 'Text Recognition:', {}
            result['model'] = model
            _write(Path(str(stem) + '.request.json'), {'model': model, 'prompt': prompt,
                'image': image.name, 'image_sha256': digest(image), 'max_tokens': maximum, **thinking})
            payload = {'model': model, 'messages': [{'role': 'user', 'content': [
                {'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,' + encoded}},
                {'type': 'text', 'text': prompt}]}],
                'temperature': 0, 'max_tokens': maximum, 'stream': False, **thinking}
            # 两个服务分别构造鉴权头，GLM 请求始终不带 Qwen 密钥。
            headers = {'Authorization': 'Bearer ' + api_key} if api_key else {}
            if fallback and self.on_qwen_call:
                self.on_qwen_call()
            response = await client.post(endpoint + '/chat/completions', json=payload,
                                         headers=headers, timeout=timeout)
            Path(str(stem) + '.response.json').write_text(response.text, encoding='utf-8')
            result['http_status'] = response.status_code
            response.raise_for_status()
            body = response.json()
            if not isinstance(body, dict):
                raise ValueError('OCR 服务返回格式无效')
            usage_callback = self.on_qwen_usage if fallback else self.on_usage
            if usage_callback:
                usage_callback(body.get('usage') or {})
            choice = body['choices'][0]
            if not isinstance(choice, dict):
                raise ValueError('OCR 服务未返回有效候选消息')
            message = choice['message']
            if not isinstance(message, dict):
                raise ValueError('OCR 服务未返回有效消息')
            text = message.get('content')
            result.update(finish_reason=choice.get('finish_reason'), response_id=body.get('id', ''))
            if choice.get('finish_reason') != 'stop' or not isinstance(text, str) or not text.strip() or message.get('tool_calls'):
                raise ValueError('OCR 未返回完整页面文字，finish_reason=' + str(choice.get('finish_reason')))
            if repeated_ocr(text):
                raise ValueError('OCR 输出出现长串循环重复，不能视为完整页面')
            result.update(status='completed', text=text.strip())
        except (httpx.HTTPError, ValueError, KeyError, IndexError, TypeError) as exc:
            result.update(text='', error=type(exc).__name__ + ': ' + str(exc)[:240])
        result['duration_ms'] = round((time.monotonic() - started) * 1000)
        logger.info('PDF 扫描页识别结束', extra={'kv': {'event': 'workflow.ocr.page',
            'page': row['page'], **{key: value for key, value in result.items() if key != 'text'}}})
        return result

    async def _ocr(self, row, directory, client):
        started = time.monotonic()
        image = directory / row['image_name']
        encoded = base64.b64encode(await asyncio.to_thread(image.read_bytes)).decode('ascii')
        attempts = [await self._ocr_attempt(row, directory, client, encoded)]
        if attempts[0]['status'] != 'completed':
            logger.warning('GLM OCR 失败，改用 Qwen 识别该页', extra={'kv': {
                'event': 'workflow.ocr.fallback', 'page': row['page'], 'error': attempts[0].get('error')}})
            attempts.append(await self._ocr_attempt(row, directory, self.qwen_client or client, encoded, fallback=True))
        final = attempts[-1]
        row.update(method=final['method'], status=final['status'], text=final.get('text', ''),
            response_id=final.get('response_id', ''), fallback_used=len(attempts) > 1,
            attempts=[{key: value for key, value in attempt.items() if key != 'text'} for attempt in attempts],
            duration_ms=round((time.monotonic() - started) * 1000))
        if final['status'] != 'completed':
            row.update(error_stage=final['method'], error=final.get('error', 'OCR 未完成'))
        else:
            row.pop('error_stage', None)
            row.pop('error', None)

    async def extract(self, materials):
        if self.client is not None:
            return await self._extract(materials, self.client)
        async with httpx.AsyncClient(trust_env=False) as client:
            return await self._extract(materials, client)

    async def _extract(self, materials, client):
        reports = []
        for number, material in enumerate(materials, 1):
            source = workspace_file(self.user, self.session, material['path'])
            sha = await asyncio.to_thread(digest, source)
            if sha != material['sha256']:
                raise ValueError('PDF 原件已变化：' + material['path'])
            signature = self.settings.cache_signature()
            key = hashlib.sha256(json.dumps([sha, signature], sort_keys=True).encode()).hexdigest()
            cache = self.cache_root / (key + '.json')
            old_prefix = self.prefix
            self.prefix = old_prefix + f'/pdf-{number}'
            directory = workspace_file(self.user, self.session, self.prefix + '/extraction.json').parent
            directory.mkdir(parents=True, exist_ok=True)
            pages = (await asyncio.to_thread(load_cache, cache, sha, signature, material['page_count'])
                     if self.cache_policy == 'auto' else None)
            reused = pages is not None
            if pages is None:
                pages = await asyncio.to_thread(prepare_pages, source, directory, self.settings)
                if len(pages) != material['page_count']:
                    raise ValueError('PDF 实际页数与下载记录不一致')
                semaphore = asyncio.Semaphore(self.settings.concurrency)
                async def recognize(row):
                    async with semaphore:
                        await self._ocr(row, directory, client)
                        if self.progress:
                            self.progress({'stage': 'pdf_ocr', 'report': number, 'page': row['page'],
                                'page_count': len(pages), 'status': row['status']})
                await asyncio.gather(*(recognize(row) for row in pages if row['status'] == 'pending'))
                if self.cache_policy != 'off' and all(row['status'] == 'completed' for row in pages):
                    # 缓存仅存原始文字和页码；本轮绑定与正文来源总是重新附加。
                    cached_pages = [{key: value for key, value in row.items() if key != 'image_name'} for row in pages]
                    await asyncio.to_thread(_write, cache, {'source_sha256': sha, 'signature': signature,
                        'pages': cached_pages, 'pages_sha256': _pages_digest(cached_pages)})
            for row in pages:
                text_path = directory / f"page-{row['page']:04d}.txt"
                text_path.write_text(row['text'], encoding='utf-8')
                row['text_path'] = self.prefix + '/' + text_path.name
                row.pop('image_name', None)
            report = {key: material.get(key) for key in ('path', 'filename', 'sha256', 'is_mock', 'binding', 'note')}
            report.update(page_count=len(pages), complete=all(row['status'] == 'completed' for row in pages),
                          reused=reused, pages=pages)
            _write(directory / 'extraction.json', report)
            for path in directory.iterdir():
                if path.is_file():
                    self.remember(path)
            reports.append(report)
            self.prefix = old_prefix
            failures = [row for row in pages if row['status'] != 'completed']
            logger.info('PDF 全页提取结束', extra={'kv': {'event': 'workflow.pdf.completed',
                'report': number, 'page_count': len(pages), 'complete': report['complete'], 'reused': reused,
                'python_pages': sum(row['method'] == 'python' and row['status'] == 'completed' for row in pages),
                'ocr_pages': sum(row['method'] in {'glm_ocr', 'qwen_ocr'} and row['status'] == 'completed' for row in pages),
                'qwen_ocr_pages': sum(row['method'] == 'qwen_ocr' and row['status'] == 'completed' for row in pages),
                'failed_pages': len(failures),
                **({'first_error': {key: failures[0].get(key) for key in ('page', 'error_stage', 'error')}}
                   if failures else {})}})
        return reports