You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
265 lines
9.3 KiB
265 lines
9.3 KiB
"""
|
|
首次运行初始化脚本
|
|
- 创建桌面快捷方式指向 main.py，用 Python 运行，图标为 icon.png
|
|
- 检查 knowledge 文件夹文件数量
|
|
- 如果 <= 3，提示用户建立知识库，并自动运行 process_knowledge_base.py
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import shutil
|
|
import json
|
|
import platform
|
|
import urllib.parse
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
# 获取脚本所在目录（项目根目录）
|
|
PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
MAIN_SCRIPT = os.path.join(PROJECT_DIR, "main.py")
|
|
ICON_ICO = os.path.join(PROJECT_DIR, "icon.ico")
|
|
ICON_PNG = os.path.join(PROJECT_DIR, "icon.png")
|
|
KNOWLEDGE_DIR = os.path.join(PROJECT_DIR, "knowledge")
|
|
PROCESS_SCRIPT = os.path.join(PROJECT_DIR, "process_knowledge_base.py")
|
|
CONFIG_PATH = os.path.join(PROJECT_DIR, "config.json")
|
|
|
|
|
|
def load_config():
|
|
if not os.path.exists(CONFIG_PATH):
|
|
return {}
|
|
try:
|
|
with open(CONFIG_PATH, "r", encoding="utf-8") as f:
|
|
return json.load(f) or {}
|
|
except Exception as exc:
|
|
print(f"✗ 读取 config.json 失败：{exc}")
|
|
return {}
|
|
|
|
|
|
def save_config(cfg):
|
|
try:
|
|
with open(CONFIG_PATH, "w", encoding="utf-8") as f:
|
|
json.dump(cfg, f, ensure_ascii=False, indent=4)
|
|
print(f"✓ 配置已写入：{CONFIG_PATH}")
|
|
except Exception as exc:
|
|
print(f"✗ 写入 config.json 失败：{exc}")
|
|
|
|
|
|
def check_network_reachable(url="https://www.google.com/generate_204", timeout_ms=2000):
|
|
"""Check connectivity via ping + lightweight HTTPS GET. Returns True if either succeeds."""
|
|
parsed = urllib.parse.urlparse(url)
|
|
host = parsed.hostname or "google.com"
|
|
is_windows = platform.system().lower().startswith("win")
|
|
if is_windows:
|
|
cmd = ["ping", "-n", "1", "-w", str(timeout_ms), host]
|
|
else:
|
|
cmd = ["ping", "-c", "1", "-W", str(int(timeout_ms / 1000)), host]
|
|
|
|
try:
|
|
result = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
if result.returncode == 0:
|
|
return True
|
|
except Exception as exc:
|
|
print(f"✗ 无法执行 ping：{exc}")
|
|
|
|
# Fallback: HTTPS request (respects proxy envs)
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=timeout_ms / 1000.0) as resp:
|
|
return 200 <= getattr(resp, "status", 200) < 400
|
|
except Exception as exc:
|
|
print(f"✗ HTTPS 连通性检测失败：{exc}")
|
|
return False
|
|
|
|
|
|
def prompt_config_if_needed():
|
|
cfg = load_config()
|
|
|
|
defaults = {
|
|
"base_url": cfg.get("base_url") or "https://api.moonshot.cn/v1",
|
|
"model": cfg.get("model") or "kimi-latest",
|
|
}
|
|
|
|
if not cfg.get("system_prompt"):
|
|
cfg["system_prompt"] = (
|
|
"You are a WorldQuant BRAIN platform expert and Consultant. "
|
|
"Your goal is to assist users with Alpha development, BRAIN API usage, and maximizing consultant income."
|
|
)
|
|
|
|
def ask(field, label, default_value=None, allow_empty=False):
|
|
current = str(cfg.get(field, "")).strip()
|
|
prompt = f"请输入 {label}" + (f" [默认: {current or default_value}]" if (current or default_value) else "") + ": "
|
|
value = input(prompt).strip()
|
|
if not value:
|
|
value = current if current else (default_value if default_value is not None else "")
|
|
if value or allow_empty:
|
|
cfg[field] = value
|
|
return value
|
|
|
|
api_key = ask("api_key", "API Key")
|
|
base_url = ask("base_url", "Base URL", defaults["base_url"], allow_empty=True)
|
|
model = ask("model", "模型名称", defaults["model"], allow_empty=True)
|
|
|
|
if api_key:
|
|
save_config(cfg)
|
|
return True
|
|
|
|
print("✗ 配置文件中的 api_key 为空，请填写后再运行本脚本。")
|
|
return False
|
|
|
|
|
|
def is_api_key_configured():
|
|
"""Check config; if api_key missing/empty, prompt user to fill and persist."""
|
|
cfg = load_config()
|
|
|
|
api_key = str(cfg.get("api_key", "")).strip()
|
|
if api_key:
|
|
return True
|
|
|
|
print("✗ 当前 config.json 中 api_key 为空。")
|
|
return prompt_config_if_needed()
|
|
|
|
def create_desktop_shortcut():
|
|
"""创建桌面快捷方式"""
|
|
try:
|
|
import win32com.client
|
|
|
|
desktop = os.path.expanduser("~\\Desktop")
|
|
shortcut_path = os.path.join(desktop, "BRAIN顾问助手.lnk")
|
|
|
|
# 获取 Python 可执行文件路径
|
|
python_exe = sys.executable
|
|
|
|
# 创建 shortcut
|
|
shell = win32com.client.Dispatch("WScript.Shell")
|
|
shortcut = shell.CreateShortcut(shortcut_path)
|
|
shortcut.TargetPath = python_exe
|
|
shortcut.Arguments = f'"{MAIN_SCRIPT}"'
|
|
shortcut.WorkingDirectory = PROJECT_DIR
|
|
|
|
# 设置图标：优先使用 .ico（Windows 标准格式），再试 .png
|
|
icon_found = False
|
|
if os.path.exists(ICON_ICO):
|
|
# .ico 文件：使用 "路径,0" 格式
|
|
shortcut.IconLocation = f"{ICON_ICO},0"
|
|
icon_found = True
|
|
print(f" 使用图标：{ICON_ICO}")
|
|
elif os.path.exists(ICON_PNG):
|
|
# .png 文件：也需要 "路径,0" 格式
|
|
shortcut.IconLocation = f"{ICON_PNG},0"
|
|
icon_found = True
|
|
print(f" 使用图标：{ICON_PNG}")
|
|
else:
|
|
# 使用 Python 可执行文件的图标作为默认
|
|
shortcut.IconLocation = f"{python_exe},0"
|
|
print(f" 使用默认图标（Python 图标）")
|
|
|
|
shortcut.Description = "BRAIN 顾问助手 - AI 驱动的交互工具"
|
|
shortcut.save()
|
|
|
|
print(f"✓ 桌面快捷方式已创建：{shortcut_path}")
|
|
return True
|
|
except ImportError:
|
|
print("⚠ pywin32 未安装，尝试使用备选方案创建快捷方式...")
|
|
try:
|
|
# 备选方案：使用 Windows API
|
|
create_shortcut_via_batch(desktop, shortcut_path, python_exe)
|
|
return True
|
|
except Exception as e:
|
|
print(f"✗ 创建快捷方式失败：{e}")
|
|
print(f" 请手动创建快捷方式，目标为：{python_exe} \"{MAIN_SCRIPT}\"")
|
|
return False
|
|
except Exception as e:
|
|
print(f"✗ 创建快捷方式失败：{e}")
|
|
return False
|
|
|
|
def check_knowledge_base():
|
|
"""检查 knowledge 文件夹中的文件数量"""
|
|
if not os.path.exists(KNOWLEDGE_DIR):
|
|
os.makedirs(KNOWLEDGE_DIR)
|
|
print(f"✓ 创建 knowledge 文件夹：{KNOWLEDGE_DIR}")
|
|
return 0
|
|
|
|
# 只计算文件，不计算目录
|
|
files = [f for f in os.listdir(KNOWLEDGE_DIR) if os.path.isfile(os.path.join(KNOWLEDGE_DIR, f))]
|
|
file_count = len(files)
|
|
print(f"知识库文件数量：{file_count}")
|
|
|
|
if files:
|
|
print(f" 已有文件：{', '.join(files)}")
|
|
|
|
return file_count
|
|
|
|
def run_process_knowledge_base():
|
|
"""运行 process_knowledge_base.py 建立知识库"""
|
|
if not os.path.exists(PROCESS_SCRIPT):
|
|
print(f"✗ 找不到 {PROCESS_SCRIPT}")
|
|
return False
|
|
|
|
print(f"\n正在运行知识库初始化脚本...")
|
|
print(f"命令：{sys.executable} \"{PROCESS_SCRIPT}\"")
|
|
|
|
try:
|
|
subprocess.run([sys.executable, PROCESS_SCRIPT], cwd=PROJECT_DIR)
|
|
print("✓ 知识库初始化完成")
|
|
return True
|
|
except Exception as e:
|
|
print(f"✗ 运行 process_knowledge_base.py 失败：{e}")
|
|
return False
|
|
|
|
def main():
|
|
print("=" * 60)
|
|
print("BRAIN 顾问助手 - 首次运行初始化")
|
|
print("=" * 60)
|
|
print()
|
|
|
|
# 0. 检查 api_key 是否已配置，未配置则直接退出
|
|
if not is_api_key_configured():
|
|
input("按 Enter 键退出...")
|
|
return
|
|
|
|
# 0.1 检查网络（ping + HTTPS），以确认代理已开启
|
|
print("[网络检查] 正在测试到 https://www.google.com 的连通性...")
|
|
if not check_network_reachable("https://www.google.com/generate_204"):
|
|
print("✗ 无法连通 google.com，请确认代理已开启，然后重新运行本脚本。")
|
|
input("按 Enter 键退出...")
|
|
return
|
|
print("✓ 网络检查通过")
|
|
|
|
# 1. 创建桌面快捷方式
|
|
print("[1/2] 创建桌面快捷方式...")
|
|
shortcut_ok = create_desktop_shortcut()
|
|
print()
|
|
|
|
# 2. 检查知识库
|
|
print("[2/2] 检查本地知识库...")
|
|
file_count = check_knowledge_base()
|
|
print()
|
|
|
|
if file_count <= 3:
|
|
print("⚠ 知识库文件较少（<= 3 个）")
|
|
print("建议建立本地知识库以增强 AI 回答效果")
|
|
print()
|
|
|
|
response = input("是否现在运行知识库初始化脚本？(y/n) [默认: y]: ").strip().lower()
|
|
if response != 'n':
|
|
run_process_knowledge_base()
|
|
else:
|
|
print(f"✓ 知识库已初始化（包含 {file_count} 个文件）")
|
|
|
|
print()
|
|
print("=" * 60)
|
|
print("初始化完成！")
|
|
print("=" * 60)
|
|
print()
|
|
|
|
if shortcut_ok:
|
|
print("📌 你现在可以从桌面快捷方式启动应用")
|
|
print(f" 或者运行：{sys.executable} \"{MAIN_SCRIPT}\"")
|
|
else:
|
|
print(f"请运行：{sys.executable} \"{MAIN_SCRIPT}\"")
|
|
print()
|
|
|
|
input("按 Enter 键退出...")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|