SourceIndex 0.1.0__py3-none-any.whl

sourceindex/search/__init__.py

@@ -0,0 +1,251 @@
+"""Two-pass search pipeline producing an enriched codebase roadmap.
+
+Pass 1a: index_search_model reads tier-1 index + query → picks relevant files.
+Pass 1b: index_search_model reads tier-2 function indexes → picks specific
+         functions.
+Enrichment: annotate each entry with its tier-2 description, extract top-level
+            imports from the selected files, group entries per file with merged
+            read offset/limit, and wrap in a roadmap preamble.
+"""
+
+import re
+import time
+from pathlib import Path
+
+from ..build import DEFAULT_MODEL, refresh_index_lazy
+from ..lib.llm import call_llm
+from ..lib.log import get_logger
+from ..lib.timing import elapsed_ms
+from .experiments import apply_experiments
+from .passes import _strip_llm_wrapper
+from .prompts import (
+    PASS_1A_PROMPT, PASS_1A_QA_PROMPT,
+    PASS_1B_PROMPT, PASS_1B_QA_PROMPT,
+)
+from .roadmap import _build_roadmap
+
+_log = get_logger(__name__)
+
+
+def _salvage_selected_path(filepath: str, store) -> str | None:
+    """Recover a Pass-1a path that doesn't match the index keys.
+
+    The search model occasionally echoes the agent's mount prefix onto paths
+    (``/app/cmd/run.go`` for the indexed ``cmd/run.go``) or inserts stray spaces
+    around separators. Detail keys are repo-relative, so normalize spacing and
+    then progressively strip leading components until a known detail matches.
+    """
+    candidate = re.sub(r"\s*/\s*", "/", filepath).strip()
+    for _ in range(4):
+        if store.exists(f"details/{candidate}.md"):
+            return candidate
+        if "/" not in candidate:
+            return None
+        candidate = candidate.split("/", 1)[1]
+    return candidate if store.exists(f"details/{candidate}.md") else None
+
+
+def search_index(
+    query: str,
+    repo_root: Path,
+    index_dir: Path | None = None,
+    index_search_model: str = DEFAULT_MODEL,
+    mode: str = "implement",
+) -> str:
+    """Run the two-pass search pipeline and return an enriched roadmap.
+
+    Args:
+        query: Natural language description of what the agent needs.
+        repo_root: Path to the repository root.
+        index_dir: Path to the index directory (default: repo_root/.sourceindex).
+        index_search_model: LLM model for passes 1a and 1b.
+        mode: "implement" (default) or "qa" — selects prompt variants and
+              roadmap constraints appropriate for the task type.
+
+    Returns:
+        Roadmap text with per-file groupings, function descriptions, and imports.
+    """
+    from ..build import _resolve_plaintext_store
+    store = _resolve_plaintext_store(repo_root, index_dir)
+    return search_index_with_store(query, repo_root, store, index_search_model, None, mode=mode)
+
+
+def search_index_with_store(
+    query: str,
+    repo_root: Path,
+    store,
+    index_search_model: str | None = None,
+    api_key: str | None = None,
+    *,
+    mode: str = "implement",
+) -> str:
+    is_qa = mode == "qa"
+    prompt_1a_tpl = PASS_1A_QA_PROMPT if is_qa else PASS_1A_PROMPT
+    prompt_1b_tpl = PASS_1B_QA_PROMPT if is_qa else PASS_1B_PROMPT
+    model = index_search_model or DEFAULT_MODEL
+    search_start = time.perf_counter()
+    _log.debug(
+        "Search starting",
+        extra={
+            "event": "search.start",
+            "query": query,
+            "query_len": len(query),
+            "repo_root": str(repo_root),
+            "model": model,
+        },
+    )
+
+    try:
+        refresh_index_lazy(repo_root, store=store, model=model, api_key=api_key)
+    except Exception as e:
+        _log.warning(
+            f"Lazy refresh failed (continuing with stale index): {e}",
+            extra={
+                "event": "lazy_refresh.failed",
+                "error_type": type(e).__name__,
+                "error": str(e),
+            },
+        )
+
+    if not store.exists("index.md"):
+        _log.error(
+            "No index found. Run `sourceindex init` first.",
+            extra={"event": "search.no_index", "index_path": str(store.path_for("index.md"))},
+        )
+        return "ERROR: No index found. Run `sourceindex init` first."
+
+    index_content = store.read_text("index.md")
+
+    from ..lib.llm import resolve_api_key
+    search_api_key = resolve_api_key(api_key)
+    assert search_api_key, (
+        "No API key found. Set SOURCEINDEX_API_KEY (an OpenRouter key, used "
+        "for both indexing and search)."
+    )
+
+    # ── Pass 1a: Pick relevant files ──────────────────────────────────────
+    _log.info("Pass 1a: Selecting relevant files...", extra={"event": "search.pass_1a.start"})
+
+    prompt_1a = prompt_1a_tpl.format(index_content=index_content, instruction=query)
+    pass_1a_t0 = time.perf_counter()
+    file_list, _ = call_llm(
+        prompt_1a,
+        model=model,
+        api_key=search_api_key,
+    )
+    if not file_list:
+        _log.error(
+            "Pass 1a (file selection) failed.",
+            extra={"event": "search.pass_1a.failed"},
+        )
+        return "ERROR: Pass 1a (file selection) failed."
+    file_list = _strip_llm_wrapper(file_list)
+
+    selected_files = [f.strip().lstrip("/") for f in file_list.strip().split("\n") if f.strip()]
+    _log.info(
+        f"Selected files:\n{file_list}",
+        extra={
+            "event": "search.pass_1a",
+            "files_selected": len(selected_files),
+            "files": selected_files,
+            "latency_ms": elapsed_ms(pass_1a_t0),
+        },
+    )
+
+    # ── Load tier-2 details for selected files via the store ──────────────
+    selected_details = ""
+    skipped: list[str] = []
+    detail_texts: dict[str, str] = {}
+    for filepath in selected_files:
+        if not filepath:
+            continue
+        if not store.exists(f"details/{filepath}.md"):
+            salvaged = _salvage_selected_path(filepath, store)
+            if salvaged is None:
+                skipped.append(filepath)
+                _log.debug(
+                    f"  Skipping (no detail): {filepath}",
+                    extra={"event": "search.detail_missing", "path": filepath},
+                )
+                continue
+            _log.info(
+                f"  Salvaged Pass-1a path: {filepath} -> {salvaged}",
+                extra={"event": "search.detail_salvaged",
+                       "from_path": filepath, "to_path": salvaged},
+            )
+            filepath = salvaged
+        text = store.read_text(f"details/{filepath}.md")
+        detail_texts[filepath] = text
+        selected_details += f"\n--- {filepath} ---\n{text}\n--- end ---\n\n"
+
+    if not selected_details:
+        _log.error(
+            "No tier-2 details found for any selected file.",
+            extra={"event": "search.no_details", "skipped": skipped},
+        )
+        return "ERROR: No tier-2 details found for any selected file."
+
+    # ── Pass 1b: Pick specific functions ──────────────────────────────────
+    _log.info(
+        "Pass 1b: Selecting specific functions...",
+        extra={"event": "search.pass_1b.start"},
+    )
+
+    prompt_1b = prompt_1b_tpl.format(selected_details=selected_details, instruction=query)
+    pass_1b_t0 = time.perf_counter()
+    function_list, _ = call_llm(
+        prompt_1b,
+        model=model,
+        api_key=search_api_key,
+    )
+    if not function_list:
+        _log.error(
+            "Pass 1b (function selection) failed.",
+            extra={"event": "search.pass_1b.failed"},
+        )
+        return "ERROR: Pass 1b (function selection) failed."
+    # Split out optional CODE LANDSCAPE block. The Pass-1b prompt asks the
+    # model to emit `=== CODE LANDSCAPE ===` followed by `=== SELECTIONS ===`.
+    # If markers are missing (older models / format drift), treat the entire
+    # response as SELECTIONS — backward compatible with the prior shape.
+    raw_output = function_list
+    landscape_block = ""
+    if "=== SELECTIONS ===" in raw_output:
+        head, _, tail = raw_output.partition("=== SELECTIONS ===")
+        if "=== CODE LANDSCAPE ===" in head:
+            _, _, landscape_block = head.partition("=== CODE LANDSCAPE ===")
+        else:
+            landscape_block = head
+        function_list = tail
+    function_list = _strip_llm_wrapper(function_list)
+
+    selected_funcs = [l.strip() for l in function_list.split("\n") if l.strip()]
+    _log.info(
+        f"Selected functions:\n{function_list}",
+        extra={
+            "event": "search.pass_1b",
+            "functions_selected": len(selected_funcs),
+            "landscape_chars": len(landscape_block.strip()),
+            "latency_ms": elapsed_ms(pass_1b_t0),
+        },
+    )
+
+    # Env-gated post-processors (SOURCEINDEX_FIX_*). No-op when flags unset.
+    function_list = apply_experiments(function_list, query, repo_root, store)
+
+    # ── Enrichment: annotate, group, extract imports, wrap preamble ──────
+    roadmap = _build_roadmap(
+        function_list, repo_root, store, detail_texts,
+        landscape=landscape_block, mode=mode, question=query,
+    )
+    _log.info(
+        f"Roadmap: {len(roadmap)} bytes.",
+        extra={
+            "event": "search.done",
+            "roadmap_bytes": len(roadmap),
+            "files_selected": len(selected_files),
+            "functions_selected": len(selected_funcs),
+            "latency_ms": elapsed_ms(search_start),
+        },
+    )
+    return roadmap

sourceindex/search/experiments.py

@@ -0,0 +1,276 @@
+"""Pass-1b post-processing experiments, env-var gated.
+
+Three independent post-processors on the Pass-1b SELECTIONS string. Each
+is a no-op unless its ``SOURCEINDEX_FIX_*`` env var is set to ``1``, so
+the production search path is byte-identical when flags are unset.
+
+  SOURCEINDEX_FIX_SIBCAP=1     cap_sibling_pulls
+  SOURCEINDEX_FIX_TESTPAIR=1   inject_test_pairs
+  SOURCEINDEX_FIX_DOCS=1       inject_docs_for_add_task
+
+Apply via ``apply_experiments`` immediately before ``_build_roadmap``.
+"""
+
+import os
+import re
+from collections import defaultdict
+from pathlib import Path
+
+# "EDIT pre_commit/store.py::Store._new_repo (L132-L171)" → tag, path, rest
+_TAG_LINE_RE = re.compile(r"^(EDIT|READ)\s+([^\s:]+(?:[^\s:]|/)*?)::(.+)$")
+_TIER2_ENTRY_RE = re.compile(r"^(.+?) \(L(\d+)-L(\d+)\):")
+
+
+def _parse_selections(function_list: str):
+    """Yield (tag, path, rest) per non-blank SELECTIONS line."""
+    for raw in function_list.split("\n"):
+        line = raw.strip()
+        if not line or "::" not in line:
+            continue
+        m = _TAG_LINE_RE.match(line)
+        if m:
+            yield m.group(1), m.group(2).lstrip("/"), m.group(3)
+        else:
+            path, _, rest = line.partition("::")
+            yield "READ", path.strip().lstrip("/"), rest
+
+
+def _sibling_groups(paths: list[str]):
+    """Yield index lists for sibling groups of ≥3 paths.
+
+    Two heuristics combined:
+      A) Same parent directory  — catches e.g. ``pre_commit/{store,clientlib,
+         repository,hook,util,yaml}.py``.
+      B) Same grand-parent dir + same basename — catches e.g.
+         ``faker/providers/company/{es_ES,fr_FR,pt_BR}/__init__.py``
+         where parent dirs differ but the basename + grand-parent are shared.
+    """
+    by_parent: dict[str, set[int]] = defaultdict(set)
+    by_grand_basename: dict[tuple[str, str], set[int]] = defaultdict(set)
+    for i, p in enumerate(paths):
+        parts = Path(p).parts
+        if len(parts) >= 1:
+            by_parent[str(Path(p).parent)].add(i)
+        if len(parts) >= 3:
+            grand = str(Path(p).parent.parent)
+            base = parts[-1]
+            by_grand_basename[(grand, base)].add(i)
+    seen: set[tuple[int, ...]] = set()
+    for idxs in list(by_parent.values()) + list(by_grand_basename.values()):
+        if len(idxs) >= 3:
+            key = tuple(sorted(idxs))
+            if key not in seen:
+                seen.add(key)
+                yield list(key)
+
+
+def cap_sibling_pulls(function_list: str, max_read_per_group: int = 2) -> str:
+    """Drop READ entries when a sibling group exceeds ``max_read_per_group``.
+
+    EDIT entries are never dropped. Within a sibling group, the first N READ
+    files (by Pass-1b emission order) are kept; the rest have *all* their
+    SELECTIONS lines stripped.
+    """
+    parsed = list(_parse_selections(function_list))
+    if not parsed:
+        return function_list
+
+    # Per-file roll-up: dominant tag, first-emission index
+    file_tag: dict[str, str] = {}
+    file_first_idx: dict[str, int] = {}
+    for i, (tag, path, _) in enumerate(parsed):
+        if path not in file_tag:
+            file_tag[path] = tag
+            file_first_idx[path] = i
+        elif tag == "EDIT":
+            file_tag[path] = "EDIT"
+
+    unique_paths = list(file_tag.keys())
+    drop_files: set[str] = set()
+    for group_idxs in _sibling_groups(unique_paths):
+        group_paths = {unique_paths[i] for i in group_idxs}
+        read_paths = [p for p in unique_paths if p in group_paths and file_tag[p] == "READ"]
+        for p in read_paths[max_read_per_group:]:
+            drop_files.add(p)
+
+    if not drop_files:
+        return function_list
+
+    kept = []
+    for raw in function_list.split("\n"):
+        line = raw.strip()
+        if line and "::" in line:
+            m = _TAG_LINE_RE.match(line)
+            if m:
+                path = m.group(2).lstrip("/")
+            else:
+                path = line.partition("::")[0].strip().lstrip("/")
+            if path in drop_files:
+                continue
+        kept.append(raw)
+    return "\n".join(kept)
+
+
+def _test_path_candidates(edit_path: str) -> list[str]:
+    """Generate candidate test-file paths for an EDIT target.
+
+    Covers conventions seen in the eval corpus:
+      pre_commit/store.py       → tests/store_test.py
+      src/jinja2/parser.py      → tests/test_parser.py
+      hermes_cli/models.py      → tests/hermes_cli/test_ai_gateway_models.py
+                                  (we won't catch the qualifier — handled
+                                  by listing both basic + dir-mirrored forms)
+      awscli/customizations/s3/subcommands.py
+                                → tests/functional/s3/test_subcommands.py
+    """
+    p = Path(edit_path)
+    basename = p.stem
+    parts = list(p.parts)
+    if parts and parts[0] == "src":
+        parts = parts[1:]
+    cands: list[str] = []
+    # Direct conventions at top of tests/ tree
+    for prefix in ("tests", "test"):
+        cands.append(f"{prefix}/test_{basename}.py")
+        cands.append(f"{prefix}/{basename}_test.py")
+    # Mirror inner package layout
+    if len(parts) >= 2:
+        subdir = "/".join(parts[1:-1])
+        if subdir:
+            for prefix in ("tests", "test"):
+                cands.append(f"{prefix}/{subdir}/test_{basename}.py")
+                cands.append(f"{prefix}/{subdir}/{basename}_test.py")
+    # awscli/django split layout
+    for sub in ("functional", "unit"):
+        cands.append(f"tests/{sub}/test_{basename}.py")
+        if len(parts) >= 2:
+            subdir = "/".join(parts[1:-1])
+            if subdir:
+                cands.append(f"tests/{sub}/{subdir}/test_{basename}.py")
+    seen = set()
+    out = []
+    for c in cands:
+        if c not in seen:
+            seen.add(c)
+            out.append(c)
+    return out
+
+
+def inject_test_pairs(
+    function_list: str, repo_root: Path, store=None, max_entries: int = 3
+) -> str:
+    """Append READ entries for test files that pair with EDIT targets.
+
+    For each EDIT target, walk a list of conventional test-file paths, pick
+    the first that exists on disk, and emit up to ``max_entries`` synthetic
+    READ lines drawn from that test file's tier-2 details (so the roadmap
+    builder produces real ``offset``/``limit`` windows). If no tier-2 is
+    cached for the file, fall back to a whole-file READ marker.
+    """
+    parsed = list(_parse_selections(function_list))
+    edit_files = [p for tag, p, _ in parsed if tag == "EDIT"]
+    listed = {p for _, p, _ in parsed}
+    seen_edit: set[str] = set()
+    extras: list[str] = []
+    for edit_path in edit_files:
+        if edit_path in seen_edit:
+            continue
+        seen_edit.add(edit_path)
+        for candidate in _test_path_candidates(edit_path):
+            if candidate in listed:
+                break
+            full = repo_root / candidate
+            if not full.is_file():
+                continue
+            count = 0
+            rel = f"details/{candidate}.md"
+            if store is not None:
+                try:
+                    if store.exists(rel):
+                        for line in store.read_text(rel).splitlines():
+                            m = _TIER2_ENTRY_RE.match(line.strip())
+                            if m:
+                                extras.append(
+                                    f"READ {candidate}::{m.group(1).strip()} "
+                                    f"(L{m.group(2)}-L{m.group(3)})"
+                                )
+                                count += 1
+                                if count >= max_entries:
+                                    break
+                except Exception:
+                    pass
+            if count == 0:
+                extras.append(f"READ {candidate}::module (L1-L1)")
+            listed.add(candidate)
+            break
+    if not extras:
+        return function_list
+    return function_list.rstrip("\n") + "\n" + "\n".join(extras) + "\n"
+
+
+_ADD_VERB_RE = re.compile(
+    r"\b(add|adds|added|adding|introduce|introduces|introducing|"
+    r"support|new\s+(?:flag|command|option|provider|locale|channel|tag))\b",
+    re.IGNORECASE,
+)
+
+
+def inject_docs_for_add_task(
+    function_list: str, query: str, repo_root: Path
+) -> str:
+    """When the task verb is add/introduce/etc., surface CHANGES + docs.
+
+    Scans the repo for the obvious documentation-update sites and appends
+    them as whole-file READ entries. Limited scan (no recursion past the
+    first level of ``docs/``) to keep the roadmap from ballooning.
+    """
+    if not _ADD_VERB_RE.search(query or ""):
+        return function_list
+    listed = {p for _, p, _ in _parse_selections(function_list)}
+    extras: list[str] = []
+
+    # Top-level CHANGES / CHANGELOG variants
+    for name in ("CHANGES.rst", "CHANGES.md", "CHANGES.txt",
+                 "CHANGELOG.md", "CHANGELOG.rst", "CHANGELOG.txt"):
+        if (repo_root / name).is_file() and name not in listed:
+            extras.append(f"READ {name}::module (L1-L1)")
+            listed.add(name)
+            break  # at most one CHANGELOG
+
+    # .changes/next-release/ (awscli convention)
+    cd = repo_root / ".changes" / "next-release"
+    if cd.is_dir():
+        for entry in sorted(cd.iterdir())[:2]:
+            if entry.is_file():
+                rel = f".changes/next-release/{entry.name}"
+                if rel not in listed:
+                    extras.append(f"READ {rel}::module (L1-L1)")
+                    listed.add(rel)
+
+    # Single docs/ entry — surface one likely-relevant doc, not all of them
+    docs = repo_root / "docs"
+    if docs.is_dir():
+        for entry in sorted(docs.iterdir()):
+            if entry.is_file() and entry.suffix in (".md", ".rst"):
+                rel = f"docs/{entry.name}"
+                if rel not in listed:
+                    extras.append(f"READ {rel}::module (L1-L1)")
+                    listed.add(rel)
+                    break
+
+    if not extras:
+        return function_list
+    return function_list.rstrip("\n") + "\n" + "\n".join(extras) + "\n"
+
+
+def apply_experiments(
+    function_list: str, query: str, repo_root: Path, store=None
+) -> str:
+    """Apply env-gated post-processors to Pass-1b SELECTIONS."""
+    if os.environ.get("SOURCEINDEX_FIX_SIBCAP") == "1":
+        function_list = cap_sibling_pulls(function_list)
+    if os.environ.get("SOURCEINDEX_FIX_TESTPAIR") == "1":
+        function_list = inject_test_pairs(function_list, repo_root, store)
+    if os.environ.get("SOURCEINDEX_FIX_DOCS") == "1":
+        function_list = inject_docs_for_add_task(function_list, query, repo_root)
+    return function_list

sourceindex/search/imports.py

@@ -0,0 +1,155 @@
+"""Top-of-file imports extraction + Python imports-preamble detection.
+
+Used by the roadmap builder to:
+- show the agent every import statement from each roadmap file (so it sees
+  the surrounding API without a full-file read), and
+- emit a separate ``(1, preamble_end)`` Read window for Python EDIT targets
+  so new imports are added without an extra round trip.
+
+Python is special-cased to use the AST because regex can't handle two
+common patterns: multi-line ``from X import (a, b, c)`` and
+``try/except ImportError`` shims with ``if TYPE_CHECKING:`` guards. Other
+languages use the per-Language ``import_pattern`` regex.
+"""
+
+import ast
+import re
+from pathlib import Path
+
+from ..lib.languages import default_languages, language_for_path
+
+
+# Languages whose imports are extracted via AST instead of the per-Language
+# `import_pattern` regex. Python is the only one — and stays special-cased
+# because (a) Python is the dominant target language for this tool, so its
+# output quality matters most, and (b) regex can't reproduce two AST-only
+# behaviors that are common in real Python code:
+#   1. Multi-line imports — `from foo import (\n a,\n b,\n)` is one logical
+#      statement; regex only matches the first line and drops `a, b`.
+#   2. `try/except ImportError:` shims and `if TYPE_CHECKING:` guards — AST
+#      sees the surrounding structure; regex sees the import lines in
+#      isolation, missing the shim that explains them.
+# `_imports_preamble_end` (used in _build_roadmap to add a `(1, preamble_end)`
+# Read window for EDIT targets) is also Python-only for the same reason: it
+# needs the AST to know where the imports block ends — regex matches don't
+# tell you that.
+_AST_IMPORT_LANGS: frozenset[str] = frozenset({"python"})
+
+_IMPORT_SCAN_MAX_LINES = 400  # cap for regex-based scan in non-Python files
+
+
+def _parse_py(source_path: Path) -> ast.Module | None:
+    """Parse a Python file to an AST, or return None on any read/parse failure."""
+    try:
+        return ast.parse(source_path.read_text(errors="replace"))
+    except (OSError, SyntaxError, ValueError):
+        return None
+
+
+def _is_import_only_try(node: ast.Try) -> bool:
+    return bool(node.body) and all(
+        isinstance(b, (ast.Import, ast.ImportFrom)) for b in node.body
+    )
+
+
+def _is_type_checking_if(node: ast.If) -> bool:
+    t = node.test
+    return (isinstance(t, ast.Name) and t.id == "TYPE_CHECKING") or (
+        isinstance(t, ast.Attribute) and t.attr == "TYPE_CHECKING"
+    )
+
+
+def _is_dunder_assign(node: ast.AST) -> bool:
+    targets = getattr(node, "targets", None) or (
+        [node.target] if isinstance(node, ast.AnnAssign) else []
+    )
+    return any(
+        isinstance(t, ast.Name) and t.id.startswith("__") and t.id.endswith("__")
+        for t in targets
+    )
+
+
+def _extract_imports_python(source_path: Path) -> str:
+    """Return top-level import statements for a Python file, one per line."""
+    tree = _parse_py(source_path)
+    if tree is None:
+        return ""
+
+    out: list[str] = []
+    for node in ast.iter_child_nodes(tree):
+        if isinstance(node, ast.Import):
+            names = ", ".join(
+                a.name + (f" as {a.asname}" if a.asname else "") for a in node.names
+            )
+            out.append(f"import {names}")
+        elif isinstance(node, ast.ImportFrom):
+            module = "." * node.level + (node.module or "")
+            names = ", ".join(
+                a.name + (f" as {a.asname}" if a.asname else "") for a in node.names
+            )
+            out.append(f"from {module} import {names}")
+    return "\n".join(out)
+
+
+def _extract_imports_regex(source_path: Path, pattern: str) -> str:
+    """Scan the top of a file for lines matching the language's import pattern.
+
+    Stops after _IMPORT_SCAN_MAX_LINES; collects every matching line verbatim
+    (trimmed). Designed for non-Python languages where AST parsing isn't
+    available — pattern correctness is best-effort, not authoritative."""
+    rx = re.compile(pattern)
+    out: list[str] = []
+    try:
+        with source_path.open(errors="replace") as f:
+            for i, raw in enumerate(f):
+                if i >= _IMPORT_SCAN_MAX_LINES:
+                    break
+                if rx.match(raw):
+                    out.append(raw.rstrip("\n").strip())
+    except OSError:
+        return ""
+    return "\n".join(out)
+
+
+def extract_imports(source_path: Path, rel_path: str) -> str:
+    """Return import statements for any supported language. Empty string if
+    the language is unknown or has no import_pattern."""
+    lang = language_for_path(rel_path, list(default_languages()))
+    if lang is None:
+        return ""
+    if lang.key in _AST_IMPORT_LANGS:
+        return _extract_imports_python(source_path)
+    if lang.import_pattern is None:
+        return ""
+    return _extract_imports_regex(source_path, lang.import_pattern)
+
+
+def _imports_preamble_end(source_path: Path) -> int:
+    """Return 1-indexed end line of the top-of-file imports preamble, or 0.
+
+    Preamble = contiguous module docstring + Import/ImportFrom +
+    try/except-ImportError shims + ``if TYPE_CHECKING:`` guards + dunder
+    assignments (__all__, __version__), stopping at the first FunctionDef,
+    ClassDef, or other real code.
+    """
+    tree = _parse_py(source_path)
+    if tree is None:
+        return 0
+
+    end = 0
+    for i, node in enumerate(ast.iter_child_nodes(tree)):
+        # A module docstring is only a docstring when it's the first statement;
+        # a string-Expr at i > 0 is a no-op, not a docstring, so reject it there.
+        ok = (
+            isinstance(node, (ast.Import, ast.ImportFrom))
+            or (i == 0 and isinstance(node, ast.Expr)
+                and isinstance(getattr(node, "value", None), ast.Constant)
+                and isinstance(node.value.value, str))
+            or (isinstance(node, (ast.Assign, ast.AnnAssign)) and _is_dunder_assign(node))
+            or (isinstance(node, ast.If) and _is_type_checking_if(node))
+            or (isinstance(node, ast.Try) and _is_import_only_try(node))
+        )
+        if not ok:
+            break
+        end = max(end, getattr(node, "end_lineno", 0) or 0)
+    return end