1
package api
1
package api
2
2
3
// This file is for Golang clients of Klex.
3
// This file is for Golang clients of Klex.
4
4
5
import (
5
import (
6
"bytes"
6
"bytes"
7
"encoding/json"
7
"encoding/json"
8
"fmt"
8
"fmt"
9
"io/ioutil"
9
"io/ioutil"
10
"log"
10
"log"
11
"net/http"
11
"net/http"
12
"sort"
12
"sort"
13
)
13
)
14
14
15
type Client struct {
15
type Client struct {
16
KlexURL string
16
KlexURL string
17
APIKey string
17
APIKey string
18
}
18
}
19
19
20
func NewClient(klexURL, apiKey string) *Client {
20
func NewClient(klexURL, apiKey string) *Client {
21
if klexURL == "" || apiKey == "" {
21
if klexURL == "" || apiKey == "" {
22
log.Printf("NewClient: missing klexURL or apiKey")
22
log.Printf("NewClient: missing klexURL or apiKey")
23
return nil
23
return nil
24
}
24
}
25
return &Client{klexURL, apiKey}
25
return &Client{klexURL, apiKey}
26
}
26
}
27
27
28
func (c *Client) call(method, path string, req, res interface{}) error {
28
func (c *Client) call(method, path string, req, res interface{}) error {
29
reqBody, err := json.Marshal(req)
29
reqBody, err := json.Marshal(req)
30
if err != nil {
30
if err != nil {
31
return fmt.Errorf("Cannot marshal request: %v", err)
31
return fmt.Errorf("Cannot marshal request: %v", err)
32
}
32
}
33
reqBytes := bytes.NewBuffer(reqBody)
33
reqBytes := bytes.NewBuffer(reqBody)
34
r, err := http.NewRequest(method, c.KlexURL + path, reqBytes)
34
r, err := http.NewRequest(method, c.KlexURL + path, reqBytes)
35
if err != nil {
35
if err != nil {
36
return fmt.Errorf("In http.NewRequest: %v", err)
36
return fmt.Errorf("In http.NewRequest: %v", err)
37
}
37
}
38
r.Header.Set("Authorization", "Bearer " + c.APIKey)
38
r.Header.Set("Authorization", "Bearer " + c.APIKey)
39
r.Header.Set("Content-Type", "application/json")
39
r.Header.Set("Content-Type", "application/json")
40
resHttp, err := http.DefaultClient.Do(r)
40
resHttp, err := http.DefaultClient.Do(r)
41
if err != nil {
41
if err != nil {
42
return fmt.Errorf("http.DefaultClient.Do: %v", err)
42
return fmt.Errorf("http.DefaultClient.Do: %v", err)
43
}
43
}
44
defer resHttp.Body.Close()
44
defer resHttp.Body.Close()
45
resBody, err := ioutil.ReadAll(resHttp.Body)
45
resBody, err := ioutil.ReadAll(resHttp.Body)
46
if err != nil {
46
if err != nil {
47
return fmt.Errorf("Response error: %v", err)
47
return fmt.Errorf("Response error: %v", err)
48
}
48
}
49
if resHttp.StatusCode != 200 && resHttp.StatusCode != 204 {
49
if resHttp.StatusCode != 200 && resHttp.StatusCode != 204 {
50
return fmt.Errorf("Status %d; response=%s", resHttp.StatusCode, resBody)
50
return fmt.Errorf("Status %d; response=%s", resHttp.StatusCode, resBody)
51
}
51
}
52
if res != nil {
52
if res != nil {
53
if err := json.Unmarshal(resBody, res); err != nil {
53
if err := json.Unmarshal(resBody, res); err != nil {
54
return fmt.Errorf("Bad response %s\nerror=%v", resBody, err)
54
return fmt.Errorf("Bad response %s\nerror=%v", resBody, err)
55
}
55
}
56
}
56
}
57
return nil
57
return nil
58
}
58
}
59
59
60
// F executes a function on one given input.
60
// F executes a function on one given input.
61
func (c *Client) F(f, in string) (string, error) {
61
func (c *Client) F(f, in string) (string, error) {
62
var res FResponse
62
var res FResponse
63
err := c.call("POST", "/f", FRequest{FName: f, In: in}, &res)
63
err := c.call("POST", "/f", FRequest{FName: f, In: in}, &res)
64
if err != nil {
64
if err != nil {
65
return "", err
65
return "", err
66
}
66
}
67
if res.Err != "" {
67
if res.Err != "" {
68
return "", fmt.Errorf(res.Err)
68
return "", fmt.Errorf(res.Err)
69
}
69
}
70
return res.Out, nil
70
return res.Out, nil
71
}
71
}
72
72
73
// Messages executes an LLM function using the Messages API.
74
func (c *Client) Messages(f string, req *MessagesRequest) (*MessagesResponse, error) {
75
in, err := json.Marshal(req)
76
if err != nil {
77
return nil, fmt.Errorf("Cannot marshal request: %v", err)
78
}
79
out, err := c.F(f, string(in))
80
if err != nil {
81
return nil, err
82
}
83
var res MessagesResponse
84
err = json.Unmarshal([]byte(out), &res)
85
if err != nil {
86
// Instead of failing, treat the whole output as text, and add an error.
87
// Let the caller figure this out.
88
res.Error = &ErrorResponse{
89
Type: "response-json",
90
Message: err.Error(),
91
}
92
res.Content = []ContentBlock{{Type: "text", Text: out}}
93
}
94
return &res, nil
95
}
96
73
// NewDataset creates a new dataset or updates an existing one.
97
// NewDataset creates a new dataset or updates an existing one.
74
// This is the simplest way, meant for datasets smaller than ~1GB.
98
// This is the simplest way, meant for datasets smaller than ~1GB.
75
func (c *Client) NewDataset(name string, data map[string]string) error {
99
func (c *Client) NewDataset(name string, data map[string]string) error {
76
// TODO: this loses key names; e t API.
100
// TODO: this loses key names; e t API.
77
req := NewDatasetRequest{Name: name, Data: nil}
101
req := NewDatasetRequest{Name: name, Data: nil}
78
keys := make([]string, 0, len(data))
102
keys := make([]string, 0, len(data))
79
for k := range data {
103
for k := range data {
80
keys = append(keys, k)
104
keys = append(keys, k)
81
}
105
}
82
sort.Strings(keys)
106
sort.Strings(keys)
83
for _, k := range keys {
107
for _, k := range keys {
84
req.Data = append(req.Data, data[k])
108
req.Data = append(req.Data, data[k])
85
}
109
}
86
110
87
var res NewDatasetResponse
111
var res NewDatasetResponse
88
err := c.call("POST", "/datasets/new", req, &res)
112
err := c.call("POST", "/datasets/new", req, &res)
89
if err != nil {
113
if err != nil {
90
return fmt.Errorf("Error POSTing to /datasets/new: %v", err)
114
return fmt.Errorf("Error POSTing to /datasets/new: %v", err)
91
}
115
}
92
if res.Name != name || res.Size != len(data) {
116
if res.Name != name || res.Size != len(data) {
93
pretty, _ := json.MarshalIndent(res, "", " ")
117
pretty, _ := json.MarshalIndent(res, "", " ")
94
return fmt.Errorf("Unexpected response from /datasets/new: %s", pretty)
118
return fmt.Errorf("Unexpected response from /datasets/new: %s", pretty)
95
}
119
}
96
return nil
120
return nil
97
}
121
}
98
122
99
// BeginNewDataset starts a new dataset upload using the v2 API.
123
// BeginNewDataset starts a new dataset upload using the v2 API.
100
// Returns the version key to use in UploadKv() and EndNewDataset().
124
// Returns the version key to use in UploadKv() and EndNewDataset().
101
// Keep the key secret until EndNewDataset() returns successfully.
125
// Keep the key secret until EndNewDataset() returns successfully.
102
func (c *Client) BeginNewDataset(name string) (string, error) {
126
func (c *Client) BeginNewDataset(name string) (string, error) {
103
req := BeginNewDatasetRequest{Name: name}
127
req := BeginNewDatasetRequest{Name: name}
104
var res BeginNewDatasetResponse
128
var res BeginNewDatasetResponse
105
err := c.call("POST", "/datasets/begin_new", req, &res)
129
err := c.call("POST", "/datasets/begin_new", req, &res)
106
if err != nil {
130
if err != nil {
107
return "", fmt.Errorf("Error POSTing to /datasets/begin_new: %v", err)
131
return "", fmt.Errorf("Error POSTing to /datasets/begin_new: %v", err)
108
}
132
}
109
return res.VersionKey, nil
133
return res.VersionKey, nil
110
}
134
}
111
135
112
// UploadKv uploads more key-value pairs of the dataset being created.
136
// UploadKv uploads more key-value pairs of the dataset being created.
113
func (c *Client) UploadKV(versionKey string, records []KV) error {
137
func (c *Client) UploadKV(versionKey string, records []KV) error {
114
req := UploadKVRequest{VersionKey: versionKey, Records: records}
138
req := UploadKVRequest{VersionKey: versionKey, Records: records}
115
err := c.call("POST", "/datasets/upload_kv", req, nil)
139
err := c.call("POST", "/datasets/upload_kv", req, nil)
116
if err != nil {
140
if err != nil {
117
return fmt.Errorf("Error POSTing to /datasets/upload_kv: %v", err)
141
return fmt.Errorf("Error POSTing to /datasets/upload_kv: %v", err)
118
}
142
}
119
return nil
143
return nil
120
}
144
}
121
145
122
// EndNewDataset commits the dataset being created.
146
// EndNewDataset commits the dataset being created.
123
func (c *Client) EndNewDataset(name, version_key string, size int) error {
147
func (c *Client) EndNewDataset(name, version_key string, size int) error {
124
req := EndNewDatasetRequest{Name: name, VersionKey: version_key, Size: size}
148
req := EndNewDatasetRequest{Name: name, VersionKey: version_key, Size: size}
125
err := c.call("POST", "/datasets/end_new", req, nil)
149
err := c.call("POST", "/datasets/end_new", req, nil)
126
if err != nil {
150
if err != nil {
127
return fmt.Errorf("Error POSTing to /datasets/end_new: %v", err)
151
return fmt.Errorf("Error POSTing to /datasets/end_new: %v", err)
128
}
152
}
129
return nil
153
return nil
130
}
154
}