beaver/beavergram
2
0 Fork 0
Code Issues 1 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

feat(backend): search finds word prefixes, earlier versions and deleted messages, with filters and cursor paging

Browse Source
This commit is contained in:
hh committed 2026-10-19 03:41:39 +00:00
1 parent 42b4e1d7e6
commit d45c96c259
5 files changed
+377 -88
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/migrations/versions/d8f1a6c3e5b7_message_search.py
+138
View File
@@ -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)")
backend/src/api/mcp/server.py
+33 -9
View File
@@ -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
backend/src/api/routers/search.py
+2 -2
View File
@@ -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)
backend/src/utils/search/models.py
+41 -4
View File
@@ -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]
backend/src/utils/search/repository.py
+163 -73
View File
@@ -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)