"""逐页提取：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