consistent error handling for pull and generate

This commit is contained in:
Jeffrey Morgan 2023-07-10 21:34:15 -07:00
parent 407a5cabf4
commit a3ec1ec2a0
3 changed files with 42 additions and 111 deletions

View file

@ -6,7 +6,6 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"io"
"net/http" "net/http"
"net/url" "net/url"
) )
@ -26,47 +25,18 @@ func NewClient(hosts ...string) *Client {
} }
} }
func StatusError(status int, message ...string) error { func (c *Client) stream(ctx context.Context, method, path string, data any, callback func([]byte) error) error {
if status < 400 { var buf *bytes.Buffer
return nil if data != nil {
bts, err := json.Marshal(data)
if err != nil {
return err
}
buf = bytes.NewBuffer(bts)
} }
if len(message) > 0 && len(message[0]) > 0 { request, err := http.NewRequestWithContext(ctx, method, c.base.JoinPath(path).String(), buf)
return fmt.Errorf("%d %s: %s", status, http.StatusText(status), message[0])
}
return fmt.Errorf("%d %s", status, http.StatusText(status))
}
type options struct {
requestBody io.Reader
responseFunc func(bts []byte) error
}
func OptionRequestBody(data any) func(*options) {
bts, err := json.Marshal(data)
if err != nil {
panic(err)
}
return func(opts *options) {
opts.requestBody = bytes.NewReader(bts)
}
}
func OptionResponseFunc(fn func([]byte) error) func(*options) {
return func(opts *options) {
opts.responseFunc = fn
}
}
func (c *Client) stream(ctx context.Context, method, path string, fns ...func(*options)) error {
var opts options
for _, fn := range fns {
fn(&opts)
}
request, err := http.NewRequestWithContext(ctx, method, c.base.JoinPath(path).String(), opts.requestBody)
if err != nil { if err != nil {
return err return err
} }
@ -80,25 +50,23 @@ func (c *Client) stream(ctx context.Context, method, path string, fns ...func(*o
} }
defer response.Body.Close() defer response.Body.Close()
if opts.responseFunc != nil { scanner := bufio.NewScanner(response.Body)
scanner := bufio.NewScanner(response.Body) for scanner.Scan() {
for scanner.Scan() { var errorResponse struct {
var errorResponse struct { Error string `json:"error"`
Error string `json:"error"` }
}
bts := scanner.Bytes() bts := scanner.Bytes()
if err := json.Unmarshal(bts, &errorResponse); err != nil { if err := json.Unmarshal(bts, &errorResponse); err != nil {
return err return fmt.Errorf("unmarshal: %w", err)
} }
if err := StatusError(response.StatusCode, errorResponse.Error); err != nil { if len(errorResponse.Error) > 0 {
return err return fmt.Errorf("stream: %s", errorResponse.Error)
} }
if err := opts.responseFunc(bts); err != nil { if err := callback(bts); err != nil {
return err return err
}
} }
} }
@ -108,36 +76,25 @@ func (c *Client) stream(ctx context.Context, method, path string, fns ...func(*o
type GenerateResponseFunc func(GenerateResponse) error type GenerateResponseFunc func(GenerateResponse) error
func (c *Client) Generate(ctx context.Context, req *GenerateRequest, fn GenerateResponseFunc) error { func (c *Client) Generate(ctx context.Context, req *GenerateRequest, fn GenerateResponseFunc) error {
return c.stream(ctx, http.MethodPost, "/api/generate", return c.stream(ctx, http.MethodPost, "/api/generate", req, func(bts []byte) error {
OptionRequestBody(req), var resp GenerateResponse
OptionResponseFunc(func(bts []byte) error { if err := json.Unmarshal(bts, &resp); err != nil {
var resp GenerateResponse return err
if err := json.Unmarshal(bts, &resp); err != nil { }
return err
}
return fn(resp) return fn(resp)
}), })
)
} }
type PullProgressFunc func(PullProgress) error type PullProgressFunc func(PullProgress) error
func (c *Client) Pull(ctx context.Context, req *PullRequest, fn PullProgressFunc) error { func (c *Client) Pull(ctx context.Context, req *PullRequest, fn PullProgressFunc) error {
return c.stream(ctx, http.MethodPost, "/api/pull", return c.stream(ctx, http.MethodPost, "/api/pull", req, func(bts []byte) error {
OptionRequestBody(req), var resp PullProgress
OptionResponseFunc(func(bts []byte) error { if err := json.Unmarshal(bts, &resp); err != nil {
var resp PullProgress return err
if err := json.Unmarshal(bts, &resp); err != nil { }
return err
}
if resp.Error.Message != "" { return fn(resp)
// couldn't pull the model from the directory, proceed anyway })
return nil
}
return fn(resp)
}),
)
} }

View file

@ -1,24 +1,5 @@
package api package api
import (
"fmt"
"net/http"
"strings"
)
type Error struct {
Code int32 `json:"code"`
Message string `json:"message"`
}
func (e Error) Error() string {
if e.Message == "" {
return fmt.Sprintf("%d %v", e.Code, strings.ToLower(http.StatusText(int(e.Code))))
}
return e.Message
}
type PullRequest struct { type PullRequest struct {
Model string `json:"model"` Model string `json:"model"`
} }
@ -27,7 +8,6 @@ type PullProgress struct {
Total int64 `json:"total"` Total int64 `json:"total"`
Completed int64 `json:"completed"` Completed int64 `json:"completed"`
Percent float64 `json:"percent"` Percent float64 `json:"percent"`
Error Error `json:"error"`
} }
type GenerateRequest struct { type GenerateRequest struct {

View file

@ -54,7 +54,7 @@ func generate(c *gin.Context) {
} }
if _, err := os.Stat(req.Model); err != nil { if _, err := os.Stat(req.Model); err != nil {
if !errors.Is(err, os.ErrNotExist) { if !errors.Is(err, os.ErrNotExist) {
c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return return
} }
req.Model = path.Join(cacheDir(), "models", req.Model+".bin") req.Model = path.Join(cacheDir(), "models", req.Model+".bin")
@ -136,7 +136,7 @@ func Serve(ln net.Listener) error {
r.POST("api/pull", func(c *gin.Context) { r.POST("api/pull", func(c *gin.Context) {
var req api.PullRequest var req api.PullRequest
if err := c.ShouldBindJSON(&req); err != nil { if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return return
} }
@ -146,16 +146,10 @@ func Serve(ln net.Listener) error {
if err := pull(req.Model, progressCh); err != nil { if err := pull(req.Model, progressCh); err != nil {
var opError *net.OpError var opError *net.OpError
if errors.As(err, &opError) { if errors.As(err, &opError) {
result := api.PullProgress{ c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
Error: api.Error{
Code: http.StatusBadGateway,
Message: "failed to get models from directory",
},
}
c.JSON(http.StatusBadGateway, result)
return return
} }
c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return return
} }
}() }()