diff --git a/backend/migrations/versions/d8f1a6c3e5b7_message_search.py b/backend/migrations/versions/d8f1a6c3e5b7_message_search.py new file mode 100644 index 0000000..cc0d6a7 --- /dev/null +++ b/backend/migrations/versions/d8f1a6c3e5b7_message_search.py @@ -0,0 +1,138 @@ +"""message search index outside the hypertable, versions searchable + +Revision ID: d8f1a6c3e5b7 +Revises: c9e2a7f4b1d8 +Create Date: 2026-10-19 03:00:00.000000 + +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "d8f1a6c3e5b7" +down_revision: str | None = "c9e2a7f4b1d8" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_VECTOR = """ +CREATE FUNCTION search_vector(body text) RETURNS tsvector +LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ + SELECT to_tsvector('russian', translate(coalesce(body, ''), 'Ёё', 'Ее')) + || to_tsvector('simple', translate(coalesce(body, ''), 'Ёё', 'Ее')) +$$ +""" + +_TABLE = """ +CREATE TABLE message_search ( + account_id integer NOT NULL, + chat_id bigint NOT NULL, + message_id bigint NOT NULL, + date timestamptz NOT NULL, + sender_id bigint, + text text, + has_media boolean NOT NULL DEFAULT false, + edited_at timestamptz, + deleted_at timestamptz, + tsv tsvector GENERATED ALWAYS AS (search_vector(text)) STORED, + PRIMARY KEY (account_id, chat_id, message_id, date) +) +""" + +_BACKFILL = """ +INSERT INTO message_search (account_id, chat_id, message_id, date, sender_id, + text, has_media, edited_at, deleted_at) +SELECT account_id, chat_id, message_id, date, sender_id, + text, has_media, edited_at, deleted_at +FROM messages +""" + +_SYNC = """ +CREATE FUNCTION message_search_sync() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN + IF TG_OP = 'DELETE' THEN + DELETE FROM message_search + WHERE account_id = OLD.account_id AND chat_id = OLD.chat_id + AND message_id = OLD.message_id AND date = OLD.date; + RETURN NULL; + END IF; + INSERT INTO message_search (account_id, chat_id, message_id, date, + sender_id, text, has_media, edited_at, + deleted_at) + VALUES (NEW.account_id, NEW.chat_id, NEW.message_id, NEW.date, + NEW.sender_id, NEW.text, NEW.has_media, NEW.edited_at, + NEW.deleted_at) + ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET + sender_id = EXCLUDED.sender_id, + text = EXCLUDED.text, + has_media = EXCLUDED.has_media, + edited_at = EXCLUDED.edited_at, + deleted_at = EXCLUDED.deleted_at; + RETURN NULL; +END; +$$ +""" + +_CHANGED = ( + "OLD.text IS DISTINCT FROM NEW.text " + "OR OLD.sender_id IS DISTINCT FROM NEW.sender_id " + "OR OLD.has_media IS DISTINCT FROM NEW.has_media " + "OR OLD.edited_at IS DISTINCT FROM NEW.edited_at " + "OR OLD.deleted_at IS DISTINCT FROM NEW.deleted_at" +) + + +def upgrade() -> None: + op.execute(_VECTOR) + op.execute(_TABLE) + op.execute(_BACKFILL) + op.execute("CREATE INDEX ix_message_search_tsv ON message_search USING gin (tsv)") + op.execute( + "CREATE INDEX ix_message_search_recent ON message_search " + "(account_id, date DESC, chat_id DESC, message_id DESC)" + ) + op.execute( + "CREATE INDEX ix_message_search_sender ON message_search " + "(account_id, sender_id, date DESC)" + ) + op.execute(_SYNC) + op.execute( + "CREATE TRIGGER messages_search_insert AFTER INSERT ON messages " + "FOR EACH ROW EXECUTE FUNCTION message_search_sync()" + ) + op.execute( + "CREATE TRIGGER messages_search_update AFTER UPDATE ON messages " + f"FOR EACH ROW WHEN ({_CHANGED}) EXECUTE FUNCTION message_search_sync()" + ) + op.execute( + "CREATE TRIGGER messages_search_delete AFTER DELETE ON messages " + "FOR EACH ROW EXECUTE FUNCTION message_search_sync()" + ) + op.execute( + "CREATE INDEX ix_message_versions_tsv ON message_versions USING gin " + "(search_vector(text))" + ) + op.execute("DROP INDEX IF EXISTS ix_media_tsv") + op.execute( + "CREATE INDEX ix_media_tsv ON media USING gin (search_vector(extracted_text))" + ) + op.execute("DROP INDEX IF EXISTS ix_messages_tsv") + + +def downgrade() -> None: + op.execute( + "CREATE INDEX ix_messages_tsv ON messages USING gin " + "(to_tsvector('russian', coalesce(text, '')))" + ) + op.execute("DROP INDEX IF EXISTS ix_media_tsv") + op.execute( + "CREATE INDEX ix_media_tsv ON media USING gin " + "(to_tsvector('russian', coalesce(extracted_text, '')))" + ) + op.execute("DROP INDEX IF EXISTS ix_message_versions_tsv") + op.execute("DROP TRIGGER IF EXISTS messages_search_delete ON messages") + op.execute("DROP TRIGGER IF EXISTS messages_search_update ON messages") + op.execute("DROP TRIGGER IF EXISTS messages_search_insert ON messages") + op.execute("DROP FUNCTION IF EXISTS message_search_sync()") + op.execute("DROP TABLE IF EXISTS message_search") + op.execute("DROP FUNCTION IF EXISTS search_vector(text)") diff --git a/backend/src/api/mcp/server.py b/backend/src/api/mcp/server.py index e302d23..31eec40 100644 --- a/backend/src/api/mcp/server.py +++ b/backend/src/api/mcp/server.py @@ -38,7 +38,7 @@ from utils.read.feed import CURSOR_PATTERN, FeedChatKind, FeedFilter, FeedKind from utils.read.models import DEFAULT_LIMIT, Page from utils.read.period import DEFAULT_TZ, MAX_DAYS, known_timezone from utils.read.read_receipts import read_info -from utils.search.models import SearchFilters +from utils.search.models import ChatKind, SearchFilters from utils.search.repository import search_messages from utils.storage import ContentAddressedStorage from utils.watch_rules import WatchKind, WatchParamsError, clean_name, normalize_params @@ -72,8 +72,9 @@ Typical flows: "📝 [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). +- Find something: `search_messages_tool` (words by prefix over message text, + STT transcripts and earlier versions of edited messages; filters by chat, + sender, chat kind, media, deleted, edited, dates, regex). - People: `get_person` gives one person in full: profile changes, shared groups, joins and leaves, deleted message count, the private chat. - Forensics: `get_feed` lists what was deleted or edited in Telegram across @@ -99,7 +100,8 @@ Watched bots: 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`. +Paging: list tools accept `limit` and `offset`; `search_messages_tool` returns +`next_cursor` to pass back as `cursor`. """ mcp: FastMCP = FastMCP("beavergram", instructions=INSTRUCTIONS) @@ -138,27 +140,49 @@ async def search_messages_tool( query: str | None = None, chat_id: int | None = None, sender_id: int | None = None, + kind: ChatKind | None = None, has_media: bool | None = None, + deleted: bool | None = None, + edited: bool | None = None, date_from: datetime | None = None, date_to: datetime | None = None, regex: str | None = None, + cursor: str | None = None, limit: int = DEFAULT_LIMIT, - offset: int = 0, -) -> list[dict[str, Any]]: - """Full-text search over message text and STT transcripts.""" +) -> dict[str, Any]: + """Search messages by words, newest first, including deleted ones. + + Words match by prefix (3+ letters) and by stem, Russian and English. + `"exact phrase"`, `-word` and `or` work as in web search. Besides the + current text it searches voice transcripts and earlier versions of edited + messages: a hit with `source="version"` matched only an old text, given in + `version_text`. Every hit carries `deleted_at` and `edited_at`. + + Filters: `chat_id`, `sender_id` (peer id), `kind` (private, group, + channel), `has_media`, `deleted` / `edited` (true for only those, false to + exclude them), `date_from` / `date_to`, `regex` over the current text. With + no `query` it lists messages that pass the filters. + + Paging: pass `next_cursor` from the previous result as `cursor`; it is null + on the last page. + """ filters = SearchFilters( account_id=account_id, query=query, chat_id=chat_id, sender_id=sender_id, + kind=kind, has_media=has_media, + deleted=deleted, + edited=edited, date_from=date_from, date_to=date_to, regex=regex, + cursor=cursor, limit=limit, - offset=offset, ) - return _dump(await search_messages(await _pool(), filters)) + page = await search_messages(await _pool(), filters) + return page.model_dump(mode="json", exclude={"terms"}) @mcp.tool diff --git a/backend/src/api/routers/search.py b/backend/src/api/routers/search.py index 613900f..a8137b1 100644 --- a/backend/src/api/routers/search.py +++ b/backend/src/api/routers/search.py @@ -4,7 +4,7 @@ import asyncpg from dishka.integrations.fastapi import DishkaRoute, FromDishka from fastapi import APIRouter, Query -from utils.search.models import SearchFilters, SearchHit +from utils.search.models import SearchFilters, SearchPage from utils.search.repository import search_messages router = APIRouter(prefix="/api/search", tags=["search"], route_class=DishkaRoute) @@ -13,5 +13,5 @@ router = APIRouter(prefix="/api/search", tags=["search"], route_class=DishkaRout @router.get("") async def search( pool: FromDishka[asyncpg.Pool], filters: Annotated[SearchFilters, Query()] -) -> list[SearchHit]: +) -> SearchPage: return await search_messages(pool, filters) diff --git a/backend/src/utils/search/models.py b/backend/src/utils/search/models.py index bf9c2b6..bf3c9c9 100644 --- a/backend/src/utils/search/models.py +++ b/backend/src/utils/search/models.py @@ -1,23 +1,48 @@ -from datetime import datetime +from datetime import UTC, datetime, timedelta from typing import Literal -from pydantic import BaseModel +from pydantic import BaseModel, field_validator DEFAULT_LIMIT = 50 MAX_LIMIT = 500 +_EPOCH = datetime(1970, 1, 1, tzinfo=UTC) +_MICROSECOND = timedelta(microseconds=1) + +ChatKind = Literal["private", "group", "channel"] +HitSource = Literal["text", "stt", "version"] + + +def encode_cursor(date: datetime, chat_id: int, message_id: int) -> str: + return f"{(date - _EPOCH) // _MICROSECOND}_{chat_id}_{message_id}" + + +def decode_cursor(cursor: str) -> tuple[datetime, int, int]: + micros, chat_id, message_id = (int(part) for part in cursor.split("_")) + return _EPOCH + micros * _MICROSECOND, chat_id, message_id + class SearchFilters(BaseModel): account_id: int query: str | None = None chat_id: int | None = None sender_id: int | None = None + kind: ChatKind | None = None has_media: bool | None = None + deleted: bool | None = None + edited: bool | None = None date_from: datetime | None = None date_to: datetime | None = None regex: str | None = None + cursor: str | None = None limit: int = DEFAULT_LIMIT - offset: int = 0 + + @field_validator("cursor") + @classmethod + def _valid_cursor(cls, value: str | None) -> str | None: + if value is not None: + decode_cursor(value) + return value class SearchHit(BaseModel): @@ -26,5 +51,17 @@ class SearchHit(BaseModel): date: datetime sender_id: int | None text: str | None - source: Literal["text", "stt"] + has_media: bool + media_kind: str | None + edited_at: datetime | None + deleted_at: datetime | None + source: HitSource extracted_text: str | None + version_text: str | None + version_at: datetime | None + + +class SearchPage(BaseModel): + hits: list[SearchHit] + next_cursor: str | None + terms: list[str] diff --git a/backend/src/utils/search/repository.py b/backend/src/utils/search/repository.py index 10b96cf..5423b32 100644 --- a/backend/src/utils/search/repository.py +++ b/backend/src/utils/search/repository.py @@ -1,88 +1,178 @@ import asyncpg -from utils.search.models import MAX_LIMIT, SearchFilters, SearchHit - -_TEXT_TSV = "to_tsvector('russian', coalesce(m.text, ''))" -_STT_TSV = "to_tsvector('russian', coalesce(md.extracted_text, ''))" -_TSQUERY = "websearch_to_tsquery('russian', ${})" - -_HITS = "{cols} FROM messages m WHERE {where}" # text branch -_STT_HITS = ( - "{cols} FROM messages m " - "JOIN media md ON md.account_id = m.account_id " - "AND md.chat_id = m.chat_id AND md.message_id = m.message_id " - "WHERE {where}" -) -_TEXT_COLS = ( - "SELECT m.chat_id, m.message_id, m.date, m.sender_id, m.text, " - "'text' AS source, NULL::text AS extracted_text" -) -_STT_COLS = ( - "SELECT m.chat_id, m.message_id, m.date, m.sender_id, m.text, " - "'stt' AS source, md.extracted_text" +from utils.search.models import ( + MAX_LIMIT, + SearchFilters, + SearchHit, + SearchPage, + decode_cursor, + encode_cursor, ) +_PARSE = """ +SELECT (regexp_replace(websearch_to_tsquery('russian', q)::text, + '''((?:[^'']|''''){3,})''', '''\\1'':*', 'g')::tsquery + || regexp_replace(websearch_to_tsquery('simple', q)::text, + '''((?:[^'']|''''){3,})''', '''\\1'':*', 'g')::tsquery + )::text AS tsquery, + tsvector_to_array(search_vector(q)) AS terms +FROM translate($1, 'Ёё', 'Ее') AS q +""" -def _common_conditions(filters: SearchFilters, params: list[object]) -> list[str]: - params.append(filters.account_id) - conds = [f"m.account_id = ${len(params)}"] - if filters.chat_id is not None: - params.append(filters.chat_id) - conds.append(f"m.chat_id = ${len(params)}") - if filters.sender_id is not None: - params.append(filters.sender_id) - conds.append(f"m.sender_id = ${len(params)}") - if filters.has_media is not None: - params.append(filters.has_media) - conds.append(f"m.has_media = ${len(params)}") - if filters.date_from is not None: - params.append(filters.date_from) - conds.append(f"m.date >= ${len(params)}") - if filters.date_to is not None: - params.append(filters.date_to) - conds.append(f"m.date <= ${len(params)}") +_CHANNELS = """ +(SELECT chat_id FROM ( + SELECT DISTINCT ON (chat_id) chat_id, + coalesce(raw->'chat'->>'type', raw->>'type') AS type + FROM chat_history + WHERE account_id = {account} + AND coalesce(raw->'chat'->>'type', raw->>'type') IS NOT NULL + ORDER BY chat_id, ts DESC +) typed WHERE type = 'ChatType.CHANNEL') +""" + +_KEY_ORDER = "s.date DESC, s.chat_id DESC, s.message_id DESC" + +_COLS = ( + "s.chat_id, s.message_id, s.date, s.sender_id, s.text, s.has_media, " + "s.edited_at, s.deleted_at" +) + +_TEXT = ( + f"SELECT {_COLS}, 0 AS rank, NULL::text AS extracted_text, " # noqa: S608 + "NULL::text AS version_text, NULL::timestamptz AS version_at " + "FROM message_search s WHERE {where}" +) + +_STT = ( + f"SELECT {_COLS}, 1 AS rank, md.extracted_text, NULL::text, " # noqa: S608 + "NULL::timestamptz " + "FROM media md JOIN message_search s ON s.account_id = md.account_id " + "AND s.chat_id = md.chat_id AND s.message_id = md.message_id " + "WHERE md.account_id = {account} " + "AND search_vector(md.extracted_text) @@ {tsquery} " + "AND {where}" +) + +_VERSION = ( + "SELECT DISTINCT ON (s.date, s.chat_id, s.message_id) " # noqa: S608 + f"{_COLS}, 2 AS rank, " + "NULL::text, v.text, v.observed_at " + "FROM message_versions v JOIN message_search s ON s.account_id = v.account_id " + "AND s.chat_id = v.chat_id AND s.message_id = v.message_id " + "WHERE v.account_id = {account} " + "AND search_vector(v.text) @@ {tsquery} " + "AND {where}" +) + +_WITH_MEDIA = ( + "SELECT h.*, md.kind AS media_kind FROM ({hits}) h " + "LEFT JOIN media md ON md.account_id = {account} " + "AND md.chat_id = h.chat_id AND md.message_id = h.message_id " + "ORDER BY h.date DESC, h.chat_id DESC, h.message_id DESC" +) + +_SOURCES = ("text", "stt", "version") + + +class _Binder: + def __init__(self) -> None: + self.params: list[object] = [] + + def __call__(self, value: object) -> str: + self.params.append(value) + return f"${len(self.params)}" + + +def _kind(filters: SearchFilters, bind: _Binder) -> str | None: + if filters.kind is None: + return None + if filters.kind == "private": + return "s.chat_id > 0" + channels = _CHANNELS.format(account=bind(filters.account_id)) + if filters.kind == "channel": + return f"s.chat_id IN {channels}" + return f"s.chat_id < 0 AND s.chat_id NOT IN {channels}" + + +def _conditions(filters: SearchFilters, bind: _Binder) -> list[str]: + conds = [f"s.account_id = {bind(filters.account_id)}"] + compared = ( + ("s.chat_id =", filters.chat_id), + ("s.sender_id =", filters.sender_id), + ("s.has_media =", filters.has_media), + ("s.date >=", filters.date_from), + ("s.date <=", filters.date_to), + ("s.text ~", filters.regex), + ) + conds += [f"{left} {bind(value)}" for left, value in compared if value is not None] + flags = (("s.deleted_at", filters.deleted), ("s.edited_at", filters.edited)) + conds += [ + f"{column} IS {'NOT ' if flag else ''}NULL" + for column, flag in flags + if flag is not None + ] + kind = _kind(filters, bind) + if kind is not None: + conds.append(kind) + if filters.cursor is not None: + date, chat_id, message_id = decode_cursor(filters.cursor) + conds.append( + "(s.date, s.chat_id, s.message_id) < " + f"({bind(date)}, {bind(chat_id)}, {bind(message_id)})" + ) return conds -async def search_messages( - pool: asyncpg.Pool, filters: SearchFilters -) -> list[SearchHit]: - query = filters.query.strip() if filters.query else None - params: list[object] = [] - common = _common_conditions(filters, params) - if filters.regex is not None: - params.append(filters.regex) - common.append(f"m.text ~ ${len(params)}") +def _branch(sql: str, limit: str, order: str = _KEY_ORDER, **parts: str) -> str: + return f"({sql.format(**parts)} ORDER BY {order} LIMIT {limit})" - q_idx: int | None = None - if query is not None: - params.append(query) - q_idx = len(params) - text_conds = [*common] - if q_idx is not None: - text_conds.append(f"{_TEXT_TSV} @@ {_TSQUERY.format(q_idx)}") +async def search_messages(pool: asyncpg.Pool, filters: SearchFilters) -> SearchPage: + query = (filters.query or "").strip() + terms: list[str] = [] + tsquery = None + if query: + parsed = await pool.fetchrow(_PARSE, query) + terms = list(parsed["terms"]) + if not parsed["tsquery"]: + return SearchPage(hits=[], next_cursor=None, terms=terms) + tsquery = parsed["tsquery"] - branches = [_HITS.format(cols=_TEXT_COLS, where=" AND ".join(text_conds))] - if q_idx is not None: - stt_conds = [*common, f"{_STT_TSV} @@ {_TSQUERY.format(q_idx)}"] - branches.append(_STT_HITS.format(cols=_STT_COLS, where=" AND ".join(stt_conds))) + bind = _Binder() + where = " AND ".join(_conditions(filters, bind)) + size = min(filters.limit, MAX_LIMIT) + limit = bind(size + 1) - params.append(min(filters.limit, MAX_LIMIT)) - limit_idx = len(params) - params.append(filters.offset) - offset_idx = len(params) - page = f"ORDER BY date DESC LIMIT ${limit_idx} OFFSET ${offset_idx}" - - if len(branches) == 1: - sql = f"SELECT * FROM ({branches[0]}) hits {page}" # noqa: S608 + account = bind(filters.account_id) + if tsquery is None: + sql = _branch(_TEXT, limit, where=where) else: - union = " UNION ALL ".join(f"({b})" for b in branches) - dedup = ( - "SELECT DISTINCT ON (chat_id, message_id) * " # noqa: S608 - f"FROM ({union}) hits ORDER BY chat_id, message_id, source DESC" + match = f"{bind(tsquery)}::tsquery" + branches = ( + _branch(_TEXT, limit, where=f"{where} AND s.tsv @@ {match}"), + _branch(_STT, limit, where=where, account=account, tsquery=match), + _branch( + _VERSION, + limit, + f"{_KEY_ORDER}, v.observed_at DESC", + where=where, + account=account, + tsquery=match, + ), + ) + sql = ( + "SELECT DISTINCT ON (date, chat_id, message_id) * FROM (" # noqa: S608 + f"{' UNION ALL '.join(branches)}) s " + f"ORDER BY {_KEY_ORDER}, rank LIMIT {limit}" ) - sql = f"SELECT * FROM ({dedup}) d {page}" # noqa: S608 - rows = await pool.fetch(sql, *params) - return [SearchHit(**dict(row)) for row in rows] + rows = await pool.fetch(_WITH_MEDIA.format(hits=sql, account=account), *bind.params) + hits = [ + SearchHit(**{**dict(row), "source": _SOURCES[row["rank"]]}) + for row in rows[:size] + ] + next_cursor = None + if len(rows) > size: + last = hits[-1] + next_cursor = encode_cursor(last.date, last.chat_id, last.message_id) + return SearchPage(hits=hits, next_cursor=next_cursor, terms=terms)