sshclient/ui/config_dialog.py

"""
AI Agent 配置对话框
支持自定义 OpenAI 兼容 API（OpenAI / DeepSeek / Moonshot / Ollama 等）
"""
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QFormLayout, QLineEdit, QPushButton,
    QComboBox, QTextEdit, QLabel, QDialogButtonBox, QMessageBox, QGroupBox,
)

from core.ai_agent import AIAgent
from core.manager import load_ai_config, save_ai_config


PRESETS = [
    ("OpenAI 官方", "https://api.openai.com/v1", "gpt-4o-mini"),
    ("DeepSeek", "https://api.deepseek.com/v1", "deepseek-chat"),
    ("Moonshot Kimi", "https://api.moonshot.cn/v1", "moonshot-v1-8k"),
    ("通义千问 (DashScope 兼容)", "https://dashscope.aliyuncs.com/compatible-mode/v1", "qwen-turbo"),
    ("智谱 GLM (BigModel)", "https://open.bigmodel.cn/api/paas/v4", "glm-4-flash"),
    ("Ollama (本地)", "http://127.0.0.1:11434/v1", "qwen2.5:7b"),
    ("自定义", "", ""),
]


class AIConfigDialog(QDialog):
    def __init__(self, agent: AIAgent, parent=None):
        super().__init__(parent)
        self.agent = agent
        self.setWindowTitle("AI Agent 设置")
        self.resize(560, 480)
        self._build()

        # 用持久化的配置优先，再用 agent 当前值填充
        cfg = load_ai_config()
        self.api_key_edit.setText(cfg.get("api_key", "") or agent.api_key)
        self.base_url_edit.setText(cfg.get("base_url", "") or agent.base_url)
        self.model_edit.setText(cfg.get("model", "") or agent.model)
        self.system_prompt_edit.setPlainText(
            cfg.get("system_prompt", "") or agent.system_prompt
        )
        # 匹配预设
        for i, (_, url, model) in enumerate(PRESETS):
            if url and url == self.base_url_edit.text() and model == self.model_edit.text():
                self.preset_combo.setCurrentIndex(i)
                break

    def _build(self):
        layout = QVBoxLayout(self)

        # 预设
        preset_box = QGroupBox("服务商预设")
        pv = QVBoxLayout(preset_box)
        row = QHBoxLayout()
        row.addWidget(QLabel("快速选择:"))
        self.preset_combo = QComboBox()
        for name, _, _ in PRESETS:
            self.preset_combo.addItem(name)
        self.preset_combo.currentIndexChanged.connect(self._on_preset_changed)
        row.addWidget(self.preset_combo, 1)
        pv.addLayout(row)
        layout.addWidget(preset_box)

        # 表单
        form_box = QGroupBox("API 配置")
        form = QFormLayout(form_box)
        self.api_key_edit = QLineEdit()
        self.api_key_edit.setEchoMode(QLineEdit.Password)
        self.api_key_edit.setPlaceholderText("sk-...")
        form.addRow("API Key:", self.api_key_edit)
        self.base_url_edit = QLineEdit()
        self.base_url_edit.setPlaceholderText("https://api.openai.com/v1")
        form.addRow("Base URL:", self.base_url_edit)
        self.model_edit = QLineEdit()
        self.model_edit.setPlaceholderText("gpt-4o-mini")
        form.addRow("Model:", self.model_edit)
        layout.addWidget(form_box)

        # 系统提示
        sp_box = QGroupBox("系统提示词（可自定义 AI 行为）")
        sv = QVBoxLayout(sp_box)
        self.system_prompt_edit = QTextEdit()
        self.system_prompt_edit.setMaximumHeight(120)
        sv.addWidget(self.system_prompt_edit)
        layout.addWidget(sp_box)

        # 按钮
        bb = QDialogButtonBox()
        self.btn_test = bb.addButton("测试连接", QDialogButtonBox.ActionRole)
        self.btn_test.clicked.connect(self._test)
        bb.addButton(QDialogButtonBox.Save).clicked.connect(self._save)
        bb.addButton(QDialogButtonBox.Cancel).clicked.connect(self.reject)
        layout.addWidget(bb)

    def _on_preset_changed(self, idx: int):
        name, url, model = PRESETS[idx]
        if url:
            self.base_url_edit.setText(url)
            self.model_edit.setText(model)

    def _save(self):
        api_key = self.api_key_edit.text().strip()
        base_url = self.base_url_edit.text().strip()
        model = self.model_edit.text().strip()
        system_prompt = self.system_prompt_edit.toPlainText().strip()
        if not base_url or not model:
            QMessageBox.warning(self, "保存失败", "Base URL 和 Model 必填")
            return
        self.agent.update_config(api_key, base_url, model, system_prompt)
        save_ai_config({
            "api_key": api_key, "base_url": base_url,
            "model": model, "system_prompt": system_prompt,
        })
        self.accept()

    def _test(self):
        # 保存到 agent 临时测一下
        api_key = self.api_key_edit.text().strip()
        base_url = self.base_url_edit.text().strip()
        model = self.model_edit.text().strip()
        if not api_key or not base_url or not model:
            QMessageBox.warning(self, "提示", "请先填好 API Key / Base URL / Model")
            return
        # 保存一份用于测试的临时 agent
        tmp = AIAgent(api_key=api_key, base_url=base_url, model=model)
        try:
            from PyQt5.QtWidgets import QApplication
            QApplication.setOverrideCursor(Qt.WaitCursor)
            try:
                tmp.chat("ping", lambda n, a: "ok")
            finally:
                QApplication.restoreOverrideCursor()
            QMessageBox.information(self, "成功", "连接成功！")
        except Exception as e:
            QMessageBox.critical(self, "失败", f"连接失败：\n{e}")