SourceIndex 0.1.0__py3-none-any.whl

sourceindex/build/indexer.py

@@ -0,0 +1,403 @@
+"""LLM dispatch, response parsing, batch concurrency, and progress UI.
+
+The indexer turns ``(content, relative_path, language)`` triples into
+``(oneliner, functions_text, usage)`` triples. Two delivery shapes:
+
+- ``index_file`` — single-file, used by sequential update fallback.
+- ``submit_and_wait_batch`` + ``index_by_hash`` — concurrent batch dispatch
+  via a ThreadPoolExecutor over ``call_llm``. Used for cold builds and
+  multi-file updates.
+
+Progress UI (``_format_eta``, ``_render_progress_bar``, ``_emit_progress``)
+and fd-capacity management (``_ensure_fd_capacity``) live here too —
+they're only invoked from the batch path and don't deserve their own files.
+"""
+
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Callable
+
+import litellm
+
+from .. import REALTIME_MAX_WORKERS
+from ..lib.languages import Language
+from ..lib.llm import LLMUsage, call_llm
+from ..lib.log import get_logger
+from .linerange import fix_functions_for_language
+from .prompts import prompt_for
+
+_log = get_logger(__name__)
+
+
+REALTIME_MAX_RETRIES = 5
+# update_index falls back to sequential index_file() below this miss count —
+# the threadpool + per-request setup cost isn't worth amortizing for a single
+# file, but pays for itself at 2+.
+BATCH_MIN_FILES = 2
+RETRY_DELAY = 1.0
+# Best price/latency/quality balance found across DS-V4-Flash, Haiku 4.5,
+# Mistral-Small, Gemini-Flash-Lite, gpt-5-nano, gpt-5.4-nano (Opus-judged
+# description quality + AST-grounded recall + line-range accuracy).
+# Reached via SOURCEINDEX_API_KEY: an sk-si bearer routes through the hosted
+# backend, an sk-or-/sk-ant-/etc. provider key auto-opts-out and hits the
+# provider directly. SOURCEINDEX_BACKEND_URL pins/disables the routing.
+DEFAULT_MODEL = "openrouter/openai/gpt-5.4-nano"
+BATCH_POLL_INTERVAL = 30
+BATCH_MAX_WAIT = 3600
+BATCH_MAX_TOKENS = 16384
+
+
+# ── Response parsing ───────────────────────────────────────────────────────
+
+
+def parse_response(response: str) -> tuple[str, str]:
+    """Parse merged LLM response into (oneliner, functions_text)."""
+    oneliner = ""
+    functions = ""
+
+    lines = response.split("\n")
+    for i, line in enumerate(lines):
+        if line.startswith("ONELINER:"):
+            oneliner = line[len("ONELINER:"):].strip()
+        elif line.startswith("FUNCTIONS:"):
+            functions = "\n".join(lines[i + 1:]).strip()
+            break
+
+    if not oneliner:
+        for line in lines:
+            s = line.strip()
+            if s and not s.startswith("#"):
+                oneliner = s[:80]
+                break
+
+    return oneliner, functions
+
+
+# ── Single-file path (sequential fallback) ─────────────────────────────────
+
+
+def index_file(
+    content: str,
+    relative_path: str,
+    language: Language,
+    model: str = DEFAULT_MODEL,
+    api_key: str | None = None,
+) -> tuple[str, str, LLMUsage]:
+    """Index a single file (tier-1 + tier-2 in one call when applicable).
+
+    Returns (oneliner, functions_text, usage). ``functions_text`` is always
+    empty when ``language.build_tier2`` is False. ``usage`` is zero-filled
+    when no LLM call succeeded (e.g. empty file, all retries failed).
+    """
+    if not content.strip():
+        return "(empty file)", "", LLMUsage()
+
+    prompt = prompt_for(language).format(
+        language_name=language.name,
+        relative_path=relative_path,
+        file_content=content,
+    )
+
+    total_usage = LLMUsage()
+    for attempt in range(2):
+        try:
+            response, usage = call_llm(prompt, model=model, api_key=api_key)
+            total_usage.add(usage)
+            if response:
+                oneliner, functions = parse_response(response)
+                functions = fix_functions_for_language(language, content, functions)
+                return oneliner, functions, total_usage
+        except litellm.AuthenticationError:
+            # Deterministic — retrying just delays the friendly CLI message.
+            raise
+        except Exception as e:
+            _log.warning(
+                f"  LLM error on {relative_path}: {e}",
+                extra={
+                    "event": "llm.retry",
+                    "path": relative_path,
+                    "attempt": attempt + 1,
+                    "error_type": type(e).__name__,
+                    "error": str(e),
+                },
+            )
+        if attempt == 0:
+            _log.info(
+                f"  Retry in {RETRY_DELAY}s...",
+                extra={"event": "llm.retry_wait", "path": relative_path, "delay_s": RETRY_DELAY},
+            )
+            time.sleep(RETRY_DELAY)
+
+    _log.error(
+        f"  LLM failed for {relative_path}, skipping",
+        extra={"event": "llm.failed", "path": relative_path},
+    )
+    return "(LLM summary failed)", "", total_usage
+
+
+# ── Batch dispatch ─────────────────────────────────────────────────────────
+
+
+def build_batch_requests(file_contents: dict[str, tuple[str, str, Language]], model: str = DEFAULT_MODEL) -> list[dict]:
+    """Build per-file LLM requests for concurrent dispatch.
+
+    file_contents: {content_hash: (file_content, relative_path, language)}
+    Uses the tier-1-only prompt for languages where build_tier2 is False.
+    """
+    requests = []
+    for h, (content, rel_path, language) in file_contents.items():
+        prompt = prompt_for(language).format(
+            language_name=language.name,
+            relative_path=rel_path,
+            file_content=content,
+        )
+        requests.append({
+            "custom_id": h,
+            "model": model,
+            "max_tokens": BATCH_MAX_TOKENS,
+            "prompt": prompt,
+        })
+    return requests
+
+
+def _ensure_fd_capacity(needed: int) -> int:
+    """Raise this process's soft fd limit to fit ``needed`` open sockets,
+    plus a small headroom for stdio/log/cache fds. Returns the effective
+    worker capacity after the raise — may be less than ``needed`` if the
+    hard limit is lower (the caller must clamp). No-op on platforms
+    without a ``resource`` module (Windows)."""
+    try:
+        import resource
+    except ImportError:
+        return needed
+    HEADROOM = 64
+    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+    target = needed + HEADROOM
+    if soft >= target:
+        return needed
+    new_soft = min(target, hard) if hard > 0 else target
+    try:
+        resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, hard))
+    except (ValueError, OSError):
+        return max(1, soft - HEADROOM)
+    return max(1, new_soft - HEADROOM)
+
+
+# ── Progress UI ────────────────────────────────────────────────────────────
+
+
+def _format_eta(seconds: float) -> str:
+    """Human-friendly elapsed/ETA: seconds under 90s, otherwise minutes."""
+    if seconds < 90:
+        return f"~{int(round(seconds))}s"
+    minutes = seconds / 60.0
+    if minutes < 10:
+        return f"~{minutes:.1f} min"
+    return f"~{int(round(minutes))} min"
+
+
+def _render_progress_bar(done: int, total: int, elapsed: float, *, width: int = 28) -> str:
+    pct = (done / total) if total else 1.0
+    filled = int(pct * width)
+    bar = "█" * filled + "░" * (width - filled)
+    rate = done / elapsed if elapsed > 0 else 0.0
+    if rate > 0 and done < total:
+        eta = _format_eta((total - done) / rate)
+    else:
+        eta = "—"
+    return (
+        f"[{bar}] {int(pct * 100):3d}% "
+        f"elapsed {_format_eta(elapsed)} ETA {eta}"
+    )
+
+
+_progress_writer: "Callable[[str], None] | None" = None
+
+
+def set_progress_writer(fn: "Callable[[str], None] | None") -> None:
+    """Register a sink that receives every progress line. Used by the
+    daemon to forward progress to a file the CLI can tail; in-process
+    callers (build_index_with_store invoked directly) can leave it unset."""
+    global _progress_writer
+    _progress_writer = fn
+
+
+def emit_status(line: str) -> None:
+    """One-shot status emit to the progress side channel. Use for phases
+    without per-item progress (e.g. post-batch finalization) so the CLI's
+    bar viewer doesn't appear frozen between the last LLM result and
+    'Index written to …'."""
+    if _progress_writer is not None:
+        try:
+            _progress_writer(line)
+        except Exception:
+            pass
+
+
+def _emit_progress(label: str, done: int, total: int, elapsed: float, *, is_tty: bool, final: bool = False) -> None:
+    line = f"[{label}] {_render_progress_bar(done, total, elapsed)}"
+    if _progress_writer is not None:
+        try:
+            _progress_writer(line)
+        except Exception:
+            pass
+    if is_tty and not final:
+        sys.stderr.write("\r\033[K" + line)
+        sys.stderr.flush()
+    else:
+        # Non-TTY: emit on milestones so logs aren't spammed but still show motion.
+        step = max(1, total // 20)
+        if final or done == total or done % step == 0:
+            print(line, file=sys.stderr)
+
+
+def submit_and_wait_batch(
+    requests: list[dict],
+    api_key: str,
+    label: str = "batch",
+    *,
+    workers: int = REALTIME_MAX_WORKERS,
+) -> tuple[dict[str, str], LLMUsage]:
+    """Fan out requests concurrently through the LiteLLM adapter.
+
+    Returns ({custom_id: response_text}, aggregated_usage). Usage covers
+    every successful call across the pool; failed-after-retries calls drop
+    out (their attempts' usage is also lost — this matches the existing
+    "results dict" behaviour).
+    """
+    requested_pool = max(1, min(workers, len(requests)))
+    fd_cap = _ensure_fd_capacity(requested_pool)
+    pool_size = min(requested_pool, fd_cap)
+    if pool_size < requested_pool:
+        _log.warning(
+            f"[{label}] fd hard limit too low for {requested_pool} workers; "
+            f"clamping to {pool_size}. Raise `ulimit -n` to use more.",
+            extra={
+                "event": "batch.fd_clamp",
+                "label": label,
+                "requested": requested_pool,
+                "clamped_to": pool_size,
+            },
+        )
+    # Heuristic: ~5s/request amortized over the worker pool. The inline bar refines this once real rates are known.
+    est_seconds = (len(requests) / max(1, pool_size)) * 5.0
+    _log.info(
+        f"[{label}] Dispatching {len(requests)} requests "
+        f"(workers={pool_size}, est build time {_format_eta(est_seconds)})...",
+        extra={
+            "event": "batch.dispatch",
+            "label": label,
+            "requests": len(requests),
+            "workers": pool_size,
+            "est_seconds": round(est_seconds, 1),
+        },
+    )
+
+    def _one(req: dict) -> tuple[str, str | None, LLMUsage]:
+        cid = req["custom_id"]
+        last_err: Exception | None = None
+        attempt_usage = LLMUsage()
+        for attempt in range(REALTIME_MAX_RETRIES + 1):
+            try:
+                text, usage = call_llm(
+                    req["prompt"],
+                    model=req["model"],
+                    api_key=api_key,
+                    max_tokens=req["max_tokens"],
+                )
+                attempt_usage.add(usage)
+                return cid, text or "", attempt_usage
+            except litellm.AuthenticationError:
+                # Deterministic; let the outer loop cancel siblings.
+                raise
+            except Exception as e:  # network/rate-limit/5xx
+                last_err = e
+                if attempt < REALTIME_MAX_RETRIES:
+                    time.sleep(RETRY_DELAY * (attempt + 1))
+        _log.error(
+            f"[{label}]   {cid}: failed after retries ({last_err})",
+            extra={
+                "event": "llm.failed",
+                "label": label,
+                "custom_id": cid,
+                "attempts": REALTIME_MAX_RETRIES + 1,
+                "error_type": type(last_err).__name__ if last_err else None,
+                "error": str(last_err) if last_err else None,
+            },
+        )
+        return cid, None, attempt_usage
+
+    results: dict[str, str] = {}
+    total_usage = LLMUsage()
+    succeeded = 0
+    failed = 0
+    total = len(requests)
+    start = time.time()
+    is_tty = sys.stderr.isatty()
+
+    with ThreadPoolExecutor(max_workers=pool_size) as pool:
+        # Materialize the futures list so we can cancel queued tasks on a
+        # fatal error. Default `with`-exit waits for everything to finish.
+        futures = [pool.submit(_one, r) for r in requests]
+        try:
+            for fut in as_completed(futures):
+                cid, text, usage = fut.result()
+                total_usage.add(usage)
+                if text is not None:
+                    results[cid] = text
+                    succeeded += 1
+                else:
+                    failed += 1
+                done = succeeded + failed
+                _emit_progress(label, done, total, time.time() - start, is_tty=is_tty)
+        except litellm.AuthenticationError:
+            # In-flight threads can't be interrupted, but cancelling drains
+            # the queue so the remaining ~(total − pool_size) requests don't
+            # fire after we know the key is rejected.
+            for f in futures:
+                f.cancel()
+            raise
+
+    if is_tty:
+        sys.stderr.write("\n")
+        sys.stderr.flush()
+    _log.info(
+        f"[{label}] Results: {succeeded} succeeded, {failed} failed",
+        extra={
+            "event": "batch.results",
+            "label": label,
+            "succeeded": succeeded,
+            "failed": failed,
+            "elapsed_s": round(time.time() - start, 3),
+            "input_tokens": total_usage.input_tokens,
+            "output_tokens": total_usage.output_tokens,
+            "cost_usd": round(total_usage.cost_usd, 6),
+        },
+    )
+    return results, total_usage
+
+
+def index_by_hash(
+    hash_to_content: dict[str, tuple[str, str, Language]],
+    *,
+    summary_cache: dict[str, dict[str, str]],
+    model: str,
+    api_key: str,
+    workers: int,
+    label: str,
+) -> tuple[dict[str, tuple[str, str]], LLMUsage]:
+    """Fan content-deduped requests out concurrently, parse + line-range-fix
+    responses, and merge into ``summary_cache``. Returns
+    ``(hash_to_parsed, usage)`` mapping each successful response's content
+    hash to its ``(oneliner, functions)`` tuple."""
+    requests = build_batch_requests(hash_to_content, model=model)
+    results, usage = submit_and_wait_batch(requests, api_key, label=label, workers=workers)
+    hash_to_parsed: dict[str, tuple[str, str]] = {}
+    for h, text in results.items():
+        oneliner, functions = parse_response(text)
+        content, _, lang = hash_to_content[h]
+        functions = fix_functions_for_language(lang, content, functions)
+        hash_to_parsed[h] = (oneliner, functions)
+        summary_cache[h] = {"oneliner": oneliner, "functions": functions}
+    return hash_to_parsed, usage

sourceindex/build/linerange/__init__.py

@@ -0,0 +1,24 @@
+"""Post-LLM line-range correction.
+
+The build pipeline calls ``fix_functions_for_language`` after the LLM
+returns its tier-2 entries. For Python we use the stdlib ``ast``; for
+other languages we fall through to tree-sitter. Both fixers replace the
+LLM's ``(Lstart-Lend)`` ranges with parser-truth ranges so the agent's
+Read calls land on real function boundaries.
+"""
+
+from ...lib.languages import Language
+from .python_ast import fix_python_line_ranges, fill_missing_python_classes
+from .treesitter import fix_line_ranges as _fix_treesitter_line_ranges
+
+
+def fix_functions_for_language(lang: Language, content: str, functions: str) -> str:
+    """Apply language-specific line-range fixes to LLM-emitted functions text.
+    No-op when ``functions`` is empty."""
+    if not functions:
+        return functions
+    if lang.key == "python":
+        functions = fix_python_line_ranges(content, functions)
+        functions = fill_missing_python_classes(content, functions)
+        return functions
+    return _fix_treesitter_line_ranges(lang.key, content, functions)

sourceindex/build/linerange/python_ast.py

@@ -0,0 +1,155 @@
+"""Python AST-grounded line-range fixer for tier-2 LLM output.
+
+Cheaper LLMs (gemini, mistral, ds-v4-flash) emit roughly correct names but
+fabricate line ranges. We replace LLM line numbers with AST ground truth
+whenever the qualified name matches. Description quality is unchanged.
+"""
+
+import ast
+import re
+
+
+_LINE_RANGE_RE = re.compile(r"\(L?\d+\s*-\s*L?\d+\)")
+_ENTRY_RE = re.compile(r"^([\w.:_<>]+)\s+\(L?\d+\s*-\s*L?\d+\)\s*:?")
+
+
+def _python_ast_truth(source: str) -> dict[str, tuple[int, int]]:
+    """{qualified_name: (start_lineno, end_lineno)} for every def in source.
+
+    Methods qualified as Class.method, nested funcs as Outer.inner AND
+    Outer.<locals>.inner so either naming convention matches."""
+    try:
+        tree = ast.parse(source)
+    except (SyntaxError, ValueError):
+        return {}
+
+    out: dict[str, tuple[int, int]] = {}
+
+    def walk(node, prefix: str = ""):
+        for child in ast.iter_child_nodes(node):
+            if isinstance(child, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
+                qname = f"{prefix}{child.name}" if prefix else child.name
+                out[qname] = (child.lineno, child.end_lineno or child.lineno)
+                walk(child, prefix=f"{qname}.")
+                if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                    walk(child, prefix=f"{qname}.<locals>.")
+
+    walk(tree)
+    return out
+
+
+def fix_python_line_ranges(source: str, functions_text: str) -> str:
+    """Replace LLM-emitted line ranges with AST ground truth for Python files.
+
+    Lines whose qualified name matches an AST def get their (Lstart-Lend)
+    rewritten to the real range. Unmatched names (likely hallucinations or
+    deeply nested helpers AST didn't enumerate) are left untouched — kept
+    so a downstream judge can flag them, not silently dropped."""
+    truth = _python_ast_truth(source)
+    if not truth:
+        return functions_text
+
+    out_lines: list[str] = []
+    for line in functions_text.splitlines():
+        m = _ENTRY_RE.match(line.strip())
+        if not m:
+            out_lines.append(line)
+            continue
+        qname = m.group(1).rstrip(".")
+        match_key = qname if qname in truth else None
+        if match_key is None:
+            for k in truth:
+                if k.endswith(qname) or qname.endswith(k):
+                    match_key = k
+                    break
+        if match_key is None:
+            out_lines.append(line)
+            continue
+        ls, le = truth[match_key]
+        new_line = _LINE_RANGE_RE.sub(f"(L{ls}-L{le})", line, count=1)
+        if match_key != qname:
+            new_line = new_line.replace(qname, match_key, 1)
+        out_lines.append(new_line)
+    return "\n".join(out_lines)
+
+
+def fill_missing_python_classes(source: str, functions_text: str) -> str:
+    """Defense-in-depth: inject class entries the LLM omitted but AST shows exist.
+
+    The current INDEX_PROMPT_TEMPLATE has class-specific examples and a
+    purpose-vs-behavior rubric, so gpt-5.4-nano + the in-tree prompt
+    produce full coverage on its own. This function is a no-op when the
+    LLM already lists every class. It's kept for two reasons: (a) cheaper
+    or differently-trained models may still drop class lines, (b) prompt
+    wording can drift in future commits and we don't want a regression to
+    silently lose class-level overview from the roadmap. Existing entries
+    are preserved verbatim — we only ADD lines, never modify."""
+    try:
+        tree = ast.parse(source)
+    except (SyntaxError, ValueError):
+        return functions_text
+
+    # Collect class defs with their method count + docstring oneliner.
+    class_defs: list[tuple[str, int, int, str]] = []  # (qname, ls, le, summary)
+
+    def walk_classes(node, prefix: str = ""):
+        for child in ast.iter_child_nodes(node):
+            if isinstance(child, ast.ClassDef):
+                qname = f"{prefix}{child.name}" if prefix else child.name
+                docstring = ast.get_docstring(child) or ""
+                first_line = (docstring.split("\n", 1)[0] or "").strip()
+                methods = sum(
+                    1 for c in child.body
+                    if isinstance(c, (ast.FunctionDef, ast.AsyncFunctionDef))
+                )
+                if first_line:
+                    summary = first_line[:90]
+                else:
+                    bases = [
+                        ast.unparse(b) if hasattr(ast, "unparse") else getattr(b, "id", "?")
+                        for b in child.bases
+                    ]
+                    base_part = f"({', '.join(bases)})" if bases else ""
+                    summary = f"Class {child.name}{base_part} with {methods} method(s)"
+                class_defs.append((
+                    qname,
+                    child.lineno,
+                    child.end_lineno or child.lineno,
+                    summary,
+                ))
+                walk_classes(child, prefix=f"{qname}.")
+            elif isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                walk_classes(child, prefix=f"{prefix}{child.name}.<locals>.")
+
+    walk_classes(tree)
+    if not class_defs:
+        return functions_text
+
+    # Find which class names already appear in the LLM output.
+    present = set()
+    for line in functions_text.splitlines():
+        m = _ENTRY_RE.match(line.strip())
+        if m:
+            present.add(m.group(1).rstrip("."))
+
+    # Inject missing classes immediately before the first method that
+    # belongs to them (so file order is preserved). Classes with no method
+    # yet listed get appended at the end.
+    out_lines = list(functions_text.splitlines())
+    for qname, ls, le, summary in class_defs:
+        if qname in present:
+            continue
+        new_entry = f"{qname} (L{ls}-L{le}): {summary}"
+        # Find first line whose qname starts with `qname.`
+        method_prefix = f"{qname}."
+        injected = False
+        for i, line in enumerate(out_lines):
+            m = _ENTRY_RE.match(line.strip())
+            if m and m.group(1).startswith(method_prefix):
+                out_lines.insert(i, new_entry)
+                injected = True
+                break
+        if not injected:
+            out_lines.append(new_entry)
+
+    return "\n".join(out_lines)