coderag-ai 0.3.0__py3-none-any.whl

coderag/parsing/python.py

@@ -0,0 +1,294 @@
+"""Tree-sitter based Python parser.
+
+Extracts module/class/function/method symbols with precise line ranges,
+signatures, docstrings, and full-text search terms. Chunk boundaries always
+follow constructs (ADR-002). Relationship extraction is added in Phase 4 via
+``extract_relationships`` on the produced tree.
+"""
+
+from __future__ import annotations
+
+import tree_sitter_python as tspython
+from tree_sitter import Language, Parser
+
+from coderag.core.text import build_search_terms, extract_identifiers, split_identifier
+from coderag.parsing.base import (
+    CALLS,
+    CLASS,
+    FUNCTION,
+    IMPORTS,
+    INHERITS,
+    METHOD,
+    MODULE,
+    LanguageParser,
+    ParsedRelationship,
+    ParsedSymbol,
+    ParseResult,
+)
+
+_DEF_TYPES = {"function_definition", "class_definition", "decorated_definition"}
+
+
+class PythonParser(LanguageParser):
+    language = "python"
+    extensions = (".py",)
+
+    def __init__(self) -> None:
+        self._lang = Language(tspython.language(), "python")
+        self._parser = Parser()
+        self._parser.set_language(self._lang)
+
+    def parse(self, module_qualified_name: str, source: str) -> ParseResult:
+        src = source.encode("utf-8")
+        tree = self._parser.parse(src)
+        root = tree.root_node
+        symbols: list[ParsedSymbol] = []
+        relationships: list[ParsedRelationship] = []
+        counter = _Counter()
+
+        def text(node) -> str:
+            return src[node.start_byte : node.end_byte].decode("utf-8", "replace")
+
+        # -- module symbol: top-level code excluding class/function bodies ----
+        header_nodes = [c for c in root.named_children if c.type not in _DEF_TYPES]
+        header_src = "\n".join(text(n) for n in header_nodes)
+        mod_end = header_nodes[-1].end_point[0] + 1 if header_nodes else 1
+        mod_name = module_qualified_name.rsplit(".", 1)[-1]
+        module_sym = ParsedSymbol(
+            local_id=counter.next(),
+            parent_local_id=None,
+            symbol_name=mod_name,
+            qualified_name=module_qualified_name,
+            symbol_type=MODULE,
+            start_line=1,
+            end_line=mod_end,
+            source_code=header_src,
+            signature=None,
+            docstring=_docstring_of(root, src),
+            search_terms=build_search_terms(
+                split_identifier(mod_name),
+                extract_identifiers(header_src),
+            ),
+        )
+        symbols.append(module_sym)
+
+        # -- module imports --------------------------------------------------
+        for node in header_nodes:
+            for target in _imports(node, src):
+                relationships.append(
+                    ParsedRelationship(module_sym.local_id, IMPORTS, target, confidence=0.9)
+                )
+
+        # -- recurse into definitions ----------------------------------------
+        def emit(node, parent_local_id: int, qual_prefix: str, enclosing_is_class: bool):
+            inner = _unwrap(node)
+            name_node = inner.child_by_field_name("name")
+            if name_node is None:
+                return
+            name = text(name_node)
+            qualified = f"{qual_prefix}.{name}"
+            is_class = inner.type == "class_definition"
+            symbol_type = CLASS if is_class else (METHOD if enclosing_is_class else FUNCTION)
+            body = inner.child_by_field_name("body")
+            signature = _signature(node, inner, body, src)
+            source_code = text(node)  # node may be the decorated wrapper
+            sym = ParsedSymbol(
+                local_id=counter.next(),
+                parent_local_id=parent_local_id,
+                symbol_name=name,
+                qualified_name=qualified,
+                symbol_type=symbol_type,
+                start_line=node.start_point[0] + 1,
+                end_line=node.end_point[0] + 1,
+                source_code=source_code,
+                signature=signature,
+                docstring=_docstring_of(inner, src),
+                search_terms=build_search_terms(
+                    split_identifier(name),
+                    [p for seg in qualified.split(".") for p in split_identifier(seg)],
+                    extract_identifiers(signature or ""),
+                    extract_identifiers(source_code),
+                ),
+            )
+            symbols.append(sym)
+
+            if is_class:
+                for base in _bases(inner, src):
+                    relationships.append(
+                        ParsedRelationship(sym.local_id, INHERITS, base, confidence=0.9)
+                    )
+            elif body is not None:
+                for target, via in _calls(body, src):
+                    relationships.append(
+                        ParsedRelationship(
+                            sym.local_id, CALLS, target,
+                            confidence=0.9 if via == "self" else 0.6,
+                            metadata={"via": via},
+                        )
+                    )
+
+            if body is not None:
+                for child in body.named_children:
+                    if child.type in _DEF_TYPES:
+                        emit(child, sym.local_id, qualified, enclosing_is_class=is_class)
+
+        for child in root.named_children:
+            if child.type in _DEF_TYPES:
+                emit(child, module_sym.local_id, module_qualified_name, enclosing_is_class=False)
+
+        return ParseResult(symbols=symbols, relationships=relationships)
+
+
+class _Counter:
+    def __init__(self) -> None:
+        self._n = -1
+
+    def next(self) -> int:
+        self._n += 1
+        return self._n
+
+
+def _unwrap(node):
+    """decorated_definition -> the inner function/class definition."""
+    if node.type == "decorated_definition":
+        for c in node.named_children:
+            if c.type in ("function_definition", "class_definition"):
+                return c
+    return node
+
+
+def _signature(outer, inner, body, src: bytes) -> str | None:
+    end = body.start_byte if body is not None else inner.end_byte
+    text = src[inner.start_byte : end].decode("utf-8", "replace").strip()
+    return text.rstrip(":").strip() or None
+
+
+def _docstring_of(scope_node, src: bytes) -> str | None:
+    """Return the cleaned docstring of a module/function/class node, if any."""
+    body = scope_node.child_by_field_name("body")
+    candidates = body.named_children if body is not None else scope_node.named_children
+    for child in candidates:
+        if child.type == "expression_statement" and child.named_children:
+            first = child.named_children[0]
+            if first.type == "string":
+                raw = src[first.start_byte : first.end_byte].decode("utf-8", "replace")
+                return _clean_docstring(raw)
+        # only the *first* statement can be a docstring
+        break
+    return None
+
+
+def _node_text(node, src: bytes) -> str:
+    return src[node.start_byte : node.end_byte].decode("utf-8", "replace")
+
+
+def _imports(node, src: bytes) -> list[str]:
+    """Return imported dotted names from an import/import-from statement."""
+    out: list[str] = []
+    if node.type == "import_statement":
+        for child in node.named_children:
+            if child.type in ("dotted_name", "aliased_import"):
+                dn = child if child.type == "dotted_name" else child.child_by_field_name("name")
+                if dn is not None:
+                    out.append(_node_text(dn, src))
+    elif node.type == "import_from_statement":
+        mod_node = node.child_by_field_name("module_name")
+        # module_name field may not be set on older grammars; fall back to first dotted_name
+        module = None
+        if mod_node is not None:
+            module = _node_text(mod_node, src)
+        else:
+            for child in node.named_children:
+                if child.type in ("dotted_name", "relative_import"):
+                    module = _node_text(child, src)
+                    break
+        # imported names are the dotted_name/aliased_import nodes after the module
+        names: list[str] = []
+        seen_module = False
+        for child in node.named_children:
+            if not seen_module:
+                if module and _node_text(child, src) == module:
+                    seen_module = True
+                continue
+            if child.type == "dotted_name":
+                names.append(_node_text(child, src))
+            elif child.type == "aliased_import":
+                nm = child.child_by_field_name("name")
+                if nm is not None:
+                    names.append(_node_text(nm, src))
+            elif child.type == "wildcard_import":
+                pass
+        if names and module:
+            out.extend(f"{module}.{n}" for n in names)
+        elif module:
+            out.append(module)
+    return out
+
+
+def _bases(class_inner, src: bytes) -> list[str]:
+    sup = class_inner.child_by_field_name("superclasses")
+    if sup is None:
+        return []
+    bases: list[str] = []
+    for child in sup.named_children:
+        if child.type == "identifier":
+            bases.append(_node_text(child, src))
+        elif child.type == "attribute":
+            attr = child.child_by_field_name("attribute")
+            if attr is not None:
+                bases.append(_node_text(attr, src))
+    return bases
+
+
+def _calls(body, src: bytes) -> list[tuple[str, str]]:
+    """Collect (target_name, via) call targets within a body.
+
+    Does not descend into nested definitions (their calls belong to them).
+    ``via`` is "self" for self.method(), else "attr" or "bare".
+    """
+    out: list[tuple[str, str]] = []
+    seen: set[tuple[str, str]] = set()
+
+    def visit(node):
+        for child in node.children:
+            if child.type in _DEF_TYPES:
+                continue  # nested symbol; handled separately
+            if child.type == "call":
+                fn = child.child_by_field_name("function")
+                hit = _call_target(fn, src)
+                if hit and hit not in seen:
+                    seen.add(hit)
+                    out.append(hit)
+            visit(child)
+
+    visit(body)
+    return out
+
+
+def _call_target(fn, src: bytes) -> tuple[str, str] | None:
+    if fn is None:
+        return None
+    if fn.type == "identifier":
+        return (_node_text(fn, src), "bare")
+    if fn.type == "attribute":
+        obj = fn.child_by_field_name("object")
+        attr = fn.child_by_field_name("attribute")
+        if attr is None:
+            return None
+        via = "self" if obj is not None and _node_text(obj, src) == "self" else "attr"
+        return (_node_text(attr, src), via)
+    return None
+
+
+def _clean_docstring(raw: str) -> str:
+    s = raw.strip()
+    # strip a leading string prefix (r, b, f, u and combinations)
+    i = 0
+    while i < len(s) and s[i] in "rbfuRBFU":
+        i += 1
+    s = s[i:]
+    for q in ('"""', "'''", '"', "'"):
+        if s.startswith(q) and s.endswith(q) and len(s) >= 2 * len(q):
+            s = s[len(q) : -len(q)]
+            break
+    return s.strip()

coderag/parsing/registry.py

@@ -0,0 +1,40 @@
+"""Parser selection by file extension, plus module-name derivation."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+
+from coderag.parsing.base import LanguageParser
+
+
+@lru_cache
+def _python_parser() -> LanguageParser:
+    from coderag.parsing.python import PythonParser
+
+    return PythonParser()
+
+
+def get_parser_for_path(path: str) -> LanguageParser | None:
+    from coderag.indexing.ignore import language_for_path
+
+    lang = language_for_path(path)
+    if lang == "python":
+        return _python_parser()
+    return None
+
+
+def module_qualified_name(rel_path: str) -> str:
+    """Derive a dotted module name from a repo-relative path.
+
+    ``services/payment_service.py`` -> ``services.payment_service``
+    ``services/__init__.py``        -> ``services``
+    ``main.py``                     -> ``main``
+    """
+    p = rel_path.replace("\\", "/")
+    dot = p.rfind(".")
+    if dot != -1:
+        p = p[:dot]
+    parts = [seg for seg in p.split("/") if seg]
+    if parts and parts[-1] == "__init__":
+        parts = parts[:-1]
+    return ".".join(parts) if parts else "root"

coderag/retrieval/base.py

@@ -0,0 +1,60 @@
+"""Retrieval value objects and the Retriever protocol.
+
+A ``Retriever`` returns ``RawHit``s (symbol id + a source-local score) ordered
+best-first. The engine fuses hits from all retrievers into ``Candidate``s, each
+of which records *why* it was retrieved (its reasons) and its per-source ranks —
+so retrieval is explainable, per the design principles.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Protocol
+
+from sqlalchemy.orm import Session
+
+# Source / reason labels.
+EXACT_SYMBOL = "exact_symbol"
+LEXICAL = "lexical"
+SEMANTIC = "semantic"
+GRAPH = "graph"  # refined to graph_call/graph_test/... in Phase 4
+
+
+@dataclass
+class RawHit:
+    symbol_id: int
+    score: float
+    # Optional richer reason (e.g. "graph_call"); defaults to the source label.
+    reason: str | None = None
+
+
+@dataclass
+class Candidate:
+    symbol_id: int
+    qualified_name: str
+    symbol_name: str
+    symbol_type: str
+    file_path: str
+    start_line: int
+    end_line: int
+    token_count: int
+    reasons: set[str] = field(default_factory=set)
+    source_ranks: dict[str, int] = field(default_factory=dict)
+    source_scores: dict[str, float] = field(default_factory=dict)
+    fused_score: float = 0.0
+    rerank_score: float | None = None  # set when a reranker runs (Phase 9)
+    # Populated lazily by the context builder when code is actually needed.
+    signature: str | None = None
+    docstring: str | None = None
+    source_code: str | None = None
+
+    def explain(self) -> str:
+        return ",".join(sorted(self.reasons))
+
+
+class Retriever(Protocol):
+    source: str
+
+    def retrieve(
+        self, session: Session, repository_id: int, query: str, limit: int
+    ) -> list[RawHit]: ...

coderag/retrieval/engine.py

@@ -0,0 +1,176 @@
+"""Hybrid retrieval engine.
+
+Runs the configured retrievers, fuses them with RRF, (optionally, from Phase 4)
+expands one graph hop, loads symbol metadata, and returns explainable
+``Candidate``s. Every query is scoped to a single ``repository_id`` — there is no
+retrieval path that ignores it (ADR-006).
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from coderag.core.config import Settings, get_settings
+from coderag.db.models import Symbol
+from coderag.retrieval.base import (
+    EXACT_SYMBOL,
+    GRAPH,
+    LEXICAL,
+    SEMANTIC,
+    Candidate,
+    RawHit,
+    Retriever,
+)
+from coderag.retrieval.fusion import reciprocal_rank_fusion
+from coderag.retrieval.lexical import LexicalRetriever
+from coderag.retrieval.symbol import SymbolRetriever
+
+
+@dataclass
+class RetrievalOutcome:
+    candidates: list[Candidate]
+    latency_ms: float
+    per_source_counts: dict[str, int]
+
+
+class RetrievalEngine:
+    def __init__(
+        self,
+        retrievers: list[Retriever],
+        settings: Settings | None = None,
+        graph_expander=None,
+        reranker=None,
+    ) -> None:
+        self.retrievers = retrievers
+        self.settings = settings or get_settings()
+        self.graph_expander = graph_expander  # set in Phase 4
+        self.reranker = reranker  # set in Phase 9 (optional)
+
+    def _weights(self) -> dict[str, float]:
+        s = self.settings
+        return {
+            EXACT_SYMBOL: s.weight_symbol,
+            LEXICAL: s.weight_lexical,
+            SEMANTIC: s.weight_semantic,
+            GRAPH: s.weight_graph,
+        }
+
+    def _source_limit(self, source: str) -> int:
+        s = self.settings
+        return {
+            EXACT_SYMBOL: s.symbol_limit,
+            LEXICAL: s.lexical_limit,
+            SEMANTIC: s.semantic_limit,
+        }.get(source, s.lexical_limit)
+
+    def search(
+        self, session: Session, repository_id: int, query: str, top_n: int | None = None
+    ) -> RetrievalOutcome:
+        started = time.perf_counter()
+        source_hits: dict[str, list[RawHit]] = {}
+        per_source_counts: dict[str, int] = {}
+        for retriever in self.retrievers:
+            hits = retriever.retrieve(
+                session, repository_id, query, self._source_limit(retriever.source)
+            )
+            source_hits[retriever.source] = hits
+            per_source_counts[retriever.source] = len(hits)
+
+        fused = reciprocal_rank_fusion(source_hits, self._weights(), self.settings.rrf_k)
+        fused = fused[: self.settings.max_candidates]
+
+        candidates = self._load_candidates(session, repository_id, fused)
+        # deterministic order BEFORE expansion so graph seeds are the true top-N
+        candidates.sort(key=lambda c: (-c.fused_score, c.symbol_id))
+
+        # Phase 4: bounded one-hop graph expansion (added candidates with graph reasons).
+        if self.graph_expander is not None and self.settings.graph_enabled:
+            candidates = self.graph_expander.expand(
+                session, repository_id, candidates, query
+            )
+
+        if self.reranker is not None:
+            candidates = self.reranker.rerank(session, repository_id, query, candidates)
+        else:
+            candidates.sort(key=lambda c: (-c.fused_score, c.symbol_id))
+        if top_n is not None:
+            candidates = candidates[:top_n]
+        latency_ms = (time.perf_counter() - started) * 1000
+        return RetrievalOutcome(
+            candidates=candidates,
+            latency_ms=latency_ms,
+            per_source_counts=per_source_counts,
+        )
+
+    def _load_candidates(self, session, repository_id, fused) -> list[Candidate]:
+        if not fused:
+            return []
+        by_id = {e.symbol_id: e for e in fused}
+        rows = session.execute(
+            select(
+                Symbol.id, Symbol.qualified_name, Symbol.symbol_name, Symbol.symbol_type,
+                Symbol.file_path, Symbol.start_line, Symbol.end_line, Symbol.token_count,
+                Symbol.signature, Symbol.docstring,
+            )
+            # repository scope is enforced here as well as in each retriever.
+            .where(Symbol.repository_id == repository_id)
+            .where(Symbol.id.in_(list(by_id)))
+        ).all()
+        candidates: list[Candidate] = []
+        for r in rows:
+            e = by_id[r.id]
+            candidates.append(
+                Candidate(
+                    symbol_id=r.id,
+                    qualified_name=r.qualified_name,
+                    symbol_name=r.symbol_name,
+                    symbol_type=r.symbol_type,
+                    file_path=r.file_path,
+                    start_line=r.start_line,
+                    end_line=r.end_line,
+                    token_count=r.token_count,
+                    reasons=set(e.reasons),
+                    source_ranks=dict(e.source_ranks),
+                    source_scores=dict(e.source_scores),
+                    fused_score=e.score,
+                    signature=r.signature,
+                    docstring=r.docstring,
+                )
+            )
+        return candidates
+
+
+def build_engine(
+    settings: Settings | None = None,
+    embedding_provider=None,
+    with_graph: bool = False,
+    with_reranker: bool = False,
+) -> RetrievalEngine:
+    """Construct an engine with the default retriever set.
+
+    The semantic retriever (Phase 3) is added when an embedding provider is
+    supplied; graph expansion (Phase 4) when ``with_graph`` is set.
+    """
+    settings = settings or get_settings()
+    retrievers: list[Retriever] = [SymbolRetriever(), LexicalRetriever()]
+    if embedding_provider is not None:
+        from coderag.retrieval.semantic import SemanticRetriever
+
+        retrievers.append(SemanticRetriever(embedding_provider, settings))
+    graph_expander = None
+    if with_graph:
+        from coderag.retrieval.graph import GraphExpander
+
+        graph_expander = GraphExpander(settings)
+    reranker = None
+    if with_reranker or settings.reranker_enabled:
+        from coderag.retrieval.reranking import get_reranker
+
+        reranker = get_reranker(settings)
+    return RetrievalEngine(
+        retrievers, settings=settings, graph_expander=graph_expander, reranker=reranker
+    )

coderag/retrieval/fusion.py

@@ -0,0 +1,41 @@
+"""Reciprocal Rank Fusion (RRF).
+
+RRF fuses ranked lists from incomparable sources without score normalisation:
+a hit at rank ``r`` (1-based) from a source contributes ``weight / (k + r)``.
+Contributions are summed across sources. ``k`` and the per-source weights come
+from configuration (ADR-003) — no magic constants here.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from coderag.retrieval.base import RawHit
+
+
+@dataclass
+class FusedEntry:
+    symbol_id: int
+    score: float = 0.0
+    reasons: set[str] = field(default_factory=set)
+    source_ranks: dict[str, int] = field(default_factory=dict)
+    source_scores: dict[str, float] = field(default_factory=dict)
+
+
+def reciprocal_rank_fusion(
+    source_hits: dict[str, list[RawHit]],
+    weights: dict[str, float],
+    k: int = 60,
+) -> list[FusedEntry]:
+    entries: dict[int, FusedEntry] = {}
+    for source, hits in source_hits.items():
+        weight = weights.get(source, 1.0)
+        for rank, hit in enumerate(hits, start=1):
+            entry = entries.setdefault(hit.symbol_id, FusedEntry(symbol_id=hit.symbol_id))
+            entry.score += weight / (k + rank)
+            entry.reasons.add(hit.reason or source)
+            # keep the best (lowest) rank per source
+            if source not in entry.source_ranks or rank < entry.source_ranks[source]:
+                entry.source_ranks[source] = rank
+                entry.source_scores[source] = hit.score
+    return sorted(entries.values(), key=lambda e: (-e.score, e.symbol_id))