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

Michael Yang's avatar
Michael Yang committed
17
	"github.com/jmorganca/ollama/format"
Michael Yang's avatar
Michael Yang committed
18
	"github.com/jmorganca/ollama/version"
19
20
)

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

var envHost = os.Getenv("OLLAMA_HOST")

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
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),
		},
65
66
	}

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

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

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

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

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

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

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

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

220
type PushProgressFunc func(ProgressResponse) error
221
222
223

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

		return fn(resp)
	})
}

233
type CreateProgressFunc func(ProgressResponse) error
234
235
236

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

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

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

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

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

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

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