feat(mcp): MCP server (5 tools) — 任务/状态/数据集/详情/今日汇总
新增 app/mcp_server.py，让 LLM (Claude Desktop / Claude Code) 通过 MCP stdio 协议查询 market_sync 的同步情况： 工具列表: 1. list_sync_tasks - 列出 13 个 sync task + 当前状态 + 7d 运行统计 2. get_task_status - 单个 task 的状态机 + 最近 N 天所有 run 3. list_datasets - 17 张 PG 表的行数 + 日期范围 4. get_dataset_info - 单表的 schema/索引/PK + 3 行抽样(白名单防注入) 5. get_today_sync_summary - 今日 KPI + 每个 task 的最新 run 传输: stdio (本地启动 Claude Desktop 直连) 启动: python -m app.mcp_server 依赖: mcp[cli]>=1.0.0 (已加入 requirements.txt) 复用层: app.core.db.ops + app.core.sync.registry 的 helper, 不绕过 ORM session (与 dashboard API 一致)。 验证: - initialize 握手 OK - tools/list 返回 5 个工具的完整 schema - 各 tool 在 stdio 协议下正常返回 JSON
This commit is contained in:
gao
@@ -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())
|
||||
Reference in New Issue
Block a user