add jina reranker

pkg/reranker/jina/config.go

@@ -0,0 +1,28 @@
+package jina
+
+import (
+	"net/http"
+)
+
+type Config struct {
+	url string
+
+	token string
+	model string
+
+	client *http.Client
+}
+
+type Option func(*Config)
+
+func WithClient(client *http.Client) Option {
+	return func(c *Config) {
+		c.client = client
+	}
+}
+
+func WithToken(token string) Option {
+	return func(c *Config) {
+		c.token = token
+	}
+}

pkg/reranker/jina/reranker.go

@@ -0,0 +1,108 @@
+package jina
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"github.com/adrianliechti/llama/pkg/reranker"
+)
+
+var _ reranker.Provider = (*Reranker)(nil)
+
+type Reranker struct {
+	*Config
+}
+
+func NewReranker(url string, options ...Option) (*Reranker, error) {
+	if url == "" {
+		url = "https://api.jina.ai"
+	}
+
+	url = strings.TrimRight(url, "/")
+	url = strings.TrimSuffix(url, "/v1")
+
+	cfg := &Config{
+		client: http.DefaultClient,
+
+		url: url,
+
+		model: "jina-reranker-v2-base-multilingual",
+	}
+
+	for _, option := range options {
+		option(cfg)
+	}
+
+	return &Reranker{
+		Config: cfg,
+	}, nil
+}
+
+func (r *Reranker) Rerank(ctx context.Context, query string, inputs []string) ([]reranker.Result, error) {
+	body := map[string]any{
+		"query": query,
+		"input": inputs,
+	}
+
+	u, _ := url.JoinPath(r.url, "/v1/rerank")
+
+	req, _ := http.NewRequestWithContext(ctx, "POST", u, jsonReader(body))
+	req.Header.Set("Content-Type", "application/json")
+
+	if r.token != "" {
+		req.Header.Set("Authorization", "Bearer "+r.token)
+	}
+
+	resp, err := r.client.Do(req)
+
+	if err != nil {
+		return nil, err
+	}
+
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, convertError(resp)
+	}
+
+	var data RerankList
+
+	if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
+		return nil, err
+	}
+
+	if len(data.Results) == 0 {
+		return nil, errors.New("no reranks found")
+	}
+
+	var result []reranker.Result
+
+	for _, r := range data.Results {
+		result = append(result, reranker.Result{
+			Content: r.Document.Text,
+			Score:   r.Score,
+		})
+	}
+
+	return result, nil
+}
+
+type RerankList struct {
+	Model   string   `json:"model"`
+	Results []Result `json:"results"`
+}
+
+type Result struct {
+	Index int `json:"index"`
+
+	Document Document `json:"document"`
+	Score    float32  `json:"relevance_score"`
+}
+
+type Document struct {
+	Text string `json:"text"`
+}

pkg/reranker/jina/util.go

@@ -0,0 +1,29 @@
+package jina
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"io"
+	"net/http"
+)
+
+func convertError(resp *http.Response) error {
+	data, _ := io.ReadAll(resp.Body)
+
+	if len(data) == 0 {
+		return errors.New(http.StatusText(resp.StatusCode))
+	}
+
+	return errors.New(string(data))
+}
+
+func jsonReader(v any) io.Reader {
+	b := new(bytes.Buffer)
+
+	enc := json.NewEncoder(b)
+	enc.SetEscapeHTML(false)
+
+	enc.Encode(v)
+	return b
+}

pkg/reranker/provider.go

@@ -0,0 +1,14 @@
+package reranker
+
+import (
+	"context"
+)
+
+type Provider interface {
+	Rerank(ctx context.Context, query string, inputs []string) ([]Result, error)
+}
+
+type Result struct {
+	Content string
+	Score   float32
+}