From f5963276d6f01ff41a5cf5d86e05d514fcfd5136 Mon Sep 17 00:00:00 2001 From: begoniezhao Date: Mon, 1 Dec 2025 18:15:39 +0800 Subject: [PATCH] Refactor code for improved readability and consistency --- mcp-server/__init__.py | 2 +- mcp-server/main.py | 65 ++-- mcp-server/run.py | 13 +- mcp-server/run_server.py | 20 +- mcp-server/setup.py | 10 +- mcp-server/test_imports.py | 10 +- mcp-server/test_module.py | 96 +++--- mcp-server/weknora_mcp_server.py | 537 +++++++++++++++++++------------ 8 files changed, 461 insertions(+), 292 deletions(-) diff --git a/mcp-server/__init__.py b/mcp-server/__init__.py index 987ea629e..3a574e121 100644 --- a/mcp-server/__init__.py +++ b/mcp-server/__init__.py @@ -11,4 +11,4 @@ __description__ = "WeKnora MCP Server - Model Context Protocol server for WeKnor from .weknora_mcp_server import WeKnoraClient, run -__all__ = ["WeKnoraClient", "run"] \ No newline at end of file +__all__ = ["WeKnoraClient", "run"] diff --git a/mcp-server/main.py b/mcp-server/main.py index 6203a276c..8d03d8460 100644 --- a/mcp-server/main.py +++ b/mcp-server/main.py @@ -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()) \ No newline at end of file + asyncio.run(main()) diff --git a/mcp-server/run.py b/mcp-server/run.py index a20dd351e..02c6f6229 100644 --- a/mcp-server/run.py +++ b/mcp-server/run.py @@ -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() \ No newline at end of file + main() diff --git a/mcp-server/run_server.py b/mcp-server/run_server.py index 1956e292f..fc1226556 100644 --- a/mcp-server/run_server.py +++ b/mcp-server/run_server.py @@ -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() \ No newline at end of file + main() diff --git a/mcp-server/setup.py b/mcp-server/setup.py index 95e2bfb4b..008d7c53d 100644 --- a/mcp-server/setup.py +++ b/mcp-server/setup.py @@ -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", -) \ No newline at end of file +) diff --git a/mcp-server/test_imports.py b/mcp-server/test_imports.py index 0317a88d0..a42975b11 100644 --- a/mcp-server/test_imports.py +++ b/mcp-server/test_imports.py @@ -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)}") \ No newline at end of file +print(f"MCP 包内容: {dir(mcp)}") diff --git a/mcp-server/test_module.py b/mcp-server/test_module.py index 3912e87fc..04d7026df 100644 --- a/mcp-server/test_module.py +++ b/mcp-server/test_module.py @@ -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) \ No newline at end of file + sys.exit(0 if success else 1) diff --git a/mcp-server/weknora_mcp_server.py b/mcp-server/weknora_mcp_server.py index c8ca42698..1f7ac94da 100644 --- a/mcp-server/weknora_mcp_server.py +++ b/mcp-server/weknora_mcp_server.py @@ -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()