Skip to content

Commit

Permalink
chore: sdkserver: update dataset methods for the rewrite
Browse files Browse the repository at this point in the history 
Signed-off-by: Grant Linville <[email protected]>
  • Loading branch information
@g-linville
g-linville committed Nov 6, 2024
1 parent d21c001 commit d1d6179
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 165 deletions.
186 changes: 23 additions & 163 deletions pkg/sdkserver/datasets.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,24 +11,21 @@ import (
)

func (s *server) getDatasetTool(req datasetRequest) string {
if req.DatasetToolRepo != "" {
return req.DatasetToolRepo
if req.DatasetTool != "" {
return req.DatasetTool
}

return s.datasetTool
}

type datasetRequest struct {
Input string `json:"input"`
WorkspaceID string `json:"workspaceID"`
DatasetToolRepo string `json:"datasetToolRepo"`
Env []string `json:"env"`
Input string `json:"input"`
DatasetTool string `json:"datasetTool"`
Env []string `json:"env"`
}

func (r datasetRequest) validate(requireInput bool) error {
if r.WorkspaceID == "" {
return fmt.Errorf("workspaceID is required")
} else if requireInput && r.Input == "" {
if requireInput && r.Input == "" {
return fmt.Errorf("input is required")
} else if len(r.Env) == 0 {
return fmt.Errorf("env is required")
Expand All @@ -38,10 +35,9 @@ func (r datasetRequest) validate(requireInput bool) error {

func (r datasetRequest) opts(o gptscript.Options) gptscript.Options {
opts := gptscript.Options{
Cache: o.Cache,
Monitor: o.Monitor,
Runner: o.Runner,
Workspace: r.WorkspaceID,
Cache: o.Cache,
Monitor: o.Monitor,
Runner: o.Runner,
}
return opts
}
Expand Down Expand Up @@ -84,148 +80,19 @@ func (s *server) listDatasets(w http.ResponseWriter, r *http.Request) {
writeResponse(logger, w, map[string]any{"stdout": result})
}

type createDatasetArgs struct {
Name string `json:"datasetName"`
Description string `json:"datasetDescription"`
}

func (a createDatasetArgs) validate() error {
if a.Name == "" {
return fmt.Errorf("datasetName is required")
}
return nil
}

func (s *server) createDataset(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 createDatasetArgs
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(), s.getDatasetTool(req), "Create Dataset", loader.Options{
Cache: g.Cache,
})

if err != nil {
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
return
}

result, err := g.Run(r.Context(), prg, req.Env, req.Input)
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 addDatasetElementArgs struct {
DatasetID string `json:"datasetID"`
ElementName string `json:"elementName"`
ElementDescription string `json:"elementDescription"`
ElementContent string `json:"elementContent"`
}

func (a addDatasetElementArgs) validate() error {
if a.DatasetID == "" {
return fmt.Errorf("datasetID is required")
}
if a.ElementName == "" {
return fmt.Errorf("elementName is required")
}
if a.ElementContent == "" {
return fmt.Errorf("elementContent is required")
}
return nil
}

func (s *server) addDatasetElement(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 addDatasetElementArgs
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(), s.getDatasetTool(req), "Add Element", loader.Options{
Cache: g.Cache,
})
if err != nil {
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
return
}

result, err := g.Run(r.Context(), prg, req.Env, req.Input)
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 addDatasetElementsArgs struct {
DatasetID string `json:"datasetID"`
Elements []struct {
Name string `json:"name"`
Description string `json:"description"`
Contents string `json:"contents"`
}
DatasetID string `json:"datasetID"`
Name string `json:"name"`
Description string `json:"description"`
Elements []struct {
Name string `json:"name"`
Description string `json:"description"`
Contents string `json:"contents"`
BinaryContents []byte `json:"binaryContents"`
} `json:"elements"`
}

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")
}
Expand Down Expand Up @@ -271,13 +138,7 @@ func (s *server) addDatasetElements(w http.ResponseWriter, r *http.Request) {
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, req.Env, fmt.Sprintf(`{"datasetID":%q, "elements":%q}`, args.DatasetID, string(elementsJSON)))
result, err := g.Run(r.Context(), prg, req.Env, req.Input)
if err != nil {
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to run program: %w", err))
return
Expand Down Expand Up @@ -347,15 +208,14 @@ func (s *server) listDatasetElements(w http.ResponseWriter, r *http.Request) {

type getDatasetElementArgs struct {
DatasetID string `json:"datasetID"`
Element string `json:"element"`
Name string `json:"name"`
}

func (a getDatasetElementArgs) validate() error {
if a.DatasetID == "" {
return fmt.Errorf("datasetID is required")
}
if a.Element == "" {
return fmt.Errorf("element is required")
} else if a.Name == "" {
return fmt.Errorf("name is required")
}
return nil
}
Expand Down Expand Up @@ -391,7 +251,7 @@ func (s *server) getDatasetElement(w http.ResponseWriter, r *http.Request) {
return
}

prg, err := loader.Program(r.Context(), s.getDatasetTool(req), "Get Element SDK", loader.Options{
prg, err := loader.Program(r.Context(), s.getDatasetTool(req), "Get Element", loader.Options{
Cache: g.Cache,
})
if err != nil {
Expand Down
2 changes: 0 additions & 2 deletions pkg/sdkserver/routes.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,10 +69,8 @@ func (s *server) addRoutes(mux *http.ServeMux) {
mux.HandleFunc("POST /credentials/delete", s.deleteCredential)

mux.HandleFunc("POST /datasets", s.listDatasets)
mux.HandleFunc("POST /datasets/create", s.createDataset)
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)
Expand Down

0 comments on commit d1d6179

Please sign in to comment.