142 lines
4.1 KiB
Python
142 lines
4.1 KiB
Python
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()
|