transcript_search.py

"""Search across timestamped transcripts with a lightweight inverted index.

This module answers the archival question "find where we discussed X":
given one or many transcripts made of timestamped segments (for example the
``.json`` sidecars produced by the transcription pipeline), it builds an
in-memory inverted index and returns every segment that matches a query,
each with its recording id and start/end timestamps.

Design notes
------------
* **Stdlib only.** No third-party dependencies; safe to vendor anywhere.
* **Duck-typed segments.** :meth:`TranscriptIndex.add` accepts any object
  exposing ``start``, ``end``, and ``text`` attributes, or an equivalent
  mapping with those keys — so it composes with the transcription sidecar's
  JSON output (via :func:`load_transcript_json`) without importing it.
* **AND semantics.** A multi-word query matches only segments containing
  *every* query term; results are scored by summed term frequency and
  sorted by score (descending), then by recording id and start time.

Known limitation (documented honestly)
--------------------------------------
Tokenization uses a Unicode word regex split on non-word boundaries, which
works well for whitespace-delimited languages. CJK text (Korean, Japanese,
Chinese) is only split on whitespace/punctuation, **not** morphologically
segmented — so a Korean query term matches only when the same
space-delimited token appears in the transcript. Proper CJK support would
require a morphological analyzer, which is out of scope for a
stdlib-only module.
"""

from __future__ import annotations

import json
import re
from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable, Mapping

__all__ = [
    "Match",
    "Segment",
    "TranscriptIndex",
    "load_transcript_json",
    "tokenize",
]

# Unicode-aware word tokenizer: runs of word characters (letters, digits,
# underscore).  Case is folded by the caller; punctuation is stripped by
# construction because it never matches ``\w+``.
_WORD_RE = re.compile(r"\w+", re.UNICODE)


def tokenize(text: str) -> list[str]:
    """Split *text* into lowercase word tokens.

    Uses a Unicode ``\\w+`` regex, so punctuation is dropped and tokens are
    case-folded (``"Hello, World!"`` -> ``["hello", "world"]``).

    Note:
        CJK text is split only on whitespace/punctuation boundaries — see
        the module docstring for the honest limitation statement.

    Args:
        text: Arbitrary text to tokenize.

    Returns:
        List of lowercase tokens (possibly empty).
    """
    lower = text.lower()
    if lower.isalnum():
        return [lower]
    return _WORD_RE.findall(lower)


@dataclass(frozen=True)
class Segment:
    """One timestamped chunk of a transcript.

    Attributes:
        start: Segment start time in seconds.
        end: Segment end time in seconds.
        text: Spoken text of the segment.
    """

    start: float
    end: float
    text: str


@dataclass(frozen=True)
class Match:
    """A single search hit inside a recording.

    Attributes:
        recording_id: Identifier of the recording the hit belongs to
            (as passed to :meth:`TranscriptIndex.add`).
        start: Start timestamp (seconds) of the matching segment.
        end: End timestamp (seconds) of the matching segment.
        text: Original (unnormalized) text of the matching segment.
        score: Relevance score — the summed term frequency of all query
            terms within the segment.  Higher is more relevant.
    """

    recording_id: str
    start: float
    end: float
    text: str
    score: int


@dataclass(frozen=True)
class _Entry:
    """Internal indexed segment: original fields plus its token counts."""

    recording_id: str
    start: float
    end: float
    text: str
    counts: Counter = field(compare=False)


def _read_attr(segment: Any, name: str) -> Any:
    """Fetch *name* from a duck-typed segment (attribute or mapping key).

    Args:
        segment: Object with ``start``/``end``/``text`` attributes, or a
            mapping with those keys.
        name: Field name to read.

    Returns:
        The field value.

    Raises:
        TypeError: If the segment exposes the field neither as an
            attribute nor as a mapping key.
    """
    if isinstance(segment, Mapping):
        try:
            return segment[name]
        except KeyError:
            raise TypeError(
                f"segment mapping is missing required key {name!r}"
            ) from None
    try:
        return getattr(segment, name)
    except AttributeError:
        raise TypeError(
            f"segment object is missing required attribute {name!r}"
        ) from None


class TranscriptIndex:
    """Inverted index over timestamped transcript segments.

    Add one or many recordings with :meth:`add`, then query with
    :meth:`search`.  The index is in-memory and append-only; re-adding a
    recording id simply indexes more segments under the same id.

    Example:
        >>> idx = TranscriptIndex()
        >>> idx.add("standup-01", [Segment(0.0, 4.0, "codec budget review")])
        >>> [m.recording_id for m in idx.search("codec")]
        ['standup-01']
    """

    def __init__(self) -> None:
        """Create an empty index."""
        # token -> set of entry positions in self._entries
        self._postings: dict[str, set[int]] = {}
        self._entries: list[_Entry] = []

    def __len__(self) -> int:
        """Return the number of indexed segments."""
        return len(self._entries)

    def add(self, recording_id: str, segments: Iterable[Any]) -> int:
        """Index the *segments* of one recording.

        Args:
            recording_id: Stable identifier for the recording (e.g. the
                source filename); echoed back on every :class:`Match`.
            segments: Iterable of duck-typed segments — each must expose
                ``start``, ``end`` and ``text`` as attributes (e.g.
                :class:`Segment` or the transcription sidecar's segment
                objects) or as mapping keys.

        Returns:
            The number of segments indexed from this call.

        Raises:
            TypeError: If a segment lacks one of the required fields.
        """
        added = 0
        for segment in segments:
            entry = _Entry(
                recording_id=recording_id,
                start=float(_read_attr(segment, "start")),
                end=float(_read_attr(segment, "end")),
                text=str(_read_attr(segment, "text")),
                counts=Counter(tokenize(str(_read_attr(segment, "text")))),
            )
            position = len(self._entries)
            self._entries.append(entry)
            for token in entry.counts:
                self._postings.setdefault(token, set()).add(position)
            added += 1
        return added

    def search(self, query: str) -> list[Match]:
        """Find segments containing **all** words of *query*.

        Matching is case-insensitive and punctuation-insensitive (both the
        query and the indexed text pass through :func:`tokenize`).
        Multi-word queries use AND semantics: only segments containing
        every query term are returned.

        Args:
            query: One or more words to look for.

        Returns:
            Matches sorted by ``score`` descending, then by
            ``recording_id`` and ``start`` ascending.  Empty list when
            nothing matches.

        Raises:
            ValueError: If the query is empty or contains no indexable
                words (e.g. punctuation only).
        """
        terms = tokenize(query)
        if not terms:
            raise ValueError("query must contain at least one word")

        # Intersect postings lists (AND semantics), rarest term first so
        # the working set shrinks as fast as possible.
        unique_terms = sorted(
            set(terms), key=lambda t: len(self._postings.get(t, ()))
        )
        candidates: set[int] | None = None
        for term in unique_terms:
            postings = self._postings.get(term)
            if not postings:
                return []
            candidates = (
                set(postings) if candidates is None else candidates & postings
            )
            if not candidates:
                return []

        matches = []
        for position in candidates or ():
            entry = self._entries[position]
            score = sum(entry.counts[term] for term in unique_terms)
            matches.append(
                Match(
                    recording_id=entry.recording_id,
                    start=entry.start,
                    end=entry.end,
                    text=entry.text,
                    score=score,
                )
            )
        matches.sort(key=lambda m: (-m.score, m.recording_id, m.start))
        return matches


def load_transcript_json(path: str | Path) -> list[Segment]:
    """Load segments from a transcription sidecar JSON file.

    Reads the sidecar shape ``{"segments": [{"start": .., "end": ..,
    "text": ..}, ...]}`` and returns :class:`Segment` objects ready for
    :meth:`TranscriptIndex.add`.  This mirrors the transcription
    pipeline's output format without importing that module, so the two
    features compose while remaining independent.

    Args:
        path: Path to the ``.json`` sidecar file.

    Returns:
        List of :class:`Segment` in file order.

    Raises:
        ValueError: If the file is not a JSON object with a ``"segments"``
            list, or a segment entry is missing ``start``/``end``/``text``.
        OSError: If the file cannot be read.
        json.JSONDecodeError: If the file is not valid JSON.
    """
    raw = json.loads(Path(path).read_text(encoding="utf-8"))
    if not isinstance(raw, dict) or not isinstance(raw.get("segments"), list):
        raise ValueError(
            f"{path}: expected a JSON object with a 'segments' list"
        )
    segments = []
    for i, item in enumerate(raw["segments"]):
        if not isinstance(item, dict):
            raise ValueError(f"{path}: segments[{i}] is not an object")
        try:
            segments.append(
                Segment(
                    start=float(item["start"]),
                    end=float(item["end"]),
                    text=str(item["text"]),
                )
            )
        except KeyError as exc:
            raise ValueError(
                f"{path}: segments[{i}] is missing key {exc.args[0]!r}"
            ) from None
    return segments