client.go 1.6 KB
Newer Older
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1
2
3
package api

import (
Jeffrey Morgan's avatar
Jeffrey Morgan committed
4
	"bufio"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
5
6
7
8
9
10
11
12
13
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
)

type Client struct {
Bruce MacDonald's avatar
Bruce MacDonald committed
14
15
	URL  string
	HTTP http.Client
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16
17
}

Michael Yang's avatar
Michael Yang committed
18
func (c *Client) stream(ctx context.Context, method string, path string, reqData any, fn func(bts []byte) error) error {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
	var reqBody io.Reader
	var data []byte
	var err error
	if reqData != nil {
		data, err = json.Marshal(reqData)
		if err != nil {
			return err
		}
		reqBody = bytes.NewReader(data)
	}

	url := fmt.Sprintf("%s%s", c.URL, path)

	req, err := http.NewRequestWithContext(ctx, method, url, reqBody)
	if err != nil {
		return err
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
37
38
39
40
41
42
43
44
45
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	res, err := c.HTTP.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

Michael Yang's avatar
Michael Yang committed
46
47
48
	scanner := bufio.NewScanner(res.Body)
	for scanner.Scan() {
		if err := fn(scanner.Bytes()); err != nil {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
49
50
51
52
			return err
		}
	}

Michael Yang's avatar
Michael Yang committed
53
54
	return nil
}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
55

Michael Yang's avatar
Michael Yang committed
56
type GenerateResponseFunc func(GenerateResponse) error
Jeffrey Morgan's avatar
Jeffrey Morgan committed
57

Michael Yang's avatar
Michael Yang committed
58
59
60
61
func (c *Client) Generate(ctx context.Context, req *GenerateRequest, fn GenerateResponseFunc) error {
	return c.stream(ctx, http.MethodPost, "/api/generate", req, func(bts []byte) error {
		var resp GenerateResponse
		if err := json.Unmarshal(bts, &resp); err != nil {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
62
63
			return err
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
64

Michael Yang's avatar
Michael Yang committed
65
66
		return fn(resp)
	})
Jeffrey Morgan's avatar
Jeffrey Morgan committed
67
}
Bruce MacDonald's avatar
Bruce MacDonald committed
68

Michael Yang's avatar
Michael Yang committed
69
70
71
72
73
74
75
type PullProgressFunc func(PullProgress) error

func (c *Client) Pull(ctx context.Context, req *PullRequest, fn PullProgressFunc) error {
	return c.stream(ctx, http.MethodPost, "/api/pull", req, func(bts []byte) error {
		var resp PullProgress
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
Bruce MacDonald's avatar
Bruce MacDonald committed
76
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
77

Michael Yang's avatar
Michael Yang committed
78
79
		return fn(resp)
	})
Bruce MacDonald's avatar
Bruce MacDonald committed
80
}