codenib.ops.expand

Classes:

Name	Description
ExpandContext	Shared resources for graph expansion operators.

Functions:

Name	Description
nodeinfo_to_queried	Convert NodeInfo list to QueriedNode list, assigning rank-based scores.
hydrate_candidate_contents	Fill missing candidate content from persisted source spans.
expand_graph_neighbors	Return graph neighbors for retrieval candidates.
expand_graph_region	Expand ranked seeds through a bounded graph region.
expand_retrieval_candidates	Apply a bounded graph expansion plan to ranked retrieval seeds.

ExpandContext dataclass

ExpandContext(
    code_graph: CodeGraph | None = None,
    lsp_provider: Any | None = None,
    default_top_k: int = 50,
    default_hops: int = 2,
    default_direction: str = "both",
    default_method: str = "bfs",
    default_damping: float = 0.85,
    filter_tests: bool = True,
)

Shared resources for graph expansion operators.

nodeinfo_to_queried

nodeinfo_to_queried(nodes: list[NodeInfo]) -> list[QueriedNode]

Convert NodeInfo list to QueriedNode list, assigning rank-based scores.

Source code in codenib/ops/expand.py

def nodeinfo_to_queried(nodes: List[NodeInfo]) -> List[QueriedNode]:
    """Convert NodeInfo list to QueriedNode list, assigning rank-based scores."""
    results: List[QueriedNode] = []
    for rank, ni in enumerate(nodes):
        score = ni.score if ni.score is not None else 1.0 / (rank + 1)
        results.append(
            QueriedNode(
                node_name=ni.node_name,
                type=ni.type,
                file=ni.file,
                node_id=ni.node_id or ni.node_name,
                start_line=ni.start_line,
                end_line=ni.end_line,
                score=score,
                content=ni.content,
            )
        )
    return results

hydrate_candidate_contents

hydrate_candidate_contents(
    candidates: Sequence[QueriedNode],
    *,
    repo_path: str | None,
    git_commit: str | None = None
) -> list[QueriedNode]

Fill missing candidate content from persisted source spans.

Source code in codenib/ops/expand.py

def hydrate_candidate_contents(
    candidates: Sequence[QueriedNode],
    *,
    repo_path: Optional[str],
    git_commit: Optional[str] = None,
) -> List[QueriedNode]:
    """Fill missing candidate content from persisted source spans."""
    hydrated: List[QueriedNode] = []
    for candidate in candidates:
        if candidate.content:
            hydrated.append(candidate)
            continue
        content = _read_span_content(
            repo_path=repo_path,
            file_path=candidate.file,
            start_line=candidate.start_line,
            end_line=candidate.end_line,
            git_commit=git_commit,
        )
        if not content:
            hydrated.append(candidate)
        elif hasattr(candidate, "model_copy"):
            hydrated.append(candidate.model_copy(update={"content": content}))
        else:
            hydrated.append(candidate.copy(update={"content": content}))
    return hydrated

expand_graph_neighbors

expand_graph_neighbors(
    context: ExpandContext,
    seeds: Sequence[QueriedNode],
    *,
    per_seed: int = 5,
    direction: str = "both",
    repo_path: str | None = None,
    include_content: bool = False,
    symbol_only: bool = True,
    require_span: bool = True,
    score_scale: float = 0.5,
    edge_types: Sequence[str] | None = (EDGE_TYPE_REFERENCE,)
) -> list[QueriedNode]

Return graph neighbors for retrieval candidates.

This is the retrieval-pipeline counterpart to the agent-facing graph navigation helpers: keep dense candidates as seeds, append compact graph neighbors, and leave ranking to the downstream reranker. Only one-hop predecessor/successor expansion is performed; no BFS/PPR is hidden here. Expansion follows reference edges by default so containment does not consume the dependency-neighbor budget.

Source code in codenib/ops/expand.py

def expand_graph_neighbors(
    context: ExpandContext,
    seeds: Sequence[QueriedNode],
    *,
    per_seed: int = 5,
    direction: str = "both",
    repo_path: Optional[str] = None,
    include_content: bool = False,
    symbol_only: bool = True,
    require_span: bool = True,
    score_scale: float = 0.5,
    edge_types: Optional[Sequence[str]] = (EDGE_TYPE_REFERENCE,),
) -> List[QueriedNode]:
    """Return graph neighbors for retrieval candidates.

    This is the retrieval-pipeline counterpart to the agent-facing graph
    navigation helpers: keep dense candidates as seeds, append compact graph
    neighbors, and leave ranking to the downstream reranker. Only one-hop
    predecessor/successor expansion is performed; no BFS/PPR is hidden here.
    Expansion follows reference edges by default so containment does not consume
    the dependency-neighbor budget.
    """
    graph = context.code_graph
    if graph is None or not seeds:
        return []

    per_seed = max(0, int(per_seed or 0))
    if per_seed == 0:
        return []
    if direction not in ("callees", "successors", "callers", "predecessors", "both"):
        raise ValueError(f"Unsupported graph expansion direction: {direction!r}")

    out: List[QueriedNode] = []
    selected_edge_types = None if edge_types is None else set(edge_types)
    for seed_rank, seed in enumerate(seeds, 1):
        canonical = _resolve_seed(graph, seed)
        if canonical is None:
            continue
        successor_ids = (
            list(graph.get_successors(canonical, edge_types=selected_edge_types))
            if direction in ("callees", "successors", "both")
            else []
        )
        predecessor_ids = (
            list(graph.get_predecessors(canonical, edge_types=selected_edge_types))
            if direction in ("callers", "predecessors", "both")
            else []
        )
        if direction == "both":
            interleaved = (
                vertex_id
                for pair in zip_longest(successor_ids, predecessor_ids)
                for vertex_id in pair
                if vertex_id is not None
            )
            neighbor_ids = list(dict.fromkeys(interleaved))
        else:
            neighbor_ids = successor_ids or predecessor_ids

        added_for_seed = 0
        for vid in neighbor_ids:
            node = _neighbor_to_queried(
                graph,
                vid,
                seed=seed,
                seed_rank=seed_rank,
                repo_path=repo_path,
                include_content=include_content,
                score_scale=score_scale,
            )
            if node is None:
                continue
            if symbol_only and not is_symbol_node(node.type):
                continue
            if require_span and (
                node.file is None or node.start_line is None or node.end_line is None
            ):
                continue
            out.append(node)
            added_for_seed += 1
            if added_for_seed >= per_seed:
                break
    return out

expand_graph_region

expand_graph_region(
    context: ExpandContext,
    seeds: Sequence[QueriedNode],
    *,
    top_k: int = 50,
    hops: int = 2,
    direction: str = "both",
    use_ppr: bool = False,
    repo_path: str | None = None,
    include_content: bool = False
) -> list[QueriedNode]

Expand ranked seeds through a bounded graph region.

Multi-hop and PPR expansion share this operator so manifest-backed runtimes and the standalone retrieval pipeline execute the same graph semantics. If no region can be resolved, the function falls back to the one-hop neighbor operator.

Source code in codenib/ops/expand.py

def expand_graph_region(
    context: ExpandContext,
    seeds: Sequence[QueriedNode],
    *,
    top_k: int = 50,
    hops: int = 2,
    direction: str = "both",
    use_ppr: bool = False,
    repo_path: Optional[str] = None,
    include_content: bool = False,
) -> List[QueriedNode]:
    """Expand ranked seeds through a bounded graph region.

    Multi-hop and PPR expansion share this operator so manifest-backed
    runtimes and the standalone retrieval pipeline execute the same graph
    semantics. If no region can be resolved, the function falls back to the
    one-hop neighbor operator.
    """

    graph = context.code_graph
    if graph is None or not seeds or top_k <= 0:
        return []

    seed_names: List[str] = []
    for seed in seeds:
        canonical = _resolve_seed(graph, seed)
        if canonical is not None and canonical not in seed_names:
            seed_names.append(canonical)
    if not seed_names:
        return []
    if not hasattr(graph, "get_graph") or not hasattr(graph, "name_to_vertex"):
        return _expand_neighbor_fallback(
            context,
            seeds,
            top_k=top_k,
            direction=direction,
            repo_path=repo_path,
            include_content=include_content,
        )

    from ..graph.roi_subgraph import ROISubgraph

    roi = ROISubgraph(graph)
    if use_ppr:
        nodes = roi.expand_ppr(
            seed_names,
            top_k=top_k,
            damping=context.default_damping,
            filter_tests=context.filter_tests,
        )
    else:
        subgraph = roi.extract_subgraph(
            seed_names,
            k_hop=max(1, int(hops)),
            direction=_roi_direction(direction),
        )
        nodes = roi.get_filtered_subgraph_nodes(
            subgraph,
            filter_tests=context.filter_tests,
        )[:top_k]

    expanded = nodeinfo_to_queried(nodes)
    if include_content:
        expanded = hydrate_candidate_contents(expanded, repo_path=repo_path)
    if expanded:
        return expanded[:top_k]

    return _expand_neighbor_fallback(
        context,
        seeds,
        top_k=top_k,
        direction=direction,
        repo_path=repo_path,
        include_content=include_content,
    )

expand_retrieval_candidates

expand_retrieval_candidates(
    context: ExpandContext,
    seeds: Sequence[QueriedNode],
    *,
    seed_top_k: int | None,
    expand_top_k: int,
    hops: int,
    direction: str,
    use_ppr: bool,
    repo_path: str | None,
    include_content: bool = True
) -> list[QueriedNode]

Apply a bounded graph expansion plan to ranked retrieval seeds.

Source code in codenib/ops/expand.py

def expand_retrieval_candidates(
    context: ExpandContext,
    seeds: Sequence[QueriedNode],
    *,
    seed_top_k: Optional[int],
    expand_top_k: int,
    hops: int,
    direction: str,
    use_ppr: bool,
    repo_path: Optional[str],
    include_content: bool = True,
) -> List[QueriedNode]:
    """Apply a bounded graph expansion plan to ranked retrieval seeds."""
    if context.code_graph is None:
        raise RuntimeError("Graph expansion requested but no code graph is loaded.")
    if expand_top_k <= 0:
        raise ValueError("expand_top_k must be positive.")

    seed_cap = seed_top_k or len(seeds)
    graph_seeds = list(seeds[:seed_cap])
    if not graph_seeds:
        return []

    expanded = expand_graph_region(
        context,
        graph_seeds,
        top_k=expand_top_k,
        hops=hops,
        direction=direction,
        use_ppr=use_ppr,
        repo_path=repo_path,
        include_content=include_content,
    )
    return dedup_queried_nodes([*graph_seeds, *expanded])[:expand_top_k]