codebase-retrieval-context-engine 2.0.0__py3-none-any.whl

corbell/core/query/formatter.py

@@ -0,0 +1,98 @@
+"""Output formatter for code retrieval results."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from corbell.core.query.graph_expander import ScoredChunk
+
+
+def format_results(
+    chunks: List["ScoredChunk"],
+    repo_paths: Dict[str, str],
+) -> str:
+    """Format scored chunks as annotated code blocks for LLM context injection.
+
+    Output format:
+        <absolute_path>#L<start>-<end>
+        <start>: <code line>
+        <start+1>: <code line>
+        ...
+        <end>: <code line>
+
+    Args:
+        chunks: Scored chunks to format (pre-sorted by score descending).
+        repo_paths: Mapping of repo_id -> absolute repo path string.
+                    Used to resolve relative file paths to absolute paths.
+
+    Returns:
+        Formatted string with all chunks, separated by blank lines.
+    """
+    if not chunks:
+        return ""
+
+    blocks: List[str] = []
+
+    for chunk in chunks:
+        abs_path = _resolve_absolute_path(chunk.file_path, chunk.repo_id, repo_paths)
+
+        # Read the actual lines for this chunk range
+        lines = _read_chunk_lines(abs_path, chunk.start_line, chunk.end_line)
+        if lines is None:
+            # File not readable — use content from chunk object
+            lines = chunk.content.splitlines()
+
+        # Build the header: path#Lstart-end
+        header = f"{abs_path}#L{chunk.start_line}-{chunk.end_line}"
+
+        # Build numbered lines
+        numbered_lines: List[str] = []
+        for i, line in enumerate(lines):
+            line_num = chunk.start_line + i
+            numbered_lines.append(f"{line_num}: {line}")
+
+        block = header + "\n" + "\n".join(numbered_lines)
+        blocks.append(block)
+
+    return "\n\n".join(blocks)
+
+
+def _resolve_absolute_path(
+    file_path: str,
+    repo_id: str,
+    repo_paths: Dict[str, str],
+) -> str:
+    """Resolve a file_path to an absolute path string.
+
+    If file_path is already absolute, returns it as-is.
+    Otherwise, joins it with the repo's root path.
+    """
+    p = Path(file_path)
+    if p.is_absolute():
+        return str(p)
+
+    repo_root = repo_paths.get(repo_id, "")
+    if repo_root:
+        return str((Path(repo_root) / file_path).resolve())
+
+    return file_path  # best effort
+
+
+def _read_chunk_lines(
+    file_path: str,
+    start_line: int,
+    end_line: int,
+) -> List[str] | None:
+    """Read lines start_line..end_line (1-based, inclusive) from a file.
+
+    Returns None if the file cannot be read.
+    """
+    try:
+        all_lines = Path(file_path).read_text(encoding="utf-8", errors="ignore").splitlines()
+        start_idx = max(0, start_line - 1)
+        end_idx = min(len(all_lines), end_line)
+        return all_lines[start_idx:end_idx]
+    except Exception:
+        return None

corbell/core/query/graph_expander.py

@@ -0,0 +1,284 @@
+"""Graph-based call-chain expansion for query results."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from corbell.core.query.diagnostics import QueryDiagnostics
+
+
+@dataclass
+class ScoredChunk:
+    """An embedding chunk with a relevance score."""
+
+    chunk_id: str
+    score: float
+    file_path: str      # absolute path
+    start_line: int
+    end_line: int
+    content: str
+    repo_id: str
+    symbol: Optional[str] = None
+    chunk_type: str = "block"
+    language: str = "python"
+
+
+def expand_via_graph(
+    embedding_results: List[ScoredChunk],
+    graph_store: Any,
+    repos: List,
+    max_depth: int = 2,
+    max_chunks: int = 30,
+    diagnostics: Optional["QueryDiagnostics"] = None,
+) -> List[ScoredChunk]:
+    """Expand embedding results by following call-chain edges in the graph.
+
+    For each embedding result that overlaps with a MethodNode (by file path +
+    line range), BFS is performed over callers (score * 0.6) and callees
+    (score * 0.5). A global cap of max_chunks expanded chunks is enforced.
+    Score floor of 0.15 stops expansion of low-confidence chains.
+
+    Args:
+        embedding_results: Initial scored chunks from embedding search.
+        graph_store: SQLiteGraphStore instance.
+        repos: List of RepoConfig objects (for resolving relative paths).
+        max_depth: BFS depth limit for expansion.
+        max_chunks: Maximum total expanded (bonus) chunks to add.
+        diagnostics: Optional diagnostics object for tracking warnings.
+
+    Returns:
+        List of bonus ScoredChunk objects (excluding the original results).
+    """
+    if not embedding_results:
+        return []
+
+    # Build repo_id → absolute path mapping
+    repo_path_map: Dict[str, Path] = {}
+    for repo in repos:
+        if repo.resolved_path:
+            repo_path_map[repo.id] = repo.resolved_path
+
+    bonus_chunks: List[ScoredChunk] = []
+    visited: Set[str] = set()  # visited method IDs
+
+    # Get all method nodes from the graph (cached as dict keyed by file+lines)
+    try:
+        all_services = graph_store.get_all_services()
+        service_ids = [s.id for s in all_services]
+    except Exception:
+        return []
+
+    for base_chunk in embedding_results:
+        if len(bonus_chunks) >= max_chunks:
+            break
+
+        # Find MethodNodes that overlap with this chunk's file+line range
+        matching_methods = _find_matching_methods(
+            base_chunk, graph_store, repo_path_map, service_ids
+        )
+
+        for method_node in matching_methods:
+            if len(bonus_chunks) >= max_chunks:
+                break
+            if method_node.id in visited:
+                continue
+
+            visited.add(method_node.id)
+            _bfs_expand(
+                method_node=method_node,
+                parent_score=base_chunk.score,
+                depth=0,
+                max_depth=max_depth,
+                graph_store=graph_store,
+                repo_path_map=repo_path_map,
+                bonus_chunks=bonus_chunks,
+                max_chunks=max_chunks,
+                visited=visited,
+                diagnostics=diagnostics,
+            )
+
+    return bonus_chunks
+
+
+def _find_matching_methods(
+    chunk: ScoredChunk,
+    graph_store: Any,
+    repo_path_map: Dict[str, Path],
+    service_ids: List[str],
+) -> List[Any]:
+    """Find MethodNodes whose file_path and line range overlap with the given chunk."""
+    results = []
+
+    try:
+        for service_id in service_ids:
+            methods = graph_store.get_methods_for_service(service_id)
+            for method in methods:
+                # Normalize method file_path to absolute
+                method_abs = Path(method.file_path)
+                chunk_abs = Path(chunk.file_path)
+
+                # Compare absolute paths
+                if not method_abs.is_absolute():
+                    repo_path = repo_path_map.get(service_id)
+                    if repo_path:
+                        method_abs = (repo_path / method.file_path).resolve()
+
+                if method_abs != chunk_abs:
+                    continue
+
+                # Check line overlap: method and chunk lines overlap
+                if method.line_end < chunk.start_line:
+                    continue
+                if method.line_start > chunk.end_line:
+                    continue
+
+                results.append(method)
+    except Exception:
+        pass
+
+    return results
+
+
+def _bfs_expand(
+    method_node: Any,
+    parent_score: float,
+    depth: int,
+    max_depth: int,
+    graph_store: Any,
+    repo_path_map: Dict[str, Path],
+    bonus_chunks: List[ScoredChunk],
+    max_chunks: int,
+    visited: Set[str],
+    diagnostics: Optional["QueryDiagnostics"],
+) -> None:
+    """BFS expansion from a method node, adding callers and callees as bonus chunks."""
+    if depth >= max_depth:
+        return
+    if len(bonus_chunks) >= max_chunks:
+        return
+
+    # Expand callers (score * 0.6)
+    try:
+        callers = graph_store.get_callers_of_method(method_node.id)
+        for caller in callers:
+            if len(bonus_chunks) >= max_chunks:
+                return
+            caller_score = parent_score * 0.6
+            if caller_score < 0.15:
+                continue
+            if caller.id in visited:
+                continue
+            visited.add(caller.id)
+
+            bonus = _method_to_scored_chunk(
+                caller, caller_score, repo_path_map, diagnostics
+            )
+            if bonus is not None:
+                bonus_chunks.append(bonus)
+                _bfs_expand(
+                    method_node=caller,
+                    parent_score=caller_score,
+                    depth=depth + 1,
+                    max_depth=max_depth,
+                    graph_store=graph_store,
+                    repo_path_map=repo_path_map,
+                    bonus_chunks=bonus_chunks,
+                    max_chunks=max_chunks,
+                    visited=visited,
+                    diagnostics=diagnostics,
+                )
+    except Exception:
+        if diagnostics:
+            diagnostics.graph_expansion_failures += 1
+
+    # Expand callees (score * 0.5)
+    try:
+        outgoing = graph_store.get_dependencies(method_node.id)
+        for edge in outgoing:
+            if edge.kind != "method_call":
+                continue
+            if len(bonus_chunks) >= max_chunks:
+                return
+            callee_score = parent_score * 0.5
+            if callee_score < 0.15:
+                continue
+
+            callee_id = edge.target_id
+            if callee_id in visited:
+                continue
+            visited.add(callee_id)
+
+            callee_node = graph_store.get_method(callee_id)
+            if callee_node is None:
+                if diagnostics:
+                    diagnostics.skipped_methods += 1
+                continue
+
+            bonus = _method_to_scored_chunk(
+                callee_node, callee_score, repo_path_map, diagnostics
+            )
+            if bonus is not None:
+                bonus_chunks.append(bonus)
+                _bfs_expand(
+                    method_node=callee_node,
+                    parent_score=callee_score,
+                    depth=depth + 1,
+                    max_depth=max_depth,
+                    graph_store=graph_store,
+                    repo_path_map=repo_path_map,
+                    bonus_chunks=bonus_chunks,
+                    max_chunks=max_chunks,
+                    visited=visited,
+                    diagnostics=diagnostics,
+                )
+    except Exception:
+        if diagnostics:
+            diagnostics.graph_expansion_failures += 1
+
+
+def _method_to_scored_chunk(
+    method_node: Any,
+    score: float,
+    repo_path_map: Dict[str, Path],
+    diagnostics: Optional["QueryDiagnostics"],
+) -> Optional[ScoredChunk]:
+    """Convert a MethodNode to a ScoredChunk by reading its source lines.
+
+    Returns None if the file doesn't exist (increments diagnostics counter).
+    """
+    file_path = Path(method_node.file_path)
+    if not file_path.is_absolute():
+        repo_path = repo_path_map.get(method_node.service_id)
+        if repo_path:
+            file_path = (repo_path / method_node.file_path).resolve()
+
+    if not file_path.exists():
+        if diagnostics:
+            diagnostics.skipped_files += 1
+        return None
+
+    try:
+        lines = file_path.read_text(encoding="utf-8", errors="ignore").splitlines()
+        start = max(0, method_node.line_start - 1)
+        end = min(len(lines), method_node.line_end)
+        content = "\n".join(lines[start:end])
+    except Exception:
+        if diagnostics:
+            diagnostics.skipped_files += 1
+        return None
+
+    return ScoredChunk(
+        chunk_id=method_node.id,
+        score=score,
+        file_path=str(file_path),
+        start_line=method_node.line_start,
+        end_line=method_node.line_end,
+        content=content,
+        repo_id=method_node.service_id,
+        symbol=method_node.method_name,
+        chunk_type="method",
+        language="python",  # best effort; MethodNode doesn't store language
+    )

corbell/core/query/merger.py

@@ -0,0 +1,171 @@
+"""Deduplication and adjacent chunk merging for query results."""
+
+from __future__ import annotations
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from corbell.core.query.graph_expander import ScoredChunk
+
+# Maximum number of lines in a merged chunk block
+_MAX_MERGED_LINES = 60
+
+
+def merge_and_dedup(chunks: List["ScoredChunk"]) -> List["ScoredChunk"]:
+    """Deduplicate by chunk_id (keep max score), group by file, merge adjacent chunks.
+
+    Steps:
+    1. Deduplicate: for duplicate chunk_ids, keep the one with the highest score.
+    2. Group chunks by file path.
+    3. Within each file, sort by start_line and merge adjacent/overlapping chunks.
+    4. Cap merged blocks at 60 lines.
+    5. Remove chunks whose line range is fully contained within another chunk in the same file.
+
+    Args:
+        chunks: List of ScoredChunk objects (may have duplicates from multiple queries).
+
+    Returns:
+        Deduplicated and merged list of ScoredChunk objects.
+    """
+    if not chunks:
+        return []
+
+    # Step 1: Dedup by chunk_id, keeping max score
+    best: Dict[str, "ScoredChunk"] = {}
+    for chunk in chunks:
+        if chunk.chunk_id not in best or chunk.score > best[chunk.chunk_id].score:
+            best[chunk.chunk_id] = chunk
+
+    deduped = list(best.values())
+
+    # Step 2: Group by file path
+    by_file: Dict[str, List["ScoredChunk"]] = {}
+    for chunk in deduped:
+        by_file.setdefault(chunk.file_path, []).append(chunk)
+
+    # Step 3+4: Sort and merge adjacent/overlapping chunks per file
+    result: List["ScoredChunk"] = []
+    for file_path, file_chunks in by_file.items():
+        merged = _merge_file_chunks(file_chunks)
+        # Step 5: Drop chunks fully contained within a larger chunk in the same file
+        merged = _drop_contained_ranges(merged)
+        result.extend(merged)
+
+    # Sort final result by score descending for output ordering
+    result.sort(key=lambda c: c.score, reverse=True)
+    return result
+
+
+def _merge_file_chunks(chunks: List["ScoredChunk"]) -> List["ScoredChunk"]:
+    """Merge adjacent/overlapping chunks within a single file.
+
+    Two chunks are merged if next.start_line <= current.end_line + 1.
+    The merged chunk gets the max score and a fresh chunk_id combining both.
+    Merged blocks are capped at 60 lines.
+    """
+    # Sort by start line
+    sorted_chunks = sorted(chunks, key=lambda c: c.start_line)
+
+    merged: List["ScoredChunk"] = []
+    if not sorted_chunks:
+        return merged
+
+    current = sorted_chunks[0]
+
+    for next_chunk in sorted_chunks[1:]:
+        # Check adjacency: next starts before or at current.end + 1
+        if next_chunk.start_line <= current.end_line + 1:
+            # Check if merging would exceed 60-line cap
+            merged_lines = next_chunk.end_line - current.start_line + 1
+            if merged_lines <= _MAX_MERGED_LINES:
+                # Merge: extend current to cover next_chunk
+                current = _merge_two(current, next_chunk)
+            else:
+                # Would exceed cap: flush current, start new
+                merged.append(current)
+                current = next_chunk
+        else:
+            # Not adjacent: flush current, start new
+            merged.append(current)
+            current = next_chunk
+
+    merged.append(current)
+    return merged
+
+
+def _drop_contained_ranges(chunks: List["ScoredChunk"]) -> List["ScoredChunk"]:
+    """Remove chunks whose line range is fully contained within another chunk's range.
+
+    A chunk B is considered contained in chunk A when:
+      A.start_line <= B.start_line and B.end_line <= A.end_line
+
+    When two chunks share the exact same range, the one with the lower score is dropped;
+    ties keep the first encountered (already deduped by chunk_id earlier).
+
+    The larger containing chunk already includes all information from the inner chunk,
+    so sending both to the reranker wastes tokens without adding signal.
+    """
+    if len(chunks) <= 1:
+        return chunks
+
+    # Sort by range width descending (widest first), break ties by score descending.
+    # This lets us efficiently check whether later (narrower) chunks are contained.
+    sorted_by_width = sorted(
+        chunks,
+        key=lambda c: (-(c.end_line - c.start_line), -c.score),
+    )
+
+    kept: List["ScoredChunk"] = []
+    for candidate in sorted_by_width:
+        contained = False
+        for keeper in kept:
+            if keeper.start_line <= candidate.start_line and candidate.end_line <= keeper.end_line:
+                contained = True
+                break
+        if not contained:
+            kept.append(candidate)
+
+    return kept
+
+
+def _merge_two(a: "ScoredChunk", b: "ScoredChunk") -> "ScoredChunk":
+    """Merge two chunks into one, combining their content and taking the max score."""
+    from corbell.core.query.graph_expander import ScoredChunk  # local import to avoid circular
+
+    new_start = min(a.start_line, b.start_line)
+    new_end = max(a.end_line, b.end_line)
+
+    # Rebuild content by merging lines if both have content
+    # We read the combined content from file to avoid duplication.
+    # If reading fails, concatenate with a separator.
+    try:
+        lines = _read_lines(a.file_path, new_start, new_end)
+        content = "\n".join(lines)
+    except Exception:
+        # Fall back: concatenate without overlap
+        content = a.content + "\n" + b.content
+
+    combined_id = f"{a.chunk_id}+{b.chunk_id}"
+
+    return ScoredChunk(
+        chunk_id=combined_id,
+        score=max(a.score, b.score),
+        file_path=a.file_path,
+        start_line=new_start,
+        end_line=new_end,
+        content=content,
+        repo_id=a.repo_id,
+        symbol=a.symbol,
+        chunk_type=a.chunk_type,
+        language=a.language,
+    )
+
+
+def _read_lines(file_path: str, start_line: int, end_line: int) -> List[str]:
+    """Read specific lines from a file (1-based, inclusive)."""
+    from pathlib import Path
+    path = Path(file_path)
+    all_lines = path.read_text(encoding="utf-8", errors="ignore").splitlines()
+    start_idx = max(0, start_line - 1)
+    end_idx = min(len(all_lines), end_line)
+    return all_lines[start_idx:end_idx]

corbell/core/query/reranker.py

@@ -0,0 +1,131 @@
+"""LLM-based result reranker for the query pipeline."""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from corbell.core.query.graph_expander import ScoredChunk
+
+
+def rerank_chunks(
+    query: str,
+    chunks: List["ScoredChunk"],
+    llm_client: Optional[Any],
+    graph_meta: Optional[Dict[str, Dict]] = None,
+) -> List[str]:
+    """Rerank and filter query results using an LLM.
+
+    Sends chunk content + metadata to the LLM. The LLM returns a JSON array
+    of 0-based chunk indices ordered by relevance, omitting irrelevant chunks.
+
+    On any failure (parse error, LLM timeout), all chunk IDs are returned
+    in their original order (graceful fallback).
+
+    Args:
+        query: The original user query.
+        chunks: Scored chunks to rerank.
+        llm_client: An LLMClient instance (or None / unconfigured).
+        graph_meta: Optional dict mapping chunk_id to graph metadata
+            (callers count, callees count, flow name). When provided,
+            metadata is included in the chunk header sent to the LLM.
+
+    Returns:
+        List of chunk_ids in reranked order (most relevant first).
+        Irrelevant chunks are excluded.
+        Falls back to original order on any failure.
+    """
+    if not chunks:
+        return []
+
+    all_ids = [c.chunk_id for c in chunks]
+
+    if llm_client is None or not getattr(llm_client, "is_configured", False):
+        return all_ids
+
+    # Build payload with code content, indexed for compact LLM output
+    entries = []
+    for i, chunk in enumerate(chunks):
+        meta = graph_meta.get(chunk.chunk_id) if graph_meta else None
+        if meta:
+            callers = meta.get("callers", 0)
+            callees = meta.get("callees", 0)
+            flow = meta.get("flow")
+            if flow:
+                meta_str = f"score={chunk.score:.2f}, callers={callers}, callees={callees}, flow={flow}"
+            else:
+                meta_str = f"score={chunk.score:.2f}, callers={callers}, callees={callees}"
+        else:
+            meta_str = f"score={chunk.score:.2f}"
+
+        content = chunk.content
+        content_lines = content.splitlines()
+        if len(content_lines) > 100:
+            truncated_count = len(content_lines) - 100
+            content = "\n".join(
+                content_lines[:50]
+                + [f"... ({truncated_count} lines truncated) ..."]
+                + content_lines[-50:]
+            )
+
+        entry = (
+            f"[{i}] {meta_str} | {chunk.file_path}:{chunk.start_line}-{chunk.end_line}"
+            f" ({chunk.chunk_type}, {chunk.symbol or 'no symbol'})\n"
+            f"{content}"
+        )
+        entries.append(entry)
+
+    system = (
+        "You are a code search relevance ranker. "
+        "Given a query and numbered code chunks with metadata (relevance score, callers count, "
+        "callees count, flow membership), return a JSON array of chunk indices ordered from most "
+        "relevant to least relevant. "
+        "OMIT chunks that are not relevant to the query. "
+        "Higher score, more callers, and flow membership indicate higher structural importance. "
+        "Return ONLY a valid JSON array of integers, e.g. [2,0,5]."
+    )
+
+    separator = "---\n"
+    chunks_text = separator.join(entries)
+
+    user = (
+        f"Query: {query}\n\n"
+        f"Chunks:\n{chunks_text}\n\n"
+        "Return JSON array of relevant chunk indices (most relevant first):"
+    )
+
+    try:
+        response = llm_client.call(
+            system, user,
+            max_tokens=200,
+            temperature=0.0,
+        )
+
+        text = response.strip()
+        if text.startswith("```"):
+            lines = text.splitlines()
+            text = "\n".join(
+                line for line in lines
+                if not line.startswith("```")
+            ).strip()
+
+        indices = json.loads(text)
+
+        if not isinstance(indices, list):
+            return all_ids
+
+        # Validate indices are ints within range
+        n = len(chunks)
+        filtered = [
+            all_ids[idx] for idx in indices
+            if isinstance(idx, int) and 0 <= idx < n
+        ]
+
+        if not filtered:
+            return all_ids
+
+        return filtered
+
+    except Exception:
+        return all_ids