155 lines
4.9 KiB
Python
155 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
真实 B 组实验（Single-Step 模式）- 使用 openai 库
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
from datetime import datetime
|
|
import openai
|
|
|
|
# 配置
|
|
base_dir = os.getcwd()
|
|
experiment_dir = os.path.join(base_dir, "tools", "experiments", "wiki-generation-compare")
|
|
source_file = os.path.join(base_dir, "raw", "呼吸之间_李谨伯", "第一编 从身体入手.md")
|
|
output_dir = os.path.join(experiment_dir, "output", "group-b-real")
|
|
|
|
# 创建输出目录
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
os.makedirs(os.path.join(output_dir, "wiki", "concepts"), exist_ok=True)
|
|
os.makedirs(os.path.join(output_dir, "wiki", "methods"), exist_ok=True)
|
|
os.makedirs(os.path.join(output_dir, "wiki", "entities"), exist_ok=True)
|
|
|
|
# 记录开始时间
|
|
start_time = datetime.now()
|
|
start_time_str = start_time.strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
print("=== B 组真实实验开始：Single-Step 模式 ===")
|
|
print(f"开始时间: {start_time_str}")
|
|
|
|
# 检查 API Key
|
|
api_key = os.environ.get('OPENAI_API_KEY')
|
|
if not api_key:
|
|
print("[ERROR] 未找到 OPENAI_API_KEY 环境变量")
|
|
exit(1)
|
|
|
|
print(f"API Key 长度: {len(api_key)}")
|
|
|
|
# 初始化 openai 客户端
|
|
client = openai.OpenAI(api_key=api_key, timeout=300)
|
|
|
|
# 直接生成
|
|
print("\n生成阶段...")
|
|
|
|
# 读取源文件
|
|
with open(source_file, 'r', encoding='utf-8') as f:
|
|
source_content = f.read()
|
|
|
|
# 读取单步 prompt
|
|
singlestep_prompt_file = os.path.join(experiment_dir, "prompts", "singlestep.md")
|
|
with open(singlestep_prompt_file, 'r', encoding='utf-8') as f:
|
|
template = f.read()
|
|
|
|
full_prompt = template.replace("{SOURCE_CONTENT}", source_content)
|
|
|
|
print(f" Prompt 长度: {len(full_prompt)} 字符")
|
|
print(" 调用 OpenAI API 生成 Wiki 页面...")
|
|
|
|
try:
|
|
response = client.chat.completions.create(
|
|
model="reasoning",
|
|
messages=[
|
|
{
|
|
"role": "system",
|
|
"content": "你是一位知识库构建专家。请直接从源文件中提取关键术语并生成 Wiki 页面。"
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": full_prompt
|
|
}
|
|
],
|
|
temperature=0.5,
|
|
max_tokens=16000
|
|
)
|
|
|
|
generation_result = response.choices[0].message.content
|
|
print(" [OK] LLM 生成完成")
|
|
print(f" 生成内容长度: {len(generation_result)} 字符")
|
|
|
|
# 保存原始生成结果
|
|
raw_output_file = os.path.join(output_dir, "raw_generation.md")
|
|
with open(raw_output_file, 'w', encoding='utf-8') as f:
|
|
f.write(generation_result)
|
|
print(f" [OK] 原始生成结果已保存到 {raw_output_file}")
|
|
|
|
except Exception as e:
|
|
print(f" [ERROR] LLM API 调用失败: {type(e).__name__}: {e}")
|
|
exit(1)
|
|
|
|
# 解析生成结果
|
|
sample_pages = []
|
|
|
|
# 解析 ---FILE: ... ---END FILE--- 块
|
|
file_blocks = re.findall(r'---FILE: (.*?)---(.*?)---END FILE---', generation_result, re.DOTALL)
|
|
|
|
if not file_blocks:
|
|
print(" [WARNING] 未找到标准文件块格式，尝试宽松匹配...")
|
|
# 尝试宽松匹配
|
|
file_blocks = re.findall(r'FILE: (.*?)\n(.*?)(?=(?:FILE:|$))', generation_result, re.DOTALL)
|
|
|
|
print(f"\n 解析到 {len(file_blocks)} 个文件块")
|
|
|
|
for i, block in enumerate(file_blocks):
|
|
if isinstance(block, tuple):
|
|
file_path = block[0].strip()
|
|
file_content = block[1].strip()
|
|
else:
|
|
# 如果没有匹配，尝试直接提取整个内容
|
|
continue
|
|
|
|
# 创建目录（如果需要）
|
|
full_path = os.path.join(output_dir, file_path)
|
|
file_dir = os.path.dirname(full_path)
|
|
os.makedirs(file_dir, exist_ok=True)
|
|
|
|
# 写入文件
|
|
with open(full_path, 'w', encoding='utf-8') as f:
|
|
f.write(file_content)
|
|
sample_pages.append(full_path)
|
|
print(f" [OK] 生成 [{i+1}/{len(file_blocks)}]: {file_path}")
|
|
|
|
# 记录结束时间
|
|
end_time = datetime.now()
|
|
end_time_str = end_time.strftime("%Y-%m-%d %H:%M:%S")
|
|
duration = (end_time - start_time).total_seconds() / 60
|
|
|
|
print(f"\n=== B 组实验完成 ===")
|
|
print(f"结束时间: {end_time_str}")
|
|
print(f"总耗时: {round(duration, 2)} 分钟")
|
|
print(f"生成文件数: {len(sample_pages)}")
|
|
|
|
# 保存元数据
|
|
metadata = {
|
|
"group": "B-real",
|
|
"mode": "Single-Step (Real API)",
|
|
"sourceFile": "raw/呼吸之间_李谨伯/第一编 从身体入手.md",
|
|
"startTime": start_time_str,
|
|
"endTime": end_time_str,
|
|
"durationMinutes": round(duration, 2),
|
|
"pagesGenerated": len(sample_pages),
|
|
"apiModel": "reasoning",
|
|
"rawContentLength": len(generation_result),
|
|
"status": "success"
|
|
}
|
|
|
|
metadata_file = os.path.join(output_dir, "metadata.json")
|
|
with open(metadata_file, 'w', encoding='utf-8') as f:
|
|
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
|
|
|
print(f"[OK] 元数据已保存到 {metadata_file}")
|
|
print(f"\n文件输出位置:")
|
|
print(f" 原始生成: {raw_output_file}")
|
|
print(f" Wiki 页面: {output_dir}/wiki/")
|
|
print(f" 元数据: {metadata_file}") |