opensource/WeKnora
3
0
Fork 0
mirror of https://github.com/Tencent/WeKnora.git synced 2026-08-29 02:04:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor code for improved readability and consistency

Browse Source
This commit is contained in:
begoniezhao
2025-12-01 18:15:39 +08:00
parent 5d6729fc0c
commit f5963276d6
8 changed files with 461 additions and 292 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mcp-server/__init__.py
+1 -1
View File
@@ -11,4 +11,4 @@ __description__ = "WeKnora MCP Server - Model Context Protocol server for WeKnor
from .weknora_mcp_server import WeKnoraClient, run
__all__ = ["WeKnoraClient", "run"]
__all__ = ["WeKnoraClient", "run"]
mcp-server/main.py
+34 -31
View File
@@ -9,12 +9,13 @@ WeKnora MCP Server 主入口点
3. weknora-mcp-server (安装后)
"""
import argparse
import asyncio
import os
import sys
import asyncio
import argparse
from pathlib import Path
def setup_environment():
"""设置环境和路径"""
# 确保当前目录在 Python 路径中
@@ -22,35 +23,39 @@ def setup_environment():
if str(current_dir) not in sys.path:
sys.path.insert(0, str(current_dir))
def check_dependencies():
"""检查依赖是否已安装"""
try:
import mcp
import requests
return True
except ImportError as e:
print(f"缺少依赖: {e}")
print("请运行: pip install -r requirements.txt")
return False
def check_environment_variables():
"""检查环境变量配置"""
base_url = os.getenv("WEKNORA_BASE_URL")
api_key = os.getenv("WEKNORA_API_KEY")
print("=== WeKnora MCP Server 环境检查 ===")
print(f"Base URL: {base_url or 'http://localhost:8080/api/v1 (默认)'}")
print(f"API Key: {'已设置' if api_key else '未设置 (警告)'}")
if not base_url:
print("提示: 可以设置 WEKNORA_BASE_URL 环境变量")
if not api_key:
print("警告: 建议设置 WEKNORA_API_KEY 环境变量")
print("=" * 40)
return True
def parse_arguments():
"""解析命令行参数"""
parser = argparse.ArgumentParser(
@@ -65,61 +70,56 @@ def parse_arguments():
环境变量:
WEKNORA_BASE_URL WeKnora API 基础 URL (默认: http://localhost:8080/api/v1)
WEKNORA_API_KEY WeKnora API 密钥
"""
""",
)
parser.add_argument(
"--check-only",
action="store_true",
help="仅检查环境配置，不启动服务器"
"--check-only", action="store_true", help="仅检查环境配置，不启动服务器"
)
parser.add_argument("--verbose", "-v", action="store_true", help="启用详细日志输出")
parser.add_argument(
"--verbose", "-v",
action="store_true",
help="启用详细日志输出"
"--version", action="version", version="WeKnora MCP Server 1.0.0"
)
parser.add_argument(
"--version",
action="version",
version="WeKnora MCP Server 1.0.0"
)
return parser.parse_args()
async def main():
"""主函数"""
args = parse_arguments()
# 设置环境
setup_environment()
# 检查依赖
if not check_dependencies():
sys.exit(1)
# 检查环境变量
check_environment_variables()
# 如果只是检查环境，则退出
if args.check_only:
print("环境检查完成。")
return
# 设置日志级别
if args.verbose:
import logging
logging.basicConfig(level=logging.DEBUG)
print("已启用详细日志模式")
try:
print("正在启动 WeKnora MCP Server...")
# 导入并运行服务器
from weknora_mcp_server import run
await run()
except ImportError as e:
print(f"导入错误: {e}")
print("请确保所有文件都在正确的位置")
@@ -130,12 +130,15 @@ async def main():
print(f"服务器运行错误: {e}")
if args.verbose:
import traceback
traceback.print_exc()
sys.exit(1)
def sync_main():
"""同步版本的主函数，用于 entry_points"""
asyncio.run(main())
if __name__ == "__main__":
asyncio.run(main())
asyncio.run(main())
mcp-server/run.py
+8 -5
View File
@@ -6,29 +6,31 @@ WeKnora MCP Server 便捷启动脚本
对于更多选项，请使用 main.py
"""
import sys
import os
import sys
from pathlib import Path
def main():
"""简单的启动函数"""
# 添加当前目录到 Python 路径
current_dir = Path(__file__).parent.absolute()
if str(current_dir) not in sys.path:
sys.path.insert(0, str(current_dir))
# 检查环境变量
base_url = os.getenv("WEKNORA_BASE_URL", "http://localhost:8080/api/v1")
api_key = os.getenv("WEKNORA_API_KEY", "")
print("WeKnora MCP Server")
print(f"Base URL: {base_url}")
print(f"API Key: {'已设置' if api_key else '未设置'}")
print("-" * 40)
try:
# 导入并运行
from main import sync_main
sync_main()
except ImportError:
print("错误: 无法导入必要模块")
@@ -40,5 +42,6 @@ def main():
print(f"错误: {e}")
sys.exit(1)
if __name__ == "__main__":
main()
main()
mcp-server/run_server.py
+13 -7
View File
@@ -3,31 +3,36 @@
WeKnora MCP Server 启动脚本
"""
import asyncio
import os
import sys
import asyncio
def check_environment():
"""检查环境配置"""
base_url = os.getenv("WEKNORA_BASE_URL")
api_key = os.getenv("WEKNORA_API_KEY")
if not base_url:
print("警告: WEKNORA_BASE_URL 环境变量未设置，使用默认值: http://localhost:8080/api/v1")
print(
"警告: WEKNORA_BASE_URL 环境变量未设置，使用默认值: http://localhost:8080/api/v1"
)
if not api_key:
print("警告: WEKNORA_API_KEY 环境变量未设置")
print(f"WeKnora Base URL: {base_url or 'http://localhost:8080/api/v1'}")
print(f"API Key: {'已设置' if api_key else '未设置'}")
def main():
"""主函数"""
print("启动 WeKnora MCP Server...")
check_environment()
try:
from weknora_mcp_server import run
asyncio.run(run())
except ImportError as e:
print(f"导入错误: {e}")
@@ -39,5 +44,6 @@ def main():
print(f"服务器运行错误: {e}")
sys.exit(1)
if __name__ == "__main__":
main()
main()
mcp-server/setup.py
+7 -3
View File
@@ -4,7 +4,7 @@ WeKnora MCP Server 安装脚本
"""
from setuptools import setup
import os
# 读取 README 文件
def read_readme():
@@ -14,14 +14,18 @@ def read_readme():
except FileNotFoundError:
return "WeKnora MCP Server - Model Context Protocol server for WeKnora API"
# 读取依赖
def read_requirements():
try:
with open("requirements.txt", "r", encoding="utf-8") as f:
return [line.strip() for line in f if line.strip() and not line.startswith("#")]
return [
line.strip() for line in f if line.strip() and not line.startswith("#")
]
except FileNotFoundError:
return ["mcp>=1.0.0", "requests>=2.31.0"]
setup(
name="weknora-mcp-server",
version="1.0.0",
@@ -58,4 +62,4 @@ setup(
("", ["README.md", "requirements.txt", "LICENSE"]),
],
keywords="mcp model-context-protocol weknora knowledge-management api-server",
)
)
mcp-server/test_imports.py
+8 -2
View File
@@ -5,34 +5,40 @@
try:
import mcp.types as types
print("✓ mcp.types 导入成功")
except ImportError as e:
print(f"✗ mcp.types 导入失败: {e}")
try:
from mcp.server import Server, NotificationOptions
from mcp.server import NotificationOptions, Server
print("✓ mcp.server 导入成功")
except ImportError as e:
print(f"✗ mcp.server 导入失败: {e}")
try:
import mcp.server.stdio
print("✓ mcp.server.stdio 导入成功")
except ImportError as e:
print(f"✗ mcp.server.stdio 导入失败: {e}")
try:
from mcp.server.models import InitializationOptions
print("✓ InitializationOptions 从 mcp.server.models 导入成功")
except ImportError:
try:
from mcp import InitializationOptions
print("✓ InitializationOptions 从 mcp 导入成功")
except ImportError as e:
print(f"✗ InitializationOptions 导入失败: {e}")
# 检查 MCP 包结构
import mcp
print(f"\nMCP 包版本: {getattr(mcp, '__version__', '未知')}")
print(f"MCP 包路径: {mcp.__file__}")
print(f"MCP 包内容: {dir(mcp)}")
print(f"MCP 包内容: {dir(mcp)}")
mcp-server/test_module.py
+54 -42
View File
@@ -6,90 +6,98 @@ WeKnora MCP Server 模组测试脚本
"""
import os
import sys
import subprocess
import importlib.util
import sys
from pathlib import Path
def test_imports():
"""测试模块导入"""
print("=== 测试模块导入 ===")
try:
# 测试基础依赖
import mcp
print("✓ mcp 模块导入成功")
import requests
print("✓ requests 模块导入成功")
# 测试主模块
import weknora_mcp_server
print("✓ weknora_mcp_server 模块导入成功")
# 测试包导入
from weknora_mcp_server import WeKnoraClient, run
print("✓ WeKnoraClient 和 run 函数导入成功")
# 测试主入口点
import main
print("✓ main 模块导入成功")
return True
except ImportError as e:
print(f"✗ 导入失败: {e}")
return False
def test_environment():
"""测试环境配置"""
print("\n=== 测试环境配置 ===")
base_url = os.getenv("WEKNORA_BASE_URL")
api_key = os.getenv("WEKNORA_API_KEY")
print(f"WEKNORA_BASE_URL: {base_url or '未设置 (将使用默认值)'}")
print(f"WEKNORA_API_KEY: {'已设置' if api_key else '未设置'}")
if not base_url:
print("提示: 可以设置环境变量 WEKNORA_BASE_URL")
if not api_key:
print("提示: 建议设置环境变量 WEKNORA_API_KEY")
return True
def test_client_creation():
"""测试客户端创建"""
print("\n=== 测试客户端创建 ===")
try:
from weknora_mcp_server import WeKnoraClient
base_url = os.getenv("WEKNORA_BASE_URL", "http://localhost:8080/api/v1")
api_key = os.getenv("WEKNORA_API_KEY", "test_key")
client = WeKnoraClient(base_url, api_key)
print("✓ WeKnoraClient 创建成功")
# 检查客户端属性
assert client.base_url == base_url
assert client.api_key == api_key
print("✓ 客户端配置正确")
return True
except Exception as e:
print(f"✗ 客户端创建失败: {e}")
return False
def test_file_structure():
"""测试文件结构"""
print("\n=== 测试文件结构 ===")
required_files = [
"__init__.py",
"main.py",
"main.py",
"run_server.py",
"weknora_mcp_server.py",
"requirements.txt",
@@ -98,9 +106,9 @@ def test_file_structure():
"README.md",
"INSTALL.md",
"LICENSE",
"MANIFEST.in"
"MANIFEST.in",
]
missing_files = []
for file in required_files:
if Path(file).exists():
@@ -108,25 +116,26 @@ def test_file_structure():
else:
print(f"✗ {file} (缺失)")
missing_files.append(file)
if missing_files:
print(f"缺失文件: {missing_files}")
return False
print("✓ 所有必需文件都存在")
return True
def test_entry_points():
"""测试入口点"""
print("\n=== 测试入口点 ===")
# 测试 main.py 的帮助选项
try:
result = subprocess.run(
[sys.executable, "main.py", "--help"],
capture_output=True,
text=True,
timeout=10
timeout=10,
)
if result.returncode == 0:
print("✓ main.py --help 工作正常")
@@ -139,14 +148,14 @@ def test_entry_points():
except Exception as e:
print(f"✗ main.py --help 错误: {e}")
return False
# 测试环境检查
try:
result = subprocess.run(
[sys.executable, "main.py", "--check-only"],
capture_output=True,
text=True,
timeout=10
timeout=10,
)
if result.returncode == 0:
print("✓ main.py --check-only 工作正常")
@@ -159,42 +168,44 @@ def test_entry_points():
except Exception as e:
print(f"✗ main.py --check-only 错误: {e}")
return False
return True
def test_package_installation():
"""测试包安装 (开发模式)"""
print("\n=== 测试包安装 ===")
try:
# 检查是否可以以开发模式安装
result = subprocess.run(
[sys.executable, "setup.py", "check"],
capture_output=True,
text=True,
timeout=30
timeout=30,
)
if result.returncode == 0:
print("✓ setup.py 检查通过")
else:
print(f"✗ setup.py 检查失败: {result.stderr}")
return False
except subprocess.TimeoutExpired:
print("✗ setup.py 检查超时")
return False
except Exception as e:
print(f"✗ setup.py 检查错误: {e}")
return False
return True
def main():
"""运行所有测试"""
print("WeKnora MCP Server 模组测试")
print("=" * 50)
tests = [
("模块导入", test_imports),
("环境配置", test_environment),
@@ -203,10 +214,10 @@ def main():
("入口点", test_entry_points),
("包安装", test_package_installation),
]
passed = 0
total = len(tests)
for test_name, test_func in tests:
try:
if test_func():
@@ -215,10 +226,10 @@ def main():
print(f"测试失败: {test_name}")
except Exception as e:
print(f"测试异常: {test_name} - {e}")
print("\n" + "=" * 50)
print(f"测试结果: {passed}/{total} 通过")
if passed == total:
print("✓ 所有测试通过！模组可以正常使用。")
return True
@@ -226,6 +237,7 @@ def main():
print("✗ 部分测试失败，请检查上述错误。")
return False
if __name__ == "__main__":
success = main()
sys.exit(0 if success else 1)
sys.exit(0 if success else 1)
mcp-server/weknora_mcp_server.py
+336 -201
View File
@@ -5,17 +5,17 @@ WeKnora MCP Server
A Model Context Protocol server that provides access to the WeKnora knowledge management API.
"""
import os
import json
import logging
from typing import Dict, List, Any, Optional
from datetime import datetime
import requests
from requests.exceptions import RequestException
import os
from typing import Any, Dict
import mcp.server.stdio
import mcp.types as types
import requests
from mcp.server import NotificationOptions, Server
from mcp.server.models import InitializationOptions
from requests.exceptions import RequestException
# Set up logging configuration for the MCP server
logging.basicConfig(level=logging.INFO)
@@ -25,9 +25,10 @@ logger = logging.getLogger(__name__)
WEKNORA_BASE_URL = os.getenv("WEKNORA_BASE_URL", "http://localhost:8080/api/v1")
WEKNORA_API_KEY = os.getenv("WEKNORA_API_KEY", "")
class WeKnoraClient:
"""Client for interacting with WeKnora API"""
def __init__(self, base_url: str, api_key: str):
"""Initialize the WeKnora API client with base URL and authentication"""
self.base_url = base_url
@@ -35,19 +36,21 @@ class WeKnoraClient:
# Create a persistent session for connection pooling and performance
self.session = requests.Session()
# Set default headers for all requests
self.session.headers.update({
"X-API-Key": api_key, # API key for authentication
"Content-Type": "application/json" # Default content type
})
self.session.headers.update(
{
"X-API-Key": api_key, # API key for authentication
"Content-Type": "application/json", # Default content type
}
)
def _request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
"""Make a request to the WeKnora API
Args:
method: HTTP method (GET, POST, PUT, DELETE)
endpoint: API endpoint path
**kwargs: Additional arguments to pass to requests
Returns:
JSON response as dictionary
"""
@@ -62,103 +65,123 @@ class WeKnoraClient:
except RequestException as e:
logger.error(f"API request failed: {e}")
raise
# Tenant Management - Methods for managing multi-tenant configurations
def create_tenant(self, name: str, description: str, business: str, retriever_engines: Dict) -> Dict:
def create_tenant(
self, name: str, description: str, business: str, retriever_engines: Dict
) -> Dict:
"""Create a new tenant with specified configuration"""
data = {
"name": name,
"description": description,
"business": business,
"retriever_engines": retriever_engines # Configuration for search engines
"retriever_engines": retriever_engines, # Configuration for search engines
}
return self._request("POST", "/tenants", json=data)
def get_tenant(self, tenant_id: str) -> Dict:
"""Get tenant information"""
return self._request("GET", f"/tenants/{tenant_id}")
def list_tenants(self) -> Dict:
"""List all tenants"""
return self._request("GET", "/tenants")
# Knowledge Base Management - Methods for managing knowledge bases
def create_knowledge_base(self, name: str, description: str, config: Dict) -> Dict:
"""Create a new knowledge base with chunking and model configuration"""
data = {
"name": name,
"description": description,
**config # Merge additional configuration (chunking, models, etc.)
**config, # Merge additional configuration (chunking, models, etc.)
}
return self._request("POST", "/knowledge-bases", json=data)
def list_knowledge_bases(self) -> Dict:
"""List all knowledge bases"""
return self._request("GET", "/knowledge-bases")
def get_knowledge_base(self, kb_id: str) -> Dict:
"""Get knowledge base details"""
return self._request("GET", f"/knowledge-bases/{kb_id}")
def update_knowledge_base(self, kb_id: str, updates: Dict) -> Dict:
"""Update knowledge base"""
return self._request("PUT", f"/knowledge-bases/{kb_id}", json=updates)
def delete_knowledge_base(self, kb_id: str) -> Dict:
"""Delete knowledge base"""
return self._request("DELETE", f"/knowledge-bases/{kb_id}")
def hybrid_search(self, kb_id: str, query: str, config: Dict) -> Dict:
"""Perform hybrid search combining vector and keyword search"""
data = {
"query_text": query,
**config # Include thresholds and match count
**config, # Include thresholds and match count
}
return self._request("GET", f"/knowledge-bases/{kb_id}/hybrid-search", json=data)
return self._request(
"GET", f"/knowledge-bases/{kb_id}/hybrid-search", json=data
)
# Knowledge Management - Methods for creating and managing knowledge entries
def create_knowledge_from_file(self, kb_id: str, file_path: str, enable_multimodel: bool = True) -> Dict:
def create_knowledge_from_file(
self, kb_id: str, file_path: str, enable_multimodel: bool = True
) -> Dict:
"""Create knowledge from a local file with optional multimodal processing"""
with open(file_path, 'rb') as f:
files = {'file': f}
data = {'enable_multimodel': str(enable_multimodel).lower()}
with open(file_path, "rb") as f:
files = {"file": f}
data = {"enable_multimodel": str(enable_multimodel).lower()}
# Temporarily remove Content-Type header for multipart/form-data request
# (requests will set it automatically with boundary)
headers = self.session.headers.copy()
del headers['Content-Type']
del headers["Content-Type"]
# Use requests.post directly instead of session to avoid header conflicts
response = requests.post(
f"{self.base_url}/knowledge-bases/{kb_id}/knowledge/file",
headers=headers,
files=files,
data=data
data=data,
)
response.raise_for_status()
return response.json()
def create_knowledge_from_url(self, kb_id: str, url: str, enable_multimodel: bool = True) -> Dict:
def create_knowledge_from_url(
self, kb_id: str, url: str, enable_multimodel: bool = True
) -> Dict:
"""Create knowledge from a web URL with optional multimodal processing"""
data = {
"url": url, # Web URL to fetch and process
"enable_multimodel": enable_multimodel # Enable image/multimodal extraction
"enable_multimodel": enable_multimodel, # Enable image/multimodal extraction
}
return self._request("POST", f"/knowledge-bases/{kb_id}/knowledge/url", json=data)
return self._request(
"POST", f"/knowledge-bases/{kb_id}/knowledge/url", json=data
)
def list_knowledge(self, kb_id: str, page: int = 1, page_size: int = 20) -> Dict:
"""List knowledge in a knowledge base"""
params = {"page": page, "page_size": page_size}
return self._request("GET", f"/knowledge-bases/{kb_id}/knowledge", params=params)
return self._request(
"GET", f"/knowledge-bases/{kb_id}/knowledge", params=params
)
def get_knowledge(self, knowledge_id: str) -> Dict:
"""Get knowledge details"""
return self._request("GET", f"/knowledge/{knowledge_id}")
def delete_knowledge(self, knowledge_id: str) -> Dict:
"""Delete knowledge"""
return self._request("DELETE", f"/knowledge/{knowledge_id}")
# Model Management - Methods for managing AI models (LLM, Embedding, Rerank)
def create_model(self, name: str, model_type: str, source: str, description: str, parameters: Dict, is_default: bool = False) -> Dict:
def create_model(
self,
name: str,
model_type: str,
source: str,
description: str,
parameters: Dict,
is_default: bool = False,
) -> Dict:
"""Create a new AI model configuration"""
data = {
"name": name,
@@ -166,40 +189,40 @@ class WeKnoraClient:
"source": source, # local, openai, etc.
"description": description,
"parameters": parameters, # API keys, base URLs, etc.
"is_default": is_default # Set as default model for this type
"is_default": is_default, # Set as default model for this type
}
return self._request("POST", "/models", json=data)
def list_models(self) -> Dict:
"""List all models"""
return self._request("GET", "/models")
def get_model(self, model_id: str) -> Dict:
"""Get model details"""
return self._request("GET", f"/models/{model_id}")
# Session Management - Methods for managing chat sessions
def create_session(self, kb_id: str, strategy: Dict) -> Dict:
"""Create a new chat session with conversation strategy"""
data = {
"knowledge_base_id": kb_id, # Knowledge base to query
"session_strategy": strategy # Conversation settings (max rounds, rewrite, etc.)
"session_strategy": strategy, # Conversation settings (max rounds, rewrite, etc.)
}
return self._request("POST", "/sessions", json=data)
def get_session(self, session_id: str) -> Dict:
"""Get session details"""
return self._request("GET", f"/sessions/{session_id}")
def list_sessions(self, page: int = 1, page_size: int = 20) -> Dict:
"""List sessions"""
params = {"page": page, "page_size": page_size}
return self._request("GET", "/sessions", params=params)
def delete_session(self, session_id: str) -> Dict:
"""Delete session"""
return self._request("DELETE", f"/sessions/{session_id}")
# Chat Functionality - Methods for conversational interactions
def chat(self, session_id: str, query: str) -> Dict:
"""Send a chat message and get AI response"""
@@ -207,22 +230,26 @@ class WeKnoraClient:
# Note: The actual API returns Server-Sent Events (SSE) stream
# This simplified version returns the complete response
return self._request("POST", f"/knowledge-chat/{session_id}", json=data)
# Chunk Management - Methods for managing knowledge chunks (text segments)
def list_chunks(self, knowledge_id: str, page: int = 1, page_size: int = 20) -> Dict:
def list_chunks(
self, knowledge_id: str, page: int = 1, page_size: int = 20
) -> Dict:
"""List text chunks of a knowledge entry with pagination"""
params = {"page": page, "page_size": page_size}
return self._request("GET", f"/chunks/{knowledge_id}", params=params)
def delete_chunk(self, knowledge_id: str, chunk_id: str) -> Dict:
"""Delete a chunk"""
return self._request("DELETE", f"/chunks/{knowledge_id}/{chunk_id}")
# Initialize MCP server instance
app = Server("weknora-server")
# Initialize WeKnora API client with configuration
client = WeKnoraClient(WEKNORA_BASE_URL, WEKNORA_API_KEY)
# Tool definitions - Register all available tools for the MCP protocol
@app.list_tools()
async def handle_list_tools() -> list[types.Tool]:
@@ -236,7 +263,10 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"name": {"type": "string", "description": "Tenant name"},
"description": {"type": "string", "description": "Tenant description"},
"description": {
"type": "string",
"description": "Tenant description",
},
"business": {"type": "string", "description": "Business type"},
"retriever_engines": {
"type": "object",
@@ -248,22 +278,21 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"retriever_type": {"type": "string"},
"retriever_engine_type": {"type": "string"}
}
}
"retriever_engine_type": {"type": "string"},
},
},
}
}
}
},
},
},
"required": ["name", "description", "business"]
}
"required": ["name", "description", "business"],
},
),
types.Tool(
name="list_tenants",
description="List all tenants",
inputSchema={"type": "object", "properties": {}}
inputSchema={"type": "object", "properties": {}},
),
# Knowledge Base Management
types.Tool(
name="create_knowledge_base",
@@ -272,17 +301,26 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"name": {"type": "string", "description": "Knowledge base name"},
"description": {"type": "string", "description": "Knowledge base description"},
"embedding_model_id": {"type": "string", "description": "Embedding model ID"},
"summary_model_id": {"type": "string", "description": "Summary model ID"}
"description": {
"type": "string",
"description": "Knowledge base description",
},
"embedding_model_id": {
"type": "string",
"description": "Embedding model ID",
},
"summary_model_id": {
"type": "string",
"description": "Summary model ID",
},
},
"required": ["name", "description"]
}
"required": ["name", "description"],
},
),
types.Tool(
name="list_knowledge_bases",
description="List all knowledge bases",
inputSchema={"type": "object", "properties": {}}
inputSchema={"type": "object", "properties": {}},
),
types.Tool(
name="get_knowledge_base",
@@ -292,8 +330,8 @@ async def handle_list_tools() -> list[types.Tool]:
"properties": {
"kb_id": {"type": "string", "description": "Knowledge base ID"}
},
"required": ["kb_id"]
}
"required": ["kb_id"],
},
),
types.Tool(
name="delete_knowledge_base",
@@ -303,8 +341,8 @@ async def handle_list_tools() -> list[types.Tool]:
"properties": {
"kb_id": {"type": "string", "description": "Knowledge base ID"}
},
"required": ["kb_id"]
}
"required": ["kb_id"],
},
),
types.Tool(
name="hybrid_search",
@@ -314,14 +352,25 @@ async def handle_list_tools() -> list[types.Tool]:
"properties": {
"kb_id": {"type": "string", "description": "Knowledge base ID"},
"query": {"type": "string", "description": "Search query"},
"vector_threshold": {"type": "number", "description": "Vector similarity threshold", "default": 0.5},
"keyword_threshold": {"type": "number", "description": "Keyword match threshold", "default": 0.3},
"match_count": {"type": "integer", "description": "Number of results to return", "default": 5}
"vector_threshold": {
"type": "number",
"description": "Vector similarity threshold",
"default": 0.5,
},
"keyword_threshold": {
"type": "number",
"description": "Keyword match threshold",
"default": 0.3,
},
"match_count": {
"type": "integer",
"description": "Number of results to return",
"default": 5,
},
},
"required": ["kb_id", "query"]
}
"required": ["kb_id", "query"],
},
),
# Knowledge Management
types.Tool(
name="create_knowledge_from_file",
@@ -330,11 +379,18 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"kb_id": {"type": "string", "description": "Knowledge base ID"},
"file_path": {"type": "string", "description": "Absolute path to the local file on the server"},
"enable_multimodel": {"type": "boolean", "description": "Enable multimodal processing", "default": True}
"file_path": {
"type": "string",
"description": "Absolute path to the local file on the server",
},
"enable_multimodel": {
"type": "boolean",
"description": "Enable multimodal processing",
"default": True,
},
},
"required": ["kb_id", "file_path"]
}
"required": ["kb_id", "file_path"],
},
),
types.Tool(
name="create_knowledge_from_url",
@@ -343,11 +399,18 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"kb_id": {"type": "string", "description": "Knowledge base ID"},
"url": {"type": "string", "description": "URL to create knowledge from"},
"enable_multimodel": {"type": "boolean", "description": "Enable multimodal processing", "default": True}
"url": {
"type": "string",
"description": "URL to create knowledge from",
},
"enable_multimodel": {
"type": "boolean",
"description": "Enable multimodal processing",
"default": True,
},
},
"required": ["kb_id", "url"]
}
"required": ["kb_id", "url"],
},
),
types.Tool(
name="list_knowledge",
@@ -356,11 +419,19 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"kb_id": {"type": "string", "description": "Knowledge base ID"},
"page": {"type": "integer", "description": "Page number", "default": 1},
"page_size": {"type": "integer", "description": "Page size", "default": 20}
"page": {
"type": "integer",
"description": "Page number",
"default": 1,
},
"page_size": {
"type": "integer",
"description": "Page size",
"default": 20,
},
},
"required": ["kb_id"]
}
"required": ["kb_id"],
},
),
types.Tool(
name="get_knowledge",
@@ -370,8 +441,8 @@ async def handle_list_tools() -> list[types.Tool]:
"properties": {
"knowledge_id": {"type": "string", "description": "Knowledge ID"}
},
"required": ["knowledge_id"]
}
"required": ["knowledge_id"],
},
),
types.Tool(
name="delete_knowledge",
@@ -381,10 +452,9 @@ async def handle_list_tools() -> list[types.Tool]:
"properties": {
"knowledge_id": {"type": "string", "description": "Knowledge ID"}
},
"required": ["knowledge_id"]
}
"required": ["knowledge_id"],
},
),
# Model Management
types.Tool(
name="create_model",
@@ -393,20 +463,42 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"name": {"type": "string", "description": "Model name"},
"type": {"type": "string", "description": "Model type (KnowledgeQA, Embedding, Rerank)"},
"source": {"type": "string", "description": "Model source", "default": "local"},
"description": {"type": "string", "description": "Model description"},
"base_url": {"type": "string", "description": "Model API base URL", "default": ""},
"api_key": {"type": "string", "description": "Model API key", "default": ""},
"is_default": {"type": "boolean", "description": "Set as default model", "default": False}
"type": {
"type": "string",
"description": "Model type (KnowledgeQA, Embedding, Rerank)",
},
"source": {
"type": "string",
"description": "Model source",
"default": "local",
},
"description": {
"type": "string",
"description": "Model description",
},
"base_url": {
"type": "string",
"description": "Model API base URL",
"default": "",
},
"api_key": {
"type": "string",
"description": "Model API key",
"default": "",
},
"is_default": {
"type": "boolean",
"description": "Set as default model",
"default": False,
},
},
"required": ["name", "type", "description"]
}
"required": ["name", "type", "description"],
},
),
types.Tool(
name="list_models",
description="List all models",
inputSchema={"type": "object", "properties": {}}
inputSchema={"type": "object", "properties": {}},
),
types.Tool(
name="get_model",
@@ -416,10 +508,9 @@ async def handle_list_tools() -> list[types.Tool]:
"properties": {
"model_id": {"type": "string", "description": "Model ID"}
},
"required": ["model_id"]
}
"required": ["model_id"],
},
),
# Session Management
types.Tool(
name="create_session",
@@ -428,13 +519,28 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"kb_id": {"type": "string", "description": "Knowledge base ID"},
"max_rounds": {"type": "integer", "description": "Maximum conversation rounds", "default": 5},
"enable_rewrite": {"type": "boolean", "description": "Enable query rewriting", "default": True},
"fallback_response": {"type": "string", "description": "Fallback response", "default": "Sorry, I cannot answer this question."},
"summary_model_id": {"type": "string", "description": "Summary model ID"}
"max_rounds": {
"type": "integer",
"description": "Maximum conversation rounds",
"default": 5,
},
"enable_rewrite": {
"type": "boolean",
"description": "Enable query rewriting",
"default": True,
},
"fallback_response": {
"type": "string",
"description": "Fallback response",
"default": "Sorry, I cannot answer this question.",
},
"summary_model_id": {
"type": "string",
"description": "Summary model ID",
},
},
"required": ["kb_id"]
}
"required": ["kb_id"],
},
),
types.Tool(
name="get_session",
@@ -444,8 +550,8 @@ async def handle_list_tools() -> list[types.Tool]:
"properties": {
"session_id": {"type": "string", "description": "Session ID"}
},
"required": ["session_id"]
}
"required": ["session_id"],
},
),
types.Tool(
name="list_sessions",
@@ -453,10 +559,18 @@ async def handle_list_tools() -> list[types.Tool]:
inputSchema={
"type": "object",
"properties": {
"page": {"type": "integer", "description": "Page number", "default": 1},
"page_size": {"type": "integer", "description": "Page size", "default": 20}
}
}
"page": {
"type": "integer",
"description": "Page number",
"default": 1,
},
"page_size": {
"type": "integer",
"description": "Page size",
"default": 20,
},
},
},
),
types.Tool(
name="delete_session",
@@ -466,10 +580,9 @@ async def handle_list_tools() -> list[types.Tool]:
"properties": {
"session_id": {"type": "string", "description": "Session ID"}
},
"required": ["session_id"]
}
"required": ["session_id"],
},
),
# Chat Functionality
types.Tool(
name="chat",
@@ -478,12 +591,11 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"session_id": {"type": "string", "description": "Session ID"},
"query": {"type": "string", "description": "User query"}
"query": {"type": "string", "description": "User query"},
},
"required": ["session_id", "query"]
}
"required": ["session_id", "query"],
},
),
# Chunk Management
types.Tool(
name="list_chunks",
@@ -492,11 +604,19 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"knowledge_id": {"type": "string", "description": "Knowledge ID"},
"page": {"type": "integer", "description": "Page number", "default": 1},
"page_size": {"type": "integer", "description": "Page size", "default": 20}
"page": {
"type": "integer",
"description": "Page number",
"default": 1,
},
"page_size": {
"type": "integer",
"description": "Page size",
"default": 20,
},
},
"required": ["knowledge_id"]
}
"required": ["knowledge_id"],
},
),
types.Tool(
name="delete_chunk",
@@ -505,31 +625,32 @@ async def handle_list_tools() -> list[types.Tool]:
"type": "object",
"properties": {
"knowledge_id": {"type": "string", "description": "Knowledge ID"},
"chunk_id": {"type": "string", "description": "Chunk ID"}
"chunk_id": {"type": "string", "description": "Chunk ID"},
},
"required": ["knowledge_id", "chunk_id"]
}
)
"required": ["knowledge_id", "chunk_id"],
},
),
]
@app.call_tool()
async def handle_call_tool(
name: str, arguments: dict | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
"""Handle tool execution requests from MCP clients
Args:
name: Name of the tool to execute
arguments: Tool arguments as dictionary
Returns:
List of content items (text, image, or embedded resources)
"""
try:
# Use empty dict if no arguments provided
args = arguments or {}
# Tenant Management - Route tenant-related operations
if name == "create_tenant":
result = client.create_tenant(
@@ -537,33 +658,43 @@ async def handle_call_tool(
args["description"],
args["business"],
# Default to postgres-based keyword and vector search if not specified
args.get("retriever_engines", {
"engines": [
{"retriever_type": "keywords", "retriever_engine_type": "postgres"},
{"retriever_type": "vector", "retriever_engine_type": "postgres"}
]
})
args.get(
"retriever_engines",
{
"engines": [
{
"retriever_type": "keywords",
"retriever_engine_type": "postgres",
},
{
"retriever_type": "vector",
"retriever_engine_type": "postgres",
},
]
},
),
)
elif name == "list_tenants":
result = client.list_tenants()
# Knowledge Base Management - Route knowledge base operations
elif name == "create_knowledge_base":
# Build configuration with defaults for chunking and models
config = {
"chunking_config": args.get("chunking_config", {
"chunk_size": 1000, # Default chunk size in characters
"chunk_overlap": 200, # Default overlap between chunks
"separators": ["."], # Default text separators
"enable_multimodal": True # Enable image processing by default
}),
"chunking_config": args.get(
"chunking_config",
{
"chunk_size": 1000, # Default chunk size in characters
"chunk_overlap": 200, # Default overlap between chunks
"separators": ["."], # Default text separators
"enable_multimodal": True, # Enable image processing by default
},
),
"embedding_model_id": args.get("embedding_model_id", ""),
"summary_model_id": args.get("summary_model_id", "")
"summary_model_id": args.get("summary_model_id", ""),
}
result = client.create_knowledge_base(
args["name"],
args["description"],
config
args["name"], args["description"], config
)
elif name == "list_knowledge_bases":
result = client.list_knowledge_bases()
@@ -574,42 +705,42 @@ async def handle_call_tool(
elif name == "hybrid_search":
# Configure hybrid search with thresholds and result count
config = {
"vector_threshold": args.get("vector_threshold", 0.5), # Minimum similarity score
"keyword_threshold": args.get("keyword_threshold", 0.3), # Minimum keyword match score
"match_count": args.get("match_count", 5) # Number of results to return
"vector_threshold": args.get(
"vector_threshold", 0.5
), # Minimum similarity score
"keyword_threshold": args.get(
"keyword_threshold", 0.3
), # Minimum keyword match score
"match_count": args.get(
"match_count", 5
), # Number of results to return
}
result = client.hybrid_search(args["kb_id"], args["query"], config)
# Knowledge Management
elif name == "create_knowledge_from_file":
result = client.create_knowledge_from_file(
args["kb_id"],
args["file_path"],
args.get("enable_multimodel", True)
args["kb_id"], args["file_path"], args.get("enable_multimodel", True)
)
elif name == "create_knowledge_from_url":
result = client.create_knowledge_from_url(
args["kb_id"],
args["url"],
args.get("enable_multimodel", True)
args["kb_id"], args["url"], args.get("enable_multimodel", True)
)
elif name == "list_knowledge":
result = client.list_knowledge(
args["kb_id"],
args.get("page", 1),
args.get("page_size", 20)
args["kb_id"], args.get("page", 1), args.get("page_size", 20)
)
elif name == "get_knowledge":
result = client.get_knowledge(args["knowledge_id"])
elif name == "delete_knowledge":
result = client.delete_knowledge(args["knowledge_id"])
# Model Management - Route model configuration operations
elif name == "create_model":
# Build model parameters (API credentials, endpoints, etc.)
parameters = {
"base_url": args.get("base_url", ""), # Model API endpoint
"api_key": args.get("api_key", "") # Model API key
"api_key": args.get("api_key", ""), # Model API key
}
result = client.create_model(
args["name"],
@@ -617,71 +748,72 @@ async def handle_call_tool(
args.get("source", "local"),
args["description"],
parameters,
args.get("is_default", False)
args.get("is_default", False),
)
elif name == "list_models":
result = client.list_models()
elif name == "get_model":
result = client.get_model(args["model_id"])
# Session Management - Route chat session operations
elif name == "create_session":
# Build session strategy with conversation settings
strategy = {
"max_rounds": args.get("max_rounds", 5), # Maximum conversation turns
"enable_rewrite": args.get("enable_rewrite", True), # Enable query rewriting
"enable_rewrite": args.get(
"enable_rewrite", True
), # Enable query rewriting
"fallback_strategy": "FIXED_RESPONSE", # Strategy when no answer found
"fallback_response": args.get("fallback_response", "Sorry, I cannot answer this question."),
"fallback_response": args.get(
"fallback_response", "Sorry, I cannot answer this question."
),
"embedding_top_k": 10, # Number of chunks to retrieve
"keyword_threshold": 0.5, # Keyword match threshold
"vector_threshold": 0.7, # Vector similarity threshold
"summary_model_id": args.get("summary_model_id", "") # Model for summarization
"summary_model_id": args.get(
"summary_model_id", ""
), # Model for summarization
}
result = client.create_session(args["kb_id"], strategy)
elif name == "get_session":
result = client.get_session(args["session_id"])
elif name == "list_sessions":
result = client.list_sessions(
args.get("page", 1),
args.get("page_size", 20)
args.get("page", 1), args.get("page_size", 20)
)
elif name == "delete_session":
result = client.delete_session(args["session_id"])
# Chat Functionality
elif name == "chat":
result = client.chat(args["session_id"], args["query"])
# Chunk Management
elif name == "list_chunks":
result = client.list_chunks(
args["knowledge_id"],
args.get("page", 1),
args.get("page_size", 20)
args["knowledge_id"], args.get("page", 1), args.get("page_size", 20)
)
elif name == "delete_chunk":
result = client.delete_chunk(args["knowledge_id"], args["chunk_id"])
else:
# Handle unknown tool names
return [types.TextContent(
type="text",
text=f"Unknown tool: {name}"
)]
return [types.TextContent(type="text", text=f"Unknown tool: {name}")]
# Return successful result as formatted JSON
return [types.TextContent(
type="text",
text=json.dumps(result, indent=2, ensure_ascii=False)
)]
return [
types.TextContent(
type="text", text=json.dumps(result, indent=2, ensure_ascii=False)
)
]
except Exception as e:
# Log and return error message
logger.error(f"Tool execution failed: {e}")
return [types.TextContent(
type="text",
text=f"Error executing {name}: {str(e)}"
)]
return [
types.TextContent(type="text", text=f"Error executing {name}: {str(e)}")
]
async def run():
"""Run the MCP server using stdio transport"""
@@ -701,11 +833,14 @@ async def run():
),
)
def main():
"""Main entry point for console_scripts"""
import asyncio
# Run the async server
asyncio.run(run())
if __name__ == "__main__":
main()