client.go 4.36 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
	"bytes"
	"context"
	"encoding/json"
8
	"fmt"
Patrick Devine's avatar
Patrick Devine committed
9
	"io"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
10
	"net/http"
Michael Yang's avatar
Michael Yang committed
11
	"net/url"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
12
13
)

Patrick Devine's avatar
Patrick Devine committed
14
15
16
17
type Client struct {
	base    url.URL
	HTTP    http.Client
	Headers http.Header
Michael Yang's avatar
Michael Yang committed
18
19
}

Patrick Devine's avatar
Patrick Devine committed
20
21
22
func checkError(resp *http.Response, body []byte) error {
	if resp.StatusCode >= 200 && resp.StatusCode < 400 {
		return nil
Michael Yang's avatar
Michael Yang committed
23
24
	}

Patrick Devine's avatar
Patrick Devine committed
25
	apiError := StatusError{StatusCode: resp.StatusCode}
Michael Yang's avatar
Michael Yang committed
26

Patrick Devine's avatar
Patrick Devine committed
27
28
29
30
31
32
33
	err := json.Unmarshal(body, &apiError)
	if err != nil {
		// Use the full body as the message if we fail to decode a response.
		apiError.Message = string(body)
	}

	return apiError
Michael Yang's avatar
Michael Yang committed
34
35
36
37
38
39
40
41
42
43
}

func NewClient(hosts ...string) *Client {
	host := "127.0.0.1:11434"
	if len(hosts) > 0 {
		host = hosts[0]
	}

	return &Client{
		base: url.URL{Scheme: "http", Host: host},
Patrick Devine's avatar
Patrick Devine committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
		HTTP: http.Client{},
	}
}

func (c *Client) do(ctx context.Context, method, path string, reqData, respData any) error {
	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 := c.base.JoinPath(path).String()

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

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	for k, v := range c.Headers {
		req.Header[k] = v
Michael Yang's avatar
Michael Yang committed
72
	}
Patrick Devine's avatar
Patrick Devine committed
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

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

	respBody, err := io.ReadAll(respObj.Body)
	if err != nil {
		return err
	}

	if err := checkError(respObj, respBody); err != nil {
		return err
	}

	if len(respBody) > 0 && respData != nil {
		if err := json.Unmarshal(respBody, respData); err != nil {
			return err
		}
	}
	return nil

Jeffrey Morgan's avatar
Jeffrey Morgan committed
96
97
}

Michael Yang's avatar
Michael Yang committed
98
func (c *Client) stream(ctx context.Context, method, path string, data any, fn func([]byte) error) error {
99
100
101
102
103
104
	var buf *bytes.Buffer
	if data != nil {
		bts, err := json.Marshal(data)
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
105

106
		buf = bytes.NewBuffer(bts)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
107
108
	}

109
	request, err := http.NewRequestWithContext(ctx, method, c.base.JoinPath(path).String(), buf)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
110
111
112
113
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
114
115
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Accept", "application/json")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
116

Michael Yang's avatar
Michael Yang committed
117
	response, err := http.DefaultClient.Do(request)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
118
119
120
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
121
	defer response.Body.Close()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
122

123
124
125
	scanner := bufio.NewScanner(response.Body)
	for scanner.Scan() {
		var errorResponse struct {
Michael Yang's avatar
Michael Yang committed
126
			Error string `json:"error,omitempty"`
127
128
129
130
131
132
133
		}

		bts := scanner.Bytes()
		if err := json.Unmarshal(bts, &errorResponse); err != nil {
			return fmt.Errorf("unmarshal: %w", err)
		}

Michael Yang's avatar
Michael Yang committed
134
135
136
137
138
139
		if response.StatusCode >= 400 {
			return StatusError{
				StatusCode: response.StatusCode,
				Status:     response.Status,
				Message:    errorResponse.Error,
			}
140
141
		}

Michael Yang's avatar
Michael Yang committed
142
		if err := fn(bts); err != nil {
143
			return err
Jeffrey Morgan's avatar
Jeffrey Morgan committed
144
145
146
		}
	}

Michael Yang's avatar
Michael Yang committed
147
148
	return nil
}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
149

Michael Yang's avatar
Michael Yang committed
150
type GenerateResponseFunc func(GenerateResponse) error
Jeffrey Morgan's avatar
Jeffrey Morgan committed
151

Michael Yang's avatar
Michael Yang committed
152
func (c *Client) Generate(ctx context.Context, req *GenerateRequest, fn GenerateResponseFunc) error {
153
154
155
156
157
158
159
160
	return c.stream(ctx, http.MethodPost, "/api/generate", req, func(bts []byte) error {
		var resp GenerateResponse
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
Jeffrey Morgan's avatar
Jeffrey Morgan committed
161
}
Bruce MacDonald's avatar
Bruce MacDonald committed
162

163
type PullProgressFunc func(ProgressResponse) error
Michael Yang's avatar
Michael Yang committed
164
165

func (c *Client) Pull(ctx context.Context, req *PullRequest, fn PullProgressFunc) error {
166
	return c.stream(ctx, http.MethodPost, "/api/pull", req, func(bts []byte) error {
167
		var resp ProgressResponse
168
169
170
171
172
173
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
Bruce MacDonald's avatar
Bruce MacDonald committed
174
}
175

176
type PushProgressFunc func(ProgressResponse) error
177
178
179

func (c *Client) Push(ctx context.Context, req *PushRequest, fn PushProgressFunc) error {
	return c.stream(ctx, http.MethodPost, "/api/push", req, func(bts []byte) error {
180
		var resp ProgressResponse
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
}

type CreateProgressFunc func(CreateProgress) error

func (c *Client) Create(ctx context.Context, req *CreateRequest, fn CreateProgressFunc) error {
	return c.stream(ctx, http.MethodPost, "/api/create", req, func(bts []byte) error {
		var resp CreateProgress
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
}
Patrick Devine's avatar
Patrick Devine committed
201
202
203
204
205
206
207
208

func (c *Client) List(ctx context.Context) (*ListResponse, error) {
	var lr ListResponse
	if err := c.do(ctx, http.MethodGet, "/api/tags", nil, &lr); err != nil {
		return nil, err
	}
	return &lr, nil
}