code.oscarkilo.com/klex-git/api/api.go

package api

// This file is for Golang clients of Klex.

import (
  "bytes"
  "encoding/json"
  "fmt"
  "io/ioutil"
  "log"
  "net/http"
  "sort"
)

type Client struct {
  KlexURL string
  APIKey string
}

func NewClient(klexURL, apiKey string) *Client {
  if klexURL == "" || apiKey == "" {
    log.Printf("NewClient: missing klexURL or apiKey")
    return nil
  }
  return &Client{klexURL, apiKey}
}

func (c *Client) call(method, path string, req, res interface{}) error {
  reqBody, err := json.Marshal(req)
  if err != nil {
    return fmt.Errorf("Cannot marshal request: %v", err)
  }
  reqBytes := bytes.NewBuffer(reqBody)
  r, err := http.NewRequest(method, c.KlexURL + path, reqBytes)
  if err != nil {
    return fmt.Errorf("In http.NewRequest: %v", err)
  }
  r.Header.Set("Authorization", "Bearer " + c.APIKey)
  r.Header.Set("Content-Type", "application/json")
  resHttp, err := http.DefaultClient.Do(r)
  if err != nil {
    return fmt.Errorf("http.DefaultClient.Do: %v", err)
  }
  defer resHttp.Body.Close()
  resBody, err := ioutil.ReadAll(resHttp.Body)
  if err != nil {
    return fmt.Errorf("Response error: %v", err)
  }
  if resHttp.StatusCode != 200 && resHttp.StatusCode != 204 {
    return fmt.Errorf("Status %d; response=%s", resHttp.StatusCode, resBody)
  }
  if res != nil {
    if err := json.Unmarshal(resBody, res); err != nil {
      return fmt.Errorf("Bad response %s\nerror=%v", resBody, err)
    }
  }
  return nil
}

// F executes a function on one given input.
// TODO: Implement F() in terms of Messagse(), rather than the otehr way around.
func (c *Client) F(f, in string) (string, error) {
  var res FResponse
  err := c.call("POST", "/f", FRequest{FName: f, In: in}, &res)
  if err != nil {
    return "", err
  }
  if res.Err != "" {
    return "", fmt.Errorf(res.Err)
  }
  return res.Out, nil
}

// Messages executes an LLM function using the Messages API.
// Set req.Model to one of the Klex LLM function names.
func (c *Client) Messages(req MessagesRequest) (*MessagesResponse, error) {
  f := req.Model
  req.Model = ""
  if f == "" {
    return nil, fmt.Errorf("MessagesRequest.Model is empty")
  }
  in, err := json.Marshal(req)
  if err != nil {
    return nil, fmt.Errorf("Cannot marshal request: %v", err)
  }
  out, err := c.F(f, string(in))
  if err != nil {
    return nil, err
  }
  var res MessagesResponse
  err = json.Unmarshal([]byte(out), &res)
  if err != nil {
    // Instead of failing, treat the whole output as text, and add an error.
    // Let the caller figure this out.
    res.Error = &ErrorResponse{
      Type: "response-json",
      Message: err.Error(),
    }
    res.Content = []ContentBlock{{Type: "text", Text: out}}
  }
  return &res, nil
}

// NewDataset creates a new dataset or updates an existing one.
// This is the simplest way, meant for datasets smaller than ~1GB.
func (c *Client) NewDataset(name string, data map[string]string) error {
  // TODO: this loses key names; get rid of this API.
  req := NewDatasetRequest{Name: name, Data: nil}
  keys := make([]string, 0, len(data))
  for k := range data {
    keys = append(keys, k)
  }
  sort.Strings(keys)
  for _, k := range keys {
    req.Data = append(req.Data, data[k])
  }

  var res NewDatasetResponse
  err := c.call("POST", "/datasets/new", req, &res)
  if err != nil {
    return fmt.Errorf("Error POSTing to /datasets/new: %v", err)
  }
  if res.Name != name || res.Size != len(data) {
    pretty, _ := json.MarshalIndent(res, "", "  ")
    return fmt.Errorf("Unexpected response from /datasets/new: %s", pretty)
  }
  return nil
}

// BeginNewDataset starts a new dataset upload using the v2 API.
// Returns the version key to use in UploadKv() and EndNewDataset().
// Keep the key secret until EndNewDataset() returns successfully.
func (c *Client) BeginNewDataset(name string) (string, error) {
  req := BeginNewDatasetRequest{Name: name}
  var res BeginNewDatasetResponse
  err := c.call("POST", "/datasets/begin_new", req, &res)
  if err != nil {
    return "", fmt.Errorf("Error POSTing to /datasets/begin_new: %v", err)
  }
  return res.VersionKey, nil
}

// UploadKv uploads more key-value pairs of the dataset being created.
func (c *Client) UploadKV(versionKey string, records []KV) error {
  req := UploadKVRequest{VersionKey: versionKey, Records: records}
  err := c.call("POST", "/datasets/upload_kv", req, nil)
  if err != nil {
    return fmt.Errorf("Error POSTing to /datasets/upload_kv: %v", err)
  }
  return nil
}

// EndNewDataset commits the dataset being created.
func (c *Client) EndNewDataset(name, version_key string, size int) error {
  req := EndNewDatasetRequest{Name: name, VersionKey: version_key, Size: size}
  err := c.call("POST", "/datasets/end_new", req, nil)
  if err != nil {
    return fmt.Errorf("Error POSTing to /datasets/end_new: %v", err)
  }
  return nil
}