ollama/api/client.go

272 lines
6.2 KiB
Go
Raw Normal View History

package api
import (
2023-07-04 04:47:00 +00:00
"bufio"
"bytes"
"context"
"encoding/json"
2023-07-07 21:04:43 +00:00
"fmt"
2023-07-18 16:09:45 +00:00
"io"
"net/http"
2023-07-06 22:02:10 +00:00
"net/url"
"os"
2023-08-22 01:24:42 +00:00
"runtime"
2023-08-17 22:20:38 +00:00
"strings"
2023-08-22 01:24:42 +00:00
"github.com/jmorganca/ollama/version"
)
const DefaultHost = "127.0.0.1:11434"
var (
envHost = os.Getenv("OLLAMA_HOST")
)
2023-07-18 16:09:45 +00:00
type Client struct {
Base url.URL
2023-07-18 16:09:45 +00:00
HTTP http.Client
Headers http.Header
2023-07-11 20:05:51 +00:00
}
2023-07-18 16:09:45 +00:00
func checkError(resp *http.Response, body []byte) error {
2023-08-27 04:55:21 +00:00
if resp.StatusCode < http.StatusBadRequest {
2023-07-18 16:09:45 +00:00
return nil
2023-07-11 20:05:51 +00:00
}
2023-07-18 16:09:45 +00:00
apiError := StatusError{StatusCode: resp.StatusCode}
2023-07-11 20:05:51 +00:00
2023-07-18 16:09:45 +00:00
err := json.Unmarshal(body, &apiError)
if err != nil {
// Use the full body as the message if we fail to decode a response.
2023-07-20 18:45:12 +00:00
apiError.ErrorMessage = string(body)
2023-07-18 16:09:45 +00:00
}
return apiError
2023-07-06 22:02:10 +00:00
}
// Host returns the default host to use for the client. It is determined in the following order:
// 1. The OLLAMA_HOST environment variable
// 2. The default host (localhost:11434)
func Host() string {
if envHost != "" {
return envHost
}
return DefaultHost
}
// FromEnv creates a new client using Host() as the host. An error is returns
// if the host is invalid.
func FromEnv() (*Client, error) {
2023-08-17 22:20:38 +00:00
h := Host()
if !strings.HasPrefix(h, "http://") && !strings.HasPrefix(h, "https://") {
h = "http://" + h
}
2023-08-17 22:20:38 +00:00
u, err := url.Parse(h)
if err != nil {
return nil, fmt.Errorf("could not parse host: %w", err)
2023-07-06 22:02:10 +00:00
}
2023-08-17 22:20:38 +00:00
if u.Port() == "" {
u.Host += ":11434"
2023-07-18 16:09:45 +00:00
}
2023-08-17 22:20:38 +00:00
return &Client{Base: *u, HTTP: http.Client{}}, nil
2023-07-18 16:09:45 +00:00
}
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)
}
2023-08-22 01:24:42 +00:00
requestURL := c.Base.JoinPath(path)
request, err := http.NewRequestWithContext(ctx, method, requestURL.String(), reqBody)
2023-07-18 16:09:45 +00:00
if err != nil {
return err
}
2023-08-22 01:24:42 +00:00
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()))
2023-07-18 16:09:45 +00:00
for k, v := range c.Headers {
2023-08-22 01:24:42 +00:00
request.Header[k] = v
2023-07-06 22:02:10 +00:00
}
2023-07-18 16:09:45 +00:00
2023-08-22 01:24:42 +00:00
respObj, err := c.HTTP.Do(request)
2023-07-18 16:09:45 +00:00
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
}
2023-07-11 20:05:51 +00:00
func (c *Client) stream(ctx context.Context, method, path string, data any, fn func([]byte) error) error {
var buf *bytes.Buffer
if data != nil {
bts, err := json.Marshal(data)
if err != nil {
return err
}
2023-07-06 23:53:14 +00:00
buf = bytes.NewBuffer(bts)
}
2023-08-22 01:24:42 +00:00
requestURL := c.Base.JoinPath(path)
request, err := http.NewRequestWithContext(ctx, method, requestURL.String(), buf)
if err != nil {
return err
}
2023-07-06 22:02:10 +00:00
request.Header.Set("Content-Type", "application/json")
request.Header.Set("Accept", "application/json")
2023-08-22 01:24:42 +00:00
request.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
2023-07-04 04:47:00 +00:00
2023-07-06 22:02:10 +00:00
response, err := http.DefaultClient.Do(request)
2023-07-04 04:47:00 +00:00
if err != nil {
return err
}
2023-07-06 22:02:10 +00:00
defer response.Body.Close()
2023-07-04 04:47:00 +00:00
scanner := bufio.NewScanner(response.Body)
for scanner.Scan() {
var errorResponse struct {
2023-07-11 20:05:51 +00:00
Error string `json:"error,omitempty"`
}
bts := scanner.Bytes()
if err := json.Unmarshal(bts, &errorResponse); err != nil {
return fmt.Errorf("unmarshal: %w", err)
}
2023-07-20 19:12:08 +00:00
if errorResponse.Error != "" {
return fmt.Errorf(errorResponse.Error)
2023-07-20 19:12:08 +00:00
}
2023-08-27 04:55:21 +00:00
if response.StatusCode >= http.StatusBadRequest {
2023-07-11 20:05:51 +00:00
return StatusError{
2023-07-20 18:45:12 +00:00
StatusCode: response.StatusCode,
Status: response.Status,
ErrorMessage: errorResponse.Error,
2023-07-11 20:05:51 +00:00
}
}
2023-07-11 20:05:51 +00:00
if err := fn(bts); err != nil {
return err
}
}
2023-07-06 21:05:55 +00:00
return nil
}
2023-07-06 21:05:55 +00:00
type GenerateResponseFunc func(GenerateResponse) error
2023-07-06 21:05:55 +00:00
func (c *Client) Generate(ctx context.Context, req *GenerateRequest, fn GenerateResponseFunc) error {
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)
})
2023-07-04 04:47:00 +00:00
}
2023-07-06 16:24:49 +00:00
2023-07-19 01:51:30 +00:00
type PullProgressFunc func(ProgressResponse) error
2023-07-06 21:05:55 +00:00
func (c *Client) Pull(ctx context.Context, req *PullRequest, fn PullProgressFunc) error {
return c.stream(ctx, http.MethodPost, "/api/pull", req, func(bts []byte) error {
2023-07-19 01:51:30 +00:00
var resp ProgressResponse
if err := json.Unmarshal(bts, &resp); err != nil {
return err
}
return fn(resp)
})
2023-07-06 16:24:49 +00:00
}
2023-07-19 01:51:30 +00:00
type PushProgressFunc func(ProgressResponse) error
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 {
2023-07-19 01:51:30 +00:00
var resp ProgressResponse
if err := json.Unmarshal(bts, &resp); err != nil {
return err
}
return fn(resp)
})
}
type CreateProgressFunc func(ProgressResponse) error
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 {
var resp ProgressResponse
if err := json.Unmarshal(bts, &resp); err != nil {
return err
}
return fn(resp)
})
}
2023-07-18 16:09:45 +00:00
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
}
2023-07-20 23:09:23 +00:00
2023-07-24 15:27:28 +00:00
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
}
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
2023-07-20 23:09:23 +00:00
}
2023-09-06 18:04:17 +00:00
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
}
func (c *Client) Heartbeat(ctx context.Context) error {
2023-08-01 18:50:38 +00:00
if err := c.do(ctx, http.MethodHead, "/", nil, nil); err != nil {
return err
}
return nil
}