main.py

"""
LLM-Local 主入口
自动检测平台并选择最优的推理引擎
"""

import sys
import argparse
from pathlib import Path

# 添加 src 到路径
sys.path.insert(0, str(Path(__file__).parent / "src"))

from src.utils.device_manager import DeviceManager
from src.utils.config_manager import ConfigManager
from src.utils.logger import setup_logger, get_logger, log_system_info


def create_inference_engine(backend: str, config: ConfigManager):
    """
    创建推理引擎
    
    Args:
        backend: 'mlx' 或 'pytorch'
        config: 配置管理器
        
    Returns:
        推理引擎实例
    """
    model_config = config.get_model_config()
    system_config = config.get_system_config()
    advanced_config = config.get("advanced", {})
    cache_dir = advanced_config.get("cache_dir", "./models")
    
    if backend == "mlx":
        try:
            from src.inference.mlx_inference import MLXInference
            
            return MLXInference(
                model_id=model_config.get("default_model"),
                max_tokens=model_config.get("max_new_tokens", 512),
                temperature=model_config.get("temperature", 0.7),
                top_p=model_config.get("top_p", 0.9),
                cache_dir=cache_dir,
            )
        except ImportError as e:
            print(f"❌ MLX 未安装: {e}")
            print("请运行: pip install -r requirements-mlx.txt")
            sys.exit(1)
    
    elif backend == "pytorch":
        try:
            from src.inference.pytorch_inference import PyTorchInference
            
            device = system_config.get("device_preference")
            if device == "auto":
                device = None
            
            return PyTorchInference(
                model_id=model_config.get("default_model"),
                torch_dtype=model_config.get("torch_dtype", "float16"),
                device=device,
                max_new_tokens=model_config.get("max_new_tokens", 512),
                temperature=model_config.get("temperature", 0.7),
                top_p=model_config.get("top_p", 0.9),
                do_sample=model_config.get("do_sample", True),
                cache_dir=cache_dir,
            )
        except ImportError as e:
            print(f"❌ PyTorch 或 Transformers 未安装: {e}")
            print("请运行: pip install -r requirements.txt")
            sys.exit(1)
    
    else:
        raise ValueError(f"不支持的后端: {backend}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="LLM-Local - 本地大语言模型部署系统",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 自动选择最优平台
  python main.py
  
  # 指定使用 MLX (Apple Silicon)
  python main.py --backend mlx
  
  # 指定使用 PyTorch
  python main.py --backend pytorch
  
  # 使用自定义模型
  python main.py --model Qwen/Qwen2.5-7B-Instruct
  
  # 查看设备信息
  python main.py --info
        """
    )
    
    parser.add_argument(
        "--backend",
        type=str,
        choices=["auto", "mlx", "pytorch"],
        default="auto",
        help="推理后端 (默认: auto)"
    )
    
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="模型 ID (覆盖配置文件)"
    )
    
    parser.add_argument(
        "--config",
        type=str,
        default=None,
        help="配置文件路径"
    )
    
    parser.add_argument(
        "--info",
        action="store_true",
        help="显示设备信息后退出"
    )
    
    parser.add_argument(
        "--show-config",
        action="store_true",
        help="显示当前配置后退出"
    )
    
    parser.add_argument(
        "--debug",
        action="store_true",
        help="启用调试模式(详细日志)"
    )
    
    args = parser.parse_args()
    
    # 加载配置
    config = ConfigManager(args.config)
    
    # 设置日志
    log_config = config.get("logging", {})
    log_level = "DEBUG" if args.debug else config.get("advanced.log_level", "INFO")
    
    setup_logger(
        log_level=log_level,
        log_dir=log_config.get("log_dir", "./logs"),
        console_output=log_config.get("console_output", True),
        file_output=log_config.get("file_output", True),
    )
    
    logger = get_logger()
    logger.info("=" * 60)
    logger.info("LLM-Local 启动")
    logger.info("=" * 60)
    
    # 记录系统信息
    log_system_info()
    
    # 显示设备信息
    if args.info:
        DeviceManager.print_device_info()
        return
    
    # 显示配置
    if args.show_config:
        config.print_config()
        return
    
    # 如果指定了模型,覆盖配置
    if args.model:
        config.set("model.default_model", args.model)
    
    # 确定使用的后端
    if args.backend == "auto":
        backend = DeviceManager.get_recommended_backend()
        logger.info(f"自动选择后端: {backend}")
        print(f"🔍 自动选择后端: {backend.upper()}")
    else:
        backend = args.backend
        logger.info(f"使用指定后端: {backend}")
        print(f"🎯 使用指定后端: {backend.upper()}")
    
    # 显示设备信息
    print()
    DeviceManager.print_device_info()
    print()
    
    try:
        # 创建推理引擎
        logger.info("开始创建推理引擎")
        inference = create_inference_engine(backend, config)
        logger.info("推理引擎创建成功")
        
        # 获取对话配置
        chat_config = config.get_chat_config()
        system_prompt = chat_config.get("system_prompt", "你是一个乐于助人的AI助手。")
        stream = chat_config.get("stream", True)
        
        logger.info("启动交互式对话")
        
        # 启动交互式对话
        if backend == "mlx":
            inference.interactive_chat(system_prompt=system_prompt, stream=stream)
        else:
            inference.interactive_chat(system_prompt=system_prompt)
        
        logger.info("程序正常退出")
        
    except KeyboardInterrupt:
        logger.info("用户中断程序")
        print("\n\n👋 再见!")
    except Exception as e:
        logger.error(f"程序异常退出: {str(e)}")
        logger.debug(f"详细错误信息", exc_info=True)
        print(f"\n❌ 发生错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()