github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2026-01-13 19:55:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support multi benchmark datasets

Browse Source
This commit is contained in:
yaoyifan-yyf
2025-10-16 14:56:25 +08:00
parent ba80df5485
commit 2a823ee25c
1 changed files with 32 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
packages/dbgpt-serve/src/dbgpt_serve/evaluate/api/endpoints.py
View File

@@ -298,32 +298,49 @@ async def benchmark_task_list(
)
)
@router.get("/benchmark/datasets", dependencies=[Depends(check_api_key)])
@router.get("/benchmark/list_datasets", dependencies=[Depends(check_api_key)])
async def list_benchmark_datasets():
manager = get_benchmark_manager(global_system_app)
info = await manager.get_table_info()
result = [
{
"name": name,
"rowCount": meta.get("row_count", 0),
"columns": meta.get("columns", []),
"dataset_id": "1",
"name": "Falcon",
"tableCount": len(info.items())
}
for name, meta in info.items()
]
return Result.succ(result)
@router.get("/benchmark/dataset/{dataset_id}", dependencies=[Depends(check_api_key)])
async def list_benchmark_dataset_tables(dataset_id: str, limit: int = 200, offset: int = 0):
if dataset_id == "1":
manager = get_benchmark_manager(global_system_app)
info = await manager.get_table_info()
result = [
{
"name": name,
"rowCount": meta.get("row_count", 0),
"columns": meta.get("columns", []),
}
for name, meta in info.items()
]
return Result.succ(result)
else:
return Result.succ("dataset not found")
@router.get("/benchmark/datasets/{table}/rows", dependencies=[Depends(check_api_key)])
async def get_benchmark_table_rows(table: str, limit: int = 10):
manager = get_benchmark_manager(global_system_app)
info = await manager.get_table_info()
if table not in info:
raise HTTPException(status_code=404, detail=f"table '{table}' not found")
sql = f'SELECT * FROM "{table}" LIMIT :limit'
rows = await manager.query(sql, {"limit": limit})
return Result.succ({"table": table, "limit": limit, "rows": rows})
@router.get("/benchmark/dataset/{dataset_id}/{table}/rows", dependencies=[Depends(check_api_key)])
async def get_benchmark_table_rows(dataset_id:str, table: str, limit: int = 10):
if dataset_id == "1":
manager = get_benchmark_manager(global_system_app)
info = await manager.get_table_info()
if table not in info:
raise HTTPException(status_code=404, detail=f"table '{table}' not found")
sql = f'SELECT * FROM "{table}" LIMIT :limit'
rows = await manager.query(sql, {"limit": limit})
return Result.succ({"table": table, "limit": limit, "rows": rows})
else:
return Result.succ("dataset not found")
@router.get("/benchmark_result_download", dependencies=[Depends(check_api_key)])
async def download_benchmark_result(