Skip to content

Commit c55a1cc

Browse files
Siddhant-K-codeona-agent
Siddhant-K-code
and
ona-agent
authored
feat(memory): add task-relevance ranking with boost tags and min relevance filter (#89)
- Add BoostTags field to RecallRequest for tag-based relevance boosting - Add TaskContext field for source-matching relevance boost - Add MinRelevance field to filter low-scoring results - Relevance scores clamped to [0, 1] - 6 new tests: boost tags, min relevance, task context, clamping, no-op Closes #78 Co-authored-by: Ona <no-reply@ona.com>
1 parent 020f678 commit c55a1cc

3 files changed

Lines changed: 224 additions & 0 deletions

File tree

‎pkg/memory/relevance_test.go‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
package memory
2+
3+
import (
4+
"context"
5+
"testing"
6+
)
7+
8+
func TestRecall_BoostTags(t *testing.T) {
9+
s := newTestStore(t)
10+
ctx := context.Background()
11+
12+
// Store two entries with different tags at nearly equal distance from query
13+
// angle 0.3 from query(0.3): auth at 0 → dist=0.045, db at 0.6 → dist=0.045
14+
_, _ = s.Store(ctx, StoreRequest{
15+
Entries: []StoreEntry{
16+
{Text: "Auth uses JWT", Embedding: makeEmbedding(0, 8), Tags: []string{"auth"}},
17+
{Text: "DB uses Postgres", Embedding: makeEmbedding(0.6, 8), Tags: []string{"database"}},
18+
},
19+
})
20+
21+
// Query equidistant from both; boost on "database" should tip the ranking
22+
recall, err := s.Recall(ctx, RecallRequest{
23+
Query: "infrastructure",
24+
QueryEmbedding: makeEmbedding(0.3, 8), // equidistant from 0 and 0.6
25+
MaxResults: 10,
26+
BoostTags: []string{"database"},
27+
})
28+
if err != nil {
29+
t.Fatalf("Recall: %v", err)
30+
}
31+
if len(recall.Memories) != 2 {
32+
t.Fatalf("expected 2 memories, got %d", len(recall.Memories))
33+
}
34+
// With the boost, database entry should be ranked first
35+
if recall.Memories[0].Tags[0] != "database" {
36+
t.Errorf("expected database entry first (boosted), got tags=%v", recall.Memories[0].Tags)
37+
}
38+
}
39+
40+
func TestRecall_MinRelevance(t *testing.T) {
41+
s := newTestStore(t)
42+
ctx := context.Background()
43+
44+
_, _ = s.Store(ctx, StoreRequest{
45+
Entries: []StoreEntry{
46+
{Text: "Highly relevant", Embedding: makeEmbedding(0, 8)},
47+
{Text: "Somewhat relevant", Embedding: makeEmbedding(0.6, 8)},
48+
{Text: "Not relevant", Embedding: makeEmbedding(2.0, 8)},
49+
},
50+
})
51+
52+
// Query with high min relevance — should filter out low-scoring entries
53+
recall, err := s.Recall(ctx, RecallRequest{
54+
Query: "test",
55+
QueryEmbedding: makeEmbedding(0, 8),
56+
MaxResults: 10,
57+
MinRelevance: 0.8,
58+
})
59+
if err != nil {
60+
t.Fatalf("Recall: %v", err)
61+
}
62+
// Only the highly relevant entry (cosine similarity ~1.0) should pass
63+
if len(recall.Memories) == 0 {
64+
t.Fatal("expected at least 1 memory above min relevance")
65+
}
66+
for _, m := range recall.Memories {
67+
if m.Relevance < 0.8 {
68+
t.Errorf("memory %s has relevance %.3f, below min 0.8", m.ID, m.Relevance)
69+
}
70+
}
71+
}
72+
73+
func TestRecall_MinRelevance_Zero_NoFilter(t *testing.T) {
74+
s := newTestStore(t)
75+
ctx := context.Background()
76+
77+
_, _ = s.Store(ctx, StoreRequest{
78+
Entries: []StoreEntry{
79+
{Text: "Entry A", Embedding: makeEmbedding(0, 8)},
80+
{Text: "Entry B", Embedding: makeEmbedding(2.0, 8)},
81+
},
82+
})
83+
84+
// MinRelevance=0 should return all entries
85+
recall, _ := s.Recall(ctx, RecallRequest{
86+
Query: "test",
87+
QueryEmbedding: makeEmbedding(0, 8),
88+
MaxResults: 10,
89+
MinRelevance: 0,
90+
})
91+
if len(recall.Memories) != 2 {
92+
t.Errorf("expected 2 memories with no min filter, got %d", len(recall.Memories))
93+
}
94+
}
95+
96+
func TestRecall_TaskContext_SourceBoost(t *testing.T) {
97+
s := newTestStore(t)
98+
ctx := context.Background()
99+
100+
// Use angles far enough apart to avoid dedup (>0.555 rad)
101+
_, _ = s.Store(ctx, StoreRequest{
102+
Entries: []StoreEntry{
103+
{Text: "JWT validation logic", Embedding: makeEmbedding(0, 8), Source: "code_review"},
104+
{Text: "JWT token format", Embedding: makeEmbedding(0.6, 8), Source: "docs"},
105+
},
106+
})
107+
108+
// Query equidistant; task context mentions "code_review" — should boost that source
109+
recall, err := s.Recall(ctx, RecallRequest{
110+
Query: "JWT",
111+
QueryEmbedding: makeEmbedding(0.3, 8), // equidistant from 0 and 0.6
112+
MaxResults: 10,
113+
TaskContext: "reviewing code_review findings",
114+
})
115+
if err != nil {
116+
t.Fatalf("Recall: %v", err)
117+
}
118+
if len(recall.Memories) < 2 {
119+
t.Fatalf("expected 2 memories, got %d", len(recall.Memories))
120+
}
121+
if recall.Memories[0].Source != "code_review" {
122+
t.Errorf("expected code_review source first (boosted), got %s", recall.Memories[0].Source)
123+
}
124+
}
125+
126+
func TestRecall_RelevanceClamped(t *testing.T) {
127+
s := newTestStore(t)
128+
ctx := context.Background()
129+
130+
_, _ = s.Store(ctx, StoreRequest{
131+
Entries: []StoreEntry{
132+
{Text: "Perfect match", Embedding: makeEmbedding(0, 8), Tags: []string{"auth"}},
133+
},
134+
})
135+
136+
// Exact embedding match + boost tag + task context = would exceed 1.0
137+
recall, _ := s.Recall(ctx, RecallRequest{
138+
Query: "auth",
139+
QueryEmbedding: makeEmbedding(0, 8),
140+
MaxResults: 10,
141+
BoostTags: []string{"auth"},
142+
TaskContext: "auth",
143+
})
144+
if len(recall.Memories) != 1 {
145+
t.Fatalf("expected 1 memory, got %d", len(recall.Memories))
146+
}
147+
if recall.Memories[0].Relevance > 1.0 {
148+
t.Errorf("relevance should be clamped to 1.0, got %.3f", recall.Memories[0].Relevance)
149+
}
150+
}
151+
152+
func TestRecall_BoostTags_Empty_NoEffect(t *testing.T) {
153+
s := newTestStore(t)
154+
ctx := context.Background()
155+
156+
_, _ = s.Store(ctx, StoreRequest{
157+
Entries: []StoreEntry{
158+
{Text: "Entry A", Embedding: makeEmbedding(0, 8), Tags: []string{"a"}},
159+
{Text: "Entry B", Embedding: makeEmbedding(0.6, 8), Tags: []string{"b"}},
160+
},
161+
})
162+
163+
// No boost tags — ranking should be purely by similarity
164+
recall, _ := s.Recall(ctx, RecallRequest{
165+
Query: "test",
166+
QueryEmbedding: makeEmbedding(0, 8),
167+
MaxResults: 10,
168+
})
169+
if len(recall.Memories) != 2 {
170+
t.Fatalf("expected 2 memories, got %d", len(recall.Memories))
171+
}
172+
// Entry A is closer to query (angle 0 vs 0.6)
173+
if recall.Memories[0].Text != "Entry A" {
174+
t.Errorf("expected Entry A first (closer), got %s", recall.Memories[0].Text)
175+
}
176+
}

‎pkg/memory/sqlite.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -352,6 +352,15 @@ func (s *SQLiteStore) Recall(ctx context.Context, req RecallRequest) (*RecallRes
352352
}
353353
_ = rows.Close()
354354

355+
// Build boost tag set for O(1) lookup
356+
boostTagSet := make(map[string]bool, len(req.BoostTags))
357+
for _, t := range req.BoostTags {
358+
boostTagSet[t] = true
359+
}
360+
361+
// Lowercase task context for substring matching
362+
taskCtxLower := strings.ToLower(req.TaskContext)
363+
355364
var candidates []scored
356365
now := time.Now()
357366

@@ -378,6 +387,36 @@ func (s *SQLiteStore) Recall(ctx context.Context, req RecallRequest) (*RecallRes
378387

379388
relevance := (1.0-recencyWeight)*similarity + recencyWeight*recency
380389

390+
// Boost for matching tags
391+
if len(boostTagSet) > 0 {
392+
for _, tag := range tags {
393+
if boostTagSet[tag] {
394+
relevance += 0.1
395+
break
396+
}
397+
}
398+
}
399+
400+
// Boost for task context match (source or text substring)
401+
if taskCtxLower != "" {
402+
if r.source != "" && strings.Contains(taskCtxLower, strings.ToLower(r.source)) {
403+
relevance += 0.05
404+
}
405+
if strings.Contains(strings.ToLower(r.text), taskCtxLower) {
406+
relevance += 0.05
407+
}
408+
}
409+
410+
// Clamp relevance to [0, 1]
411+
if relevance > 1.0 {
412+
relevance = 1.0
413+
}
414+
415+
// Apply minimum relevance filter
416+
if req.MinRelevance > 0 && relevance < req.MinRelevance {
417+
continue
418+
}
419+
381420
candidates = append(candidates, scored{
382421
memory: RecalledMemory{
383422
ID: r.id,

‎pkg/memory/store.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,15 @@ type RecallRequest struct {
9696
MaxResults int `json:"max_results,omitempty"`
9797
RecencyWeight float64 `json:"recency_weight,omitempty"`
9898
IncludeExpired bool `json:"include_expired,omitempty"`
99+
// TaskContext provides additional context about the current task.
100+
// When set, memories with matching tags or source are boosted.
101+
TaskContext string `json:"task_context,omitempty"`
102+
// BoostTags are tags that receive a relevance boost during ranking.
103+
// Useful for prioritizing domain-specific memories for the current task.
104+
BoostTags []string `json:"boost_tags,omitempty"`
105+
// MinRelevance filters out memories below this relevance score (0-1).
106+
// Default: 0 (no filtering).
107+
MinRelevance float64 `json:"min_relevance,omitempty"`
99108
}
100109

101110
// RecallResult is the output of a recall operation.

0 commit comments

Comments
 (0)