client.go 6.38 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

	"github.com/jmorganca/ollama/version"
18
19
)

20
21
22
23
const DefaultHost = "127.0.0.1:11434"

var envHost = os.Getenv("OLLAMA_HOST")

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

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

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

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

	return apiError
Michael Yang's avatar
Michael Yang committed
43
44
}

Michael Yang's avatar
Michael Yang committed
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
func ClientFromEnvironment() (*Client, error) {
	scheme, hostport, ok := strings.Cut(os.Getenv("OLLAMA_HOST"), "://")
	if !ok {
		scheme, hostport = "http", os.Getenv("OLLAMA_HOST")
	}

	host, port, err := net.SplitHostPort(hostport)
	if err != nil {
		host, port = "127.0.0.1", "11434"
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
			host = ip.String()
		}
	}

	client := Client{
		base: &url.URL{
			Scheme: scheme,
			Host:   net.JoinHostPort(host, port),
		},
64
65
	}

Michael Yang's avatar
Michael Yang committed
66
67
68
	mockRequest, err := http.NewRequest("HEAD", client.base.String(), nil)
	if err != nil {
		return nil, err
69
70
	}

Michael Yang's avatar
Michael Yang committed
71
	proxyURL, err := http.ProxyFromEnvironment(mockRequest)
72
	if err != nil {
Michael Yang's avatar
Michael Yang committed
73
		return nil, err
Michael Yang's avatar
Michael Yang committed
74
75
	}

Michael Yang's avatar
Michael Yang committed
76
77
78
79
	client.http = http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyURL(proxyURL),
		},
Patrick Devine's avatar
Patrick Devine committed
80
	}
81

Michael Yang's avatar
Michael Yang committed
82
	return &client, nil
Patrick Devine's avatar
Patrick Devine committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
}

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

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 * 1000 // 512KB
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

206
type PullProgressFunc func(ProgressResponse) error
Michael Yang's avatar
Michael Yang committed
207
208

func (c *Client) Pull(ctx context.Context, req *PullRequest, fn PullProgressFunc) error {
209
	return c.stream(ctx, http.MethodPost, "/api/pull", req, func(bts []byte) error {
210
		var resp ProgressResponse
211
212
213
214
215
216
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
Bruce MacDonald's avatar
Bruce MacDonald committed
217
}
218

219
type PushProgressFunc func(ProgressResponse) error
220
221
222

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 {
223
		var resp ProgressResponse
224
225
226
227
228
229
230
231
		if err := json.Unmarshal(bts, &resp); err != nil {
			return err
		}

		return fn(resp)
	})
}

232
type CreateProgressFunc func(ProgressResponse) error
233
234
235

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

		return fn(resp)
	})
}
Patrick Devine's avatar
Patrick Devine committed
244
245
246
247
248
249
250
251

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

Patrick Devine's avatar
Patrick Devine committed
253
254
255
256
257
258
259
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
}

260
261
262
263
264
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
265
}
266

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

275
func (c *Client) Heartbeat(ctx context.Context) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
276
	if err := c.do(ctx, http.MethodHead, "/", nil, nil); err != nil {
277
278
279
280
		return err
	}
	return nil
}