feat(brain_answer): Qwen3-Reranker cross-encoder filter (opt-in)
Adds an opt-in cross-encoder rerank step between BM25 retrieval and LLM synthesis. With BRAIN_RERANKER_URL set, brain_answer retrieves BM25 top-20, scores each excerpt against the query via Qwen3-Reranker on Ollama, drops the "no" answers, and forwards up to 5 surviving sources to the LLM. Unset, behaviour is unchanged (BM25 top-10 → LLM). The reranker is a *filter*, not a re-ranker: Qwen3-Reranker emits a binary yes/no token under its native chat template, and ties within the "yes" set are broken by BM25 rank — what got retrieved first stays ahead. New package ingestion/internal/reranker: - Client with URL, Model, HTTP fields. - New(url, model) returns nil on empty url so callers can treat "feature disabled" as a single nil check. - Score(ctx, query, docs) issues one /api/generate call per doc using the Qwen3-Reranker yes/no chat template (verbatim, because the model was trained on this exact wording). Parses the first non-think token. Wiring: - mcp.Server gains a WithReranker fluent setter to keep NewServer signature stable. - brain_answer's BM25 limit jumps to 20 only when a reranker is wired, to give the filter something to do. - cmd/server/main.go reads BRAIN_RERANKER_URL (+ optional BRAIN_RERANKER_MODEL, default dengcao/Qwen3-Reranker-0.6B:F16). Tests cover: nil-on-empty-url, ordered yes/no scoring, request shape (model, prompt contents, yes/no template), ambiguous response → 0, empty doc slice, upstream-error propagation, plus an end-to-end brain_answer integration that proves only the relevant note reaches the LLM when noise.md is rejected. Closes hyperguild#7.
This commit is contained in:
119
ingestion/internal/reranker/reranker_test.go
Normal file
119
ingestion/internal/reranker/reranker_test.go
Normal file
@@ -0,0 +1,119 @@
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user