hyperguild/ingestion/internal/reranker/reranker_test.go

package reranker_test

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/mathiasbq/hyperguild/ingestion/internal/reranker"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// fakeOllama responds to /api/generate based on a per-document
// {needle → answer} map: if the prompt contains the needle, returns
// the mapped answer.
type fakeOllama struct {
	t        *testing.T
	answers  map[string]string // needle → "yes" or "no"
	calls    int
	lastBody map[string]any
}

func (f *fakeOllama) handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(f.t, http.MethodPost, r.Method)
		require.Equal(f.t, "/api/generate", r.URL.Path)
		body, err := io.ReadAll(r.Body)
		require.NoError(f.t, err)
		var p map[string]any
		require.NoError(f.t, json.Unmarshal(body, &p))
		f.calls++
		f.lastBody = p
		prompt := p["prompt"].(string)
		answer := "no"
		for needle, a := range f.answers {
			if strings.Contains(prompt, needle) {
				answer = a
				break
			}
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{
			"model": p["model"], "response": answer, "done": true,
		})
	})
}

func TestNew_EmptyURLReturnsNil(t *testing.T) {
	assert.Nil(t, reranker.New("", "model"))
}

func TestScore_YesAndNoOrdered(t *testing.T) {
	f := &fakeOllama{t: t, answers: map[string]string{
		"alpha doc": "yes",
		"beta doc":  "no",
		"gamma doc": "yes",
	}}
	srv := httptest.NewServer(f.handler())
	defer srv.Close()
	c := reranker.New(srv.URL, "test-model")
	require.NotNil(t, c)

	scores, err := c.Score(context.Background(), "what is alpha",
		[]string{"alpha doc body", "beta doc body", "gamma doc body"})
	require.NoError(t, err)
	require.Len(t, scores, 3)
	assert.Equal(t, 1.0, scores[0])
	assert.Equal(t, 0.0, scores[1])
	assert.Equal(t, 1.0, scores[2])
	assert.Equal(t, 3, f.calls)
}

func TestScore_SendsCorrectShape(t *testing.T) {
	f := &fakeOllama{t: t, answers: map[string]string{"hello": "yes"}}
	srv := httptest.NewServer(f.handler())
	defer srv.Close()
	c := reranker.New(srv.URL, "qwen3-rerank")
	_, err := c.Score(context.Background(), "greeting", []string{"hello world"})
	require.NoError(t, err)
	assert.Equal(t, "qwen3-rerank", f.lastBody["model"])
	prompt := f.lastBody["prompt"].(string)
	assert.Contains(t, prompt, "greeting")
	assert.Contains(t, prompt, "hello world")
	assert.Contains(t, prompt, `"yes" or "no"`)
}

func TestScore_HandlesAmbiguousResponse(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{"response": "maybe — unclear", "done": true})
	}))
	defer srv.Close()
	c := reranker.New(srv.URL, "m")
	scores, err := c.Score(context.Background(), "q", []string{"d"})
	require.NoError(t, err)
	// Anything that does not start with "yes" (case-insensitive, after
	// whitespace/think trim) is treated as "no" = 0.
	assert.Equal(t, []float64{0}, scores)
}

func TestScore_EmptyDocsReturnsEmpty(t *testing.T) {
	c := reranker.New("http://127.0.0.1:1", "m")
	scores, err := c.Score(context.Background(), "q", nil)
	require.NoError(t, err)
	assert.Empty(t, scores)
}

func TestScore_UpstreamErrorPropagates(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	defer srv.Close()
	c := reranker.New(srv.URL, "m")
	_, err := c.Score(context.Background(), "q", []string{"d"})
	require.Error(t, err)
}