Skip to content

Commit aca5591

Browse files
kerthcetkerthcet
authored
Optimize dataset page (#199)
* update the graphql api for datasets Signed-off-by: kerthcet <kerthcet@gmail.com> * Add page for datasets Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent 843afc7 commit aca5591

24 files changed

Lines changed: 1490 additions & 770 deletions

‎alphatrion/server/graphql/resolvers.py‎

Lines changed: 74 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
from alphatrion import envs
99
from alphatrion.artifact import artifact
10+
from alphatrion.server.graphql.types import ArtifactFile
1011
from alphatrion.storage import runtime
1112
from alphatrion.storage.sql_models import (
1213
FINISHED_STATUS,
@@ -354,9 +355,52 @@ async def list_artifact_tags(
354355
arf = artifact.Artifact(team_id=team_id, insecure=True)
355356
return [ArtifactTag(name=tag) for tag in arf.list_versions(repo_name)]
356357

358+
@staticmethod
359+
async def list_artifact_files(
360+
team_id: str, tag: str, repo_name: str
361+
) -> list[ArtifactFile]:
362+
"""List files in an artifact without loading content."""
363+
364+
try:
365+
arf = artifact.Artifact(team_id=team_id, insecure=True)
366+
file_paths = arf.pull(repo_name=repo_name, version=tag)
367+
368+
if not file_paths:
369+
return []
370+
371+
files = []
372+
for file_path in file_paths:
373+
filename = os.path.basename(file_path)
374+
file_size = os.path.getsize(file_path)
375+
376+
# Determine content type based on file extension
377+
if filename.endswith(".json"):
378+
content_type = "application/json"
379+
elif (
380+
filename.endswith(".txt")
381+
or filename.endswith(".log")
382+
or filename.endswith((".py", ".js", ".ts", ".tsx", ".jsx"))
383+
):
384+
content_type = "text/plain"
385+
else:
386+
content_type = "text/plain"
387+
388+
files.append(
389+
ArtifactFile(
390+
filename=filename, size=file_size, content_type=content_type
391+
)
392+
)
393+
394+
return files
395+
except Exception as e:
396+
raise RuntimeError(f"Failed to list artifact files: {e}") from e
397+
357398
@staticmethod
358399
async def get_artifact_content(
359-
team_id: str, tag: str, repo_name: str | None = None
400+
team_id: str,
401+
tag: str,
402+
repo_name: str | None = None,
403+
filename: str | None = None,
360404
) -> ArtifactContent:
361405
"""Get artifact content from registry."""
362406
try:
@@ -365,33 +409,44 @@ async def get_artifact_content(
365409

366410
# Pull the artifact - ORAS will manage temp directory
367411
# Returns absolute paths to files in ORAS temp directory
368-
# Note: One potential issue is if we download too many large files,
369-
# it may fill up disk space. For now we assume artifacts are
370-
# reasonably sized and/or users will manage their registry storage.
371412
file_paths = arf.pull(repo_name=repo_name, version=tag)
372413

373414
if not file_paths:
374415
raise RuntimeError("No files found in artifact")
375416

376-
# Read first file content (file_paths now contains absolute paths)
377-
file_path = file_paths[0]
417+
# Find the requested file or use first file
418+
file_path = None
419+
if filename:
420+
for path in file_paths:
421+
if os.path.basename(path) == filename:
422+
file_path = path
423+
break
424+
if not file_path:
425+
raise RuntimeError(f"File '{filename}' not found in artifact")
426+
else:
427+
file_path = file_paths[0]
428+
429+
# Read file content
378430
with open(file_path, encoding="utf-8") as f:
379431
content = f.read()
380432

381433
# Get filename from path
382-
filename = os.path.basename(file_path)
434+
actual_filename = os.path.basename(file_path)
383435

384436
# Determine content type based on file extension
385-
# TODO: for multiple files, this is not right.
386-
if filename.endswith(".json"):
437+
if actual_filename.endswith(".json"):
387438
content_type = "application/json"
388-
elif filename.endswith(".txt") or filename.endswith(".log"):
439+
elif (
440+
actual_filename.endswith(".txt")
441+
or actual_filename.endswith(".log")
442+
or actual_filename.endswith((".py", ".js", ".ts", ".tsx", ".jsx"))
443+
):
389444
content_type = "text/plain"
390445
else:
391446
content_type = "text/plain"
392447

393448
return ArtifactContent(
394-
filename=filename, content=content, content_type=content_type
449+
filename=actual_filename, content=content, content_type=content_type
395450
)
396451
except Exception as e:
397452
raise RuntimeError(f"Failed to get artifact content: {e}") from e
@@ -635,14 +690,18 @@ def get_daily_token_usage(
635690
@staticmethod
636691
def list_datasets(
637692
team_id: strawberry.ID,
693+
experiment_id: strawberry.ID | None = None,
694+
run_id: strawberry.ID | None = None,
638695
page: int = 0,
639696
page_size: int = 20,
640697
order_by: str = "created_at",
641698
order_desc: bool = True,
642699
) -> list[Dataset]:
643700
metadb = runtime.storage_runtime().metadb
644-
datasets = metadb.list_datasets_by_team_id(
645-
team_id=uuid.UUID(team_id),
701+
datasets = metadb.list_datasets(
702+
team_id=team_id,
703+
experiment_id=experiment_id,
704+
run_id=run_id,
646705
page=page,
647706
page_size=page_size,
648707
order_by=order_by,
@@ -653,6 +712,7 @@ def list_datasets(
653712
id=d.uuid,
654713
name=d.name,
655714
description=d.description,
715+
path=d.path,
656716
meta=d.meta,
657717
team_id=d.team_id,
658718
experiment_id=d.experiment_id,
@@ -673,6 +733,7 @@ def get_dataset(id: strawberry.ID) -> Dataset | None:
673733
id=dataset.uuid,
674734
name=dataset.name,
675735
description=dataset.description,
736+
path=dataset.path,
676737
meta=dataset.meta,
677738
team_id=dataset.team_id,
678739
experiment_id=dataset.experiment_id,
@@ -683,38 +744,6 @@ def get_dataset(id: strawberry.ID) -> Dataset | None:
683744
)
684745
return None
685746

686-
@staticmethod
687-
def list_datasets_by_experiment(
688-
experiment_id: strawberry.ID,
689-
page: int = 0,
690-
page_size: int = 20,
691-
order_by: str = "created_at",
692-
order_desc: bool = True,
693-
) -> list[Dataset]:
694-
metadb = runtime.storage_runtime().metadb
695-
datasets = metadb.list_datasets_by_exp_id(
696-
exp_id=uuid.UUID(experiment_id),
697-
page=page,
698-
page_size=page_size,
699-
order_by=order_by,
700-
order_desc=order_desc,
701-
)
702-
return [
703-
Dataset(
704-
id=d.uuid,
705-
name=d.name,
706-
description=d.description,
707-
meta=d.meta,
708-
team_id=d.team_id,
709-
experiment_id=d.experiment_id,
710-
run_id=d.run_id,
711-
user_id=d.user_id,
712-
created_at=d.created_at,
713-
updated_at=d.updated_at,
714-
)
715-
for d in datasets
716-
]
717-
718747

719748
class GraphQLMutations:
720749
@staticmethod

‎alphatrion/server/graphql/schema.py‎

Lines changed: 18 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from alphatrion.server.graphql.types import (
55
AddUserToTeamInput,
66
ArtifactContent,
7+
ArtifactFile,
78
ArtifactRepository,
89
ArtifactTag,
910
CreateTeamInput,
@@ -95,52 +96,50 @@ async def artifact_tags(
9596
) -> list[ArtifactTag]:
9697
return await GraphQLResolvers.list_artifact_tags(str(team_id), repo_name)
9798

99+
@strawberry.field
100+
async def artifact_files(
101+
self,
102+
team_id: strawberry.ID,
103+
tag: str,
104+
repo_name: str,
105+
) -> list[ArtifactFile]:
106+
return await GraphQLResolvers.list_artifact_files(str(team_id), tag, repo_name)
107+
98108
@strawberry.field
99109
async def artifact_content(
100110
self,
101111
team_id: strawberry.ID,
102112
tag: str,
103113
repo_name: str,
114+
filename: str | None = None,
104115
) -> ArtifactContent:
105-
return await GraphQLResolvers.get_artifact_content(str(team_id), tag, repo_name)
116+
return await GraphQLResolvers.get_artifact_content(
117+
str(team_id), tag, repo_name, filename
118+
)
106119

107120
# Dataset queries
108121
@strawberry.field
109122
def datasets(
110123
self,
111124
team_id: strawberry.ID,
125+
experiment_id: strawberry.ID | None = None,
126+
run_id: strawberry.ID | None = None,
112127
page: int = 0,
113128
page_size: int = 20,
114129
order_by: str = "created_at",
115130
order_desc: bool = True,
116131
) -> list[Dataset]:
117132
return GraphQLResolvers.list_datasets(
118133
team_id=team_id,
119-
page=page,
120-
page_size=page_size,
121-
order_by=order_by,
122-
order_desc=order_desc,
123-
)
124-
125-
dataset: Dataset | None = strawberry.field(resolver=GraphQLResolvers.get_dataset)
126-
127-
@strawberry.field
128-
def datasets_by_experiment(
129-
self,
130-
experiment_id: strawberry.ID,
131-
page: int = 0,
132-
page_size: int = 20,
133-
order_by: str = "created_at",
134-
order_desc: bool = True,
135-
) -> list[Dataset]:
136-
return GraphQLResolvers.list_datasets_by_experiment(
137134
experiment_id=experiment_id,
135+
run_id=run_id,
138136
page=page,
139137
page_size=page_size,
140138
order_by=order_by,
141139
order_desc=order_desc,
142140
)
143141

142+
dataset: Dataset | None = strawberry.field(resolver=GraphQLResolvers.get_dataset)
144143

145144
@strawberry.type
146145
class Mutation:

‎alphatrion/server/graphql/types.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ class Dataset:
208208
id: strawberry.ID
209209
name: str
210210
description: str | None
211+
path: str
211212
meta: JSON | None
212213
team_id: strawberry.ID
213214
experiment_id: strawberry.ID | None
@@ -252,7 +253,6 @@ class RemoveUserFromTeamInput:
252253
user_id: strawberry.ID
253254
team_id: strawberry.ID
254255

255-
256256
# Artifact types
257257
@strawberry.type
258258
class ArtifactRepository:
@@ -264,6 +264,13 @@ class ArtifactTag:
264264
name: str
265265

266266

267+
@strawberry.type
268+
class ArtifactFile:
269+
filename: str
270+
size: int
271+
content_type: str
272+
273+
267274
@strawberry.type
268275
class ArtifactContent:
269276
filename: str

‎dashboard/src/App.tsx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { ExperimentDetailPage } from './pages/experiments/[id]';
1212
import { ExperimentComparePage } from './pages/experiments/compare';
1313
import { RunsPage } from './pages/runs';
1414
import { RunDetailPage } from './pages/runs/[id]';
15+
import { DatasetsPage } from './pages/datasets';
1516
import { ArtifactsPage } from './pages/artifacts';
1617
import type { Team } from './types';
1718

@@ -141,6 +142,7 @@ function App() {
141142
<Route index element={<RunsPage />} />
142143
<Route path=":id" element={<RunDetailPage />} />
143144
</Route>
145+
<Route path="datasets" element={<DatasetsPage />} />
144146
<Route path="artifacts" element={<ArtifactsPage />} />
145147
</Route>
146148
</Routes>

‎dashboard/src/components/dashboard/daily-token-usage-chart.tsx‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -80,27 +80,27 @@ export function DailyTokenUsageChart({ data, timeRange }: DailyTokenUsageChartPr
8080
}, [chartData]);
8181

8282
return (
83-
<div className="space-y-2">
83+
<div className="space-y-1.5">
8484
<div className="flex items-center justify-between">
85-
<h3 className="text-sm font-semibold">Token Usage</h3>
86-
<div className="text-xs text-muted-foreground">
85+
<h3 className="text-xs font-semibold">Token Usage</h3>
86+
<div className="text-[10px] text-muted-foreground">
8787
Total: {totalTokens.toLocaleString()} ({totalInputTokens.toLocaleString()}↓ {totalOutputTokens.toLocaleString()}↑)
8888
</div>
8989
</div>
9090

91-
<ResponsiveContainer width="100%" height={240}>
92-
<LineChart data={chartData} margin={{ left: 10, right: 15, top: 10, bottom: 5 }}>
91+
<ResponsiveContainer width="100%" height={200}>
92+
<LineChart data={chartData} margin={{ left: 5, right: 10, top: 5, bottom: 0 }}>
9393
<CartesianGrid strokeDasharray="3 3" stroke="#e2e8f0" opacity={0.5} />
9494
<XAxis
9595
dataKey="displayDate"
96-
tick={{ fontSize: 10 }}
96+
tick={{ fontSize: 9 }}
9797
angle={-45}
9898
textAnchor="end"
99-
height={50}
99+
height={40}
100100
/>
101101
<YAxis
102-
tick={{ fontSize: 10 }}
103-
width={50}
102+
tick={{ fontSize: 9 }}
103+
width={35}
104104
tickFormatter={(value) =>
105105
value >= 1000000
106106
? `${(value / 1000000).toFixed(1)}M`
@@ -112,8 +112,8 @@ export function DailyTokenUsageChart({ data, timeRange }: DailyTokenUsageChartPr
112112
value: 'Tokens',
113113
angle: -90,
114114
position: 'insideLeft',
115-
offset: -5,
116-
style: { textAnchor: 'middle', fontSize: 11 }
115+
offset: 0,
116+
style: { textAnchor: 'middle', fontSize: 9 }
117117
}}
118118
/>
119119
<Tooltip
@@ -151,11 +151,11 @@ export function DailyTokenUsageChart({ data, timeRange }: DailyTokenUsageChartPr
151151
}}
152152
/>
153153
<Legend
154-
wrapperStyle={{ fontSize: '11px', paddingTop: '2px' }}
154+
wrapperStyle={{ fontSize: '10px', paddingTop: '2px' }}
155155
iconType="circle"
156-
iconSize={8}
156+
iconSize={6}
157157
verticalAlign="bottom"
158-
height={25}
158+
height={20}
159159
/>
160160
<Line
161161
type="monotone"

0 commit comments

Comments
 (0)