client.go 7.39 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"
Michael Yang's avatar
Michael Yang committed
10
	"net"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
11
	"net/http"
Michael Yang's avatar
Michael Yang committed
12
	"net/url"
13
	"os"
Michael Yang's avatar
Michael Yang committed
14
	"runtime"
15
	"strings"
Michael Yang's avatar
Michael Yang committed
16

17
18
	"github.com/ollama/ollama/format"
	"github.com/ollama/ollama/version"
19
20
)

Patrick Devine's avatar
Patrick Devine committed
21
type Client struct {
Michael Yang's avatar
Michael Yang committed
22
	base *url.URL
Michael Yang's avatar
Michael Yang committed
23
	http *http.Client
Michael Yang's avatar
Michael Yang committed
24
25
}

Patrick Devine's avatar
Patrick Devine committed
26
func checkError(resp *http.Response, body []byte) error {
Michael Yang's avatar
Michael Yang committed
27
	if resp.StatusCode < http.StatusBadRequest {
Patrick Devine's avatar
Patrick Devine committed
28
		return nil
Michael Yang's avatar
Michael Yang committed
29
30
	}

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

Patrick Devine's avatar
Patrick Devine committed
33
34
35
	err := json.Unmarshal(body, &apiError)
	if err != nil {
		// Use the full body as the message if we fail to decode a response.
36
		apiError.ErrorMessage = string(body)
Patrick Devine's avatar
Patrick Devine committed
37
38
39
	}

	return apiError
Michael Yang's avatar
Michael Yang committed
40
41
}

Michael Yang's avatar
Michael Yang committed
42
func ClientFromEnvironment() (*Client, error) {
Michael Yang's avatar
Michael Yang committed
43
44
	defaultPort := "11434"

Michael Yang's avatar
Michael Yang committed
45
	scheme, hostport, ok := strings.Cut(os.Getenv("OLLAMA_HOST"), "://")
Michael Yang's avatar
Michael Yang committed
46
47
	switch {
	case !ok:
Michael Yang's avatar
Michael Yang committed
48
		scheme, hostport = "http", os.Getenv("OLLAMA_HOST")
Michael Yang's avatar
Michael Yang committed
49
50
51
52
	case scheme == "http":
		defaultPort = "80"
	case scheme == "https":
		defaultPort = "443"
Michael Yang's avatar
Michael Yang committed
53
54
	}

Michael Yang's avatar
Michael Yang committed
55
56
57
	// trim trailing slashes
	hostport = strings.TrimRight(hostport, "/")

Michael Yang's avatar
Michael Yang committed
58
59
	host, port, err := net.SplitHostPort(hostport)
	if err != nil {
Michael Yang's avatar
Michael Yang committed
60
		host, port = "127.0.0.1", defaultPort
Michael Yang's avatar
Michael Yang committed
61
		if ip := net.ParseIP(strings.Trim(hostport, "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
62
			host = ip.String()
Michael Yang's avatar
Michael Yang committed
63
64
		} else if hostport != "" {
			host = hostport
Michael Yang's avatar
Michael Yang committed
65
66
67
		}
	}

Michael Yang's avatar
Michael Yang committed
68
	return &Client{
Michael Yang's avatar
Michael Yang committed
69
70
71
72
		base: &url.URL{
			Scheme: scheme,
			Host:   net.JoinHostPort(host, port),
		},
Michael Yang's avatar
Michael Yang committed
73
74
		http: http.DefaultClient,
	}, nil
Patrick Devine's avatar
Patrick Devine committed
75
76
77
78
79
80
}

func (c *Client) do(ctx context.Context, method, path string, reqData, respData any) error {
	var reqBody io.Reader
	var data []byte
	var err error
Michael Yang's avatar
Michael Yang committed
81
82
83
84
85
86
87
88

	switch reqData := reqData.(type) {
	case io.Reader:
		// reqData is already an io.Reader
		reqBody = reqData
	case nil:
		// noop
	default:
Patrick Devine's avatar
Patrick Devine committed
89
90
91
92
		data, err = json.Marshal(reqData)
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
93

Patrick Devine's avatar
Patrick Devine committed
94
95
96
		reqBody = bytes.NewReader(data)
	}

Michael Yang's avatar
Michael Yang committed
97
	requestURL := c.base.JoinPath(path)
Michael Yang's avatar
Michael Yang committed
98
	request, err := http.NewRequestWithContext(ctx, method, requestURL.String(), reqBody)
Patrick Devine's avatar
Patrick Devine committed
99
100
101
102
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
103
104
105
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Accept", "application/json")
	request.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
Patrick Devine's avatar
Patrick Devine committed
106

Michael Yang's avatar
Michael Yang committed
107
	respObj, err := c.http.Do(request)
Patrick Devine's avatar
Patrick Devine committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
	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
128
129
}

Michael Yang's avatar
Michael Yang committed
130
const maxBufferSize = 512 * format.KiloByte
131

Michael Yang's avatar
Michael Yang committed
132
func (c *Client) stream(ctx context.Context, method, path string, data any, fn func([]byte) error) error {
133
134
135
136
137
138
	var buf *bytes.Buffer
	if data != nil {
		bts, err := json.Marshal(data)
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
139

140
		buf = bytes.NewBuffer(bts)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
141
142
	}

Michael Yang's avatar
Michael Yang committed
143
	requestURL := c.base.JoinPath(path)
Michael Yang's avatar
Michael Yang committed
144
	request, err := http.NewRequestWithContext(ctx, method, requestURL.String(), buf)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
145
146
147
148
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
149
	request.Header.Set("Content-Type", "application/json")
150
	request.Header.Set("Accept", "application/x-ndjson")
Michael Yang's avatar
Michael Yang committed
151
	request.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
Jeffrey Morgan's avatar
Jeffrey Morgan committed
152

Michael Yang's avatar
Michael Yang committed
153
	response, err := c.http.Do(request)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
154
155
156
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
157
	defer response.Body.Close()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
158

159
	scanner := bufio.NewScanner(response.Body)
160
161
162
	// increase the buffer size to avoid running out of space
	scanBuf := make([]byte, 0, maxBufferSize)
	scanner.Buffer(scanBuf, maxBufferSize)
163
164
	for scanner.Scan() {
		var errorResponse struct {
Michael Yang's avatar
Michael Yang committed
165
			Error string `json:"error,omitempty"`
166
167
168
169
170
171
172
		}

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

Michael Yang's avatar
Michael Yang committed
173
		if errorResponse.Error != "" {
174
			return fmt.Errorf(errorResponse.Error)
Michael Yang's avatar
Michael Yang committed
175
176
		}

Michael Yang's avatar
Michael Yang committed
177
		if response.StatusCode >= http.StatusBadRequest {
Michael Yang's avatar
Michael Yang committed
178
			return StatusError{
179
180
181
				StatusCode:   response.StatusCode,
				Status:       response.Status,
				ErrorMessage: errorResponse.Error,
Michael Yang's avatar
Michael Yang committed
182
			}
183
184
		}

Michael Yang's avatar
Michael Yang committed
185
		if err := fn(bts); err != nil {
186
			return err
Jeffrey Morgan's avatar
Jeffrey Morgan committed
187
188
189
		}
	}

Michael Yang's avatar
Michael Yang committed
190
191
	return nil
}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
192

Michael Yang's avatar
Michael Yang committed
193
type GenerateResponseFunc func(GenerateResponse) error
Jeffrey Morgan's avatar
Jeffrey Morgan committed
194

Michael Yang's avatar
Michael Yang committed
195
func (c *Client) Generate(ctx context.Context, req *GenerateRequest, fn GenerateResponseFunc) error {
196
197
198
199
200
201
202
203
	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
204
}
Bruce MacDonald's avatar
Bruce MacDonald committed
205

Bruce MacDonald's avatar
Bruce MacDonald committed
206
207
208
209
210
211
212
213
214
215
216
217
218
type ChatResponseFunc func(ChatResponse) error

func (c *Client) Chat(ctx context.Context, req *ChatRequest, fn ChatResponseFunc) error {
	return c.stream(ctx, http.MethodPost, "/api/chat", req, func(bts []byte) error {
		var resp ChatResponse
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
}

219
type PullProgressFunc func(ProgressResponse) error
Michael Yang's avatar
Michael Yang committed
220
221

func (c *Client) Pull(ctx context.Context, req *PullRequest, fn PullProgressFunc) error {
222
	return c.stream(ctx, http.MethodPost, "/api/pull", req, func(bts []byte) error {
223
		var resp ProgressResponse
224
225
226
227
228
229
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
Bruce MacDonald's avatar
Bruce MacDonald committed
230
}
231

232
type PushProgressFunc func(ProgressResponse) error
233
234
235

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 {
236
		var resp ProgressResponse
237
238
239
240
241
242
243
244
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
}

245
type CreateProgressFunc func(ProgressResponse) error
246
247
248

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 {
249
		var resp ProgressResponse
250
251
252
253
254
255
256
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
}
Patrick Devine's avatar
Patrick Devine committed
257
258
259
260
261
262
263
264

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
}
265

Patrick Devine's avatar
Patrick Devine committed
266
267
268
269
270
271
272
func (c *Client) Copy(ctx context.Context, req *CopyRequest) error {
	if err := c.do(ctx, http.MethodPost, "/api/copy", req, nil); err != nil {
		return err
	}
	return nil
}

273
274
275
276
277
func (c *Client) Delete(ctx context.Context, req *DeleteRequest) error {
	if err := c.do(ctx, http.MethodDelete, "/api/delete", req, nil); err != nil {
		return err
	}
	return nil
278
}
279

Patrick Devine's avatar
Patrick Devine committed
280
281
282
283
284
285
286
287
func (c *Client) Show(ctx context.Context, req *ShowRequest) (*ShowResponse, error) {
	var resp ShowResponse
	if err := c.do(ctx, http.MethodPost, "/api/show", req, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

288
func (c *Client) Heartbeat(ctx context.Context) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
289
	if err := c.do(ctx, http.MethodHead, "/", nil, nil); err != nil {
290
291
292
293
		return err
	}
	return nil
}
Brian Murray's avatar
Brian Murray committed
294
295
296
297
298
299
300
func (c *Client) Embeddings(ctx context.Context, req *EmbeddingRequest) (*EmbeddingResponse, error) {
	var resp EmbeddingResponse
	if err := c.do(ctx, http.MethodPost, "/api/embeddings", req, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}
Michael Yang's avatar
Michael Yang committed
301

Michael Yang's avatar
Michael Yang committed
302
func (c *Client) CreateBlob(ctx context.Context, digest string, r io.Reader) error {
303
	return c.do(ctx, http.MethodPost, fmt.Sprintf("/api/blobs/%s", digest), r, nil)
Michael Yang's avatar
Michael Yang committed
304
}
Michael Yang's avatar
Michael Yang committed
305
306
307
308
309
310
311
312
313
314
315
316

func (c *Client) Version(ctx context.Context) (string, error) {
	var version struct {
		Version string `json:"version"`
	}

	if err := c.do(ctx, http.MethodGet, "/api/version", nil, &version); err != nil {
		return "", err
	}

	return version.Version, nil
}