refactor: chats method shebangs; drop chats index (D14). (#11)
Tests / Test (push) Skipped
Tests / Release (semver) (push) Skipped
Tests / Test (push) Skipped
Tests / Release (semver) (push) Skipped
Parsers and commands live in internal/chats. bin/chats/{sync,import,facts,apply}.go
are tagged shebang mains. Brain ingest is not a chats command.
This commit is contained in:
committed by
GitHub
GitHub
co-authored by
GitHub
GitHub
parent
46310f8773
commit
cec0161ff6
@@ -0,0 +1,424 @@
|
||||
package chats
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type MCPClient struct {
|
||||
cmd *exec.Cmd
|
||||
stdin *bufio.Writer
|
||||
stdout *bufio.Scanner
|
||||
msgID int
|
||||
}
|
||||
|
||||
type mcpRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID int `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params interface{} `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
type mcpResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID int `json:"id"`
|
||||
Result json.RawMessage `json:"result,omitempty"`
|
||||
Error *struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type mcpToolResult struct {
|
||||
Content []struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
IsError bool `json:"isError,omitempty"`
|
||||
}
|
||||
|
||||
type ListChatsResult struct {
|
||||
ChatID int64 `json:"chat_id"`
|
||||
Title string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
Username string `json:"username,omitempty"`
|
||||
}
|
||||
|
||||
type listChatsEnvelope struct {
|
||||
Results []ListChatsResult `json:"results"`
|
||||
}
|
||||
|
||||
type historyEnvelope struct {
|
||||
Results []GetHistoryResult `json:"results"`
|
||||
}
|
||||
|
||||
type GetHistoryResult struct {
|
||||
ID int `json:"id"`
|
||||
Sender string `json:"sender"`
|
||||
Date string `json:"date"`
|
||||
Text string `json:"text"`
|
||||
Media string `json:"media,omitempty"`
|
||||
Out bool `json:"out,omitempty"`
|
||||
}
|
||||
|
||||
func NewMCPClient(ctx context.Context, apiID int, apiHash, phone, sessionString, mcpDir string) (*MCPClient, error) {
|
||||
env := os.Environ()
|
||||
env = append(env,
|
||||
fmt.Sprintf("TELEGRAM_API_ID=%d", apiID),
|
||||
fmt.Sprintf("TELEGRAM_API_HASH=%s", apiHash),
|
||||
fmt.Sprintf("TELEGRAM_PHONE=%s", phone),
|
||||
fmt.Sprintf("TELEGRAM_SESSION_STRING=%s", sessionString),
|
||||
"MCP_TRANSPORT=stdio",
|
||||
)
|
||||
|
||||
serverPath := filepath.Join(mcpDir, ".venv", "bin", "python3")
|
||||
mainPath := filepath.Join(mcpDir, "main.py")
|
||||
|
||||
cmd := exec.CommandContext(ctx, serverPath, mainPath)
|
||||
cmd.Env = env
|
||||
cmd.Dir = mcpDir
|
||||
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stdin pipe: %w", err)
|
||||
}
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stdout pipe: %w", err)
|
||||
}
|
||||
cmd.Stderr = os.Stderr
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, fmt.Errorf("start mcp: %w", err)
|
||||
}
|
||||
|
||||
c := &MCPClient{
|
||||
cmd: cmd,
|
||||
stdin: bufio.NewWriter(stdin),
|
||||
stdout: bufio.NewScanner(stdout),
|
||||
msgID: 0,
|
||||
}
|
||||
c.stdout.Buffer(make([]byte, 1<<20), 1<<20)
|
||||
|
||||
if err := c.initialize(ctx); err != nil {
|
||||
c.Close()
|
||||
return nil, fmt.Errorf("initialize: %w", err)
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) nextID() int {
|
||||
c.msgID++
|
||||
return c.msgID
|
||||
}
|
||||
|
||||
func (c *MCPClient) sendRequest(ctx context.Context, method string, params interface{}) (json.RawMessage, error) {
|
||||
id := c.nextID()
|
||||
req := mcpRequest{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Method: method,
|
||||
Params: params,
|
||||
}
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal: %w", err)
|
||||
}
|
||||
|
||||
if _, err := c.stdin.Write(body); err != nil {
|
||||
return nil, fmt.Errorf("write: %w", err)
|
||||
}
|
||||
if err := c.stdin.WriteByte('\n'); err != nil {
|
||||
return nil, fmt.Errorf("write newline: %w", err)
|
||||
}
|
||||
if err := c.stdin.Flush(); err != nil {
|
||||
return nil, fmt.Errorf("flush: %w", err)
|
||||
}
|
||||
|
||||
for c.stdout.Scan() {
|
||||
line := c.stdout.Text()
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
var resp mcpResponse
|
||||
if err := json.Unmarshal([]byte(line), &resp); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal response: %w\nline: %s", err, line[:min(len(line), 500)])
|
||||
}
|
||||
if resp.Error != nil {
|
||||
return nil, fmt.Errorf("rpc error %d: %s", resp.Error.Code, resp.Error.Message)
|
||||
}
|
||||
return resp.Result, nil
|
||||
}
|
||||
return nil, fmt.Errorf("no response: %w", c.stdout.Err())
|
||||
}
|
||||
|
||||
func (c *MCPClient) initialize(ctx context.Context) error {
|
||||
params := map[string]interface{}{
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": map[string]interface{}{},
|
||||
"clientInfo": map[string]string{
|
||||
"name": "chats-sync",
|
||||
"version": "0.1.0",
|
||||
},
|
||||
}
|
||||
_, err := c.sendRequest(ctx, "initialize", params)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *MCPClient) ListChats(ctx context.Context, chatType string, limit int) ([]ListChatsResult, error) {
|
||||
args := map[string]interface{}{
|
||||
"chat_type": chatType,
|
||||
"limit": limit,
|
||||
}
|
||||
result, err := c.sendRequest(ctx, "tools/call", map[string]interface{}{
|
||||
"name": "list_chats",
|
||||
"arguments": args,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var toolRes mcpToolResult
|
||||
if err := json.Unmarshal(result, &toolRes); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal tool result: %w", err)
|
||||
}
|
||||
if toolRes.IsError {
|
||||
msg := "unknown"
|
||||
if len(toolRes.Content) > 0 {
|
||||
msg = toolRes.Content[0].Text
|
||||
}
|
||||
return nil, fmt.Errorf("list_chats error: %s", msg)
|
||||
}
|
||||
if len(toolRes.Content) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
text := toolRes.Content[0].Text
|
||||
if text == "" || text == "No chats found matching the criteria." {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var env listChatsEnvelope
|
||||
if err := json.Unmarshal([]byte(text), &env); err != nil {
|
||||
var arr []ListChatsResult
|
||||
if err2 := json.Unmarshal([]byte(text), &arr); err2 != nil {
|
||||
return nil, fmt.Errorf("parse chats: %w (also tried array: %v)\nbody: %s", err, err2, text[:min(len(text), 500)])
|
||||
}
|
||||
return arr, nil
|
||||
}
|
||||
return env.Results, nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) GetHistory(ctx context.Context, chatID int64, limit int) ([]GetHistoryResult, error) {
|
||||
args := map[string]interface{}{
|
||||
"chat_id": chatID,
|
||||
"limit": limit,
|
||||
}
|
||||
result, err := c.sendRequest(ctx, "tools/call", map[string]interface{}{
|
||||
"name": "get_history",
|
||||
"arguments": args,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var toolRes mcpToolResult
|
||||
if err := json.Unmarshal(result, &toolRes); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal tool result: %w", err)
|
||||
}
|
||||
if toolRes.IsError {
|
||||
msg := "unknown"
|
||||
if len(toolRes.Content) > 0 {
|
||||
msg = toolRes.Content[0].Text
|
||||
}
|
||||
return nil, fmt.Errorf("get_history error: %s", msg)
|
||||
}
|
||||
if len(toolRes.Content) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
text := toolRes.Content[0].Text
|
||||
if text == "" || text == "No messages found for this page." || text == "No messages found matching the criteria." {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var env historyEnvelope
|
||||
if err := json.Unmarshal([]byte(text), &env); err != nil {
|
||||
var arr []GetHistoryResult
|
||||
if err2 := json.Unmarshal([]byte(text), &arr); err2 != nil {
|
||||
return nil, fmt.Errorf("parse history: %w (also tried array: %v)\nbody: %s", err, err2, text[:min(len(text), 500)])
|
||||
}
|
||||
return arr, nil
|
||||
}
|
||||
return env.Results, nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) Close() error {
|
||||
if c.stdin != nil {
|
||||
c.stdin.Flush()
|
||||
}
|
||||
if c.cmd != nil && c.cmd.Process != nil {
|
||||
c.cmd.Process.Kill()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type TelegramMCPSource struct {
|
||||
mcpDir string
|
||||
apiID int
|
||||
apiHash string
|
||||
phone string
|
||||
sessionStr string
|
||||
limit int
|
||||
}
|
||||
|
||||
func NewTelegramMCPSource(apiID int, apiHash, phone, sessionString, mcpDir string) *TelegramMCPSource {
|
||||
return &TelegramMCPSource{
|
||||
mcpDir: mcpDir,
|
||||
apiID: apiID,
|
||||
apiHash: apiHash,
|
||||
phone: phone,
|
||||
sessionStr: sessionString,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TelegramMCPSource) Name() string { return "telegram" }
|
||||
|
||||
func (s *TelegramMCPSource) Sync(ctx context.Context, outDir string, limit int) error {
|
||||
if limit > 0 {
|
||||
s.limit = limit
|
||||
}
|
||||
|
||||
client, err := NewMCPClient(ctx, s.apiID, s.apiHash, s.phone, s.sessionStr, s.mcpDir)
|
||||
if err != nil {
|
||||
return fmt.Errorf("mcp client: %w", err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
chats, err := client.ListChats(ctx, "user", 100)
|
||||
if err != nil {
|
||||
return fmt.Errorf("list chats: %w", err)
|
||||
}
|
||||
if len(chats) == 0 {
|
||||
fmt.Println("chats: no personal chats found")
|
||||
return nil
|
||||
}
|
||||
fmt.Printf("chats: found %d personal chats\n", len(chats))
|
||||
|
||||
var filtered []ListChatsResult
|
||||
for _, c := range chats {
|
||||
if strings.Contains(strings.ToLower(c.Username), "bot") {
|
||||
continue
|
||||
}
|
||||
if c.ChatID == 777000 { // Telegram service
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, c)
|
||||
}
|
||||
fmt.Printf("chats: %d after filter (bots excluded)\n", len(filtered))
|
||||
|
||||
for _, chat := range filtered {
|
||||
chatID := fmt.Sprintf("user_%d", chat.ChatID)
|
||||
chatName := chat.Title
|
||||
if chatName == "" {
|
||||
chatName = chatID
|
||||
}
|
||||
|
||||
chatDir := filepath.Join(outDir, "telegram", chatID)
|
||||
if err := os.MkdirAll(chatDir, 0755); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "chats: mkdir %s: %v\n", chatDir, err)
|
||||
continue
|
||||
}
|
||||
|
||||
jsonlPath := filepath.Join(chatDir, "messages.jsonl")
|
||||
f, err := os.Create(jsonlPath)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "chats: create %s: %v\n", jsonlPath, err)
|
||||
continue
|
||||
}
|
||||
|
||||
msgLimit := 100
|
||||
if s.limit > 0 {
|
||||
msgLimit = s.limit
|
||||
}
|
||||
|
||||
msgs, err := client.GetHistory(ctx, chat.ChatID, msgLimit)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "chats: get_history for %s: %v\n", chatName, err)
|
||||
f.Close()
|
||||
continue
|
||||
}
|
||||
|
||||
enc := json.NewEncoder(f)
|
||||
written := 0
|
||||
for _, m := range msgs {
|
||||
if m.Out {
|
||||
continue
|
||||
}
|
||||
text := m.Text
|
||||
if text == "" && m.Media != "" {
|
||||
text = fmt.Sprintf("[%s]", m.Media)
|
||||
}
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
sender := cleanSender(m.Sender)
|
||||
ts := m.Date
|
||||
if t, err := time.Parse(time.RFC3339, m.Date); err == nil {
|
||||
ts = t.UTC().Format(time.RFC3339)
|
||||
}
|
||||
|
||||
chatMsg := Message{
|
||||
ID: fmt.Sprintf("tg_%d_%d", chat.ChatID, m.ID),
|
||||
Timestamp: ts,
|
||||
From: sender,
|
||||
Text: text,
|
||||
Platform: "telegram",
|
||||
}
|
||||
if m.Media != "" {
|
||||
desc := fmt.Sprintf("[%s]", m.Media)
|
||||
chatMsg.Media = &desc
|
||||
}
|
||||
if err := enc.Encode(chatMsg); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "chats: encode msg: %v\n", err)
|
||||
continue
|
||||
}
|
||||
written++
|
||||
}
|
||||
f.Close()
|
||||
|
||||
if written > 0 {
|
||||
fmt.Printf("chats: synced %s (%s) — %d messages\n", chatName, chatID, written)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func cleanSender(sender string) string {
|
||||
if idx := strings.Index(sender, " ("); idx > 0 {
|
||||
sender = sender[:idx]
|
||||
} else if idx := strings.Index(sender, " @"); idx > 0 {
|
||||
sender = sender[:idx]
|
||||
}
|
||||
if idx := strings.Index(sender, " ["); idx > 0 {
|
||||
sender = sender[:idx]
|
||||
}
|
||||
return sender
|
||||
}
|
||||
|
||||
func min(a, b int) int {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
Reference in New Issue
Block a user