client.go 7.68 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"
Michael Yang's avatar
Michael Yang committed
8
	"errors"
9
	"fmt"
Patrick Devine's avatar
Patrick Devine committed
10
	"io"
Michael Yang's avatar
Michael Yang committed
11
	"net"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
12
	"net/http"
Michael Yang's avatar
Michael Yang committed
13
	"net/url"
14
	"os"
Michael Yang's avatar
Michael Yang committed
15
	"runtime"
16
	"strings"
Michael Yang's avatar
Michael Yang committed
17

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

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

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

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

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

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

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

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

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

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

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

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
82
83
84
85
86
87
88
89

	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
90
91
92
93
		data, err = json.Marshal(reqData)
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
94

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

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

Michael Yang's avatar
Michael Yang committed
104
105
106
	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
107

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
150
	request.Header.Set("Content-Type", "application/json")
151
	request.Header.Set("Accept", "application/x-ndjson")
Michael Yang's avatar
Michael Yang committed
152
	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
153

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

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

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

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

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

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

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

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

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

Bruce MacDonald's avatar
Bruce MacDonald committed
207
208
209
210
211
212
213
214
215
216
217
218
219
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)
	})
}

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

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

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

233
type PushProgressFunc func(ProgressResponse) error
234
235
236

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

		return fn(resp)
	})
}

246
type CreateProgressFunc func(ProgressResponse) error
247
248
249

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

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

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

Patrick Devine's avatar
Patrick Devine committed
267
268
269
270
271
272
273
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
}

274
275
276
277
278
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
279
}
280

Patrick Devine's avatar
Patrick Devine committed
281
282
283
284
285
286
287
288
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
}

289
func (c *Client) Heartbeat(ctx context.Context) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
290
	if err := c.do(ctx, http.MethodHead, "/", nil, nil); err != nil {
291
292
293
294
		return err
	}
	return nil
}
Brian Murray's avatar
Brian Murray committed
295
296
297
298
299
300
301
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
302

Michael Yang's avatar
Michael Yang committed
303
304
func (c *Client) CreateBlob(ctx context.Context, digest string, r io.Reader) error {
	if err := c.do(ctx, http.MethodHead, fmt.Sprintf("/api/blobs/%s", digest), nil, nil); err != nil {
Michael Yang's avatar
Michael Yang committed
305
306
		var statusError StatusError
		if !errors.As(err, &statusError) || statusError.StatusCode != http.StatusNotFound {
Michael Yang's avatar
Michael Yang committed
307
			return err
Michael Yang's avatar
Michael Yang committed
308
309
		}

Michael Yang's avatar
Michael Yang committed
310
311
		if err := c.do(ctx, http.MethodPost, fmt.Sprintf("/api/blobs/%s", digest), r, nil); err != nil {
			return err
Michael Yang's avatar
Michael Yang committed
312
313
314
		}
	}

Michael Yang's avatar
Michael Yang committed
315
	return nil
Michael Yang's avatar
Michael Yang committed
316
}
Michael Yang's avatar
Michael Yang committed
317
318
319
320
321
322
323
324
325
326
327
328

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
}