diff --git a/src/bot.py b/src/bot.py index 3ba3551..439b65a 100644 --- a/src/bot.py +++ b/src/bot.py @@ -15,6 +15,7 @@ from config import config from google_auth import generate_auth_url, TokenStorage from timebox import build_timebox_handler +from search import search_command import drive_handler # Configure logging @@ -38,6 +39,7 @@ async def start_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> N "You are authenticated with Google Drive.\n" "Send me any message and I'll save it to your Drive.\n\n" "Commands:\n" + "/search - Search your captured notes\n" "/timebox - Plan a timeboxed schedule for tomorrow\n" "/status - Check connection status\n" "/logout - Disconnect Google Drive\n" @@ -69,6 +71,7 @@ async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> No "Edit a message here and it updates in Drive too.\n\n" "Commands:\n" "/authenticate - Connect your Google Drive\n" + "/search - Search your captured notes\n" "/timebox - Plan a timeboxed schedule for tomorrow\n" "/status - Check connection status\n" "/logout - Disconnect Google Drive\n" @@ -309,6 +312,7 @@ def register_handlers(application: Application) -> None: application.add_handler(CommandHandler("authenticate", authenticate_command)) application.add_handler(CommandHandler("status", status_command)) application.add_handler(CommandHandler("logout", logout_command)) + application.add_handler(CommandHandler("search", search_command)) application.add_handler(build_timebox_handler()) diff --git a/src/drive_handler.py b/src/drive_handler.py index bc55924..6f88518 100644 --- a/src/drive_handler.py +++ b/src/drive_handler.py @@ -250,6 +250,26 @@ def _replace_message_content( return file_content[:header_start] + new_block + file_content[content_end:] +def list_markdown_files(service, folder_id: str) -> list[dict]: + """Return the markdown files in the folder as [{id, name}, ...] (read-only).""" + try: + results = service.files().list( + q=f"mimeType='{MARKDOWN_MIME_TYPE}' and trashed=false and parents='{folder_id}'", + spaces='drive', + fields='files(id, name)', + pageSize=1000, + ).execute() + return results.get('files', []) + except Exception as e: + logger.error(f"Failed to list markdown files: {e}") + return [] + + +def read_file(service, file_id: str) -> str | None: + """Public read of a Drive file's text content, or None on failure.""" + return _download_file_content(service, file_id) + + def _download_file_content(service, file_id: str) -> str | None: """Download a file's content from Drive.""" try: diff --git a/src/search.py b/src/search.py new file mode 100644 index 0000000..3ef2a4b --- /dev/null +++ b/src/search.py @@ -0,0 +1,107 @@ +"""/search: keyword search across the captured daily notes in Drive. + +Thin Telegram glue. The matching/parsing helpers are pure functions so they +can be unit-tested without Drive. Scans the most recent daily files only, so a +long history stays cheap. +""" + +import logging +import re + +from telegram import Update +from telegram.ext import ContextTypes + +from config import config +import drive_handler + +logger = logging.getLogger(__name__) + +_BLOCK_RE = re.compile(r"\n") +_DAILY_RE = re.compile(r"^(\d{4}-\d{2}-\d{2})\.md$") + +_MAX_FILES = 60 # how many recent daily files to scan +_MAX_HITS = 20 # how many matches to show +_SNIPPET_LEN = 160 # per-hit snippet length +_MAX_REPLY = 3900 + + +def parse_blocks(content: str) -> list[str]: + """Captured note texts in order, stripped of the msg_id markers.""" + parts = _BLOCK_RE.split(content) + return [p.strip() for p in parts[1:] if p.strip()] + + +def find_matches(content: str, query: str) -> list[str]: + """Notes in `content` that contain `query` (case-insensitive substring).""" + q = query.lower() + return [n for n in parse_blocks(content) if q in n.lower()] + + +def _snippet(note: str) -> str: + """Collapse a note to a single trimmed line for compact display.""" + one_line = " ".join(note.split()) + if len(one_line) > _SNIPPET_LEN: + one_line = one_line[:_SNIPPET_LEN].rstrip() + "…" + return one_line + + +def _date_of(file_name: str) -> str: + """The YYYY-MM-DD date encoded in a daily file name, or the raw name.""" + m = _DAILY_RE.match(file_name) + return m.group(1) if m else file_name + + +async def search_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Handle /search - find captured notes matching the query.""" + token_storage = context.bot_data.get("token_storage") + user_id = update.effective_user.id + if not token_storage or not token_storage.is_authenticated(user_id): + await update.message.reply_text( + "Please authenticate with Google Drive first using /authenticate" + ) + return + + query = " ".join(context.args).strip() if context.args else "" + if not query: + await update.message.reply_text("Usage: /search \ne.g. /search dentist") + return + + service = drive_handler.get_drive_service(user_id, token_storage) + if not service: + await update.message.reply_text( + "Your Google Drive session expired — please /logout then /authenticate." + ) + return + + folder_id = drive_handler.get_or_create_folder(service, config.drive_folder_name) + if not folder_id: + await update.message.reply_text( + "Couldn't reach your Drive folder — please try again later." + ) + return + + files = drive_handler.list_markdown_files(service, folder_id) + # Newest day first; cap the scan so a long history stays cheap. + files = sorted(files, key=lambda f: f.get("name", ""), reverse=True)[:_MAX_FILES] + + hits: list[tuple[str, str]] = [] # (date, snippet) + for f in files: + if len(hits) >= _MAX_HITS: + break + content = drive_handler.read_file(service, f["id"]) or "" + for note in find_matches(content, query): + hits.append((_date_of(f.get("name", "")), _snippet(note))) + if len(hits) >= _MAX_HITS: + break + + if not hits: + await update.message.reply_text(f'No notes found for "{query}".') + return + + lines = [f'šŸ”Ž {len(hits)} match(es) for "{query}":', ""] + for d, snip in hits: + lines.append(f"[{d}] {snip}") + reply = "\n".join(lines) + if len(reply) > _MAX_REPLY: + reply = reply[:_MAX_REPLY] + "\n…(truncated)" + await update.message.reply_text(reply) diff --git a/tests/test_search.py b/tests/test_search.py new file mode 100644 index 0000000..e4aca3c --- /dev/null +++ b/tests/test_search.py @@ -0,0 +1,37 @@ +from search import parse_blocks, find_matches, _snippet, _date_of + +CONTENT = ( + "# Telegram Messages\n\n" + "\nBook dentist appointment\n" + "\nGroceries: milk, eggs\n" + "\nDENTIST follow-up next month\n" +) + + +def test_parse_blocks(): + assert parse_blocks(CONTENT) == [ + "Book dentist appointment", + "Groceries: milk, eggs", + "DENTIST follow-up next month", + ] + + +def test_find_matches_is_case_insensitive(): + hits = find_matches(CONTENT, "dentist") + assert hits == ["Book dentist appointment", "DENTIST follow-up next month"] + + +def test_find_matches_none(): + assert find_matches(CONTENT, "passport") == [] + + +def test_snippet_collapses_and_truncates(): + assert _snippet("line one\n line two") == "line one line two" + long = "x" * 300 + out = _snippet(long) + assert out.endswith("…") and len(out) <= 161 + + +def test_date_of(): + assert _date_of("2026-06-25.md") == "2026-06-25" + assert _date_of("weird.md") == "weird.md"