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.

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