458 lines
15 KiB
Python
458 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
语义搜索MCP服务器
|
||
基于embedding向量进行语义相似度搜索
|
||
参考multi_keyword_search_server.py的实现方式
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import pickle
|
||
import sys
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
import numpy as np
|
||
from sentence_transformers import SentenceTransformer, util
|
||
|
||
# 延迟加载模型
|
||
embedder = None
|
||
|
||
def get_model(model_name_or_path='sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'):
|
||
"""获取模型实例（延迟加载）
|
||
|
||
Args:
|
||
model_name_or_path (str): 模型名称或本地路径
|
||
- 可以是 HuggingFace 模型名称
|
||
- 可以是本地模型路径
|
||
"""
|
||
global embedder
|
||
if embedder is None:
|
||
# 优先使用本地模型路径
|
||
local_model_path = "./models/paraphrase-multilingual-MiniLM-L12-v2"
|
||
|
||
# 检查本地模型是否存在
|
||
if os.path.exists(local_model_path):
|
||
print(f"使用本地模型: {local_model_path}")
|
||
embedder = SentenceTransformer(local_model_path, device='cpu')
|
||
else:
|
||
print(f"本地模型不存在，使用HuggingFace模型: {model_name_or_path}")
|
||
embedder = SentenceTransformer(model_name_or_path, device='cpu')
|
||
|
||
return embedder
|
||
|
||
|
||
def validate_file_path(file_path: str, allowed_dir: str) -> str:
|
||
"""验证文件路径是否在允许的目录内"""
|
||
# 转换为绝对路径
|
||
if not os.path.isabs(file_path):
|
||
file_path = os.path.abspath(file_path)
|
||
|
||
allowed_dir = os.path.abspath(allowed_dir)
|
||
|
||
# 检查路径是否在允许的目录内
|
||
if not file_path.startswith(allowed_dir):
|
||
raise ValueError(f"访问被拒绝: 路径 {file_path} 不在允许的目录 {allowed_dir} 内")
|
||
|
||
# 检查路径遍历攻击
|
||
if ".." in file_path:
|
||
raise ValueError(f"访问被拒绝: 检测到路径遍历攻击尝试")
|
||
|
||
return file_path
|
||
|
||
|
||
def get_allowed_directory():
|
||
"""获取允许访问的目录"""
|
||
# 从环境变量读取项目数据目录
|
||
project_dir = os.getenv("PROJECT_DATA_DIR", "./projects")
|
||
return os.path.abspath(project_dir)
|
||
|
||
|
||
def semantic_search(query: str, embeddings_file: str, top_k: int = 20) -> Dict[str, Any]:
|
||
"""执行语义搜索"""
|
||
if not query.strip():
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": "错误：查询不能为空"
|
||
}
|
||
]
|
||
}
|
||
|
||
# 处理项目目录限制
|
||
project_data_dir = get_allowed_directory()
|
||
|
||
# 验证embeddings文件路径
|
||
try:
|
||
# 解析相对路径
|
||
if not os.path.isabs(embeddings_file):
|
||
# 移除 projects/ 前缀（如果存在）
|
||
clean_path = embeddings_file
|
||
if clean_path.startswith('projects/'):
|
||
clean_path = clean_path[9:] # 移除 'projects/' 前缀
|
||
elif clean_path.startswith('./projects/'):
|
||
clean_path = clean_path[11:] # 移除 './projects/' 前缀
|
||
|
||
# 尝试在项目目录中查找文件
|
||
full_path = os.path.join(project_data_dir, clean_path.lstrip('./'))
|
||
if not os.path.exists(full_path):
|
||
# 如果直接路径不存在，尝试递归查找
|
||
found = find_file_in_project(clean_path, project_data_dir)
|
||
if found:
|
||
embeddings_file = found
|
||
else:
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"错误：在项目目录 {project_data_dir} 中未找到embeddings文件 {embeddings_file}"
|
||
}
|
||
]
|
||
}
|
||
else:
|
||
embeddings_file = full_path
|
||
else:
|
||
if not embeddings_file.startswith(project_data_dir):
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"错误：embeddings文件路径必须在项目目录 {project_data_dir} 内"
|
||
}
|
||
]
|
||
}
|
||
if not os.path.exists(embeddings_file):
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"错误：embeddings文件 {embeddings_file} 不存在"
|
||
}
|
||
]
|
||
}
|
||
except Exception as e:
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"错误：embeddings文件路径验证失败 - {str(e)}"
|
||
}
|
||
]
|
||
}
|
||
|
||
try:
|
||
# 加载嵌入数据
|
||
with open(embeddings_file, 'rb') as f:
|
||
embedding_data = pickle.load(f)
|
||
|
||
# 兼容新旧数据结构
|
||
if 'chunks' in embedding_data:
|
||
# 新的数据结构（使用chunks）
|
||
sentences = embedding_data['chunks']
|
||
sentence_embeddings = embedding_data['embeddings']
|
||
# 从embedding_data中获取模型路径（如果有的话）
|
||
model_path = embedding_data.get('model_path', 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
|
||
model = get_model(model_path)
|
||
else:
|
||
# 旧的数据结构（使用sentences）
|
||
sentences = embedding_data['sentences']
|
||
sentence_embeddings = embedding_data['embeddings']
|
||
model = get_model()
|
||
|
||
# 编码查询
|
||
query_embedding = model.encode(query, convert_to_tensor=True)
|
||
|
||
# 计算相似度
|
||
cos_scores = util.cos_sim(query_embedding, sentence_embeddings)[0]
|
||
|
||
# 获取top_k结果
|
||
top_results = np.argsort(-cos_scores.cpu().numpy())[:top_k]
|
||
|
||
# 格式化结果
|
||
results = []
|
||
for i, idx in enumerate(top_results):
|
||
sentence = sentences[idx]
|
||
score = cos_scores[idx].item()
|
||
results.append({
|
||
'rank': i + 1,
|
||
'content': sentence,
|
||
'similarity_score': score,
|
||
'file_path': embeddings_file
|
||
})
|
||
|
||
if not results:
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": "未找到匹配的结果"
|
||
}
|
||
]
|
||
}
|
||
|
||
# 格式化输出
|
||
formatted_output = "\n".join([
|
||
f"#{result['rank']} [similarity:{result['similarity_score']:.4f}]: {result['content']}"
|
||
for result in results
|
||
])
|
||
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": formatted_output
|
||
}
|
||
]
|
||
}
|
||
|
||
except FileNotFoundError:
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"错误：找不到embeddings文件 {embeddings_file}"
|
||
}
|
||
]
|
||
}
|
||
except Exception as e:
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"搜索时出错：{str(e)}"
|
||
}
|
||
]
|
||
}
|
||
|
||
|
||
def find_file_in_project(filename: str, project_dir: str) -> Optional[str]:
|
||
"""在项目目录中递归查找文件"""
|
||
for root, dirs, files in os.walk(project_dir):
|
||
if filename in files:
|
||
return os.path.join(root, filename)
|
||
return None
|
||
|
||
|
||
def get_model_info() -> Dict[str, Any]:
|
||
"""获取当前模型信息"""
|
||
try:
|
||
# 检查本地模型路径
|
||
local_model_path = "./models/paraphrase-multilingual-MiniLM-L12-v2"
|
||
|
||
if os.path.exists(local_model_path):
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"✅ 使用本地模型: {local_model_path}\n"
|
||
f"模型状态: 已加载\n"
|
||
f"设备: CPU\n"
|
||
f"说明: 避免从HuggingFace下载，提高响应速度"
|
||
}
|
||
]
|
||
}
|
||
else:
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"⚠️ 本地模型不存在: {local_model_path}\n"
|
||
f"将使用HuggingFace模型: sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2\n"
|
||
f"建议: 下载模型到本地以提高响应速度\n"
|
||
f"设备: CPU"
|
||
}
|
||
]
|
||
}
|
||
except Exception as e:
|
||
return {
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"❌ 获取模型信息失败: {str(e)}"
|
||
}
|
||
]
|
||
}
|
||
|
||
|
||
async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""Handle MCP request"""
|
||
try:
|
||
method = request.get("method")
|
||
params = request.get("params", {})
|
||
request_id = request.get("id")
|
||
|
||
if method == "initialize":
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"result": {
|
||
"protocolVersion": "2024-11-05",
|
||
"capabilities": {
|
||
"tools": {}
|
||
},
|
||
"serverInfo": {
|
||
"name": "semantic-search",
|
||
"version": "1.0.0"
|
||
}
|
||
}
|
||
}
|
||
|
||
elif method == "ping":
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"result": {
|
||
"pong": True
|
||
}
|
||
}
|
||
|
||
elif method == "tools/list":
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"result": {
|
||
"tools": [
|
||
{
|
||
"name": "semantic_search",
|
||
"description": "语义搜索工具，基于向量嵌入进行相似度搜索。格式：#[排名] [相似度分数]: [匹配内容]",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {
|
||
"type": "string",
|
||
"description": "搜索查询文本"
|
||
},
|
||
"embeddings_file": {
|
||
"type": "string",
|
||
"description": "embeddings pickle文件路径"
|
||
},
|
||
"top_k": {
|
||
"type": "integer",
|
||
"description": "返回结果的最大数量，默认20",
|
||
"default": 20
|
||
}
|
||
},
|
||
"required": ["query", "embeddings_file"]
|
||
}
|
||
},
|
||
{
|
||
"name": "get_model_info",
|
||
"description": "获取当前使用的模型信息，包括模型路径、加载状态等",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {},
|
||
"required": []
|
||
}
|
||
}
|
||
]
|
||
}
|
||
}
|
||
|
||
elif method == "tools/call":
|
||
tool_name = params.get("name")
|
||
arguments = params.get("arguments", {})
|
||
|
||
if tool_name == "semantic_search":
|
||
query = arguments.get("query", "")
|
||
embeddings_file = arguments.get("embeddings_file", "")
|
||
top_k = arguments.get("top_k", 20)
|
||
|
||
result = semantic_search(query, embeddings_file, top_k)
|
||
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"result": result
|
||
}
|
||
|
||
elif tool_name == "get_model_info":
|
||
result = get_model_info()
|
||
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"result": result
|
||
}
|
||
|
||
else:
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"error": {
|
||
"code": -32601,
|
||
"message": f"Unknown tool: {tool_name}"
|
||
}
|
||
}
|
||
|
||
else:
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"error": {
|
||
"code": -32601,
|
||
"message": f"Unknown method: {method}"
|
||
}
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": request.get("id"),
|
||
"error": {
|
||
"code": -32603,
|
||
"message": f"Internal error: {str(e)}"
|
||
}
|
||
}
|
||
|
||
|
||
async def main():
|
||
"""Main entry point."""
|
||
try:
|
||
while True:
|
||
# Read from stdin
|
||
line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline)
|
||
if not line:
|
||
break
|
||
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
|
||
try:
|
||
request = json.loads(line)
|
||
response = await handle_request(request)
|
||
|
||
# Write to stdout
|
||
sys.stdout.write(json.dumps(response) + "\n")
|
||
sys.stdout.flush()
|
||
|
||
except json.JSONDecodeError:
|
||
error_response = {
|
||
"jsonrpc": "2.0",
|
||
"error": {
|
||
"code": -32700,
|
||
"message": "Parse error"
|
||
}
|
||
}
|
||
sys.stdout.write(json.dumps(error_response) + "\n")
|
||
sys.stdout.flush()
|
||
|
||
except Exception as e:
|
||
error_response = {
|
||
"jsonrpc": "2.0",
|
||
"error": {
|
||
"code": -32603,
|
||
"message": f"Internal error: {str(e)}"
|
||
}
|
||
}
|
||
sys.stdout.write(json.dumps(error_response) + "\n")
|
||
sys.stdout.flush()
|
||
|
||
except KeyboardInterrupt:
|
||
pass
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|