MCP 服务：ds 余额查询 + ros 邮件/QQ/服务器状态

2026-08-27 21:43:28 +08:00
commit 74402b5e55
4 changed files with 1209 additions and 0 deletions
@@ -0,0 +1,141 @@
+import json
+import os
+import sys
+import requests
+
+BASE_URL = 'https://api.deepseek.com'
+
+
+# ============================================================
+# 内部辅助
+# ============================================================
+
+def resolve_api_key(arguments):
+    """解析 API Key，优先级：参数 > DEEPSEEK_API_KEY 环境变量"""
+    if arguments.get('api_key'):
+        return arguments['api_key']
+
+    key = os.environ.get('DEEPSEEK_API_KEY', '')
+    if key:
+        return key
+
+    return None
+
+
+# ============================================================
+# 工具函数
+# ============================================================
+
+def get_user_balance(api_key=None):
+    """查询 DeepSeek 账户余额"""
+    key = resolve_api_key({'api_key': api_key} if api_key else {})
+    if not key:
+        raise Exception(
+            '未找到 API Key，请通过以下任一方式提供：\n'
+            '1. 调用工具时传入 api_key 参数\n'
+            '2. 设置环境变量 DEEPSEEK_API_KEY'
+        )
+
+    headers = {
+        'Authorization': f'Bearer {key}',
+        'Accept': 'application/json',
+    }
+    resp = requests.get(f'{BASE_URL}/user/balance', headers=headers, timeout=15)
+    if resp.status_code != 200:
+        raise Exception(f'API 请求失败 (HTTP {resp.status_code}): {resp.text}')
+    return resp.json()
+
+
+# ============================================================
+# MCP JSON-RPC 处理
+# ============================================================
+
+TOOLS = [
+    {
+        "name": "get_user_balance",
+        "description": "查询 DeepSeek API 账户余额，返回各币种的可用余额、赠送余额和充值余额。支持从参数或环境变量 DEEPSEEK_API_KEY 读取 API Key。",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "api_key": {
+                    "type": "string",
+                    "description": "DeepSeek API Key（可选，不传则从环境变量 DEEPSEEK_API_KEY 读取）"
+                }
+            },
+            "required": []
+        }
+    }
+]
+
+TOOL_MAP = {
+    "get_user_balance": get_user_balance
+}
+
+
+def handle_request(req):
+    method = req.get("method", "")
+    req_id = req.get("id")
+
+    if method == "initialize":
+        return {
+            "jsonrpc": "2.0",
+            "id": req_id,
+            "result": {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {
+                    "tools": {}
+                },
+                "serverInfo": {
+                    "name": "ds-balance-inquiry",
+                    "version": "1.0.0"
+                }
+            }
+        }
+
+    if method == "ping":
+        return {"jsonrpc": "2.0", "id": req_id, "result": {}}
+
+    if method == "notifications/initialized":
+        return None
+
+    if method == "tools/list":
+        return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}}
+
+    if method == "tools/call":
+        tool_name = req["params"]["name"]
+        arguments = req["params"].get("arguments", {})
+        tool_func = TOOL_MAP.get(tool_name)
+        if not tool_func:
+            return {"jsonrpc": "2.0", "id": req_id,
+                    "error": {"code": -32601, "message": f"未知工具: {tool_name}"}}
+        try:
+            result = tool_func(**arguments)
+            text = json.dumps(result, ensure_ascii=False, indent=2)
+            return {"jsonrpc": "2.0", "id": req_id,
+                    "result": {"content": [{"type": "text", "text": text}]}}
+        except Exception as e:
+            return {"jsonrpc": "2.0", "id": req_id,
+                    "error": {"code": -32000, "message": str(e)}}
+
+    return {"jsonrpc": "2.0", "id": req_id,
+            "error": {"code": -32601, "message": f"未知方法: {method}"}}
+
+
+def main():
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            req = json.loads(line)
+            resp = handle_request(req)
+            if resp is None:
+                continue
+            sys.stdout.write(json.dumps(resp, ensure_ascii=False) + '\n')
+            sys.stdout.flush()
+        except json.JSONDecodeError:
+            pass
+
+
+if __name__ == '__main__':
+    main()