diff --git a/api/__init__.py b/api/__init__.py index 5606afa..864a92f 100644 --- a/api/__init__.py +++ b/api/__init__.py @@ -67,7 +67,17 @@ def create_app() -> 'Flask': static_folder = os.path.join(PROJECT_ROOT, 'chat-ui') app = Flask(__name__, static_folder=static_folder, static_url_path='') - CORS(app) + + # CORS 配置:生产环境限制来源,开发环境允许全部 + if IS_PROD: + cors_origins = os.environ.get('CORS_ORIGINS', '').split(',') if os.environ.get('CORS_ORIGINS') else [] + if cors_origins: + CORS(app, origins=cors_origins) + else: + CORS(app) # 未配置时仍允许全部,但记录警告 + logger.warning("生产环境未配置 CORS_ORIGINS,CORS 允许所有来源") + else: + CORS(app) # ==================== Repository 依赖注入 ==================== @@ -227,14 +237,14 @@ def _validate_production_config() -> None: - DASHSCOPE_API_KEY: 大模型调用必需 Raises: - AssertionError: 缺少必需的配置项 + RuntimeError: 缺少必需的配置项 """ import os from config import DASHSCOPE_API_KEY # 检查环境变量或配置文件中的 API Key has_key = os.getenv("DASHSCOPE_API_KEY") or os.environ.get("DASHSCOPE_API_KEY") or DASHSCOPE_API_KEY - assert has_key and has_key != "", \ - "Missing DASHSCOPE_API_KEY in production environment" + if not has_key or has_key == "": + raise RuntimeError("Missing DASHSCOPE_API_KEY in production environment") logger.info("Configuration validated") diff --git a/api/audit_routes.py b/api/audit_routes.py index e9b038c..9fefe4e 100644 --- a/api/audit_routes.py +++ b/api/audit_routes.py @@ -97,7 +97,8 @@ def get_audit_logs(): return jsonify({"logs": logs, "total": total}) except Exception as e: - return jsonify({"error": str(e), "logs": [], "total": 0}), 500 + logger.error(f"审计查询异常: {e}") + return jsonify({"error": "查询失败", "logs": [], "total": 0}), 500 def log_audit_event(user_id: str, username: str, action: str, diff --git a/api/chat_routes.py b/api/chat_routes.py index ec0adb5..7a389bd 100644 --- a/api/chat_routes.py +++ b/api/chat_routes.py @@ -41,6 +41,7 @@ from auth.gateway import require_gateway_auth from auth.security import validate_query, filter_response from config import RAG_CHAT_MODEL from core.llm_utils import call_llm, call_llm_stream +from core.prompt_guard import detect_injection, sanitize_user_input def _get_vlm_cache(image_path: str) -> Optional[str]: @@ -1953,11 +1954,11 @@ def rag(): yield f"data: {json.dumps(finish_event, ensure_ascii=False)}\n\n" except Exception as e: - import traceback + import logging as _logging + _logging.getLogger(__name__).error(f"[SSE] RAG 流异常: {e}", exc_info=True) error_event = { "type": "error", - "message": str(e), - "traceback": traceback.format_exc() + "message": "服务内部错误,请稍后重试" } yield f"data: {json.dumps(error_event, ensure_ascii=False)}\n\n" @@ -1986,12 +1987,27 @@ def search(): """ data = request.json or {} query = data.get('query', '') + query = sanitize_user_input(query) + injection_matches = detect_injection(query) + if injection_matches: + logger.warning(f"[Chat] 检测到可疑注入: {injection_matches}") top_k = data.get('top_k', 5) collections = data.get('collections') # 后端传入的知识库列表 if not query: return jsonify({'error': 'query is required'}), 400 + # 输入安全校验(注入检测、违禁词、长度限制) + is_valid, reason = validate_query(query) + if not is_valid: + return jsonify({'error': reason}), 400 + + # top_k 范围校验 + try: + top_k = max(1, min(int(top_k), 50)) + except (ValueError, TypeError): + top_k = 5 + # 如果没有指定 collections,使用默认的公开库 if not collections: collections = ['public_kb'] diff --git a/api/document_routes.py b/api/document_routes.py index 5860c1e..0aae6e2 100644 --- a/api/document_routes.py +++ b/api/document_routes.py @@ -96,6 +96,24 @@ def safe_filename(filename: str) -> str: return safe_name + +def _validate_doc_path(doc_path: str, documents_path: str) -> str: + """ + 校验文档路径是否在允许目录内,防止路径遍历攻击。 + + Returns: + 解析后的安全绝对路径 + Raises: + ValueError: 路径不合法 + """ + filepath = os.path.join(documents_path, doc_path) + real_path = os.path.realpath(filepath) + real_base = os.path.realpath(documents_path) + if not real_path.startswith(real_base + os.sep) and real_path != real_base: + raise ValueError("非法路径") + return real_path + + # 延迟初始化缓存 _kb_manager = None _kb_checked = False @@ -160,6 +178,10 @@ def serve_document_file(doc_path: str) -> Tuple[Any, int]: from flask import send_from_directory filepath = os.path.join(DOCUMENTS_PATH, doc_path) + try: + filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH) + except ValueError: + return jsonify({"error": "非法路径"}), 403 if not os.path.exists(filepath): return jsonify({"error": "文件不存在"}), 404 @@ -329,7 +351,8 @@ def upload_document() -> Tuple[Any, int]: sync_service.process_change(change) sync_status = "已保存并添加到向量库" except Exception as e: - sync_status = f"已保存,向量化失败: {str(e)}" + logger.warning(f"向量化失败: {e}") + sync_status = "已保存,向量化失败" return success_response( data={ @@ -404,6 +427,18 @@ def batch_upload_documents() -> Tuple[Any, int]: }) continue + # 文件大小校验 + file.seek(0, 2) + file_size = file.tell() + file.seek(0) + if file_size > MAX_FILE_SIZE: + results.append({ + "filename": file.filename, + "status": "error", + "message": f"文件过大(最大 {MAX_FILE_SIZE // 1024 // 1024}MB)" + }) + continue + try: original_filename = file.filename ext = os.path.splitext(original_filename)[1].lower() @@ -464,10 +499,11 @@ def batch_upload_documents() -> Tuple[Any, int]: "replaced": replaced }) except Exception as e: + logger.error(f"上传处理异常: {e}") results.append({ "filename": file.filename, "status": "error", - "message": str(e) + "message": "上传处理失败" }) return success_response( @@ -644,12 +680,27 @@ def update_document(doc_path: str) -> Tuple[Any, int]: if file.filename == '': return jsonify({"error": "没有选择文件"}), 400 + # 文件类型校验 + ext = os.path.splitext(file.filename)[1].lower() + if ext not in ALLOWED_EXTENSIONS: + return jsonify({"error": f"不支持的文件类型: {ext}"}), 400 + + # 文件大小校验 + file.seek(0, 2) # 跳到文件末尾获取大小 + file_size = file.tell() + file.seek(0) # 回到文件开头 + if file_size > MAX_FILE_SIZE: + return jsonify({"error": f"文件过大(最大 {MAX_FILE_SIZE // 1024 // 1024}MB)"}), 400 + # 解析路径 parts = doc_path.split('/') if len(parts) < 2: return jsonify({"error": "无效的文档路径"}), 400 - filepath = os.path.join(DOCUMENTS_PATH, doc_path) + try: + filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH) + except ValueError: + return jsonify({"error": "非法路径"}), 403 if not os.path.exists(filepath): return jsonify({"error": "文件不存在"}), 404 @@ -711,7 +762,10 @@ def delete_document(doc_path: str) -> Tuple[Any, int]: # 目录名即向量库名 collection = subdir - filepath = os.path.join(DOCUMENTS_PATH, doc_path) + try: + filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH) + except ValueError: + return jsonify({"error": "非法路径"}), 403 if not os.path.exists(filepath): return jsonify({"error": "文件不存在"}), 404 @@ -730,7 +784,8 @@ def delete_document(doc_path: str) -> Tuple[Any, int]: }) except Exception as e: - return jsonify({"error": f"删除失败: {str(e)}"}), 500 + logger.error(f"删除文档异常: {e}") + return jsonify({"error": "删除失败,请稍后重试"}), 500 @document_bp.route('/documents//chunks', methods=['GET']) @@ -817,7 +872,10 @@ def preview_document(doc_path: str) -> Tuple[Any, int]: # 查询参数 chunk_index_str = request.args.get('chunk_index') - context_count = int(request.args.get('context', 2)) + try: + context_count = max(0, min(int(request.args.get('context', 2)), 10)) + except (ValueError, TypeError): + context_count = 2 # 获取所有切片 all_chunks = kb_manager.get_document_chunks(collection, filename) @@ -1082,5 +1140,5 @@ def delete_chunks_by_source() -> Tuple[Any, int]: "message": f"已删除 {deleted_count} 个切片" }) except Exception as e: - logger.error(f"批量删除切片失败: {e}") - return jsonify({"error": str(e)}), 500 + logger.error(f"操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 diff --git a/api/feedback_routes.py b/api/feedback_routes.py index 5608450..5b677f0 100644 --- a/api/feedback_routes.py +++ b/api/feedback_routes.py @@ -33,6 +33,9 @@ from flask import Blueprint, request, jsonify from auth.gateway import require_gateway_auth, require_role +import logging + +logger = logging.getLogger(__name__) feedback_bp = Blueprint('feedback', __name__) @@ -95,7 +98,8 @@ def submit_feedback(): "suggestion_id": result.get('suggestion_id') }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/feedback/stats', methods=['GET']) @@ -113,7 +117,8 @@ def get_feedback_stats(): "stats": stats }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/feedback/list', methods=['GET']) @@ -141,7 +146,8 @@ def get_feedback_list(): "total": len(feedbacks) }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/reports/weekly', methods=['GET']) @@ -156,7 +162,8 @@ def get_weekly_report(): "report": report.to_dict() }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/reports/monthly', methods=['GET']) @@ -171,7 +178,8 @@ def get_monthly_report(): "report": report.to_dict() }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/faq', methods=['GET']) @@ -190,7 +198,8 @@ def get_faq_list(): "total": len(faqs) }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/faq', methods=['POST']) @@ -233,7 +242,8 @@ def create_faq(): "message": "FAQ已创建,请通过 /faq//approve 接口确认后生效" }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/faq//approve', methods=['POST']) @@ -276,7 +286,8 @@ def approve_faq(faq_id): "message": "FAQ已批准并同步到知识库" }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/faq/', methods=['PUT']) @@ -326,7 +337,8 @@ def update_faq(faq_id): "sync_status": sync_status }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/faq/', methods=['DELETE']) @@ -358,7 +370,8 @@ def delete_faq(faq_id): "message": "FAQ删除成功" if deleted else "FAQ不存在" }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/faq/suggestions', methods=['GET']) @@ -378,7 +391,8 @@ def get_faq_suggestions(): "total": len(suggestions) }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/faq/suggestions//approve', methods=['POST']) @@ -413,7 +427,8 @@ def approve_faq_suggestion(suggestion_id): else: return jsonify({"error": result.get('error', '批准失败')}), 400 except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/faq/suggestions//reject', methods=['POST']) @@ -429,7 +444,8 @@ def reject_faq_suggestion(suggestion_id): "message": "FAQ建议已拒绝" if rejected else "建议不存在" }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 # ==================== Bad Case 分析接口 ==================== @@ -472,7 +488,8 @@ def get_bad_cases(): ] }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 @feedback_bp.route('/feedback/blacklist', methods=['GET']) @@ -497,4 +514,5 @@ def get_chunk_blacklist(): "usage": "在检索时过滤这些来源以提升回答质量" }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"反馈操作异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 diff --git a/api/image_routes.py b/api/image_routes.py index 824345b..ae8b2f1 100644 --- a/api/image_routes.py +++ b/api/image_routes.py @@ -9,8 +9,11 @@ """ import os +import logging from flask import Blueprint, send_file, jsonify, current_app +logger = logging.getLogger(__name__) + image_bp = Blueprint('images', __name__) @@ -70,7 +73,8 @@ def get_image(image_id: str): return send_file(image_path, mimetype=mimetype) except Exception as e: - return jsonify({"error": f"读取图片失败: {str(e)}"}), 500 + logger.error(f"读取图片异常: {e}") + return jsonify({"error": "读取图片失败"}), 500 return jsonify({"error": "图片不存在", "image_id": image_id}), 404 @@ -122,7 +126,8 @@ def get_image_info(image_id: str): "url": f"/images/{image_id}" }) except Exception as e: - return jsonify({"error": f"读取图片信息失败: {str(e)}"}), 500 + logger.error(f"读取图片信息异常: {e}") + return jsonify({"error": "读取图片信息失败"}), 500 return jsonify({"error": "图片不存在", "image_id": image_id}), 404 @@ -179,7 +184,8 @@ def list_images(): }) except Exception as e: - return jsonify({"error": f"列出图片失败: {str(e)}"}), 500 + logger.error(f"列出图片异常: {e}") + return jsonify({"error": "列出图片失败"}), 500 @image_bp.route('/images/stats', methods=['GET']) @@ -222,4 +228,5 @@ def image_stats(): }) except Exception as e: - return jsonify({"error": f"获取统计信息失败: {str(e)}"}), 500 + logger.error(f"获取统计信息异常: {e}") + return jsonify({"error": "获取统计信息失败"}), 500 diff --git a/api/kb_routes.py b/api/kb_routes.py index 4b365d7..9f5f886 100644 --- a/api/kb_routes.py +++ b/api/kb_routes.py @@ -391,10 +391,11 @@ def sync_documents() -> Tuple[Any, int]: } }) except Exception as e: + logger.error(f"知识库操作异常: {e}") results.append({ "collection": "all", "status": "error", - "message": str(e) + "message": "操作失败" }) else: # 没有 sync_service,返回提示 @@ -462,7 +463,8 @@ def debug_scan() -> Tuple[Any, int]: result["scanned_count"] = len(scanned) result["scanned_ids"] = list(scanned.keys())[:10] except Exception as e: - result["scan_error"] = str(e) + logger.error(f"扫描异常: {e}") + result["scan_error"] = "扫描失败" return jsonify(result) @@ -499,9 +501,11 @@ def reindex_collection(kb_name: str) -> Tuple[Any, int]: from data.db import get_connection with get_connection("knowledge") as conn: cursor = conn.cursor() + # 转义 LIKE 通配符,防止 kb_name 中的 % 或 _ 导致非预期匹配 + escaped_kb = kb_name.replace('%', '\\%').replace('_', '\\_') # 删除以 "{kb_name}/" 或 "{kb_name}\" 开头的文档哈希(兼容 Windows 和 Linux) - cursor.execute("DELETE FROM document_hashes WHERE document_id LIKE ? OR document_id LIKE ?", - (f"{kb_name}/%", f"{kb_name}\\%")) + cursor.execute("DELETE FROM document_hashes WHERE document_id LIKE ? ESCAPE '\\' OR document_id LIKE ? ESCAPE '\\'", + (f"{escaped_kb}/%", f"{escaped_kb}\\%")) deleted = cursor.rowcount logger.info(f"已清除 {deleted} 条哈希记录: {kb_name}") except Exception as e: @@ -520,7 +524,8 @@ def reindex_collection(kb_name: str) -> Tuple[Any, int]: "errors": result.errors }) except Exception as e: - return jsonify({"error": str(e)}), 500 + logger.error(f"reindex 异常: {e}") + return jsonify({"error": "操作失败,请稍后重试"}), 500 else: return jsonify({"error": "同步服务不可用"}), 503 @@ -633,7 +638,8 @@ def deprecate_document(kb_name: str, filename: str) -> Tuple[Any, int]: ) return jsonify(result) except Exception as e: - return jsonify({"success": False, "error": str(e)}), 500 + logger.error(f"操作异常: {e}") + return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500 @kb_bp.route('/collections//documents//restore', methods=['POST']) @@ -664,7 +670,8 @@ def restore_document(kb_name: str, filename: str) -> Tuple[Any, int]: result = kb_manager.restore_document(kb_name, filename) return jsonify(result) except Exception as e: - return jsonify({"success": False, "error": str(e)}), 500 + logger.error(f"操作异常: {e}") + return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500 @kb_bp.route('/collections//documents//versions', methods=['GET']) @@ -720,7 +727,8 @@ def get_document_versions(kb_name: str, filename: str) -> Tuple[Any, int]: "total": len(versions_data) }) except Exception as e: - return jsonify({"success": False, "error": str(e)}), 500 + logger.error(f"操作异常: {e}") + return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500 @kb_bp.route('/collections//update-image-descriptions', methods=['POST']) @@ -754,7 +762,8 @@ def update_image_descriptions(kb_name: str) -> Tuple[Any, int]: result = kb_manager.update_image_descriptions(kb_name) return jsonify(result) except Exception as e: - return jsonify({"success": False, "error": str(e)}), 500 + logger.error(f"操作异常: {e}") + return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500 @kb_bp.route('/collections/sync-vlm-cache', methods=['POST']) diff --git a/api/session_routes.py b/api/session_routes.py index 40a79bd..864075e 100644 --- a/api/session_routes.py +++ b/api/session_routes.py @@ -33,6 +33,8 @@ def get_sessions(): } """ session_manager = current_app.config['SESSION_MANAGER'] + if session_manager is None: + return jsonify({"error": "会话服务不可用"}), 503 user_id = request.current_user["user_id"] sessions = session_manager.get_user_sessions(user_id, limit=20) @@ -62,6 +64,8 @@ def get_history(session_id): } """ session_manager = current_app.config['SESSION_MANAGER'] + if session_manager is None: + return jsonify({"error": "会话服务不可用"}), 503 user_id = request.current_user["user_id"] # 验证会话归属 @@ -81,6 +85,8 @@ def get_history(session_id): def delete_session(session_id): """删除会话""" session_manager = current_app.config['SESSION_MANAGER'] + if session_manager is None: + return jsonify({"error": "会话服务不可用"}), 503 user_id = request.current_user["user_id"] # 验证会话归属 @@ -100,6 +106,8 @@ def delete_session(session_id): def clear_history(session_id): """清空会话历史(保留会话)""" session_manager = current_app.config['SESSION_MANAGER'] + if session_manager is None: + return jsonify({"error": "会话服务不可用"}), 503 user_id = request.current_user["user_id"] # 验证会话归属 diff --git a/api/sync_routes.py b/api/sync_routes.py index 8deced2..55dce3c 100644 --- a/api/sync_routes.py +++ b/api/sync_routes.py @@ -112,10 +112,11 @@ def trigger_sync() -> Tuple[Any, int]: message="同步完成" ) except Exception as e: + logger.error(f"同步操作异常: {e}") return error_response( error="SYNC_ERROR", error_code=SYNC_ERROR, - message=str(e), + message="操作失败", http_status=500 ) @@ -164,11 +165,12 @@ def get_sync_status() -> Tuple[Any, int]: return jsonify(status) except Exception as e: + logger.error(f"状态查询异常: {e}") return jsonify({ "status": "failed", "status_code": INTERNAL_ERROR, "enabled": True, - "error": str(e) + "error": "操作失败" }) @@ -196,10 +198,11 @@ def get_sync_history() -> Tuple[Any, int]: history = service.get_sync_history(limit=limit) if hasattr(service, 'get_sync_history') else [] return jsonify({"history": history}) except Exception as e: + logger.error(f"同步操作异常: {e}") return error_response( error="SYNC_ERROR", error_code=SYNC_ERROR, - message=str(e), + message="操作失败", http_status=500 ) @@ -230,10 +233,11 @@ def get_change_logs() -> Tuple[Any, int]: changes = service.get_change_logs(limit=limit, collection=collection) if hasattr(service, 'get_change_logs') else [] return jsonify({"changes": changes}) except Exception as e: + logger.error(f"同步操作异常: {e}") return error_response( error="SYNC_ERROR", error_code=SYNC_ERROR, - message=str(e), + message="操作失败", http_status=500 ) @@ -275,10 +279,11 @@ def start_sync_monitor() -> Tuple[Any, int]: else: return jsonify({"status": "success", "message": "文件监控功能不可用"}) except Exception as e: + logger.error(f"同步操作异常: {e}") return error_response( error="SYNC_ERROR", error_code=SYNC_ERROR, - message=str(e), + message="操作失败", http_status=500 ) @@ -303,9 +308,10 @@ def stop_sync_monitor() -> Tuple[Any, int]: service.stop() return jsonify({"status": "success", "status_code": SYNC_SUCCESS, "message": "文件监控已停止"}) except Exception as e: + logger.error(f"同步操作异常: {e}") return error_response( error="SYNC_ERROR", error_code=SYNC_ERROR, - message=str(e), + message="操作失败", http_status=500 ) diff --git a/core/bm25_index.py b/core/bm25_index.py index 1a9b2fd..0bbfe74 100644 --- a/core/bm25_index.py +++ b/core/bm25_index.py @@ -14,6 +14,7 @@ BM25 关键词检索索引 import os import pickle +import threading import numpy as np from rank_bm25 import BM25Okapi import jieba @@ -105,24 +106,27 @@ class BM25Index: # ==================== 全局 BM25 索引管理器 ==================== _bm25_indexer: BM25Index = None +_bm25_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件 def get_bm25_indexer() -> BM25Index: """ - 获取全局 BM25 索引器实例 + 获取全局 BM25 索引器实例(双重检查锁定,线程安全) Returns: BM25Index 实例 """ global _bm25_indexer if _bm25_indexer is None: - _bm25_indexer = BM25Index() + with _bm25_lock: + if _bm25_indexer is None: + _bm25_indexer = BM25Index() return _bm25_indexer def init_bm25_indexer(ids=None, documents=None, metadatas=None) -> BM25Index: """ - 初始化 BM25 索引器并添加文档 + 初始化 BM25 索引器并添加文档(加锁,线程安全) Args: ids: 文档 ID 列表 @@ -133,7 +137,8 @@ def init_bm25_indexer(ids=None, documents=None, metadatas=None) -> BM25Index: 初始化后的 BM25Index 实例 """ global _bm25_indexer - _bm25_indexer = BM25Index() - if ids and documents: - _bm25_indexer.add_documents(ids, documents, metadatas or []) + with _bm25_lock: + _bm25_indexer = BM25Index() + if ids and documents: + _bm25_indexer.add_documents(ids, documents, metadatas or []) return _bm25_indexer diff --git a/core/engine.py b/core/engine.py index 5092344..b4468b6 100644 --- a/core/engine.py +++ b/core/engine.py @@ -31,6 +31,7 @@ import os import gc import time import logging +import threading import numpy as np from typing import List, Tuple, Optional, Dict, Any @@ -50,6 +51,7 @@ except ImportError: # 延迟导入,防止循环依赖 _engine_instance = None +_engine_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件 try: from config import ( @@ -327,10 +329,12 @@ class RAGEngine: @classmethod def get_instance(cls) -> 'RAGEngine': - """获取引擎单例实例""" + """获取引擎单例实例(双重检查锁定,线程安全)""" global _engine_instance if _engine_instance is None: - _engine_instance = cls() + with _engine_lock: + if _engine_instance is None: + _engine_instance = cls() return _engine_instance def __init__(self) -> None: diff --git a/core/intent_analyzer.py b/core/intent_analyzer.py index 2ca3a17..5f2dea2 100644 --- a/core/intent_analyzer.py +++ b/core/intent_analyzer.py @@ -18,6 +18,7 @@ import json import logging import hashlib +import threading from dataclasses import dataclass, asdict from typing import List, Dict, Optional, Tuple @@ -503,12 +504,15 @@ class IntentAnalyzer: # ==================== 便捷函数 ==================== _analyzer = None +_analyzer_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件 def get_intent_analyzer() -> IntentAnalyzer: - """获取意图分析器单例""" + """获取意图分析器单例(双重检查锁定,线程安全)""" global _analyzer if _analyzer is None: - _analyzer = IntentAnalyzer() + with _analyzer_lock: + if _analyzer is None: + _analyzer = IntentAnalyzer() return _analyzer diff --git a/core/llm_utils.py b/core/llm_utils.py index 0473bba..d4af40f 100644 --- a/core/llm_utils.py +++ b/core/llm_utils.py @@ -237,6 +237,76 @@ def parse_json_list_from_response(content: str) -> Optional[List[dict]]: return None +def extract_json_object(content: str) -> Optional[dict]: + """ + 多策略从 LLM 响应中提取 JSON 对象(增强版) + + 在 parse_json_from_response 基础上增加 fallback 策略: + 1. 先调用 parse_json_from_response(markdown 代码块 → 直接解析) + 2. 失败后 fallback 到正则匹配最外层 {...} 块 + + Args: + content: LLM 返回的原始内容 + + Returns: + 解析后的字典,全部策略失败返回 None + """ + if not content: + return None + + # 策略1+2:markdown 代码块提取 + 直接 json.loads + result = parse_json_from_response(content) + if result is not None and isinstance(result, dict): + return result + + # 策略3(fallback):正则匹配最外层 JSON 对象 {...} + brace_match = re.search(r'\{[\s\S]*\}', content) + if brace_match: + try: + parsed = json.loads(brace_match.group(0)) + if isinstance(parsed, dict): + return parsed + except (json.JSONDecodeError, TypeError, ValueError): + pass + + return None + + +def extract_json_list(content: str) -> Optional[list]: + """ + 多策略从 LLM 响应中提取 JSON 数组(增强版) + + 在 parse_json_list_from_response 基础上增加 fallback 策略: + 1. 先调用 parse_json_list_from_response(markdown 代码块 → 直接解析 → 嵌套提取) + 2. 失败后 fallback 到正则匹配最外层 [...] 块 + + Args: + content: LLM 返回的原始内容 + + Returns: + 解析后的列表,全部策略失败返回 None + """ + if not content: + return None + + # 策略1+2:markdown 代码块提取 + 直接 json.loads + 嵌套 key 提取 + result = parse_json_list_from_response(content) + if result is not None: + return result + + # 策略3(fallback):正则匹配最外层 JSON 数组 [...] + bracket_match = re.search(r'\[[\s\S]*\]', content) + if bracket_match: + try: + parsed = json.loads(bracket_match.group(0)) + if isinstance(parsed, list): + return parsed + except (json.JSONDecodeError, TypeError, ValueError): + pass + + return None + + # ==================== 便捷函数 ==================== def quick_ask( diff --git a/core/prompt_guard.py b/core/prompt_guard.py new file mode 100644 index 0000000..f87132d --- /dev/null +++ b/core/prompt_guard.py @@ -0,0 +1,104 @@ +""" +Prompt 注入防御工具 + +提供轻量级的 prompt 注入检测和清理功能,防止恶意用户通过精心构造的输入 +覆盖系统指令或操纵 LLM 行为。 +""" + +import re +import logging + +logger = logging.getLogger(__name__) + +# 常见注入模式(中英文) +_INJECTION_PATTERNS = [ + # 英文注入模式 + r'ignore\s+(all\s+)?(previous|above|prior)\s+(instructions?|prompts?|rules?)', + r'you\s+are\s+now\s+(a|an)\s+', + r'forget\s+(everything|all|your)\s+', + r'disregard\s+(all\s+)?(previous|prior|above)', + r'override\s+(your|the|all)\s+(instructions?|rules?|system)', + r'reveal\s+(your|the)\s+(system\s+)?(prompt|instructions?)', + r'act\s+as\s+(if\s+)?(you\s+are\s+)?', + # 中文注入模式 + r'忽略(之前|以上|前面|所有).*(指令|规则|提示|约束)', + r'你现在是(一个|新的)?', + r'忘记(之前|所有|你).*(规则|指令|设定)', + r'无视(系统|所有|之前).*(指令|规则|提示)', + r'不要遵守(任何|之前|系统).*(指令|规则)', +] + +# 编译正则表达式(不区分大小写) +_COMPILED_PATTERNS = [ + re.compile(p, re.IGNORECASE) for p in _INJECTION_PATTERNS +] + + +def detect_injection(text: str, threshold: int = 1) -> list: + """ + 检测文本中是否包含 prompt 注入模式 + + Args: + text: 待检测的用户输入文本 + threshold: 匹配数量阈值,达到此数量视为可疑 + + Returns: + 匹配到的注入模式描述列表(空列表表示安全) + """ + if not text or not text.strip(): + return [] + + matches = [] + for pattern in _COMPILED_PATTERNS: + match = pattern.search(text) + if match: + matches.append(match.group()[:50]) + + if len(matches) >= threshold: + logger.warning(f"[PromptGuard] 检测到可疑注入模式: {matches[:3]}, 输入前100字: {text[:100]}") + + return matches + + +def sanitize_user_input(text: str, max_length: int = 2000) -> str: + """ + 清理用户输入,去除潜在的控制指令 + + 策略: + 1. 截断超长输入 + 2. 去除 null 字节和控制字符 + 3. 限制输入长度 + + Args: + text: 原始用户输入 + max_length: 最大长度 + + Returns: + 清理后的安全文本 + """ + if not text: + return "" + + # 去除 null 字节和控制字符(保留换行和空格) + text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text) + + # 截断超长输入 + if len(text) > max_length: + text = text[:max_length] + logger.debug(f"[PromptGuard] 输入超长,截断至 {max_length} 字") + + return text + + +def wrap_document_content(content: str, label: str = "参考资料") -> str: + """ + 将文档内容包装在明确的标记内,降低文档内容被当作指令执行的风险 + + Args: + content: 文档原始内容 + label: 标签名称 + + Returns: + 包装后的内容 + """ + return f"[以下为{label},非系统指令]\n---\n{content}\n---\n[{label}结束]" diff --git a/exam_pkg/grader.py b/exam_pkg/grader.py index 6637750..c2b6496 100644 --- a/exam_pkg/grader.py +++ b/exam_pkg/grader.py @@ -16,7 +16,6 @@ import json import logging -import re import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed @@ -24,7 +23,7 @@ from functools import wraps from typing import List, Dict, Any, Optional # 导入 LLM 工具函数 -from core.llm_utils import call_llm +from core.llm_utils import call_llm, extract_json_object logger = logging.getLogger(__name__) @@ -378,47 +377,12 @@ class AnswerGrader: return result def _extract_json(self, response: str) -> dict: - """ - 多策略从 LLM 响应中提取 JSON 对象 - - 策略优先级: - 1. markdown 代码块提取 ```json ... ``` - 2. 直接 json.loads - 3. 正则匹配第一个 {...} 块 - """ + """多策略从 LLM 响应中提取 JSON 对象(使用共享工具)""" + result = extract_json_object(response) + if result is not None: + return result if not response: raise ValueError("LLM 返回为空") - - # 策略1:提取 markdown 代码块 - json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', response) - if json_match: - json_str = json_match.group(1).strip() - try: - result = json.loads(json_str) - if isinstance(result, dict): - return result - except json.JSONDecodeError: - pass # 继续下一策略 - - # 策略2:直接解析整个响应 - try: - result = json.loads(response.strip()) - if isinstance(result, dict): - return result - except json.JSONDecodeError: - pass # 继续下一策略 - - # 策略3:正则匹配最外层 JSON 对象 - brace_match = re.search(r'\{[\s\S]*\}', response) - if brace_match: - try: - result = json.loads(brace_match.group(0)) - if isinstance(result, dict): - return result - except json.JSONDecodeError: - pass - - # 全部策略失败 raise json.JSONDecodeError( f"无法从 LLM 响应中提取有效 JSON,响应前300字: {response[:300]}", response, 0 diff --git a/main.py b/main.py index f2cf666..5b98914 100644 --- a/main.py +++ b/main.py @@ -28,7 +28,7 @@ def main(): parser = argparse.ArgumentParser(description='RAG API 服务') parser.add_argument('--host', default='0.0.0.0', help='监听地址(默认 0.0.0.0)') parser.add_argument('--port', type=int, default=5001, help='监听端口(默认 5001)') - parser.add_argument('--debug', action='store_true', default=True, help='调试模式') + parser.add_argument('--debug', action='store_true', default=False, help='调试模式') parser.add_argument('--no-debug', action='store_true', help='关闭调试模式') args = parser.parse_args() diff --git a/parsers/excel_parser.py b/parsers/excel_parser.py index 30a1a31..9442556 100644 --- a/parsers/excel_parser.py +++ b/parsers/excel_parser.py @@ -24,6 +24,9 @@ logger = logging.getLogger(__name__) # 大表切片阈值 MAX_ROWS_PER_CHUNK = 200 +# 文件大小限制 +MAX_FILE_SIZE = 50 * 1024 * 1024 # 50MB + @dataclass class UnifiedChunk: @@ -67,6 +70,13 @@ def parse_excel( } """ filepath = Path(filepath) + + # 检查文件大小 + if filepath.exists(): + file_size = filepath.stat().st_size + if file_size > MAX_FILE_SIZE: + raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB,最大允许 {MAX_FILE_SIZE / 1024 / 1024:.0f}MB") + if not filepath.exists(): raise FileNotFoundError(f"文件不存在: {filepath}") diff --git a/parsers/mineru_parser.py b/parsers/mineru_parser.py index a83f67c..62c3bfc 100644 --- a/parsers/mineru_parser.py +++ b/parsers/mineru_parser.py @@ -46,6 +46,9 @@ import logging logger = logging.getLogger(__name__) +# 文件大小限制 +MAX_PDF_SIZE = 100 * 1024 * 1024 # 100MB + # 支持的文件格式 SUPPORTED_FORMATS = { '.pdf': 'PDF 文档', @@ -176,6 +179,11 @@ def parse_with_mineru_online( if not file_path.exists(): raise FileNotFoundError(f"文件不存在: {file_path}") + # 检查文件大小 + file_size = file_path.stat().st_size + if file_size > MAX_PDF_SIZE: + raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB,最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB") + logger.info(f"使用 MinerU 在线 API 解析: {file_path.name}") # 读取文件内容 @@ -600,6 +608,11 @@ def parse_with_mineru( if not file_path.exists(): raise FileNotFoundError(f"文件不存在: {file_path}") + # 检查文件大小 + file_size = file_path.stat().st_size + if file_size > MAX_PDF_SIZE: + raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB,最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB") + # 检查文件格式 suffix = file_path.suffix.lower() if suffix not in SUPPORTED_FORMATS: @@ -629,8 +642,17 @@ def parse_with_mineru( if not mineru_exe.exists(): mineru_exe = "mineru" # 回退到系统 PATH + # 参数白名单校验,防止注入非法参数 + ALLOWED_BACKENDS = {'auto', 'pipeline', 'vlm', 'vlm-sglang', 'vlm-auto-engine', 'hybrid-auto-engine', 'ocr'} + ALLOWED_LANGS = {'ch', 'en', 'ch_lite', 'en_lite', 'formula', 'table'} + if backend not in ALLOWED_BACKENDS: + backend = 'pipeline' + if lang not in ALLOWED_LANGS: + lang = 'ch' + cmd = [ str(mineru_exe), + "--", "-p", str(file_path), "-o", str(output_dir), "-m", "auto", @@ -1458,6 +1480,11 @@ def parse_with_mineru_persistent( if not file_path.exists(): raise FileNotFoundError(f"文件不存在: {file_path}") + # 检查文件大小 + file_size = file_path.stat().st_size + if file_size > MAX_PDF_SIZE: + raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB,最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB") + # 计算文件 hash,用于隔离输出目录 file_hash = compute_file_hash(str(file_path)) output_dir = Path(output_base) / file_hash diff --git a/parsers/txt_parser.py b/parsers/txt_parser.py index f91b05a..1a507b0 100644 --- a/parsers/txt_parser.py +++ b/parsers/txt_parser.py @@ -5,12 +5,25 @@ TXT 文本解析器 """ import logging +import os logger = logging.getLogger(__name__) +# 文件大小限制 +MAX_TXT_SIZE = 20 * 1024 * 1024 # 20MB + def extract_text_from_txt(filepath): """从TXT提取文本""" + # 检查文件大小 + try: + file_size = os.path.getsize(filepath) + if file_size > MAX_TXT_SIZE: + logger.error(f"TXT文件过大: {file_size / 1024 / 1024:.1f}MB") + return "" + except OSError: + return "" + try: with open(filepath, 'r', encoding='utf-8') as f: return f.read()