diff --git a/app/mcp_server.py b/app/mcp_server.py new file mode 100644 index 0000000..cb9828f --- /dev/null +++ b/app/mcp_server.py @@ -0,0 +1,308 @@ +"""market_sync MCP Server。 + +提供 5 类工具(tool),让 LLM 能查 market_sync 的同步情况: + +1. list_sync_tasks — 列出所有同步任务 +2. get_task_status — 查单个任务的状态机 + 最近 run +3. list_datasets — 列出所有数据集(含数据规模) +4. get_dataset_info — 单个数据集的详细范围(日期范围/数据量/列) +5. get_today_sync_summary — 今日同步汇总 + +传输: stdio(本地启动 Claude Desktop / Claude Code 直连)。 + +启动: + python -m app.mcp_server + # 或作为 MCP server 注册到客户端 + +架构: + - MCP 协议层: mcp.server.Server + stdio_server + - 业务层: 复用 app.core.db.ops + app.core.sync.registry 的 helper + - 不直接操作 PG,通过 ORM session(一致的事务/连接管理) +""" +from __future__ import annotations + +import asyncio +import json +import sys +from datetime import datetime +from pathlib import Path + +_PROJECT_ROOT = Path(__file__).resolve().parent.parent +if str(_PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(_PROJECT_ROOT)) + + +# ── MCP server bootstrap ──────────────────────────────────────────── +from mcp.server import Server +from mcp.server.stdio import stdio_server +from mcp.types import Tool, TextContent + +server = Server("market_sync") + + +# ── helpers ───────────────────────────────────────────────────────── + +def _serialize(obj): + """JSON-friendly: datetime/date/Decimal → str.""" + if isinstance(obj, dict): + return {k: _serialize(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_serialize(v) for v in obj] + if isinstance(obj, (datetime,)): + return obj.isoformat() + try: + # handle date + import datetime as _dt + if isinstance(obj, _dt.date): + return obj.isoformat() + except Exception: + pass + return obj + + +def _ok(data) -> list[TextContent]: + """Standard success response: JSON dump.""" + return [TextContent( + type="text", + text=json.dumps(_serialize(data), ensure_ascii=False, indent=2, default=str), + )] + + +def _err(msg: str) -> list[TextContent]: + return [TextContent(type="text", text=f"❌ {msg}")] + + +# ── Tool implementations ───────────────────────────────────────────── + +def list_sync_tasks() -> list[dict]: + from app.tasks import TASKS + from app.core.db import ops as db_ops + from app.core.sync import get_registry_status + + registry = {r["dataset_id"]: r for r in get_registry_status()} + history = db_ops.list_sync_history(days=7, limit=1000) + last_by_task: dict[str, dict] = {} + for h in history: + if h["dataset_id"] not in last_by_task: + last_by_task[h["dataset_id"]] = h + + out = [] + for ds_id, cls in TASKS.items(): + reg = registry.get(ds_id, {}) + last = last_by_task.get(ds_id) + out.append({ + "dataset_id": ds_id, + "name": reg.get("name", ""), + "description": reg.get("description", ""), + "source": reg.get("source", ""), + "sort_order": reg.get("sort_order", 0), + "current_status": reg.get("status", "unknown"), + "last_run_at": (last or {}).get("started_at"), + "last_run_status": (last or {}).get("status"), + "last_run_rows": (last or {}).get("rows_written"), + "history_count_7d": sum(1 for h in history if h["dataset_id"] == ds_id), + }) + return sorted(out, key=lambda x: x["sort_order"]) + + +def get_task_status(dataset_id: str, days: int = 7) -> dict: + from app.core.db import ops as db_ops + from app.core.sync import get_registry_status + + reg = next((r for r in get_registry_status() if r["dataset_id"] == dataset_id), None) + if not reg: + return {"error": f"unknown dataset_id: {dataset_id}"} + history = db_ops.list_sync_history(dataset_id=dataset_id, days=days, limit=50) + return { + "dataset_id": dataset_id, + "current_state": reg, + "recent_runs": history, + } + + +def list_datasets() -> list[dict]: + """每个数据集的当前规模(行数 + 日期范围 + 来源)。""" + from sqlalchemy import func, select + from app.core.db.models import ( + KlineStock, KlineIndex, Kline5Min, Moneyflow, Share, TickTrade, + SectorIndices, SectorFeaturesDaily, MarketRegimeDaily, + LonghubangDaily, LonghubangSeat, Stock, Industry, + NodeCategory, Node, StockNodeMap, KlineStockMADaily, + ) + from app.core.db.orm import SessionLocal + + targets = [ + ("stocks", Stock), + ("kline_stock", KlineStock), + ("kline_index", KlineIndex), + ("kline_5min", Kline5Min), + ("kline_stock_ma_daily", KlineStockMADaily), + ("moneyflow", Moneyflow), + ("share", Share), + ("tick_trade", TickTrade), + ("sector_indices", SectorIndices), + ("sector_features_daily", SectorFeaturesDaily), + ("market_regime_daily", MarketRegimeDaily), + ("longhubang_daily", LonghubangDaily), + ("longhubang_seat", LonghubangSeat), + ("industry", Industry), + ("node_categories", NodeCategory), + ("nodes", Node), + ("stock_node_map", StockNodeMap), + ] + out = [] + with SessionLocal() as s: + for name, model in targets: + try: + count = int(s.execute(select(func.count()).select_from(model)).scalar() or 0) + dr = None + if hasattr(model, "trade_date"): + mn = s.execute(select(func.min(model.trade_date))).scalar() + mx = s.execute(select(func.max(model.trade_date))).scalar() + dr = {"min": str(mn), "max": str(mx)} if mn else None + out.append({"table": name, "rows": count, "date_range": dr}) + except Exception as e: + out.append({"table": name, "rows": -1, "error": str(e)[:100]}) + return out + + +def get_dataset_info(table_name: str) -> dict: + """单表 schema + 数据范围 + 抽样行。""" + from sqlalchemy import inspect + from app.core.db.orm import SessionLocal + + # 表名白名单(避免任意 SQL 注入) + allowed = { + "stocks", "kline_stock", "kline_index", "kline_5min", + "kline_stock_ma_daily", "moneyflow", "share", "tick_trade", + "sector_indices", "sector_features_daily", "market_regime_daily", + "longhubang_daily", "longhubang_seat", "industry", + "node_categories", "nodes", "stock_node_map", + "dataset_registry", "sync_history", "config", + } + if table_name not in allowed: + return {"error": f"unknown table: {table_name}", "allowed": sorted(allowed)} + + with SessionLocal() as s: + insp = inspect(s.get_bind()) + cols = [{"name": c["name"], "type": str(c["type"]), "nullable": c.get("nullable", True)} + for c in insp.get_columns(table_name, schema="market_data")] + pk = insp.get_pk_constraint(table_name, schema="market_data") + idx = [i["name"] for i in insp.get_indexes(table_name, schema="market_data")] + sample = s.execute( + __import__("sqlalchemy").text( + f'SELECT * FROM market_data."{table_name}" LIMIT 3' + ) + ).mappings().all() + return { + "table": table_name, + "schema": "market_data", + "primary_key": pk.get("constrained_columns", []), + "indexes": idx, + "columns": cols, + "sample_rows": [dict(r) for r in sample], + } + + +def get_today_sync_summary() -> dict: + from datetime import date + from app.core.db import ops as db_ops + + today = date.today().isoformat() + summary_rows = db_ops.daily_sync_summary(days=1) + today_summary = next((r for r in summary_rows if r["date"] == today), None) or { + "date": today, "ok": 0, "warning": 0, "error": 0, "blocked": 0, "total": 0, + } + history = db_ops.list_sync_history(days=1, limit=500) + return { + "date": today, + "summary": today_summary, + "tasks_today": history, + } + + +# ── MCP tool 注册 ────────────────────────────────────────────────── + +@server.list_tools() +async def list_tools(): + return [ + Tool( + name="list_sync_tasks", + description="列出 market_sync 项目的所有同步任务(dataset_id + 名称 + 当前状态 + 最近 7 天运行统计)。", + inputSchema={"type": "object", "properties": {}}, + ), + Tool( + name="get_task_status", + description="查询单个同步任务的详细状态机 + 最近 N 天所有 run 历史。", + inputSchema={ + "type": "object", + "properties": { + "dataset_id": {"type": "string", "description": "例如 kline_daily / moneyflow"}, + "days": {"type": "integer", "description": "查最近 N 天, 默认 7", "default": 7}, + }, + "required": ["dataset_id"], + }, + ), + Tool( + name="list_datasets", + description="列出所有 PG 数据集(含 market_data schema 下 17 张业务表)的行数 + 日期范围。", + inputSchema={"type": "object", "properties": {}}, + ), + Tool( + name="get_dataset_info", + description="查询单个数据集的表结构(列名 + 类型 + 主键 + 索引)+ 3 行抽样数据。", + inputSchema={ + "type": "object", + "properties": { + "table_name": {"type": "string", "description": "如 kline_stock / moneyflow / tick_trade"}, + }, + "required": ["table_name"], + }, + ), + Tool( + name="get_today_sync_summary", + description="今日同步情况汇总(OK / Warning / Error / Blocked 计数 + 每个任务的最新 run)。", + inputSchema={"type": "object", "properties": {}}, + ), + ] + + +@server.call_tool() +async def call_tool(name: str, arguments: dict): + try: + if name == "list_sync_tasks": + return _ok(list_sync_tasks()) + elif name == "get_task_status": + ds_id = arguments.get("dataset_id", "") + days = int(arguments.get("days", 7)) + return _ok(get_task_status(ds_id, days)) + elif name == "list_datasets": + return _ok(list_datasets()) + elif name == "get_dataset_info": + return _ok(get_dataset_info(arguments.get("table_name", ""))) + elif name == "get_today_sync_summary": + return _ok(get_today_sync_summary()) + else: + return _err(f"unknown tool: {name}") + except Exception as e: + import traceback + return _err(f"{type(e).__name__}: {e}\n{traceback.format_exc(limit=3)}") + + +# ── main ──────────────────────────────────────────────────────────── + +async def main(): + # 启动时注册数据源(让 datasource.health_check 之类方法可用) + from app.core.datasource.registry import build_default_registry + build_default_registry() + + async with stdio_server() as (read_stream, write_stream): + await server.run( + read_stream, + write_stream, + server.create_initialization_options(), + ) + + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 147cd33..4ca50b3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,9 @@ uvicorn[standard]>=0.30.0 pydantic>=2.0.0 pydantic-settings>=2.0.0 +# MCP server (stdio 传输, 给 Claude Desktop / Claude Code 用) +mcp[cli]>=1.0.0 + # Data pandas>=2.2.0 numpy>=1.26.0