180 lines
4.4 KiB
Python
180 lines
4.4 KiB
Python
"""
|
|
工具函数模块
|
|
提供配置加载、日志设置等通用功能
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, Any
|
|
from loguru import logger
|
|
import sys
|
|
|
|
|
|
def load_config(config_path: str = "config/config.json") -> Dict[str, Any]:
|
|
"""
|
|
加载配置文件
|
|
|
|
Args:
|
|
config_path: 配置文件路径
|
|
|
|
Returns:
|
|
配置字典
|
|
"""
|
|
try:
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
config = json.load(f)
|
|
|
|
# 从环境变量获取 API Key
|
|
if 'OPENROUTER_API_KEY' in os.environ:
|
|
config['openrouter']['api_key'] = os.environ['OPENROUTER_API_KEY']
|
|
|
|
return config
|
|
except FileNotFoundError:
|
|
raise FileNotFoundError(f"配置文件未找到: {config_path}")
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"配置文件格式错误: {e}")
|
|
|
|
|
|
def load_prompts(prompts_path: str = "config/prompts.json") -> Dict[str, str]:
|
|
"""
|
|
加载提示词模板
|
|
|
|
Args:
|
|
prompts_path: 提示词文件路径
|
|
|
|
Returns:
|
|
提示词字典
|
|
"""
|
|
try:
|
|
with open(prompts_path, 'r', encoding='utf-8') as f:
|
|
return json.load(f)
|
|
except FileNotFoundError:
|
|
raise FileNotFoundError(f"提示词文件未找到: {prompts_path}")
|
|
|
|
|
|
def setup_logging(config: Dict[str, Any]) -> None:
|
|
"""
|
|
设置日志配置
|
|
|
|
Args:
|
|
config: 配置字典
|
|
"""
|
|
log_config = config.get('logging', {})
|
|
|
|
# 移除默认处理器
|
|
logger.remove()
|
|
|
|
# 添加控制台输出
|
|
logger.add(
|
|
sys.stdout,
|
|
level=log_config.get('level', 'INFO'),
|
|
format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
|
|
)
|
|
|
|
# 添加文件输出
|
|
if 'file' in log_config:
|
|
log_file = log_config['file']
|
|
# 确保日志目录存在
|
|
Path(log_file).parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
logger.add(
|
|
log_file,
|
|
level=log_config.get('level', 'INFO'),
|
|
rotation=log_config.get('rotation', '10 MB'),
|
|
retention=log_config.get('retention', '7 days'),
|
|
encoding='utf-8',
|
|
format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}"
|
|
)
|
|
|
|
|
|
def ensure_output_dir(output_dir: str) -> Path:
|
|
"""
|
|
确保输出目录存在
|
|
|
|
Args:
|
|
output_dir: 输出目录路径
|
|
|
|
Returns:
|
|
输出目录的 Path 对象
|
|
"""
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
return output_path
|
|
|
|
|
|
def sanitize_filename(filename: str) -> str:
|
|
"""
|
|
清理文件名，移除非法字符
|
|
|
|
Args:
|
|
filename: 原始文件名
|
|
|
|
Returns:
|
|
清理后的文件名
|
|
"""
|
|
import re
|
|
# 移除或替换非法字符
|
|
filename = re.sub(r'[<>:"/\\|?*]', '_', filename)
|
|
# 移除多余的空格和点
|
|
filename = re.sub(r'\s+', ' ', filename).strip('. ')
|
|
return filename
|
|
|
|
|
|
def format_file_size(size_bytes: int) -> str:
|
|
"""
|
|
格式化文件大小显示
|
|
|
|
Args:
|
|
size_bytes: 字节数
|
|
|
|
Returns:
|
|
格式化的大小字符串
|
|
"""
|
|
if size_bytes == 0:
|
|
return "0B"
|
|
|
|
size_names = ["B", "KB", "MB", "GB"]
|
|
import math
|
|
i = int(math.floor(math.log(size_bytes, 1024)))
|
|
p = math.pow(1024, i)
|
|
s = round(size_bytes / p, 2)
|
|
return f"{s} {size_names[i]}"
|
|
|
|
|
|
def estimate_tokens(text: str) -> int:
|
|
"""
|
|
估算文本的 token 数量
|
|
|
|
Args:
|
|
text: 输入文本
|
|
|
|
Returns:
|
|
估算的 token 数量
|
|
"""
|
|
# 简单估算：英文约 4 字符/token，中文约 1.5 字符/token
|
|
import re
|
|
|
|
# 分离中英文
|
|
chinese_chars = len(re.findall(r'[\u4e00-\u9fff]', text))
|
|
other_chars = len(text) - chinese_chars
|
|
|
|
# 估算 tokens
|
|
estimated_tokens = chinese_chars / 1.5 + other_chars / 4
|
|
return int(estimated_tokens)
|
|
|
|
|
|
def truncate_text(text: str, max_length: int = 100) -> str:
|
|
"""
|
|
截断文本用于显示
|
|
|
|
Args:
|
|
text: 原始文本
|
|
max_length: 最大长度
|
|
|
|
Returns:
|
|
截断后的文本
|
|
"""
|
|
if len(text) <= max_length:
|
|
return text
|
|
return text[:max_length-3] + "..." |