beavergram/backend/src/api/mcp/format.py

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