hack/tools/scripts/llm-key-hunter/results/_cache/content/2692341798_load-balanced-online-oj.json

{"repo": "2692341798/load-balanced-online-oj", "files": {"scripts/add_test_data.py": {"81849c36ce54040e93d629b9e95fff0b08073f12": {"t": 1786068248.9991286, "n": 6646, "h": "3744b25b77b693b3", "c": "import os\nimport json\nimport time\nimport re\nimport pymysql\nimport requests\nimport logging\n\n# 配置日志格式\nlogging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n\n# MySQL 数据库配置\nDB_CONFIG = {\n    'host': '127.0.0.1',\n    'user': 'oj_client',\n    'password': '123456',\n    'database': 'oj',\n    'port': 3306,\n    'charset': 'utf8mb4',\n    'cursorclass': pymysql.cursors.DictCursor\n}\n\n# DeepSeek API 配置\nDEEPSEEK_API_KEY = os.getenv('DEEPSEEK_API_KEY', '')\nDEEPSEEK_API_URL = os.getenv('DEEPSEEK_API_URL', 'https://api.deepseek.com/chat/completions').rstrip('/')\n\n# 进度保存文件与相关配置\nPROGRESS_FILE = 'scripts/add_test_progress.json'\nMAX_RETRIES = 3\nRATE_LIMIT_DELAY = 2  # 每次请求后休眠2秒\n\ndef load_progress():\n    \"\"\"加载已处理过的题目编号列表\"\"\"\n    if os.path.exists(PROGRESS_FILE):\n        try:\n            with open(PROGRESS_FILE, 'r', encoding='utf-8') as f:\n                return set(json.load(f))\n        except Exception as e:\n            logging.error(f\"无法加载进度文件: {e}\")\n    return set()\n\ndef save_progress(progress_set):\n    \"\"\"保存已处理的题目编号列表到文件\"\"\"\n    os.makedirs(os.path.dirname(PROGRESS_FILE), exist_ok=True)\n    with open(PROGRESS_FILE, 'w', encoding='utf-8') as f:\n        json.dump(list(progress_set), f)\n\ndef extract_json_from_text(text):\n    \"\"\"通过正则提取并解析大模型返回文本中的 JSON 数组\"\"\"\n    # 尝试匹配 markdown 代码块中的 JSON 数组\n    match = re.search(r'```(?:json)?\\s*(\\[.*?\\])\\s*```', text, re.DOTALL)\n    json_str = \"\"\n    if match:\n        json_str = match.group(1)\n    else:\n        # 降级：直接尝试寻找被 [] 包裹的内容\n        match = re.search(r'\\[.*\\]', text, re.DOTALL)\n        if match:\n            json_str = match.group(0)\n    \n    if not json_str:\n        logging.error(\"未能从返回结果中匹配到 JSON 数组。\")\n        return None\n\n    try:\n        return json.loads(json_str)\n    except json.JSONDecodeError as e:\n        logging.error(f\"JSON 解析失败: {e}\\n原文提取为:\\n{json_str}\")\n        return None\n\ndef call_deepseek_api(description):\n    \"\"\"调用 DeepSeek API 生成测试用例，包含重试机制\"\"\"\n    if not DEEPSEEK_API_KEY:\n        logging.error(\"未配置 DEEPSEEK_API_KEY 环境变量。\")\n        return None\n\n    headers = {\n        'Content-Type': 'application/json',\n        'Authorization': f'Bearer {DEEPSEEK_API_KEY}'\n    }\n    \n    prompt = (\n        \"根据以下编程题目的描述，生成 3 到 5 个测试用例。\\n\"\n        \"必须严格按照以下 JSON 数组格式输出，不要包含任何其他解释文字或 markdown 格式：\\n\"\n        '[{\"input\": \"输入样例1\", \"expect\": \"预期输出1\"}, {\"input\": \"输入样例2\", \"expect\": \"预期输出2\"}]\\n\\n'\n        f\"题目描述：\\n{description}\"\n    )\n\n    data = {\n        \"model\": \"deepseek-chat\",\n        \"messages\": [\n            {\"role\": \"system\", \"content\": \"你是一个严谨的编程测试用例生成助手。只输出合法的 JSON 数组，不附加任何其他文本。\"},\n            {\"role\": \"user\", \"content\": prompt}\n        ],\n        \"temperature\": 0.3\n    }\n\n    for attempt in range(1, MAX_RETRIES + 1):\n        try:\n            response = requests.post(DEEPSEEK_API_URL, headers=headers, json=data, timeout=30)\n            response.raise_for_status()\n            result = response.json()\n            \n            content = result['choices'][0]['message']['content']\n            parsed_json = extract_json_from_text(content)\n            \n            if parsed_json is not None and isinstance(parsed_json, list):\n                return parsed_json\n            else:\n                logging.warning(f\"第 {attempt} 次尝试：返回的数据格式无效或不是数组。\")\n        except Exception as e:\n            logging.warning(f\"第 {attempt} 次尝试请求失败: {e}\")\n        \n        if attempt < MAX_RETRIES:\n            # 失败后指数退避重试\n            time.sleep(RATE_LIMIT_DELAY * attempt)\n    \n    return None\n\ndef merge_and_deduplicate(existing_tail_code, new_cases):\n    \"\"\"解析现有 test cases 并与新生成的 cases 合并，基于 input 去重\"\"\"\n    existing_cases = []\n    if existing_tail_code and str(existing_tail_code).strip():\n        try:\n            existing_cases = json.loads(existing_tail_code)\n            if not isinstance(existing_cases, list):\n                existing_cases = []\n        except json.JSONDecodeError:\n            logging.warning(\"原有的 tail_code 不是有效的 JSON 数组，将忽略原数据。\")\n            existing_cases = []\n\n    seen_inputs = set()\n    merged_cases = []\n    \n    # 优先保留原有的测试用例\n    for case in existing_cases + new_cases:\n        if isinstance(case, dict) and 'input' in case and 'expect' in case:\n            # 标准化 input 用于去重比较\n            input_str = str(case['input']).strip()\n            if input_str not in seen_inputs:\n                seen_inputs.add(input_str)\n                merged_cases.append(case)\n                \n    return json.dumps(merged_cases, ensure_ascii=False)\n\ndef main():\n    processed_numbers = load_progress()\n    logging.info(f\"已加载 {len(processed_numbers)} 个已处理题目的进度记录。\")\n\n    # 连接数据库\n    try:\n        connection = pymysql.connect(**DB_CONFIG)\n        cursor = connection.cursor()\n        logging.info(\"成功连接到 MySQL 数据库。\")\n    except Exception as e:\n        logging.error(f\"连接数据库失败: {e}\")\n        return\n\n    try:\n        cursor.execute(\"SELECT number, description, tail_code FROM oj_questions\")\n        questions = cursor.fetchall()\n        logging.info(f\"数据库中共有 {len(questions)} 道题目。\")\n\n        for q in questions:\n            number = str(q['number'])\n            \n            # 如果已经处理过，直接跳过（支持断点续传）\n            if number in processed_numbers:\n                logging.debug(f\"跳过已处理题目: {number}\")\n                continue\n\n            logging.info(f\"开始处理题目 {number} ...\")\n            \n            description = q['description']\n            existing_tail_code = q['tail_code']\n\n            # 请求大模型生成用例\n            new_cases = call_deepseek_api(description)\n            \n            if new_cases:\n                # 合并并去重\n                updated_tail_code = merge_and_deduplicate(existing_tail_code, new_cases)\n                \n                try:\n                    # 更新回数据库\n                    cursor.execute(\n                        \"UPDATE oj_questions SET tail_code = %s WHERE number = %s\",\n                        (updated_tail_code, number)\n                    )\n                    connection.commit()\n                    logging.info(f\"成功更新题目 {number} 的测试用例。\")\n                    \n                    # 记录进度\n                    processed_numbers.add(number)\n                    save_progress(processed_numbers)\n                except Exception as e:\n                    connection.rollback()\n                    logging.error(f\"题目 {number} 数据库更新失败: {e}\")\n            else:\n                logging.error(f\"题目 {number} 测试用例生成失败，重试 {MAX_RETRIES} 次后放弃。\")\n\n            # 速率限制：每次处理完一题后暂停，避免触发 API 限流\n            time.sleep(RATE_LIMIT_DELAY)\n\n    except Exception as e:\n        logging.error(f\"运行过程中发生错误: {e}\")\n    finally:\n        cursor.close()\n        connection.close()\n        logging.info(\"数据库连接已关闭。\")\n\nif __name__ == '__main__':\n    main()\n"}}}}