qmd_py.search.fts¶

def normalize_cjk_for_fts(text: str) -> str:
    """Postgres's built-in text search configs don't segment CJK runs into
    words (same gap FTS5's unicode61 tokenizer has) - space out every
    character in a CJK run so it can be treated as an adjacency phrase."""
    return _CJK_RUN_PATTERN.sub(lambda m: " " + " ".join(m.group(0)) + " ", text)

def contains_cjk(text: str) -> bool:
    """Whether `text` holds any character from the covered CJK blocks.

    Covers the common BMP blocks only, not every rare extension - see
    `_CJK_RANGES`.
    """
    return _CJK_CHAR_PATTERN.search(text) is not None

def sanitize_fts_term(term: str) -> str:
    """Strip a term down to what is safe inside a `to_tsquery` expression.

    Keeps letters, digits, apostrophes and underscores; drops everything
    else and lowercases the rest. That removal is what makes the query
    injection-safe, since the result is interpolated into tsquery syntax
    rather than bound as a parameter. Leading/trailing apostrophes are
    stripped too: in tsquery input `'` is the lexeme-quote character, so a
    term starting with one opens a quoted lexeme that never closes
    (`'n:*` is a syntax error) - and a lexeme can't start or end with an
    apostrophe anyway. Interior ones (`don't`) are legal and kept.

    Returns:
        The sanitized term, possibly empty - callers must drop empties
        rather than emit a bare operator.
    """
    sanitized = "".join(ch for ch in term if ch in ("'", "_") or _is_word_char(ch)).lower()
    return sanitized.strip("'")

def is_hyphenated_token(token: str) -> bool:
    """True for compound words like multi-agent, DEC-0054, gpt-4: starts and
    ends with a letter/number, contains at least one internal hyphen, and
    every character is a letter/number/apostrophe/hyphen."""
    if "-" not in token or not token:
        return False
    if not _is_word_char(token[0]) or not _is_word_char(token[-1]):
        return False
    return all(_is_word_char(ch) or ch in ("'", "-") for ch in token)

def is_dotted_token(token: str) -> bool:
    """True for version-like strings (2026.4.10, 3.14.0): splitting on dots
    yields >= 2 non-empty word/number/underscore parts."""
    parts = token.split(".")
    if len(parts) < 2:
        return False
    return all(p and all(_is_word_char(ch) or ch == "_" for ch in p) for p in parts)

def build_ts_query(query: str) -> str | None:
    """Parse lex query syntax into a Postgres `to_tsquery` expression.

    Supports the same surface syntax as the TS reference's FTS5-era parser:
    - Quoted phrases: "exact phrase" -> word <-> word (adjacency)
    - Negation: -term or -"phrase" -> tsquery's unary `!`
    - Hyphenated tokens: multi-agent, DEC-0054 -> phrase (not prefix) match
    - Dotted tokens: 2026.4.10 -> the whole dotted string, prefix-matched
      (Postgres's tokenizer keeps these as one lexeme, unlike SQLite FTS5,
      so - unlike the TS version - the parts are NOT split and ANDed)
    - Plain terms: term -> term:* (prefix match)

    Negative-only queries return None (no results), matching the original's
    restriction even though tsquery's `!` is technically usable standalone.
    """
    positive: list[str] = []
    negative: list[str] = []

    s = query.strip()
    i = 0
    n = len(s)

    while i < n:
        while i < n and s[i].isspace():
            i += 1
        if i >= n:
            break

        negated = s[i] == "-"
        if negated:
            i += 1

        if i < n and s[i] == '"':
            start = i + 1
            i += 1
            while i < n and s[i] != '"':
                i += 1
            phrase = s[start:i].strip()
            i += 1  # skip closing quote
            if phrase:
                words = [
                    w
                    for t in normalize_cjk_for_fts(phrase).split()
                    if (w := sanitize_fts_term(t))
                ]
                phrase_query = _to_phrase(words)
                if phrase_query:
                    (negative if negated else positive).append(phrase_query)
        else:
            start = i
            while i < n and not s[i].isspace() and s[i] != '"':
                i += 1
            term = s[start:i]

            if is_hyphenated_token(term):
                words = [sanitize_fts_term(t) for t in term.split("-") if sanitize_fts_term(t)]
                phrase_query = _to_phrase(words)
                if phrase_query:
                    (negative if negated else positive).append(phrase_query)
            elif is_dotted_token(term):
                sanitized = "".join(
                    ch for ch in term.lower() if _is_word_char(ch) or ch in (".", "_")
                )
                if sanitized:
                    (negative if negated else positive).append(f"{sanitized}:*")
            elif contains_cjk(term):
                words = [
                    w
                    for t in normalize_cjk_for_fts(term).split()
                    if (w := sanitize_fts_term(t))
                ]
                phrase_query = _to_phrase(words)
                if phrase_query:
                    (negative if negated else positive).append(phrase_query)
            else:
                sanitized = sanitize_fts_term(term)
                if sanitized:
                    (negative if negated else positive).append(f"{sanitized}:*")

    if not positive:
        return None

    result = " & ".join(positive)
    for neg in negative:
        result = f"{result} & !{neg}"
    return result

def validate_semantic_query(query: str) -> str | None:
    """Check a `vec:`/`hyde:` sub-query for lex-only syntax.

    Returns:
        An actionable message, or None if the query is fine. Only
        applied to sub-queries a caller spelled out - never to
        `expand_query()`'s generated variants; see
        `hybrid.validate_typed_queries()`.
    """
    if re.search(r'(^|\s)-[\w"]', query):
        return "Negation (-term) is not supported in vec/hyde queries. Use lex for exclusions."
    return None

def validate_lex_query(query: str) -> str | None:
    """Check a `lex:` sub-query for syntax `build_ts_query` can't honour.

    Catches a newline (each lex line is parsed separately) and an
    unmatched double quote, which would otherwise run to end of input
    and silently swallow the rest of the query as one phrase.

    Returns:
        An actionable message, or None if the query is fine.
    """
    if re.search(r"[\r\n]", query):
        return (
            "Lex queries must be a single line. Remove newline characters or split into "
            "separate lex: lines."
        )
    if query.count('"') % 2 == 1:
        return 'Lex query has an unmatched double quote ("). Add the closing quote or remove it.'
    return None

async def update_document_search_vector(session: AsyncSession, document_id: int) -> None:
    """Recompute a document's `search_vector` from its current title,
    display path, and body. Not a Postgres `GENERATED` column (the body
    lives in `Content.doc`, joined via hash - a generated column's
    expression can't reach across tables) - called wherever a document's
    title/path/body changes (see store.py's insert_document/update_document).

    Weights mirror the original SQLite FTS5 `bm25(documents_fts, 1.5, 4.0,
    1.0)` weighting (filepath, title, body): title matters most, filepath
    next, body is the baseline - setweight()'s A/B/C zones encode that
    order; the actual ratio is applied at query time via ts_rank's weights
    array (see FTS_RANK_WEIGHTS below).
    """
    row = (
        await session.execute(
            select(
                col(Document.title), col(Document.path), col(Content.doc), col(Collection.name)
            )
            .join(Content, col(Content.hash) == col(Document.hash))
            .join(Collection, col(Collection.id) == col(Document.collection_id))
            .where(col(Document.id) == document_id, col(Document.active))
        )
    ).one_or_none()

    if row is None:
        await session.execute(
            sa_update(Document).where(col(Document.id) == document_id).values(search_vector=None)
        )
        return

    title, path, body, collection_name = row
    # setweight()'s zone argument is Postgres's single-byte "char" type, not
    # varchar/text - passed as a bound VARCHAR parameter, no setweight
    # overload matches. Embed it as a SQL literal instead.
    title_tsv = func.setweight(
        func.to_tsvector("english", normalize_cjk_for_fts(title)), literal_column("'A'")
    )
    path_tsv = func.setweight(
        func.to_tsvector("english", normalize_cjk_for_fts(f"{collection_name}/{path}")),
        literal_column("'B'"),
    )
    body_tsv = func.setweight(
        func.to_tsvector("english", normalize_cjk_for_fts(body)), literal_column("'C'")
    )

    await session.execute(
        sa_update(Document)
        .where(col(Document.id) == document_id)
        .values(search_vector=title_tsv.op("||")(path_tsv).op("||")(body_tsv))
    )

def get_docid(hash_: str) -> str:
    """The short document id shown in output and accepted as `#abc123`.

    Six hex chars of the content hash. Short enough to collide in
    principle; `find_document()` resolves collisions deterministically
    rather than uniquely.
    """
    return hash_[:6]

async def get_context_for_path(
    session: AsyncSession, user: CurrentUser, collection_id: int, path: str
) -> str | None:
    """Hierarchical context: the user's global context (`context add /`)
    first, then all matching per-path CollectionContext rows for this
    collection, shortest/most-general prefix first - port of the TS
    reference's `getContextForPath`."""
    contexts: list[str] = []

    global_context = (
        await session.execute(select(col(User.global_context)).where(col(User.id) == user.id))
    ).scalar_one_or_none()
    if global_context:
        contexts.append(global_context)

    rows = (
        await session.execute(
            select(col(CollectionContext.path_prefix), col(CollectionContext.context)).where(
                col(CollectionContext.collection_id) == collection_id
            )
        )
    ).all()

    normalized_path = path if path.startswith("/") else f"/{path}"
    matching = []
    for prefix, context in rows:
        normalized_prefix = prefix if prefix.startswith("/") else f"/{prefix}"
        if normalized_path.startswith(normalized_prefix):
            matching.append((normalized_prefix, context))
    matching.sort(key=lambda pc: len(pc[0]))
    contexts.extend(context for _, context in matching)

    return "\n\n".join(contexts) if contexts else None

async def search_fts(
    session: AsyncSession,
    user: CurrentUser,
    query: str,
    limit: int = 20,
    collection_name: str | None = None,
) -> list[SearchResult]:
    """Rank documents lexically with Postgres full-text search.

    Args:
        query: Lex syntax - prefix terms, `"quoted phrases"`, `-negation`.
            Parsed by `build_ts_query()`.
        limit: Maximum hits.
        collection_name: Restrict to one collection; None searches every
            collection the user can read.

    Returns:
        Hits ordered by descending rank, `source="fts"`. Empty when the
        query has no positive term (a negation-only query), when no
        collection is accessible, or when nothing matches - none of which
        is an error.

    Note:
        Fetches one context per hit (an N+1). Bounded by `limit`, so it
        is a fixed small cost per search rather than a scaling one.
    """
    ts_query = build_ts_query(query)
    if ts_query is None:
        return []

    collection_ids = await resolve_collection_ids(session, user, collection_name)
    if not collection_ids:
        return []

    tsq_expr = func.to_tsquery("english", ts_query)
    weights: ColumnElement[Any] = literal_column(f"'{FTS_RANK_WEIGHTS}'::float4[]")
    rank_expr = func.ts_rank(weights, col(Document.search_vector), tsq_expr, 32)

    stmt = (
        select(
            col(Document.collection_id),
            col(Document.path),
            col(Document.title),
            col(Document.hash),
            col(Document.modified_at),
            col(Content.doc).label("body"),
            rank_expr.label("rank_score"),
        )
        .select_from(Document)
        .join(Content, col(Content.hash) == col(Document.hash))
        .where(
            col(Document.active),
            col(Document.search_vector).op("@@")(tsq_expr),
            col(Document.collection_id).in_(collection_ids),
        )
        .order_by(rank_expr.desc())
        .limit(limit)
    )

    rows = (await session.execute(stmt)).all()
    if not rows:
        return []

    collection_names = await collection_names_by_id(
        session, {row.collection_id for row in rows}
    )

    results = []
    for row in rows:
        name = collection_names[row.collection_id]
        context = await get_context_for_path(session, user, row.collection_id, row.path)
        results.append(
            SearchResult(
                filepath=f"marq://{name}/{row.path}",
                display_path=f"{name}/{row.path}",
                title=row.title,
                hash=row.hash,
                docid=get_docid(row.hash),
                collection_name=name,
                modified_at=row.modified_at,
                body_length=len(row.body),
                body=row.body,
                context=context,
                score=row.rank_score,
                source="fts",
            )
        )
    return results