codebase-index 1.6.0__py3-none-any.whl

codebase_index/parsers/line_chunker.py

@@ -0,0 +1,39 @@
+"""Fallback chunker: overlapping fixed-size line windows."""
+
+from __future__ import annotations
+
+from .base import Chunk
+
+_CHARS_PER_TOKEN = 4
+
+
+def estimate_tokens(text: str) -> int:
+    return max(1, round(len(text) / _CHARS_PER_TOKEN))
+
+
+def chunk_text(text: str, *, window_lines: int, overlap_lines: int) -> list[Chunk]:
+    if not text or not text.strip():
+        return []
+    if overlap_lines >= window_lines:
+        overlap_lines = window_lines - 1
+    stride = window_lines - overlap_lines
+
+    lines = text.splitlines()
+    chunks: list[Chunk] = []
+    start = 0
+    while start < len(lines):
+        end = min(start + window_lines, len(lines))
+        body = "\n".join(lines[start:end])
+        chunks.append(
+            Chunk(
+                line_start=start + 1,
+                line_end=end,
+                content=body,
+                token_est=estimate_tokens(body),
+                kind="window",
+            )
+        )
+        if end >= len(lines):
+            break
+        start += stride
+    return chunks

codebase_index/parsers/symbol_chunks.py

@@ -0,0 +1,62 @@
+"""Symbol-aligned chunking with fallback line windows."""
+
+from __future__ import annotations
+
+from .base import Chunk, Symbol
+from .line_chunker import chunk_text, estimate_tokens
+
+_GAP_WINDOW = 80
+_GAP_OVERLAP = 0
+
+
+def build_chunks(text: str, symbols: list[Symbol]) -> list[Chunk]:
+    if not text.strip():
+        return []
+    if not symbols:
+        return chunk_text(text, window_lines=80, overlap_lines=10)
+
+    lines = text.splitlines()
+    top = sorted(
+        [s for s in symbols if s.parent_index is None],
+        key=lambda s: s.line_start,
+    )
+
+    chunks: list[Chunk] = []
+    cursor = 1
+    for symbol in top:
+        symbol_index = symbols.index(symbol)
+        if symbol.line_start > cursor:
+            chunks.extend(_gap(lines, cursor, symbol.line_start - 1))
+        body = "\n".join(lines[symbol.line_start - 1 : symbol.line_end])
+        chunks.append(
+            Chunk(
+                line_start=symbol.line_start,
+                line_end=symbol.line_end,
+                content=body,
+                token_est=estimate_tokens(body),
+                kind="symbol_body",
+                symbol_index=symbol_index,
+            )
+        )
+        cursor = max(cursor, symbol.line_end + 1)
+    if cursor <= len(lines):
+        chunks.extend(_gap(lines, cursor, len(lines)))
+    return chunks
+
+
+def _gap(lines: list[str], start: int, end: int) -> list[Chunk]:
+    segment = "\n".join(lines[start - 1 : end])
+    if not segment.strip():
+        return []
+    out: list[Chunk] = []
+    for chunk in chunk_text(segment, window_lines=_GAP_WINDOW, overlap_lines=_GAP_OVERLAP):
+        out.append(
+            Chunk(
+                line_start=start + chunk.line_start - 1,
+                line_end=start + chunk.line_end - 1,
+                content=chunk.content,
+                token_est=chunk.token_est,
+                kind="window",
+            )
+        )
+    return out

codebase_index/parsers/treesitter.py

@@ -0,0 +1,439 @@
+"""Tree-sitter parsing: text -> symbols, intra-file call edges, and chunks."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from tree_sitter import Parser, Query, QueryCursor
+from tree_sitter_language_pack import get_language
+
+from .base import Edge, ParseResult, Symbol
+from .languages import CONTAINER_KINDS, has_grammar, spec_for
+from .symbol_chunks import build_chunks
+
+
+class UnsupportedLanguage(Exception):
+    pass
+
+
+def parse_file(lang: str, text: str) -> ParseResult:
+    spec = spec_for(lang)
+    # Tier A (hand-tuned spec) or Tier B (any loadable grammar). Only raise when no grammar
+    # exists at all, so "any language with a grammar" produces symbols, not a silent fallback.
+    ts_name = spec.ts_name if spec is not None else lang
+    if spec is None and not has_grammar(lang):
+        raise UnsupportedLanguage(lang)
+
+    grammar = get_language(ts_name)
+    parser = Parser(grammar)
+    source = text.encode("utf-8")
+    tree = parser.parse(source)
+    if tree is None:
+        raise ValueError("tree-sitter parser returned no tree")
+    root = tree.root_node
+
+    if spec is not None:
+        symbols = _extract_symbols(root, lang, source)
+    else:
+        symbols = _extract_symbols_generic(root, source)
+    edges = _extract_edges(root, symbols, source)
+    if spec is not None:
+        edges.extend(_extract_graph_edges(spec, grammar, root, symbols))
+    del grammar
+    chunks = build_chunks(text, symbols)
+    return ParseResult(chunks=chunks, symbols=symbols, edges=edges)
+
+
+def _row(point) -> int:
+    return point.row if hasattr(point, "row") else point[0]
+
+
+def _text(node) -> str:
+    raw = getattr(node, "text", None)
+    if callable(raw):
+        raw = raw()
+    if isinstance(raw, bytes):
+        return raw.decode("utf-8", errors="ignore")
+    return raw if isinstance(raw, str) else ""
+
+
+class _Sym:
+    __slots__ = ("symbol", "start_byte", "end_byte")
+
+    def __init__(self, symbol: Symbol, def_node) -> None:
+        self.symbol = symbol
+        self.start_byte, self.end_byte = _byte_range(def_node)
+
+
+def _extract_symbols(root, lang: str, source: bytes) -> list[Symbol]:
+    raw: list[_Sym] = []
+    for def_node in _walk(root):
+        kind = _definition_kind(def_node, lang)
+        if kind is None:
+            continue
+        name_node = _name_node(def_node)
+        if name_node is None:
+            continue
+        raw.append(
+            _Sym(
+                Symbol(
+                    name=_node_text(name_node, source),
+                    kind=kind,
+                    line_start=_row(_start_point(def_node)) + 1,
+                    line_end=_row(_end_point(def_node)) + 1,
+                    signature=_signature(def_node, source),
+                    docstring=_python_docstring(def_node, source) if lang == "python" else None,
+                ),
+                def_node,
+            )
+        )
+
+    raw.sort(key=lambda item: (item.start_byte, -(item.end_byte - item.start_byte)))
+    for item in raw:
+        parent = _enclosing(raw, item)
+        if parent is None:
+            item.symbol.qualified = item.symbol.name
+            continue
+        item.symbol.parent_index = raw.index(parent)
+        if item.symbol.kind == "function" and parent.symbol.kind in CONTAINER_KINDS:
+            item.symbol.kind = "method"
+        item.symbol.qualified = f"{parent.symbol.qualified or parent.symbol.name}.{item.symbol.name}"
+    return [item.symbol for item in raw]
+
+
+# Maps a tree-sitter node `type` to a coarse symbol kind. Node types are largely unique across
+# grammars, so a single table covers Java/Go/Rust/C/C++/C#/Ruby/PHP/Kotlin/JS/TS at once.
+_DEF_KINDS: dict[str, str] = {
+    # functions
+    "function_declaration": "function",  # go, kotlin, js
+    "function_definition": "function",  # c, cpp, php  (python handled separately below)
+    "function_item": "function",  # rust
+    "function_signature_item": "function",  # rust trait method signatures
+    # methods
+    "method_declaration": "method",  # go, java, csharp
+    "method_definition": "method",  # js/ts
+    "constructor_declaration": "method",  # java, csharp
+    "method": "method",  # ruby
+    # classes / type-like containers
+    "class_declaration": "class",  # java, csharp, php, kotlin, js/ts
+    "class_specifier": "class",  # cpp
+    "class": "class",  # ruby
+    "object_declaration": "class",  # kotlin
+    "record_declaration": "record",  # java
+    "struct_item": "struct",  # rust
+    "struct_specifier": "struct",  # c, cpp
+    "struct_declaration": "struct",  # csharp
+    "interface_declaration": "interface",  # java, csharp, php, ts
+    "trait_item": "trait",  # rust
+    "trait_declaration": "trait",  # php
+    "enum_declaration": "enum",  # java, csharp, ts
+    "enum_item": "enum",  # rust
+    "enum_specifier": "enum",  # c/cpp
+    "impl_item": "impl",  # rust
+    # modules / namespaces (NOT containers — a function inside stays a function, not a method)
+    "mod_item": "module",  # rust
+    "module": "module",  # ruby
+    "namespace_definition": "module",  # cpp
+    "namespace_declaration": "module",  # csharp
+    "type_alias_declaration": "type",  # ts
+}
+
+
+def _definition_kind(node, lang: str) -> Optional[str]:
+    kind = _kind(node)
+    if lang == "python":
+        if kind == "function_definition":
+            return "function"
+        if kind == "class_definition":
+            return "class"
+        return None
+    if kind == "type_spec":  # go: refine struct/interface from the underlying type
+        underlying = _field(node, "type")
+        u = _kind(underlying) if underlying is not None else ""
+        if u == "struct_type":
+            return "struct"
+        if u == "interface_type":
+            return "interface"
+        return "type"
+    mapped = _DEF_KINDS.get(kind)
+    if mapped is not None:
+        return mapped
+    if kind == "variable_declarator":
+        value = _field(node, "value")
+        if value is not None and _kind(value) in {"arrow_function", "function_expression"}:
+            return "function"
+    return None
+
+
+# Identifier-like node types that can carry a definition's name across grammars.
+_NAME_NODE_TYPES = {
+    "identifier",
+    "type_identifier",
+    "field_identifier",
+    "property_identifier",
+    "simple_identifier",
+    "constant",
+    "name",
+    "namespace_identifier",
+}
+
+
+def _name_node(def_node):
+    """Find the name node for a definition, tolerating grammars without a "name" field.
+
+    Handles: field "name" (most), Rust `impl_item` (field "type"), C/C++ function definitions
+    (name nested under the declarator), and fieldless grammars (Kotlin) via a child scan.
+    """
+    named = _field(def_node, "name")
+    if named is not None:
+        return named
+    kind = _kind(def_node)
+    if kind == "impl_item":
+        return _field(def_node, "type")
+    if kind == "function_definition":  # c / cpp: descend the declarator chain
+        decl = _field(def_node, "declarator")
+        return _declarator_identifier(decl) if decl is not None else None
+    for child in _named_children(def_node):
+        if _kind(child) in _NAME_NODE_TYPES:
+            return child
+    return None
+
+
+def _declarator_identifier(node):
+    if node is None:
+        return None
+    if _kind(node) in {"identifier", "field_identifier"}:
+        return node
+    inner = _field(node, "declarator")
+    if inner is not None:
+        return _declarator_identifier(inner)
+    for child in _named_children(node):
+        found = _declarator_identifier(child)
+        if found is not None:
+            return found
+    return None
+
+
+def _extract_symbols_generic(root, source: bytes) -> list[Symbol]:
+    """Tier B: harvest definition-like nodes from an untuned grammar.
+
+    Any node whose `type` ends in declaration/definition/_item/_specifier and that has an
+    identifier-like named child is treated as a symbol; kind is a coarse keyword mapping.
+    """
+    raw: list[_Sym] = []
+    for node in _walk(root):
+        ntype = _kind(node)
+        if not ntype.endswith(("declaration", "definition", "_item", "_specifier")):
+            continue
+        name_node = _name_node(node)
+        if name_node is None:
+            continue
+        raw.append(
+            _Sym(
+                Symbol(
+                    name=_node_text(name_node, source),
+                    kind=_generic_kind(ntype),
+                    line_start=_row(_start_point(node)) + 1,
+                    line_end=_row(_end_point(node)) + 1,
+                    signature=_signature(node, source),
+                ),
+                node,
+            )
+        )
+    raw.sort(key=lambda item: (item.start_byte, -(item.end_byte - item.start_byte)))
+    for item in raw:
+        parent = _enclosing(raw, item)
+        if parent is None:
+            item.symbol.qualified = item.symbol.name
+            continue
+        item.symbol.parent_index = raw.index(parent)
+        if item.symbol.kind == "function" and parent.symbol.kind in CONTAINER_KINDS:
+            item.symbol.kind = "method"
+        item.symbol.qualified = (
+            f"{parent.symbol.qualified or parent.symbol.name}.{item.symbol.name}"
+        )
+    return [item.symbol for item in raw]
+
+
+def _generic_kind(ntype: str) -> str:
+    low = ntype.lower()
+    for key in ("class", "struct", "enum", "interface", "trait", "module", "namespace"):
+        if key in low:
+            return "struct" if key == "struct" else key
+    return "function"
+
+
+def _signature(def_node, source: bytes) -> str:
+    return _node_text(def_node, source).splitlines()[0].strip().rstrip("{").strip()
+
+
+def _python_docstring(def_node, source: bytes) -> Optional[str]:
+    body = _field(def_node, "body")
+    if body is None:
+        return None
+    for stmt in _named_children(body):
+        if _kind(stmt) == "string":
+            return _node_text(stmt, source).strip().strip('"').strip("'").strip()
+        if _kind(stmt) == "expression_statement":
+            children = _named_children(stmt)
+            if children and _kind(children[0]) == "string":
+                return _node_text(children[0], source).strip().strip('"').strip("'").strip()
+        break
+    return None
+
+
+def _enclosing(raw: list[_Sym], child: _Sym) -> Optional[_Sym]:
+    best: Optional[_Sym] = None
+    for other in raw:
+        if other is child:
+            continue
+        if other.start_byte <= child.start_byte and other.end_byte >= child.end_byte:
+            other_span = other.end_byte - other.start_byte
+            child_span = child.end_byte - child.start_byte
+            if other_span <= child_span:
+                continue
+            if best is None or other_span < best.end_byte - best.start_byte:
+                best = other
+    return best
+
+
+def _extract_edges(root, symbols: list[Symbol], source: bytes) -> list[Edge]:
+    edges: list[Edge] = []
+    for node in _walk(root):
+        callee = _callee_node(node)
+        if callee is None:
+            continue
+        line = _row(_start_point(callee)) + 1
+        edges.append(
+            Edge(
+                edge_type="call",
+                callee_name=_node_text(callee, source),
+                line=line,
+                src_symbol_index=_enclosing_symbol_index(symbols, line),
+            )
+        )
+    return edges
+
+
+_EDGE_PREFIXES = {"import.": "import", "extends.": "extends", "implements.": "implements"}
+
+
+def _extract_graph_edges(spec, grammar, root, symbols) -> list[Edge]:
+    if not spec.imports_query:
+        return []
+    query = Query(grammar, spec.imports_query)
+    cursor = QueryCursor(query)
+    edges: list[Edge] = []
+    for _pattern_idx, captures in cursor.matches(root):
+        for capture_name, nodes in captures.items():
+            for node in nodes:
+                edge_type = next(
+                    (et for pfx, et in _EDGE_PREFIXES.items() if capture_name.startswith(pfx)),
+                    None,
+                )
+                if edge_type is None:
+                    continue
+                line = _row(node.start_point) + 1
+                src_idx = None if edge_type == "import" else _enclosing_symbol_index(symbols, line)
+                edges.append(Edge(
+                    edge_type=edge_type,
+                    callee_name=_text(node).strip().strip('"').strip("'"),
+                    line=line,
+                    src_symbol_index=src_idx,
+                ))
+    return edges
+
+
+def _enclosing_symbol_index(symbols: list[Symbol], line: int) -> Optional[int]:
+    best_idx: Optional[int] = None
+    best_span: Optional[int] = None
+    for idx, symbol in enumerate(symbols):
+        if symbol.line_start <= line <= symbol.line_end:
+            span = symbol.line_end - symbol.line_start
+            if best_span is None or span < best_span:
+                best_idx = idx
+                best_span = span
+    return best_idx
+
+
+_CALLEE_LEAVES = {"identifier", "property_identifier", "field_identifier", "simple_identifier"}
+
+
+def _callee_node(node):
+    kind = _kind(node)
+    if kind == "method_invocation":  # java: obj.method(...) / method(...)
+        return _field(node, "name")
+    if kind == "macro_invocation":  # rust: name!(...)
+        return _field(node, "macro")
+    if kind not in {"call", "call_expression", "invocation_expression", "function_call_expression"}:
+        return None
+    # ruby `call` uses field "method"; everything else uses field "function".
+    fn = _field(node, "function") or _field(node, "method")
+    if fn is None:
+        return None
+    if _kind(fn) in _CALLEE_LEAVES:
+        return fn
+    # member / selector / scoped / field access: take the trailing identifier.
+    attr = (
+        _field(fn, "attribute")
+        or _field(fn, "property")
+        or _field(fn, "field")
+        or _field(fn, "name")
+    )
+    if attr is not None and _kind(attr) in _CALLEE_LEAVES:
+        return attr
+    return None
+
+
+def _kind(node) -> str:
+    value = getattr(node, "type", None)
+    if value is None:
+        value = getattr(node, "kind", None)
+    resolved = value() if callable(value) else value
+    return resolved if isinstance(resolved, str) else ""
+
+
+def _field(node, name: str):
+    return node.child_by_field_name(name)
+
+
+def _start_point(node):
+    value = getattr(node, "start_point", None)
+    if value is None:
+        value = getattr(node, "start_position", None)
+    return value() if callable(value) else value
+
+
+def _end_point(node):
+    value = getattr(node, "end_point", None)
+    if value is None:
+        value = getattr(node, "end_position", None)
+    return value() if callable(value) else value
+
+
+def _byte_range(node) -> tuple[int, int]:
+    start = getattr(node, "start_byte", None)
+    end = getattr(node, "end_byte", None)
+    if start is not None and end is not None:
+        return (start() if callable(start) else start, end() if callable(end) else end)
+    br = node.byte_range()
+    return br.start, br.end
+
+
+def _node_text(node, source: bytes) -> str:
+    start, end = _byte_range(node)
+    return source[start:end].decode("utf-8", errors="ignore")
+
+
+def _named_children(node) -> list[object]:
+    children = getattr(node, "named_children", None)
+    if children is not None:
+        return list(children() if callable(children) else children)
+    count = node.named_child_count() if callable(node.named_child_count) else node.named_child_count
+    return [node.named_child(i) for i in range(count)]
+
+
+def _walk(node):
+    yield node
+    for child in _named_children(node):
+        yield from _walk(child)

codebase_index/retrieval/__init__.py

@@ -0,0 +1,9 @@
+"""Hybrid retrieval engine. See docs/RETRIEVAL.md for the full pipeline.
+
+intent.py    : classify the query into an Intent + retriever weights + graph strategy.
+searchers.py : path / symbol / fts / vector searchers -> uniform Candidate lists.
+fusion.py    : Reciprocal Rank Fusion across retriever lists (rrf_k, per-intent weights).
+rerank.py    : feature-based reordering (symbol-kind, path proximity, centrality, recency) +
+               produces the human-readable `reason` per result.
+budget.py    : greedy token-budgeted assembly of snippets vs. recommended_reads; secret redaction.
+"""

codebase_index/retrieval/budget.py

@@ -0,0 +1,82 @@
+"""Greedy token budgeting (RETRIEVAL.md §6).
+
+Metadata for every result is always emitted (cheap). Snippets are attached to the
+highest-ranked results until the budget is hit; the remainder become
+recommended_reads. All snippet text is secret-redacted before emission.
+
+A result is added to recommended_reads when:
+  - it has no snippet (budget exceeded or no content), OR
+  - its snippet is below _MIN_USEFUL_TOKENS (e.g. a bare function signature).
+    Claude still gets the short preview but also receives the read plan.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Optional
+
+from ..output.redact import redact_snippet
+from .skeleton import Compacted
+from .types import Candidate
+
+# Snippets shorter than this threshold are treated as previews only; the result
+# is still added to recommended_reads so Claude knows where to read the full body.
+_MIN_USEFUL_TOKENS = 40
+
+
+def _meta(c: Candidate) -> dict:
+    return {
+        "path": c.path,
+        "line_start": c.line_start,
+        "line_end": c.line_end,
+        "symbols": [c.symbol] if c.symbol else [],
+        "score": round(c.score, 4),
+        "reason": c.reason if c.reason else c.source,
+        "token_est": c.token_est,
+    }
+
+
+def apply_budget(
+    candidates: list[Candidate],
+    *,
+    token_budget: int,
+    compactor: Optional[Callable[[Candidate], Compacted]] = None,
+) -> tuple[list[dict], list[dict]]:
+    results: list[dict] = []
+    recommended: list[dict] = []
+    spent = 0
+
+    for rank, c in enumerate(candidates, start=1):
+        meta = _meta(c)
+        meta["rank"] = rank
+        meta["skeletonized"] = False
+        meta["elided_lines"] = 0
+
+        # Resolve the snippet text + cost. A compactor only changes anything
+        # when it returns a real skeleton; otherwise we keep today's raw path
+        # byte-for-byte (uses c.content / c.token_est).
+        text = c.content
+        cost = c.token_est
+        if compactor is not None and c.content:
+            comp = compactor(c)
+            if comp.skeletonized:
+                text = comp.text
+                cost = comp.token_est
+                meta["skeletonized"] = True
+                meta["elided_lines"] = comp.elided_lines
+
+        snippet = None
+        snippet_is_useful = False
+        if text and spent + cost <= token_budget:
+            snippet = redact_snippet(text)
+            spent += cost
+            meta["token_est"] = cost
+            snippet_is_useful = cost >= _MIN_USEFUL_TOKENS
+
+        if not snippet_is_useful:
+            recommended.append(
+                {"path": c.path, "line_start": c.line_start, "line_end": c.line_end}
+            )
+        meta["snippet"] = snippet
+        results.append(meta)
+
+    return results, recommended

codebase_index/retrieval/fusion.py

@@ -0,0 +1,62 @@
+"""Reciprocal Rank Fusion across per-source ranked candidate lists.
+
+RRF(d) = Σ_r  w_r · k / (k + rank_r(d))   — robust to incomparable raw scores.
+
+Two deliberate departures from the textbook formula:
+
+* Scaled by k. Raw RRF tops out at w/k (≈0.017 for k=60), an order of magnitude
+  below the bounded bonuses the reranker layers on top, so rerank would silently
+  become the primary ranker and RRF a mere tiebreak. Multiplying by k is a pure
+  monotonic rescale (fusion order is identical) that lifts the top contribution to
+  ≈w, putting fused scores and rerank bonuses on the same O(1) scale.
+* Fused on a coarse (path, line-bucket) key, not (path, start, end). Different
+  retrievers report different line ranges for the same place; an exact key almost
+  never coincides across sources, so cross-source agreement — RRF's whole point —
+  would never fire. `agreeing_sources` is therefore counted at file granularity.
+
+On merge, the candidate carrying the most signal (symbol > fts > path) is kept as
+the representative so downstream rerank/snippet logic has the richest fields.
+"""
+
+from __future__ import annotations
+
+from dataclasses import replace as _replace
+
+from .types import Candidate
+
+_SOURCE_RICHNESS = {"symbol": 3, "fts": 2, "vector": 2, "path": 1}
+
+
+def _richer(a: Candidate, b: Candidate) -> Candidate:
+    return a if _SOURCE_RICHNESS.get(a.source, 0) >= _SOURCE_RICHNESS.get(b.source, 0) else b
+
+
+def fuse(
+    lists: dict[str, list[Candidate]],
+    *,
+    weights: dict[str, float],
+    k: int,
+) -> list[Candidate]:
+    accum: dict[tuple, float] = {}
+    rep: dict[tuple, Candidate] = {}
+    seen: set[tuple] = set()
+    file_sources: dict[str, set[str]] = {}
+
+    for source, candidates in lists.items():
+        w = weights.get(source, 0.0)
+        if w <= 0.0:
+            continue
+        for rank, cand in enumerate(candidates):
+            file_sources.setdefault(cand.path, set()).add(source)
+            key = cand.fuse_key()
+            # One contribution per source per locator: a file matching three FTS
+            # chunks in the same bucket is one lexical signal, not three.
+            if (source, key) in seen:
+                continue
+            seen.add((source, key))
+            accum[key] = accum.get(key, 0.0) + w * k / (k + rank)
+            rep[key] = _richer(rep[key], cand) if key in rep else cand
+
+    fused = [_replace(rep[key], score=score) for key, score in accum.items()]
+    fused.sort(key=lambda c: c.score, reverse=True)
+    return [_replace(c, agreeing_sources=len(file_sources[c.path])) for c in fused]