mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-12 04:55:18 +00:00
Release dec 2025 (#15)
* Resolves issue #13 - Switched model to bge-small-en-v1.5 - Added lazy re-embedding - Added model version tracking per vector - Added conversion of vectors to the new model * Add lfs support to the workflow. * Implements importance scoring with decay + voting #6 * Resolves issue #5 by marking observations as superseeded and scheduled for deletion * Implement pattern detection #7 * Improve injections and observations accuracy - Session start: Recent observations for project context (recency-based) - User prompt: Semantically relevant observations (similarity-based with threshold) * Added two stage retrieval with bi and cross encoder #8 * Implement query expansion and reformulation #9 * Knowledge graph and relationships ( resolves #4 ) - File Overlap Detection: Detects relationships when observations modify/read the same files - Concept Overlap Detection: Detects relationships based on shared semantic concepts - Type Progression Detection: Infers relationships from natural observation type progressions (e.g., discovery → bugfix = "fixes") - Temporal Proximity Detection: Detects relationships between observations in the same session within 5 minutes - Narrative Mention Detection: Detects explicit relationship language in narratives (e.g., "fixes", "depends on", "supersedes") * Add visualisation of the relations to the dashboard. * fixup! Add visualisation of the relations to the dashboard. * Update documentation with new settings and screenshots.
This commit is contained in:
@@ -0,0 +1,14 @@
|
||||
// Package reranking provides cross-encoder reranking for search results.
|
||||
package reranking
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
)
|
||||
|
||||
// Cross-encoder model and tokenizer files - embedded for all platforms
|
||||
//
|
||||
//go:embed assets/model.onnx
|
||||
var crossEncoderModelData []byte
|
||||
|
||||
//go:embed assets/tokenizer.json
|
||||
var crossEncoderTokenizerData []byte
|
||||
Binary file not shown.
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,57 @@
|
||||
{
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "[PAD]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"100": {
|
||||
"content": "[UNK]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"101": {
|
||||
"content": "[CLS]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"102": {
|
||||
"content": "[SEP]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"103": {
|
||||
"content": "[MASK]",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
}
|
||||
},
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"cls_token": "[CLS]",
|
||||
"do_basic_tokenize": true,
|
||||
"do_lower_case": true,
|
||||
"mask_token": "[MASK]",
|
||||
"model_max_length": 512,
|
||||
"never_split": null,
|
||||
"pad_token": "[PAD]",
|
||||
"sep_token": "[SEP]",
|
||||
"strip_accents": null,
|
||||
"tokenize_chinese_chars": true,
|
||||
"tokenizer_class": "BertTokenizer",
|
||||
"unk_token": "[UNK]"
|
||||
}
|
||||
@@ -0,0 +1,380 @@
|
||||
// Package reranking provides cross-encoder reranking for search results.
|
||||
// Uses MS-MARCO MiniLM L6 v2 cross-encoder model for relevance scoring.
|
||||
package reranking
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"sync"
|
||||
|
||||
"github.com/rs/zerolog/log"
|
||||
"github.com/sugarme/tokenizer"
|
||||
"github.com/sugarme/tokenizer/pretrained"
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
)
|
||||
|
||||
const (
|
||||
// ModelName is the human-readable name for the cross-encoder model
|
||||
ModelName = "ms-marco-MiniLM-L6-v2"
|
||||
// ModelVersion is the short version string for identification
|
||||
ModelVersion = "msmarco-v2"
|
||||
// MaxSequenceLength is the maximum combined query+document token length
|
||||
MaxSequenceLength = 512
|
||||
// DefaultCandidateLimit is the default number of candidates to rerank
|
||||
DefaultCandidateLimit = 100
|
||||
// DefaultResultLimit is the default number of results to return after reranking
|
||||
DefaultResultLimit = 10
|
||||
)
|
||||
|
||||
// Candidate represents a search result candidate for reranking.
|
||||
type Candidate struct {
|
||||
ID string // Document ID
|
||||
Content string // Document text content for scoring
|
||||
Score float64 // Original bi-encoder similarity score
|
||||
Metadata map[string]any // Preserved metadata
|
||||
RerankInfo map[string]float64 // Reranking debug info (optional)
|
||||
}
|
||||
|
||||
// RerankResult represents a reranked search result.
|
||||
type RerankResult struct {
|
||||
ID string // Document ID
|
||||
Content string // Document text content
|
||||
OriginalScore float64 // Original bi-encoder score
|
||||
RerankScore float64 // Cross-encoder relevance score
|
||||
CombinedScore float64 // Weighted combination of scores
|
||||
Metadata map[string]any // Preserved metadata
|
||||
OriginalRank int // Position before reranking (1-indexed)
|
||||
RerankRank int // Position after reranking (1-indexed)
|
||||
RankImprovement int // How much the rank improved (positive = moved up)
|
||||
}
|
||||
|
||||
// Service provides cross-encoder reranking functionality.
|
||||
type Service struct {
|
||||
tk *tokenizer.Tokenizer
|
||||
session *ort.DynamicAdvancedSession
|
||||
mu sync.Mutex
|
||||
|
||||
// Weight for combining scores: combined = alpha*rerank + (1-alpha)*original
|
||||
// Default 0.7 favors cross-encoder score
|
||||
Alpha float64
|
||||
}
|
||||
|
||||
// Config holds configuration for the reranking service.
|
||||
type Config struct {
|
||||
// Alpha is the weight for combining scores (0.0-1.0)
|
||||
// Higher values favor cross-encoder scores, lower values favor bi-encoder scores
|
||||
Alpha float64
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults for reranking.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Alpha: 0.7, // Favor cross-encoder by default
|
||||
}
|
||||
}
|
||||
|
||||
// NewService creates a new cross-encoder reranking service.
|
||||
// Note: ONNX runtime must be initialized before calling this (via embedding.NewService).
|
||||
func NewService(cfg Config) (*Service, error) {
|
||||
// Load tokenizer from embedded data
|
||||
tk, err := pretrained.FromReader(bytes.NewReader(crossEncoderTokenizerData))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load cross-encoder tokenizer: %w", err)
|
||||
}
|
||||
|
||||
// Configure tokenizer for sequence classification (pairs)
|
||||
tk.WithTruncation(&tokenizer.TruncationParams{
|
||||
MaxLength: MaxSequenceLength,
|
||||
Strategy: tokenizer.LongestFirst,
|
||||
Stride: 0,
|
||||
})
|
||||
|
||||
// Cross-encoder outputs a single logit for relevance scoring
|
||||
inputNames := []string{"input_ids", "attention_mask", "token_type_ids"}
|
||||
outputNames := []string{"logits"}
|
||||
|
||||
session, err := ort.NewDynamicAdvancedSessionWithONNXData(
|
||||
crossEncoderModelData,
|
||||
inputNames,
|
||||
outputNames,
|
||||
nil,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create cross-encoder ONNX session: %w", err)
|
||||
}
|
||||
|
||||
alpha := cfg.Alpha
|
||||
if alpha <= 0 || alpha > 1 {
|
||||
alpha = 0.7
|
||||
}
|
||||
|
||||
return &Service{
|
||||
tk: tk,
|
||||
session: session,
|
||||
Alpha: alpha,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Rerank reranks candidates using the cross-encoder model.
|
||||
// Takes a query and list of candidates, returns reranked results.
|
||||
func (s *Service) Rerank(query string, candidates []Candidate, limit int) ([]RerankResult, error) {
|
||||
if len(candidates) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if limit <= 0 {
|
||||
limit = DefaultResultLimit
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
// Score all query-document pairs
|
||||
scores, err := s.scoreAll(query, candidates)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("score candidates: %w", err)
|
||||
}
|
||||
|
||||
// Build results with combined scores
|
||||
results := make([]RerankResult, len(candidates))
|
||||
for i, c := range candidates {
|
||||
// Normalize cross-encoder score to 0-1 range using sigmoid
|
||||
normalizedRerank := sigmoid(scores[i])
|
||||
|
||||
results[i] = RerankResult{
|
||||
ID: c.ID,
|
||||
Content: c.Content,
|
||||
OriginalScore: c.Score,
|
||||
RerankScore: normalizedRerank,
|
||||
CombinedScore: s.Alpha*normalizedRerank + (1-s.Alpha)*c.Score,
|
||||
Metadata: c.Metadata,
|
||||
OriginalRank: i + 1,
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by combined score (descending)
|
||||
sort.Slice(results, func(i, j int) bool {
|
||||
return results[i].CombinedScore > results[j].CombinedScore
|
||||
})
|
||||
|
||||
// Assign rerank positions and calculate improvement
|
||||
for i := range results {
|
||||
results[i].RerankRank = i + 1
|
||||
results[i].RankImprovement = results[i].OriginalRank - results[i].RerankRank
|
||||
}
|
||||
|
||||
// Apply limit
|
||||
if len(results) > limit {
|
||||
results = results[:limit]
|
||||
}
|
||||
|
||||
log.Debug().
|
||||
Int("candidates", len(candidates)).
|
||||
Int("returned", len(results)).
|
||||
Float64("alpha", s.Alpha).
|
||||
Msg("Cross-encoder reranking completed")
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// RerankByScore reranks candidates and returns sorted by pure cross-encoder score.
|
||||
// Useful when you want to completely replace bi-encoder ranking.
|
||||
func (s *Service) RerankByScore(query string, candidates []Candidate, limit int) ([]RerankResult, error) {
|
||||
if len(candidates) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if limit <= 0 {
|
||||
limit = DefaultResultLimit
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
scores, err := s.scoreAll(query, candidates)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("score candidates: %w", err)
|
||||
}
|
||||
|
||||
results := make([]RerankResult, len(candidates))
|
||||
for i, c := range candidates {
|
||||
normalizedRerank := sigmoid(scores[i])
|
||||
results[i] = RerankResult{
|
||||
ID: c.ID,
|
||||
Content: c.Content,
|
||||
OriginalScore: c.Score,
|
||||
RerankScore: normalizedRerank,
|
||||
CombinedScore: normalizedRerank, // Use pure rerank score
|
||||
Metadata: c.Metadata,
|
||||
OriginalRank: i + 1,
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by rerank score only
|
||||
sort.Slice(results, func(i, j int) bool {
|
||||
return results[i].RerankScore > results[j].RerankScore
|
||||
})
|
||||
|
||||
for i := range results {
|
||||
results[i].RerankRank = i + 1
|
||||
results[i].RankImprovement = results[i].OriginalRank - results[i].RerankRank
|
||||
}
|
||||
|
||||
if len(results) > limit {
|
||||
results = results[:limit]
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// scoreAll scores all query-document pairs using the cross-encoder.
|
||||
// Returns raw logits (before sigmoid normalization).
|
||||
func (s *Service) scoreAll(query string, candidates []Candidate) ([]float64, error) {
|
||||
batchSize := len(candidates)
|
||||
|
||||
// Tokenize all query-document pairs
|
||||
pairs := make([]tokenizer.EncodeInput, batchSize)
|
||||
for i, c := range candidates {
|
||||
// Cross-encoder takes query and document as a pair
|
||||
pairs[i] = tokenizer.NewDualEncodeInput(
|
||||
tokenizer.NewRawInputSequence(query),
|
||||
tokenizer.NewRawInputSequence(c.Content),
|
||||
)
|
||||
}
|
||||
|
||||
encodings, err := s.tk.EncodeBatch(pairs, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("tokenize pairs: %w", err)
|
||||
}
|
||||
|
||||
// Find max sequence length
|
||||
seqLength := 0
|
||||
for _, enc := range encodings {
|
||||
if len(enc.Ids) > seqLength {
|
||||
seqLength = len(enc.Ids)
|
||||
}
|
||||
}
|
||||
if seqLength > MaxSequenceLength {
|
||||
seqLength = MaxSequenceLength
|
||||
}
|
||||
|
||||
inputShape := ort.NewShape(int64(batchSize), int64(seqLength))
|
||||
|
||||
// Create input tensors
|
||||
inputIdsData := make([]int64, batchSize*seqLength)
|
||||
attentionMaskData := make([]int64, batchSize*seqLength)
|
||||
tokenTypeIdsData := make([]int64, batchSize*seqLength)
|
||||
|
||||
for b := 0; b < batchSize; b++ {
|
||||
copyLen := len(encodings[b].Ids)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
inputIdsData[b*seqLength+i] = int64(encodings[b].Ids[i])
|
||||
}
|
||||
|
||||
copyLen = len(encodings[b].AttentionMask)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
attentionMaskData[b*seqLength+i] = int64(encodings[b].AttentionMask[i])
|
||||
}
|
||||
|
||||
copyLen = len(encodings[b].TypeIds)
|
||||
if copyLen > seqLength {
|
||||
copyLen = seqLength
|
||||
}
|
||||
for i := 0; i < copyLen; i++ {
|
||||
tokenTypeIdsData[b*seqLength+i] = int64(encodings[b].TypeIds[i])
|
||||
}
|
||||
}
|
||||
|
||||
inputIdsTensor, err := ort.NewTensor(inputShape, inputIdsData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create input_ids tensor: %w", err)
|
||||
}
|
||||
defer inputIdsTensor.Destroy()
|
||||
|
||||
attentionMaskTensor, err := ort.NewTensor(inputShape, attentionMaskData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create attention_mask tensor: %w", err)
|
||||
}
|
||||
defer attentionMaskTensor.Destroy()
|
||||
|
||||
tokenTypeIdsTensor, err := ort.NewTensor(inputShape, tokenTypeIdsData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create token_type_ids tensor: %w", err)
|
||||
}
|
||||
defer tokenTypeIdsTensor.Destroy()
|
||||
|
||||
// Cross-encoder outputs [batch, 1] logits
|
||||
outputShape := ort.NewShape(int64(batchSize), 1)
|
||||
outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create output tensor: %w", err)
|
||||
}
|
||||
defer outputTensor.Destroy()
|
||||
|
||||
// Run inference
|
||||
inputTensors := []ort.Value{inputIdsTensor, attentionMaskTensor, tokenTypeIdsTensor}
|
||||
outputTensors := []ort.Value{outputTensor}
|
||||
|
||||
if err := s.session.Run(inputTensors, outputTensors); err != nil {
|
||||
return nil, fmt.Errorf("run cross-encoder inference: %w", err)
|
||||
}
|
||||
|
||||
// Extract scores
|
||||
flatOutput := outputTensor.GetData()
|
||||
scores := make([]float64, batchSize)
|
||||
for i := 0; i < batchSize; i++ {
|
||||
scores[i] = float64(flatOutput[i])
|
||||
}
|
||||
|
||||
return scores, nil
|
||||
}
|
||||
|
||||
// Score scores a single query-document pair.
|
||||
// Returns the raw cross-encoder logit and normalized score.
|
||||
func (s *Service) Score(query, document string) (rawScore, normalizedScore float64, err error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
scores, err := s.scoreAll(query, []Candidate{{Content: document}})
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
rawScore = scores[0]
|
||||
normalizedScore = sigmoid(rawScore)
|
||||
return rawScore, normalizedScore, nil
|
||||
}
|
||||
|
||||
// Close releases model resources.
|
||||
func (s *Service) Close() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.session != nil {
|
||||
if err := s.session.Destroy(); err != nil {
|
||||
return fmt.Errorf("destroy cross-encoder session: %w", err)
|
||||
}
|
||||
s.session = nil
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// sigmoid applies the sigmoid function to normalize scores to 0-1 range.
|
||||
func sigmoid(x float64) float64 {
|
||||
if x > 20 {
|
||||
return 1.0
|
||||
}
|
||||
if x < -20 {
|
||||
return 0.0
|
||||
}
|
||||
return 1.0 / (1.0 + math.Exp(-x))
|
||||
}
|
||||
@@ -0,0 +1,448 @@
|
||||
package reranking
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/lukaszraczylo/claude-mnemonic/internal/embedding"
|
||||
)
|
||||
|
||||
// initONNX initializes the ONNX runtime via the embedding service.
|
||||
// Must be called before creating reranking service.
|
||||
func initONNX(t *testing.T) func() {
|
||||
t.Helper()
|
||||
embSvc, err := embedding.NewService()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to initialize ONNX via embedding service: %v", err)
|
||||
}
|
||||
return func() {
|
||||
embSvc.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// TestSigmoid tests the sigmoid normalization function.
|
||||
func TestSigmoid(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input float64
|
||||
wantMin float64
|
||||
wantMax float64
|
||||
}{
|
||||
{"positive large", 10, 0.9999, 1.0},
|
||||
{"positive small", 1, 0.7, 0.8},
|
||||
{"zero", 0, 0.4999, 0.5001},
|
||||
{"negative small", -1, 0.2, 0.3},
|
||||
{"negative large", -10, 0, 0.0001},
|
||||
{"very positive", 25, 0.999999, 1.0},
|
||||
{"very negative", -25, 0, 0.000001},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := sigmoid(tt.input)
|
||||
if got < tt.wantMin || got > tt.wantMax {
|
||||
t.Errorf("sigmoid(%v) = %v, want in range [%v, %v]",
|
||||
tt.input, got, tt.wantMin, tt.wantMax)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultConfig tests the default configuration values.
|
||||
func TestDefaultConfig(t *testing.T) {
|
||||
cfg := DefaultConfig()
|
||||
|
||||
if cfg.Alpha < 0 || cfg.Alpha > 1 {
|
||||
t.Errorf("DefaultConfig().Alpha = %v, want in range [0, 1]", cfg.Alpha)
|
||||
}
|
||||
if cfg.Alpha != 0.7 {
|
||||
t.Errorf("DefaultConfig().Alpha = %v, want 0.7", cfg.Alpha)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewService tests service creation.
|
||||
func TestNewService(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
if svc == nil {
|
||||
t.Fatal("NewService() returned nil")
|
||||
}
|
||||
|
||||
if svc.Alpha != cfg.Alpha {
|
||||
t.Errorf("Service.Alpha = %v, want %v", svc.Alpha, cfg.Alpha)
|
||||
}
|
||||
}
|
||||
|
||||
// TestScore tests single pair scoring.
|
||||
func TestScore(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
query := "What is the capital of France?"
|
||||
relevant := "Paris is the capital and largest city of France."
|
||||
irrelevant := "Dogs are popular pets known for their loyalty."
|
||||
|
||||
// Score relevant document
|
||||
_, relevantNorm, err := svc.Score(query, relevant)
|
||||
if err != nil {
|
||||
t.Fatalf("Score(relevant) error = %v", err)
|
||||
}
|
||||
|
||||
// Score irrelevant document
|
||||
_, irrelevantNorm, err := svc.Score(query, irrelevant)
|
||||
if err != nil {
|
||||
t.Fatalf("Score(irrelevant) error = %v", err)
|
||||
}
|
||||
|
||||
// Relevant document should score higher
|
||||
if relevantNorm <= irrelevantNorm {
|
||||
t.Errorf("Expected relevant (%v) > irrelevant (%v)",
|
||||
relevantNorm, irrelevantNorm)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRerank tests the reranking functionality.
|
||||
func TestRerank(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
query := "How to handle errors in Go?"
|
||||
candidates := []Candidate{
|
||||
{
|
||||
ID: "1",
|
||||
Content: "Python exception handling with try/except blocks.",
|
||||
Score: 0.8, // High bi-encoder score but irrelevant
|
||||
},
|
||||
{
|
||||
ID: "2",
|
||||
Content: "Go error handling uses explicit return values. Functions return error as the last value.",
|
||||
Score: 0.6, // Lower bi-encoder score but relevant
|
||||
},
|
||||
{
|
||||
ID: "3",
|
||||
Content: "JavaScript uses Promise.catch for async error handling.",
|
||||
Score: 0.7,
|
||||
},
|
||||
}
|
||||
|
||||
results, err := svc.Rerank(query, candidates, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank() error = %v", err)
|
||||
}
|
||||
|
||||
if len(results) != 3 {
|
||||
t.Fatalf("Rerank() returned %d results, want 3", len(results))
|
||||
}
|
||||
|
||||
// The Go error handling document should rank higher after reranking
|
||||
var goRank int
|
||||
for i, r := range results {
|
||||
if r.ID == "2" {
|
||||
goRank = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if goRank == 0 {
|
||||
t.Error("Go document not found in results")
|
||||
}
|
||||
|
||||
// Verify all results have required fields populated
|
||||
for i, r := range results {
|
||||
if r.ID == "" {
|
||||
t.Errorf("Result %d has empty ID", i)
|
||||
}
|
||||
if r.Content == "" {
|
||||
t.Errorf("Result %d has empty Content", i)
|
||||
}
|
||||
if r.RerankRank != i+1 {
|
||||
t.Errorf("Result %d has RerankRank %d, want %d", i, r.RerankRank, i+1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRerankEmpty tests reranking with empty input.
|
||||
func TestRerankEmpty(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
results, err := svc.Rerank("test query", nil, 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank(nil) error = %v", err)
|
||||
}
|
||||
|
||||
if results != nil {
|
||||
t.Errorf("Rerank(nil) = %v, want nil", results)
|
||||
}
|
||||
|
||||
results, err = svc.Rerank("test query", []Candidate{}, 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank([]) error = %v", err)
|
||||
}
|
||||
|
||||
if results != nil {
|
||||
t.Errorf("Rerank([]) = %v, want nil", results)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRerankLimit tests that limit is respected.
|
||||
func TestRerankLimit(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
candidates := make([]Candidate, 20)
|
||||
for i := range candidates {
|
||||
candidates[i] = Candidate{
|
||||
ID: string(rune('A' + i)),
|
||||
Content: "Test document content for ranking.",
|
||||
Score: 0.5,
|
||||
}
|
||||
}
|
||||
|
||||
results, err := svc.Rerank("test query", candidates, 5)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank() error = %v", err)
|
||||
}
|
||||
|
||||
if len(results) != 5 {
|
||||
t.Errorf("Rerank() returned %d results, want 5", len(results))
|
||||
}
|
||||
}
|
||||
|
||||
// TestRerankByScore tests pure cross-encoder ranking.
|
||||
func TestRerankByScore(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
query := "machine learning algorithms"
|
||||
candidates := []Candidate{
|
||||
{
|
||||
ID: "1",
|
||||
Content: "Cooking recipes for Italian pasta dishes.",
|
||||
Score: 0.9, // High original score
|
||||
},
|
||||
{
|
||||
ID: "2",
|
||||
Content: "Neural networks are a type of machine learning algorithm.",
|
||||
Score: 0.3, // Low original score
|
||||
},
|
||||
}
|
||||
|
||||
results, err := svc.RerankByScore(query, candidates, 2)
|
||||
if err != nil {
|
||||
t.Fatalf("RerankByScore() error = %v", err)
|
||||
}
|
||||
|
||||
// Document 2 should rank first since it's about ML
|
||||
if results[0].ID != "2" {
|
||||
t.Errorf("Expected ML document to rank first, got %v", results[0].ID)
|
||||
}
|
||||
|
||||
// CombinedScore should equal RerankScore when using RerankByScore
|
||||
for _, r := range results {
|
||||
if r.CombinedScore != r.RerankScore {
|
||||
t.Errorf("RerankByScore: CombinedScore (%v) != RerankScore (%v)",
|
||||
r.CombinedScore, r.RerankScore)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRankImprovement tests that rank improvement is calculated correctly.
|
||||
func TestRankImprovement(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
// Create candidates where we know the expected reranking
|
||||
candidates := []Candidate{
|
||||
{ID: "A", Content: "Unrelated content about weather forecasting.", Score: 0.9},
|
||||
{ID: "B", Content: "How to fix memory leaks in Go programs.", Score: 0.8},
|
||||
{ID: "C", Content: "More unrelated content about gardening tips.", Score: 0.7},
|
||||
}
|
||||
|
||||
results, err := svc.Rerank("debugging memory issues in Go", candidates, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank() error = %v", err)
|
||||
}
|
||||
|
||||
for _, r := range results {
|
||||
// RankImprovement = OriginalRank - RerankRank
|
||||
// Positive means moved up, negative means moved down
|
||||
expectedImprovement := r.OriginalRank - r.RerankRank
|
||||
if r.RankImprovement != expectedImprovement {
|
||||
t.Errorf("ID %s: RankImprovement = %d, want %d (orig=%d, new=%d)",
|
||||
r.ID, r.RankImprovement, expectedImprovement,
|
||||
r.OriginalRank, r.RerankRank)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentRerank tests concurrent reranking calls.
|
||||
func TestConcurrentRerank(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
candidates := []Candidate{
|
||||
{ID: "1", Content: "Test document one.", Score: 0.5},
|
||||
{ID: "2", Content: "Test document two.", Score: 0.5},
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errors := make(chan error, 10)
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
_, err := svc.Rerank("concurrent test query", candidates, 2)
|
||||
if err != nil {
|
||||
errors <- err
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
close(errors)
|
||||
|
||||
for err := range errors {
|
||||
t.Errorf("Concurrent Rerank error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestClose tests service cleanup.
|
||||
func TestClose(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
|
||||
err = svc.Close()
|
||||
if err != nil {
|
||||
t.Errorf("Close() error = %v", err)
|
||||
}
|
||||
|
||||
// Double close should not panic
|
||||
err = svc.Close()
|
||||
if err != nil {
|
||||
t.Errorf("Close() on closed service error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMetadataPreserved tests that metadata is preserved through reranking.
|
||||
func TestMetadataPreserved(t *testing.T) {
|
||||
cleanup := initONNX(t)
|
||||
defer cleanup()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
svc, err := NewService(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewService() error = %v", err)
|
||||
}
|
||||
defer svc.Close()
|
||||
|
||||
candidates := []Candidate{
|
||||
{
|
||||
ID: "1",
|
||||
Content: "Test content.",
|
||||
Score: 0.5,
|
||||
Metadata: map[string]any{"custom": "value1", "num": 42},
|
||||
},
|
||||
{
|
||||
ID: "2",
|
||||
Content: "Another test.",
|
||||
Score: 0.5,
|
||||
Metadata: map[string]any{"custom": "value2"},
|
||||
},
|
||||
}
|
||||
|
||||
results, err := svc.Rerank("query", candidates, 2)
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank() error = %v", err)
|
||||
}
|
||||
|
||||
for _, r := range results {
|
||||
if r.Metadata == nil {
|
||||
t.Errorf("Result %s has nil metadata", r.ID)
|
||||
continue
|
||||
}
|
||||
|
||||
// Find original candidate
|
||||
var original *Candidate
|
||||
for i := range candidates {
|
||||
if candidates[i].ID == r.ID {
|
||||
original = &candidates[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if original == nil {
|
||||
t.Errorf("Could not find original for result %s", r.ID)
|
||||
continue
|
||||
}
|
||||
|
||||
// Check metadata preserved
|
||||
if original.Metadata["custom"] != r.Metadata["custom"] {
|
||||
t.Errorf("Metadata not preserved for %s", r.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user