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

@@ -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{{