feat: enrich search with msglib display entities
This commit is contained in:
108
internal/tools/display_entities.go
Normal file
108
internal/tools/display_entities.go
Normal 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}
|
||||
}
|
||||
@@ -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
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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{{
|
||||
|
||||
Reference in New Issue
Block a user