start_mcp_server.py

#!/usr/bin/env python3
"""
Obsidian RAG MCP 服务器启动脚本
用于自动启动 MCP 服务器并处理依赖检查
"""

import os
import sys
import subprocess
import json
import logging
from pathlib import Path

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def check_python_version():
    """检查 Python 版本"""
    version = sys.version_info
    if version.major < 3 or (version.major == 3 and version.minor < 8):
        logger.error(f"需要 Python 3.8+，当前版本: {version.major}.{version.minor}")
        sys.exit(1)
    logger.info(f"Python 版本检查通过: {version.major}.{version.minor}.{version.micro}")

def check_virtual_env():
    """检查虚拟环境"""
    venv_path = Path("venv")
    if not venv_path.exists():
        logger.error("虚拟环境不存在，请先运行 ./setup.sh")
        sys.exit(1)

    # 检查是否在虚拟环境中
    if sys.prefix == sys.base_prefix:
        logger.info("激活虚拟环境...")
        if sys.platform == "win32":
            activate_script = venv_path / "Scripts" / "activate.bat"
            if activate_script.exists():
                logger.info("请在 Windows 中手动激活虚拟环境或使用 run.bat")
                sys.exit(1)
        else:
            activate_script = venv_path / "bin" / "activate"
            if activate_script.exists():
                # 重新在虚拟环境中执行此脚本
                os.execv(str(activate_script), [str(activate_script), "python", __file__] + sys.argv[1:])

    logger.info("虚拟环境检查通过")

def check_dependencies():
    """检查依赖"""
    try:
        import mcp
        logger.info("MCP 库检查通过")
    except ImportError:
        logger.error("MCP 库未安装，正在安装...")
        subprocess.run([sys.executable, "-m", "pip", "install", "mcp>=1.0.0"], check=True)
        logger.info("MCP 库安装完成")

def check_config():
    """检查配置文件"""
    env_file = Path(".env")
    if not env_file.exists():
        logger.error(".env 文件不存在，请从 .env.example 复制并配置")
        sys.exit(1)

    # 检查必要的环境变量
    with open(env_file, 'r', encoding='utf-8') as f:
        content = f.read()
        if "OBSIDIAN_VAULT_PATH" not in content:
            logger.error("请在 .env 文件中设置 OBSIDIAN_VAULT_PATH")
            sys.exit(1)

    logger.info("配置文件检查通过")

def check_index():
    """检查向量索引是否存在"""
    try:
        from src.config import get_config
        config = get_config()
        chroma_path = Path(config.chroma_persist_directory)

        if not chroma_path.exists():
            logger.warning("向量索引不存在，将自动构建...")
            build_index()
        else:
            logger.info("向量索引检查通过")
    except Exception as e:
        logger.error(f"检查索引失败: {e}")
        logger.info("尝试构建索引...")
        build_index()

def build_index():
    """构建向量索引"""
    try:
        logger.info("开始构建向量索引...")
        result = subprocess.run([
            sys.executable, "src/build_index.py"
        ], capture_output=True, text=True, encoding='utf-8')

        if result.returncode == 0:
            logger.info("向量索引构建完成")
        else:
            logger.error(f"构建索引失败: {result.stderr}")
            logger.error("请手动运行: ./run.sh build")
            sys.exit(1)
    except Exception as e:
        logger.error(f"构建索引时出错: {e}")
        sys.exit(1)

def create_claude_config():
    """创建 Claude Code MCP 配置"""
    config_dir = Path.home() / ".claude"
    plugins_dir = config_dir / "plugins"

    # 创建目录
    plugins_dir.mkdir(parents=True, exist_ok=True)

    config_file = plugins_dir / "config.json"

    # 当前项目路径
    project_path = Path.cwd()

    mcp_config = {
        "mcpServers": {
            "obsidian-rag": {
                "command": str(project_path / "venv" / "bin" / "python"),
                "args": [str(project_path / "src" / "mcp_server.py")],
                "cwd": str(project_path),
                "env": {
                    "PYTHONPATH": str(project_path / "src"),
                    "OBSIDIAN_VAULT_PATH": str(project_path / "test_vault"),
                    "EMBEDDING_MODEL": "paraphrase-multilingual-MiniLM-L12-v2",
                    "EMBEDDING_DEVICE": "cpu",
                    "CHROMA_PERSIST_DIRECTORY": "./data/chroma",
                    "CHROMA_COLLECTION_NAME": "obsidian_notes",
                    "CHUNK_SIZE": "500",
                    "CHUNK_OVERLAP": "50",
                    "TOP_K_RETRIEVAL": "5",
                    "SIMILARITY_THRESHOLD": "0.7",
                    "LOG_LEVEL": "INFO",
                    "LOG_FILE": "./logs/obsidian_rag.log",
                    "LLM_PROVIDER": "mock"
                }
            }
        }
    }

    # 读取现有配置
    existing_config = {}
    if config_file.exists():
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                existing_config = json.load(f)
        except Exception as e:
            logger.warning(f"读取现有配置失败: {e}")

    # 更新配置
    if "mcpServers" not in existing_config:
        existing_config["mcpServers"] = {}

    existing_config["mcpServers"]["obsidian-rag"] = mcp_config["mcpServers"]["obsidian-rag"]

    # 写入配置
    try:
        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(existing_config, f, indent=2, ensure_ascii=False)

        logger.info(f"Claude Code MCP 配置已更新: {config_file}")
        return True
    except Exception as e:
        logger.error(f"更新 Claude Code 配置失败: {e}")
        return False

def start_mcp_server():
    """启动 MCP 服务器"""
    try:
        logger.info("启动 Obsidian RAG MCP 服务器...")

        # 启动服务器进程
        process = subprocess.Popen([
            sys.executable, "src/mcp_server.py"
        ], stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
           text=True, encoding='utf-8', bufsize=1, universal_newlines=True)

        logger.info(f"MCP 服务器已启动 (PID: {process.pid})")
        logger.info("服务器正在运行，按 Ctrl+C 停止...")

        # 实时输出日志
        for line in iter(process.stdout.readline, ''):
            if line.strip():
                print(f"[MCP] {line.strip()}")

        # 等待进程结束
        process.wait()

    except KeyboardInterrupt:
        logger.info("收到停止信号，正在关闭 MCP 服务器...")
        if 'process' in locals():
            process.terminate()
            process.wait()
        logger.info("MCP 服务器已停止")
    except Exception as e:
        logger.error(f"启动 MCP 服务器失败: {e}")
        sys.exit(1)

def main():
    """主函数"""
    print("=" * 60)
    print("Obsidian RAG MCP 服务器启动器")
    print("=" * 60)

    # 检查环境
    logger.info("检查运行环境...")
    check_python_version()
    check_virtual_env()
    check_dependencies()
    check_config()
    check_index()

    # 创建 Claude Code 配置
    logger.info("配置 Claude Code MCP 集成...")
    if create_claude_config():
        logger.info("Claude Code MCP 配置成功")
    else:
        logger.warning("Claude Code MCP 配置失败，请手动配置")

    print("\n" + "=" * 60)
    print("环境检查完成，准备启动 MCP 服务器")
    print("=" * 60)
    print("\n可用的 MCP 工具:")
    print("- vector_search: 向量检索")
    print("- rag_chat: RAG 对话")
    print("- get_document_info: 获取文档信息")
    print("- build_index: 构建索引")
    print("- search_suggestions: 搜索建议")
    print("- list_documents: 列出文档")
    print("\n现在可以在 Claude Code 中使用这些工具!")
    print("=" * 60)

    # 启动服务器
    start_mcp_server()

if __name__ == "__main__":
    main()