"""工作目录管理器 - 管理每个项目的独立工作目录"""
from config import (
BATCH_TASK_DIR_NAME,
LOGS_DIR_NAME,
PROMPTS_DIR_NAME,
OUTPUT_DIR_NAME,
TASKS_FILE_NAME,
DEFAULT_PROMPT_FILE_NAME,
)
from pathlib import Path
from typing import Optional
import os
import sys
# 导入配置常量
sys.path.insert(0, str(Path(__file__).parent.parent))
class WorkspaceManager:
"""工作目录管理器，支持为每个项目设置独立的工作目录"""
def __init__(self, batch_task_dir_name: str = None):
"""
初始化工作目录管理器
Args:
batch_task_dir_name: 批处理任务目录名，默认使用 config.BATCH_TASK_DIR_NAME
"""
if batch_task_dir_name is None:
# 使用配置文件中的值（已从环境变量读取）
batch_task_dir_name = BATCH_TASK_DIR_NAME
self.batch_task_dir_name = batch_task_dir_name
self._current_work_dir: Optional[Path] = None
def set_work_dir(self, work_dir: str | Path) -> dict:
"""
设置工作目录
Args:
work_dir: 工作目录路径
Returns:
设置结果
"""
work_path = Path(work_dir).resolve()
if not work_path.exists():
return {
"success": False,
"message": f"❌ 工作目录不存在: {work_path}"
}
if not work_path.is_dir():
return {
"success": False,
"message": f"❌ 路径不是目录: {work_path}"
}
self._current_work_dir = work_path
# 初始化目录结构
batch_dir = self.get_batch_dir()
prompts_dir = self.get_prompts_dir()
tasks_file = self.get_tasks_file()
output_dir = self.get_output_dir() # 输出目录在 .batch_task 中
logs_dir = self.get_logs_dir() # 日志目录在 .batch_task 中
# 创建必要的目录
batch_dir.mkdir(parents=True, exist_ok=True)
prompts_dir.mkdir(parents=True, exist_ok=True)
output_dir.mkdir(parents=True, exist_ok=True)
logs_dir.mkdir(parents=True, exist_ok=True)
# 创建默认提示词文件
default_prompt = self.get_default_prompt_file()
if not default_prompt.exists():
default_prompt.write_text(
"# 小说改写提示词\n\n在这里编写你的改写提示词...",
encoding='utf-8'
)
return {
"success": True,
"message": f"✅ 工作目录已设置: {work_path}\n"
f"📁 批处理目录: {batch_dir}\n"
f"📝 提示词目录: {prompts_dir}\n"
f"📂 输出目录: {output_dir}",
"work_dir": str(work_path),
"batch_dir": str(batch_dir),
"prompts_dir": str(prompts_dir),
"output_dir": str(output_dir)
}
def get_work_dir(self) -> Path:
"""获取当前工作目录"""
if self._current_work_dir is None:
# 如果没有设置，使用环境变量或默认路径
default_dir = os.getenv('BATCH_WORK_DIR')
if default_dir:
self._current_work_dir = Path(default_dir).resolve()
else:
# 使用当前工作目录
self._current_work_dir = Path.cwd()
return self._current_work_dir
def get_batch_dir(self) -> Path:
"""获取批处理任务目录"""
return self.get_work_dir() / self.batch_task_dir_name
def get_prompts_dir(self) -> Path:
"""获取提示词目录"""
return self.get_batch_dir() / PROMPTS_DIR_NAME
def get_tasks_file(self) -> Path:
"""获取任务文件路径"""
return self.get_batch_dir() / TASKS_FILE_NAME
def get_default_prompt_file(self) -> Path:
"""获取默认提示词文件路径"""
return self.get_prompts_dir() / DEFAULT_PROMPT_FILE_NAME
def get_output_dir(self) -> Path:
"""获取输出目录（在 .batch_task 目录中）"""
return self.get_batch_dir() / OUTPUT_DIR_NAME
def get_logs_dir(self) -> Path:
"""获取日志目录（在 .batch_task 目录中）"""
return self.get_batch_dir() / LOGS_DIR_NAME
def get_info(self) -> dict:
"""获取当前工作目录信息"""
work_dir = self.get_work_dir()
batch_dir = self.get_batch_dir()
prompts_dir = self.get_prompts_dir()
tasks_file = self.get_tasks_file()
output_dir = self.get_output_dir()
return {
"work_dir": str(work_dir),
"batch_dir": str(batch_dir),
"prompts_dir": str(prompts_dir),
"tasks_file": str(tasks_file),
"output_dir": str(output_dir),
"batch_dir_exists": batch_dir.exists(),
"prompts_dir_exists": prompts_dir.exists(),
"tasks_file_exists": tasks_file.exists(),
"output_dir_exists": output_dir.exists(),
}
# 创建全局工作目录管理器实例
workspace_manager = WorkspaceManager()
__all__ = ['WorkspaceManager', 'workspace_manager']
