From e9a242c25f9c540b6eae567ed59ec1be3bfd6a47 Mon Sep 17 00:00:00 2001 From: mileslai Date: Tue, 26 May 2026 14:48:41 +0800 Subject: [PATCH] feat(mcp-server): add multi-transport support (stdio / SSE / HTTP) --- .env.example | 6 + docker-compose.yml | 16 + mcp-server/Dockerfile | 16 + mcp-server/main.py | 30 +- mcp-server/requirements.txt | 6 +- mcp-server/weknora_mcp_server.py | 537 ++++++++++++++++++++++++++----- 6 files changed, 531 insertions(+), 80 deletions(-) create mode 100644 mcp-server/Dockerfile diff --git a/.env.example b/.env.example index b3adb532..f8578cf9 100644 --- a/.env.example +++ b/.env.example @@ -3,6 +3,12 @@ # 2. 替换所有占位符为实际值 # 3. 确保 .env 文件不会被提交到版本控制系统 +# ========== MCP Server ========== +# WeKnora API Key — 在 WeKnora 前端 "设置 → API Keys" 生成，供 MCP Server 访问 REST API +# WEKNORA_API_KEY= +# 对外暴露的 MCP Server 端口（默认 8082） +# MCP_PORT=8082 + # ========== 镜像版本 ========== # WeKnora 镜像版本标签，可选值: latest(稳定版), main(最新开发版) # WEKNORA_VERSION=latest diff --git a/docker-compose.yml b/docker-compose.yml index 1a734d14..1a78bd6a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -772,6 +772,22 @@ services: - langfuse - full + mcp: + build: + context: ./mcp-server + container_name: WeKnora-mcp + ports: + - "${MCP_PORT:-8082}:8000" + environment: + - WEKNORA_BASE_URL=http://app:8080/api/v1 + - WEKNORA_API_KEY=${WEKNORA_API_KEY:-} + depends_on: + app: + condition: service_healthy + networks: + - WeKnora-network + restart: unless-stopped + networks: WeKnora-network: driver: bridge diff --git a/mcp-server/Dockerfile b/mcp-server/Dockerfile new file mode 100644 index 00000000..5b76a614 --- /dev/null +++ b/mcp-server/Dockerfile @@ -0,0 +1,16 @@ +FROM python:3.11-slim + +WORKDIR /app +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY . . +RUN pip install --no-cache-dir -e . + +ENV MCP_HOST=0.0.0.0 +ENV MCP_PORT=8000 +ENV WEKNORA_BASE_URL=http://app:8080/api/v1 + +EXPOSE 8000 + +CMD ["weknora-mcp-server", "--transport", "http", "--host", "0.0.0.0", "--port", "8000"] diff --git a/mcp-server/main.py b/mcp-server/main.py index 8d03d846..ffa2de8d 100644 --- a/mcp-server/main.py +++ b/mcp-server/main.py @@ -83,6 +83,24 @@ def parse_arguments(): "--version", action="version", version="WeKnora MCP Server 1.0.0" ) + parser.add_argument( + "--transport", + choices=["stdio", "sse", "http"], + default=os.getenv("MCP_TRANSPORT", "stdio"), + help="Transport type: stdio (default), sse, or http", + ) + parser.add_argument( + "--host", + default=os.getenv("MCP_HOST", "0.0.0.0"), + help="Bind host for network transports (default: 0.0.0.0)", + ) + parser.add_argument( + "--port", + type=int, + default=int(os.getenv("MCP_PORT", "8000")), + help="Bind port for network transports (default: 8000)", + ) + return parser.parse_args() @@ -113,12 +131,16 @@ async def main(): print("已启用详细日志模式") try: - print("正在启动 WeKnora MCP Server...") + print(f"正在启动 WeKnora MCP Server (transport={args.transport})...") - # 导入并运行服务器 - from weknora_mcp_server import run + from weknora_mcp_server import run_stdio, run_sse, run_http - await run() + if args.transport == "stdio": + await run_stdio() + elif args.transport == "sse": + await run_sse(args.host, args.port) + elif args.transport == "http": + await run_http(args.host, args.port) except ImportError as e: print(f"导入错误: {e}") diff --git a/mcp-server/requirements.txt b/mcp-server/requirements.txt index 28264a17..23b36714 100644 --- a/mcp-server/requirements.txt +++ b/mcp-server/requirements.txt @@ -1,2 +1,4 @@ -mcp>=1.0.0 -requests>=2.31.0 \ No newline at end of file +mcp>=1.0.0 +requests>=2.31.0 +starlette>=0.27.0 +uvicorn>=0.24.0 \ No newline at end of file diff --git a/mcp-server/weknora_mcp_server.py b/mcp-server/weknora_mcp_server.py index afc99d79..6feb2cb2 100644 --- a/mcp-server/weknora_mcp_server.py +++ b/mcp-server/weknora_mcp_server.py @@ -5,11 +5,16 @@ WeKnora MCP Server A Model Context Protocol server that provides access to the WeKnora knowledge management API. """ +import argparse +import asyncio +import functools import json import logging import os +import re from typing import Any, Dict +import urllib3 import mcp.server.stdio import mcp.types as types import requests @@ -24,6 +29,8 @@ logger = logging.getLogger(__name__) # Configuration - Load from environment variables with defaults WEKNORA_BASE_URL = os.getenv("WEKNORA_BASE_URL", "http://localhost:8080/api/v1") WEKNORA_API_KEY = os.getenv("WEKNORA_API_KEY", "") +# Chat SSE read timeout in seconds. LLM responses can be slow; default 300s. +WEKNORA_CHAT_TIMEOUT = int(os.getenv("WEKNORA_CHAT_TIMEOUT", "300")) class WeKnoraClient: @@ -33,8 +40,13 @@ class WeKnoraClient: """Initialize the WeKnora API client with base URL and authentication""" self.base_url = base_url self.api_key = api_key + # SSL verification: set WEKNORA_VERIFY_SSL=true to enable (default off for self-signed certs) + self.verify_ssl = os.getenv("WEKNORA_VERIFY_SSL", "false").lower() == "true" + if not self.verify_ssl: + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) # Create a persistent session for connection pooling and performance self.session = requests.Session() + self.session.verify = self.verify_ssl # Set default headers for all requests self.session.headers.update( { @@ -113,6 +125,58 @@ class WeKnoraClient: """Delete knowledge base""" return self._request("DELETE", f"/knowledge-bases/{kb_id}") + # ── UUID pattern (8-4-4-4-12 hex) ────────────────────────────────────── + _UUID_RE = re.compile( + r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", + re.IGNORECASE, + ) + + def resolve_agent_id(self, agent_id_or_name: str) -> str: + """Resolve an agent name to its UUID if needed. + + If *agent_id_or_name* is already a UUID it is returned unchanged. + Otherwise all agents are listed and the first one whose + ``name`` matches (case-insensitive) is returned. + Raises ValueError when no match is found. + """ + if self._UUID_RE.match(agent_id_or_name): + return agent_id_or_name + resp = self._request("GET", "/agents") + agents = resp.get("data", resp) if isinstance(resp, dict) else resp + if isinstance(agents, dict): + agents = agents.get("list", agents.get("items", [])) + needle = agent_id_or_name.lower() + for agent in (agents or []): + if isinstance(agent, dict) and agent.get("name", "").lower() == needle: + return agent["id"] + raise ValueError( + f"Agent {agent_id_or_name!r} not found. " + "Use list_agents to see available agent IDs and names." + ) + + def resolve_kb_id(self, kb_id_or_name: str) -> str: + """Resolve a knowledge base name to its UUID if needed. + + If *kb_id_or_name* is already a UUID it is returned unchanged. + Otherwise all knowledge bases are listed and the first one whose + ``name`` matches (case-insensitive) is returned. + Raises ValueError when no match is found. + """ + if self._UUID_RE.match(kb_id_or_name): + return kb_id_or_name + resp = self.list_knowledge_bases() + kbs = resp.get("data", resp) if isinstance(resp, dict) else resp + if isinstance(kbs, dict): + kbs = kbs.get("list", kbs.get("items", [])) + needle = kb_id_or_name.lower() + for kb in (kbs or []): + if isinstance(kb, dict) and kb.get("name", "").lower() == needle: + return kb["id"] + raise ValueError( + f"Knowledge base {kb_id_or_name!r} not found. " + "Use list_knowledge_bases to see available IDs and names." + ) + def hybrid_search(self, kb_id: str, query: str, config: Dict) -> Dict: """Perform hybrid search combining vector and keyword search""" data = { @@ -202,11 +266,11 @@ class WeKnoraClient: 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""" + def create_session(self, title: str = "", description: str = "") -> Dict: + """Create a new chat session (sessions are knowledge-base-independent)""" data = { - "knowledge_base_id": kb_id, # Knowledge base to query - "session_strategy": strategy, # Conversation settings (max rounds, rewrite, etc.) + "title": title, + "description": description, } return self._request("POST", "/sessions", json=data) @@ -224,12 +288,122 @@ class WeKnoraClient: 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""" - data = {"query": query} - # 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) + def _consume_sse_stream(self, url: str, body: Dict[str, Any]) -> Dict: + """POST to *url* with *body*, consume the SSE stream, and return the assembled result. + + Centralised helper used by both chat() and agent_chat(). + Timeout: (10s connect, WEKNORA_CHAT_TIMEOUT read) — configurable via env var. + """ + try: + response = self.session.post( + url, json=body, stream=True, + timeout=(10, WEKNORA_CHAT_TIMEOUT), + ) + response.raise_for_status() + + answer_chunks: list = [] + references: list = [] + debug_events: list = [] + + for raw_line in response.iter_lines(): + if not raw_line: + continue + if isinstance(raw_line, bytes): + raw_line = raw_line.decode("utf-8") + if not raw_line.startswith("data:"): + continue + payload = raw_line[5:].lstrip(" ") + try: + event_data = json.loads(payload) + except json.JSONDecodeError: + continue + + response_type = event_data.get("response_type", "") + debug_events.append({"type": response_type, "content": event_data.get("content", "")[:80]}) + + if response_type == "answer": + chunk = event_data.get("content", "") + if chunk: + answer_chunks.append(chunk) + elif response_type == "references": + references = event_data.get("knowledge_references") or [] + elif response_type == "error": + raise RequestException( + f"Server error: {event_data.get('content', 'unknown error')}" + ) + elif response_type == "complete": + break + + return { + "answer": "".join(answer_chunks), + "references": references, + "_debug_events": debug_events, + } + except RequestException as e: + logger.error(f"SSE stream request failed ({url}): {e}") + raise + + def chat( + self, + session_id: str, + query: str, + knowledge_base_ids: list = None, + web_search_enabled: bool = False, + enable_memory: bool = False, + ) -> Dict: + """Send a message to the RAG pipeline (knowledge-chat) and return the assembled answer. + + Provide *knowledge_base_ids* (UUID or name) so the backend can retrieve + relevant chunks before summarising with the LLM. + For agentic tool-calling use agent_chat() instead. + """ + url = f"{self.base_url}/knowledge-chat/{session_id}" + body: Dict[str, Any] = {"query": query, "channel": "api"} + if knowledge_base_ids: + body["knowledge_base_ids"] = knowledge_base_ids + if web_search_enabled: + body["web_search_enabled"] = True + if enable_memory: + body["enable_memory"] = True + result = self._consume_sse_stream(url, body) + result["session_id"] = session_id + return result + + def agent_chat( + self, + session_id: str, + query: str, + agent_id: str, + knowledge_base_ids: list = None, + web_search_enabled: bool = False, + enable_memory: bool = False, + ) -> Dict: + """Send a message to the agentic pipeline (agent-chat) and return the assembled answer. + + *agent_id* is required — the backend uses the CustomAgent config for + tool selection (knowledge_search, web_search, SQL, etc.). + The agent autonomously decides which knowledge bases to query; + pass *knowledge_base_ids* to override or supplement the agent's default KBs. + """ + url = f"{self.base_url}/agent-chat/{session_id}" + body: Dict[str, Any] = {"query": query, "agent_id": agent_id, "channel": "api"} + if knowledge_base_ids: + body["knowledge_base_ids"] = knowledge_base_ids + if web_search_enabled: + body["web_search_enabled"] = True + if enable_memory: + body["enable_memory"] = True + result = self._consume_sse_stream(url, body) + result["session_id"] = session_id + return result + + def list_agents(self, page: int = 1, page_size: int = 50) -> Dict: + """List all custom agents available to the current tenant.""" + return self._request("GET", "/agents", params={"page": page, "page_size": page_size}) + + def get_agent(self, agent_id: str) -> Dict: + """Get full config of a single agent by UUID.""" + return self._request("GET", f"/agents/{agent_id}") # Chunk Management - Methods for managing knowledge chunks (text segments) def list_chunks( @@ -371,7 +545,10 @@ async def handle_list_tools() -> list[types.Tool]: inputSchema={ "type": "object", "properties": { - "kb_id": {"type": "string", "description": "Knowledge base ID"}, + "kb_id": { + "type": "string", + "description": "Knowledge base UUID (e.g. 'a1b2c3d4-e5f6-7890-abcd-ef1234567890') OR name (e.g. 'my-knowledge-base'). Use list_knowledge_bases to discover available knowledge bases.", + }, "query": {"type": "string", "description": "Search query"}, "vector_threshold": { "type": "number", @@ -535,32 +712,14 @@ async def handle_list_tools() -> list[types.Tool]: # Session Management types.Tool( name="create_session", - description="Create a new chat session", + description="Create a new chat session. Sessions are knowledge-base-independent; pass knowledge_base_ids when calling chat.", inputSchema={ "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", - }, + "title": {"type": "string", "description": "Session title (optional)"}, + "description": {"type": "string", "description": "Session description (optional)"}, }, - "required": ["kb_id"], + "required": [], }, ), types.Tool( @@ -607,16 +766,89 @@ async def handle_list_tools() -> list[types.Tool]: # Chat Functionality types.Tool( name="chat", - description="Send a chat message to a session", + description=( + "RAG pipeline chat: retrieve relevant chunks from knowledge bases, then summarise with LLM. " + "ALWAYS provide knowledge_base_ids (names like 'my-knowledge-base' or UUIDs) so retrieval can run — " + "without them the answer is based on LLM knowledge only. " + "Use list_knowledge_bases to discover available knowledge bases. " + "For multi-step reasoning or tool-calling use agent_chat instead." + ), inputSchema={ "type": "object", "properties": { - "session_id": {"type": "string", "description": "Session ID"}, + "session_id": {"type": "string", "description": "Session ID (from create_session or list_sessions)"}, "query": {"type": "string", "description": "User query"}, + "knowledge_base_ids": { + "type": "array", + "items": {"type": "string"}, + "description": "Knowledge base names OR UUIDs to search. REQUIRED for RAG. E.g. ['my-knowledge-base'] or ['a1b2c3d4-...']. Use list_knowledge_bases to find them.", + }, + "web_search_enabled": {"type": "boolean", "description": "Enable web search alongside KB retrieval.", "default": False}, + "enable_memory": {"type": "boolean", "description": "Enable cross-session memory.", "default": False}, }, "required": ["session_id", "query"], }, ), + types.Tool( + name="agent_chat", + description=( + "Agentic pipeline chat: the agent autonomously calls tools (knowledge_search, web_search, SQL, etc.) " + "to answer the query. Use this for complex multi-step questions or comparative analysis. " + "REQUIRED: agent_id (name or UUID) — use list_agents to discover agents. " + "IMPORTANT: many agents have KBSelectionMode=none and NO built-in knowledge bases. " + "In that case you MUST pass knowledge_base_ids, otherwise the agent will fail with " + "'no search targets available'. " + "Use get_agent to inspect an agent's kb_selection_mode and knowledge_bases before calling. " + "If kb_selection_mode is 'none' or 'selected' with an empty list, always provide knowledge_base_ids." + ), + inputSchema={ + "type": "object", + "properties": { + "session_id": {"type": "string", "description": "Session ID (from create_session or list_sessions)"}, + "query": {"type": "string", "description": "User query"}, + "agent_id": { + "type": "string", + "description": "REQUIRED. Custom agent UUID or name. Use list_agents to discover agents. Use get_agent to check its kb_selection_mode.", + }, + "knowledge_base_ids": { + "type": "array", + "items": {"type": "string"}, + "description": "Names or UUIDs of knowledge bases to search. REQUIRED when the agent's kb_selection_mode is 'none' or 'selected' with no built-in KBs. Use list_knowledge_bases to find them.", + }, + "web_search_enabled": {"type": "boolean", "description": "Enable web search.", "default": False}, + "enable_memory": {"type": "boolean", "description": "Enable cross-session memory.", "default": False}, + }, + "required": ["session_id", "query", "agent_id"], + }, + ), + types.Tool( + name="list_agents", + description="List all custom agents available to the current tenant. Use this to discover agent IDs, names, and their KB selection mode before calling agent_chat.", + inputSchema={ + "type": "object", + "properties": { + "page": {"type": "integer", "description": "Page number", "default": 1}, + "page_size": {"type": "integer", "description": "Page size", "default": 50}, + }, + "required": [], + }, + ), + types.Tool( + name="get_agent", + description=( + "Get full configuration of a single agent by UUID or name. " + "Check kb_selection_mode and knowledge_bases fields: " + "if kb_selection_mode is 'none' or 'selected' with an empty knowledge_bases list, " + "you MUST pass knowledge_base_ids when calling agent_chat." + ), + inputSchema={ + "type": "object", + "properties": { + "agent_id": {"type": "string", "description": "Agent UUID or name"}, + }, + "required": ["agent_id"], + }, + ), # Chunk Management types.Tool( name="list_chunks", @@ -785,7 +1017,8 @@ async def handle_call_tool( "match_count", 5 ), # Number of results to return } - result = client.hybrid_search(args["kb_id"], args["query"], config) + kb_id = client.resolve_kb_id(args["kb_id"]) + result = client.hybrid_search(kb_id, args["query"], config) # Knowledge Management elif name == "create_knowledge_from_file": @@ -827,24 +1060,10 @@ async def handle_call_tool( # 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 - "fallback_strategy": "FIXED_RESPONSE", # Strategy when no answer found - "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 - } - result = client.create_session(args["kb_id"], strategy) + result = client.create_session( + title=args.get("title", ""), + description=args.get("description", ""), + ) elif name == "get_session": result = client.get_session(args["session_id"]) elif name == "list_sessions": @@ -856,7 +1075,74 @@ async def handle_call_tool( # Chat Functionality elif name == "chat": - result = client.chat(args["session_id"], args["query"]) + # Resolve KB names → UUIDs + raw_kb_ids = args.get("knowledge_base_ids") or [] + kb_ids = [client.resolve_kb_id(k) for k in raw_kb_ids] if raw_kb_ids else None + fn = functools.partial( + client.chat, + args["session_id"], + args["query"], + knowledge_base_ids=kb_ids, + web_search_enabled=args.get("web_search_enabled", False), + enable_memory=args.get("enable_memory", False), + ) + result = await asyncio.get_event_loop().run_in_executor(None, fn) + + elif name == "agent_chat": + # Resolve required agent name → UUID + agent_id = client.resolve_agent_id(args["agent_id"]) + # Resolve optional KB overrides + raw_kb_ids = args.get("knowledge_base_ids") or [] + kb_ids = [client.resolve_kb_id(k) for k in raw_kb_ids] if raw_kb_ids else None + # Pre-check: if no KB IDs provided, inspect agent config to detect + # kb_selection_mode=none/selected-empty so we fail fast with a clear message + # instead of the cryptic backend error "no search targets available". + if not kb_ids: + try: + agent_info = client.get_agent(agent_id) + cfg = (agent_info.get("data") or agent_info).get("config") or {} + mode = cfg.get("kb_selection_mode", "selected") + built_in_kbs = cfg.get("knowledge_bases") or [] + needs_kbs = (mode == "none") or (mode in ("selected", "") and not built_in_kbs) + if needs_kbs: + kb_list = client.list_knowledge_bases() + kbs = (kb_list.get("data") or kb_list) + if isinstance(kbs, dict): + kbs = kbs.get("list", kbs.get("items", [])) + kb_summary = ", ".join( + f"{kb.get('name')} ({kb.get('id')})" + for kb in (kbs or [])[:10] + if isinstance(kb, dict) + ) + raise ValueError( + f"Agent '{args['agent_id']}' has kb_selection_mode='{mode}' with no built-in " + f"knowledge bases. You must provide knowledge_base_ids. " + f"Available knowledge bases: [{kb_summary}]" + ) + except ValueError: + raise + except Exception as preflight_err: + logger.warning(f"agent_chat preflight KB check failed (non-fatal): {preflight_err}") + fn = functools.partial( + client.agent_chat, + args["session_id"], + args["query"], + agent_id, + knowledge_base_ids=kb_ids, + web_search_enabled=args.get("web_search_enabled", False), + enable_memory=args.get("enable_memory", False), + ) + result = await asyncio.get_event_loop().run_in_executor(None, fn) + + elif name == "list_agents": + result = client.list_agents( + page=args.get("page", 1), + page_size=args.get("page_size", 50), + ) + + elif name == "get_agent": + resolved_id = client.resolve_agent_id(args["agent_id"]) + result = client.get_agent(resolved_id) # Chunk Management elif name == "list_chunks": @@ -897,31 +1183,134 @@ async def handle_call_tool( ] -async def run(): +def _init_options() -> InitializationOptions: + """Build MCP InitializationOptions (shared across all transports)""" + return InitializationOptions( + server_name="weknora-server", + server_version="1.0.0", + capabilities=app.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ) + + +async def run_stdio(): """Run the MCP server using stdio transport""" - # Create stdio streams for communication with MCP client async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): - # Run the server with initialization options - await app.run( - read_stream, - write_stream, - InitializationOptions( - server_name="weknora-server", - server_version="1.0.0", - capabilities=app.get_capabilities( - notification_options=NotificationOptions(), - experimental_capabilities={}, - ), - ), - ) + await app.run(read_stream, write_stream, _init_options()) + + +async def run_sse(host: str, port: int): + """Run the MCP server using SSE transport (legacy MCP clients)""" + try: + from mcp.server.sse import SseServerTransport + from starlette.applications import Starlette + from starlette.requests import Request + from starlette.routing import Mount, Route + import uvicorn + except ImportError as e: + raise ImportError( + f"SSE transport requires 'starlette' and 'uvicorn': pip install starlette uvicorn\n{e}" + ) from e + + sse = SseServerTransport("/messages/") + + async def handle_sse(request: Request): + async with sse.connect_sse( + request.scope, request.receive, request._send + ) as streams: + await app.run(streams[0], streams[1], _init_options()) + + starlette_app = Starlette( + routes=[ + Route("/sse", endpoint=handle_sse), + Mount("/messages/", app=sse.handle_post_message), + ] + ) + + logger.info("Starting SSE MCP server on %s:%d", host, port) + logger.info("SSE endpoint: http://%s:%d/sse", host, port) + config = uvicorn.Config(starlette_app, host=host, port=port, log_level="info") + server = uvicorn.Server(config) + await server.serve() + + +async def run_http(host: str, port: int): + """Run the MCP server using Streamable HTTP transport (MCP 2025-03-26 spec)""" + try: + from contextlib import asynccontextmanager + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + from starlette.applications import Starlette + from starlette.routing import Mount + import uvicorn + except ImportError as e: + raise ImportError( + f"HTTP transport requires 'starlette' and 'uvicorn': pip install starlette uvicorn\n{e}" + ) from e + + session_manager = StreamableHTTPSessionManager( + app=app, + event_store=None, + json_response=False, + stateless=True, + ) + + @asynccontextmanager + async def lifespan(_app): + async with session_manager.run(): + yield + + starlette_app = Starlette( + routes=[Mount("/", app=session_manager.handle_request)], + lifespan=lifespan, + ) + + logger.info("Starting Streamable HTTP MCP server on %s:%d", host, port) + logger.info("MCP endpoint: http://%s:%d/mcp", host, port) + config = uvicorn.Config(starlette_app, host=host, port=port, log_level="info") + server = uvicorn.Server(config) + await server.serve() + + +# Backward-compatible alias used by run_server.py +run = run_stdio def main(): - """Main entry point for console_scripts""" - import asyncio + """Main entry point — supports stdio, sse, and http transports. - # Run the async server - asyncio.run(run()) + Transport selection (in priority order): + 1. --transport CLI flag + 2. MCP_TRANSPORT environment variable + 3. Default: stdio + """ + parser = argparse.ArgumentParser(description="WeKnora MCP Server") + parser.add_argument( + "--transport", + choices=["stdio", "sse", "http"], + default=os.getenv("MCP_TRANSPORT", "stdio"), + help="Transport type: stdio (default), sse, or http", + ) + parser.add_argument( + "--host", + default=os.getenv("MCP_HOST", "0.0.0.0"), + help="Bind host for network transports (default: 0.0.0.0)", + ) + parser.add_argument( + "--port", + type=int, + default=int(os.getenv("MCP_PORT", "8000")), + help="Bind port for network transports (default: 8000)", + ) + args = parser.parse_args() + + if args.transport == "stdio": + asyncio.run(run_stdio()) + elif args.transport == "sse": + asyncio.run(run_sse(args.host, args.port)) + elif args.transport == "http": + asyncio.run(run_http(args.host, args.port)) if __name__ == "__main__":