779 lines
24 KiB
Python
779 lines
24 KiB
Python
import httpx
|
||
from mcp.server.fastmcp import FastMCP
|
||
import os
|
||
import json
|
||
|
||
def load_config():
|
||
config_path = os.path.join(os.path.dirname(__file__), "../config/config.json")
|
||
default_config = {
|
||
"DATASET_API_BASE_URL": "https://www.ai4mats.com",
|
||
"DATASET_DEFAULT_USERNAME": "fanshuai",
|
||
"DATASET_DEFAULT_PASSWORD": "h1n2x3j4y5@",
|
||
"MCP_TRANSPORT": "stdio"
|
||
}
|
||
|
||
try:
|
||
if os.path.exists(config_path):
|
||
with open(config_path, 'r', encoding='utf-8') as f:
|
||
config = json.load(f)
|
||
default_config.update(config)
|
||
print(f"✅ 已加载配置文件: {config_path}")
|
||
else:
|
||
print(f"⚠️ 配置文件不存在，使用默认配置: {config_path}")
|
||
except Exception as e:
|
||
print(f"❌ 读取配置文件失败: {e}，使用默认配置")
|
||
|
||
return default_config
|
||
|
||
config = load_config()
|
||
API_BASE_URL = config.get("DATASET_API_BASE_URL")
|
||
DEFAULT_USERNAME = config.get("DATASET_DEFAULT_USERNAME")
|
||
DEFAULT_PASSWORD = config.get("DATASET_DEFAULT_PASSWORD")
|
||
TRANSPORT_MODE = config.get("MCP_TRANSPORT")
|
||
|
||
mcp = FastMCP("Dataset-Service")
|
||
|
||
@mcp.tool()
|
||
async def login(
|
||
username: str = "",
|
||
password: str = ""
|
||
) -> str:
|
||
"""
|
||
获取访问token
|
||
|
||
参数说明:
|
||
- username: 用户名 (可选，默认使用配置文件中的用户名)
|
||
- password: 密码 (可选，默认使用配置文件中的密码)
|
||
|
||
返回: 包含access_token和expires_in的JSON字符串。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/auth/login"
|
||
|
||
if not username:
|
||
username = DEFAULT_USERNAME
|
||
if not password:
|
||
password = DEFAULT_PASSWORD
|
||
|
||
payload = {
|
||
"username": username,
|
||
"password": password
|
||
}
|
||
|
||
headers = {"Content-Type": "application/json"}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.post(api_url, json=payload, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"登录 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接认证服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"登录时发生未知错误: {str(e)}"
|
||
|
||
def build_auth_headers(token: str) -> dict:
|
||
"""构建包含token的请求头"""
|
||
return {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {token}"
|
||
}
|
||
|
||
@mcp.tool()
|
||
async def add_dataset(
|
||
name: str,
|
||
token: str,
|
||
preview_pic: str = "",
|
||
dataset_source: str = "add",
|
||
data_type: str = "通用数据",
|
||
data_tag: str = "",
|
||
is_public: bool = False,
|
||
is_hot_stone: bool = False
|
||
) -> str:
|
||
"""
|
||
新增数据集
|
||
|
||
参数说明:
|
||
- name: 数据集名称
|
||
- token: 访问令牌
|
||
- preview_pic: 预览图片URL (可选)
|
||
- dataset_source: 数据集来源 (默认: add)
|
||
- data_type: 数据类型 (默认: 通用数据)
|
||
- data_tag: 数据标签 (可选)
|
||
- is_public: 是否公开 (默认: False)
|
||
- is_hot_stone: 是否热门 (默认: False)
|
||
|
||
返回: 新增数据集的详细信息。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/addDataset"
|
||
|
||
payload = {
|
||
"name": name,
|
||
"preview_pic": preview_pic,
|
||
"dataset_source": dataset_source,
|
||
"data_type": data_type,
|
||
"data_tag": data_tag,
|
||
"is_public": is_public,
|
||
"is_hot_stone": is_hot_stone
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.post(api_url, json=payload, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"新增数据集 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"新增数据集时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def get_asset_icon(
|
||
token: str,
|
||
page: int = 0,
|
||
size: int = 10000,
|
||
category_id: int = 1
|
||
) -> str:
|
||
"""
|
||
查询数据集分类接口
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- page: 页码 (默认: 0)
|
||
- size: 每页数量 (默认: 10000)
|
||
- category_id: 分类ID (默认: 1)
|
||
|
||
返回: 数据集分类列表，包含一级分类和二级分类信息。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/assetIcon"
|
||
|
||
params = {
|
||
"page": page,
|
||
"size": size,
|
||
"category_id": category_id
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.get(api_url, params=params, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"查询数据集分类 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"查询数据集分类时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def get_data_types(token: str) -> str:
|
||
"""
|
||
获取可用的数据集类型列表
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
|
||
返回: 可用的数据类型列表，格式为[{"id": 138, "name": "通用数据"}, ...]
|
||
"""
|
||
try:
|
||
asset_icon_result = await get_asset_icon(token)
|
||
asset_icon_data = json.loads(asset_icon_result)
|
||
|
||
if asset_icon_data.get("code") != 200:
|
||
return f"获取数据集分类失败: {asset_icon_data.get('msg', '未知错误')}"
|
||
|
||
data = asset_icon_data.get("data", [])
|
||
data_types = []
|
||
|
||
for category in data:
|
||
second_list = category.get("second_asset_icon_list", [])
|
||
if second_list:
|
||
for item in second_list:
|
||
data_types.append({
|
||
"id": item.get("id"),
|
||
"name": item.get("name"),
|
||
"parent_id": item.get("parent_id"),
|
||
"path": item.get("path")
|
||
})
|
||
else:
|
||
data_types.append({
|
||
"id": category.get("id"),
|
||
"name": category.get("name"),
|
||
"parent_id": category.get("parent_id"),
|
||
"path": category.get("path")
|
||
})
|
||
|
||
return json.dumps({
|
||
"code": 200,
|
||
"msg": "操作成功",
|
||
"data": data_types
|
||
}, indent=2, ensure_ascii=False)
|
||
|
||
except json.JSONDecodeError as e:
|
||
return f"解析数据集分类失败: {str(e)}"
|
||
except Exception as e:
|
||
return f"获取数据类型列表时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def upload_chunk(
|
||
token: str,
|
||
chunkNumber: int,
|
||
chunkSize: int,
|
||
currentChunkSize: int,
|
||
totalSize: int,
|
||
identifier: str,
|
||
filename: str,
|
||
relativePath: str,
|
||
totalChunks: int
|
||
) -> str:
|
||
"""
|
||
上传版本文件（分片上传）
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- chunkNumber: 当前分片序号
|
||
- chunkSize: 分片大小
|
||
- currentChunkSize: 当前分片实际大小
|
||
- totalSize: 总文件大小
|
||
- identifier: 文件唯一标识
|
||
- filename: 文件名
|
||
- relativePath: 相对路径
|
||
- totalChunks: 总分片数
|
||
|
||
返回: 上传结果，包含文件location用于新增版本。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/uploader/chunk"
|
||
|
||
params = {
|
||
"chunkNumber": chunkNumber,
|
||
"chunkSize": chunkSize,
|
||
"currentChunkSize": currentChunkSize,
|
||
"totalSize": totalSize,
|
||
"identifier": identifier,
|
||
"filename": filename,
|
||
"relativePath": relativePath,
|
||
"totalChunks": totalChunks
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.get(api_url, params=params, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"上传文件 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"上传文件时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def add_version(
|
||
token: str,
|
||
git_id: int,
|
||
id: int,
|
||
identifier: str,
|
||
is_public: bool,
|
||
owner: str,
|
||
name: str,
|
||
version: str,
|
||
version_desc: str,
|
||
dataset_source: str = "add",
|
||
dataset_version_vos: list = None
|
||
) -> str:
|
||
"""
|
||
新增数据集版本
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- git_id: Git仓库ID
|
||
- id: 数据集ID
|
||
- identifier: 数据集标识
|
||
- is_public: 是否公开
|
||
- owner: 所有者
|
||
- name: 数据集名称
|
||
- version: 版本号 (例如: v1)
|
||
- version_desc: 版本描述
|
||
- dataset_source: 数据集来源 (默认: add)
|
||
- dataset_version_vos: 版本文件列表，格式: [{"file_name":"xxx","file_size":xxx,"url":"xxx"}]
|
||
|
||
返回: 新增版本结果。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/addVersion"
|
||
|
||
payload = {
|
||
"git_id": git_id,
|
||
"id": id,
|
||
"identifier": identifier,
|
||
"is_public": is_public,
|
||
"owner": owner,
|
||
"name": name,
|
||
"version": version,
|
||
"version_desc": version_desc,
|
||
"dataset_source": dataset_source,
|
||
"dataset_version_vos": dataset_version_vos if dataset_version_vos else []
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.post(api_url, json=payload, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"新增版本 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"新增版本时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def update_dataset(
|
||
token: str,
|
||
id: int,
|
||
name: str,
|
||
identifier: str,
|
||
description: str = "",
|
||
is_public: bool = False,
|
||
data_type: str = "通用数据",
|
||
data_tag: str = "",
|
||
praises_count: int = 0,
|
||
praised: bool = False,
|
||
create_by: str = "",
|
||
update_time: str = "",
|
||
owner: str = "",
|
||
dataset_source: str = "add",
|
||
relative_paths: str = "",
|
||
is_hot_stone: bool = False,
|
||
git_id: int = 0,
|
||
preview_pic: str = "",
|
||
type: int = 0
|
||
) -> str:
|
||
"""
|
||
修改数据集
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- id: 数据集ID
|
||
- name: 数据集名称
|
||
- identifier: 数据集标识
|
||
- description: 描述 (可选)
|
||
- is_public: 是否公开 (默认: False)
|
||
- data_type: 数据类型 (默认: 通用数据)
|
||
- data_tag: 数据标签 (可选)
|
||
- praises_count: 点赞数 (默认: 0)
|
||
- praised: 是否已点赞 (默认: False)
|
||
- create_by: 创建者 (可选)
|
||
- update_time: 更新时间 (可选)
|
||
- owner: 所有者 (可选)
|
||
- dataset_source: 数据集来源 (默认: add)
|
||
- relative_paths: 相对路径 (可选)
|
||
- is_hot_stone: 是否热门 (默认: False)
|
||
- git_id: Git仓库ID (默认: 0)
|
||
- preview_pic: 预览图片URL (可选)
|
||
- type: 类型 (默认: 0)
|
||
|
||
返回: 修改后的数据集信息。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/updateDataset"
|
||
|
||
payload = {
|
||
"id": id,
|
||
"name": name,
|
||
"identifier": identifier,
|
||
"description": description,
|
||
"is_public": is_public,
|
||
"data_type": data_type,
|
||
"data_tag": data_tag,
|
||
"praises_count": praises_count,
|
||
"praised": praised,
|
||
"create_by": create_by,
|
||
"update_time": update_time,
|
||
"owner": owner,
|
||
"dataset_source": dataset_source,
|
||
"relative_paths": relative_paths,
|
||
"is_hot_stone": is_hot_stone,
|
||
"git_id": git_id,
|
||
"preview_pic": preview_pic,
|
||
"type": type
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.put(api_url, json=payload, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"修改数据集 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"修改数据集时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def update_desc(
|
||
token: str,
|
||
git_id: int,
|
||
identifier: str,
|
||
description: str
|
||
) -> str:
|
||
"""
|
||
编辑数据集简介
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- git_id: Git仓库ID
|
||
- identifier: 数据集标识
|
||
- description: 数据集简介
|
||
|
||
返回: 操作结果。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/updateDesc"
|
||
|
||
payload = {
|
||
"git_id": git_id,
|
||
"identifier": identifier,
|
||
"description": description
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.put(api_url, json=payload, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"编辑简介 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"编辑简介时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def publish_dataset(
|
||
token: str,
|
||
id: int,
|
||
name: str
|
||
) -> str:
|
||
"""
|
||
发布数据集
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- id: 数据集ID
|
||
- name: 数据集名称
|
||
|
||
返回: 发布后的数据集信息。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/publish"
|
||
|
||
payload = {
|
||
"id": id,
|
||
"name": name
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.post(api_url, json=payload, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"发布数据集 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"发布数据集时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def download_all_files(
|
||
token: str,
|
||
name: str,
|
||
git_id: int,
|
||
version: str,
|
||
identifier: str,
|
||
owner: str,
|
||
is_public: bool
|
||
) -> str:
|
||
"""
|
||
当前版本所有文件打包下载
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- name: 数据集名称
|
||
- git_id: Git仓库ID
|
||
- version: 版本号
|
||
- identifier: 数据集标识
|
||
- owner: 所有者
|
||
- is_public: 是否公开
|
||
|
||
返回: 文件下载链接或文件内容。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/downloadAllFiles"
|
||
|
||
params = {
|
||
"name": name,
|
||
"git_id": git_id,
|
||
"version": version,
|
||
"identifier": identifier,
|
||
"owner": owner,
|
||
"is_public": is_public
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
response = await client.get(api_url, params=params, headers=headers)
|
||
response.raise_for_status()
|
||
|
||
content_type = response.headers.get("content-type", "")
|
||
if "application/json" in content_type:
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
else:
|
||
return f"文件下载成功，内容长度: {len(response.content)} bytes"
|
||
except httpx.HTTPStatusError as e:
|
||
return f"下载文件 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"下载文件时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def download_single_file(
|
||
token: str,
|
||
url: str
|
||
) -> str:
|
||
"""
|
||
当前版本选中文件下载
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- url: 文件路径
|
||
|
||
返回: 文件下载链接或文件内容。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/downloadSingleFile"
|
||
|
||
params = {"url": url}
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
response = await client.get(api_url, params=params, headers=headers)
|
||
response.raise_for_status()
|
||
|
||
content_type = response.headers.get("content-type", "")
|
||
if "application/json" in content_type:
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
else:
|
||
return f"文件下载成功，内容长度: {len(response.content)} bytes"
|
||
except httpx.HTTPStatusError as e:
|
||
return f"下载单个文件 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"下载单个文件时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def delete_version(
|
||
token: str,
|
||
git_id: int,
|
||
owner: str,
|
||
identifier: str,
|
||
relative_paths: str,
|
||
version: str
|
||
) -> str:
|
||
"""
|
||
删除当前版本
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- git_id: Git仓库ID
|
||
- owner: 所有者
|
||
- identifier: 数据集标识
|
||
- relative_paths: 相对路径
|
||
- version: 版本号
|
||
|
||
返回: 操作结果。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/deleteDatasetVersion"
|
||
|
||
params = {
|
||
"git_id": git_id,
|
||
"owner": owner,
|
||
"identifier": identifier,
|
||
"relative_paths": relative_paths,
|
||
"version": version
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.delete(api_url, params=params, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"删除版本 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"删除版本时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def delete_dataset(
|
||
token: str,
|
||
id: int
|
||
) -> str:
|
||
"""
|
||
删除数据集
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- id: 数据集ID
|
||
|
||
返回: 操作结果。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/deleteDataset/{id}"
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.delete(api_url, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"删除数据集 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"删除数据集时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def query_datasets(
|
||
token: str,
|
||
page: int = 0,
|
||
size: int = 20,
|
||
is_public: bool = None,
|
||
data_type: str = "",
|
||
is_hot_stone: bool = None
|
||
) -> str:
|
||
"""
|
||
查询数据集列表
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- page: 页码 (默认: 0)
|
||
- size: 每页数量 (默认: 20)
|
||
- is_public: 是否公开 (可选)
|
||
- data_type: 数据类型 (可选)
|
||
- is_hot_stone: 是否热门 (可选)
|
||
|
||
返回: 数据集列表。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/queryDatasets"
|
||
|
||
params = {
|
||
"page": page,
|
||
"size": size
|
||
}
|
||
|
||
if is_public is not None:
|
||
params["is_public"] = is_public
|
||
if data_type:
|
||
params["data_type"] = data_type
|
||
if is_hot_stone is not None:
|
||
params["is_hot_stone"] = is_hot_stone
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.get(api_url, params=params, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"查询数据集列表 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"查询数据集列表时发生未知错误: {str(e)}"
|
||
|
||
@mcp.tool()
|
||
async def get_dataset_detail(
|
||
token: str,
|
||
git_id: int,
|
||
owner: str,
|
||
name: str,
|
||
identifier: str,
|
||
is_public: bool
|
||
) -> str:
|
||
"""
|
||
查询数据集详情
|
||
|
||
参数说明:
|
||
- token: 访问令牌
|
||
- git_id: Git仓库ID
|
||
- owner: 所有者
|
||
- name: 数据集名称
|
||
- identifier: 数据集标识
|
||
- is_public: 是否公开
|
||
|
||
返回: 数据集详细信息。
|
||
"""
|
||
api_url = f"{API_BASE_URL}/api/mmp/newdataset/getDatasetDetail"
|
||
|
||
params = {
|
||
"git_id": git_id,
|
||
"owner": owner,
|
||
"name": name,
|
||
"identifier": identifier,
|
||
"is_public": is_public
|
||
}
|
||
|
||
headers = build_auth_headers(token)
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||
response = await client.get(api_url, params=params, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||
except httpx.HTTPStatusError as e:
|
||
return f"查询数据集详情 API 返回错误 ({e.response.status_code}): {e.response.text}"
|
||
except httpx.RequestError as e:
|
||
return f"无法连接数据集服务: {str(e)}"
|
||
except Exception as e:
|
||
return f"查询数据集详情时发生未知错误: {str(e)}"
|
||
|
||
if __name__ == "__main__":
|
||
mcp.run(transport=TRANSPORT_MODE) |