"""Memory Gateway MCP Server. 通用 Memory Gateway 服务，为 AI agent / harness 提供统一的 OpenViking 记忆检索、总结和知识沉淀入口。 """ import asyncio import hashlib import json import logging import re import tempfile from datetime import datetime, timezone from contextlib import asynccontextmanager from pathlib import Path from typing import Any, Optional from fastapi import APIRouter, Depends, FastAPI, File, Form, Header, HTTPException, Request, UploadFile, status from fastapi.responses import JSONResponse from fastapi.middleware.cors import CORSMiddleware from mcp.server import Server from mcp.types import TextContent, Tool from sse_starlette import EventSourceResponse from .config import get_config, set_config, Config from .openviking_client import get_openviking_client, close_openviking_client from .document_ingest import convert_file_to_markdown, save_markdown_to_obsidian, slugify from .llm import LLMConfigurationError, LLMSummaryError, summarize_with_llm from .types import SearchRequest, AddMemoryRequest, AddResourceRequest, CommitSummaryRequest # 配置日志 logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) # 创建 MCP Server mcp_server = Server("memory-gateway") @mcp_server.list_tools() async def list_tools() -> list[Tool]: """列出可用的 MCP 工具""" return [ Tool( name="search", description="语义搜索记忆和资源", inputSchema={ "type": "object", "properties": { "query": {"type": "string", "description": "搜索查询"}, "namespace": {"type": "string", "description": "命名空间（可选）"}, "limit": {"type": "integer", "description": "返回结果数量（默认10）"}, "uri": {"type": "string", "description": "资源 URI（可选）"}, }, "required": ["query"], }, ), Tool( name="add_memory", description="添加新记忆", inputSchema={ "type": "object", "properties": { "content": {"type": "string", "description": "记忆内容"}, "namespace": {"type": "string", "description": "命名空间（可选）"}, "memory_type": {"type": "string", "description": "记忆类型（默认general）"}, }, "required": ["content"], }, ), Tool( name="add_resource", description="添加资源", inputSchema={ "type": "object", "properties": { "uri": {"type": "string", "description": "资源 URI"}, "content": {"type": "string", "description": "资源内容"}, "resource_type": {"type": "string", "description": "资源类型（默认text）"}, }, "required": ["uri", "content"], }, ), Tool( name="commit_summary", description="总结一段通用内容并按需沉淀为 OpenViking memory/resource", inputSchema={ "type": "object", "properties": { "content": {"type": "string", "description": "需要总结和沉淀的原文内容"}, "title": {"type": "string", "description": "标题（可选）"}, "summary": {"type": "string", "description": "人工提供的摘要（可选）"}, "namespace": {"type": "string", "description": "OpenViking memory namespace（可选）"}, "memory_type": {"type": "string", "description": "记忆类型，默认 summary"}, "tags": {"type": "array", "items": {"type": "string"}, "description": "标签列表"}, "source": {"type": "string", "description": "来源说明或外部链接"}, "resource_uri": {"type": "string", "description": "写入 resource 的 URI（可选）"}, "resource_type": {"type": "string", "description": "资源类型，默认 json"}, "persist_as": {"type": "string", "enum": ["memory", "resource", "both", "none"], "description": "沉淀方式"}, "max_summary_chars": {"type": "integer", "description": "摘要最大长度"}, }, "required": ["content"], }, ), Tool( name="get_status", description="检查系统状态", inputSchema={ "type": "object", "properties": {}, }, ), Tool( name="list_memories", description="列出已存储的记忆", inputSchema={ "type": "object", "properties": { "namespace": {"type": "string", "description": "命名空间（可选）"}, "memory_type": {"type": "string", "description": "记忆类型（可选）"}, "limit": {"type": "integer", "description": "返回数量（默认10）"}, }, }, ), Tool( name="list_resources", description="列出已存储的资源", inputSchema={ "type": "object", "properties": { "namespace": {"type": "string", "description": "命名空间（可选）"}, "limit": {"type": "integer", "description": "返回数量（默认10）"}, }, }, ), ] @mcp_server.call_tool() async def call_tool(name: str, arguments: Any) -> list[TextContent]: """调用 MCP 工具""" try: ov_client = await get_openviking_client() if name == "search": result = await ov_client.search( query=arguments.get("query"), namespace=arguments.get("namespace"), limit=arguments.get("limit"), uri=arguments.get("uri"), ) return [TextContent(type="text", text=str(result.results))] elif name == "add_memory": result = await ov_client.add_memory( content=arguments.get("content"), namespace=arguments.get("namespace"), memory_type=arguments.get("memory_type", "general"), ) return [TextContent(type="text", text=str(result))] elif name == "add_resource": result = await ov_client.add_resource( uri=arguments.get("uri"), content=arguments.get("content"), resource_type=arguments.get("resource_type", "text"), ) return [TextContent(type="text", text=str(result))] elif name == "commit_summary": request = CommitSummaryRequest(**arguments) result = await commit_summary_to_openviking(request) return [TextContent(type="text", text=json.dumps(result, ensure_ascii=False))] elif name == "get_status": ov_status = await ov_client.health_check() return [TextContent(type="text", text=f"Memory Gateway: OK\nOpenViking: {ov_status}")] elif name == "list_memories": memories = await ov_client.list_memories( namespace=arguments.get("namespace"), memory_type=arguments.get("memory_type"), limit=arguments.get("limit"), ) return [TextContent(type="text", text=str([m.model_dump() for m in memories]))] elif name == "list_resources": resources = await ov_client.list_resources( namespace=arguments.get("namespace"), limit=arguments.get("limit"), ) return [TextContent(type="text", text=str([r.model_dump() for r in resources]))] else: raise ValueError(f"Unknown tool: {name}") except Exception as e: logger.error(f"工具执行失败: {e}") return [TextContent(type="text", text=f"Error: {str(e)}")] @asynccontextmanager async def lifespan(app: FastAPI): """应用生命周期管理""" logger.info("Memory Gateway 启动中...") config = get_config() logger.info(f"配置加载完成: {config.server.host}:{config.server.port}") logger.info(f"OpenViking 后端: {config.openviking.url}") # 测试 OpenViking 连接 try: ov_client = await get_openviking_client() status = await ov_client.health_check() logger.info(f"OpenViking 连接状态: {status}") except Exception as e: logger.warning(f"OpenViking 连接失败: {e}") yield logger.info("Memory Gateway 关闭中...") await close_openviking_client() def verify_api_key(x_api_key: Optional[str] = Header(default=None)) -> None: """在配置了 API Key 时校验请求头。""" expected_key = get_config().server.api_key if not expected_key: return if x_api_key != expected_key: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or missing API key", ) _SENTENCE_RE = re.compile(r"(?<=[。！？.!?])\s+") _WORD_RE = re.compile(r"[^a-zA-Z0-9\u4e00-\u9fff_-]+") def _normalize_whitespace(value: str) -> str: return re.sub(r"\s+", " ", value).strip() def _slugify(value: str, fallback: str) -> str: slug = _WORD_RE.sub("-", value.lower()).strip("-") slug = re.sub(r"-+", "-", slug)[:80].strip("-") return slug or fallback def _derive_title(content: str, title: Optional[str]) -> str: if title and title.strip(): return title.strip() for line in content.splitlines(): line = line.strip("# -*\t") if line: return line[:120] return "Untitled summary" def _derive_summary(content: str, provided: Optional[str], max_chars: int) -> str: if provided and provided.strip(): return provided.strip()[:max_chars] normalized = _normalize_whitespace(content) if not normalized: return "" sentences = [part.strip() for part in _SENTENCE_RE.split(normalized) if part.strip()] if not sentences: return normalized[:max_chars] summary = " ".join(sentences[:3]) return summary[:max_chars] def _extract_key_points(content: str, limit: int = 8) -> list[str]: points: list[str] = [] for raw_line in content.splitlines(): line = raw_line.strip() if not line: continue stripped = re.sub(r"^(?:[-*•]\s*|\d+[.、)]\s*)", "", line).strip() if not stripped: continue is_structured = line.startswith(("-", "*", "•")) or re.match(r"^\d+[.、)]\s+", line) has_signal = any(token in stripped.lower() for token in [ "verdict", "result", "finding", "evidence", "action", "risk", "ioc", "结论", "结果", "证据", "建议", "动作", "风险", "命中", "关联", ]) if is_structured or has_signal: point = _normalize_whitespace(stripped) if point and point not in points: points.append(point[:240]) if len(points) >= limit: break if points: return points summary = _derive_summary(content, None, 500) return [summary] if summary else [] def _render_memory_text(artifact: dict[str, Any]) -> str: lines = [ f"Title: {artifact['title']}", f"Summary: {artifact['summary']}", ] if artifact.get("tags"): lines.append("Tags: " + ", ".join(artifact["tags"])) if artifact.get("source"): lines.append("Source: " + artifact["source"]) if artifact.get("key_points"): lines.append("Key points:") lines.extend(f"- {point}" for point in artifact["key_points"]) return "\n".join(lines) def _default_summary_resource_uri(request: CommitSummaryRequest, title: str) -> str: namespace = (request.namespace or get_config().memory.default_namespace or "general").strip("/") memory_type = (request.memory_type or "summary").strip("/") digest = hashlib.sha1(request.content.encode("utf-8")).hexdigest()[:12] slug = _slugify(title, digest) return f"viking://resources/{namespace}/{memory_type}/{slug}-{digest}.json" async def build_summary_artifact(request: CommitSummaryRequest) -> dict[str, Any]: max_chars = max(120, min(request.max_summary_chars, 4000)) llm_result = await summarize_with_llm( request.content, title=request.title, summary_hint=request.summary, tags=request.tags, max_summary_chars=max_chars, purpose=request.purpose or "generic knowledge memory", ) title = llm_result.get("title") or _derive_title(request.content, request.title) return { "schema_version": "memory-gateway.summary.v1", "id": hashlib.sha1(request.content.encode("utf-8")).hexdigest()[:16], "title": title, "summary": llm_result.get("summary", ""), "key_points": llm_result.get("key_points", []), "tags": llm_result.get("tags", request.tags), "source": request.source, "namespace": request.namespace or get_config().memory.default_namespace, "memory_type": request.memory_type or "summary", "created_at": datetime.now(timezone.utc).isoformat(), "content": request.content, "llm": llm_result.get("llm"), } async def commit_summary_to_openviking(request: CommitSummaryRequest) -> dict[str, Any]: artifact = await build_summary_artifact(request) ov_client = await get_openviking_client() memory_result: Optional[dict[str, Any]] = None resource_result: Optional[dict[str, Any]] = None if request.persist_as in {"memory", "both"}: memory_result = await ov_client.add_memory( content=_render_memory_text(artifact), namespace=artifact["namespace"], memory_type=artifact["memory_type"], ) if request.persist_as in {"resource", "both"}: resource_uri = request.resource_uri or _default_summary_resource_uri(request, artifact["title"]) artifact["resource_uri"] = resource_uri resource_result = await ov_client.add_resource( uri=resource_uri, content=json.dumps(artifact, ensure_ascii=False, indent=2), resource_type=request.resource_type or "json", ) return { "status": "ok", "artifact": artifact, "memory_result": memory_result, "resource_result": resource_result, } # FastAPI 应用 app = FastAPI(title="Memory Gateway", version="0.1.0", lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @app.get("/health", dependencies=[Depends(verify_api_key)]) async def health_check(): """健康检查""" try: ov_client = await get_openviking_client() ov_status = await ov_client.health_check() return { "status": "ok", "gateway": "memory-gateway", "openviking": ov_status, } except Exception as e: return { "status": "degraded", "gateway": "memory-gateway", "error": str(e), } mcp_router = APIRouter() async def mcp_server_events(request: Request, _: None = Depends(verify_api_key)): """MCP Server-Sent Events 端点 - 使用 stdio 模式模拟""" async def event_generator(): # 发送初始化消息 yield {"event": "initialize", "data": json.dumps({"protocolVersion": "2024-11-05"})} # 保持连接 try: while True: await asyncio.sleep(30) yield {"event": "ping", "data": ""} except asyncio.CancelledError: pass return EventSourceResponse(event_generator()) mcp_router.add_api_route("/sse", mcp_server_events, methods=["GET"]) # MCP JSON-RPC 端点（简化实现） async def mcp_rpc(request: Request, _: None = Depends(verify_api_key)): """处理 MCP JSON-RPC 请求""" body = await request.json() method = body.get("method") params = body.get("params", {}) msg_id = body.get("id") try: if method == "tools/list": tools = await list_tools() result = { "tools": [ { "name": t.name, "description": t.description, "inputSchema": t.inputSchema, } for t in tools ] } elif method == "tools/call": tool_name = params.get("name") tool_args = params.get("arguments", {}) result_content = await call_tool_tool(tool_name, tool_args) result = {"content": [c.model_dump() for c in result_content]} else: return JSONResponse( status_code=400, content={"jsonrpc": "2.0", "error": {"code": -32601, "message": f"Method not found: {method}"}, "id": msg_id} ) return {"jsonrpc": "2.0", "result": result, "id": msg_id} except Exception as e: logger.error(f"MCP RPC 错误: {e}") return JSONResponse( status_code=500, content={"jsonrpc": "2.0", "error": {"code": -32603, "message": str(e)}, "id": msg_id} ) async def call_tool_tool(name: str, arguments: dict) -> list[TextContent]: """调用工具的内部函数""" return await call_tool(name, arguments) mcp_router.add_api_route("/rpc", mcp_rpc, methods=["POST"]) # 注册 MCP 路由 app.include_router(mcp_router, prefix="/mcp", tags=["mcp"]) @app.post("/api/search", dependencies=[Depends(verify_api_key)]) async def api_search(request: SearchRequest): """REST API: 搜索""" ov_client = await get_openviking_client() result = await ov_client.search( query=request.query, namespace=request.namespace or get_config().memory.default_namespace, limit=request.limit or get_config().memory.search_limit, uri=request.uri, ) return {"results": result.results, "total": result.total} @app.post("/api/memory", dependencies=[Depends(verify_api_key)]) async def api_add_memory(request: AddMemoryRequest): """REST API: 添加记忆""" ov_client = await get_openviking_client() result = await ov_client.add_memory( content=request.content, namespace=request.namespace or get_config().memory.default_namespace, memory_type=request.memory_type, ) return result @app.post("/api/resource", dependencies=[Depends(verify_api_key)]) async def api_add_resource(request: AddResourceRequest): """REST API: 添加资源""" ov_client = await get_openviking_client() result = await ov_client.add_resource( uri=request.uri, content=request.content, resource_type=request.resource_type, ) return result @app.post("/api/summary", dependencies=[Depends(verify_api_key)]) async def api_commit_summary(request: CommitSummaryRequest): """REST API: 通用内容 LLM 总结与记忆沉淀。""" try: return await commit_summary_to_openviking(request) except LLMConfigurationError as exc: raise HTTPException(status_code=503, detail=str(exc)) from exc except (LLMSummaryError, Exception) as exc: if isinstance(exc, HTTPException): raise raise HTTPException(status_code=502, detail=f"LLM summary failed: {exc}") from exc def _parse_tags(tags: str | None) -> list[str]: if not tags: return [] return [tag.strip() for tag in re.split(r"[,\n]", tags) if tag.strip()] def _default_knowledge_uri(namespace: str, knowledge_type: str, title: str, content: str) -> str: digest = hashlib.sha1(content.encode("utf-8")).hexdigest()[:12] return f"viking://resources/{namespace.strip('/')}/knowledge/{knowledge_type.strip('/')}/{slugify(title, digest)}-{digest}.json" @app.post("/api/knowledge/upload", dependencies=[Depends(verify_api_key)]) async def api_upload_knowledge( file: UploadFile = File(...), title: Optional[str] = Form(default=None), namespace: str = Form(default="memory-gateway"), knowledge_type: str = Form(default="knowledge"), tags: str = Form(default=""), source: Optional[str] = Form(default=None), obsidian_dir: Optional[str] = Form(default=None), resource_uri: Optional[str] = Form(default=None), persist_as: str = Form(default="resource"), max_summary_chars: int = Form(default=1000), ): """Upload a document, convert it to Markdown, save to Obsidian, summarize with LLM, and commit to OpenViking.""" if persist_as not in {"memory", "resource", "both", "none"}: raise HTTPException(status_code=422, detail="persist_as must be one of memory/resource/both/none") original_name = file.filename or "uploaded-document" suffix = Path(original_name).suffix or ".bin" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(await file.read()) tmp_path = Path(tmp.name) try: markdown = await asyncio.to_thread(convert_file_to_markdown, tmp_path) except RuntimeError as exc: tmp_path.unlink(missing_ok=True) raise HTTPException(status_code=500, detail=str(exc)) from exc except Exception as exc: # noqa: BLE001 tmp_path.unlink(missing_ok=True) raise HTTPException(status_code=500, detail=f"Document conversion failed: {exc}") from exc finally: tmp_path.unlink(missing_ok=True) parsed_tags = _parse_tags(tags) effective_title = title or Path(original_name).stem or "Uploaded knowledge" request = CommitSummaryRequest( content=markdown, title=effective_title, namespace=namespace, memory_type=knowledge_type, tags=parsed_tags, source=source or original_name, persist_as="none", max_summary_chars=max_summary_chars, purpose=f"knowledge upload: {knowledge_type}", ) try: artifact = await build_summary_artifact(request) except LLMConfigurationError as exc: raise HTTPException(status_code=503, detail=str(exc)) from exc except Exception as exc: # noqa: BLE001 raise HTTPException(status_code=502, detail=f"LLM summary failed: {exc}") from exc config = get_config() relative_dir = obsidian_dir or getattr(config.obsidian, "knowledge_dir", "01_Knowledge/Uploaded") obsidian_path = save_markdown_to_obsidian( vault_path=config.obsidian.vault_path, relative_dir=relative_dir, title=artifact["title"], markdown=markdown, source_filename=original_name, tags=artifact.get("tags", []), knowledge_type=knowledge_type, summary=artifact.get("summary"), ) artifact.update( { "schema_version": "memory-gateway.knowledge_upload.v1", "knowledge_type": knowledge_type, "source_filename": original_name, "obsidian_path": str(obsidian_path), "obsidian_relative_path": str(obsidian_path.relative_to(config.obsidian.vault_path)), "markdown_content": markdown, } ) ov_client = await get_openviking_client() memory_result: Optional[dict[str, Any]] = None resource_result: Optional[dict[str, Any]] = None if persist_as in {"memory", "both"}: memory_result = await ov_client.add_memory( content=_render_memory_text(artifact), namespace=namespace, memory_type=knowledge_type, ) if persist_as in {"resource", "both"}: final_uri = resource_uri or _default_knowledge_uri(namespace, knowledge_type, artifact["title"], markdown) artifact["resource_uri"] = final_uri resource_result = await ov_client.add_resource( uri=final_uri, content=json.dumps(artifact, ensure_ascii=False, indent=2), resource_type="json", ) return { "status": "ok", "artifact": artifact, "markdown_chars": len(markdown), "obsidian_path": str(obsidian_path), "memory_result": memory_result, "resource_result": resource_result, } def create_app(config: Optional[Config] = None) -> FastAPI: """创建 FastAPI 应用""" if config: set_config(config) return app # 入口点 def main(): """主入口""" import argparse import uvicorn parser = argparse.ArgumentParser(description="Memory Gateway MCP Server") parser.add_argument("--config", default="config.yaml", help="配置文件路径") parser.add_argument("--host", default=None, help="监听地址") parser.add_argument("--port", type=int, default=None, help="监听端口") args = parser.parse_args() # 加载配置 from .config import load_config as load config = load(args.config) if args.host: config.server.host = args.host if args.port: config.server.port = args.port set_config(config) # 启动服务 uvicorn.run( app, host=config.server.host, port=config.server.port, log_level=config.logging.level.lower(), ) if __name__ == "__main__": main()