msg-summarizer 0.1.0__py3-none-any.whl

msg_summarizer/contacts.py

@@ -0,0 +1,76 @@
+"""Resolve iMessage handles to real names via the macOS AddressBook (Contacts).
+
+Reads the AddressBook SQLite stores directly (under ~/Library/Application
+Support/AddressBook), which existing Full Disk Access already covers — no extra
+permission prompt. Fails soft: returns {} if Contacts can't be read, so the
+rest of the app still works with raw handles.
+"""
+
+import glob
+import pathlib
+import re
+import sqlite3
+
+_AB_DIR = pathlib.Path.home() / "Library" / "Application Support" / "AddressBook"
+
+
+def normalize_handle(handle: str) -> str:
+    """Reduce a handle to a match key: lowercased email, or a phone's last 10 digits."""
+    s = (handle or "").strip().lower()
+    if "@" in s:
+        return s
+    digits = re.sub(r"\D", "", s)
+    return digits[-10:] if len(digits) >= 10 else (digits or s)
+
+
+def _addressbook_paths() -> list[str]:
+    paths = glob.glob(str(_AB_DIR / "Sources" / "*" / "AddressBook-v22.abcddb"))
+    top = _AB_DIR / "AddressBook-v22.abcddb"
+    if top.exists():
+        paths.append(str(top))
+    return paths
+
+
+def _record_name(first, last, org) -> str:
+    parts = [p for p in (first, last) if p]
+    if parts:
+        return " ".join(parts)
+    return org or ""
+
+
+def _harvest(conn: sqlite3.Connection, mapping: dict, source: str) -> None:
+    for table, col in (("ZABCDPHONENUMBER", "ZFULLNUMBER"), ("ZABCDEMAILADDRESS", "ZADDRESS")):
+        query = (
+            f"SELECT x.{col}, r.Z_PK, r.ZFIRSTNAME, r.ZLASTNAME, r.ZORGANIZATION "
+            f"FROM {table} x JOIN ZABCDRECORD r ON x.ZOWNER = r.Z_PK"
+        )
+        try:
+            rows = conn.execute(query).fetchall()
+        except sqlite3.Error:
+            continue  # schema variation in this source; skip this table
+        for value, pk, first, last, org in rows:
+            if not value:
+                continue
+            name = _record_name(first, last, org)
+            if not name:
+                continue
+            mapping.setdefault(normalize_handle(value), (f"ab:{source}:{pk}", name))
+
+
+def load_contacts() -> dict[str, tuple[str, str]]:
+    """Map normalized handle -> (contact_key, display_name).
+
+    The contact_key is stable per contact card, so a person's phone and email
+    collapse to the same key (and thus the same merged Person).
+    """
+    mapping: dict[str, tuple[str, str]] = {}
+    for i, path in enumerate(_addressbook_paths()):
+        try:
+            conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
+        except sqlite3.OperationalError:
+            continue
+        try:
+            _harvest(conn, mapping, source=str(i))
+        finally:
+            conn.close()
+    return mapping

msg_summarizer/db.py

@@ -0,0 +1,401 @@
+import datetime
+import pathlib
+import re
+import sqlite3
+
+from typedstream.stream import TypedStreamReader
+
+from .contacts import load_contacts, normalize_handle
+from .models import Chat, Conversation, Message, Person
+
+DB_PATH = pathlib.Path.home() / "Library" / "Messages" / "chat.db"
+
+# Seconds between Unix epoch (1970-01-01) and Apple Core Data epoch (2001-01-01)
+_APPLE_EPOCH_OFFSET = 978307200
+
+# Base query; date-bound clauses and ORDER BY are appended by _fetch().
+# Messages on macOS Ventura+ usually leave `text` NULL and store the body in
+# `attributedBody`, so we select both and resolve in _rows_to_messages().
+_MESSAGES_BASE = """
+SELECT
+    m.ROWID           AS message_rowid,
+    m.date            AS apple_date,
+    h.id              AS handle_id,
+    m.is_from_me      AS is_from_me,
+    m.text            AS text,
+    m.attributedBody  AS attributed_body
+FROM message m
+LEFT JOIN handle h ON m.handle_id = h.ROWID
+JOIN chat_message_join cmj ON m.ROWID = cmj.message_id
+JOIN chat c ON cmj.chat_id = c.ROWID
+WHERE
+    (
+        LOWER(h.id)                 LIKE :contact_pattern
+        OR LOWER(c.display_name)    LIKE :contact_pattern
+        OR LOWER(c.chat_identifier) LIKE :contact_pattern
+    )
+    AND (
+        (m.text IS NOT NULL AND m.text != '')
+        OR m.attributedBody IS NOT NULL
+    )
+"""
+
+
+def _apple_ts(dt: datetime.datetime) -> int:
+    """Convert a datetime to an Apple nanosecond timestamp."""
+    return int((dt.timestamp() - _APPLE_EPOCH_OFFSET) * 1_000_000_000)
+
+
+def _from_apple_ts(apple_ts: int) -> datetime.datetime:
+    unix = (apple_ts / 1_000_000_000) + _APPLE_EPOCH_OFFSET
+    return datetime.datetime.fromtimestamp(unix, tz=datetime.timezone.utc)
+
+
+def _connect(db_path: pathlib.Path) -> sqlite3.Connection:
+    """Open *db_path* read-only, or raise a helpful PermissionError."""
+    try:
+        return sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
+    except sqlite3.OperationalError as exc:
+        raise PermissionError(
+            f"Cannot open {db_path}.\n"
+            "Grant Full Disk Access to your terminal app in:\n"
+            "  System Settings > Privacy & Security > Full Disk Access"
+        ) from exc
+
+
+def _decode_attributed_body(blob: bytes | None) -> str | None:
+    """Extract the message text from a Messages `attributedBody` blob.
+
+    macOS Ventura+ leaves the `text` column NULL and stores the body as a
+    serialized NSAttributedString (Apple's "typedstream" format). The message
+    text is the first string payload in the stream, ahead of the attribute-run
+    metadata (e.g. ``__kIMMessagePartAttributeName``).
+    """
+    if not blob:
+        return None
+    try:
+        for event in TypedStreamReader.from_data(blob):
+            if isinstance(event, bytes):
+                return event.decode("utf-8", errors="replace")
+            if isinstance(event, str):
+                return event
+    except Exception:
+        return None
+    return None
+
+
+def _rows_to_messages(rows: list[tuple]) -> list[Message]:
+    messages: list[Message] = []
+    for rowid, apple_date, handle_id, is_from_me, text, attributed_body in rows:
+        resolved = text if text else _decode_attributed_body(attributed_body)
+        if not resolved or not resolved.strip():
+            continue
+        messages.append(
+            Message(
+                timestamp=_from_apple_ts(apple_date),
+                sender="me" if is_from_me else (handle_id or "them"),
+                text=resolved.strip(),
+                is_from_me=bool(is_from_me),
+                rowid=rowid,
+            )
+        )
+    return messages
+
+
+_LIST_CONTACTS_QUERY = """
+SELECT
+    h.id                         AS handle_id,
+    COALESCE(c.display_name, '') AS display_name,
+    COUNT(DISTINCT m.ROWID)      AS msg_count
+FROM handle h
+JOIN chat_handle_join chj ON h.ROWID = chj.handle_id
+JOIN chat c ON chj.chat_id = c.ROWID
+LEFT JOIN chat_message_join cmj ON c.ROWID = cmj.chat_id
+LEFT JOIN message m ON cmj.message_id = m.ROWID
+GROUP BY h.id, c.display_name
+ORDER BY msg_count DESC
+LIMIT :limit
+"""
+
+
+def list_contacts(
+    db_path: pathlib.Path = DB_PATH,
+    limit: int = 50,
+) -> list[tuple[str, str, int]]:
+    """Return (handle_id, display_name, msg_count) for contacts in the DB."""
+    conn = _connect(db_path)
+    with conn:
+        rows = conn.execute(_LIST_CONTACTS_QUERY, {"limit": limit}).fetchall()
+
+    return [(handle_id, display_name, msg_count) for handle_id, display_name, msg_count in rows]
+
+
+def fetch_messages(
+    contact: str,
+    db_path: pathlib.Path = DB_PATH,
+    start: datetime.datetime | None = None,
+    end: datetime.datetime | None = None,
+) -> Conversation:
+    """Fetch messages matching *contact*, optionally bounded by *start*/*end* (UTC).
+
+    With no bounds, returns the full message history.
+    """
+    query = _MESSAGES_BASE
+    params: dict[str, object] = {"contact_pattern": f"%{contact.lower()}%"}
+    if start is not None:
+        query += "    AND m.date >= :start_ts\n"
+        params["start_ts"] = _apple_ts(start)
+    if end is not None:
+        query += "    AND m.date <= :end_ts\n"
+        params["end_ts"] = _apple_ts(end)
+    query += "ORDER BY m.date ASC"
+
+    conn = _connect(db_path)
+    with conn:
+        rows = conn.execute(query, params).fetchall()
+
+    return Conversation(contact=contact, messages=_rows_to_messages(rows))
+
+
+# --- Chat- and person-oriented model -------------------------------------
+
+_CHATS_QUERY = """
+SELECT
+    c.ROWID                        AS chat_id,
+    c.chat_identifier              AS identifier,
+    COALESCE(c.display_name, '')   AS display_name,
+    COUNT(DISTINCT cmj.message_id) AS msg_count,
+    MAX(m.date)                    AS last_date
+FROM chat c
+LEFT JOIN chat_message_join cmj ON c.ROWID = cmj.chat_id
+LEFT JOIN message m ON cmj.message_id = m.ROWID
+GROUP BY c.ROWID
+"""
+
+_CHAT_PARTICIPANTS_QUERY = """
+SELECT chj.chat_id, h.id
+FROM chat_handle_join chj
+JOIN handle h ON chj.handle_id = h.ROWID
+"""
+
+_EPOCH_MIN = datetime.datetime.min.replace(tzinfo=datetime.timezone.utc)
+
+# How many resolved names to show before collapsing to "& N more" in a group name.
+_GROUP_NAME_MAX = 4
+
+
+def _load_raw_chats(db_path: pathlib.Path) -> tuple[list[tuple], dict[int, list[str]]]:
+    """Load raw chat rows + a chat_id -> [handle ids] map. Shared by list_chats/people."""
+    conn = _connect(db_path)
+    with conn:
+        chat_rows = conn.execute(_CHATS_QUERY).fetchall()
+        part_rows = conn.execute(_CHAT_PARTICIPANTS_QUERY).fetchall()
+    participants: dict[int, list[str]] = {}
+    for chat_id, handle_id in part_rows:
+        participants.setdefault(chat_id, []).append(handle_id)
+    return chat_rows, participants
+
+
+def _build_people(
+    chat_rows: list[tuple],
+    participants: dict[int, list[str]],
+    contacts: dict[str, tuple[str, str]],
+) -> list[Person]:
+    """Derive Person objects from raw chats, merging handles via Contacts."""
+    people: dict[str, Person] = {}
+    for chat_id, _identifier, _display_name, msg_count, last_date in chat_rows:
+        parts = participants.get(chat_id, [])
+        is_group = len(parts) > 1
+        last = _from_apple_ts(last_date) if last_date else None
+        for handle_id in parts:
+            norm = normalize_handle(handle_id)
+            entry = contacts.get(norm)
+            key, name = entry if entry else (norm, handle_id)
+            person = people.get(key)
+            if person is None:
+                person = Person(key=key, name=name)
+                people[key] = person
+            if handle_id not in person.handles:
+                person.handles.append(handle_id)
+            if is_group:
+                person.group_chat_ids.append(chat_id)
+            else:
+                person.dm_chat_ids.append(chat_id)
+                person.message_count += msg_count or 0
+            if last and (person.last_date is None or last > person.last_date):
+                person.last_date = last
+    out = list(people.values())
+    out.sort(key=lambda p: p.last_date or _EPOCH_MIN, reverse=True)
+    return out
+
+
+def _format_group_name(names: list[str]) -> str:
+    if not names:
+        return "Group"
+    if len(names) <= _GROUP_NAME_MAX:
+        return ", ".join(names)
+    return ", ".join(names[:_GROUP_NAME_MAX]) + f" & {len(names) - _GROUP_NAME_MAX} more"
+
+
+def _build_chats(
+    chat_rows: list[tuple],
+    participants: dict[int, list[str]],
+    handle_to_person: dict[str, Person],
+) -> list[Chat]:
+    chats: list[Chat] = []
+    for chat_id, identifier, display_name, msg_count, last_date in chat_rows:
+        parts = participants.get(chat_id, [])
+        is_group = len(parts) > 1
+        if display_name:
+            name = display_name
+        elif not is_group and parts:
+            person = handle_to_person.get(parts[0])
+            name = person.name if person else parts[0]
+        elif is_group:
+            resolved = sorted(
+                (handle_to_person[h].name if h in handle_to_person else h) for h in parts
+            )
+            name = _format_group_name(resolved)
+        else:
+            name = identifier or "Unknown"
+        chats.append(
+            Chat(
+                chat_id=chat_id,
+                identifier=identifier or "",
+                display_name=display_name or "",
+                name=name,
+                is_group=is_group,
+                participants=parts,
+                message_count=msg_count or 0,
+                last_date=_from_apple_ts(last_date) if last_date else None,
+            )
+        )
+    chats.sort(key=lambda c: c.last_date or _EPOCH_MIN, reverse=True)
+    return chats
+
+
+def list_chats(db_path: pathlib.Path = DB_PATH) -> list[Chat]:
+    """Return every chat thread (1:1 and group), newest activity first.
+
+    Names are resolved: 1:1s use the matched Contact's name; unnamed groups use
+    the alphabetized list of participant Contact names.
+    """
+    chat_rows, participants = _load_raw_chats(db_path)
+    contacts = load_contacts()
+    people = _build_people(chat_rows, participants, contacts)
+    handle_to_person = {h: p for p in people for h in p.handles}
+    return _build_chats(chat_rows, participants, handle_to_person)
+
+
+def _fetch_by_chats(
+    chat_ids: list[int],
+    label: str,
+    db_path: pathlib.Path,
+    start: datetime.datetime | None,
+    end: datetime.datetime | None,
+) -> Conversation:
+    if not chat_ids:
+        return Conversation(contact=label, messages=[])
+
+    placeholders = ",".join("?" for _ in chat_ids)
+    query = f"""
+        SELECT m.ROWID, m.date, h.id, m.is_from_me, m.text, m.attributedBody
+        FROM message m
+        LEFT JOIN handle h ON m.handle_id = h.ROWID
+        JOIN chat_message_join cmj ON m.ROWID = cmj.message_id
+        WHERE cmj.chat_id IN ({placeholders})
+          AND ((m.text IS NOT NULL AND m.text != '') OR m.attributedBody IS NOT NULL)
+    """
+    params: list[object] = list(chat_ids)
+    if start is not None:
+        query += " AND m.date >= ?"
+        params.append(_apple_ts(start))
+    if end is not None:
+        query += " AND m.date <= ?"
+        params.append(_apple_ts(end))
+    query += " ORDER BY m.date ASC"
+
+    conn = _connect(db_path)
+    with conn:
+        rows = conn.execute(query, params).fetchall()
+    return Conversation(contact=label, messages=_rows_to_messages(rows))
+
+
+def fetch_chat_messages(
+    chat_id: int,
+    label: str | None = None,
+    db_path: pathlib.Path = DB_PATH,
+    start: datetime.datetime | None = None,
+    end: datetime.datetime | None = None,
+) -> Conversation:
+    """Fetch the messages of a single chat thread."""
+    return _fetch_by_chats([chat_id], label or str(chat_id), db_path, start, end)
+
+
+def fetch_person_messages(
+    person: Person,
+    db_path: pathlib.Path = DB_PATH,
+    start: datetime.datetime | None = None,
+    end: datetime.datetime | None = None,
+) -> Conversation:
+    """Fetch a person's direct (1:1) messages, merged across their DM chats."""
+    return _fetch_by_chats(person.dm_chat_ids, person.name, db_path, start, end)
+
+
+def list_people(db_path: pathlib.Path = DB_PATH) -> list[Person]:
+    """Derive people from chats, merging handles via Contacts and normalized ids.
+
+    A person's headline is their 1:1 DM(s); group chats they're in are tracked
+    separately so the UI can link to them.
+    """
+    chat_rows, participants = _load_raw_chats(db_path)
+    return _build_people(chat_rows, participants, load_contacts())
+
+
+def list_conversations(db_path: pathlib.Path = DB_PATH) -> list[dict]:
+    """Unified Chats list: one row per person DM (merged across handles) plus
+    one row per group chat. Sorted newest first.
+
+    Each row is shaped as a click target:
+      * person DM:  {"kind": "person", "key": ..., "name", ...}
+      * group:      {"kind": "chat",   "chat_id": ..., "name", ...}
+    """
+    chat_rows, participants = _load_raw_chats(db_path)
+    contacts = load_contacts()
+    people = _build_people(chat_rows, participants, contacts)
+    handle_to_person = {h: p for p in people for h in p.handles}
+    chats = _build_chats(chat_rows, participants, handle_to_person)
+
+    rows: list[dict] = []
+    for person in people:
+        if not person.dm_chat_ids:
+            continue
+        rows.append(
+            {
+                "kind": "person",
+                "key": person.key,
+                "name": person.name,
+                "is_group": False,
+                "message_count": person.message_count,
+                "last_date": person.last_date.isoformat() if person.last_date else None,
+                "participant_count": 1,
+                "handle_count": len(person.handles),
+            }
+        )
+    for chat in chats:
+        if not chat.is_group:
+            continue
+        rows.append(
+            {
+                "kind": "chat",
+                "chat_id": chat.chat_id,
+                "name": chat.name,
+                "is_group": True,
+                "message_count": chat.message_count,
+                "last_date": chat.last_date.isoformat() if chat.last_date else None,
+                "participant_count": len(chat.participants),
+                "handle_count": len(chat.participants),
+            }
+        )
+    rows.sort(key=lambda r: r["last_date"] or "", reverse=True)
+    return rows

msg_summarizer/embeddings.py

@@ -0,0 +1,152 @@
+"""On-device embedding for semantic search over message windows.
+
+Uses BAAI/bge-small-en-v1.5 via fastembed (ONNX runtime, no PyTorch).
+The model file is downloaded once on first use and cached locally.
+"""
+
+from dataclasses import dataclass
+from datetime import datetime
+
+import numpy as np
+
+from .models import Message
+
+_MODEL_NAME = "BAAI/bge-small-en-v1.5"
+_DIM = 384
+_WINDOW_SIZE = 5
+
+_model = None
+
+
+def _load_model():
+    """Lazy-load the embedding model on first use (slow once; cached after)."""
+    global _model
+    if _model is None:
+        from fastembed import TextEmbedding  # heavy import; defer
+        _model = TextEmbedding(model_name=_MODEL_NAME)
+    return _model
+
+
+def embedding_dim() -> int:
+    return _DIM
+
+
+@dataclass
+class Window:
+    first_rowid: int
+    last_rowid: int
+    first_ts: datetime
+    last_ts: datetime
+    text: str
+
+
+def window_messages(messages: list[Message], n: int = _WINDOW_SIZE) -> list[Window]:
+    """Group *messages* (in chronological order) into windows of up to n messages each.
+
+    Each window is one semantic unit — short context bundles tend to embed better
+    than single messages like "lol" or "ok".
+    """
+    out: list[Window] = []
+    for i in range(0, len(messages), n):
+        chunk = messages[i : i + n]
+        if not chunk:
+            continue
+        text = "\n".join(
+            f"{'me' if m.is_from_me else m.sender}: {m.text}" for m in chunk
+        )
+        out.append(
+            Window(
+                first_rowid=chunk[0].rowid,
+                last_rowid=chunk[-1].rowid,
+                first_ts=chunk[0].timestamp,
+                last_ts=chunk[-1].timestamp,
+                text=text,
+            )
+        )
+    return out
+
+
+def embed_documents(texts: list[str], batch_size: int = 64) -> np.ndarray:
+    """Embed document texts. Returns N x D float32 (L2-normalized)."""
+    if not texts:
+        return np.zeros((0, _DIM), dtype=np.float32)
+    model = _load_model()
+    vecs = list(model.embed(texts, batch_size=batch_size))
+    return np.stack([np.asarray(v, dtype=np.float32) for v in vecs])
+
+
+def embed_query(text: str) -> np.ndarray:
+    """Embed a query. Returns a 1-D float32 vector (L2-normalized)."""
+    model = _load_model()
+    # query_embed applies the BGE query-side prefix expected by this model.
+    return np.asarray(next(model.query_embed([text])), dtype=np.float32)
+
+
+def cosine_topk(
+    query_vec: np.ndarray, doc_matrix: np.ndarray, k: int = 20
+) -> list[tuple[int, float]]:
+    """Top-K cosine similarities (vectors assumed L2-normalized)."""
+    if doc_matrix.size == 0:
+        return []
+    scores = doc_matrix @ query_vec  # (N,)
+    k = min(k, len(scores))
+    top = np.argpartition(-scores, k - 1)[:k]
+    top = top[np.argsort(-scores[top])]
+    return [(int(i), float(scores[i])) for i in top]
+
+
+# Per-target in-memory cache: (target_kind, target_id) -> (meta, matrix, last_window_id)
+_index_cache: dict[tuple[str, str], tuple[list[dict], np.ndarray, int]] = {}
+
+
+def get_index_matrix(target_kind: str, target_id: str):
+    """Return (meta, matrix) for *target*, reloading from disk if it's grown."""
+    from .store import get_index_progress, load_windows  # avoid cycle at import
+
+    _, last_window_id = get_index_progress(target_kind, target_id)
+    key = (target_kind, target_id)
+    cached = _index_cache.get(key)
+    if cached is not None and cached[2] == last_window_id:
+        return cached[0], cached[1]
+    meta, vec_bytes = load_windows(target_kind, target_id)
+    matrix = np.frombuffer(vec_bytes, dtype=np.float32).reshape(-1, _DIM).copy()
+    _index_cache[key] = (meta, matrix, last_window_id)
+    return meta, matrix
+
+
+def invalidate_index(target_kind: str, target_id: str) -> None:
+    _index_cache.pop((target_kind, target_id), None)
+
+
+def search_index(
+    target_kind: str,
+    target_id: str,
+    query: str,
+    k: int = 20,
+    start_ts: str | None = None,
+    end_ts: str | None = None,
+) -> list[dict]:
+    """Return top-K window hits for *query* in *target*'s index.
+
+    If *start_ts*/*end_ts* (ISO 8601 UTC strings) are given, only windows whose
+    time range overlaps [start_ts, end_ts] are eligible — so callers can scope
+    retrieval to the same date range the user picked elsewhere in the UI.
+    """
+    meta, matrix = get_index_matrix(target_kind, target_id)
+    if len(meta) == 0:
+        return []
+    q_vec = embed_query(query)
+    scores = matrix @ q_vec
+    if start_ts is not None or end_ts is not None:
+        for i, m in enumerate(meta):
+            if start_ts is not None and m["last_ts"] < start_ts:
+                scores[i] = -np.inf
+            if end_ts is not None and m["first_ts"] > end_ts:
+                scores[i] = -np.inf
+    eligible = int(np.sum(scores > -1e30))
+    if eligible == 0:
+        return []
+    k = min(k, eligible)
+    top = np.argpartition(-scores, k - 1)[:k]
+    top = top[np.argsort(-scores[top])]
+    return [{**meta[i], "score": float(scores[i])} for i in top]