Skip to content

Commit

Permalink
test: e2e /reranker endpoint (#2211)
Browse files Browse the repository at this point in the history
Create a simple e2e test for the /reranker api \\ go mod tidy

Signed-off-by: Dave Lee <dave@gray101.com>
  • Loading branch information
dave-gray101 committed Jun 7, 2024
1 parent 3b7a78a commit 219078a
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 2 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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: |
Expand Down
3 changes: 2 additions & 1 deletion tests/e2e-aio/e2e_suite_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
47 changes: 47 additions & 0 deletions tests/e2e-aio/e2e_test.go
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
package e2e_test

import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"

"github.com/go-skynet/LocalAI/core/schema"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/sashabaranov/go-openai"
Expand Down Expand Up @@ -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))
})
})
})
})

Expand Down

0 comments on commit 219078a

Please sign in to comment.