Skip to content

Commit c053d20

Browse files
committed
Notebook structure enhanced
1 parent e6e2196 commit c053d20

File tree

4 files changed

+197
-21
lines changed

4 files changed

+197
-21
lines changed

modules/notebook_module.go

Lines changed: 160 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package modules
22

33
import (
44
"context"
5+
"database/sql"
56
"errors"
67
"strconv"
78
"time"
@@ -104,28 +105,172 @@ func (m *NotebookModule) ListNotebooks(ctx context.Context, filters map[string]s
104105
return notebooks, nil
105106
}
106107

107-
// GetNotebookByID retrieves a single notebook by ID.
108+
// GetNotebookByID retrieves a single notebook by ID, including its cells, outputs, and evolution data.
108109
func (m *NotebookModule) GetNotebookByID(ctx context.Context, id string) (*models.Notebook, error) {
110+
notebookUUID, err := uuid.Parse(id)
111+
if err != nil {
112+
return nil, errors.New("invalid notebook ID format")
113+
}
114+
109115
query := `
110-
SELECT id, title, context_minio_url, problem_statement_id, created_at, last_modified_at
111-
FROM notebooks
112-
WHERE id = $1;
116+
SELECT
117+
n.id, n.title, n.context_minio_url, n.problem_statement_id, n.created_at, n.last_modified_at,
118+
c.id, c.notebook_id, c.cell_index, c.cell_type, c.source, c.execution_count,
119+
co.id, co.cell_id, co.output_index, co.type, co.data_json, co.minio_url, co.execution_count,
120+
er.id, er.source_cell_id, er.start_time, er.end_time, er.status,
121+
cv.id, cv.evolution_run_id, cv.code, cv.metric, cv.is_best, cv.generation, cv.parent_variant_id
122+
FROM
123+
notebooks n
124+
LEFT JOIN
125+
cells c ON n.id = c.notebook_id
126+
LEFT JOIN
127+
cell_outputs co ON c.id = co.cell_id
128+
LEFT JOIN
129+
evolution_runs er ON c.id = er.source_cell_id
130+
LEFT JOIN
131+
cell_variations cv ON er.id = cv.evolution_run_id
132+
WHERE
133+
n.id = $1
134+
ORDER BY
135+
c.cell_index ASC,
136+
co.output_index ASC,
137+
er.start_time ASC,
138+
cv.generation ASC;
113139
`
114-
row := db.Pool.QueryRow(ctx, query, id)
115140

116-
var nb models.Notebook
117-
if err := row.Scan(
118-
&nb.ID,
119-
&nb.Title,
120-
&nb.ContextMinioURL,
121-
&nb.ProblemStatementID,
122-
&nb.CreatedAt,
123-
&nb.LastModifiedAt,
124-
); err != nil {
141+
rows, err := db.Pool.Query(ctx, query, notebookUUID)
142+
if err != nil {
125143
return nil, err
126144
}
145+
defer rows.Close()
127146

128-
return &nb, nil
147+
var notebook models.Notebook
148+
cellMap := make(map[uuid.UUID]*models.Cell)
149+
runMap := make(map[uuid.UUID]*models.EvolutionRun)
150+
var orderedCellIDs []uuid.UUID
151+
152+
for rows.Next() {
153+
var (
154+
// Notebook fields are scanned into the notebook struct directly
155+
// Nullable types for fields from LEFT JOINs
156+
cellID uuid.NullUUID
157+
cellNotebookID uuid.NullUUID
158+
cellIndex sql.NullInt32
159+
cellType sql.NullString
160+
cellSource sql.NullString
161+
cellExecCount sql.NullInt32
162+
outputID uuid.NullUUID
163+
outputCellID uuid.NullUUID
164+
outputIndex sql.NullInt32
165+
outputType sql.NullString
166+
outputDataJSON []byte
167+
outputMinioURL sql.NullString
168+
outputExecCount sql.NullInt32
169+
erID uuid.NullUUID
170+
erSourceCellID uuid.NullUUID
171+
erStartTime sql.NullTime
172+
erEndTime sql.NullTime
173+
erStatus sql.NullString
174+
cvID uuid.NullUUID
175+
cvEvolutionRunID uuid.NullUUID
176+
cvCode sql.NullString
177+
cvMetric sql.NullFloat64
178+
cvIsBest sql.NullBool
179+
cvGeneration sql.NullInt32
180+
cvParentVariantID uuid.NullUUID
181+
)
182+
183+
if err := rows.Scan(
184+
&notebook.ID, &notebook.Title, &notebook.ContextMinioURL, &notebook.ProblemStatementID, &notebook.CreatedAt, &notebook.LastModifiedAt,
185+
&cellID, &cellNotebookID, &cellIndex, &cellType, &cellSource, &cellExecCount,
186+
&outputID, &outputCellID, &outputIndex, &outputType, &outputDataJSON, &outputMinioURL, &outputExecCount,
187+
&erID, &erSourceCellID, &erStartTime, &erEndTime, &erStatus,
188+
&cvID, &cvEvolutionRunID, &cvCode, &cvMetric, &cvIsBest, &cvGeneration, &cvParentVariantID,
189+
); err != nil {
190+
return nil, err
191+
}
192+
193+
if cellID.Valid {
194+
if _, exists := cellMap[cellID.UUID]; !exists {
195+
cellMap[cellID.UUID] = &models.Cell{
196+
ID: cellID.UUID,
197+
NotebookID: cellNotebookID.UUID,
198+
CellIndex: int(cellIndex.Int32),
199+
CellType: cellType.String,
200+
Source: cellSource.String,
201+
ExecutionCount: int(cellExecCount.Int32),
202+
Outputs: []models.CellOutput{},
203+
EvolutionRuns: []models.EvolutionRun{},
204+
}
205+
orderedCellIDs = append(orderedCellIDs, cellID.UUID)
206+
}
207+
}
208+
209+
if outputID.Valid {
210+
if cell, exists := cellMap[outputCellID.UUID]; exists {
211+
cell.Outputs = append(cell.Outputs, models.CellOutput{
212+
ID: outputID.UUID,
213+
CellID: outputCellID.UUID,
214+
OutputIndex: int(outputIndex.Int32),
215+
Type: outputType.String,
216+
DataJSON: outputDataJSON,
217+
MinioURL: outputMinioURL.String,
218+
ExecutionCount: int(outputExecCount.Int32),
219+
})
220+
}
221+
}
222+
223+
if erID.Valid {
224+
if _, exists := runMap[erID.UUID]; !exists {
225+
run := models.EvolutionRun{
226+
ID: erID.UUID,
227+
SourceCellID: erSourceCellID.UUID,
228+
StartTime: erStartTime.Time,
229+
Status: erStatus.String,
230+
Variations: []models.CellVariation{},
231+
}
232+
if erEndTime.Valid {
233+
run.EndTime = &erEndTime.Time
234+
}
235+
runMap[erID.UUID] = &run
236+
}
237+
}
238+
239+
if cvID.Valid {
240+
if run, exists := runMap[cvEvolutionRunID.UUID]; exists {
241+
variation := models.CellVariation{
242+
ID: cvID.UUID,
243+
EvolutionRunID: cvEvolutionRunID.UUID,
244+
Code: cvCode.String,
245+
Metric: cvMetric.Float64,
246+
IsBest: cvIsBest.Bool,
247+
Generation: int(cvGeneration.Int32),
248+
}
249+
if cvParentVariantID.Valid {
250+
variation.ParentVariantID = &cvParentVariantID.UUID
251+
}
252+
run.Variations = append(run.Variations, variation)
253+
}
254+
}
255+
}
256+
257+
if notebook.ID == "" {
258+
return nil, errors.New("notebook not found")
259+
}
260+
261+
// Assemble the nested structure
262+
for _, run := range runMap {
263+
if cell, exists := cellMap[run.SourceCellID]; exists {
264+
cell.EvolutionRuns = append(cell.EvolutionRuns, *run)
265+
}
266+
}
267+
268+
notebook.Cells = make([]models.Cell, len(orderedCellIDs))
269+
for i, cellID := range orderedCellIDs {
270+
notebook.Cells[i] = *cellMap[cellID]
271+
}
272+
273+
return &notebook, nil
129274
}
130275

131276
// UpdateNotebook updates a notebook’s title or context_minio_url.

pkg/models/cell_model.go

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,14 @@ import (
88

99
// Cell represents a single cell within a notebook.
1010
type Cell struct {
11-
ID uuid.UUID `json:"id"`
12-
NotebookID uuid.UUID `json:"notebook_id"`
13-
CellIndex int `json:"cell_index"`
14-
CellType string `json:"cell_type"`
15-
Source string `json:"source"`
16-
ExecutionCount int `json:"execution_count"`
11+
ID uuid.UUID `json:"id"`
12+
NotebookID uuid.UUID `json:"notebook_id"`
13+
CellIndex int `json:"cell_index"`
14+
CellType string `json:"cell_type"`
15+
Source string `json:"source"`
16+
ExecutionCount int `json:"execution_count"`
17+
Outputs []CellOutput `json:"outputs,omitempty"`
18+
EvolutionRuns []EvolutionRun `json:"evolution_runs,omitempty"`
1719
}
1820

1921
// CellOutput represents the output of a cell execution.

pkg/models/evolution_model.go

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package models
2+
3+
import (
4+
"time"
5+
6+
"github.com/google/uuid"
7+
)
8+
9+
// EvolutionRun represents an evolution run for a cell.
10+
type EvolutionRun struct {
11+
ID uuid.UUID `json:"id"`
12+
SourceCellID uuid.UUID `json:"source_cell_id"`
13+
StartTime time.Time `json:"start_time"`
14+
EndTime *time.Time `json:"end_time,omitempty"`
15+
Status string `json:"status"`
16+
Variations []CellVariation `json:"variations,omitempty"`
17+
}
18+
19+
// CellVariation represents a single variation within an evolution run.
20+
type CellVariation struct {
21+
ID uuid.UUID `json:"id"`
22+
EvolutionRunID uuid.UUID `json:"evolution_run_id"`
23+
Code string `json:"code"`
24+
Metric float64 `json:"metric"`
25+
IsBest bool `json:"is_best"`
26+
Generation int `json:"generation"`
27+
ParentVariantID *uuid.UUID `json:"parent_variant_id,omitempty"`
28+
}

pkg/models/notebook_model.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ type Notebook struct {
1010
ProblemStatementID *string `json:"problem_statement_id,omitempty"`
1111
CreatedAt time.Time `json:"created_at"`
1212
LastModifiedAt time.Time `json:"last_modified_at"`
13+
Cells []Cell `json:"cells,omitempty"`
1314
}
1415

1516
// CreateNotebookRequest is the payload to create a notebook.

0 commit comments

Comments
 (0)