diff --git a/README.md b/README.md index 86e47cb..71d4fa1 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,7 @@ The Model Context Protocol (MCP) is a system that lets AI apps, like Claude Desk - [x] Get current user data - [x] Get the list of dialogs (chats, channels, groups) -- [ ] Get the list of (unread) messages in the given dialog +- [x] Get the list of (unread) messages in the given dialog - [ ] Mark chanel as read - [ ] Retrieve messages by date and time - [ ] Get the list of contacts diff --git a/Taskfile.yml b/Taskfile.yml index 0330c70..7b79d85 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -49,7 +49,7 @@ tasks: tag: desc: Create a new tag cmds: - - git tag -a v0.1.6 -m "Dialog fix" - - git push origin v0.1.6 + - git tag -a v0.1.7 + - git push origin v0.1.7 diff --git a/internal/tg/dialogs.go b/internal/tg/dialogs.go index 77e8548..e7d9943 100644 --- a/internal/tg/dialogs.go +++ b/internal/tg/dialogs.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" "fmt" - "sort" "strings" "time" @@ -30,34 +29,50 @@ const ( // nolint:lll type DialogsArguments struct { - //WithLastMessages bool `json:"with_last_messages,omitempty" jsonschema:"description=Include last messages in response"` + Offset string `json:"offset,omitempty" jsonschema:"description=Offset for continuation"` } type MessageInfo struct { - Who string `json:"who"` + Who string `json:"who,omitempty"` When string `json:"when"` - Text string `json:"text"` + Text string `json:"text,omitempty"` IsUnread bool `json:"is_unread,omitempty"` ts int } type DialogInfo struct { - ID int64 `json:"id"` + Name string `json:"name,omitempty"` Type string `json:"type"` - Name string `json:"name"` + Title string `json:"title"` LastMessage *MessageInfo `json:"last_message,omitempty"` + Empty bool `json:"empty,omitempty"` +} + +type DialogsResponse struct { + Dialogs []DialogInfo `json:"dialogs"` + Offset DialogsOffset `json:"offset"` } // GetDialogs returns a list of dialogs (chats, channels, groups) func (c *Client) GetDialogs(args DialogsArguments) (*mcp.ToolResponse, error) { - var result []DialogInfo + var offset DialogsOffset + if args.Offset != "" { + if err := offset.UnmarshalJSON([]byte(args.Offset)); err != nil { + return nil, errors.Wrap(err, "failed to unmarshal offset") + } + } + if offset.Peer == nil { + offset.Peer = &tg.InputPeerEmpty{} + } var dc tg.MessagesDialogsClass client := c.T() if err := client.Run(context.Background(), func(ctx context.Context) (err error) { api := client.API() dc, err = api.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{ - OffsetPeer: &tg.InputPeerEmpty{}, + OffsetPeer: offset.Peer, + OffsetID: offset.MsgID, + OffsetDate: offset.Date, }) if err != nil { return fmt.Errorf("failed to get dialogs: %w", err) @@ -77,13 +92,12 @@ func (c *Client) GetDialogs(args DialogsArguments) (*mcp.ToolResponse, error) { return nil, errors.Wrap(err, "failed to get dialogs") } - info := d.Info() + rsp := DialogsResponse{ + Dialogs: d.Info(), + Offset: d.Offset(), + } - sort.Slice(result, func(i, j int) bool { - return info[i].LastMessage.ts > result[j].LastMessage.ts - }) - - jsonData, err := json.Marshal(info) + jsonData, err := json.Marshal(rsp) if err != nil { return nil, errors.Wrap(err, "failed to marshal response") } @@ -166,7 +180,7 @@ func (d *dialogs) Info() []DialogInfo { continue } - if info.Name == "" { + if info.Title == "" { continue } @@ -176,6 +190,31 @@ func (d *dialogs) Info() []DialogInfo { return ds } +func (d *dialogs) Offset() DialogsOffset { + for i := len(d.Dialogs) - 1; i >= 0; i-- { + dialogItem, ok := d.Dialogs[i].(*tg.Dialog) + if !ok { + continue + } + if dialogItem.Peer == nil { + continue + } + + msg, ok := d.messages[getPeerID(dialogItem.Peer)] + if !ok { + continue + } + + return DialogsOffset{ + MsgID: msg.ID, + Date: msg.Date, + Peer: getInputPeerID(dialogItem.Peer), + } + } + + return DialogsOffset{} +} + func (d *dialogs) processDialog(rawD tg.DialogClass) (DialogInfo, error) { dialogItem, ok := rawD.(*tg.Dialog) if !ok { @@ -216,48 +255,51 @@ func (d *dialogs) processDialog(rawD tg.DialogClass) (DialogInfo, error) { } var err error - info.Name, info.ID, err = d.getNameID(dialogItem.Peer) + info.Title, info.Name, err = d.getNameID(dialogItem.Peer) if err != nil { return DialogInfo{}, err } info.Type = string(d.getType(dialogItem)) + if info.LastMessage == nil { + info.Empty = true + } + return info, nil } -func (d *dialogs) getNameID(pC tg.PeerClass) (string, int64, error) { - var name string - var id int64 +func (d *dialogs) getNameID(pC tg.PeerClass) (string, string, error) { + var name, username string switch p := pC.(type) { case *tg.PeerUser: - id = p.GetUserID() - u, ok := d.users[id] + u, ok := d.users[p.GetUserID()] if !ok { - return "", 0, errors.Errorf("peerid(%d): invalid message user", id) + return "", "", errors.Errorf("peerid(%d): invalid message user", p.GetUserID()) } - name = getName(u) + name = getTitle(u) + username = getUsername(u) case *tg.PeerChannel: - id = p.GetChannelID() - channel, ok := d.channels[id] + channel, ok := d.channels[p.GetChannelID()] if !ok { - return "", 0, errors.Errorf("peerid(%d): invalid message channel", id) + return "", "", errors.Errorf("peerid(%d): invalid message channel", p.GetChannelID()) } - name = getName(channel) + name = getTitle(channel) + username = getUsername(channel) case *tg.PeerChat: - id = p.GetChatID() - chat, ok := d.chats[id] + chat, ok := d.chats[p.GetChatID()] if !ok { - return "", 0, errors.Errorf("peerid(%d): invalid message chat", id) + return "", "", errors.Errorf("peerid(%d): invalid message chat", p.GetChatID()) } - name = getName(chat) + name = getTitle(chat) + username = getUsername(chat) default: - return "", 0, fmt.Errorf("chose author(%T): invalid dialog peer", p) + return "", "", fmt.Errorf("chose author(%T): invalid dialog peer", p) } - return name, id, nil + return name, username, nil } func (d *dialogs) getType(rawD *tg.Dialog) DialogType { diff --git a/internal/tg/dialogs_offset.go b/internal/tg/dialogs_offset.go new file mode 100644 index 0000000..568c77a --- /dev/null +++ b/internal/tg/dialogs_offset.go @@ -0,0 +1,99 @@ +package tg + +import ( + "encoding/json" + "fmt" + "strconv" + "strings" + + "github.com/gotd/td/tg" +) + +type DialogsOffset struct { + MsgID int `json:"msg_id"` + Date int `json:"offset_date"` + Peer tg.InputPeerClass +} + +func getInputPeerID(p tg.PeerClass) tg.InputPeerClass { + switch v := p.(type) { + case *tg.PeerUser: + return &tg.InputPeerUser{UserID: v.UserID} + case *tg.PeerChannel: + return &tg.InputPeerChannel{ChannelID: v.ChannelID} + case *tg.PeerChat: + return &tg.InputPeerChat{ChatID: v.ChatID} + default: + return &tg.InputPeerEmpty{} + } +} + +func (o DialogsOffset) MarshalJSON() ([]byte, error) { + return json.Marshal(o.String()) +} + +func (o *DialogsOffset) String() string { + + var id int64 + var peerType string + switch p := o.Peer.(type) { + case *tg.InputPeerUser: + peerType = "user" + id = p.UserID + case *tg.InputPeerChannel: + peerType = "chan" + id = p.ChannelID + case *tg.InputPeerChat: + peerType = "chat" + id = p.ChatID + default: + peerType = "unknown" + } + + if id == 0 { + return "end" + } + + return fmt.Sprintf("%s-%d-%d-%d", peerType, id, o.MsgID, o.Date) +} + +func (o *DialogsOffset) UnmarshalJSON(data []byte) error { + parts := strings.Split(string(data), "-") + if len(parts) != 4 { + return fmt.Errorf("invalid format") + } + + var err error + switch parts[0] { + case "user": + var userID int64 + userID, err = strconv.ParseInt(parts[1], 10, 64) + o.Peer = &tg.InputPeerUser{UserID: userID} + case "chan": + var channelID int64 + channelID, err = strconv.ParseInt(parts[1], 10, 64) + o.Peer = &tg.InputPeerChannel{ChannelID: channelID} + case "chat": + var chatID int64 + chatID, err = strconv.ParseInt(parts[1], 10, 64) + o.Peer = &tg.InputPeerChat{ChatID: chatID} + default: + return fmt.Errorf("unknown peer type") + } + + if err != nil { + return fmt.Errorf("invalid peer: %w", err) + } + + o.MsgID, err = strconv.Atoi(parts[2]) + if err != nil { + return fmt.Errorf("invalid message ID: %w", err) + } + + o.Date, err = strconv.Atoi(parts[3]) + if err != nil { + return fmt.Errorf("invalid date: %w", err) + } + + return nil +} diff --git a/internal/tg/helpers.go b/internal/tg/helpers.go index f3c983a..fc76bfa 100644 --- a/internal/tg/helpers.go +++ b/internal/tg/helpers.go @@ -2,12 +2,13 @@ package tg import ( "encoding/json" + "fmt" "github.com/gotd/td/tg" "github.com/tidwall/gjson" ) -func getName(source any) string { +func getTitle(source any) string { var name string switch u := source.(type) { case *tg.User: @@ -16,22 +17,29 @@ func getName(source any) string { name += " " + u.LastName } - if username, ok := u.GetUsername(); ok && username != "" { - name += " @" + username + "" - } case *tg.Chat: name = u.Title case *tg.Channel: name = u.Title - - if username, ok := u.GetUsername(); ok && username != "" { - name += " @" + username + "" - } } return name } +func getUsername(source any) string { + var username string + switch u := source.(type) { + case *tg.User: + username = u.Username + case *tg.Chat: + username = fmt.Sprintf("chat(%d)", u.ID) + case *tg.Channel: + username = u.Username + } + + return username +} + // cleanJSON removes empty/default fields from JSON func cleanJSON(data []byte) []byte { result := gjson.ParseBytes(data) diff --git a/internal/tg/history.go b/internal/tg/history.go new file mode 100644 index 0000000..3afcaf5 --- /dev/null +++ b/internal/tg/history.go @@ -0,0 +1,146 @@ +package tg + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/gotd/td/telegram/message" + "github.com/gotd/td/tg" + mcp "github.com/metoro-io/mcp-golang" + "github.com/pkg/errors" +) + +type HistoryArguments struct { + Name string `json:"name" jsonschema:"required,description=Name of the dialog"` + Offset int `json:"offset,omitempty" jsonschema:"description=Offset for continuation"` +} + +type HistoryResponse struct { + Messages []MessageInfo `json:"messages"` + Offset int `json:"offset,omitempty"` +} + +func (c *Client) GetHistory(args HistoryArguments) (*mcp.ToolResponse, error) { + var messagesClass tg.MessagesMessagesClass + client := c.T() + if err := client.Run(context.Background(), func(ctx context.Context) (err error) { + api := client.API() + + sender := message.NewSender(api) + inputPeer, err := sender.Resolve(args.Name).AsInputPeer(ctx) + if err != nil { + return fmt.Errorf("failed to resolve name: %w", err) + } + + messagesClass, err = api.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{ + Peer: inputPeer, + OffsetID: args.Offset, + }) + if err != nil { + return fmt.Errorf("failed to get history: %w", err) + } + + //Debug + //jsonData, _ := json.Marshal(messagesClass) + //log.Info().RawJSON("history", cleanJSON(jsonData)).Msg("history") + + return nil + }); err != nil { + return nil, errors.Wrap(err, "failed to get history") + } + + h, err := newHistory(messagesClass) + if err != nil { + return nil, errors.Wrap(err, "failed to process history") + } + + rsp := HistoryResponse{ + Messages: h.Info(), + Offset: h.Offset(), + } + + jsonData, err := json.Marshal(rsp) + if err != nil { + return nil, errors.Wrap(err, "failed to marshal response") + } + + return mcp.NewToolResponse(mcp.NewTextContent(string(jsonData))), nil +} + +type history struct { + tg.MessagesMessages + users map[int64]*tg.User +} + +func newHistory(raw tg.MessagesMessagesClass) (*history, error) { + var h history + switch m := raw.(type) { + case *tg.MessagesMessages: + h = history{MessagesMessages: *m} + case *tg.MessagesMessagesSlice: + h = history{MessagesMessages: tg.MessagesMessages{ + Messages: m.Messages, + Users: m.Users, + Chats: m.Chats, + }} + case *tg.MessagesChannelMessages: + h = history{MessagesMessages: tg.MessagesMessages{ + Messages: m.Messages, + Users: m.Users, + Chats: m.Chats, + }} + default: + return nil, fmt.Errorf("unexpected type: %T", raw) + } + + h.users = make(map[int64]*tg.User) + for _, u := range h.Users { + if user, ok := u.(*tg.User); ok { + h.users[user.ID] = user + } + } + + return &h, nil +} + +func (h *history) Info() []MessageInfo { + messages := make([]MessageInfo, 0, len(h.Messages)) + + for _, msg := range h.Messages { + m, ok := msg.(*tg.Message) + if !ok { + continue + } + + var who string + if m.FromID != nil { + switch from := m.FromID.(type) { + case *tg.PeerUser: + if user, ok := h.users[from.UserID]; ok { + who = getUsername(user) + } + } + } + + messages = append(messages, MessageInfo{ + Who: who, + When: time.Unix(int64(m.Date), 0).Format(time.DateTime), + Text: m.Message, + ts: m.Date, + }) + } + + return messages +} + +func (h *history) Offset() int { + for i := len(h.Messages) - 1; i >= 0; i-- { + if msg, ok := h.Messages[i].(*tg.Message); ok { + return msg.ID + } + } + + return 0 +} diff --git a/serve.go b/serve.go index 8b82252..bf94b10 100644 --- a/serve.go +++ b/serve.go @@ -41,13 +41,20 @@ func serve(ctx context.Context, cmd *cli.Command) error { log.Info().RawJSON("answer", data).Msg("Check GetMe: OK") - answer, err = client.GetDialogs(tg.DialogsArguments{}) + answer, err = client.GetDialogs(tg.DialogsArguments{Offset: ""}) if err != nil { return fmt.Errorf("get dialogs: %w", err) } log.Info().RawJSON("answer", []byte(answer.Content[0].TextContent.Text)).Msg("Check GetDialogs: OK") + answer, err = client.GetHistory(tg.HistoryArguments{Name: "lalal", Offset: 5574}) + if err != nil { + return fmt.Errorf("get histore: %w", err) + } + + log.Info().RawJSON("answer", []byte(answer.Content[0].TextContent.Text)).Msg("Check GetHistory: OK") + return nil } @@ -56,7 +63,12 @@ func serve(ctx context.Context, cmd *cli.Command) error { return fmt.Errorf("register tool: %w", err) } - err = server.RegisterTool("tg_dialogs", "Get list of telegram dialogs (chats, channels, groups)", client.GetDialogs) + err = server.RegisterTool("tg_dialogs", "Get list of telegram dialogs (chats, channels, users)", client.GetDialogs) + if err != nil { + return fmt.Errorf("register dialogs tool: %w", err) + } + + err = server.RegisterTool("tg_dialog", "Get messages of telegram dialog (channel, user)", client.GetHistory) if err != nil { return fmt.Errorf("register dialogs tool: %w", err) }