345 lines
8 KiB
Go
345 lines
8 KiB
Go
package api
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"runtime"
|
|
"strings"
|
|
|
|
"github.com/jmorganca/ollama/format"
|
|
"github.com/jmorganca/ollama/version"
|
|
)
|
|
|
|
type Client struct {
|
|
base *url.URL
|
|
http http.Client
|
|
}
|
|
|
|
func checkError(resp *http.Response, body []byte) error {
|
|
if resp.StatusCode < http.StatusBadRequest {
|
|
return nil
|
|
}
|
|
|
|
apiError := StatusError{StatusCode: resp.StatusCode}
|
|
|
|
err := json.Unmarshal(body, &apiError)
|
|
if err != nil {
|
|
// Use the full body as the message if we fail to decode a response.
|
|
apiError.ErrorMessage = string(body)
|
|
}
|
|
|
|
return apiError
|
|
}
|
|
|
|
func ClientFromEnvironment() (*Client, error) {
|
|
defaultPort := "11434"
|
|
|
|
scheme, hostport, ok := strings.Cut(os.Getenv("OLLAMA_HOST"), "://")
|
|
switch {
|
|
case !ok:
|
|
scheme, hostport = "http", os.Getenv("OLLAMA_HOST")
|
|
case scheme == "http":
|
|
defaultPort = "80"
|
|
case scheme == "https":
|
|
defaultPort = "443"
|
|
}
|
|
|
|
// trim trailing slashes
|
|
hostport = strings.TrimRight(hostport, "/")
|
|
|
|
host, port, err := net.SplitHostPort(hostport)
|
|
if err != nil {
|
|
host, port = "127.0.0.1", defaultPort
|
|
if ip := net.ParseIP(strings.Trim(hostport, "[]")); ip != nil {
|
|
host = ip.String()
|
|
} else if hostport != "" {
|
|
host = hostport
|
|
}
|
|
}
|
|
|
|
client := Client{
|
|
base: &url.URL{
|
|
Scheme: scheme,
|
|
Host: net.JoinHostPort(host, port),
|
|
},
|
|
}
|
|
|
|
mockRequest, err := http.NewRequest(http.MethodHead, client.base.String(), nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
proxyURL, err := http.ProxyFromEnvironment(mockRequest)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
client.http = http.Client{
|
|
Transport: &http.Transport{
|
|
Proxy: http.ProxyURL(proxyURL),
|
|
},
|
|
}
|
|
|
|
return &client, nil
|
|
}
|
|
|
|
func (c *Client) do(ctx context.Context, method, path string, reqData, respData any) error {
|
|
var reqBody io.Reader
|
|
var data []byte
|
|
var err error
|
|
|
|
switch reqData := reqData.(type) {
|
|
case io.Reader:
|
|
// reqData is already an io.Reader
|
|
reqBody = reqData
|
|
case nil:
|
|
// noop
|
|
default:
|
|
data, err = json.Marshal(reqData)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
reqBody = bytes.NewReader(data)
|
|
}
|
|
|
|
requestURL := c.base.JoinPath(path)
|
|
request, err := http.NewRequestWithContext(ctx, method, requestURL.String(), reqBody)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
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()))
|
|
|
|
respObj, err := c.http.Do(request)
|
|
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
|
|
}
|
|
|
|
const maxBufferSize = 512 * format.KiloByte
|
|
|
|
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
|
|
}
|
|
|
|
buf = bytes.NewBuffer(bts)
|
|
}
|
|
|
|
requestURL := c.base.JoinPath(path)
|
|
request, err := http.NewRequestWithContext(ctx, method, requestURL.String(), buf)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
request.Header.Set("Content-Type", "application/json")
|
|
request.Header.Set("Accept", "application/x-ndjson")
|
|
request.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
|
|
|
|
response, err := c.http.Do(request)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer response.Body.Close()
|
|
|
|
scanner := bufio.NewScanner(response.Body)
|
|
// increase the buffer size to avoid running out of space
|
|
scanBuf := make([]byte, 0, maxBufferSize)
|
|
scanner.Buffer(scanBuf, maxBufferSize)
|
|
for scanner.Scan() {
|
|
var errorResponse struct {
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
bts := scanner.Bytes()
|
|
if err := json.Unmarshal(bts, &errorResponse); err != nil {
|
|
return fmt.Errorf("unmarshal: %w", err)
|
|
}
|
|
|
|
if errorResponse.Error != "" {
|
|
return fmt.Errorf(errorResponse.Error)
|
|
}
|
|
|
|
if response.StatusCode >= http.StatusBadRequest {
|
|
return StatusError{
|
|
StatusCode: response.StatusCode,
|
|
Status: response.Status,
|
|
ErrorMessage: errorResponse.Error,
|
|
}
|
|
}
|
|
|
|
if err := fn(bts); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
type GenerateResponseFunc func(GenerateResponse) error
|
|
|
|
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)
|
|
})
|
|
}
|
|
|
|
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)
|
|
})
|
|
}
|
|
|
|
type PullProgressFunc func(ProgressResponse) error
|
|
|
|
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 {
|
|
var resp ProgressResponse
|
|
if err := json.Unmarshal(bts, &resp); err != nil {
|
|
return err
|
|
}
|
|
|
|
return fn(resp)
|
|
})
|
|
}
|
|
|
|
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 {
|
|
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)
|
|
})
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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 {
|
|
if err := c.do(ctx, http.MethodHead, "/", nil, nil); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
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
|
|
}
|
|
|
|
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 {
|
|
var statusError StatusError
|
|
if !errors.As(err, &statusError) || statusError.StatusCode != http.StatusNotFound {
|
|
return err
|
|
}
|
|
|
|
if err := c.do(ctx, http.MethodPost, fmt.Sprintf("/api/blobs/%s", digest), r, nil); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
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
|
|
}
|