feat: harden search ranking and deduplication

This commit is contained in:
zhaoyilun
2026-07-10 09:39:18 +08:00
parent 1d6dd2f504
commit 6ca6c08b62
13 changed files with 368 additions and 56 deletions

View File

@@ -30,24 +30,32 @@ func SearchContactsFromMessages(messages []Message, query SearchContactsQuery) S
addMessageContact(unique, message.ReceiverID)
}
queryText := strings.ToLower(strings.TrimSpace(query.Query))
ids := make([]string, 0, len(unique))
for id := range unique {
if queryText == "" || strings.Contains(strings.ToLower(id), queryText) || strings.Contains(strings.ToLower(unique[id].DisplayName), queryText) {
ids = append(ids, id)
queryText := normalizedSearchText(query.Query)
contacts := make([]Contact, 0, len(unique))
for _, contact := range unique {
if contactMatchesQuery(contact, queryText) {
contacts = append(contacts, contact)
}
}
sort.Strings(ids)
sort.Slice(contacts, func(i, j int) bool {
leftRank := contactSearchRank(contacts[i], queryText)
rightRank := contactSearchRank(contacts[j], queryText)
if leftRank != rightRank {
return leftRank < rightRank
}
leftID := strings.ToLower(contacts[i].ContactID)
rightID := strings.ToLower(contacts[j].ContactID)
if leftID != rightID {
return leftID < rightID
}
return contacts[i].ContactID < contacts[j].ContactID
})
limit := query.Limit
if limit <= 0 || limit > len(ids) {
limit = len(ids)
if limit <= 0 || limit > len(contacts) {
limit = len(contacts)
}
contacts := make([]Contact, 0, limit)
for _, id := range ids[:limit] {
contacts = append(contacts, unique[id])
}
return SearchContactsResult{Contacts: contacts}
return SearchContactsResult{Contacts: contacts[:limit]}
}
func addMessageContact(unique map[string]Contact, id string) {
@@ -55,10 +63,11 @@ func addMessageContact(unique map[string]Contact, id string) {
if trimmed == "" {
return
}
if _, exists := unique[trimmed]; exists {
key := strings.ToLower(trimmed)
if _, exists := unique[key]; exists {
return
}
unique[trimmed] = Contact{
unique[key] = Contact{
ContactID: trimmed,
DisplayName: trimmed,
Account: trimmed,
@@ -67,3 +76,38 @@ func addMessageContact(unique map[string]Contact, id string) {
RawRef: "message_jid",
}
}
func normalizedSearchText(value string) string {
return strings.ToLower(strings.TrimSpace(value))
}
func contactMatchesQuery(contact Contact, queryText string) bool {
if queryText == "" {
return true
}
return strings.Contains(strings.ToLower(contact.ContactID), queryText) ||
strings.Contains(strings.ToLower(contact.DisplayName), queryText) ||
strings.Contains(strings.ToLower(contact.Account), queryText)
}
func contactSearchRank(contact Contact, queryText string) int {
if queryText == "" {
return 0
}
values := []string{
strings.ToLower(contact.ContactID),
strings.ToLower(contact.DisplayName),
strings.ToLower(contact.Account),
}
for _, value := range values {
if value == queryText {
return 0
}
}
for _, value := range values {
if strings.HasPrefix(value, queryText) {
return 1
}
}
return 2
}

View File

@@ -20,3 +20,19 @@ func TestSearchContactsFromMessagesMatchesBareJIDs(t *testing.T) {
t.Fatalf("SearchContactsFromMessages() = %#v, want %#v", got, want)
}
}
func TestSearchContactsRanksExactAndDeduplicates(t *testing.T) {
messages := []Message{
{ID: "msg-1", SenderID: "zzz-target@imopenfire1-lanzhou", ReceiverID: "target@imopenfire1-lanzhou"},
{ID: "msg-2", SenderID: "TARGET@imopenfire1-lanzhou", ReceiverID: "other@imopenfire1-lanzhou"},
}
got := SearchContactsFromMessages(messages, SearchContactsQuery{Query: "target@imopenfire1-lanzhou", Limit: 10})
want := SearchContactsResult{Contacts: []Contact{
{ContactID: "target@imopenfire1-lanzhou", DisplayName: "target@imopenfire1-lanzhou", Account: "target@imopenfire1-lanzhou", Source: "local_readonly", Confidence: 0.6, RawRef: "message_jid"},
{ContactID: "zzz-target@imopenfire1-lanzhou", DisplayName: "zzz-target@imopenfire1-lanzhou", Account: "zzz-target@imopenfire1-lanzhou", Source: "local_readonly", Confidence: 0.6, RawRef: "message_jid"},
}}
if !reflect.DeepEqual(got, want) {
t.Fatalf("SearchContactsFromMessages() = %#v, want %#v", got, want)
}
}

View File

@@ -38,24 +38,32 @@ func SearchGroupsFromMessages(messages []Message, query SearchGroupsQuery) Searc
}
}
queryText := strings.ToLower(strings.TrimSpace(query.Query))
ids := make([]string, 0, len(unique))
for id := range unique {
if queryText == "" || strings.Contains(strings.ToLower(id), queryText) || strings.Contains(strings.ToLower(unique[id].DisplayName), queryText) {
ids = append(ids, id)
queryText := normalizedSearchText(query.Query)
groups := make([]Group, 0, len(unique))
for _, group := range unique {
if groupMatchesQuery(group, queryText) {
groups = append(groups, group)
}
}
sort.Strings(ids)
sort.Slice(groups, func(i, j int) bool {
leftRank := groupSearchRank(groups[i], queryText)
rightRank := groupSearchRank(groups[j], queryText)
if leftRank != rightRank {
return leftRank < rightRank
}
leftID := strings.ToLower(groups[i].GroupID)
rightID := strings.ToLower(groups[j].GroupID)
if leftID != rightID {
return leftID < rightID
}
return groups[i].GroupID < groups[j].GroupID
})
limit := query.Limit
if limit <= 0 || limit > len(ids) {
limit = len(ids)
if limit <= 0 || limit > len(groups) {
limit = len(groups)
}
groups := make([]Group, 0, limit)
for _, id := range ids[:limit] {
groups = append(groups, unique[id])
}
return SearchGroupsResult{Groups: groups}
return SearchGroupsResult{Groups: groups[:limit]}
}
func addMessageGroup(unique map[string]Group, id string) {
@@ -63,10 +71,11 @@ func addMessageGroup(unique map[string]Group, id string) {
if trimmed == "" || !isLikelyGroupJID(trimmed) {
return
}
if _, exists := unique[trimmed]; exists {
key := strings.ToLower(trimmed)
if _, exists := unique[key]; exists {
return
}
unique[trimmed] = Group{
unique[key] = Group{
GroupID: trimmed,
DisplayName: trimmed,
Source: "local_readonly",
@@ -79,3 +88,32 @@ func isLikelyGroupJID(id string) bool {
lower := strings.ToLower(strings.TrimSpace(id))
return strings.Contains(lower, "conference") || strings.Contains(lower, "muc") || strings.Contains(lower, "group")
}
func groupMatchesQuery(group Group, queryText string) bool {
if queryText == "" {
return true
}
return strings.Contains(strings.ToLower(group.GroupID), queryText) ||
strings.Contains(strings.ToLower(group.DisplayName), queryText)
}
func groupSearchRank(group Group, queryText string) int {
if queryText == "" {
return 0
}
values := []string{
strings.ToLower(group.GroupID),
strings.ToLower(group.DisplayName),
}
for _, value := range values {
if value == queryText {
return 0
}
}
for _, value := range values {
if strings.HasPrefix(value, queryText) {
return 1
}
}
return 2
}

View File

@@ -20,3 +20,20 @@ func TestSearchGroupsFromMessagesMatchesGroupchatJIDs(t *testing.T) {
t.Fatalf("SearchGroupsFromMessages() = %#v, want %#v", got, want)
}
}
func TestSearchGroupsRanksExactAndDeduplicates(t *testing.T) {
messages := []Message{
{ID: "msg-1", ConversationType: "groupchat", SenderID: "zzz-project@conference.imopenfire1-lanzhou", ReceiverID: "member@imopenfire1-lanzhou"},
{ID: "msg-2", ConversationType: "groupchat", SenderID: "project@conference.imopenfire1-lanzhou", ReceiverID: "member@imopenfire1-lanzhou"},
{ID: "msg-3", ConversationType: "groupchat", SenderID: "PROJECT@conference.imopenfire1-lanzhou", ReceiverID: "member@imopenfire1-lanzhou"},
}
got := SearchGroupsFromMessages(messages, SearchGroupsQuery{Query: "project@conference.imopenfire1-lanzhou", Limit: 10})
want := SearchGroupsResult{Groups: []Group{
{GroupID: "project@conference.imopenfire1-lanzhou", DisplayName: "project@conference.imopenfire1-lanzhou", Source: "local_readonly", Confidence: 0.6, RawRef: "groupchat_jid"},
{GroupID: "zzz-project@conference.imopenfire1-lanzhou", DisplayName: "zzz-project@conference.imopenfire1-lanzhou", Source: "local_readonly", Confidence: 0.6, RawRef: "groupchat_jid"},
}}
if !reflect.DeepEqual(got, want) {
t.Fatalf("SearchGroupsFromMessages() = %#v, want %#v", got, want)
}
}

View File

@@ -2,6 +2,7 @@ package tools
import (
"context"
"sort"
"strings"
"isphere-ai-bridge/internal/isphere"
@@ -73,7 +74,7 @@ func groupsFromDisplayEntities(entities []msglib.DisplayEntity) []isphere.Group
return groups
}
func mergeContacts(primary []isphere.Contact, fallback []isphere.Contact, limit int) isphere.SearchContactsResult {
func mergeContacts(primary []isphere.Contact, fallback []isphere.Contact, query string, limit int) isphere.SearchContactsResult {
seen := map[string]bool{}
contacts := make([]isphere.Contact, 0, len(primary)+len(fallback))
for _, contact := range append(primary, fallback...) {
@@ -83,14 +84,15 @@ func mergeContacts(primary []isphere.Contact, fallback []isphere.Contact, limit
}
seen[strings.ToLower(id)] = true
contacts = append(contacts, contact)
if limit > 0 && len(contacts) >= limit {
break
}
}
sortContactsForQuery(contacts, query)
if limit > 0 && len(contacts) > limit {
contacts = contacts[:limit]
}
return isphere.SearchContactsResult{Contacts: contacts}
}
func mergeGroups(primary []isphere.Group, fallback []isphere.Group, limit int) isphere.SearchGroupsResult {
func mergeGroups(primary []isphere.Group, fallback []isphere.Group, query string, limit int) isphere.SearchGroupsResult {
seen := map[string]bool{}
groups := make([]isphere.Group, 0, len(primary)+len(fallback))
for _, group := range append(primary, fallback...) {
@@ -100,9 +102,87 @@ func mergeGroups(primary []isphere.Group, fallback []isphere.Group, limit int) i
}
seen[strings.ToLower(id)] = true
groups = append(groups, group)
if limit > 0 && len(groups) >= limit {
break
}
}
sortGroupsForQuery(groups, query)
if limit > 0 && len(groups) > limit {
groups = groups[:limit]
}
return isphere.SearchGroupsResult{Groups: groups}
}
func sortContactsForQuery(contacts []isphere.Contact, query string) {
queryText := strings.ToLower(strings.TrimSpace(query))
sort.Slice(contacts, func(i, j int) bool {
leftRank := contactRankForQuery(contacts[i], queryText)
rightRank := contactRankForQuery(contacts[j], queryText)
if leftRank != rightRank {
return leftRank < rightRank
}
leftID := strings.ToLower(contacts[i].ContactID)
rightID := strings.ToLower(contacts[j].ContactID)
if leftID != rightID {
return leftID < rightID
}
return contacts[i].ContactID < contacts[j].ContactID
})
}
func contactRankForQuery(contact isphere.Contact, queryText string) int {
if queryText == "" {
return 0
}
values := []string{
strings.ToLower(contact.ContactID),
strings.ToLower(contact.DisplayName),
strings.ToLower(contact.Account),
}
for _, value := range values {
if value == queryText {
return 0
}
}
for _, value := range values {
if strings.HasPrefix(value, queryText) {
return 1
}
}
return 2
}
func sortGroupsForQuery(groups []isphere.Group, query string) {
queryText := strings.ToLower(strings.TrimSpace(query))
sort.Slice(groups, func(i, j int) bool {
leftRank := groupRankForQuery(groups[i], queryText)
rightRank := groupRankForQuery(groups[j], queryText)
if leftRank != rightRank {
return leftRank < rightRank
}
leftID := strings.ToLower(groups[i].GroupID)
rightID := strings.ToLower(groups[j].GroupID)
if leftID != rightID {
return leftID < rightID
}
return groups[i].GroupID < groups[j].GroupID
})
}
func groupRankForQuery(group isphere.Group, queryText string) int {
if queryText == "" {
return 0
}
values := []string{
strings.ToLower(group.GroupID),
strings.ToLower(group.DisplayName),
}
for _, value := range values {
if value == queryText {
return 0
}
}
for _, value := range values {
if strings.HasPrefix(value, queryText) {
return 1
}
}
return 2
}

View File

@@ -52,7 +52,7 @@ func RegisterISphereContactToolsWithDisplayEntities(server *mcp.Server, source R
if err != nil {
return nil, nil, err
}
contacts = mergeContacts(contactsFromDisplayEntities(entities), logContacts.Contacts, input.Limit)
contacts = mergeContacts(contactsFromDisplayEntities(entities), logContacts.Contacts, input.Query, input.Limit)
}
return nil, searchContactsResultToMap(contacts, started, time.Now().UTC()), nil
})

View File

@@ -131,6 +131,56 @@ func TestISphereSearchContactsToolUsesInjectedDisplayEntities(t *testing.T) {
}
}
func TestISphereSearchContactsToolRanksExactAndDeduplicatesDisplayEntities(t *testing.T) {
fakeMessages := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{
ID: "msg-1",
SenderID: "target@imopenfire1-lanzhou",
ReceiverID: "other@imopenfire1-lanzhou",
}}},
}
fakeDisplay := &fakeDisplayEntitySource{entities: []msglib.DisplayEntity{
{EntityType: msglib.EntityTypeContacts, SourceTable: "TD_CustomEffigy", JID: "zzz-target@imopenfire1-lanzhou", DisplayName: "ZZZ Target", Confidence: 0.8},
{EntityType: msglib.EntityTypeContacts, SourceTable: "TD_CustomEffigy", JID: "target@imopenfire1-lanzhou", DisplayName: "Target", Confidence: 0.9},
}}
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
RegisterISphereContactToolsWithDisplayEntities(server, fakeMessages, fakeDisplay)
})
defer cleanup()
callResult, err := session.CallTool(context.Background(), &mcp.CallToolParams{
Name: ToolNameSearchContacts,
Arguments: map[string]any{"query": "target@imopenfire1-lanzhou", "limit": 10},
})
if err != nil {
t.Fatalf("call %s: %v", ToolNameSearchContacts, err)
}
if callResult.IsError {
t.Fatalf("call result is error: %+v", callResult)
}
var decoded struct {
Contacts []struct {
ContactID string `json:"contact_id"`
Source string `json:"source"`
RawRef string `json:"raw_ref"`
} `json:"contacts"`
}
payload, _ := json.Marshal(callResult.StructuredContent)
if err := json.Unmarshal(payload, &decoded); err != nil {
t.Fatalf("decode structured content %s: %v", payload, err)
}
if len(decoded.Contacts) != 2 {
t.Fatalf("contacts = %+v, want exact and substring only", decoded.Contacts)
}
if decoded.Contacts[0].ContactID != "target@imopenfire1-lanzhou" || decoded.Contacts[0].Source != "msglib_readonly" || decoded.Contacts[0].RawRef != "msglib:TD_CustomEffigy" {
t.Fatalf("first contact should be exact MsgLib match with source metadata, got %+v", decoded.Contacts[0])
}
if decoded.Contacts[1].ContactID != "zzz-target@imopenfire1-lanzhou" {
t.Fatalf("second contact should be substring match, got %+v", decoded.Contacts[1])
}
}
func TestISphereSearchContactsToolValidatesContractArgs(t *testing.T) {
fake := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{

View File

@@ -52,7 +52,7 @@ func RegisterISphereGroupToolsWithDisplayEntities(server *mcp.Server, source Rec
if err != nil {
return nil, nil, err
}
groups = mergeGroups(groupsFromDisplayEntities(entities), logGroups.Groups, input.Limit)
groups = mergeGroups(groupsFromDisplayEntities(entities), logGroups.Groups, input.Query, input.Limit)
}
return nil, searchGroupsResultToMap(groups, started, time.Now().UTC()), nil
})

View File

@@ -135,6 +135,57 @@ func TestISphereSearchGroupsToolUsesInjectedDisplayEntities(t *testing.T) {
}
}
func TestISphereSearchGroupsToolRanksExactAndDeduplicatesDisplayEntities(t *testing.T) {
fakeMessages := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{
ID: "msg-1",
ConversationType: "groupchat",
SenderID: "target@conference.imopenfire1-lanzhou",
ReceiverID: "member@imopenfire1-lanzhou",
}}},
}
fakeDisplay := &fakeDisplayEntitySource{entities: []msglib.DisplayEntity{
{EntityType: msglib.EntityTypeGroups, SourceTable: "TD_WorkGroupAuth", JID: "zzz-target@conference.imopenfire1-lanzhou", DisplayName: "ZZZ Target Group", Confidence: 0.8},
{EntityType: msglib.EntityTypeGroups, SourceTable: "TD_WorkGroupAuth", JID: "target@conference.imopenfire1-lanzhou", DisplayName: "Target Group", Confidence: 0.9},
}}
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
RegisterISphereGroupToolsWithDisplayEntities(server, fakeMessages, fakeDisplay)
})
defer cleanup()
callResult, err := session.CallTool(context.Background(), &mcp.CallToolParams{
Name: ToolNameSearchGroups,
Arguments: map[string]any{"query": "target@conference.imopenfire1-lanzhou", "limit": 10},
})
if err != nil {
t.Fatalf("call %s: %v", ToolNameSearchGroups, err)
}
if callResult.IsError {
t.Fatalf("call result is error: %+v", callResult)
}
var decoded struct {
Groups []struct {
GroupID string `json:"group_id"`
Source string `json:"source"`
RawRef string `json:"raw_ref"`
} `json:"groups"`
}
payload, _ := json.Marshal(callResult.StructuredContent)
if err := json.Unmarshal(payload, &decoded); err != nil {
t.Fatalf("decode structured content %s: %v", payload, err)
}
if len(decoded.Groups) != 2 {
t.Fatalf("groups = %+v, want exact and substring only", decoded.Groups)
}
if decoded.Groups[0].GroupID != "target@conference.imopenfire1-lanzhou" || decoded.Groups[0].Source != "msglib_readonly" || decoded.Groups[0].RawRef != "msglib:TD_WorkGroupAuth" {
t.Fatalf("first group should be exact MsgLib match with source metadata, got %+v", decoded.Groups[0])
}
if decoded.Groups[1].GroupID != "zzz-target@conference.imopenfire1-lanzhou" {
t.Fatalf("second group should be substring match, got %+v", decoded.Groups[1])
}
}
func TestISphereSearchGroupsToolValidatesContractArgs(t *testing.T) {
fake := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{