feat: harden search ranking and deduplication
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
})
|
||||
|
||||
@@ -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{{
|
||||
|
||||
@@ -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
|
||||
})
|
||||
|
||||
@@ -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{{
|
||||
|
||||
Reference in New Issue
Block a user