Skip to content

Commit 6ec5178

Browse files
g-linvilleg-linville
authored
feat: add dataset operations to sdkserver (#869)
Signed-off-by: Grant Linville <[email protected]>
1 parent df259f9 commit 6ec5178

File tree

2 files changed

+338
-0
lines changed

2 files changed

+338
-0
lines changed

‎pkg/sdkserver/datasets.go

Lines changed: 332 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,332 @@
1+
package sdkserver
2+
3+
import (
4+
"encoding/json"
5+
"fmt"
6+
"net/http"
7+
8+
gcontext "github.com/gptscript-ai/gptscript/pkg/context"
9+
"github.com/gptscript-ai/gptscript/pkg/gptscript"
10+
"github.com/gptscript-ai/gptscript/pkg/loader"
11+
)
12+
13+
type datasetRequest struct {
14+
Input string `json:"input"`
15+
Workspace string `json:"workspace"`
16+
DatasetToolRepo string `json:"datasetToolRepo"`
17+
}
18+
19+
func (r datasetRequest) validate(requireInput bool) error {
20+
if r.Workspace == "" {
21+
return fmt.Errorf("workspace is required")
22+
} else if requireInput && r.Input == "" {
23+
return fmt.Errorf("input is required")
24+
}
25+
return nil
26+
}
27+
28+
func (r datasetRequest) opts(o gptscript.Options) gptscript.Options {
29+
opts := gptscript.Options{
30+
Cache: o.Cache,
31+
Monitor: o.Monitor,
32+
Runner: o.Runner,
33+
Workspace: r.Workspace,
34+
}
35+
return opts
36+
}
37+
38+
func (r datasetRequest) getToolRepo() string {
39+
if r.DatasetToolRepo != "" {
40+
return r.DatasetToolRepo
41+
}
42+
return "github.com/gptscript-ai/datasets"
43+
}
44+
45+
func (s *server) listDatasets(w http.ResponseWriter, r *http.Request) {
46+
logger := gcontext.GetLogger(r.Context())
47+
48+
var req datasetRequest
49+
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
50+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to decode request body: %w", err))
51+
return
52+
}
53+
54+
if err := req.validate(false); err != nil {
55+
writeError(logger, w, http.StatusBadRequest, err)
56+
return
57+
}
58+
59+
g, err := gptscript.New(r.Context(), req.opts(s.gptscriptOpts))
60+
if err != nil {
61+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to initialize gptscript: %w", err))
62+
return
63+
}
64+
65+
prg, err := loader.Program(r.Context(), "List Datasets from "+req.getToolRepo(), "", loader.Options{
66+
Cache: g.Cache,
67+
})
68+
69+
if err != nil {
70+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
71+
return
72+
}
73+
74+
result, err := g.Run(r.Context(), prg, s.gptscriptOpts.Env, req.Input)
75+
if err != nil {
76+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to run program: %w", err))
77+
return
78+
}
79+
80+
writeResponse(logger, w, map[string]any{"stdout": result})
81+
}
82+
83+
type createDatasetArgs struct {
84+
Name string `json:"datasetName"`
85+
Description string `json:"datasetDescription"`
86+
}
87+
88+
func (a createDatasetArgs) validate() error {
89+
if a.Name == "" {
90+
return fmt.Errorf("datasetName is required")
91+
}
92+
return nil
93+
}
94+
95+
func (s *server) createDataset(w http.ResponseWriter, r *http.Request) {
96+
logger := gcontext.GetLogger(r.Context())
97+
98+
var req datasetRequest
99+
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
100+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to decode request body: %w", err))
101+
return
102+
}
103+
104+
if err := req.validate(true); err != nil {
105+
writeError(logger, w, http.StatusBadRequest, err)
106+
return
107+
}
108+
109+
g, err := gptscript.New(r.Context(), req.opts(s.gptscriptOpts))
110+
if err != nil {
111+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to initialize gptscript: %w", err))
112+
return
113+
}
114+
115+
var args createDatasetArgs
116+
if err := json.Unmarshal([]byte(req.Input), &args); err != nil {
117+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to unmarshal input: %w", err))
118+
return
119+
}
120+
121+
if err := args.validate(); err != nil {
122+
writeError(logger, w, http.StatusBadRequest, err)
123+
return
124+
}
125+
126+
prg, err := loader.Program(r.Context(), "Create Dataset from "+req.getToolRepo(), "", loader.Options{
127+
Cache: g.Cache,
128+
})
129+
130+
if err != nil {
131+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
132+
return
133+
}
134+
135+
result, err := g.Run(r.Context(), prg, s.gptscriptOpts.Env, req.Input)
136+
if err != nil {
137+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to run program: %w", err))
138+
return
139+
}
140+
141+
writeResponse(logger, w, map[string]any{"stdout": result})
142+
}
143+
144+
type addDatasetElementArgs struct {
145+
DatasetID string `json:"datasetID"`
146+
ElementName string `json:"elementName"`
147+
ElementDescription string `json:"elementDescription"`
148+
ElementContent string `json:"elementContent"`
149+
}
150+
151+
func (a addDatasetElementArgs) validate() error {
152+
if a.DatasetID == "" {
153+
return fmt.Errorf("datasetID is required")
154+
}
155+
if a.ElementName == "" {
156+
return fmt.Errorf("elementName is required")
157+
}
158+
if a.ElementContent == "" {
159+
return fmt.Errorf("elementContent is required")
160+
}
161+
return nil
162+
}
163+
164+
func (s *server) addDatasetElement(w http.ResponseWriter, r *http.Request) {
165+
logger := gcontext.GetLogger(r.Context())
166+
167+
var req datasetRequest
168+
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
169+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to decode request body: %w", err))
170+
return
171+
}
172+
173+
if err := req.validate(true); err != nil {
174+
writeError(logger, w, http.StatusBadRequest, err)
175+
return
176+
}
177+
178+
g, err := gptscript.New(r.Context(), req.opts(s.gptscriptOpts))
179+
if err != nil {
180+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to initialize gptscript: %w", err))
181+
return
182+
}
183+
184+
var args addDatasetElementArgs
185+
if err := json.Unmarshal([]byte(req.Input), &args); err != nil {
186+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to unmarshal input: %w", err))
187+
return
188+
}
189+
190+
if err := args.validate(); err != nil {
191+
writeError(logger, w, http.StatusBadRequest, err)
192+
return
193+
}
194+
195+
prg, err := loader.Program(r.Context(), "Add Element from "+req.getToolRepo(), "", loader.Options{
196+
Cache: g.Cache,
197+
})
198+
if err != nil {
199+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
200+
return
201+
}
202+
203+
result, err := g.Run(r.Context(), prg, s.gptscriptOpts.Env, req.Input)
204+
if err != nil {
205+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to run program: %w", err))
206+
return
207+
}
208+
209+
writeResponse(logger, w, map[string]any{"stdout": result})
210+
}
211+
212+
type listDatasetElementsArgs struct {
213+
DatasetID string `json:"datasetID"`
214+
}
215+
216+
func (a listDatasetElementsArgs) validate() error {
217+
if a.DatasetID == "" {
218+
return fmt.Errorf("datasetID is required")
219+
}
220+
return nil
221+
}
222+
223+
func (s *server) listDatasetElements(w http.ResponseWriter, r *http.Request) {
224+
logger := gcontext.GetLogger(r.Context())
225+
226+
var req datasetRequest
227+
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
228+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to decode request body: %w", err))
229+
return
230+
}
231+
232+
if err := req.validate(true); err != nil {
233+
writeError(logger, w, http.StatusBadRequest, err)
234+
return
235+
}
236+
237+
g, err := gptscript.New(r.Context(), req.opts(s.gptscriptOpts))
238+
if err != nil {
239+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to initialize gptscript: %w", err))
240+
return
241+
}
242+
243+
var args listDatasetElementsArgs
244+
if err := json.Unmarshal([]byte(req.Input), &args); err != nil {
245+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to unmarshal input: %w", err))
246+
return
247+
}
248+
249+
if err := args.validate(); err != nil {
250+
writeError(logger, w, http.StatusBadRequest, err)
251+
return
252+
}
253+
254+
prg, err := loader.Program(r.Context(), "List Elements from "+req.getToolRepo(), "", loader.Options{
255+
Cache: g.Cache,
256+
})
257+
if err != nil {
258+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
259+
return
260+
}
261+
262+
result, err := g.Run(r.Context(), prg, s.gptscriptOpts.Env, req.Input)
263+
if err != nil {
264+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to run program: %w", err))
265+
return
266+
}
267+
268+
writeResponse(logger, w, map[string]any{"stdout": result})
269+
}
270+
271+
type getDatasetElementArgs struct {
272+
DatasetID string `json:"datasetID"`
273+
Element string `json:"element"`
274+
}
275+
276+
func (a getDatasetElementArgs) validate() error {
277+
if a.DatasetID == "" {
278+
return fmt.Errorf("datasetID is required")
279+
}
280+
if a.Element == "" {
281+
return fmt.Errorf("element is required")
282+
}
283+
return nil
284+
}
285+
286+
func (s *server) getDatasetElement(w http.ResponseWriter, r *http.Request) {
287+
logger := gcontext.GetLogger(r.Context())
288+
289+
var req datasetRequest
290+
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
291+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to decode request body: %w", err))
292+
return
293+
}
294+
295+
if err := req.validate(true); err != nil {
296+
writeError(logger, w, http.StatusBadRequest, err)
297+
return
298+
}
299+
300+
g, err := gptscript.New(r.Context(), req.opts(s.gptscriptOpts))
301+
if err != nil {
302+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to initialize gptscript: %w", err))
303+
return
304+
}
305+
306+
var args getDatasetElementArgs
307+
if err := json.Unmarshal([]byte(req.Input), &args); err != nil {
308+
writeError(logger, w, http.StatusBadRequest, fmt.Errorf("failed to unmarshal input: %w", err))
309+
return
310+
}
311+
312+
if err := args.validate(); err != nil {
313+
writeError(logger, w, http.StatusBadRequest, err)
314+
return
315+
}
316+
317+
prg, err := loader.Program(r.Context(), "Get Element from "+req.getToolRepo(), "", loader.Options{
318+
Cache: g.Cache,
319+
})
320+
if err != nil {
321+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
322+
return
323+
}
324+
325+
result, err := g.Run(r.Context(), prg, s.gptscriptOpts.Env, req.Input)
326+
if err != nil {
327+
writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to run program: %w", err))
328+
return
329+
}
330+
331+
writeResponse(logger, w, map[string]any{"stdout": result})
332+
}

‎pkg/sdkserver/routes.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,12 @@ func (s *server) addRoutes(mux *http.ServeMux) {
6666
mux.HandleFunc("POST /credentials/create", s.createCredential)
6767
mux.HandleFunc("POST /credentials/reveal", s.revealCredential)
6868
mux.HandleFunc("POST /credentials/delete", s.deleteCredential)
69+
70+
mux.HandleFunc("POST /datasets", s.listDatasets)
71+
mux.HandleFunc("POST /datasets/create", s.createDataset)
72+
mux.HandleFunc("POST /datasets/list-elements", s.listDatasetElements)
73+
mux.HandleFunc("POST /datasets/get-element", s.getDatasetElement)
74+
mux.HandleFunc("POST /datasets/add-element", s.addDatasetElement)
6975
}
7076

7177
// health just provides an endpoint for checking whether the server is running and accessible.

0 commit comments

Comments
 (0)