from collections.abc import Sequence from datetime import datetime from typing import Any import asyncpg from fastmcp import FastMCP from fastmcp.exceptions import ToolError from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext from fastmcp.tools.base import ToolResult from mcp.types import CallToolRequestParams, TextContent from pydantic import BaseModel from api.mcp.download import ( MAX_WAIT_SECONDS, describe, fetch_media, inline_content, share_media, ) from api.mcp.format import build_transcript, load_notes, load_photos, resolve_names from dependencies.container import container from utils.files import resolve_mime from utils.read import annotations, chats, media, peers, presence, social, watches from utils.read.accounts import self_user_id, unknown_account_hint from utils.read.models import DEFAULT_LIMIT, Page from utils.search.models import SearchFilters from utils.search.repository import search_messages from utils.storage import ContentAddressedStorage INSTRUCTIONS = """\ beavergram archives Telegram data (chats, messages, media, presence, stories, peer history) into Postgres and exposes it read-only over these tools. Account scoping: - Every tool takes `account_id`. It selects which archived Telegram account to read. Unless the user names a different one, always pass `account_id=1`. - `account_id` is beavergram's own small serial (1, 2, ...), never the archived account's Telegram user id. Passing a Telegram id returns an error naming the right `account_id`. Identifiers: - `chat_id` and `peer_id` are Telegram IDs (negative for groups/channels, positive for users/bots). Discover them with `list_chats` before calling tools that need a specific chat. Typical flows: - Read a chat: `list_chats` -> pick a `chat_id` -> `get_chat_history` for a human-readable transcript (names, timestamps, voice STT, inline photos), or `get_chat_history_raw` for structured JSON. Both return messages oldest first; default gives the latest `limit`. Page to older messages with `before_id` (first id you saw), to newer with `after_id` (last id you saw); begin a full forward walk with `after_id=0`. Each transcript line starts with its `#message_id`, and photos show as `[photo #message_id]` — pass that id to `download_media` or to annotation tools. Lines marked "📝 [your private note, NOT in Telegram]" are the user's own annotations attached locally via the web UI; they never existed in Telegram. Treat them as private notes from the user to you, not as chat content. - Find something: `search_messages_tool` (full-text over message text and STT transcripts; supports chat/sender/date filters and regex). - Forensics: `get_deleted_messages` and `get_message_versions` recover content removed or edited in Telegram but kept in the archive. - Files: `download_media` gets any message attachment (photo, video, voice, document, ...). It pulls the file from Telegram if the archive lacks it and returns a public link valid 24h that needs no auth (`curl -O `), plus the content itself for small images and text files. `get_media` returns metadata only. - Monitoring: `set_watch` creates a local rule, `list_watches` lists rules, and `list_alerts` reads what those rules fired. Watched bots: - a bot whose token the user added is just another archived account: it has its own `account_id` and its chats, messages and media read through the same tools as a user account. Nothing bot-specific to call. Writes: everything is read-only except `set_watch` and the share links `download_media` creates (visible and revocable in the web UI). Paging: list/search tools accept `limit` and `offset`. """ mcp: FastMCP = FastMCP("beavergram", instructions=INSTRUCTIONS) async def _pool() -> asyncpg.Pool: return await container.get(asyncpg.Pool) class AccountGuard(Middleware): """Reject tool calls scoped to an account_id that does not exist.""" async def on_call_tool( self, context: MiddlewareContext[CallToolRequestParams], call_next: CallNext[CallToolRequestParams, ToolResult], ) -> ToolResult: account_id = (context.message.arguments or {}).get("account_id") if isinstance(account_id, int) and not isinstance(account_id, bool): hint = await unknown_account_hint(await _pool(), account_id) if hint is not None: raise ToolError(hint) return await call_next(context) mcp.add_middleware(AccountGuard()) def _dump(items: Sequence[BaseModel]) -> list[dict[str, Any]]: return [item.model_dump(mode="json") for item in items] @mcp.tool async def search_messages_tool( account_id: int, query: str | None = None, chat_id: int | None = None, sender_id: int | None = None, has_media: bool | None = None, date_from: datetime | None = None, date_to: datetime | None = None, regex: str | None = None, limit: int = DEFAULT_LIMIT, offset: int = 0, ) -> list[dict[str, Any]]: """Full-text search over message text and STT transcripts.""" filters = SearchFilters( account_id=account_id, query=query, chat_id=chat_id, sender_id=sender_id, has_media=has_media, date_from=date_from, date_to=date_to, regex=regex, limit=limit, offset=offset, ) return _dump(await search_messages(await _pool(), filters)) @mcp.tool async def list_chats( account_id: int, limit: int = DEFAULT_LIMIT, offset: int = 0 ) -> list[dict[str, Any]]: """List archived chats with message counts and last activity.""" page = Page(limit=limit, offset=offset) return _dump(await chats.list_chats(await _pool(), account_id, page)) @mcp.tool(output_schema=None) async def get_chat_history( account_id: int, chat_id: int, limit: int = DEFAULT_LIMIT, before_id: int | None = None, after_id: int | None = None, include_deleted: bool = True, include_images: bool = True, max_images: int = 20, ) -> list[Any]: """Read a chat as a readable transcript, oldest first. Renders "Name (time): text", voice notes as "(Voice message, STT Content: ...)", and inlines downloaded photos as images. Default returns the latest `limit` messages; page to older messages with `before_id` (the first id you saw) or to newer ones with `after_id` (the last id you saw). """ pool = await _pool() views = await chats.get_chat_history( pool, account_id, chat_id, Page(limit=limit), include_deleted=include_deleted, before_id=before_id, after_id=after_id, ) if after_id is None: views = list(reversed(views)) names = await resolve_names(pool, account_id, views) notes = await load_notes(pool, account_id, chat_id, views) self_id = await self_user_id(pool, account_id) photos: list[tuple[int, bytes, str]] = [] truncated = False if include_images: storage = await container.get(ContentAddressedStorage) photos, truncated = await load_photos( pool, storage, account_id, views, limit=max_images ) return build_transcript( views, names, self_id, photos, notes=notes, truncated=truncated ) @mcp.tool async def get_chat_history_raw( account_id: int, chat_id: int, limit: int = DEFAULT_LIMIT, before_id: int | None = None, after_id: int | None = None, include_deleted: bool = True, ) -> list[dict[str, Any]]: """Structured chat messages as JSON, oldest first. Default returns the latest `limit` messages. Walk the whole chat forward with `after_id` set to the last returned message_id (begin at `after_id=0`), or backward with `before_id` set to the first returned message_id. """ views = await chats.get_chat_history( await _pool(), account_id, chat_id, Page(limit=limit), include_deleted=include_deleted, before_id=before_id, after_id=after_id, ) if after_id is None: views = list(reversed(views)) return _dump(views) @mcp.tool async def get_deleted_messages( account_id: int, chat_id: int | None = None, limit: int = DEFAULT_LIMIT, offset: int = 0, ) -> list[dict[str, Any]]: """List messages that were deleted in Telegram but kept in the archive.""" return _dump( await chats.get_deleted_messages( await _pool(), account_id, Page(limit=limit, offset=offset), chat_id=chat_id ) ) @mcp.tool async def get_message_versions( account_id: int, chat_id: int, message_id: int ) -> list[dict[str, Any]]: """Get the edit history of a message.""" return _dump( await chats.get_message_versions(await _pool(), account_id, chat_id, message_id) ) @mcp.tool async def get_media( account_id: int, chat_id: int, message_id: int ) -> dict[str, Any] | None: """Get media metadata for a message; use download_media to get the file.""" item = await media.get_message_media(await _pool(), account_id, chat_id, message_id) return item.model_dump(mode="json") if item else None @mcp.tool(output_schema=None) async def download_media( account_id: int, chat_id: int, message_id: int, wait_seconds: int = 45 ) -> list[Any]: """Download the file attached to a message: any photo, video, voice, document. Pulls it from Telegram first if the archive doesn't have it yet, waiting up to `wait_seconds`. Returns a public download link (valid 24h, no auth) and inlines small images and text files. """ pool = await _pool() wait = min(max(wait_seconds, 0), MAX_WAIT_SECONDS) item = await fetch_media(pool, account_id, chat_id, message_id, wait) if item is None or item.storage_key is None: return [ TextContent( type="text", text="Still downloading from Telegram; call download_media again.", ) ] mime = resolve_mime(item.kind, item.mime, item.file_name) share = await share_media(pool, item, mime) storage = await container.get(ContentAddressedStorage) return [ TextContent(type="text", text=describe(item, mime, share)), *inline_content(storage, item.storage_key, mime), ] @mcp.tool async def get_callbacks( account_id: int, chat_id: int, message_id: int ) -> list[dict[str, Any]]: """Get bot inline-button callback data (hex) for a message.""" items = await social.get_callbacks(await _pool(), account_id, chat_id, message_id) return _dump(items) @mcp.tool async def presence_history( account_id: int, peer_id: int, date_from: datetime | None = None, date_to: datetime | None = None, limit: int = DEFAULT_LIMIT, offset: int = 0, ) -> list[dict[str, Any]]: """Get online/offline status history of a peer.""" return _dump( await presence.presence_history( await _pool(), account_id, peer_id, Page(limit=limit, offset=offset), date_from=date_from, date_to=date_to, ) ) @mcp.tool async def get_peer_history(account_id: int, peer_id: int) -> list[dict[str, Any]]: """Get name/username/avatar change history of a contact.""" return _dump(await peers.get_peer_history(await _pool(), account_id, peer_id)) @mcp.tool async def get_stories( account_id: int, peer_id: int | None = None, limit: int = DEFAULT_LIMIT, offset: int = 0, ) -> list[dict[str, Any]]: """List archived stories of contacts.""" return _dump( await peers.get_stories( await _pool(), account_id, Page(limit=limit, offset=offset), peer_id=peer_id ) ) @mcp.tool async def get_annotations( account_id: int, chat_id: int | None = None, message_id: int | None = None, limit: int = DEFAULT_LIMIT, offset: int = 0, ) -> list[dict[str, Any]]: """Read user annotations on messages (read-only via MCP).""" return _dump( await annotations.list_annotations( await _pool(), account_id, Page(limit=limit, offset=offset), chat_id=chat_id, message_id=message_id, ) ) @mcp.tool async def set_watch( account_id: int, kind: str, params: dict[str, Any] | None = None, enabled: bool = True, ) -> dict[str, Any]: """Create a local watch rule (the only local write MCP is allowed).""" watch = await watches.create_watch( await _pool(), account_id, kind, params or {}, enabled=enabled ) return watch.model_dump(mode="json") @mcp.tool async def list_watches(account_id: int) -> list[dict[str, Any]]: """List local watch rules.""" return _dump(await watches.list_watches(await _pool(), account_id)) @mcp.tool async def list_alerts( account_id: int, seen: bool | None = None, limit: int = DEFAULT_LIMIT, offset: int = 0, ) -> list[dict[str, Any]]: """List fired alerts from watch rules.""" return _dump( await watches.list_alerts( await _pool(), account_id, Page(limit=limit, offset=offset), seen=seen ) )