feat: better dialogs

internal/tg/dialogs.go

@@ -11,30 +11,26 @@ import (
 	"github.com/gotd/td/tg"
 	mcp "github.com/metoro-io/mcp-golang"
 	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
 )
 
 // DialogType represents the type of dialog for filtering
 type DialogType string
 
 const (
-	// DialogTypeAll represents all types of dialogs
+	DialogTypeUnknown DialogType = "unknown"
 	DialogTypeAll     DialogType = ""
-	// DialogTypeUser represents user chats
 	DialogTypeUser    DialogType = "user"
-	// DialogTypeChat represents group chats
+	DialogTypeBot     DialogType = "bot"
 	DialogTypeChat    DialogType = "chat"
-	// DialogTypeChannel represents channels
 	DialogTypeChannel DialogType = "channel"
 
-	// DefaultDialogsLimit is the default limit for dialogs
 	DefaultDialogsLimit = 100
 )
 
 // nolint:lll
 type DialogsArguments struct {
-	Type             DialogType `json:"type,omitempty" jsonschema:"description=Filter dialogs by type (user, chat, channel or empty for all),enum=,enum=user,enum=chat,enum=channel"`
+	//WithLastMessages bool `json:"with_last_messages,omitempty" jsonschema:"description=Include last messages in response"`
-	Limit            int        `json:"limit,omitempty" jsonschema:"description=Maximum number of dialogs to return (max: 100),default=100"`
-	WithLastMessages bool       `json:"with_last_messages,omitempty" jsonschema:"description=Include last messages in response"`
 }
 
 type MessageInfo struct {
@@ -42,99 +38,160 @@ type MessageInfo struct {
 	When     string `json:"when"`
 	Text     string `json:"text"`
 	IsUnread bool   `json:"is_unread,omitempty"`
+	ts       int
 }
 
 type DialogInfo struct {
 	ID          int64        `json:"id"`
 	Type        string       `json:"type"`
-	Title         string        `json:"title"`
+	Name        string       `json:"name"`
-	UnreadCount   int           `json:"unread_count"`
+	LastMessage *MessageInfo `json:"last_message,omitempty"`
-	LastMessageID int           `json:"last_message_id"`
-	IsVerified    bool          `json:"is_verified,omitempty"`
-	LastMessages  []MessageInfo `json:"last_messages,omitempty"`
 }
 
 // GetDialogs returns a list of dialogs (chats, channels, groups)
 func (c *Client) GetDialogs(args DialogsArguments) (*mcp.ToolResponse, error) {
 	var result []DialogInfo
 
-	if args.Limit <= 0 || args.Limit > DefaultDialogsLimit {
+	var dc tg.MessagesDialogsClass
-		args.Limit = DefaultDialogsLimit
-	}
-
-	if args.Type == "" {
-		args.Type = DialogTypeAll
-	}
-
 	client := c.T()
-	if err := client.Run(context.Background(), func(ctx context.Context) error {
+	if err := client.Run(context.Background(), func(ctx context.Context) (err error) {
 		api := client.API()
-		dialogsClass, err := api.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{
+		dc, err = api.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{
 			OffsetPeer: &tg.InputPeerEmpty{},
-			Limit:      20,
 		})
 		if err != nil {
 			return fmt.Errorf("failed to get dialogs: %w", err)
 		}
 
 		// Debug
-		// jsonData, _ := json.Marshal(dialogsClass)
+		//jsonData, _ := json.Marshal(dc)
 		//log.Info().RawJSON("dialogs", cleanJSON(jsonData)).Msg("dialogs")
 
-		var dialogs *tg.MessagesDialogs
+		return nil
-		switch d := dialogsClass.(type) {
+	}); err != nil {
+		return nil, errors.Wrap(err, "failed to get dialogs")
+	}
+
+	d, err := newDialogs(dc)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to get dialogs")
+	}
+
+	info := d.Info()
+
+	sort.Slice(result, func(i, j int) bool {
+		return info[i].LastMessage.ts > result[j].LastMessage.ts
+	})
+
+	jsonData, err := json.Marshal(info)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to marshal response")
+	}
+
+	return mcp.NewToolResponse(mcp.NewTextContent(string(jsonData))), nil
+}
+
+type dialogs struct {
+	tg.MessagesDialogs
+
+	// chat id key
+	messages map[int64]*tg.Message
+	users    map[int64]*tg.User
+	//dialogs  map[string]*tg.Dialog
+	chats    map[int64]*tg.Chat
+	channels map[int64]*tg.Channel
+}
+
+func newDialogs(rawD tg.MessagesDialogsClass) (*dialogs, error) {
+	var d dialogs
+	switch dT := rawD.(type) {
 	case *tg.MessagesDialogs:
-			dialogs = d
+		d = dialogs{MessagesDialogs: *dT}
 	case *tg.MessagesDialogsSlice:
-			dialogs = &tg.MessagesDialogs{
+		d = dialogs{MessagesDialogs: tg.MessagesDialogs{
-				Dialogs:  d.Dialogs,
+			Dialogs:  dT.Dialogs,
-				Messages: d.Messages,
+			Messages: dT.Messages,
-				Chats:    d.Chats,
+			Chats:    dT.Chats,
-				Users:    d.Users,
+			Users:    dT.Users,
-			}
+		}}
+	case *tg.MessagesDialogsNotModified:
 	default:
-			return errors.New("unexpected dialogs response type")
 	}
 
-		messageMap := make(map[string][]*tg.Message)
+	d.messages = make(map[int64]*tg.Message)
-		for _, m := range dialogs.Messages {
+	for _, m := range d.Messages {
-			msg, ok := m.(*tg.Message)
+		switch mT := m.(type) {
+		case *tg.Message:
+			d.messages[getPeerID(mT.PeerID)] = mT
+		case *tg.MessageService, *tg.MessageEmpty:
+		default:
+		}
+	}
+	delete(d.messages, 0)
+
+	d.users = make(map[int64]*tg.User)
+	for _, uc := range d.Users {
+		u, ok := uc.(*tg.User)
 		if !ok {
+			log.Debug().Msgf("newDialogs(%+v): invalid message user", uc)
 			continue
 		}
 
-			if msg.PeerID == nil {
+		d.users[u.GetID()] = u
+	}
+
+	d.chats = make(map[int64]*tg.Chat)
+	d.channels = make(map[int64]*tg.Channel)
+	for _, c := range d.Chats {
+		switch cT := c.(type) {
+		case *tg.Chat:
+			d.chats[cT.ID] = cT
+		case *tg.Channel:
+			d.channels[cT.ID] = cT
+		case *tg.ChatForbidden, *tg.ChannelForbidden, *tg.ChatEmpty:
+		default:
+		}
+	}
+
+	return &d, nil
+}
+
+func (d *dialogs) Info() []DialogInfo {
+	ds := make([]DialogInfo, 0, len(d.Dialogs))
+
+	for _, dItem := range d.Dialogs {
+		info, err := d.processDialog(dItem)
+		if err != nil {
+			log.Debug().Err(err).Str("dialog", dItem.String()).Msg("failed process dialog")
 			continue
 		}
 
-			messageMap[msg.PeerID.String()] = append(messageMap[msg.PeerID.String()], msg)
+		if info.Name == "" {
+			continue
 		}
 
-		usersMap := make(map[string]tg.UserClass)
+		ds = append(ds, info)
-		for _, u := range dialogs.Users {
-			usersMap["Peer"+u.String()] = u
 	}
 
-		result = make([]DialogInfo, 0, len(dialogs.Dialogs))
+	return ds
+}
 
-		for _, dialog := range dialogs.Dialogs {
+func (d *dialogs) processDialog(rawD tg.DialogClass) (DialogInfo, error) {
-			dialogItem, ok := dialog.(*tg.Dialog)
+	dialogItem, ok := rawD.(*tg.Dialog)
 	if !ok {
-				continue
+		return DialogInfo{}, errors.Errorf("newDialogs(%T): invalid dialog type", rawD)
 	}
 
 	var info DialogInfo
-			info.UnreadCount = dialogItem.UnreadCount
-			info.LastMessageID = dialogItem.TopMessage
 
-			if args.WithLastMessages {
+	if msg, ok := d.messages[getPeerID(dialogItem.Peer)]; ok {
-				msgs := messageMap[dialogItem.Peer.String()]
-				for _, msg := range msgs {
 		var who string
 		if msg.FromID != nil {
-						if u, ok := usersMap[msg.FromID.String()]; ok {
+			name, _, err := d.getNameID(msg.FromID)
-							who = u.String()
+			if err != nil {
+				return DialogInfo{}, errors.Wrap(err, "failed to get dialog name")
 			}
+			who = name
 		}
 
 		// Limit message to 20 words
@@ -144,98 +201,103 @@ func (c *Client) GetDialogs(args DialogsArguments) (*mcp.ToolResponse, error) {
 			text = strings.Join(words[:20], " ") + "..."
 		}
 
-					info.LastMessages = append(info.LastMessages, MessageInfo{
+		info.LastMessage = &MessageInfo{
 			Who:      who,
 			When:     time.Unix(int64(msg.Date), 0).Format(time.DateTime),
+			ts:       msg.Date,
 			Text:     text,
 			IsUnread: dialogItem.UnreadCount > 0,
-					})
-				}
 		}
 
-			switch peer := dialogItem.Peer.(type) {
-			case *tg.PeerUser:
-				if args.Type != DialogTypeAll && args.Type != DialogTypeUser {
-					continue
 	}
 
-				for _, userItem := range dialogs.Users {
+	if dialogItem.Peer == nil {
-					user, ok := userItem.(*tg.User)
+		return DialogInfo{}, fmt.Errorf("no peer: %s", dialogItem.String())
-					if !ok || user.ID != peer.UserID {
-						continue
 	}
 
-					info.ID = user.ID
+	var err error
-					info.Type = "user"
+	info.Name, info.ID, err = d.getNameID(dialogItem.Peer)
-					info.Title = getUserName(user)
-					info.IsVerified = user.Verified
-
-					result = append(result, info)
-
-					break
-				}
-
-			case *tg.PeerChat:
-				if args.Type != DialogTypeAll && args.Type != DialogTypeChat {
-					continue
-				}
-
-				for _, chatItem := range dialogs.Chats {
-					chat, ok := chatItem.(*tg.Chat)
-					if !ok || chat.ID != peer.ChatID {
-						continue
-					}
-
-					info.ID = chat.ID
-					info.Type = "chat"
-					info.Title = chat.Title
-
-					result = append(result, info)
-
-					break
-				}
-
-			case *tg.PeerChannel:
-				if args.Type != DialogTypeAll && args.Type != DialogTypeChannel {
-					continue
-				}
-
-				for _, channelItem := range dialogs.Chats {
-					channel, ok := channelItem.(*tg.Channel)
-					if !ok || channel.ID != peer.ChannelID {
-						continue
-					}
-
-					info.ID = channel.ID
-					info.Type = "channel"
-					info.Title = channel.Title
-					info.IsVerified = channel.Verified
-
-					result = append(result, info)
-
-					break
-				}
-			}
-		}
-
-		return nil
-	}); err != nil {
-		return nil, errors.Wrap(err, "failed to get dialogs")
-	}
-
-	sort.Slice(result, func(i, j int) bool {
-		return result[i].LastMessageID > result[j].LastMessageID
-	})
-
-	if len(result) > args.Limit {
-		result = result[:args.Limit]
-	}
-
-	jsonData, err := json.Marshal(result)
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to marshal response")
+		return DialogInfo{}, err
+	}
+
+	info.Type = string(d.getType(dialogItem))
+
+	return info, nil
+}
+
+func (d *dialogs) getNameID(pC tg.PeerClass) (string, int64, error) {
+	var name string
+	var id int64
+	switch p := pC.(type) {
+	case *tg.PeerUser:
+		id = p.GetUserID()
+		u, ok := d.users[id]
+		if !ok {
+			return "", 0, errors.Errorf("peerid(%d): invalid message user", id)
+		}
+		name = getName(u)
+	case *tg.PeerChannel:
+		id = p.GetChannelID()
+		channel, ok := d.channels[id]
+		if !ok {
+			return "", 0, errors.Errorf("peerid(%d): invalid message channel", id)
+		}
+
+		name = getName(channel)
+	case *tg.PeerChat:
+		id = p.GetChatID()
+		chat, ok := d.chats[id]
+		if !ok {
+			return "", 0, errors.Errorf("peerid(%d): invalid message chat", id)
+		}
+
+		name = getName(chat)
+	default:
+		return "", 0, fmt.Errorf("chose author(%T): invalid dialog peer", p)
+	}
+
+	return name, id, nil
+}
+
+func (d *dialogs) getType(rawD *tg.Dialog) DialogType {
+	switch v := rawD.Peer.(type) {
+	case *tg.PeerChannel:
+		return DialogTypeChannel
+	case *tg.PeerChat:
+		return DialogTypeChat
+	case *tg.PeerUser:
+		u, ok := d.users[getPeerID(rawD.Peer)]
+		if !ok {
+			log.Debug().Msgf("getType(%+v): user not found", v)
+			return DialogTypeUser
+		}
+
+		if u.Bot {
+			return DialogTypeBot
+		}
+
+		return DialogTypeUser
+	default:
+		log.Debug().Msgf("getType(%+v): unknown dialog type", v)
+		return DialogTypeUnknown
+	}
+}
+
+func getPeerID(p tg.PeerClass) int64 {
+	if p == nil {
+		return 0
+	}
+
+	switch v := p.(type) {
+	case *tg.PeerChannel:
+		return v.ChannelID
+	case *tg.PeerChat:
+		return v.ChatID
+	case *tg.PeerUser:
+		return v.UserID
+	default:
+		return 0
 	}
 
-	cleanedData := cleanJSON(jsonData)
-	return mcp.NewToolResponse(mcp.NewTextContent(string(cleanedData))), nil
 }

internal/tg/helpers.go

@@ -7,14 +7,26 @@ import (
 	"github.com/tidwall/gjson"
 )
 
-func getUserName(user *tg.User) string {
+func getName(source any) string {
-	if username, ok := user.GetUsername(); ok && username != "" {
+	var name string
-		return "@" + username
+	switch u := source.(type) {
+	case *tg.User:
+		name = u.FirstName
+		if u.LastName != "" {
+			name += " " + u.LastName
 		}
 
-	name := user.FirstName
+		if username, ok := u.GetUsername(); ok && username != "" {
-	if user.LastName != "" {
+			name += " @" + username + ""
-		name += " " + user.LastName
+		}
+	case *tg.Chat:
+		name = u.Title
+	case *tg.Channel:
+		name = u.Title
+
+		if username, ok := u.GetUsername(); ok && username != "" {
+			name += " @" + username + ""
+		}
 	}
 
 	return name

resourse.go

@@ -1,47 +0,0 @@
-package main
-
-import (
-	"encoding/json"
-	"fmt"
-
-	mcp "github.com/metoro-io/mcp-golang"
-)
-
-func sampleResource() (*mcp.ResourceResponse, error) {
-	type Chat struct {
-		ID          int64  `json:"id,omitempty"`
-		Type        string `json:"type"`
-		Title       string `json:"title"`
-		UnreadCount int    `json:"unread_count"`
-	}
-
-	chats := []Chat{
-		{
-			ID:          123456789,
-			Type:        "channel",
-			Title:       "Sample Channel",
-			UnreadCount: 5,
-		},
-		{
-			ID:          987654321,
-			Type:        "group",
-			Title:       "Test Group",
-			UnreadCount: 2,
-		},
-	}
-
-	rss := make([]*mcp.EmbeddedResource, 0, len(chats))
-	for _, chat := range chats {
-		chat.ID = 0
-		uri := fmt.Sprintf("telegram://chats/%d", chat.ID)
-
-		content, err := json.Marshal(chat)
-		if err != nil {
-			return nil, err
-		}
-
-		rss = append(rss, mcp.NewTextEmbeddedResource(uri, string(content), "application/json"))
-	}
-
-	return mcp.NewResourceResponse(rss...), nil
-}

serve.go

@@ -41,36 +41,26 @@ func serve(ctx context.Context, cmd *cli.Command) error {
 
 		log.Info().RawJSON("answer", data).Msg("Check GetMe: OK")
 
-		answer, err = client.GetDialogs(tg.DialogsArguments{WithLastMessages: true})
+		answer, err = client.GetDialogs(tg.DialogsArguments{})
 		if err != nil {
 			return fmt.Errorf("get dialogs: %w", err)
 		}
 
-		data, err = json.MarshalIndent(answer, "", "  ")
+		log.Info().RawJSON("answer", []byte(answer.Content[0].TextContent.Text)).Msg("Check GetDialogs: OK")
-		if err != nil {
-			return fmt.Errorf("marshal: %w", err)
-		}
-
-		log.Info().RawJSON("answer", data).Msg("Check GetDialogs: OK")
 
 		return nil
 	}
 
-	err = server.RegisterTool("tg_me", "Get current account info", client.GetMe)
+	err = server.RegisterTool("tg_me", "Get current telegram account info", client.GetMe)
 	if err != nil {
 		return fmt.Errorf("register tool: %w", err)
 	}
 
-	err = server.RegisterTool("tg_dialogs", "Get list of dialogs (chats, channels, groups)", client.GetDialogs)
+	err = server.RegisterTool("tg_dialogs", "Get list of telegram dialogs (chats, channels, groups)", client.GetDialogs)
 	if err != nil {
 		return fmt.Errorf("register dialogs tool: %w", err)
 	}
 
-	err = server.RegisterResource("telegram://chats", "tg_chats", "List of telegram chats", "application/json", sampleResource)
-	if err != nil {
-		return fmt.Errorf("register chats resource: %w", err)
-	}
-
 	if err := server.Serve(); err != nil {
 		return fmt.Errorf("serve: %w", err)
 	}