client.go 13.2 KB
Newer Older
1
2
// Package api implements the client-side API for code wishing to interact
// with the ollama service. The methods of the [Client] type correspond to
3
// the ollama REST API as described in [the API documentation].
4
5
// The ollama command-line client itself uses this package to interact with
// the backend service.
6
7
8
9
10
11
12
//
// # Examples
//
// Several examples of using this package are available [in the GitHub
// repository].
//
// [the API documentation]: https://github.com/ollama/ollama/blob/main/docs/api.md
13
// [in the GitHub repository]: https://github.com/ollama/ollama/tree/main/api/examples
Jeffrey Morgan's avatar
Jeffrey Morgan committed
14
15
16
package api

import (
Jeffrey Morgan's avatar
Jeffrey Morgan committed
17
	"bufio"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
18
19
20
	"bytes"
	"context"
	"encoding/json"
Michael Yang's avatar
lint  
Michael Yang committed
21
	"errors"
22
	"fmt"
Patrick Devine's avatar
Patrick Devine committed
23
	"io"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
24
	"net/http"
Michael Yang's avatar
Michael Yang committed
25
	"net/url"
Michael Yang's avatar
Michael Yang committed
26
	"runtime"
27
28
	"strconv"
	"time"
Michael Yang's avatar
Michael Yang committed
29

30
	"github.com/ollama/ollama/auth"
31
	"github.com/ollama/ollama/envconfig"
32
33
	"github.com/ollama/ollama/format"
	"github.com/ollama/ollama/version"
34
35
)

36
37
// Client encapsulates client state for interacting with the ollama
// service. Use [ClientFromEnvironment] to create new Clients.
Patrick Devine's avatar
Patrick Devine committed
38
type Client struct {
Michael Yang's avatar
Michael Yang committed
39
	base *url.URL
Michael Yang's avatar
Michael Yang committed
40
	http *http.Client
Michael Yang's avatar
Michael Yang committed
41
42
}

Patrick Devine's avatar
Patrick Devine committed
43
func checkError(resp *http.Response, body []byte) error {
Michael Yang's avatar
Michael Yang committed
44
	if resp.StatusCode < http.StatusBadRequest {
Patrick Devine's avatar
Patrick Devine committed
45
		return nil
Michael Yang's avatar
Michael Yang committed
46
47
	}

48
49
50
51
52
53
	if resp.StatusCode == http.StatusUnauthorized {
		authError := AuthorizationError{StatusCode: resp.StatusCode}
		json.Unmarshal(body, &authError)
		return authError
	}

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

Patrick Devine's avatar
Patrick Devine committed
56
57
58
	err := json.Unmarshal(body, &apiError)
	if err != nil {
		// Use the full body as the message if we fail to decode a response.
59
		apiError.ErrorMessage = string(body)
Patrick Devine's avatar
Patrick Devine committed
60
61
62
	}

	return apiError
Michael Yang's avatar
Michael Yang committed
63
64
}

65
66
// ClientFromEnvironment creates a new [Client] using configuration from the
// environment variable OLLAMA_HOST, which points to the network host and
67
// port on which the ollama service is listening. The format of this variable
68
69
70
71
72
73
// is:
//
//	<scheme>://<host>:<port>
//
// If the variable is not specified, a default ollama host and port will be
// used.
Michael Yang's avatar
Michael Yang committed
74
func ClientFromEnvironment() (*Client, error) {
75
	return &Client{
Michael Yang's avatar
host  
Michael Yang committed
76
		base: envconfig.Host(),
77
78
79
80
		http: http.DefaultClient,
	}, nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
81
82
83
84
85
86
87
func NewClient(base *url.URL, http *http.Client) *Client {
	return &Client{
		base: base,
		http: http,
	}
}

88
89
90
91
92
93
94
95
func getAuthorizationToken(ctx context.Context, challenge string) (string, error) {
	token, err := auth.Sign(ctx, []byte(challenge))
	if err != nil {
		return "", err
	}
	return token, nil
}

Patrick Devine's avatar
Patrick Devine committed
96
97
98
99
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
100
101
102
103
104
105
106
107

	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
108
109
110
111
		data, err = json.Marshal(reqData)
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
112

Patrick Devine's avatar
Patrick Devine committed
113
114
115
		reqBody = bytes.NewReader(data)
	}

Michael Yang's avatar
Michael Yang committed
116
	requestURL := c.base.JoinPath(path)
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

	var token string
	if envconfig.UseAuth() || c.base.Hostname() == "ollama.com" {
		now := strconv.FormatInt(time.Now().Unix(), 10)
		chal := fmt.Sprintf("%s,%s?ts=%s", method, path, now)
		token, err = getAuthorizationToken(ctx, chal)
		if err != nil {
			return err
		}

		q := requestURL.Query()
		q.Set("ts", now)
		requestURL.RawQuery = q.Encode()
	}

Michael Yang's avatar
Michael Yang committed
132
	request, err := http.NewRequestWithContext(ctx, method, requestURL.String(), reqBody)
Patrick Devine's avatar
Patrick Devine committed
133
134
135
136
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
137
138
139
	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
140

141
142
143
144
	if token != "" {
		request.Header.Set("Authorization", token)
	}

Michael Yang's avatar
Michael Yang committed
145
	respObj, err := c.http.Do(request)
Patrick Devine's avatar
Patrick Devine committed
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
	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
166
167
}

Michael Yang's avatar
Michael Yang committed
168
const maxBufferSize = 512 * format.KiloByte
169

Michael Yang's avatar
Michael Yang committed
170
func (c *Client) stream(ctx context.Context, method, path string, data any, fn func([]byte) error) error {
171
	var buf io.Reader
172
173
174
175
176
	if data != nil {
		bts, err := json.Marshal(data)
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
177

178
		buf = bytes.NewBuffer(bts)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
179
180
	}

Michael Yang's avatar
Michael Yang committed
181
	requestURL := c.base.JoinPath(path)
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197

	var token string
	if envconfig.UseAuth() || c.base.Hostname() == "ollama.com" {
		var err error
		now := strconv.FormatInt(time.Now().Unix(), 10)
		chal := fmt.Sprintf("%s,%s?ts=%s", method, path, now)
		token, err = getAuthorizationToken(ctx, chal)
		if err != nil {
			return err
		}

		q := requestURL.Query()
		q.Set("ts", now)
		requestURL.RawQuery = q.Encode()
	}

Michael Yang's avatar
Michael Yang committed
198
	request, err := http.NewRequestWithContext(ctx, method, requestURL.String(), buf)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
199
200
201
202
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
203
	request.Header.Set("Content-Type", "application/json")
204
	request.Header.Set("Accept", "application/x-ndjson")
Michael Yang's avatar
Michael Yang committed
205
	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
206

207
208
209
210
	if token != "" {
		request.Header.Set("Authorization", token)
	}

Michael Yang's avatar
Michael Yang committed
211
	response, err := c.http.Do(request)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
212
213
214
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
215
	defer response.Body.Close()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
216

217
	scanner := bufio.NewScanner(response.Body)
218
219
220
	// increase the buffer size to avoid running out of space
	scanBuf := make([]byte, 0, maxBufferSize)
	scanner.Buffer(scanBuf, maxBufferSize)
221
222
	for scanner.Scan() {
		var errorResponse struct {
223
224
			Error     string `json:"error,omitempty"`
			SigninURL string `json:"signin_url,omitempty"`
225
226
227
228
229
230
231
		}

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

232
233
234
235
		if response.StatusCode == http.StatusUnauthorized {
			return AuthorizationError{
				StatusCode: response.StatusCode,
				Status:     response.Status,
236
				SigninURL:  errorResponse.SigninURL,
237
238
			}
		} else if response.StatusCode >= http.StatusBadRequest {
Michael Yang's avatar
Michael Yang committed
239
			return StatusError{
240
241
242
				StatusCode:   response.StatusCode,
				Status:       response.Status,
				ErrorMessage: errorResponse.Error,
Michael Yang's avatar
Michael Yang committed
243
			}
244
245
		}

246
247
248
249
		if errorResponse.Error != "" {
			return errors.New(errorResponse.Error)
		}

Michael Yang's avatar
Michael Yang committed
250
		if err := fn(bts); err != nil {
251
			return err
Jeffrey Morgan's avatar
Jeffrey Morgan committed
252
253
254
		}
	}

Michael Yang's avatar
Michael Yang committed
255
256
	return nil
}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
257

258
259
260
// GenerateResponseFunc is a function that [Client.Generate] invokes every time
// a response is received from the service. If this function returns an error,
// [Client.Generate] will stop generating and return this error.
Michael Yang's avatar
Michael Yang committed
261
type GenerateResponseFunc func(GenerateResponse) error
Jeffrey Morgan's avatar
Jeffrey Morgan committed
262

263
264
265
// Generate generates a response for a given prompt. The req parameter should
// be populated with prompt details. fn is called for each response (there may
// be multiple responses, e.g. in case streaming is enabled).
Michael Yang's avatar
Michael Yang committed
266
func (c *Client) Generate(ctx context.Context, req *GenerateRequest, fn GenerateResponseFunc) error {
267
268
269
270
271
272
273
274
	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
275
}
Bruce MacDonald's avatar
Bruce MacDonald committed
276

277
278
279
// ChatResponseFunc is a function that [Client.Chat] invokes every time
// a response is received from the service. If this function returns an error,
// [Client.Chat] will stop generating and return this error.
Bruce MacDonald's avatar
Bruce MacDonald committed
280
281
type ChatResponseFunc func(ChatResponse) error

282
283
284
285
// Chat generates the next message in a chat. [ChatRequest] may contain a
// sequence of messages which can be used to maintain chat history with a model.
// fn is called for each response (there may be multiple responses, e.g. if case
// streaming is enabled).
Bruce MacDonald's avatar
Bruce MacDonald committed
286
287
288
289
290
291
292
293
294
295
296
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)
	})
}

297
298
299
// PullProgressFunc is a function that [Client.Pull] invokes every time there
// is progress with a "pull" request sent to the service. If this function
// returns an error, [Client.Pull] will stop the process and return this error.
300
type PullProgressFunc func(ProgressResponse) error
Michael Yang's avatar
Michael Yang committed
301

302
303
304
// Pull downloads a model from the ollama library. fn is called each time
// progress is made on the request and can be used to display a progress bar,
// etc.
Michael Yang's avatar
Michael Yang committed
305
func (c *Client) Pull(ctx context.Context, req *PullRequest, fn PullProgressFunc) error {
306
	return c.stream(ctx, http.MethodPost, "/api/pull", req, func(bts []byte) error {
307
		var resp ProgressResponse
308
309
310
311
312
313
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
Bruce MacDonald's avatar
Bruce MacDonald committed
314
}
315

316
317
318
// PushProgressFunc is a function that [Client.Push] invokes when progress is
// made.
// It's similar to other progress function types like [PullProgressFunc].
319
type PushProgressFunc func(ProgressResponse) error
320

321
322
323
// Push uploads a model to the model library; requires registering for ollama.ai
// and adding a public key first. fn is called each time progress is made on
// the request and can be used to display a progress bar, etc.
324
325
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 {
326
		var resp ProgressResponse
327
328
329
330
331
332
333
334
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
}

335
336
337
// CreateProgressFunc is a function that [Client.Create] invokes when progress
// is made.
// It's similar to other progress function types like [PullProgressFunc].
338
type CreateProgressFunc func(ProgressResponse) error
339

340
341
342
343
// Create creates a model from a [Modelfile]. fn is a progress function that
// behaves similarly to other methods (see [Client.Pull]).
//
// [Modelfile]: https://github.com/ollama/ollama/blob/main/docs/modelfile.md
344
345
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 {
346
		var resp ProgressResponse
347
348
349
350
351
352
353
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
}
Patrick Devine's avatar
Patrick Devine committed
354

355
// List lists models that are available locally.
Patrick Devine's avatar
Patrick Devine committed
356
357
358
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 {
359
360
361
362
363
		return nil, err
	}
	return &lr, nil
}

364
// ListRunning lists running models.
365
366
func (c *Client) ListRunning(ctx context.Context) (*ProcessResponse, error) {
	var lr ProcessResponse
367
	if err := c.do(ctx, http.MethodGet, "/api/ps", nil, &lr); err != nil {
Patrick Devine's avatar
Patrick Devine committed
368
369
370
371
		return nil, err
	}
	return &lr, nil
}
372

373
374
// Copy copies a model - creating a model with another name from an existing
// model.
Patrick Devine's avatar
Patrick Devine committed
375
376
377
378
379
380
381
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
}

382
// Delete deletes a model and its data.
383
384
385
386
387
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
388
}
389

390
// Show obtains model information, including details, modelfile, license etc.
Patrick Devine's avatar
Patrick Devine committed
391
392
393
394
395
396
397
398
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
}

399
// Heartbeat checks if the server has started and is responsive; if yes, it
400
// returns nil, otherwise an error.
401
func (c *Client) Heartbeat(ctx context.Context) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
402
	if err := c.do(ctx, http.MethodHead, "/", nil, nil); err != nil {
403
404
405
406
		return err
	}
	return nil
}
407

408
409
410
411
412
413
414
415
416
417
// Embed generates embeddings from a model.
func (c *Client) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) {
	var resp EmbedResponse
	if err := c.do(ctx, http.MethodPost, "/api/embed", req, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

// Embeddings generates an embedding from a model.
Brian Murray's avatar
Brian Murray committed
418
419
420
421
422
423
424
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
425

426
427
// CreateBlob creates a blob from a file on the server. digest is the
// expected SHA256 digest of the file, and r represents the file.
Michael Yang's avatar
Michael Yang committed
428
func (c *Client) CreateBlob(ctx context.Context, digest string, r io.Reader) error {
429
	return c.do(ctx, http.MethodPost, fmt.Sprintf("/api/blobs/%s", digest), r, nil)
Michael Yang's avatar
Michael Yang committed
430
}
Michael Yang's avatar
Michael Yang committed
431

432
// Version returns the Ollama server version as a string.
Michael Yang's avatar
Michael Yang committed
433
434
435
436
437
438
439
440
441
442
443
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
}
444

445
446
447
448
449
450
451
// Signout will signout a client for a local ollama server.
func (c *Client) Signout(ctx context.Context) error {
	return c.do(ctx, http.MethodPost, "/api/signout", nil, nil)
}

// Disconnect will disconnect an ollama instance from ollama.com.
func (c *Client) Disconnect(ctx context.Context, encodedKey string) error {
452
453
454
455
456
457
458
459
460
461
	return c.do(ctx, http.MethodDelete, fmt.Sprintf("/api/user/keys/%s", encodedKey), nil, nil)
}

func (c *Client) Whoami(ctx context.Context) (*UserResponse, error) {
	var resp UserResponse
	if err := c.do(ctx, http.MethodPost, "/api/me", nil, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}