contextual-engine 0.1.0__py3-none-any.whl

contextual/docs/retrieval.py

@@ -0,0 +1,654 @@
+"""Docs retrieval layer for Contextual Phase 2.
+
+Three public functions consumed directly by the MCP tools:
+
+  docs_search()       — hybrid BM25 + vector search over doc chunks
+  docs_get_section()  — fetch exact section by 5-tier heading path resolution
+  docs_list_files()   — enumerate indexed doc files with optional heading tree
+
+All functions accept a raw ``sqlite3.Connection`` and an ``Embedder`` so they
+are fully testable without a live MCP context.
+
+FTS5 BM25 column weights for docs:
+  (heading_path ×3, content ×1)
+  Technical heading terms are rare/high-IDF → boost them hard.
+
+RRF k = 60 (standard), bm25_weight = 0.4, dense_weight = 0.6
+(prose benefits more from semantic similarity than exact term match).
+"""
+from __future__ import annotations
+
+import fnmatch
+import json
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from sqlite3 import Connection
+from typing import Any
+
+import structlog
+from rapidfuzz import fuzz
+
+from contextual.embedding.embedder import CodeEmbedder
+
+logger = structlog.get_logger(__name__)
+
+# ---------------------------------------------------------------------------
+# RRF / scoring constants
+# ---------------------------------------------------------------------------
+RRF_K: int = 60
+BM25_WEIGHT: float = 0.4
+DENSE_WEIGHT: float = 0.6
+DOC_SOURCE_TYPES: tuple[str, ...] = ("docs", "docs_code")
+
+# Fuzzy heading match threshold (0-100)
+FUZZY_MATCH_THRESHOLD: int = 72
+
+# ---------------------------------------------------------------------------
+# Result data models
+# ---------------------------------------------------------------------------
+
+@dataclass(slots=True)
+class DocSearchHit:
+    """A single search result from docs_search."""
+    chunk_id: int
+    file_path: str
+    content: str
+    heading_path: str
+    heading_level: int
+    chunk_type: str           # "docs" | "docs_code"
+    start_line: int
+    end_line: int
+    score: float
+    bm25_rank: int | None = None
+    vector_rank: int | None = None
+
+
+@dataclass(slots=True)
+class DocSection:
+    """A resolved document section from docs_get_section."""
+    file_path: str
+    heading_path: str
+    heading_level: int
+    content: str
+    start_line: int
+    end_line: int
+    subsections: list["DocSection"] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class DocFileInfo:
+    """Metadata about an indexed doc file from docs_list_files."""
+    file_path: str
+    chunk_count: int
+    headings: list[str] = field(default_factory=list)   # only when include_headings=True
+
+
+# ---------------------------------------------------------------------------
+# Source filter helpers
+# ---------------------------------------------------------------------------
+
+def _build_source_filter_clause(
+    source_filter: dict[str, Any] | None,
+) -> tuple[str, list[Any]]:
+    """Convert source_filter dict into a SQL WHERE fragment + params.
+
+    Supported keys:
+      repo         list[str]  — match file_path prefix
+      file         list[str]  — exact or glob file_path match
+      heading      str        — SQL LIKE pattern on heading_path
+      source_type  list[str]  — restrict to 'docs', 'docs_code', or both
+    """
+    clauses: list[str] = []
+    params: list[Any] = []
+
+    # Always restrict to doc source types
+    clauses.append("source_type IN ('docs','docs_code')")
+
+    if not source_filter:
+        return " AND ".join(clauses), params
+
+    if st := source_filter.get("source_type"):
+        safe = [s for s in st if s in DOC_SOURCE_TYPES]
+        if safe:
+            ph = ",".join("?" * len(safe))
+            clauses[-1] = f"source_type IN ({ph})"   # override default
+            params.extend(safe)
+
+    if repos := source_filter.get("repo"):
+        repo_clauses = " OR ".join("file_path LIKE ?" for _ in repos)
+        clauses.append(f"({repo_clauses})")
+        params.extend(f"{r}/%" for r in repos)
+
+    if files := source_filter.get("file"):
+        file_clauses = " OR ".join("file_path = ?" for _ in files)
+        clauses.append(f"({file_clauses})")
+        params.extend(files)
+
+    if heading := source_filter.get("heading"):
+        clauses.append("heading_path LIKE ?")
+        params.append(f"%{heading}%")
+
+    return " AND ".join(clauses), params
+
+
+# ---------------------------------------------------------------------------
+# BM25 search (FTS5)
+# ---------------------------------------------------------------------------
+
+def _bm25_docs_search(
+    conn: Connection,
+    query: str,
+    k: int,
+    filter_clause: str,
+    filter_params: list[Any],
+) -> list[tuple[int, float]]:
+    """Run FTS5 BM25 search over docs chunks.
+
+    Returns list of (chunk_id, bm25_score) sorted by score descending.
+    BM25 scores from FTS5 are negative — we negate to get positive rank values.
+    Column weights: heading_path ×3, content ×1.
+    """
+    # Sanitise query: strip FTS5 special chars that cause parse errors
+    safe_q = re.sub(r'["\'\(\)\*\:\^]', " ", query).strip()
+    if not safe_q:
+        return []
+
+    sql = f"""
+        SELECT c.id, -bm25(fts_chunks) AS score
+        FROM fts_chunks
+        JOIN chunks c ON fts_chunks.chunk_id = c.id
+        WHERE fts_chunks MATCH ? AND {filter_clause}
+        ORDER BY score DESC
+        LIMIT ?
+    """
+    try:
+        rows = conn.execute(sql, [safe_q, *filter_params, k * 3]).fetchall()
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("FTS5 docs search failed", query=safe_q, error=str(exc))
+        return []
+
+    return [(row[0], float(row[1])) for row in rows]
+
+
+# ---------------------------------------------------------------------------
+# Vector search
+# ---------------------------------------------------------------------------
+
+def _vector_docs_search(
+    conn: Connection,
+    query_embedding: list[float],
+    k: int,
+    filter_clause: str,
+    filter_params: list[Any],
+) -> list[tuple[int, float]]:
+    """KNN search on vec_chunks, post-filtered to docs source_type.
+
+    vec0 doesn't support filter pushdown via JOIN, so we over-fetch (k×5)
+    and apply the filter in Python — standard workaround for sqlite-vec 0.1.x.
+    """
+    import json as _json
+
+    emb_json = _json.dumps(query_embedding)
+    try:
+        # Over-fetch to survive post-filter attrition (W3: cap at 500)
+        fetch_count = min(k * 5, 500)
+        raw = conn.execute(
+            """
+            SELECT chunk_id, distance
+            FROM vec_chunks
+            WHERE embedding MATCH ?
+            ORDER BY distance
+            LIMIT ?
+            """,
+            (emb_json, fetch_count),
+        ).fetchall()
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("Vector docs search failed", error=str(exc))
+        return []
+
+    if not raw:
+        return []
+
+    # Post-filter: restrict to doc chunks
+    raw_ids = [r[0] for r in raw]
+    dist_map = {r[0]: r[1] for r in raw}
+
+    ph = ",".join("?" * len(raw_ids))
+    filtered = conn.execute(
+        f"SELECT id FROM chunks WHERE id IN ({ph}) AND {filter_clause}",
+        [*raw_ids, *filter_params],
+    ).fetchall()
+    filtered_ids = {row[0] for row in filtered}
+
+    results = [
+        (cid, 1.0 - dist_map[cid])     # convert distance → similarity score
+        for cid in raw_ids
+        if cid in filtered_ids
+    ]
+    return results[:k]
+
+
+# ---------------------------------------------------------------------------
+# RRF fusion
+# ---------------------------------------------------------------------------
+
+def _rrf_fuse(
+    bm25: list[tuple[int, float]],
+    vector: list[tuple[int, float]],
+    k: int = RRF_K,
+    bm25_w: float = BM25_WEIGHT,
+    dense_w: float = DENSE_WEIGHT,
+    top_k: int = 10,
+) -> list[tuple[int, float, int | None, int | None]]:
+    """Reciprocal Rank Fusion with weights.
+
+    Returns:
+        List of (chunk_id, rrf_score, bm25_rank, vector_rank) sorted by score desc.
+    """
+    bm25_ranks = {cid: rank for rank, (cid, _) in enumerate(bm25, start=1)}
+    vec_ranks  = {cid: rank for rank, (cid, _) in enumerate(vector, start=1)}
+    all_ids    = set(bm25_ranks) | set(vec_ranks)
+
+    scores: list[tuple[int, float, int | None, int | None]] = []
+    for cid in all_ids:
+        br = bm25_ranks.get(cid)
+        vr = vec_ranks.get(cid)
+        score = 0.0
+        if br is not None:
+            score += bm25_w / (k + br)
+        if vr is not None:
+            score += dense_w / (k + vr)
+        scores.append((cid, score, br, vr))
+
+    scores.sort(key=lambda x: x[1], reverse=True)
+    return scores[:top_k]
+
+
+# ---------------------------------------------------------------------------
+# Chunk hydration
+# ---------------------------------------------------------------------------
+
+def _hydrate_chunks(
+    conn: Connection, chunk_ids: list[int]
+) -> dict[int, dict[str, Any]]:
+    """Fetch full chunk metadata for a list of chunk IDs."""
+    if not chunk_ids:
+        return {}
+    ph = ",".join("?" * len(chunk_ids))
+    rows = conn.execute(
+        f"""
+        SELECT id, file_path, content, heading_path, heading_level,
+               chunk_type, start_line, end_line
+        FROM chunks
+        WHERE id IN ({ph})
+        """,
+        chunk_ids,
+    ).fetchall()
+    return {
+        row[0]: {
+            "file_path": row[1],
+            "content": row[2],
+            "heading_path": row[3],
+            "heading_level": row[4],
+            "chunk_type": row[5],
+            "start_line": row[6],
+            "end_line": row[7],
+        }
+        for row in rows
+    }
+
+
+# ---------------------------------------------------------------------------
+# PUBLIC: docs_search
+# ---------------------------------------------------------------------------
+
+def docs_search(
+    query: str,
+    conn: Connection,
+    embedder: CodeEmbedder,
+    k: int = 10,
+    source_filter: dict[str, Any] | None = None,
+) -> list[DocSearchHit]:
+    """Hybrid BM25 + vector search over indexed documentation.
+
+    Args:
+        query:         Natural-language or keyword search string.
+        conn:          Active SQLite reader connection.
+        embedder:      Embedder singleton (jina 768-dim).
+        k:             Number of results to return.
+        source_filter: Optional dict with keys: repo, file, heading, source_type.
+
+    Returns:
+        List of ``DocSearchHit`` objects sorted by RRF score descending.
+    """
+    if not query.strip():
+        return []
+
+    filter_clause, filter_params = _build_source_filter_clause(source_filter)
+
+    # Embed query
+    try:
+        query_vec = embedder.embed([query])[0]
+    except Exception as exc:  # noqa: BLE001
+        logger.error("Query embedding failed", error=str(exc))
+        query_vec = None
+
+    # BM25
+    bm25_results = _bm25_docs_search(conn, query, k * 3, filter_clause, filter_params)
+
+    # Vector (skip if embedder failed)
+    vector_results: list[tuple[int, float]] = []
+    if query_vec is not None:
+        vector_results = _vector_docs_search(conn, query_vec, k * 3, filter_clause, filter_params)
+
+    # RRF
+    fused = _rrf_fuse(bm25_results, vector_results, top_k=k)
+
+    if not fused:
+        return []
+
+    chunk_ids = [cid for cid, *_ in fused]
+    meta = _hydrate_chunks(conn, chunk_ids)
+
+    hits: list[DocSearchHit] = []
+    for cid, score, br, vr in fused:
+        if cid not in meta:
+            continue
+        m = meta[cid]
+        hits.append(DocSearchHit(
+            chunk_id=cid,
+            file_path=m["file_path"],
+            content=m["content"],
+            heading_path=m["heading_path"],
+            heading_level=m["heading_level"],
+            chunk_type=m["chunk_type"],
+            start_line=m["start_line"],
+            end_line=m["end_line"],
+            score=round(score, 6),
+            bm25_rank=br,
+            vector_rank=vr,
+        ))
+
+    logger.debug(
+        "docs_search complete",
+        query=query[:60],
+        hits=len(hits),
+        bm25_candidates=len(bm25_results),
+        vector_candidates=len(vector_results),
+    )
+    return hits
+
+
+# ---------------------------------------------------------------------------
+# PUBLIC: docs_get_section
+# ---------------------------------------------------------------------------
+
+def docs_get_section(
+    file: str,
+    heading_path: str,
+    conn: Connection,
+    include_subsections: bool = False,
+    fuzzy: bool = True,
+) -> DocSection | dict[str, Any]:
+    """Retrieve an exact document section by heading path.
+
+    5-tier resolution strategy (stops at first hit):
+      1. Exact match:            heading_path == input
+      2. GitHub-slug match:      slugify(heading_path) == slugify(input)
+      3. Suffix match:           heading_path ends with input (leaf-only query)
+      4. RapidFuzz WRatio ≥ 72:  fuzzy string match across all headings in file
+      5. FTS5 trigram fallback:  content MATCH on query terms
+
+    Args:
+        file:               repo-relative file path (e.g. "docs/README.md").
+        heading_path:       Path like "Installation > macOS > Homebrew" or
+                            "installation/macos/homebrew" or just "Homebrew".
+        conn:               Active SQLite reader connection.
+        include_subsections: Append child sections' content to the result.
+        fuzzy:              Enable tiers 3-5 (set False for strict exact-only).
+
+    Returns:
+        ``DocSection`` on success, or dict ``{"error": ..., "did_you_mean": [...]}``
+        on failure.
+    """
+    # Normalise input
+    normalised = _normalise_heading_path(heading_path)
+
+    # Fetch all chunks for this file
+    rows = conn.execute(
+        """
+        SELECT id, heading_path, heading_level, content, start_line, end_line
+        FROM chunks
+        WHERE file_path = ? AND source_type = 'docs'
+        ORDER BY start_line
+        """,
+        (file,),
+    ).fetchall()
+
+    if not rows:
+        return {"error": "file_not_found", "file": file, "did_you_mean": []}
+
+    def _row_to_section(row: tuple) -> DocSection:
+        return DocSection(
+            file_path=file,
+            heading_path=row[1],
+            heading_level=row[2],
+            content=row[3],
+            start_line=row[4],
+            end_line=row[5],
+        )
+
+    all_headings = [row[1] for row in rows]
+
+    # ---- Tier 1: Exact ----
+    for row in rows:
+        if row[1] == normalised:
+            sec = _row_to_section(row)
+            if include_subsections:
+                _attach_subsections(sec, rows, row[2])
+            return sec
+
+    if not fuzzy:
+        suggestions = _top_suggestions(normalised, all_headings, n=3)
+        return {"error": "heading_not_found", "did_you_mean": suggestions}
+
+    # ---- Tier 2: GitHub slug match ----
+    slug_input = _gh_slug(normalised)
+    for row in rows:
+        if _gh_slug(row[1]) == slug_input:
+            sec = _row_to_section(row)
+            if include_subsections:
+                _attach_subsections(sec, rows, row[2])
+            return sec
+
+    # ---- Tier 3: Suffix / leaf match ----
+    leaf = normalised.split(">")[-1].strip().lower()
+    for row in rows:
+        parts = [p.strip().lower() for p in row[1].split(">")]
+        if parts and parts[-1] == leaf:
+            sec = _row_to_section(row)
+            if include_subsections:
+                _attach_subsections(sec, rows, row[2])
+            return sec
+
+    # ---- Tier 4: RapidFuzz WRatio ----
+    best_score = 0.0
+    best_row = None
+    for row in rows:
+        score = fuzz.WRatio(normalised.lower(), row[1].lower())
+        if score > best_score:
+            best_score = score
+            best_row = row
+    if best_row is not None and best_score >= FUZZY_MATCH_THRESHOLD:
+        logger.debug(
+            "Heading resolved via fuzzy match",
+            input=normalised,
+            matched=best_row[1],
+            score=best_score,
+        )
+        sec = _row_to_section(best_row)
+        if include_subsections:
+            _attach_subsections(sec, rows, best_row[2])
+        return sec
+
+    # ---- Tier 5: FTS5 trigram fallback ----
+    safe_q = re.sub(r'["\'\(\)\*\:\^]', " ", heading_path).strip()
+    try:
+        fts_row = conn.execute(
+            """
+            SELECT c.id, c.heading_path, c.heading_level, c.content,
+                   c.start_line, c.end_line
+            FROM fts_chunks
+            JOIN chunks c ON fts_chunks.chunk_id = c.id
+            WHERE fts_chunks MATCH ? AND c.file_path = ? AND c.source_type = 'docs'
+            ORDER BY -bm25(fts_chunks)
+            LIMIT 1
+            """,
+            (safe_q, file),
+        ).fetchone()
+        if fts_row:
+            sec = DocSection(
+                file_path=file,
+                heading_path=fts_row[1],
+                heading_level=fts_row[2],
+                content=fts_row[3],
+                start_line=fts_row[4],
+                end_line=fts_row[5],
+            )
+            if include_subsections:
+                _attach_subsections(sec, (fts_row,), fts_row[2])
+            return sec
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("FTS5 fallback failed", error=str(exc))
+
+    suggestions = _top_suggestions(normalised, all_headings, n=3)
+    return {
+        "error": "heading_not_found",
+        "file": file,
+        "query": heading_path,
+        "did_you_mean": suggestions,
+    }
+
+
+# ---------------------------------------------------------------------------
+# PUBLIC: docs_list_files
+# ---------------------------------------------------------------------------
+
+def docs_list_files(
+    conn: Connection,
+    repo: str | None = None,
+    glob: str | None = None,
+    include_headings: bool = False,
+) -> list[DocFileInfo]:
+    """List all indexed documentation files.
+
+    Args:
+        conn:             Active SQLite reader connection.
+        repo:             Filter by repo prefix (e.g. "src/docs").
+        glob:             fnmatch glob pattern on file_path (e.g. "**/*.md").
+        include_headings: If True, attach heading_path list to each file.
+
+    Returns:
+        List of ``DocFileInfo`` sorted by file_path.
+    """
+    sql = """
+        SELECT c.file_path, COUNT(*) AS chunk_count
+        FROM chunks c
+        WHERE c.source_type = 'docs'
+    """
+    params: list[Any] = []
+
+    if repo:
+        sql += " AND c.file_path LIKE ?"
+        params.append(f"{repo}/%")
+
+    sql += " GROUP BY c.file_path ORDER BY c.file_path"
+
+    rows = conn.execute(sql, params).fetchall()
+
+    results: list[DocFileInfo] = []
+    for file_path, chunk_count in rows:
+        # Apply glob filter in Python (fnmatch glob is not SQL)
+        if glob and not fnmatch.fnmatch(file_path, glob):
+            continue
+
+        headings: list[str] = []
+        if include_headings:
+            h_rows = conn.execute(
+                """
+                SELECT DISTINCT heading_path
+                FROM chunks
+                WHERE file_path = ? AND source_type = 'docs'
+                  AND heading_path != ''
+                ORDER BY start_line
+                """,
+                (file_path,),
+            ).fetchall()
+            headings = [r[0] for r in h_rows]
+
+        results.append(DocFileInfo(
+            file_path=file_path,
+            chunk_count=chunk_count,
+            headings=headings,
+        ))
+
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Heading helpers
+# ---------------------------------------------------------------------------
+
+def _normalise_heading_path(raw: str) -> str:
+    """Normalise heading path: strip whitespace, convert '/' to ' > '."""
+    # Accept both "Installation > macOS" and "installation/macos" styles
+    normalised = raw.replace("/", " > ").strip()
+    # Collapse multiple spaces
+    normalised = re.sub(r"\s{2,}", " ", normalised)
+    return normalised
+
+
+def _gh_slug(heading: str) -> str:
+    """GitHub-flavored markdown heading → anchor slug."""
+    slug = heading.lower()
+    slug = re.sub(r"[^\w\s-]", "", slug)
+    slug = re.sub(r"\s+", "-", slug.strip())
+    return slug
+
+
+def _top_suggestions(
+    query: str, candidates: list[str], n: int = 3
+) -> list[str]:
+    """Return top-N fuzzy matches from candidates for a not-found heading."""
+    scored = sorted(
+        candidates,
+        key=lambda h: fuzz.WRatio(query.lower(), h.lower()),
+        reverse=True,
+    )
+    return scored[:n]
+
+
+def _attach_subsections(
+    parent: DocSection,
+    all_rows: tuple | list,
+    parent_level: int,
+) -> None:
+    """Attach child sections (deeper heading level) to parent.subsections."""
+    capturing = False
+    for row in all_rows:
+        hp, lvl, content, sl, el = row[1], row[2], row[3], row[4], row[5]
+        if hp == parent.heading_path and lvl == parent_level:
+            capturing = True
+            continue
+        if capturing:
+            if lvl <= parent_level:
+                break
+            parent.subsections.append(DocSection(
+                file_path=parent.file_path,
+                heading_path=hp,
+                heading_level=lvl,
+                content=content,
+                start_line=sl,
+                end_line=el,
+            ))