diff --git a/server/routers/base_router.py b/server/routers/base_router.py index f81bfed5..ce090c7f 100644 --- a/server/routers/base_router.py +++ b/server/routers/base_router.py @@ -1,5 +1,7 @@ import os import yaml +import asyncio +import requests from pathlib import Path from fastapi import Request, Body, Depends, HTTPException from fastapi import APIRouter @@ -133,4 +135,111 @@ async def reload_info_config(): logger.error(f"重新加载信息配置失败: {e}") raise HTTPException(status_code=500, detail="重新加载信息配置失败") +@base.get("/ocr/health") +async def check_ocr_services_health(current_user: User = Depends(get_admin_user)): + """ + 检查所有OCR服务的健康状态 + 返回各个OCR服务的可用性信息 + """ + health_status = { + "rapid_ocr": {"status": "unknown", "message": ""}, + "mineru_ocr": {"status": "unknown", "message": ""}, + "paddlex_ocr": {"status": "unknown", "message": ""} + } + + # 检查 RapidOCR (ONNX) 模型 + try: + model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR") + det_model_path = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx") + rec_model_path = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx") + + if os.path.exists(model_dir) and os.path.exists(det_model_path) and os.path.exists(rec_model_path): + # 尝试初始化RapidOCR + from rapidocr_onnxruntime import RapidOCR + test_ocr = RapidOCR(det_box_thresh=0.3, det_model_path=det_model_path, rec_model_path=rec_model_path) + health_status["rapid_ocr"]["status"] = "healthy" + health_status["rapid_ocr"]["message"] = "RapidOCR模型已加载" + else: + health_status["rapid_ocr"]["status"] = "unavailable" + health_status["rapid_ocr"]["message"] = f"模型文件不存在: {model_dir}" + except Exception as e: + health_status["rapid_ocr"]["status"] = "error" + health_status["rapid_ocr"]["message"] = f"RapidOCR初始化失败: {str(e)}" + + # 检查 MinerU OCR 服务 + try: + mineru_uri = os.getenv("MINERU_OCR_URI", "http://localhost:30000") + health_url = f"{mineru_uri}/health" + + response = requests.get(health_url, timeout=5) + if response.status_code == 200: + health_status["mineru_ocr"]["status"] = "healthy" + health_status["mineru_ocr"]["message"] = f"MinerU服务运行正常 ({mineru_uri})" + else: + health_status["mineru_ocr"]["status"] = "unhealthy" + health_status["mineru_ocr"]["message"] = f"MinerU服务响应异常: {response.status_code}" + except requests.exceptions.ConnectionError: + health_status["mineru_ocr"]["status"] = "unavailable" + health_status["mineru_ocr"]["message"] = "MinerU服务无法连接,请检查服务是否启动" + except requests.exceptions.Timeout: + health_status["mineru_ocr"]["status"] = "timeout" + health_status["mineru_ocr"]["message"] = "MinerU服务连接超时" + except Exception as e: + health_status["mineru_ocr"]["status"] = "error" + health_status["mineru_ocr"]["message"] = f"MinerU服务检查失败: {str(e)}" + + # 检查 PaddleX OCR 服务 + try: + paddlex_uri = os.getenv("PADDLEX_URI", "http://localhost:8080") + health_url = f"{paddlex_uri}/health" + + response = requests.get(health_url, timeout=5) + if response.status_code == 200: + health_status["paddlex_ocr"]["status"] = "healthy" + health_status["paddlex_ocr"]["message"] = f"PaddleX服务运行正常 ({paddlex_uri})" + else: + health_status["paddlex_ocr"]["status"] = "unhealthy" + health_status["paddlex_ocr"]["message"] = f"PaddleX服务响应异常: {response.status_code}" + except requests.exceptions.ConnectionError: + health_status["paddlex_ocr"]["status"] = "unavailable" + health_status["paddlex_ocr"]["message"] = "PaddleX服务无法连接,请检查服务是否启动" + except requests.exceptions.Timeout: + health_status["paddlex_ocr"]["status"] = "timeout" + health_status["paddlex_ocr"]["message"] = "PaddleX服务连接超时" + except Exception as e: + health_status["paddlex_ocr"]["status"] = "error" + health_status["paddlex_ocr"]["message"] = f"PaddleX服务检查失败: {str(e)}" + + # 计算整体健康状态 + overall_status = "healthy" if any(svc["status"] == "healthy" for svc in health_status.values()) else "unhealthy" + + return { + "overall_status": overall_status, + "services": health_status, + "message": "OCR服务健康检查完成" + } + +@base.get("/ocr/stats") +async def get_ocr_stats(current_user: User = Depends(get_admin_user)): + """ + 获取OCR服务使用统计信息 + 返回各个OCR服务的处理统计和性能指标 + """ + try: + from src.plugins._ocr import get_ocr_stats + stats = get_ocr_stats() + + return { + "status": "success", + "stats": stats, + "message": "OCR统计信息获取成功" + } + except Exception as e: + logger.error(f"获取OCR统计信息失败: {str(e)}") + return { + "status": "error", + "stats": {}, + "message": f"获取OCR统计信息失败: {str(e)}" + } + diff --git a/src/__init__.py b/src/__init__.py index bffd33b5..97c00169 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,3 +1,4 @@ +import os from dotenv import load_dotenv load_dotenv("src/.env", override=True) @@ -8,8 +9,33 @@ executor = ThreadPoolExecutor() from src.config import Config # noqa: E402 config = Config() -from src.core.lightrag_based_kb import LightRagBasedKB # noqa: E402 -knowledge_base = LightRagBasedKB() +# 导入知识库相关模块 +from src.core.kb_factory import KnowledgeBaseFactory # noqa: E402 +from src.core.kb_manager import KnowledgeBaseManager # noqa: E402 +from src.core.lightrag_kb import LightRagKB # noqa: E402 +from src.core.chroma_kb import ChromaKB # noqa: E402 +from src.core.milvus_kb import MilvusKB # noqa: E402 + +# 注册知识库类型 +KnowledgeBaseFactory.register("lightrag", LightRagKB, { + "description": "基于图检索的知识库,支持实体关系构建和复杂查询" +}) + +KnowledgeBaseFactory.register("chroma", ChromaKB, { + "chunk_size": 1000, + "chunk_overlap": 200, + "description": "基于 ChromaDB 的轻量级向量知识库,适合开发和小规模部署" +}) + +KnowledgeBaseFactory.register("milvus", MilvusKB, { + "chunk_size": 1000, + "chunk_overlap": 200, + "description": "基于 Milvus 的生产级向量知识库,适合大规模高性能部署" +}) + +# 创建知识库管理器 +work_dir = os.path.join(config.save_dir, "knowledge_base_data") +knowledge_base = KnowledgeBaseManager(work_dir) from src.core import GraphDatabase # noqa: E402 graph_base = GraphDatabase() diff --git a/src/core/indexing.py b/src/core/indexing.py index 904d26fb..63c36fee 100644 --- a/src/core/indexing.py +++ b/src/core/indexing.py @@ -122,23 +122,53 @@ def plainreader(file_path): return text def parse_pdf(file, params=None): + """ + 解析PDF文件,支持多种OCR方式 + + Args: + file: PDF文件路径 + params: 参数字典,包含enable_ocr设置 + + Returns: + str: 解析得到的文本 + + Raises: + OCRServiceException: OCR服务不可用时抛出 + """ + from src.plugins._ocr import OCRServiceException + params = params or {} opt_ocr = params.get("enable_ocr", "disable") - if opt_ocr == "onnx_rapid_ocr": - from src.plugins import ocr - return ocr.process_pdf(file) - - elif opt_ocr == "mineru_ocr": - from src.plugins import ocr - return ocr.process_pdf_mineru(file) - - elif opt_ocr == "paddlex_ocr": - from src.plugins import ocr - return ocr.process_pdf_paddlex(file) - - else: + if opt_ocr == "disable": return pdfreader(file, params=params) + try: + if opt_ocr == "onnx_rapid_ocr": + from src.plugins import ocr + return ocr.process_pdf(file) + + elif opt_ocr == "mineru_ocr": + from src.plugins import ocr + return ocr.process_pdf_mineru(file) + + elif opt_ocr == "paddlex_ocr": + from src.plugins import ocr + return ocr.process_pdf_paddlex(file) + + else: + return pdfreader(file, params=params) + + except OCRServiceException as e: + logger.error(f"OCR service failed: {e.service_name} - {str(e)}") + raise + except Exception as e: + logger.error(f"PDF parsing failed: {str(e)}") + raise OCRServiceException( + f"PDF解析失败: {str(e)}", + opt_ocr, + "parsing_failed" + ) + async def parse_pdf_async(file, params=None): return await asyncio.to_thread(parse_pdf, file, params=params) diff --git a/src/core/knowledge_base.py b/src/core/knowledge_base.py index ac83f36c..bff4020a 100644 --- a/src/core/knowledge_base.py +++ b/src/core/knowledge_base.py @@ -390,7 +390,7 @@ class KnowledgeBase(ABC): # 使用 OCR 处理 PDF from src.core.indexing import parse_pdf_async text = await parse_pdf_async(str(file_path_obj), params=params) - return f"Using OCR to process {file_path_obj.name}\n\n{text}" + return f"# {file_path_obj.name}\n\n{text}" elif file_ext in ['.txt', '.md']: # 直接读取文本文件 diff --git a/src/core/lightrag_kb.py b/src/core/lightrag_kb.py index 9add7219..0a4b563d 100644 --- a/src/core/lightrag_kb.py +++ b/src/core/lightrag_kb.py @@ -213,10 +213,13 @@ class LightRagKB(KnowledgeBase): file_record['status'] = "done" except Exception as e: - logger.error(f"处理{content_type} {item} 失败: {e}, {traceback.format_exc()}") + error_msg = str(e) + logger.error(f"处理{content_type} {item} 失败: {error_msg}, {traceback.format_exc()}") self.files_meta[file_id]["status"] = "failed" + self.files_meta[file_id]["error"] = error_msg self._save_metadata() file_record['status'] = "failed" + file_record['error'] = error_msg processed_items_info.append(file_record) diff --git a/src/plugins/_ocr.py b/src/plugins/_ocr.py index 8e01b23d..66fc9358 100644 --- a/src/plugins/_ocr.py +++ b/src/plugins/_ocr.py @@ -1,7 +1,9 @@ import os import uuid +import time from pathlib import Path from argparse import ArgumentParser +from collections import defaultdict import fitz # fitz就是pip install PyMuPDF import numpy as np # Added import for numpy @@ -14,6 +16,53 @@ from src.utils import logger, is_text_pdf GOLBAL_STATE = {} +# OCR服务监控统计 +OCR_STATS = { + "requests": defaultdict(int), + "failures": defaultdict(int), + "service_status": defaultdict(str) +} + + +def log_ocr_request(service_name: str, file_path: str, success: bool, processing_time: float, error_msg: str = None): + """记录OCR请求统计信息""" + # 更新统计 + OCR_STATS["requests"][service_name] += 1 + + if not success: + OCR_STATS["failures"][service_name] += 1 + OCR_STATS["service_status"][service_name] = "error" + logger.error(f"OCR失败 - {service_name}: {os.path.basename(file_path)} - {error_msg}") + else: + OCR_STATS["service_status"][service_name] = "healthy" + logger.info(f"OCR成功 - {service_name}: {os.path.basename(file_path)}") + + +def get_ocr_stats(): + """获取OCR服务统计信息""" + stats = {} + for service in OCR_STATS["requests"]: + success_count = OCR_STATS["requests"][service] - OCR_STATS["failures"][service] + success_rate = (success_count / OCR_STATS["requests"][service]) if OCR_STATS["requests"][service] > 0 else 0 + + stats[service] = { + "total_requests": OCR_STATS["requests"][service], + "success_count": success_count, + "failure_count": OCR_STATS["failures"][service], + "success_rate": f"{success_rate:.2%}", + "status": OCR_STATS["service_status"][service] + } + + return stats + + +class OCRServiceException(Exception): + """OCR服务异常""" + def __init__(self, message, service_name=None, status_code=None): + super().__init__(message) + self.service_name = service_name + self.status_code = status_code + class OCRPlugin: """OCR 插件""" @@ -22,18 +71,59 @@ class OCRPlugin: self.ocr = None self.det_box_thresh = kwargs.get('det_box_thresh', 0.3) + def _check_rapid_ocr_availability(self): + """检查RapidOCR模型是否可用""" + try: + model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR") + det_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx") + rec_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx") + + if not os.path.exists(model_dir): + raise OCRServiceException( + f"模型目录不存在: {model_dir}。请下载 SWHL/RapidOCR 模型", + "rapid_ocr", + "model_not_found" + ) + + if not os.path.exists(det_model_dir) or not os.path.exists(rec_model_dir): + raise OCRServiceException( + f"模型文件缺失。请确认模型文件完整: {det_model_dir}, {rec_model_dir}", + "rapid_ocr", + "model_incomplete" + ) + + return True + + except Exception as e: + if isinstance(e, OCRServiceException): + raise + else: + raise OCRServiceException( + f"RapidOCR模型检查失败: {str(e)}", + "rapid_ocr", + "check_failed" + ) + def load_model(self): """加载 OCR 模型""" logger.info("加载 OCR 模型,仅在第一次调用时加载") + + # 先检查模型可用性 + self._check_rapid_ocr_availability() + model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR") det_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx") rec_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx") - assert os.path.exists(model_dir), ( - f"模型文件不存在,请下载 SWHL/RapidOCR 到 {model_dir}," - "并确认是否在 docker-compose.dev.yml 中添加 MODEL_DIR 环境变量" - ) - self.ocr = RapidOCR(det_box_thresh=0.3, det_model_path=det_model_dir, rec_model_path=rec_model_dir) - logger.info(f"OCR Plugin for det_box_thresh = {self.det_box_thresh} loaded.") + + try: + self.ocr = RapidOCR(det_box_thresh=0.3, det_model_path=det_model_dir, rec_model_path=rec_model_dir) + logger.info(f"OCR Plugin for det_box_thresh = {self.det_box_thresh} loaded.") + except Exception as e: + raise OCRServiceException( + f"RapidOCR模型加载失败: {str(e)}", + "rapid_ocr", + "load_failed" + ) def process_image(self, image): """ @@ -64,7 +154,9 @@ class OCRPlugin: image_path = self._create_temp_image_file(image) # 执行 OCR + start_time = time.time() result, _ = self.ocr(image_path) + processing_time = time.time() - start_time # 清理临时文件 if is_temp_file and os.path.exists(image_path): @@ -73,14 +165,17 @@ class OCRPlugin: # 提取文本 if result: text = '\n'.join([line[1] for line in result]) + log_ocr_request("rapid_ocr", image_path, True, processing_time) return text else: - logger.warning("OCR未能识别出文本内容") + log_ocr_request("rapid_ocr", image_path, False, processing_time, "OCR未能识别出文本内容") return "" except Exception as e: - logger.error(f"OCR处理失败: {str(e)}") - raise + error_msg = f"OCR处理失败: {str(e)}" + log_ocr_request("rapid_ocr", image_path, False, 0, error_msg) + logger.error(error_msg) + raise OCRServiceException(error_msg, "rapid_ocr", "processing_failed") def _create_temp_image_file(self, image): """ @@ -155,23 +250,60 @@ class OCRPlugin: """ import requests from .mineru import parse_doc + mineru_ocr_uri = os.getenv("MINERU_OCR_URI", "http://localhost:30000") mineru_ocr_uri_health = f"{mineru_ocr_uri}/health" - health_check_response = requests.get(mineru_ocr_uri_health, timeout=5) - if health_check_response.status_code != 200: - logger.error(f"Mineru OCR service health check failed with {mineru_ocr_uri_health}: {health_check_response.json()}") - raise RuntimeError("Mineru OCR service health check failed. Please check the log use `docker logs mineru-api`") + try: + # 健康检查 + health_check_response = requests.get(mineru_ocr_uri_health, timeout=5) + if health_check_response.status_code != 200: + error_detail = "Unknown error" + try: + error_detail = health_check_response.json() + except: + error_detail = health_check_response.text - pdf_path_list = [pdf_path] - output_dir = os.path.join(os.getcwd(), "tmp", "mineru_ocr") + raise OCRServiceException( + f"MinerU OCR服务健康检查失败: {error_detail}", + "mineru_ocr", + "health_check_failed" + ) - pdf_text = parse_doc(pdf_path_list, output_dir, - backend="vlm-sglang-client", - server_url=mineru_ocr_uri)[0] + except Exception as e: + if isinstance(e, OCRServiceException): + raise + raise OCRServiceException( + f"MinerU OCR服务检查失败: {str(e)}", + "mineru_ocr", + "service_error" + ) - logger.debug(f"Mineru OCR result: {pdf_text[:50]}(...) total {len(pdf_text)} characters.") - return pdf_text + try: + start_time = time.time() + pdf_path_list = [pdf_path] + output_dir = os.path.join(os.getcwd(), "tmp", "mineru_ocr") + + pdf_text = parse_doc(pdf_path_list, output_dir, + backend="vlm-sglang-client", + server_url=mineru_ocr_uri)[0] + + processing_time = time.time() - start_time + log_ocr_request("mineru_ocr", pdf_path, True, processing_time) + + logger.debug(f"Mineru OCR result: {pdf_text[:50]}(...) total {len(pdf_text)} characters.") + return pdf_text + + except Exception as e: + processing_time = time.time() - start_time + error_msg = f"MinerU OCR处理失败: {str(e)}" + log_ocr_request("mineru_ocr", pdf_path, False, processing_time, error_msg) + + raise OCRServiceException( + error_msg, + "mineru_ocr", + "processing_failed" + ) def process_pdf_paddlex(self, pdf_path): """ @@ -182,18 +314,61 @@ class OCRPlugin: from .paddlex import analyze_document, check_paddlex_health paddlex_uri = os.getenv("PADDLEX_URI", "http://localhost:8080") - health_check_response = check_paddlex_health(paddlex_uri) - if not health_check_response.ok: - logger.error(f"Paddlex OCR service health check failed with {paddlex_uri}: {health_check_response.json()}") - raise RuntimeError("Paddlex OCR service health check failed. Please check the log use `docker logs paddlex`") - result = analyze_document(pdf_path, base_url=paddlex_uri) + try: + # 健康检查 + health_check_response = check_paddlex_health(paddlex_uri) + if not health_check_response.ok: + error_detail = "Unknown error" + try: + error_detail = health_check_response.json() + except: + error_detail = health_check_response.text - if not result["success"]: - logger.error(f"Paddlex OCR failed: {result['error']}") - return "" + raise OCRServiceException( + f"PaddleX OCR服务健康检查失败: {error_detail}", + "paddlex_ocr", + "health_check_failed" + ) + except Exception as e: + if isinstance(e, OCRServiceException): + raise + raise OCRServiceException( + f"PaddleX OCR服务检查失败: {str(e)}", + "paddlex_ocr", + "service_error" + ) - return result["full_text"] + try: + start_time = time.time() + result = analyze_document(pdf_path, base_url=paddlex_uri) + processing_time = time.time() - start_time + + if not result["success"]: + error_msg = f"PaddleX OCR处理失败: {result['error']}" + log_ocr_request("paddlex_ocr", pdf_path, False, processing_time, error_msg) + + raise OCRServiceException( + error_msg, + "paddlex_ocr", + "processing_failed" + ) + + log_ocr_request("paddlex_ocr", pdf_path, True, processing_time) + return result["full_text"] + + except Exception as e: + if isinstance(e, OCRServiceException): + raise + processing_time = time.time() - start_time if 'start_time' in locals() else 0 + error_msg = f"PaddleX OCR处理失败: {str(e)}" + log_ocr_request("paddlex_ocr", pdf_path, False, processing_time, error_msg) + + raise OCRServiceException( + error_msg, + "paddlex_ocr", + "processing_failed" + ) def get_state(task_id): return GOLBAL_STATE.get(task_id, {}) diff --git a/web/src/apis/admin_api.js b/web/src/apis/admin_api.js index e4e6202c..4a7220ff 100644 --- a/web/src/apis/admin_api.js +++ b/web/src/apis/admin_api.js @@ -390,3 +390,24 @@ export const adminApi = { return apiPost(url, data, {}, true) }, } + +// OCR服务管理API +export const ocrApi = { + /** + * 检查OCR服务健康状态 + * @returns {Promise} - OCR服务健康状态信息 + */ + checkHealth: async () => { + checkAdminPermission() + return apiGet('/api/ocr/health', {}, true) + }, + + /** + * 获取OCR服务使用统计 + * @returns {Promise} - OCR服务统计信息 + */ + getStats: async () => { + checkAdminPermission() + return apiGet('/api/ocr/stats', {}, true) + } +} diff --git a/web/src/views/DataBaseInfoView.vue b/web/src/views/DataBaseInfoView.vue index c60f96a1..96713ddc 100644 --- a/web/src/views/DataBaseInfoView.vue +++ b/web/src/views/DataBaseInfoView.vue @@ -97,9 +97,34 @@