add API create/copy handlers (#1541)
This commit is contained in:
parent
f728738427
commit
86b0dd4b16
1 changed files with 75 additions and 12 deletions
|
@ -1,10 +1,12 @@
|
||||||
package server
|
package server
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"os"
|
"os"
|
||||||
|
@ -31,6 +33,20 @@ func Test_Routes(t *testing.T) {
|
||||||
Setup func(t *testing.T, req *http.Request)
|
Setup func(t *testing.T, req *http.Request)
|
||||||
Expected func(t *testing.T, resp *http.Response)
|
Expected func(t *testing.T, resp *http.Response)
|
||||||
}
|
}
|
||||||
|
var tempModelFile string
|
||||||
|
|
||||||
|
createTestModel := func(t *testing.T, name string) {
|
||||||
|
f, err := os.CreateTemp("", "ollama-model")
|
||||||
|
assert.Nil(t, err)
|
||||||
|
defer os.RemoveAll(f.Name())
|
||||||
|
|
||||||
|
modelfile := strings.NewReader(fmt.Sprintf("FROM %s", f.Name()))
|
||||||
|
commands, err := parser.Parse(modelfile)
|
||||||
|
assert.Nil(t, err)
|
||||||
|
fn := func(resp api.ProgressResponse) {}
|
||||||
|
err = CreateModel(context.TODO(), name, "", commands, fn)
|
||||||
|
assert.Nil(t, err)
|
||||||
|
}
|
||||||
|
|
||||||
testCases := []testCase{
|
testCases := []testCase{
|
||||||
{
|
{
|
||||||
|
@ -70,16 +86,7 @@ func Test_Routes(t *testing.T) {
|
||||||
Method: http.MethodGet,
|
Method: http.MethodGet,
|
||||||
Path: "/api/tags",
|
Path: "/api/tags",
|
||||||
Setup: func(t *testing.T, req *http.Request) {
|
Setup: func(t *testing.T, req *http.Request) {
|
||||||
f, err := os.CreateTemp("", "ollama-modelfile")
|
createTestModel(t, "test-model")
|
||||||
assert.Nil(t, err)
|
|
||||||
defer os.RemoveAll(f.Name())
|
|
||||||
|
|
||||||
modelfile := strings.NewReader(fmt.Sprintf("FROM %s", f.Name()))
|
|
||||||
commands, err := parser.Parse(modelfile)
|
|
||||||
assert.Nil(t, err)
|
|
||||||
fn := func(resp api.ProgressResponse) {}
|
|
||||||
err = CreateModel(context.TODO(), "test", "", commands, fn)
|
|
||||||
assert.Nil(t, err)
|
|
||||||
},
|
},
|
||||||
Expected: func(t *testing.T, resp *http.Response) {
|
Expected: func(t *testing.T, resp *http.Response) {
|
||||||
contentType := resp.Header.Get("Content-Type")
|
contentType := resp.Header.Get("Content-Type")
|
||||||
|
@ -88,12 +95,66 @@ func Test_Routes(t *testing.T) {
|
||||||
assert.Nil(t, err)
|
assert.Nil(t, err)
|
||||||
|
|
||||||
var modelList api.ListResponse
|
var modelList api.ListResponse
|
||||||
|
|
||||||
err = json.Unmarshal(body, &modelList)
|
err = json.Unmarshal(body, &modelList)
|
||||||
assert.Nil(t, err)
|
assert.Nil(t, err)
|
||||||
|
|
||||||
assert.Equal(t, 1, len(modelList.Models))
|
assert.Equal(t, 1, len(modelList.Models))
|
||||||
assert.Equal(t, modelList.Models[0].Name, "test:latest")
|
assert.Equal(t, modelList.Models[0].Name, "test-model:latest")
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "Create Model Handler",
|
||||||
|
Method: http.MethodPost,
|
||||||
|
Path: "/api/create",
|
||||||
|
Setup: func(t *testing.T, req *http.Request) {
|
||||||
|
f, err := os.CreateTemp("", "ollama-model")
|
||||||
|
assert.Nil(t, err)
|
||||||
|
tempModelFile = f.Name()
|
||||||
|
|
||||||
|
stream := false
|
||||||
|
createReq := api.CreateRequest{
|
||||||
|
Name: "t-bone",
|
||||||
|
Modelfile: fmt.Sprintf("FROM %s", f.Name()),
|
||||||
|
Stream: &stream,
|
||||||
|
}
|
||||||
|
jsonData, err := json.Marshal(createReq)
|
||||||
|
assert.Nil(t, err)
|
||||||
|
|
||||||
|
req.Body = io.NopCloser(bytes.NewReader(jsonData))
|
||||||
|
},
|
||||||
|
Expected: func(t *testing.T, resp *http.Response) {
|
||||||
|
os.RemoveAll(tempModelFile)
|
||||||
|
|
||||||
|
contentType := resp.Header.Get("Content-Type")
|
||||||
|
assert.Equal(t, "application/json", contentType)
|
||||||
|
_, err := io.ReadAll(resp.Body)
|
||||||
|
assert.Nil(t, err)
|
||||||
|
assert.Equal(t, resp.StatusCode, 200)
|
||||||
|
|
||||||
|
model, err := GetModel("t-bone")
|
||||||
|
assert.Nil(t, err)
|
||||||
|
assert.Equal(t, "t-bone:latest", model.ShortName)
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "Copy Model Handler",
|
||||||
|
Method: http.MethodPost,
|
||||||
|
Path: "/api/copy",
|
||||||
|
Setup: func(t *testing.T, req *http.Request) {
|
||||||
|
createTestModel(t, "hamshank")
|
||||||
|
copyReq := api.CopyRequest{
|
||||||
|
Source: "hamshank",
|
||||||
|
Destination: "beefsteak",
|
||||||
|
}
|
||||||
|
jsonData, err := json.Marshal(copyReq)
|
||||||
|
assert.Nil(t, err)
|
||||||
|
|
||||||
|
req.Body = io.NopCloser(bytes.NewReader(jsonData))
|
||||||
|
},
|
||||||
|
Expected: func(t *testing.T, resp *http.Response) {
|
||||||
|
model, err := GetModel("beefsteak")
|
||||||
|
assert.Nil(t, err)
|
||||||
|
assert.Equal(t, "beefsteak:latest", model.ShortName)
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
@ -121,11 +182,13 @@ func Test_Routes(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
resp, err := httpSrv.Client().Do(req)
|
resp, err := httpSrv.Client().Do(req)
|
||||||
|
defer resp.Body.Close()
|
||||||
assert.Nil(t, err)
|
assert.Nil(t, err)
|
||||||
|
|
||||||
if tc.Expected != nil {
|
if tc.Expected != nil {
|
||||||
tc.Expected(t, resp)
|
tc.Expected(t, resp)
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in a new issue