Initial commit

2026-01-19 09:51:07 +08:00
commit 9ef82393be
174 changed files with 22285 additions and 0 deletions
@@ -0,0 +1,180 @@
+"""
+工具函数模块
+提供配置加载、日志设置等通用功能
+"""
+
+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] + "..."