msgsearch 0.1.0__py3-none-any.whl

msgsearch/__init__.py

@@ -0,0 +1,18 @@
+"""Search your iMessage history by meaning, entirely on your own machine.
+
+The pipeline, in the order data flows through it:
+
+    extract   ->  chat.db rows become readable messages, decoding the
+                  `attributedBody` blob that holds ~86% of them
+    chunk     ->  messages become conversation windows, then the smaller
+                  passages that actually get embedded
+    index     ->  passages become vectors; windows become a keyword index
+    search    ->  a query becomes ranked windows, via keyword search, vector
+                  search, rank fusion and optional reranking
+
+Nothing here sends message text anywhere. Both models run locally.
+"""
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__"]

msgsearch/attributed_body.py

@@ -0,0 +1,59 @@
+"""Decode the `message.attributedBody` blob (Apple NSArchiver 'typedstream').
+
+~86% of rows in chat.db have text IS NULL and carry their content here instead.
+
+Layout, per hexdump of a real row:
+
+    04 0b "streamtyped" 81 e8 03 ... "NSString" 01 94 84 01 2b <len> <utf8 bytes>
+                                                          '+'  ^ varint length
+
+The 0x2b ('+') marker introduces a length-prefixed byte string. Length is one
+byte when < 0x80, otherwise 0x81 => uint16 LE, 0x82 => uint32 LE.
+"""
+
+import struct
+
+_MARKER = b"NSString"
+
+
+def _read_len(buf, i):
+    """Return (length, next_index) for the varint length at buf[i]."""
+    n = buf[i]
+    if n < 0x80:
+        return n, i + 1
+    if n == 0x81:
+        return struct.unpack_from("<H", buf, i + 1)[0], i + 3
+    if n == 0x82:
+        return struct.unpack_from("<I", buf, i + 1)[0], i + 5
+    return None, i + 1
+
+
+def decode(blob):
+    """Extract the message text from an attributedBody blob, or None."""
+    if not blob:
+        return None
+    buf = bytes(blob)
+
+    start = buf.find(_MARKER)
+    if start == -1:
+        return None
+    i = buf.find(b"\x2b", start + len(_MARKER))
+    if i == -1:
+        return None
+
+    length, i = _read_len(buf, i + 1)
+    if not length or i + length > len(buf):
+        return None
+
+    raw = buf[i : i + length]
+    # Runs of the string are UTF-8; a leading BOM-ish 0xff 0xfe marks UTF-16.
+    if raw[:2] == b"\xff\xfe":
+        return raw[2:].decode("utf-16-le", errors="replace")
+    return raw.decode("utf-8", errors="replace")
+
+
+def message_text(text, blob):
+    """Prefer the plain `text` column, fall back to decoding the blob."""
+    if text is not None:
+        return text
+    return decode(blob)

msgsearch/chunk.py

@@ -0,0 +1,268 @@
+"""Group messages into conversation windows.
+
+Embedding messages one at a time does not work for this problem. The thing you
+want to find is often a reply whose text has nothing in common with your query:
+someone asks for a login, and two minutes later a message arrives containing an
+email address and a password but neither the word "login" nor the name of the
+service. Searched alone, that message is unreachable. Searched as part of the
+exchange it belongs to, it is easy to find.
+
+So the unit of retrieval is a window: a run of messages in one conversation with
+no long pause in it. Windows are cut where a real pause happened, and only split
+further when they are too long for the embedding model to read.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Iterable, Iterator, Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from itertools import pairwise
+
+from . import config, tagging
+from .extract import Message
+
+
+@dataclass(frozen=True)
+class Window:
+    window_id: str
+    chat_id: int
+    chat_label: str
+    start: datetime
+    end: datetime
+    messages: tuple[Message, ...]
+
+    # What the embedding model reads. Link-only messages are left out because
+    # a bare URL has no meaning for the model to work with.
+    embed_text: str
+
+    # What keyword search reads, and what is shown to the user. Nothing is left
+    # out here, so a URL someone sent is still findable by typing part of it.
+    search_text: str
+
+    # (message rowid, start offset, end offset) into search_text, so a result can
+    # highlight the message that matched while displaying its whole context.
+    offsets: tuple[tuple[int, int, int], ...]
+
+    speakers: tuple[str, ...]
+    tags: frozenset
+
+    @property
+    def is_one_sided(self) -> bool:
+        """True if only one person spoke, which is usually a link dump."""
+        return len(set(self.speakers)) <= 1
+
+
+def _line_cost(message: Message) -> int:
+    """Characters a message contributes once rendered as `Speaker: text`."""
+    return len(message.speaker) + 2 + len(message.text) + 1
+
+
+def _rendered_cost(messages: Sequence[Message]) -> int:
+    """Length of these messages once rendered, including the date headers.
+
+    The headers have to be counted here. Leaving them out lets a window creep
+    past the budget by a few characters per day it spans.
+    """
+    total = 0
+    day = None
+    for message in messages:
+        this_day = message.timestamp.strftime("%Y-%m-%d")
+        if this_day != day:
+            total += len(f"[{this_day}]\n")
+            day = this_day
+        total += _line_cost(message)
+    return total
+
+
+def _render(messages: Sequence[Message]) -> tuple[str, tuple[tuple[int, int, int], ...]]:
+    """Render messages as dated, speaker-labelled lines.
+
+    Returns the text and, for each message, where its body starts and ends inside
+    that text.
+    """
+    parts: list[str] = []
+    offsets: list[tuple[int, int, int]] = []
+    position = 0
+    current_day = None
+
+    for message in messages:
+        day = message.timestamp.strftime("%Y-%m-%d")
+        if day != current_day:
+            header = f"[{day}]\n"
+            parts.append(header)
+            position += len(header)
+            current_day = day
+
+        prefix = f"{message.speaker}: "
+        line = f"{prefix}{message.text}\n"
+        body_start = position + len(prefix)
+        parts.append(line)
+        position += len(line)
+        offsets.append((message.rowid, body_start, position - 1))
+
+    return "".join(parts), tuple(offsets)
+
+
+def _split_on_pauses(
+    messages: Sequence[Message], gap_seconds: int
+) -> list[list[Message]]:
+    """Cut a conversation wherever nobody spoke for longer than `gap_seconds`."""
+    groups: list[list[Message]] = []
+    current: list[Message] = [messages[0]]
+
+    for previous, message in pairwise(messages):
+        if (message.timestamp - previous.timestamp).total_seconds() > gap_seconds:
+            groups.append(current)
+            current = []
+        current.append(message)
+
+    groups.append(current)
+    return groups
+
+
+def _split_oversized(
+    messages: Sequence[Message], budget_chars: int, overlap: int
+) -> list[list[Message]]:
+    """Break a window that is too long for the embedding model to read.
+
+    Only windows over the budget are touched, which on real data is around 1% of
+    them. Each piece after the first repeats the last couple of messages of the
+    previous piece, so an exchange sitting on a boundary stays intact in one of
+    the two pieces.
+    """
+    if _rendered_cost(messages) <= budget_chars:
+        return [list(messages)]
+
+    pieces: list[list[Message]] = []
+    current: list[Message] = []
+
+    for message in messages:
+        # Cost is recomputed against the real rendering rather than tracked
+        # incrementally, because a date header appears only on the first message
+        # of each day and that depends on where the piece boundaries fall. These
+        # windows are small, so the repeated work is not worth optimising away.
+        if current and _rendered_cost([*current, message]) > budget_chars:
+            pieces.append(current)
+
+            # Carry a little context forward, but never so much that the overlap
+            # alone fills the next piece.
+            tail = list(current[-overlap:]) if overlap else []
+            if _rendered_cost(tail) > budget_chars // 2:
+                tail = []
+            current = tail
+
+        current.append(message)
+
+    if current:
+        pieces.append(current)
+    return pieces
+
+
+def build_window(messages: Sequence[Message]) -> Window:
+    search_text, offsets = _render(messages)
+
+    # Link-only messages are dropped from the embedded text only. They stay in
+    # search_text, in the offsets and in the displayed result.
+    speaking = [m for m in messages if not tagging.is_bare_url(m.text)]
+    embed_text, _ = _render(speaking) if speaking else ("", ())
+
+    first = messages[0]
+    return Window(
+        window_id=f"{first.chat_id}:{first.rowid}",
+        chat_id=first.chat_id,
+        chat_label=first.chat_label,
+        start=first.timestamp,
+        end=messages[-1].timestamp,
+        messages=tuple(messages),
+        embed_text=embed_text,
+        search_text=search_text,
+        offsets=offsets,
+        speakers=tuple(m.speaker for m in messages),
+        tags=frozenset().union(*(tagging.tags(m.text) for m in messages)),
+    )
+
+
+@dataclass(frozen=True)
+class Passage:
+    """A small slice of a window, and the thing that actually gets embedded.
+
+    A window is the right unit to show a person, because it carries the context
+    that makes a result make sense. It is the wrong unit to embed, because a
+    single vector for thirty messages on eight topics represents none of them
+    well. Passages slide across the window in small overlapping steps so that
+    every message sits near the middle of at least one of them, and each passage
+    remembers which window it belongs to.
+    """
+
+    passage_id: str
+    window_id: str
+    chat_id: int
+    text: str
+    rowids: tuple[int, ...]
+
+
+def passages(
+    window: Window, size: int | None = None, stride: int | None = None
+) -> list[Passage]:
+    """Slice a window into overlapping passages for embedding."""
+    size = size if size is not None else config.PASSAGE_MESSAGES
+    stride = stride if stride is not None else config.PASSAGE_STRIDE
+
+    # Link-only messages are dropped here for the same reason they are dropped
+    # from embed_text: a bare URL gives the model nothing to work with.
+    speaking = [m for m in window.messages if not tagging.is_bare_url(m.text)]
+    if not speaking:
+        return []
+
+    if len(speaking) <= size:
+        starts = [0]
+    else:
+        starts = list(range(0, len(speaking) - size + 1, stride))
+        # Make sure the final messages are covered even when the stride does not
+        # divide the window evenly.
+        last = len(speaking) - size
+        if starts[-1] != last:
+            starts.append(last)
+
+    out = []
+    for start in starts:
+        selected = speaking[start : start + size]
+        text, _ = _render(selected)
+        out.append(
+            Passage(
+                passage_id=f"{window.window_id}#{start}",
+                window_id=window.window_id,
+                chat_id=window.chat_id,
+                text=text,
+                rowids=tuple(m.rowid for m in selected),
+            )
+        )
+    return out
+
+
+def windows(
+    messages: Iterable[Message],
+    gap_seconds: int | None = None,
+    token_budget: int | None = None,
+    overlap: int | None = None,
+) -> Iterator[Window]:
+    """Group messages into conversation windows, oldest first."""
+    gap_seconds = gap_seconds if gap_seconds is not None else config.WINDOW_GAP_SECONDS
+    token_budget = (
+        token_budget if token_budget is not None else config.WINDOW_TOKEN_BUDGET
+    )
+    overlap = overlap if overlap is not None else config.WINDOW_OVERLAP_MESSAGES
+    budget_chars = token_budget * config.CHARS_PER_TOKEN
+
+    by_chat: dict[int, list[Message]] = defaultdict(list)
+    for message in messages:
+        by_chat[message.chat_id].append(message)
+
+    for chat_id in sorted(by_chat):
+        conversation = sorted(by_chat[chat_id], key=lambda m: m.timestamp)
+        for group in _split_on_pauses(conversation, gap_seconds):
+            for piece in _split_oversized(group, budget_chars, overlap):
+                if piece:
+                    yield build_window(piece)

msgsearch/cli.py

@@ -0,0 +1,304 @@
+"""Command line interface.
+
+All argument parsing lives here so the other modules stay importable libraries
+with no opinion about how they are invoked. One command with subcommands, rather
+than a directory of scripts, because it gives people a single thing to remember
+and a single place to find help:
+
+    msgsearch doctor               check this machine is set up correctly
+    msgsearch explore              what is in your database
+    msgsearch index                build the search index
+    msgsearch search "a query"     search it
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from . import __version__, config
+
+
+def _add_index_command(subparsers) -> None:
+    parser = subparsers.add_parser(
+        "index",
+        help="build the search index",
+        description="Read the message database, group messages into conversation "
+        "windows, embed the passages inside them, and write the index.",
+    )
+    parser.add_argument(
+        "--chat",
+        default=config.TESTBED_CHAT,
+        metavar="ID",
+        help="restrict to one conversation, by phone number or email "
+        "(default: $MSGSEARCH_TESTBED_CHAT, or all conversations)",
+    )
+    parser.add_argument("--index-dir", metavar="DIR", help="where to write the index")
+    parser.add_argument("--batch-size", type=int, default=32, metavar="N")
+    parser.add_argument(
+        "--rebuild",
+        action="store_true",
+        help="re-embed everything instead of reusing unchanged passages",
+    )
+    parser.set_defaults(handler=_run_index)
+
+
+def _run_index(args) -> int:
+    from .index import build
+
+    build(
+        chat_identifier=args.chat,
+        index_dir=args.index_dir,
+        batch_size=args.batch_size,
+        rebuild=args.rebuild,
+    )
+    return 0
+
+
+def _add_search_command(subparsers) -> None:
+    parser = subparsers.add_parser(
+        "search",
+        help="search the index",
+        description="Find conversations by meaning as well as by keyword.",
+    )
+    parser.add_argument("query", nargs="+", help="what you are looking for")
+    parser.add_argument("--limit", type=int, default=config.DEFAULT_LIMIT, metavar="N")
+    parser.add_argument(
+        "--chat", metavar="TEXT", help="restrict to matching conversations"
+    )
+    parser.add_argument(
+        "--from", dest="from_", metavar="WHO", help="restrict to a speaker"
+    )
+    parser.add_argument("--after", metavar="YYYY-MM-DD")
+    parser.add_argument("--before", metavar="YYYY-MM-DD")
+    parser.add_argument(
+        "--type",
+        metavar="TAG",
+        help="restrict by shape: credential, credential_talk, email, phone, url, address",
+    )
+    parser.add_argument("--index-dir", metavar="DIR")
+    parser.add_argument("--snippet", type=int, default=1200, metavar="CHARS")
+    parser.add_argument(
+        "--full", action="store_true", help="show the whole window, not just the match"
+    )
+    parser.add_argument(
+        "--rerank",
+        dest="no_rerank",
+        action="store_false",
+        default=not config.RERANK_ENABLED,
+        help="run the cross-encoder (off by default; it measurably hurts)",
+    )
+    parser.add_argument("--no-rerank", dest="no_rerank", action="store_true")
+    parser.add_argument("--no-dense", action="store_true", help="keyword search only")
+    parser.add_argument("--no-bm25", action="store_true", help="vector search only")
+    parser.set_defaults(handler=_run_search)
+
+
+def _run_search(args) -> int:
+    from .search import format_result, search
+
+    query = " ".join(args.query)
+    try:
+        results = search(query, args)
+    except FileNotFoundError as error:
+        print(error, file=sys.stderr)
+        return 1
+
+    if not results:
+        print("no results")
+        return 0
+
+    print(f"{len(results)} result(s) for {query!r}")
+    for position, result in enumerate(results, start=1):
+        print(format_result(position, result, args.snippet, full=args.full))
+    return 0
+
+
+def _add_contacts_command(subparsers) -> None:
+    parser = subparsers.add_parser(
+        "contacts",
+        help="map phone numbers and emails to names",
+        description="Speakers appear by name rather than by phone number, which "
+        "reads better and embeds better. Names come from an alias file you "
+        "control; this command creates and inspects it.",
+    )
+    parser.add_argument(
+        "--import",
+        dest="import_path",
+        metavar="FILE.vcf",
+        help="merge names from a vCard export (Contacts.app: File -> Export)",
+    )
+    parser.add_argument(
+        "--from-addressbook",
+        action="store_true",
+        help="merge names from macOS Contacts (needs Contacts permission)",
+    )
+    parser.add_argument(
+        "--template",
+        type=int,
+        nargs="?",
+        const=25,
+        metavar="N",
+        help="write an alias file stub for the N busiest unnamed handles, for you "
+        "to fill in (default 25)",
+    )
+    parser.add_argument(
+        "--unresolved",
+        type=int,
+        nargs="?",
+        const=25,
+        metavar="N",
+        help="list the N busiest handles that still have no name (default 25)",
+    )
+    parser.set_defaults(handler=_run_contacts)
+
+
+def _run_contacts(args) -> int:
+    from . import contacts as contacts_module
+    from .extract import connect
+
+    path = contacts_module.aliases_path()
+    mapping = contacts_module.load_mapping()
+
+    imported = {}
+    if args.import_path:
+        imported = contacts_module.read_vcard_file(Path(args.import_path))
+        print(f"read {len(imported):,} handles from {args.import_path}")
+    elif args.from_addressbook:
+        try:
+            imported = contacts_module.read_addressbook()
+        except (PermissionError, FileNotFoundError) as error:
+            print(error, file=sys.stderr)
+            return 1
+        print(f"read {len(imported):,} handles from macOS Contacts")
+
+    if imported:
+        added = {k: v for k, v in imported.items() if k not in mapping}
+        mapping.update(imported)
+        contacts_module.save(mapping, path)
+        print(f"wrote {len(mapping):,} names to {path} ({len(added):,} new)")
+
+    # Report each source separately, because "no names" has two very different
+    # causes: the Contacts permission is not granted, or it is granted and simply
+    # has no entry for these handles.
+    try:
+        from_contacts = contacts_module.read_addressbook()
+        contacts_status = f"{len(from_contacts):,} entries"
+    except PermissionError:
+        from_contacts = {}
+        contacts_status = (
+            "not permitted — grant Contacts access to the app running msgsearch "
+            "under System Settings > Privacy & Security > Contacts"
+        )
+    except (FileNotFoundError, OSError) as error:
+        from_contacts = {}
+        contacts_status = str(error)
+
+    lookup = contacts_module.Contacts.load()
+    db = connect()
+    try:
+        volumes = contacts_module.handle_volumes(db)
+    finally:
+        db.close()
+
+    named = [(h, n) for h, n in volumes if lookup.name(h)]
+    unnamed = [(h, n) for h, n in volumes if not lookup.name(h)]
+    covered = sum(n for _, n in named)
+    total = sum(n for _, n in volumes) or 1
+
+    overrides = {k: v for k, v in mapping.items() if v}
+    print(f"\nmacOS Contacts : {contacts_status}")
+    print(f"alias file     : {len(overrides):,} names in {path}")
+    print(
+        f"resolved       : {len(named):,} of {len(volumes):,} handles, "
+        f"covering {100 * covered / total:.1f}% of received messages"
+    )
+
+    if args.template:
+        stub = dict(mapping)
+        for handle, _ in unnamed[: args.template]:
+            stub.setdefault(handle, "")
+        contacts_module.save(stub, path)
+        print(
+            f"\nwrote a stub for {min(args.template, len(unnamed))} handles to {path}"
+            "\nFill in the names and re-run 'msgsearch index'. Entries left empty "
+            "are ignored."
+        )
+        return 0
+
+    if args.unresolved:
+        print(f"\nbusiest handles with no name (top {args.unresolved}):")
+        for handle, count in unnamed[: args.unresolved]:
+            print(f"  {count:>7,}  {handle}")
+        print(f"\nAdd them to {path} as a JSON object of handle -> name.")
+    return 0
+
+
+def _add_doctor_command(subparsers) -> None:
+    parser = subparsers.add_parser(
+        "doctor",
+        help="check that this machine is set up correctly",
+        description="Verify the interpreter, PyTorch, the message database, the "
+        "embedding model and the index, and say how to fix whatever is wrong. "
+        "Prints no message content.",
+    )
+    parser.set_defaults(handler=_run_doctor)
+
+
+def _run_doctor(args) -> int:
+    from .doctor import run
+
+    return run()
+
+
+def _add_explore_command(subparsers) -> None:
+    parser = subparsers.add_parser(
+        "explore",
+        help="report what is in a message database",
+        description="Structural reconnaissance: message counts, how many rows hide "
+        "their text in attributedBody, the date range, and the busiest chats. "
+        "Prints no message content.",
+    )
+    parser.set_defaults(handler=_run_explore)
+
+
+def _run_explore(args) -> int:
+    from .explore import main as explore_main
+
+    explore_main()
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="msgsearch",
+        description="Search your iMessage history by meaning, entirely on your "
+        "own machine.",
+        epilog="Run 'msgsearch <command> --help' for details of a command.",
+    )
+    parser.add_argument("--version", action="version", version=f"msgsearch {__version__}")
+
+    subparsers = parser.add_subparsers(dest="command", metavar="<command>")
+    _add_contacts_command(subparsers)
+    _add_doctor_command(subparsers)
+    _add_explore_command(subparsers)
+    _add_index_command(subparsers)
+    _add_search_command(subparsers)
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    if not getattr(args, "handler", None):
+        parser.print_help()
+        return 1
+    try:
+        return args.handler(args)
+    except KeyboardInterrupt:
+        return 130
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())