feat: add get dialog tool

This commit is contained in:
John Doe 2025-04-05 03:19:22 +03:00
parent 32c3ca5644
commit 8a13b640c5
7 changed files with 353 additions and 46 deletions

View file

@ -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 current user data
- [x] Get the list of dialogs (chats, channels, groups) - [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 - [ ] Mark chanel as read
- [ ] Retrieve messages by date and time - [ ] Retrieve messages by date and time
- [ ] Get the list of contacts - [ ] Get the list of contacts

View file

@ -49,7 +49,7 @@ tasks:
tag: tag:
desc: Create a new tag desc: Create a new tag
cmds: cmds:
- git tag -a v0.1.6 -m "Dialog fix" - git tag -a v0.1.7
- git push origin v0.1.6 - git push origin v0.1.7

View file

@ -4,7 +4,6 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"sort"
"strings" "strings"
"time" "time"
@ -30,34 +29,50 @@ const (
// nolint:lll // nolint:lll
type DialogsArguments struct { 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 { type MessageInfo struct {
Who string `json:"who"` Who string `json:"who,omitempty"`
When string `json:"when"` When string `json:"when"`
Text string `json:"text"` Text string `json:"text,omitempty"`
IsUnread bool `json:"is_unread,omitempty"` IsUnread bool `json:"is_unread,omitempty"`
ts int ts int
} }
type DialogInfo struct { type DialogInfo struct {
ID int64 `json:"id"` Name string `json:"name,omitempty"`
Type string `json:"type"` Type string `json:"type"`
Name string `json:"name"` Title string `json:"title"`
LastMessage *MessageInfo `json:"last_message,omitempty"` 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) // GetDialogs returns a list of dialogs (chats, channels, groups)
func (c *Client) GetDialogs(args DialogsArguments) (*mcp.ToolResponse, error) { 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 var dc tg.MessagesDialogsClass
client := c.T() client := c.T()
if err := client.Run(context.Background(), func(ctx context.Context) (err error) { if err := client.Run(context.Background(), func(ctx context.Context) (err error) {
api := client.API() api := client.API()
dc, err = api.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{ dc, err = api.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{
OffsetPeer: &tg.InputPeerEmpty{}, OffsetPeer: offset.Peer,
OffsetID: offset.MsgID,
OffsetDate: offset.Date,
}) })
if err != nil { if err != nil {
return fmt.Errorf("failed to get dialogs: %w", err) 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") 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 { jsonData, err := json.Marshal(rsp)
return info[i].LastMessage.ts > result[j].LastMessage.ts
})
jsonData, err := json.Marshal(info)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "failed to marshal response") return nil, errors.Wrap(err, "failed to marshal response")
} }
@ -166,7 +180,7 @@ func (d *dialogs) Info() []DialogInfo {
continue continue
} }
if info.Name == "" { if info.Title == "" {
continue continue
} }
@ -176,6 +190,31 @@ func (d *dialogs) Info() []DialogInfo {
return ds 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) { func (d *dialogs) processDialog(rawD tg.DialogClass) (DialogInfo, error) {
dialogItem, ok := rawD.(*tg.Dialog) dialogItem, ok := rawD.(*tg.Dialog)
if !ok { if !ok {
@ -216,48 +255,51 @@ func (d *dialogs) processDialog(rawD tg.DialogClass) (DialogInfo, error) {
} }
var err 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 { if err != nil {
return DialogInfo{}, err return DialogInfo{}, err
} }
info.Type = string(d.getType(dialogItem)) info.Type = string(d.getType(dialogItem))
if info.LastMessage == nil {
info.Empty = true
}
return info, nil return info, nil
} }
func (d *dialogs) getNameID(pC tg.PeerClass) (string, int64, error) { func (d *dialogs) getNameID(pC tg.PeerClass) (string, string, error) {
var name string var name, username string
var id int64
switch p := pC.(type) { switch p := pC.(type) {
case *tg.PeerUser: case *tg.PeerUser:
id = p.GetUserID() u, ok := d.users[p.GetUserID()]
u, ok := d.users[id]
if !ok { 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: case *tg.PeerChannel:
id = p.GetChannelID() channel, ok := d.channels[p.GetChannelID()]
channel, ok := d.channels[id]
if !ok { 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: case *tg.PeerChat:
id = p.GetChatID() chat, ok := d.chats[p.GetChatID()]
chat, ok := d.chats[id]
if !ok { 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: 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 { func (d *dialogs) getType(rawD *tg.Dialog) DialogType {

View file

@ -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
}

View file

@ -2,12 +2,13 @@ package tg
import ( import (
"encoding/json" "encoding/json"
"fmt"
"github.com/gotd/td/tg" "github.com/gotd/td/tg"
"github.com/tidwall/gjson" "github.com/tidwall/gjson"
) )
func getName(source any) string { func getTitle(source any) string {
var name string var name string
switch u := source.(type) { switch u := source.(type) {
case *tg.User: case *tg.User:
@ -16,22 +17,29 @@ func getName(source any) string {
name += " " + u.LastName name += " " + u.LastName
} }
if username, ok := u.GetUsername(); ok && username != "" {
name += " @" + username + ""
}
case *tg.Chat: case *tg.Chat:
name = u.Title name = u.Title
case *tg.Channel: case *tg.Channel:
name = u.Title name = u.Title
if username, ok := u.GetUsername(); ok && username != "" {
name += " @" + username + ""
}
} }
return name 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 // cleanJSON removes empty/default fields from JSON
func cleanJSON(data []byte) []byte { func cleanJSON(data []byte) []byte {
result := gjson.ParseBytes(data) result := gjson.ParseBytes(data)

146
internal/tg/history.go Normal file
View file

@ -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
}

View file

@ -41,13 +41,20 @@ func serve(ctx context.Context, cmd *cli.Command) error {
log.Info().RawJSON("answer", data).Msg("Check GetMe: OK") 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 { if err != nil {
return fmt.Errorf("get dialogs: %w", err) return fmt.Errorf("get dialogs: %w", err)
} }
log.Info().RawJSON("answer", []byte(answer.Content[0].TextContent.Text)).Msg("Check GetDialogs: OK") 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 return nil
} }
@ -56,7 +63,12 @@ func serve(ctx context.Context, cmd *cli.Command) error {
return fmt.Errorf("register tool: %w", err) 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 { if err != nil {
return fmt.Errorf("register dialogs tool: %w", err) return fmt.Errorf("register dialogs tool: %w", err)
} }