diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 19bf3ccd..3799d501 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -178,7 +178,7 @@ jobs: submodules: true - name: Build images run: | - docker build --build-arg FFMPEG=true --build-arg IMAGE_TYPE=core --build-arg MAKEFLAGS="--jobs=5 --output-sync=target" -t local-ai:tests -f Dockerfile . + docker build --build-arg FFMPEG=true --build-arg IMAGE_TYPE=extras --build-arg EXTRA_BACKENDS=rerankers --build-arg MAKEFLAGS="--jobs=5 --output-sync=target" -t local-ai:tests -f Dockerfile . BASE_IMAGE=local-ai:tests DOCKER_AIO_IMAGE=local-ai-aio:test make docker-aio - name: Test run: | diff --git a/tests/e2e-aio/e2e_suite_test.go b/tests/e2e-aio/e2e_suite_test.go index 0aa68230..680bd3a5 100644 --- a/tests/e2e-aio/e2e_suite_test.go +++ b/tests/e2e-aio/e2e_suite_test.go @@ -40,7 +40,8 @@ var _ = BeforeSuite(func() { if apiEndpoint == "" { startDockerImage() defaultConfig = openai.DefaultConfig(apiKey) - defaultConfig.BaseURL = "http://localhost:" + apiPort + "/v1" + apiEndpoint = "http://localhost:" + apiPort + "/v1" // So that other tests can reference this value safely. + defaultConfig.BaseURL = apiEndpoint } else { fmt.Println("Default ", apiEndpoint) defaultConfig = openai.DefaultConfig(apiKey) diff --git a/tests/e2e-aio/e2e_test.go b/tests/e2e-aio/e2e_test.go index 670b3465..a03d70ca 100644 --- a/tests/e2e-aio/e2e_test.go +++ b/tests/e2e-aio/e2e_test.go @@ -1,6 +1,7 @@ package e2e_test import ( + "bytes" "context" "encoding/json" "fmt" @@ -8,6 +9,7 @@ import ( "net/http" "os" + "github.com/go-skynet/LocalAI/core/schema" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/sashabaranov/go-openai" @@ -227,6 +229,51 @@ var _ = Describe("E2E test", func() { Expect(resp.Text).To(ContainSubstring("This is the"), fmt.Sprint(resp.Text)) }) }) + + Context("reranker", func() { + It("correctly", func() { + modelName := "jina-reranker-v1-base-en" + + req := schema.JINARerankRequest{ + Model: modelName, + Query: "Organic skincare products for sensitive skin", + Documents: []string{ + "Eco-friendly kitchenware for modern homes", + "Biodegradable cleaning supplies for eco-conscious consumers", + "Organic cotton baby clothes for sensitive skin", + "Natural organic skincare range for sensitive skin", + "Tech gadgets for smart homes: 2024 edition", + "Sustainable gardening tools and compost solutions", + "Sensitive skin-friendly facial cleansers and toners", + "Organic food wraps and storage solutions", + "All-natural pet food for dogs with allergies", + "Yoga mats made from recycled materials", + }, + TopN: 3, + } + + serialized, err := json.Marshal(req) + Expect(err).To(BeNil()) + Expect(serialized).ToNot(BeNil()) + + rerankerEndpoint := apiEndpoint + "/rerank" + resp, err := http.Post(rerankerEndpoint, "application/json", bytes.NewReader(serialized)) + Expect(err).To(BeNil()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(200)) + + body, err := io.ReadAll(resp.Body) + Expect(err).To(BeNil()) + Expect(body).ToNot(BeNil()) + + deserializedResponse := schema.JINARerankResponse{} + err = json.Unmarshal(body, &deserializedResponse) + Expect(err).To(BeNil()) + Expect(deserializedResponse).ToNot(BeZero()) + Expect(deserializedResponse.Model).To(Equal(modelName)) + Expect(len(deserializedResponse.Results)).To(BeNumerically(">", 0)) + }) + }) }) })