72 lines
3.1 KiB
Python
72 lines
3.1 KiB
Python
|
# planning_agent/main.py
|
|||
|
|
import argparse
|
|||
|
|
import sys
|
|||
|
|
import os
|
|||
|
|
from pathlib import Path
|
|||
|
|
# 注意：这里使用了相对导入，必须以模块方式运行：python -m planner.planning_agent.main
|
|||
|
|
from .planner import build_plan_from_text
|
|||
|
|
from planner.mineru_client import extract_texts_with_mineru
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
# 尝试强制设置标准输出编码为 utf-8，防止部分环境打印异常
|
|||
|
|
if sys.stdout and hasattr(sys.stdout, 'reconfigure'):
|
|||
|
|
try:
|
|||
|
|
sys.stdout.reconfigure(encoding='utf-8')
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
parser = argparse.ArgumentParser(description="Testing Task Planning Agent (Natural Language Mode)")
|
|||
|
|
parser.add_argument("--input", required=True, help="输入文件或文件夹路径")
|
|||
|
|
parser.add_argument("--output", required=True, help="输出目录路径")
|
|||
|
|
# 新增手动参数接收
|
|||
|
|
parser.add_argument("--manual_args", required=False, default="{}", help="手动填写的参数（JSON字符串）")
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
input_path = Path(args.input)
|
|||
|
|
output_dir = Path(args.output)
|
|||
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
manual_args_str = args.manual_args
|
|||
|
|
|
|||
|
|
# 获取 planner 目录 (当前脚本的父目录的父目录)
|
|||
|
|
# planning_agent/main.py -> planning_agent -> planner
|
|||
|
|
current_file = Path(__file__).resolve()
|
|||
|
|
planner_dir = current_file.parent.parent
|
|||
|
|
mineru_save_dir = planner_dir / "mineru_result"
|
|||
|
|
|
|||
|
|
# 定义需要 MinerU 解析的文档类型
|
|||
|
|
doc_extensions = ['.pdf', '.docx', '.doc', '.xlsx', '.xls', '.pptx']
|
|||
|
|
|
|||
|
|
# 判断逻辑：如果是目录，或者后缀在支持列表中，则使用 MinerU 解析
|
|||
|
|
if input_path.is_dir() or (input_path.is_file() and input_path.suffix.lower() in doc_extensions):
|
|||
|
|
print(f"[INFO] 检测到文档/文件夹输入，将使用 MinerU API 解析文件...")
|
|||
|
|
try:
|
|||
|
|
# 传入 save_dir 到 planner/mineru_result
|
|||
|
|
merged_text_path = extract_texts_with_mineru(str(input_path), save_dir=str(mineru_save_dir))
|
|||
|
|
with open(merged_text_path, "r", encoding="utf-8") as f:
|
|||
|
|
requirement_text = f.read()
|
|||
|
|
source_doc = merged_text_path
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[ERROR] 文档解析失败: {e}")
|
|||
|
|
return
|
|||
|
|
else:
|
|||
|
|
# 否则默认当作普通 UTF-8 文本文件读取 (如 .txt, .md)
|
|||
|
|
print(f"[INFO] 检测到文本文件输入，直接读取...")
|
|||
|
|
try:
|
|||
|
|
with open(input_path, "r", encoding="utf-8") as f:
|
|||
|
|
requirement_text = f.read()
|
|||
|
|
source_doc = str(input_path)
|
|||
|
|
except UnicodeDecodeError:
|
|||
|
|
print(f"[ERROR] 文件读取编码错误: {input_path} 不是有效的 UTF-8 文本文件。如果是文档(如docx)，请确保 MinerU 流程被正确触发。")
|
|||
|
|
return
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[ERROR] 读取文件失败: {e}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 将 manual_args_str 传递给 planner
|
|||
|
|
output_path = build_plan_from_text(requirement_text, source_doc, str(output_dir), manual_args_str)
|
|||
|
|
print(f"[SUCCESS] 测试任务规划生成完成: {output_path}")
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|