add Add Elements tool to dataset sdk

Signed-off-by: Grant Linville <grant@acorn.io>

Author: g-linville

pkg/sdkserver/datasets.go

@@ -12,25 +12,21 @@ import (
 
 type datasetRequest struct {
 	Input           string `json:"input"`
-	Workspace       string `json:"workspace"`
 	DatasetToolRepo string `json:"datasetToolRepo"`
 }
 
 func (r datasetRequest) validate(requireInput bool) error {
-	if r.Workspace == "" {
-		return fmt.Errorf("workspace is required")
-	} else if requireInput && r.Input == "" {
+	if requireInput && r.Input == "" {
 		return fmt.Errorf("input is required")
 	}
 	return nil
 }
 
 func (r datasetRequest) opts(o gptscript.Options) gptscript.Options {
 	opts := gptscript.Options{
-		Cache:     o.Cache,
-		Monitor:   o.Monitor,
-		Runner:    o.Runner,
-		Workspace: r.Workspace,
+		Cache:   o.Cache,
+		Monitor: o.Monitor,
+		Runner:  o.Runner,
 	}
 	return opts
 }
@@ -209,6 +205,79 @@ func (s *server) addDatasetElement(w http.ResponseWriter, r *http.Request) {
 	writeResponse(logger, w, map[string]any{"stdout": result})
 }
 
+type addDatasetElementsArgs struct {
+	DatasetID string `json:"datasetID"`
+	Elements  []struct {
+		Name        string `json:"name"`
+		Description string `json:"description"`
+		Contents    string `json:"contents"`
+	}
+}
+
+func (a addDatasetElementsArgs) validate() error {
+	if a.DatasetID == "" {
+		return fmt.Errorf("datasetID is required")
+	}
+	if len(a.Elements) == 0 {
+		return fmt.Errorf("elements is required")
+	}
+	return nil
+}
+
+func (s *server) addDatasetElements(w http.ResponseWriter, r *http.Request) {
+	logger := gcontext.GetLogger(r.Context())
+
+	var req datasetRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to decode request body: %w", err))
+		return
+	}
+
+	if err := req.validate(true); err != nil {
+		writeError(logger, w, http.StatusBadRequest, err)
+		return
+	}
+
+	g, err := gptscript.New(r.Context(), req.opts(s.gptscriptOpts))
+	if err != nil {
+		writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to initialize gptscript: %w", err))
+		return
+	}
+
+	var args addDatasetElementsArgs
+	if err := json.Unmarshal([]byte(req.Input), &args); err != nil {
+		writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to unmarshal input: %w", err))
+		return
+	}
+
+	if err := args.validate(); err != nil {
+		writeError(logger, w, http.StatusBadRequest, err)
+		return
+	}
+
+	prg, err := loader.Program(r.Context(), req.getToolRepo(), "Add Elements", loader.Options{
+		Cache: g.Cache,
+	})
+	if err != nil {
+		writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
+		return
+	}
+
+	elementsJSON, err := json.Marshal(args.Elements)
+	if err != nil {
+		writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to marshal elements: %w", err))
+		return
+	}
+
+	result, err := g.Run(r.Context(), prg, s.gptscriptOpts.Env, fmt.Sprintf(`{"datasetID":%q, elements:%s}`, args.DatasetID, elementsJSON))
+	if err != nil {
+		writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to run program: %w", err))
+		return
+	}
+
+	writeResponse(logger, w, map[string]any{"stdout": result})
+}
+
 type listDatasetElementsArgs struct {
 	DatasetID string `json:"datasetID"`
 }

pkg/sdkserver/routes.go

@@ -73,6 +73,7 @@ func (s *server) addRoutes(mux *http.ServeMux) {
 	mux.HandleFunc("POST /datasets/list-elements", s.listDatasetElements)
 	mux.HandleFunc("POST /datasets/get-element", s.getDatasetElement)
 	mux.HandleFunc("POST /datasets/add-element", s.addDatasetElement)
+	mux.HandleFunc("POST /datasets/add-elements", s.addDatasetElements)
 
 	mux.HandleFunc("POST /workspaces/create", s.createWorkspace)
 	mux.HandleFunc("POST /workspaces/delete", s.deleteWorkspace)