feat: enrich search with msglib display entities

This commit is contained in:
zhaoyilun
2026-07-10 03:09:37 +08:00
parent 8be4a6a6c5
commit 1277316663
11 changed files with 400 additions and 30 deletions

View File

@@ -0,0 +1,108 @@
package tools
import (
"context"
"strings"
"isphere-ai-bridge/internal/isphere"
"isphere-ai-bridge/internal/msglib"
)
type DisplayEntitySource interface {
DisplayEntities(ctx context.Context, opts msglib.DisplayEntitiesOptions) ([]msglib.DisplayEntity, error)
}
func displayEntityQueryLimit(limit int) int {
if limit <= 0 {
return msglib.DefaultDisplayEntityLimit
}
if limit > msglib.MaxDisplayEntityLimit {
return msglib.MaxDisplayEntityLimit
}
return limit
}
func contactsFromDisplayEntities(entities []msglib.DisplayEntity) []isphere.Contact {
contacts := make([]isphere.Contact, 0, len(entities))
for _, entity := range entities {
id := strings.TrimSpace(entity.JID)
name := strings.TrimSpace(entity.DisplayName)
if id == "" {
id = name
}
if name == "" {
name = id
}
if id == "" {
continue
}
contacts = append(contacts, isphere.Contact{
ContactID: id,
DisplayName: name,
Account: strings.TrimSpace(entity.JID),
Source: "msglib_readonly",
Confidence: entity.Confidence,
RawRef: "msglib:" + entity.SourceTable,
})
}
return contacts
}
func groupsFromDisplayEntities(entities []msglib.DisplayEntity) []isphere.Group {
groups := make([]isphere.Group, 0, len(entities))
for _, entity := range entities {
id := strings.TrimSpace(entity.JID)
name := strings.TrimSpace(entity.DisplayName)
if id == "" {
id = name
}
if name == "" {
name = id
}
if id == "" {
continue
}
groups = append(groups, isphere.Group{
GroupID: id,
DisplayName: name,
Source: "msglib_readonly",
Confidence: entity.Confidence,
RawRef: "msglib:" + entity.SourceTable,
})
}
return groups
}
func mergeContacts(primary []isphere.Contact, fallback []isphere.Contact, limit int) isphere.SearchContactsResult {
seen := map[string]bool{}
contacts := make([]isphere.Contact, 0, len(primary)+len(fallback))
for _, contact := range append(primary, fallback...) {
id := strings.TrimSpace(contact.ContactID)
if id == "" || seen[strings.ToLower(id)] {
continue
}
seen[strings.ToLower(id)] = true
contacts = append(contacts, contact)
if limit > 0 && len(contacts) >= limit {
break
}
}
return isphere.SearchContactsResult{Contacts: contacts}
}
func mergeGroups(primary []isphere.Group, fallback []isphere.Group, limit int) isphere.SearchGroupsResult {
seen := map[string]bool{}
groups := make([]isphere.Group, 0, len(primary)+len(fallback))
for _, group := range append(primary, fallback...) {
id := strings.TrimSpace(group.GroupID)
if id == "" || seen[strings.ToLower(id)] {
continue
}
seen[strings.ToLower(id)] = true
groups = append(groups, group)
if limit > 0 && len(groups) >= limit {
break
}
}
return isphere.SearchGroupsResult{Groups: groups}
}

View File

@@ -7,6 +7,7 @@ import (
"github.com/modelcontextprotocol/go-sdk/mcp"
"isphere-ai-bridge/internal/isphere"
"isphere-ai-bridge/internal/msglib"
)
const ToolNameSearchContacts = "isphere_search_contacts"
@@ -20,6 +21,10 @@ type SearchContactsArgs struct {
}
func RegisterISphereContactTools(server *mcp.Server, source ReceiveMessagesSource) {
RegisterISphereContactToolsWithDisplayEntities(server, source, nil)
}
func RegisterISphereContactToolsWithDisplayEntities(server *mcp.Server, source ReceiveMessagesSource, displaySource DisplayEntitySource) {
if source == nil {
source = isphere.EncryptedPacketLogSource{}
}
@@ -36,7 +41,19 @@ func RegisterISphereContactTools(server *mcp.Server, source ReceiveMessagesSourc
if err != nil {
return nil, nil, err
}
contacts := isphere.SearchContactsFromMessages(messages.Messages, isphere.SearchContactsQuery{Query: input.Query, Limit: input.Limit})
logContacts := isphere.SearchContactsFromMessages(messages.Messages, isphere.SearchContactsQuery{Query: input.Query, Limit: input.Limit})
contacts := logContacts
if displaySource != nil {
entities, err := displaySource.DisplayEntities(ctx, msglib.DisplayEntitiesOptions{
EntityType: msglib.EntityTypeContacts,
Query: input.Query,
Limit: displayEntityQueryLimit(input.Limit),
})
if err != nil {
return nil, nil, err
}
contacts = mergeContacts(contactsFromDisplayEntities(entities), logContacts.Contacts, input.Limit)
}
return nil, searchContactsResultToMap(contacts, started, time.Now().UTC()), nil
})
}

View File

@@ -8,6 +8,7 @@ import (
"github.com/modelcontextprotocol/go-sdk/mcp"
"isphere-ai-bridge/internal/isphere"
"isphere-ai-bridge/internal/msglib"
)
func TestISphereSearchContactsToolReturnsJIDContacts(t *testing.T) {
@@ -75,6 +76,61 @@ func TestISphereSearchContactsToolReturnsJIDContacts(t *testing.T) {
}
}
func TestISphereSearchContactsToolUsesInjectedDisplayEntities(t *testing.T) {
fakeMessages := &fakeReceiveMessagesSource{}
fakeDisplay := &fakeDisplayEntitySource{entities: []msglib.DisplayEntity{{
EntityType: msglib.EntityTypeContacts,
SourceTable: "TD_CustomEffigy",
JID: "alice@example",
DisplayName: "Alice Zhang",
Confidence: 0.9,
MatchedColumns: []string{"PersonJid", "PersonName"},
}}}
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": "alice", "limit": 5},
})
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"`
DisplayName string `json:"display_name"`
Account string `json:"account"`
Source string `json:"source"`
Confidence float64 `json:"confidence"`
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) != 1 {
t.Fatalf("contacts = %+v, want one", decoded.Contacts)
}
contact := decoded.Contacts[0]
if contact.ContactID != "alice@example" || contact.DisplayName != "Alice Zhang" || contact.Account != "alice@example" {
t.Fatalf("unexpected MsgLib contact identity: %+v", contact)
}
if contact.Source != "msglib_readonly" || contact.Confidence != 0.9 || contact.RawRef != "msglib:TD_CustomEffigy" {
t.Fatalf("unexpected MsgLib contact metadata: %+v", contact)
}
if len(fakeDisplay.queries) != 1 || fakeDisplay.queries[0].EntityType != msglib.EntityTypeContacts || fakeDisplay.queries[0].Query != "alice" || fakeDisplay.queries[0].Limit != 5 {
t.Fatalf("display queries = %+v", fakeDisplay.queries)
}
}
func TestISphereSearchContactsToolValidatesContractArgs(t *testing.T) {
fake := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{
@@ -121,3 +177,13 @@ func TestISphereSearchContactsToolValidatesContractArgs(t *testing.T) {
t.Fatalf("non-empty cursor was accepted before pagination exists")
}
}
type fakeDisplayEntitySource struct {
queries []msglib.DisplayEntitiesOptions
entities []msglib.DisplayEntity
}
func (f *fakeDisplayEntitySource) DisplayEntities(_ context.Context, opts msglib.DisplayEntitiesOptions) ([]msglib.DisplayEntity, error) {
f.queries = append(f.queries, opts)
return f.entities, nil
}

View File

@@ -7,6 +7,7 @@ import (
"github.com/modelcontextprotocol/go-sdk/mcp"
"isphere-ai-bridge/internal/isphere"
"isphere-ai-bridge/internal/msglib"
)
const ToolNameSearchGroups = "isphere_search_groups"
@@ -20,6 +21,10 @@ type SearchGroupsArgs struct {
}
func RegisterISphereGroupTools(server *mcp.Server, source ReceiveMessagesSource) {
RegisterISphereGroupToolsWithDisplayEntities(server, source, nil)
}
func RegisterISphereGroupToolsWithDisplayEntities(server *mcp.Server, source ReceiveMessagesSource, displaySource DisplayEntitySource) {
if source == nil {
source = isphere.EncryptedPacketLogSource{}
}
@@ -36,7 +41,19 @@ func RegisterISphereGroupTools(server *mcp.Server, source ReceiveMessagesSource)
if err != nil {
return nil, nil, err
}
groups := isphere.SearchGroupsFromMessages(messages.Messages, isphere.SearchGroupsQuery{Query: input.Query, Limit: input.Limit})
logGroups := isphere.SearchGroupsFromMessages(messages.Messages, isphere.SearchGroupsQuery{Query: input.Query, Limit: input.Limit})
groups := logGroups
if displaySource != nil {
entities, err := displaySource.DisplayEntities(ctx, msglib.DisplayEntitiesOptions{
EntityType: msglib.EntityTypeGroups,
Query: input.Query,
Limit: displayEntityQueryLimit(input.Limit),
})
if err != nil {
return nil, nil, err
}
groups = mergeGroups(groupsFromDisplayEntities(entities), logGroups.Groups, input.Limit)
}
return nil, searchGroupsResultToMap(groups, started, time.Now().UTC()), nil
})
}

View File

@@ -8,6 +8,7 @@ import (
"github.com/modelcontextprotocol/go-sdk/mcp"
"isphere-ai-bridge/internal/isphere"
"isphere-ai-bridge/internal/msglib"
)
func TestISphereSearchGroupsToolReturnsJIDGroups(t *testing.T) {
@@ -80,6 +81,60 @@ func TestISphereSearchGroupsToolReturnsJIDGroups(t *testing.T) {
}
}
func TestISphereSearchGroupsToolUsesInjectedDisplayEntities(t *testing.T) {
fakeMessages := &fakeReceiveMessagesSource{}
fakeDisplay := &fakeDisplayEntitySource{entities: []msglib.DisplayEntity{{
EntityType: msglib.EntityTypeGroups,
SourceTable: "TD_WorkGroupAuth",
JID: "project-room@conference",
DisplayName: "Project Room",
Confidence: 0.9,
MatchedColumns: []string{"GroupJID", "GroupName"},
}}}
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": "project", "limit": 5},
})
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"`
DisplayName string `json:"display_name"`
Source string `json:"source"`
Confidence float64 `json:"confidence"`
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) != 1 {
t.Fatalf("groups = %+v, want one", decoded.Groups)
}
group := decoded.Groups[0]
if group.GroupID != "project-room@conference" || group.DisplayName != "Project Room" {
t.Fatalf("unexpected MsgLib group identity: %+v", group)
}
if group.Source != "msglib_readonly" || group.Confidence != 0.9 || group.RawRef != "msglib:TD_WorkGroupAuth" {
t.Fatalf("unexpected MsgLib group metadata: %+v", group)
}
if len(fakeDisplay.queries) != 1 || fakeDisplay.queries[0].EntityType != msglib.EntityTypeGroups || fakeDisplay.queries[0].Query != "project" || fakeDisplay.queries[0].Limit != 5 {
t.Fatalf("display queries = %+v", fakeDisplay.queries)
}
}
func TestISphereSearchGroupsToolValidatesContractArgs(t *testing.T) {
fake := &fakeReceiveMessagesSource{
result: isphere.ReceiveMessagesResult{Messages: []isphere.Message{{