codenib.ops.retrieve

Classes:

Name	Description
RetrievalStage	Minimal branch contract consumed by retrieval-plan execution.
RetrieveContext	Handles backing indexes for retrieval operations.

Functions:

Name	Description
to_queried_nodes	Normalize retrieval results to QueriedNode list.
execute_retrieval_stages	Execute and fuse a declarative set of retrieval branches.
merge_hybrid	Merge multiple retrieval branches with weighted scoring or RRF.
merge_file_rankings	Fuse ranked branches by unique file using weighted RRF.
queried_node_key	Stable identity for deduping candidate nodes across pipeline stages.
queried_file_key	Stable file identity across relative graph and absolute index paths.
dedup_queried_nodes	Deduplicate nodes while preserving first-seen rank order.
dedup_queried_files	Keep the first representative node for each stable file identity.

RetrievalStage

Bases: Protocol

Minimal branch contract consumed by retrieval-plan execution.

RetrieveContext dataclass

RetrieveContext(
    bm25: BM25CodeIndexer | None = None,
    vector_store: CodeVectorStore | None = None,
    regex_index: RegexNodeIndex | None = None,
    default_top_k: int = 10,
    default_level: str = "l2",
    masks: dict[str, set[str]] = dict(),
)

Handles backing indexes for retrieval operations.

to_queried_nodes

to_queried_nodes(results: Sequence[object]) -> list[QueriedNode]

Normalize retrieval results to QueriedNode list.

Source code in codenib/ops/retrieve.py

def to_queried_nodes(results: Sequence[object]) -> List[QueriedNode]:
    """Normalize retrieval results to QueriedNode list."""
    converted: List[QueriedNode] = []
    for rank, item in enumerate(results):
        if isinstance(item, QueriedNode):
            converted.append(item)
            continue
        if isinstance(item, NodeInfo):
            data = _dump_model(item)
            score = data.get("score") or 0.0
            if not score:
                score = 1.0 / (rank + 1)
            data["score"] = score
            converted.append(QueriedNode(**data))
            continue
        if isinstance(item, dict):
            data = dict(item)
            score = data.get("score") or 0.0
            if not score:
                score = 1.0 / (rank + 1)
            data["score"] = score
            data.setdefault("node_name", data.get("name", ""))
            data.setdefault("content", data.get("content"))
            converted.append(QueriedNode(**data))
            continue
        raise TypeError(f"Unsupported result type for normalization: {type(item)}")
    return converted

execute_retrieval_stages

execute_retrieval_stages(
    query: str,
    stages: Sequence[_StageT],
    execute_stage: Callable[[str, _StageT], Sequence[object]],
    *,
    top_k: int,
    fusion: str = "weighted",
    rrf_k: int = 60
) -> list[QueriedNode]

Execute and fuse a declarative set of retrieval branches.

The caller supplies backend-specific branch execution. Stage iteration, result normalization, weighting, and fusion stay here so standalone pipelines, agent runtimes, and protocol adapters share one implementation.

Source code in codenib/ops/retrieve.py

def execute_retrieval_stages(
    query: str,
    stages: Sequence[_StageT],
    execute_stage: Callable[[str, _StageT], Sequence[object]],
    *,
    top_k: int,
    fusion: str = "weighted",
    rrf_k: int = 60,
) -> List[QueriedNode]:
    """Execute and fuse a declarative set of retrieval branches.

    The caller supplies backend-specific branch execution. Stage iteration,
    result normalization, weighting, and fusion stay here so standalone
    pipelines, agent runtimes, and protocol adapters share one implementation.
    """
    if not stages:
        raise ValueError("At least one retrieval stage is required.")
    if top_k <= 0:
        raise ValueError("top_k must be positive.")

    branches = [to_queried_nodes(execute_stage(query, stage)) for stage in stages]
    if len(branches) == 1:
        return branches[0][:top_k]
    return merge_hybrid(
        branches,
        weights=[stage.weight for stage in stages],
        top_k=top_k,
        fusion=fusion,
        rrf_k=rrf_k,
    )

merge_hybrid

merge_hybrid(
    branches: list[list[QueriedNode]],
    weights: list[float] | None = None,
    top_k: int | None = None,
    *,
    fusion: str = "weighted",
    rrf_k: int = 60
) -> list[QueriedNode]

Merge multiple retrieval branches with weighted scoring or RRF.

Source code in codenib/ops/retrieve.py

def merge_hybrid(
    branches: List[List[QueriedNode]],
    weights: Optional[List[float]] = None,
    top_k: Optional[int] = None,
    *,
    fusion: str = "weighted",
    rrf_k: int = 60,
) -> List[QueriedNode]:
    """Merge multiple retrieval branches with weighted scoring or RRF."""
    return _merge_rankings(
        branches,
        weights=weights,
        top_k=top_k,
        fusion=fusion,
        rrf_k=rrf_k,
        key_fn=queried_node_key,
    )

merge_file_rankings

merge_file_rankings(
    branches: list[list[QueriedNode]],
    weights: list[float] | None = None,
    top_k: int | None = None,
    *,
    rrf_k: int = 60
) -> list[QueriedNode]

Fuse ranked branches by unique file using weighted RRF.

Each branch contributes at most once per file. Stable relative file names embedded in node_id take precedence over machine-specific absolute paths, which keeps graph and vector artifacts comparable across snapshots.

Source code in codenib/ops/retrieve.py

def merge_file_rankings(
    branches: List[List[QueriedNode]],
    weights: Optional[List[float]] = None,
    top_k: Optional[int] = None,
    *,
    rrf_k: int = 60,
) -> List[QueriedNode]:
    """Fuse ranked branches by unique file using weighted RRF.

    Each branch contributes at most once per file. Stable relative file names
    embedded in ``node_id`` take precedence over machine-specific absolute
    paths, which keeps graph and vector artifacts comparable across snapshots.
    """
    unique_branches = [dedup_queried_files(branch) for branch in branches]
    return _merge_rankings(
        unique_branches,
        weights=weights,
        top_k=top_k,
        fusion="rrf",
        rrf_k=rrf_k,
        key_fn=queried_file_key,
    )

queried_node_key

queried_node_key(item: QueriedNode) -> tuple[Any, ...]

Stable identity for deduping candidate nodes across pipeline stages.

Source code in codenib/ops/retrieve.py

def queried_node_key(item: QueriedNode) -> Tuple[Any, ...]:
    """Stable identity for deduping candidate nodes across pipeline stages."""
    node_id = (item.node_id or "").strip()
    if node_id:
        return ("node_id", node_id)
    return (
        "span",
        item.file,
        item.node_name,
        item.start_line,
        item.end_line,
    )

queried_file_key

queried_file_key(item: QueriedNode) -> tuple[Any, ...]

Stable file identity across relative graph and absolute index paths.

Source code in codenib/ops/retrieve.py

def queried_file_key(item: QueriedNode) -> Tuple[Any, ...]:
    """Stable file identity across relative graph and absolute index paths."""
    file_path = (item.file or "").strip().replace("\\", "/")
    node_id = (item.node_id or "").strip().replace("\\", "/")
    if node_id and ":" in node_id:
        node_file = node_id.split(":", 1)[0].removeprefix("./")
        if (
            node_file
            and file_path
            and (file_path == node_file or file_path.endswith("/" + node_file))
        ):
            return ("file", node_file)
    if file_path:
        return ("file", file_path)
    return ("node", *queried_node_key(item))

dedup_queried_nodes

dedup_queried_nodes(nodes: Sequence[QueriedNode]) -> list[QueriedNode]

Deduplicate nodes while preserving first-seen rank order.

Retrieval pipelines often concatenate dense hits with graph-expanded neighbors. The dense order is the ranking signal; duplicates should not consume multiple top-k slots. If a later duplicate carries code content and the first copy does not, keep the first rank but attach the content.

Source code in codenib/ops/retrieve.py

def dedup_queried_nodes(nodes: Sequence[QueriedNode]) -> List[QueriedNode]:
    """Deduplicate nodes while preserving first-seen rank order.

    Retrieval pipelines often concatenate dense hits with graph-expanded
    neighbors. The dense order is the ranking signal; duplicates should not
    consume multiple top-k slots. If a later duplicate carries code content and
    the first copy does not, keep the first rank but attach the content.
    """
    by_key: Dict[Tuple[Any, ...], int] = {}
    out: List[QueriedNode] = []
    for node in nodes:
        key = queried_node_key(node)
        existing_idx = by_key.get(key)
        if existing_idx is None:
            by_key[key] = len(out)
            out.append(node)
            continue
        existing = out[existing_idx]
        if not existing.content and node.content:
            out[existing_idx] = _with_score(
                existing,
                existing.score or 0.0,
                content_override=node.content,
            )
    return out

dedup_queried_files

dedup_queried_files(nodes: Sequence[QueriedNode]) -> list[QueriedNode]

Keep the first representative node for each stable file identity.

Source code in codenib/ops/retrieve.py

def dedup_queried_files(nodes: Sequence[QueriedNode]) -> List[QueriedNode]:
    """Keep the first representative node for each stable file identity."""
    seen: Set[Tuple[Any, ...]] = set()
    out: List[QueriedNode] = []
    for node in nodes:
        key = queried_file_key(node)
        if key in seen:
            continue
        seen.add(key)
        out.append(node)
    return out