from datetime import datetime from typing import Any import asyncpg from fastmcp.utilities.types import Image from mcp.types import TextContent from utils.read import peers from utils.read.models import MediaRef, MessageView from utils.storage import ContentAddressedStorage _VOICE_KINDS = {"voice", "video_note"} def _ts(value: datetime) -> str: return value.strftime("%Y-%m-%d %H:%M") def _name(sender_id: int | None, names: dict[int, str], self_id: int | None) -> str: if sender_id is None: return "Unknown" if sender_id == self_id: return "Me" return names.get(sender_id) or str(sender_id) def _media_note(media: list[MediaRef]) -> list[str]: notes: list[str] = [] for item in media: if item.kind in _VOICE_KINDS: if item.extracted_text: notes.append(f"(Voice message, STT Content: {item.extracted_text})") else: notes.append("(Voice message, not transcribed)") elif item.kind == "photo": state = "" if item.downloaded else ", not downloaded" notes.append(f"[photo #{item.message_id}{state}]") else: notes.append(f"[{item.kind}]") return notes def _line( view: MessageView, names: dict[int, str], self_id: int | None, notes: dict[int, list[str]], ) -> str: parts: list[str] = [] if view.reply and (view.reply.sender_name or view.reply.text): ref = view.reply.sender_name or str(view.reply.sender_id or "?") parts.append(f"(reply to {ref})") if view.text: parts.append(view.text) parts.extend(_media_note(view.media)) suffix = "" if view.edited_at: suffix += " (edited)" if view.deleted_at: suffix += " (deleted)" body = " ".join(part for part in parts if part) or "(no text)" name = _name(view.sender_id, names, self_id) line = f"#{view.message_id} {name} ({_ts(view.date)}): {body}{suffix}" for note in notes.get(view.message_id, []): line += f"\n 📝 [your private note, NOT in Telegram]: {note}" return line async def load_notes( pool: asyncpg.Pool, account_id: int, chat_id: int, views: list[MessageView] ) -> dict[int, list[str]]: ids = [view.message_id for view in views] if not ids: return {} rows = await pool.fetch( "SELECT message_id, text FROM annotations " "WHERE account_id = $1 AND chat_id = $2 AND message_id = ANY($3::bigint[]) " "ORDER BY created_at", account_id, chat_id, ids, ) notes: dict[int, list[str]] = {} for row in rows: notes.setdefault(row["message_id"], []).append(row["text"]) return notes async def resolve_names( pool: asyncpg.Pool, account_id: int, views: list[MessageView] ) -> dict[int, str]: ids = list({view.sender_id for view in views if view.sender_id is not None}) found = await peers.get_peers(pool, account_id, ids) names: dict[int, str] = {} for peer in found: name = ( " ".join(part for part in (peer.first_name, peer.last_name) if part) or peer.username ) if name: names[peer.peer_id] = name return names async def load_photos( pool: asyncpg.Pool, storage: ContentAddressedStorage, account_id: int, views: list[MessageView], *, limit: int, ) -> tuple[list[tuple[int, bytes, str]], bool]: refs = [ (item.message_id, item.id) for view in views for item in view.media if item.kind == "photo" and item.downloaded and item.id is not None ] truncated = len(refs) > limit refs = refs[-limit:] if not refs: return [], truncated rows = await pool.fetch( "SELECT id, storage_key, mime FROM media " "WHERE account_id = $1 AND id = ANY($2::bigint[])", account_id, [media_id for _, media_id in refs], ) by_id = {row["id"]: row for row in rows} out: list[tuple[int, bytes, str]] = [] for message_id, media_id in refs: row = by_id.get(media_id) if row is None or not row["storage_key"]: continue try: data = storage.get(row["storage_key"]) except OSError: continue fmt = (row["mime"] or "image/jpeg").split("/")[-1] out.append((message_id, data, fmt)) return out, truncated def build_transcript( views: list[MessageView], names: dict[int, str], self_id: int | None, photos: list[tuple[int, bytes, str]], *, notes: dict[int, list[str]] | None = None, truncated: bool = False, ) -> list[Any]: if not views: return [TextContent(type="text", text="No messages.")] notes = notes or {} header = f"{len(views)} messages (oldest first)" body = f"{header}\n\n" + "\n".join( _line(view, names, self_id, notes) for view in views ) blocks: list[Any] = [TextContent(type="text", text=body)] if truncated: blocks.append( TextContent( type="text", text="(images truncated to the most recent; narrow the range for more)", ) ) for message_id, data, fmt in photos: blocks.append( TextContent(type="text", text=f"Image attached to message #{message_id}:") ) blocks.append(Image(data=data, format=fmt)) return blocks