From 2a823ee25cad4da1c2217820fd953cdff0a37903 Mon Sep 17 00:00:00 2001
From: yaoyifan-yyf <yaoyifan.yyf@antgroup.com>
Date: Thu, 16 Oct 2025 14:56:25 +0800
Subject: [PATCH] feat: support multi benchmark datasets

---
 .../src/dbgpt_serve/evaluate/api/endpoints.py | 47 +++++++++++++------
 1 file changed, 32 insertions(+), 15 deletions(-)

diff --git a/packages/dbgpt-serve/src/dbgpt_serve/evaluate/api/endpoints.py b/packages/dbgpt-serve/src/dbgpt_serve/evaluate/api/endpoints.py
index 9c6915ead..4567ac9fc 100644
--- a/packages/dbgpt-serve/src/dbgpt_serve/evaluate/api/endpoints.py
+++ b/packages/dbgpt-serve/src/dbgpt_serve/evaluate/api/endpoints.py
@@ -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(