From 6926e02fbf5437a16354411788978268e4da13ee Mon Sep 17 00:00:00 2001
From: Geoff Wilson <geoff@gr-oss.io>
Date: Thu, 1 Aug 2024 15:13:51 -0400
Subject: [PATCH] Fix parsing request from aim UI

---
 pkg/api/aim/api/request/run.go  | 14 +++++++-------
 pkg/api/aim/api/response/run.go | 16 ++++++++++------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/pkg/api/aim/api/request/run.go b/pkg/api/aim/api/request/run.go
index 5c13119be..896bead75 100644
--- a/pkg/api/aim/api/request/run.go
+++ b/pkg/api/aim/api/request/run.go
@@ -101,13 +101,13 @@ type SearchAlignedMetricsRequest struct {
 // SearchArtifactsRequest is a request struct for `POST /runs/search/image` endpoint.
 type SearchArtifactsRequest struct {
 	BaseSearchRequest
-	Query         string `params:"q"`
-	SkipSystem    bool   `params:"skip_system"`
-	RecordDensity int    `params:"record_density"`
-	IndexDensity  int    `params:"index_density"`
-	RecordRange   string `params:"record_range"`
-	IndexRange    string `params:"index_range"`
-	CalcRanges    bool   `params:"calc_ranges"`
+	Query         string `json:"q"`
+	SkipSystem    bool   `json:"skip_system"`
+	RecordDensity any    `json:"record_density"`
+	IndexDensity  any    `json:"index_density"`
+	RecordRange   string `json:"record_range"`
+	IndexRange    string `json:"index_range"`
+	CalcRanges    bool   `json:"calc_ranges"`
 }
 
 // DeleteRunRequest is a request struct for `DELETE /runs/:id` endpoint.
diff --git a/pkg/api/aim/api/response/run.go b/pkg/api/aim/api/response/run.go
index 774083ea6..3c5b3a797 100644
--- a/pkg/api/aim/api/response/run.go
+++ b/pkg/api/aim/api/response/run.go
@@ -592,7 +592,7 @@ func NewStreamArtifactsResponse(ctx *fiber.Ctx, rows *sql.Rows, runs map[string]
 					"width":    img.Width,
 					"format":   img.Format,
 					"iter":     img.Iter,
-					"index":    img.Index,
+					"index":    imgIndex, // calculated
 					"step":     img.Step,
 				}
 
@@ -632,17 +632,21 @@ func NewStreamArtifactsResponse(ctx *fiber.Ctx, rows *sql.Rows, runs map[string]
 				return w.Flush()
 			}
 			includeStep := func(img models.Artifact) bool {
-				if req.RecordDensity == 0 {
+				switch v := req.RecordDensity.(type) {
+				case int:
+					return img.Step%int64(v) == 0
+				default:
 					return true
 				}
-				return img.Step%int64(req.RecordDensity) == 0
 			}
 			includeIndex := func(img models.Artifact) bool {
-				if req.IndexDensity == 0 {
+				switch v := req.IndexDensity.(type) {
+				case int:
+					interval := summary.StepImageCount(img.RunID, img.Name, int(img.Step)) / v
+					return imgIndex%interval == 0
+				default:
 					return true
 				}
-				interval := summary.StepImageCount(img.RunID, img.Name, int(img.Step)) / req.IndexDensity
-				return imgIndex%interval == 0
 			}
 			hasRows := false
 			for rows.Next() {