third-party-mirrors/ollama
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

registry: move req/resp types to registry/apitype

Browse Source
This commit is contained in:
Blake Mizerany 2024-03-31 12:23:10 -07:00
parent eb2c442a01
commit 48c60c01e2
4 changed files with 13 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
registry/apitypes.go → registry/apitype/apitype.go
View File

@ -1,4 +1,4 @@
package registry
package apitype
import "encoding/json"

5
registry/client.go
View File

@ -6,6 +6,7 @@ import (
"net/http"
"bllamo.com/client/ollama"
"bllamo.com/registry/apitype"
)
type Client struct {
@ -18,9 +19,9 @@ func (c *Client) oclient() *ollama.Client {
}
// Push pushes a manifest to the server.
func (c *Client) Push(ctx context.Context, ref string, manifest []byte) ([]Requirement, error) {
func (c *Client) Push(ctx context.Context, ref string, manifest []byte) ([]apitype.Requirement, error) {
// TODO(bmizerany): backoff
v, err := ollama.Do[PushResponse](ctx, c.oclient(), "POST", "/v1/push", &PushRequest{
v, err := ollama.Do[apitype.PushResponse](ctx, c.oclient(), "POST", "/v1/push", &apitype.PushRequest{
Ref: ref,
Manifest: manifest,
})

11
registry/server.go
View File

@ -13,6 +13,7 @@ import (
"bllamo.com/client/ollama"
"bllamo.com/oweb"
"bllamo.com/registry/apitype"
"github.com/minio/minio-go/v7"
"github.com/minio/minio-go/v7/pkg/credentials"
)
@ -44,7 +45,7 @@ func (s *Server) serveHTTP(w http.ResponseWriter, r *http.Request) error {
}
func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error {
pr, err := oweb.DecodeUserJSON[PushRequest]("", r.Body)
pr, err := oweb.DecodeUserJSON[apitype.PushRequest]("", r.Body)
if err != nil {
return err
}
@ -54,13 +55,13 @@ func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error {
Secure: false,
})
m, err := oweb.DecodeUserJSON[Manifest]("manifest", bytes.NewReader(pr.Manifest))
m, err := oweb.DecodeUserJSON[apitype.Manifest]("manifest", bytes.NewReader(pr.Manifest))
if err != nil {
return err
}
// TODO(bmizerany): parallelize
var requirements []Requirement
var requirements []apitype.Requirement
for _, l := range m.Layers {
if l.Size == 0 {
continue
@ -76,7 +77,7 @@ func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error {
if err != nil {
return err
}
requirements = append(requirements, Requirement{
requirements = append(requirements, apitype.Requirement{
Digest: l.Digest,
Size: l.Size,
@ -89,7 +90,7 @@ func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error {
// TODO(bmizerany): commit to db
// ref, _ := strings.CutPrefix(r.URL.Path, "/v1/push/")
return oweb.EncodeJSON(w, &PushResponse{Requirements: requirements})
return oweb.EncodeJSON(w, &apitype.PushResponse{Requirements: requirements})
}
func (s *Server) handlePull(w http.ResponseWriter, r *http.Request) error {

5
registry/server_test.go
View File

@ -8,6 +8,7 @@ import (
"testing"
"time"
"bllamo.com/registry/apitype"
"github.com/kr/pretty"
"github.com/minio/minio-go/v7"
"github.com/minio/minio-go/v7/pkg/credentials"
@ -35,11 +36,11 @@ func TestPush(t *testing.T) {
t.Fatal(err)
}
diff.Test(t, t.Errorf, got, []Requirement{
diff.Test(t, t.Errorf, got, []apitype.Requirement{
{Digest: "sha256-1", Size: 1},
{Digest: "sha256-2", Size: 2},
{Digest: "sha256-3", Size: 3},
}, diff.ZeroFields[Requirement]("URL"))
}, diff.ZeroFields[apitype.Requirement]("URL"))
for _, r := range got {
body := strings.NewReader(strings.Repeat("x", int(r.Size)))