codenib.model

Model module for CodeNib.

Modules:

Name	Description
agentless_pipeline
bm25_retrieve_pipeline
dense_graph_expand_rerank_pipeline
embedding_retrieve_pipeline
graph_augmented_rerank_pipeline
graph_retrieve_pipeline
hybrid_retrieve_pipeline	Cascade-style hybrid retrieval: BM25 → embedding rerank (no agent loop).
retrieval_planner	Query- and budget-aware retrieval path planner.
retrieve_rerank_pipeline

Classes:

Name	Description
AgentlessPipeline	The Agentless approach from the `"Agentless: Demystifying LLM-based
BM25RetrievePipeline	BM25-only retrieval pipeline.
DenseGraphExpandRerankPipeline	Dense-first graph expansion retrieval pipeline.
EmbeddingRetrievePipeline	Embedding-only retrieval pipeline.
GraphAugmentedRerankPipeline	Compatibility alias for :class:DenseGraphExpandRerankPipeline.
GraphRetrievePipeline	Compatibility alias for :class:SparseSeededGraphRetrievePipeline.
SparseSeededGraphRetrievePipeline	Sparse-seeded graph-first retrieval pipeline.
HybridRetrievePipeline	Two-stage cascade retrieval: BM25 → embedding rerank.
GraphExpansionPlan	Graph expansion controls for structural retrieval paths.
RetrievalBudget	Coarse retrieval budget used for path selection.
RetrievalCapabilities	Backends that a planner may use for this pipeline instance.
RetrievalPathPlan	Concrete path selected for a query.
RetrievalPlanner	Select one of CodeNib's retrieval paths for a query and budget.
RetrievalStagePlan	Declarative retrieval branch selected by the planner.
RetrieveRerankPipeline	Composable code retrieval pipeline built from retrieval + rerank ops.
RetrieveStageConfig	Declarative configuration for an individual retrieval branch.

Functions:

Name	Description
build_retrieve_plan	Convenience helper for common retrieval plan presets.

AgentlessPipeline

AgentlessPipeline(
    repo_path: str,
    repo_commit: str,
    llm_model: str = "gpt-4o",
    llm_temperature: float = 0.0,
    llm_max_tokens: int = 2048,
    top_n_files: int = 3,
    languages: list[str] | None = None,
    cache_dir: str | None = None,
    instance_id: str | None = None,
)

The Agentless approach from the "Agentless: Demystifying LLM-based Software Engineering Agents" <https://arxiv.org/abs/2407.01489>_ paper.

:class:~codenib.model.AgentlessPipeline localizes code locations that need modification by employing a simplistic three-phase process:

Phase 1 - File Localization: Identifies relevant files using repository structure and LLM reasoning:

.. math:: \mathcal{F} = \text{LLM}(\text{problem}, \text{structure}(\mathcal{R}))

where :math:\mathcal{R} is the repository and :math:\mathcal{F} is the set of candidate files.

TODO: add dense retrieval to the file localization phase

Phase 2 - Node Localization: Narrows down to specific code entities (functions, classes, methods) within the identified files:

.. math:: \mathcal{N} = \text{LLM}(\text{problem}, {\text{structure}(f) \mid f \in \mathcal{F}})

where :math:\mathcal{N} is the set of candidate nodes (symbols).

Phase 3 - Node Refinement: Validates and refines the final set of nodes using full code content:

.. math:: \mathcal{N}^* = \text{LLM}(\text{problem}, {\text{content}(n) \mid n \in \mathcal{N}})

where :math:\mathcal{N}^* is the refined set of nodes requiring changes.

The localization results can be used for downstream tasks such as automated program repair, code synthesis, or interactive code editing via :meth:~codenib.model.AgentlessPipeline.query.

Parameters:

Name	Type	Description	Default
repo_path	str	Path to the repository to analyze.	required
repo_commit	str	Git commit hash for version identification.	required
llm_model	str	Full litellm model identifier (e.g., "gpt-4o", "vertex_ai/gemini-2.5-flash"). (default: :obj:"gpt-4o")	'gpt-4o'
llm_temperature	float	Temperature for LLM sampling. (default: :obj:0.0)	0.0
llm_max_tokens	int	Maximum tokens for LLM responses. (default: :obj:2048)	2048
top_n_files	int	Number of top files to consider in Stage 1. (default: :obj:3)	3
languages	list[str]	Programming languages to index. If set to :obj:None, defaults to :obj:["python"]. (default: :obj:None)	None
cache_dir	str	Directory for caching SCIP indices and graphs. If set to :obj:None, defaults to :obj:~/.codenib. (default: :obj:None)	None
instance_id	str	Instance identifier for organizing cache. If provided, uses :obj:cache_dir/instance_id as the cache path. Otherwise, falls back to :obj:cache_dir/scip_{repo_name}@{commit}. (default: :obj:None)	None

Source code in codenib/model/agentless_pipeline.py

def __init__(
    self,
    # Repo config
    repo_path: str,
    repo_commit: str,
    # LLM configs
    llm_model: str = "gpt-4o",
    llm_temperature: float = 0.0,
    llm_max_tokens: int = 2048,
    # Localization configs
    top_n_files: int = 3,
    # Repo processing config
    languages: Optional[List[str]] = None,
    # Cache
    cache_dir: Optional[str] = None,
    instance_id: Optional[str] = None,
):
    # Attributes
    self.repo_path = os.path.abspath(repo_path)
    self.repo_commit = repo_commit.strip()
    self.top_n_files = top_n_files
    self.languages = languages or ["python"]
    self._non_test_file_nodes_cache: Optional[List[str]] = None

    # Validate repo
    if not os.path.exists(self.repo_path):
        raise ValueError(f"Repository path does not exist: {self.repo_path}")
    if not os.path.isdir(self.repo_path):
        raise ValueError(f"Repository path is not a directory: {self.repo_path}")
    if not self.repo_commit or not isinstance(self.repo_commit, str):
        raise ValueError("repo_commit must be provided as a non-empty string")

    # Initialize cache directory
    cache_dir = Path(cache_dir) if cache_dir else user_state_dir()
    cache_dir.mkdir(parents=True, exist_ok=True)

    # Prepare repo identifiers
    repo_name = os.path.basename(self.repo_path)
    commit_identifier = self.repo_commit[:12]

    # Create SCIP output directory based on instance_id or repo name
    if instance_id:
        scip_output_dir = cache_dir / instance_id
    else:
        scip_output_dir = cache_dir / f"scip_{repo_name}@{commit_identifier}"

    scip_output_dir.mkdir(parents=True, exist_ok=True)

    scip_indexer = LSIndexer(self.repo_path, output_dir=str(scip_output_dir))
    self.code_graph: CodeGraph = scip_indexer.run_pipeline(
        project_name=f"{repo_name}@{commit_identifier}",
        skip_level="graph",  # Enable cache: load from graph.pkl if exists
    )

    if not self.code_graph:
        raise ValueError("Failed to build code graph")

    # Initialize LLM — model string should be a full litellm identifier
    # (e.g., "gpt-4o", "vertex_ai/gemini-2.5-flash").
    self.llm = LiteLLMChat(
        model=llm_model,
        max_tokens=llm_max_tokens,
        temperature=llm_temperature,
    )

    # Load prompt templates from files
    prompt_dir = Path(__file__).parent / "prompts"
    self.stage1_prompt = (prompt_dir / "agentless_stage1.txt").read_text()
    self.stage2_prompt = (prompt_dir / "agentless_stage2.txt").read_text()
    self.stage3_prompt = (prompt_dir / "agentless_stage3.txt").read_text()

    # Create structured LLMs for each stage
    self.structured_llm_stage1 = self.llm.with_structured_output(Stage1Result)
    self.structured_llm_stage2 = self.llm.with_structured_output(Stage2Result)
    self.structured_llm_stage3 = self.llm.with_structured_output(Stage3Result)

    logger.info(
        f"AgentlessPipeline initialized: "
        f"repo={repo_name}@{commit_identifier}, "
        f"llm={llm_model}, "
        f"top_n_files={top_n_files}, "
        f"cache_dir={cache_dir}"
    )

BM25RetrievePipeline

BM25RetrievePipeline(
    repo_path: str,
    index_path: str,
    *,
    top_k: int = 50,
    project_name: str | None = None,
    language: str = "python",
    languages: list[str] | None = None,
    graph_route: str = "active"
)

BM25-only retrieval pipeline.

Builds a code graph via SCIP indexing, then retrieves top-K nodes using BM25 sparse search. Useful as a baseline to isolate whether poor graph retrieval performance is due to bad BM25 seeds or bad graph expansion.

Parameters:

Name	Type	Description	Default
repo_path	str	Repository root to index.	required
index_path	str	Directory used for SCIP index caches.	required
top_k	int	Number of results to retrieve (default: 50).	50
project_name	str | None	Project name for SCIP indexing; defaults to index_path basename.	None
language	str	Backward-compatible single graph-indexing language.	'python'
languages	list[str] | None	Graph-indexing languages. When provided, overrides language and builds a merged multi-language graph.	None
graph_route	str	"active" for public graph routes, "lsp" for backend comparison, or "scip-candidate" to explicitly evaluate candidate SCIP routes.	'active'

Methods:

Name	Description
query	Retrieve top-K nodes using BM25 sparse search.
close	Release pipeline resources.

Source code in codenib/model/bm25_retrieve_pipeline.py

def __init__(
    self,
    repo_path: str,
    index_path: str,
    *,
    top_k: int = 50,
    project_name: Optional[str] = None,
    language: str = "python",
    languages: Optional[List[str]] = None,
    graph_route: str = "active",
) -> None:
    self.top_k = top_k

    pname = project_name or Path(index_path).name
    index_languages = languages or [language]

    # Build CodeGraph via the registry-backed graph indexer.
    self.code_graph = build_graph_for_languages(
        repo_path,
        index_path,
        languages=index_languages,
        project_name=pname,
        skip_level="graph",
        graph_route=graph_route,
    )
    if self.code_graph is None:
        raise RuntimeError(f"Failed to build code graph for {repo_path!r}")

    # Build BM25 index from graph
    self.bm25_index = BM25CodeIndexer(max_k=top_k, language="english")
    self.bm25_index.build_index_from_graph(self.code_graph)

query

query(query: str, top_k: int | None = None) -> list[QueriedNode]

Retrieve top-K nodes using BM25 sparse search.

Parameters:

Name	Type	Description	Default
query	str	The search query (e.g., problem statement).	required
top_k	int | None	Number of results; overrides the instance default if provided.	None

Returns:

Type	Description
list[QueriedNode]	List of QueriedNode sorted by BM25 score.

Source code in codenib/model/bm25_retrieve_pipeline.py

def query(self, query: str, top_k: Optional[int] = None) -> List[QueriedNode]:
    """Retrieve top-K nodes using BM25 sparse search.

    Args:
        query: The search query (e.g., problem statement).
        top_k: Number of results; overrides the instance default if provided.

    Returns:
        List of QueriedNode sorted by BM25 score.
    """
    k = top_k if top_k is not None else self.top_k
    bm25_results = self.bm25_index.search(query, top_k=k)
    return [
        QueriedNode(
            node_name=n.node_name,
            type=n.type,
            file=n.file,
            node_id=n.node_id or n.node_name,
            start_line=n.start_line,
            end_line=n.end_line,
            score=n.score,
            content=n.content,
        )
        for n in bm25_results
    ]

close

close() -> None

Release pipeline resources.

Source code in codenib/model/bm25_retrieve_pipeline.py

def close(self) -> None:
    """Release pipeline resources."""
    self.bm25_index = None
    self.code_graph = None

DenseGraphExpandRerankPipeline

DenseGraphExpandRerankPipeline(
    retriever: EmbeddingRetrievePipeline,
    code_graph: CodeGraph | None = None,
    *,
    rerank_context: RerankContext | None = None,
    repo_path: str | None = None,
    dense_top_k: int = 20,
    graph_seed_top_k: int | None = None,
    graph_neighbors_per_seed: int = 5,
    graph_direction: str = "both",
    graph_edge_types: Sequence[str] | None = (EDGE_TYPE_REFERENCE,),
    graph_fusion_weight: float = 0.0,
    graph_fusion_rrf_k: int = 60,
    final_top_k: int = 10,
    candidate_top_k: int | None = None,
    candidate_filters: Sequence[CandidatePredicate] | None = None,
    include_graph_content: bool = True,
    graph_symbol_only: bool = True,
    graph_require_span: bool = True,
    close_retriever: bool = False
)

Dense-first graph expansion retrieval pipeline.

Stage 1: retrieve dense top-K candidates from the full vector index. Stage 2: append one-hop graph neighbors for those dense candidates. Stage 3: optionally filter and deduplicate the combined candidate set. Stage 4: optionally rerank only the assembled candidates.

This pipeline is intended for the "embedding baseline plus graph augmentation" experiment family. It is intentionally different from :class:SparseSeededGraphRetrievePipeline, which uses sparse/BM25 seeds and ranks inside the graph-expanded set.

Methods:

Name	Description
load_index	Hot-swap the backing dense retrieval index.
load_graph	Hot-swap the graph used for neighbor expansion.
query	Retrieve dense candidates, append graph neighbors, dedup, and rerank.
close	Release owned resources.

Source code in codenib/model/dense_graph_expand_rerank_pipeline.py

def __init__(
    self,
    retriever: EmbeddingRetrievePipeline,
    code_graph: Optional[CodeGraph] = None,
    *,
    rerank_context: Optional[RerankContext] = None,
    repo_path: Optional[str] = None,
    dense_top_k: int = 20,
    graph_seed_top_k: Optional[int] = None,
    graph_neighbors_per_seed: int = 5,
    graph_direction: str = "both",
    graph_edge_types: Optional[Sequence[str]] = (EDGE_TYPE_REFERENCE,),
    graph_fusion_weight: float = 0.0,
    graph_fusion_rrf_k: int = 60,
    final_top_k: int = 10,
    candidate_top_k: Optional[int] = None,
    candidate_filters: Optional[Sequence[CandidatePredicate]] = None,
    include_graph_content: bool = True,
    graph_symbol_only: bool = True,
    graph_require_span: bool = True,
    close_retriever: bool = False,
) -> None:
    if dense_top_k <= 0:
        raise ValueError("dense_top_k must be positive.")
    if final_top_k <= 0:
        raise ValueError("final_top_k must be positive.")
    if graph_neighbors_per_seed < 0:
        raise ValueError("graph_neighbors_per_seed must be non-negative.")
    if graph_seed_top_k is not None and graph_seed_top_k <= 0:
        raise ValueError("graph_seed_top_k must be positive when provided.")
    if graph_fusion_weight < 0:
        raise ValueError("graph_fusion_weight must be non-negative.")
    if graph_fusion_rrf_k <= 0:
        raise ValueError("graph_fusion_rrf_k must be positive.")
    if graph_fusion_weight and rerank_context is None:
        raise ValueError("graph fusion requires a rerank_context.")

    self.retriever = retriever
    self.expand_context = ExpandContext(code_graph=code_graph)
    self.rerank_context = rerank_context
    self.repo_path = repo_path
    self.dense_top_k = dense_top_k
    self.graph_seed_top_k = graph_seed_top_k
    self.graph_neighbors_per_seed = graph_neighbors_per_seed
    self.graph_direction = graph_direction
    self.graph_edge_types = (
        None if graph_edge_types is None else tuple(graph_edge_types)
    )
    self.graph_fusion_weight = graph_fusion_weight
    self.graph_fusion_rrf_k = graph_fusion_rrf_k
    self.final_top_k = final_top_k
    self.candidate_top_k = candidate_top_k
    self.candidate_filters = list(candidate_filters or [])
    self.include_graph_content = include_graph_content
    self.graph_symbol_only = graph_symbol_only
    self.graph_require_span = graph_require_span
    self.close_retriever = close_retriever

    self.last_dense_results: List[QueriedNode] = []
    self.last_graph_results: List[QueriedNode] = []
    self.last_filtered_candidates: List[QueriedNode] = []
    self.last_candidates: List[QueriedNode] = []
    self.last_semantic_results: List[QueriedNode] = []
    self.last_reranked_results: List[QueriedNode] = []

load_index

load_index(index_path: str) -> None

Hot-swap the backing dense retrieval index.

Source code in codenib/model/dense_graph_expand_rerank_pipeline.py

def load_index(self, index_path: str) -> None:
    """Hot-swap the backing dense retrieval index."""
    self.retriever.load_index(index_path)

load_graph

load_graph(code_graph: CodeGraph | None, *, repo_path: str | None = None) -> None

Hot-swap the graph used for neighbor expansion.

Source code in codenib/model/dense_graph_expand_rerank_pipeline.py

def load_graph(
    self, code_graph: Optional[CodeGraph], *, repo_path: Optional[str] = None
) -> None:
    """Hot-swap the graph used for neighbor expansion."""
    self.expand_context.code_graph = code_graph
    if repo_path is not None:
        self.repo_path = repo_path

query

query(
    query: str,
    *,
    top_k: int | None = None,
    dense_top_k: int | None = None,
    profiler: Profiler | None = None
) -> list[QueriedNode]

Retrieve dense candidates, append graph neighbors, dedup, and rerank.

Source code in codenib/model/dense_graph_expand_rerank_pipeline.py

def query(
    self,
    query: str,
    *,
    top_k: Optional[int] = None,
    dense_top_k: Optional[int] = None,
    profiler: Optional[Profiler] = None,
) -> List[QueriedNode]:
    """Retrieve dense candidates, append graph neighbors, dedup, and rerank."""
    with self._query_embedding_scope():
        return self._query_impl(
            query,
            top_k=top_k,
            dense_top_k=dense_top_k,
            profiler=profiler,
        )

close

close() -> None

Release owned resources.

Source code in codenib/model/dense_graph_expand_rerank_pipeline.py

def close(self) -> None:
    """Release owned resources."""
    if self.close_retriever:
        self.retriever.close()

EmbeddingRetrievePipeline

EmbeddingRetrievePipeline(
    repo_path: str | None = None,
    index_path: str | None = None,
    *,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_provider: str = "huggingface",
    embedding_dimension: int = 768,
    languages: list[str] | None = None,
    max_lines_per_chunk: int = 300,
    top_k: int = 50,
    embedding_kwargs: dict | None = None,
    force_rebuild: bool = False,
    index_type: str = "flat",
    ivf_nlist: int = 100,
    ivf_nprobe: int = 8
)

Embedding-only retrieval pipeline.

Typical single-instance usage (loads or builds index immediately)::

pipeline = EmbeddingRetrievePipeline(
    repo_path=repo_path,
    index_path=index_path,
    embedding_model="Qwen/Qwen3-Embedding-0.6B",
    embedding_dimension=1024,
)
results = pipeline.query(problem_statement)
pipeline.close()

Multi-instance / hot-swap usage (load model once, swap index per instance)::

pipeline = EmbeddingRetrievePipeline(
    embedding_model="Qwen/Qwen3-Embedding-0.6B",
    embedding_dimension=1024,
)
for instance in dataset:
    pipeline.load_index(index_path_for_instance)
    results = pipeline.query(instance["problem_statement"])
pipeline.close()

Parameters:

Name	Type	Description	Default
repo_path	str | None	Repository root to index. Only required when force_rebuild=True; ignored when loading a pre-built index.	None
index_path	str | None	Directory used for vector index caches. When None, the pipeline is initialised with an empty index so that :meth:load_index can be called later (hot-swap pattern).	None
embedding_model	str	Embedding model name.	'nomic-ai/CodeRankEmbed'
embedding_provider	str	Embedding provider ("huggingface" or "openai").	'huggingface'
embedding_dimension	int	Embedding vector dimension.	768
languages	list[str] | None	Languages to chunk for indexing (default: ["python"]).	None
max_lines_per_chunk	int	Maximum lines per chunk.	300
top_k	int	Default number of results to retrieve.	50
embedding_kwargs	dict | None	Extra kwargs forwarded to the embedding wrapper.	None
force_rebuild	bool	Re-chunk and re-embed even if a cached index exists.	False

Methods:

Name	Description
load_index	Hot-swap the FAISS index without reloading the embedding model.
query	Retrieve top-K nodes using embedding similarity search.
close	Release all resources including the embedding model.

Source code in codenib/model/embedding_retrieve_pipeline.py

def __init__(
    self,
    repo_path: Optional[str] = None,
    index_path: Optional[str] = None,
    *,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_provider: str = "huggingface",
    embedding_dimension: int = 768,
    languages: Optional[List[str]] = None,
    max_lines_per_chunk: int = 300,
    top_k: int = 50,
    embedding_kwargs: Optional[dict] = None,
    force_rebuild: bool = False,
    index_type: str = "flat",
    ivf_nlist: int = 100,
    ivf_nprobe: int = 8,
) -> None:
    self.top_k = top_k
    _embedding_kwargs = {"model_kwargs": {"trust_remote_code": True}}
    if embedding_kwargs:
        _embedding_kwargs.update(embedding_kwargs)

    if index_path is None:
        # Model-only initialisation — no index loaded yet.
        # Callers should follow up with load_index().
        self.vector_store = CodeVectorStore(
            embedding_model=embedding_model,
            embedding_provider=embedding_provider,
            dimension=embedding_dimension,
            index_metric="ip",
            index_type=index_type,
            ivf_nlist=ivf_nlist,
            ivf_nprobe=ivf_nprobe,
            store_path=None,
            **_embedding_kwargs,
        )
    else:
        self.vector_store: CodeVectorStore = build_hierarchical_vector_store(
            repo_path=repo_path or "",
            index_path=index_path,
            plan_name=None,
            languages=languages or ["python"],
            max_lines_per_chunk=max_lines_per_chunk,
            build_levels=["l2"],
            embedding_model=embedding_model,
            embedding_provider=embedding_provider,
            embedding_dimension=embedding_dimension,
            embedding_kwargs=_embedding_kwargs,
            index_metric="ip",
            index_type=index_type,
            ivf_nlist=ivf_nlist,
            ivf_nprobe=ivf_nprobe,
            force_rebuild=force_rebuild,
        )

load_index

load_index(index_path: str) -> None

Hot-swap the FAISS index without reloading the embedding model.

Frees the current index and loads the pre-built index at index_path. The embedding model stays in memory, so this is much faster than creating a new :class:EmbeddingRetrievePipeline for each instance.

Source code in codenib/model/embedding_retrieve_pipeline.py

def load_index(self, index_path: str) -> None:
    """Hot-swap the FAISS index without reloading the embedding model.

    Frees the current index and loads the pre-built index at *index_path*.
    The embedding model stays in memory, so this is much faster than
    creating a new :class:`EmbeddingRetrievePipeline` for each instance.
    """
    self.vector_store.swap_index(index_path)

query

query(query: str, top_k: int | None = None) -> list[QueriedNode]

Retrieve top-K nodes using embedding similarity search.

Parameters:

Name	Type	Description	Default
query	str	The search query (e.g., problem statement).	required
top_k	int | None	Number of results; overrides the instance default if provided.	None

Returns:

Type	Description
list[QueriedNode]	List of QueriedNode sorted by similarity score.

Source code in codenib/model/embedding_retrieve_pipeline.py

def query(self, query: str, top_k: Optional[int] = None) -> List[QueriedNode]:
    """Retrieve top-K nodes using embedding similarity search.

    Args:
        query: The search query (e.g., problem statement).
        top_k: Number of results; overrides the instance default if provided.

    Returns:
        List of QueriedNode sorted by similarity score.
    """
    k = top_k if top_k is not None else self.top_k
    search_results = self.vector_store.search_with_content(query, top_k=k)
    return [
        QueriedNode(
            node_name=n.node_name,
            type=n.type,
            file=n.file,
            node_id=n.node_id,
            start_line=n.start_line,
            end_line=n.end_line,
            score=n.score,
            content=n.content,
        )
        for n in search_results
    ]

close

close() -> None

Release all resources including the embedding model.

Source code in codenib/model/embedding_retrieve_pipeline.py

def close(self) -> None:
    """Release all resources including the embedding model."""
    if self.vector_store is not None:
        self.vector_store.close()
        self.vector_store = None

GraphAugmentedRerankPipeline

GraphAugmentedRerankPipeline(
    retriever: EmbeddingRetrievePipeline,
    code_graph: CodeGraph | None = None,
    *,
    rerank_context: RerankContext | None = None,
    repo_path: str | None = None,
    dense_top_k: int = 20,
    graph_seed_top_k: int | None = None,
    graph_neighbors_per_seed: int = 5,
    graph_direction: str = "both",
    graph_edge_types: Sequence[str] | None = (EDGE_TYPE_REFERENCE,),
    graph_fusion_weight: float = 0.0,
    graph_fusion_rrf_k: int = 60,
    final_top_k: int = 10,
    candidate_top_k: int | None = None,
    candidate_filters: Sequence[CandidatePredicate] | None = None,
    include_graph_content: bool = True,
    graph_symbol_only: bool = True,
    graph_require_span: bool = True,
    close_retriever: bool = False
)

Bases: DenseGraphExpandRerankPipeline

Compatibility alias for :class:DenseGraphExpandRerankPipeline.

The old name was ambiguous: the pipeline is specifically dense-first, graph-expansion-augmented, then reranked. New code should import DenseGraphExpandRerankPipeline.

Source code in codenib/model/dense_graph_expand_rerank_pipeline.py

def __init__(
    self,
    retriever: EmbeddingRetrievePipeline,
    code_graph: Optional[CodeGraph] = None,
    *,
    rerank_context: Optional[RerankContext] = None,
    repo_path: Optional[str] = None,
    dense_top_k: int = 20,
    graph_seed_top_k: Optional[int] = None,
    graph_neighbors_per_seed: int = 5,
    graph_direction: str = "both",
    graph_edge_types: Optional[Sequence[str]] = (EDGE_TYPE_REFERENCE,),
    graph_fusion_weight: float = 0.0,
    graph_fusion_rrf_k: int = 60,
    final_top_k: int = 10,
    candidate_top_k: Optional[int] = None,
    candidate_filters: Optional[Sequence[CandidatePredicate]] = None,
    include_graph_content: bool = True,
    graph_symbol_only: bool = True,
    graph_require_span: bool = True,
    close_retriever: bool = False,
) -> None:
    if dense_top_k <= 0:
        raise ValueError("dense_top_k must be positive.")
    if final_top_k <= 0:
        raise ValueError("final_top_k must be positive.")
    if graph_neighbors_per_seed < 0:
        raise ValueError("graph_neighbors_per_seed must be non-negative.")
    if graph_seed_top_k is not None and graph_seed_top_k <= 0:
        raise ValueError("graph_seed_top_k must be positive when provided.")
    if graph_fusion_weight < 0:
        raise ValueError("graph_fusion_weight must be non-negative.")
    if graph_fusion_rrf_k <= 0:
        raise ValueError("graph_fusion_rrf_k must be positive.")
    if graph_fusion_weight and rerank_context is None:
        raise ValueError("graph fusion requires a rerank_context.")

    self.retriever = retriever
    self.expand_context = ExpandContext(code_graph=code_graph)
    self.rerank_context = rerank_context
    self.repo_path = repo_path
    self.dense_top_k = dense_top_k
    self.graph_seed_top_k = graph_seed_top_k
    self.graph_neighbors_per_seed = graph_neighbors_per_seed
    self.graph_direction = graph_direction
    self.graph_edge_types = (
        None if graph_edge_types is None else tuple(graph_edge_types)
    )
    self.graph_fusion_weight = graph_fusion_weight
    self.graph_fusion_rrf_k = graph_fusion_rrf_k
    self.final_top_k = final_top_k
    self.candidate_top_k = candidate_top_k
    self.candidate_filters = list(candidate_filters or [])
    self.include_graph_content = include_graph_content
    self.graph_symbol_only = graph_symbol_only
    self.graph_require_span = graph_require_span
    self.close_retriever = close_retriever

    self.last_dense_results: List[QueriedNode] = []
    self.last_graph_results: List[QueriedNode] = []
    self.last_filtered_candidates: List[QueriedNode] = []
    self.last_candidates: List[QueriedNode] = []
    self.last_semantic_results: List[QueriedNode] = []
    self.last_reranked_results: List[QueriedNode] = []

GraphRetrievePipeline

GraphRetrievePipeline(
    repo_path: str,
    index_path: str,
    *,
    stage1_topk: int = 5,
    stage2_topk: int = 50,
    k_hop: int = 2,
    use_ppr: bool = False,
    ppr_damping: float = 0.85,
    use_embedding_rerank: bool = False,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_provider: str = "huggingface",
    embedding_dimension: int = 768,
    languages: list[str] | None = None,
    graph_route: str = "active",
    max_lines_per_chunk: int = 300,
    project_name: str | None = None,
    profiler: Profiler | None = None
)

Bases: SparseSeededGraphRetrievePipeline

Compatibility alias for :class:SparseSeededGraphRetrievePipeline.

Source code in codenib/model/graph_retrieve_pipeline.py

def __init__(
    self,
    repo_path: str,
    index_path: str,
    *,
    stage1_topk: int = 5,
    stage2_topk: int = 50,
    k_hop: int = 2,
    use_ppr: bool = False,
    ppr_damping: float = 0.85,
    use_embedding_rerank: bool = False,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_provider: str = "huggingface",
    embedding_dimension: int = 768,
    languages: Optional[List[str]] = None,
    graph_route: str = "active",
    max_lines_per_chunk: int = 300,
    project_name: Optional[str] = None,
    profiler: Optional[Profiler] = None,
) -> None:
    self.stage1_topk = stage1_topk
    self.stage2_topk = stage2_topk
    self.k_hop = k_hop
    self.use_ppr = use_ppr
    self.ppr_damping = ppr_damping
    self.use_embedding_rerank = use_embedding_rerank
    # Populated by query(): BM25 stage-1 results from the most recent
    # query. Surfaced so callers can compute correctness signals such as
    # bm25_seed_recall@k without re-running BM25.
    self.last_bm25_seeds: List[Any] = []

    pname = project_name or Path(index_path).name
    index_languages = languages or ["python"]

    # Build CodeGraph via the registry-backed graph indexer.
    with _section(profiler, "index.graph_build", {"repo": pname}):
        self.code_graph = build_graph_for_languages(
            repo_path,
            index_path,
            languages=index_languages,
            project_name=pname,
            skip_level="graph",
            profiler=profiler,
            graph_route=graph_route,
        )
        if self.code_graph is None:
            raise RuntimeError(f"Failed to build code graph for {repo_path!r}")

    # Build BM25 index from graph
    with _section(profiler, "index.bm25_build", {"max_k": stage1_topk}):
        self.bm25_index = BM25CodeIndexer(max_k=stage1_topk, language="english")
        self.bm25_index.build_index_from_graph(self.code_graph)

    # Build vector store for Stage 3 if requested
    self.vector_store: Optional[CodeVectorStore] = None
    if use_embedding_rerank:
        with _section(
            profiler,
            "index.vector_store_build",
            {"model": embedding_model, "dim": embedding_dimension},
        ):
            self.vector_store = build_hierarchical_vector_store(
                repo_path=repo_path,
                index_path=index_path,
                plan_name=None,
                languages=index_languages,
                max_lines_per_chunk=max_lines_per_chunk,
                build_levels=["l2"],
                embedding_model=embedding_model,
                embedding_provider=embedding_provider,
                embedding_dimension=embedding_dimension,
                embedding_kwargs={
                    "model_kwargs": {"trust_remote_code": True},
                    "encode_kwargs": {"batch_size": 4},
                },
                index_metric="ip",
                profiler=profiler,
            )

SparseSeededGraphRetrievePipeline

SparseSeededGraphRetrievePipeline(
    repo_path: str,
    index_path: str,
    *,
    stage1_topk: int = 5,
    stage2_topk: int = 50,
    k_hop: int = 2,
    use_ppr: bool = False,
    ppr_damping: float = 0.85,
    use_embedding_rerank: bool = False,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_provider: str = "huggingface",
    embedding_dimension: int = 768,
    languages: list[str] | None = None,
    graph_route: str = "active",
    max_lines_per_chunk: int = 300,
    project_name: str | None = None,
    profiler: Profiler | None = None
)

Sparse-seeded graph-first retrieval pipeline.

Stage 1: BM25 sparse search selects a small set of seed nodes. Stage 2: Graph expansion via k-hop BFS or Personalized PageRank. Stage 3 (optional): Embedding rerank within the expanded set.

This is a graph-first baseline. It should not be used as the direct comparison point for dense embedding retrieval unless that sparse seeding design is the intended ablation.

Parameters:

Name	Type	Description	Default
repo_path	str	Repository root to index.	required
index_path	str	Directory used for graph and vector index caches.	required
stage1_topk	int	Number of BM25 seed nodes (default: 5).	5
stage2_topk	int	Max nodes after graph expansion (default: 50).	50
k_hop	int	BFS expansion depth (default: 2). Ignored when use_ppr=True.	2
use_ppr	bool	Use Personalized PageRank instead of BFS for Stage 2.	False
ppr_damping	float	PPR damping factor (default: 0.85).	0.85
use_embedding_rerank	bool	Enable embedding rerank in Stage 3.	False
embedding_model	str	Embedding model for Stage 3.	'nomic-ai/CodeRankEmbed'
embedding_provider	str	Embedding provider for Stage 3.	'huggingface'
embedding_dimension	int	Embedding vector dimension for Stage 3.	768
languages	list[str] | None	Languages to index. Multiple entries build per-language graphs and merge them before BM25/graph expansion. The full list is also used for Stage 3 chunking.	None
graph_route	str	"active" for public graph routes, "lsp" for backend comparison, or "scip-candidate" to explicitly evaluate candidate SCIP routes.	'active'
max_lines_per_chunk	int	Maximum lines per chunk.	300
project_name	str | None	Project name for SCIP indexing; defaults to index_path basename.	None

Methods:

Name	Description
query	Run the three-stage graph retrieval pipeline.
close	Release pipeline resources.

Source code in codenib/model/graph_retrieve_pipeline.py

def __init__(
    self,
    repo_path: str,
    index_path: str,
    *,
    stage1_topk: int = 5,
    stage2_topk: int = 50,
    k_hop: int = 2,
    use_ppr: bool = False,
    ppr_damping: float = 0.85,
    use_embedding_rerank: bool = False,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_provider: str = "huggingface",
    embedding_dimension: int = 768,
    languages: Optional[List[str]] = None,
    graph_route: str = "active",
    max_lines_per_chunk: int = 300,
    project_name: Optional[str] = None,
    profiler: Optional[Profiler] = None,
) -> None:
    self.stage1_topk = stage1_topk
    self.stage2_topk = stage2_topk
    self.k_hop = k_hop
    self.use_ppr = use_ppr
    self.ppr_damping = ppr_damping
    self.use_embedding_rerank = use_embedding_rerank
    # Populated by query(): BM25 stage-1 results from the most recent
    # query. Surfaced so callers can compute correctness signals such as
    # bm25_seed_recall@k without re-running BM25.
    self.last_bm25_seeds: List[Any] = []

    pname = project_name or Path(index_path).name
    index_languages = languages or ["python"]

    # Build CodeGraph via the registry-backed graph indexer.
    with _section(profiler, "index.graph_build", {"repo": pname}):
        self.code_graph = build_graph_for_languages(
            repo_path,
            index_path,
            languages=index_languages,
            project_name=pname,
            skip_level="graph",
            profiler=profiler,
            graph_route=graph_route,
        )
        if self.code_graph is None:
            raise RuntimeError(f"Failed to build code graph for {repo_path!r}")

    # Build BM25 index from graph
    with _section(profiler, "index.bm25_build", {"max_k": stage1_topk}):
        self.bm25_index = BM25CodeIndexer(max_k=stage1_topk, language="english")
        self.bm25_index.build_index_from_graph(self.code_graph)

    # Build vector store for Stage 3 if requested
    self.vector_store: Optional[CodeVectorStore] = None
    if use_embedding_rerank:
        with _section(
            profiler,
            "index.vector_store_build",
            {"model": embedding_model, "dim": embedding_dimension},
        ):
            self.vector_store = build_hierarchical_vector_store(
                repo_path=repo_path,
                index_path=index_path,
                plan_name=None,
                languages=index_languages,
                max_lines_per_chunk=max_lines_per_chunk,
                build_levels=["l2"],
                embedding_model=embedding_model,
                embedding_provider=embedding_provider,
                embedding_dimension=embedding_dimension,
                embedding_kwargs={
                    "model_kwargs": {"trust_remote_code": True},
                    "encode_kwargs": {"batch_size": 4},
                },
                index_metric="ip",
                profiler=profiler,
            )

query

query(query: str, profiler: Profiler | None = None) -> list[QueriedNode]

Run the three-stage graph retrieval pipeline.

Parameters:

Name	Type	Description	Default
query	str	The search query (e.g., problem statement).	required
profiler	Profiler | None	Optional profiler. When provided, stages are recorded as query.bm25_seed, query.graph_expand, query.embedding_rerank.	None

Returns:

Type	Description
list[QueriedNode]	List of QueriedNode results.

Source code in codenib/model/graph_retrieve_pipeline.py

def query(
    self,
    query: str,
    profiler: Optional[Profiler] = None,
) -> List[QueriedNode]:
    """Run the three-stage graph retrieval pipeline.

    Args:
        query: The search query (e.g., problem statement).
        profiler: Optional profiler. When provided, stages are recorded as
            ``query.bm25_seed``, ``query.graph_expand``,
            ``query.embedding_rerank``.

    Returns:
        List of QueriedNode results.
    """
    # Stage 1: BM25 seed selection
    with _section(profiler, "query.bm25_seed", {"top_k": self.stage1_topk}):
        bm25_results = self.bm25_index.search(query, top_k=self.stage1_topk)
        seed_names = [r.node_name for r in bm25_results]
    self.last_bm25_seeds = list(bm25_results)
    logger.info("Stage 1: %d seed nodes from BM25", len(seed_names))

    # Stage 2: Graph expansion
    with _section(
        profiler,
        "query.graph_expand",
        {"mode": "ppr" if self.use_ppr else "bfs", "top_k": self.stage2_topk},
    ):
        roi = ROISubgraph(self.code_graph)
        if self.use_ppr:
            expanded_nodes = roi.expand_ppr(
                seed_names,
                top_k=self.stage2_topk,
                damping=self.ppr_damping,
                filter_tests=True,
            )
            logger.info(
                "Stage 2: %d nodes after PPR expansion (damping=%.2f)",
                len(expanded_nodes),
                self.ppr_damping,
            )
        else:
            subgraph = roi.extract_subgraph(
                seed_names, k_hop=self.k_hop, direction="both"
            )
            expanded_nodes = roi.get_filtered_subgraph_nodes(
                subgraph, exclude_nodes=None, filter_tests=True
            )
            expanded_nodes = expanded_nodes[: self.stage2_topk]
            logger.info(
                "Stage 2: %d nodes after %d-hop BFS expansion",
                len(expanded_nodes),
                self.k_hop,
            )

    # Stage 3 (optional): Embedding rerank within expanded set only
    if self.use_embedding_rerank and expanded_nodes and self.vector_store:
        with _section(
            profiler,
            "query.embedding_rerank",
            {"candidates": len(expanded_nodes)},
        ):
            mask_ids: set = set()
            for node in expanded_nodes:
                if node.node_name:
                    mask_ids.add(node.node_name)
                if node.node_id:
                    mask_ids.add(node.node_id)
            # Search ONLY within the graph-expanded node set — do NOT search
            # the full FAISS index globally.  This is the key design: graph
            # expansion reduces the corpus, then embedding ranks within it.
            reranked = self.vector_store.search_within_ids(
                query, mask_node_ids=mask_ids, top_k=self.stage2_topk
            )
            results = [
                QueriedNode(
                    node_name=n.node_name,
                    type=n.type,
                    file=n.file,
                    node_id=n.node_id or n.node_name,
                    start_line=n.start_line,
                    end_line=n.end_line,
                    score=n.score,
                    content=n.content,
                )
                for n in reranked
            ]
        logger.info("Stage 3: %d nodes after embedding rerank", len(results))
    else:
        # ROISubgraph returns node_id=None; use node_name as node_id since
        # it has the "file:symbol()" format that extract_predictions expects.
        results = [
            QueriedNode(
                node_name=n.node_name,
                type=n.type,
                file=n.file,
                node_id=n.node_id or n.node_name,
                start_line=n.start_line,
                end_line=n.end_line,
                score=n.score,
                content=n.content,
            )
            for n in expanded_nodes
        ]

    return results

close

close() -> None

Release pipeline resources.

Source code in codenib/model/graph_retrieve_pipeline.py

def close(self) -> None:
    """Release pipeline resources."""
    if self.vector_store is not None:
        self.vector_store.close()
        self.vector_store = None

HybridRetrievePipeline

HybridRetrievePipeline(
    repo_path: str,
    *,
    index_cache_dir: str | None = None,
    stage1_topk: int = 128,
    stage2_topk: int = 50,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_dimension: int = 768,
    languages: Sequence[str] = ("python",),
    embedding_batch_size: int = 8
)

Two-stage cascade retrieval: BM25 → embedding rerank.

Parameters:

Name	Type	Description	Default
repo_path	str	Repository root to index.	required
index_cache_dir	str | None	Directory for cached indexes (default: /.codenib_cache).	None
stage1_topk	int	Number of BM25 candidates to retrieve (default: 128).	128
stage2_topk	int	Final number of results after embedding rerank (default: 50).	50
embedding_model	str	Embedding model name (default: nomic-ai/CodeRankEmbed).	'nomic-ai/CodeRankEmbed'
embedding_dimension	int	Embedding vector dimension (default: 768).	768
languages	Sequence[str]	Languages to index (default: ["python"]).	('python',)
embedding_batch_size	int	Batch size for embedding inference (default: 8 to avoid CUDA OOM).	8

Example

pipeline = HybridRetrievePipeline( ... repo_path="/path/to/repo", ... index_cache_dir="~/.codenib/cache", ... ) results = pipeline.query("how to parse CSV files", top_k=10) for node in results[:5]: ... print(f"{node.file}:{node.node_name} (score={node.score:.3f})")

Methods:

Name	Description
query	Run cascade retrieval: BM25 → embedding rerank.
close	Release pipeline resources.

Source code in codenib/model/hybrid_retrieve_pipeline.py

def __init__(
    self,
    repo_path: str,
    *,
    index_cache_dir: Optional[str] = None,
    stage1_topk: int = 128,
    stage2_topk: int = 50,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_dimension: int = 768,
    languages: Sequence[str] = ("python",),
    embedding_batch_size: int = 8,
):
    self.repo_path = repo_path
    self.stage1_topk = stage1_topk
    self.stage2_topk = stage2_topk

    # Load skill registry (required before build_skill_contexts)
    import os as _os

    from ..agent.skills.loader import SkillLoader
    from ..agent.skills.registry import SkillRegistry

    skills_dir = _os.path.join(
        _os.path.dirname(_os.path.dirname(__file__)), "agent", "skills"
    )
    registry = SkillRegistry()
    loader = SkillLoader()
    loaded_skills = loader.load_all(skills_dir, contexts={}, registry=registry)
    logger.debug("Loaded %d skills into registry", len(loaded_skills))

    # Build indexes via compiler (ADR-0001: always-chunks path)
    logger.info(
        "Building indexes for BM25 + embedding (cascade mode) at %s",
        index_cache_dir or "<repo>/.codenib_cache",
    )
    contexts = build_skill_contexts(
        repo_path=repo_path,
        skill_ids=["bm25_search", "embedding_search"],
        languages=list(languages),
        cache_dir=index_cache_dir,
        skill_registry=registry,  # Pass loaded registry
        embedding_model=embedding_model,
        embedding_dimension=embedding_dimension,
        default_top_k=stage1_topk,  # BM25 max_k
        default_level="l2",  # embedding level
        trust_remote_code=True,  # Required for nomic-ai/CodeRankEmbed
        embedding_batch_size=embedding_batch_size,  # Avoid CUDA OOM
    )

    retrieve_ctx = contexts.get("retrieve")
    if retrieve_ctx is None:
        raise RuntimeError(
            "build_skill_contexts did not return a 'retrieve' context. "
            "Check skill registry and index_requirements."
        )

    self.bm25 = retrieve_ctx.bm25
    self.vector_store = retrieve_ctx.vector_store

    if self.bm25 is None:
        raise RuntimeError("BM25 index not built (bm25_search skill failed).")
    if self.vector_store is None:
        raise RuntimeError(
            "Embedding index not built (embedding_search skill failed)."
        )

    logger.info(
        "HybridRetrievePipeline ready: BM25 (top-%d) → embedding rerank (top-%d)",
        stage1_topk,
        stage2_topk,
    )

query

query(query: str, top_k: int | None = None) -> list[QueriedNode]

Run cascade retrieval: BM25 → embedding rerank.

Parameters:

Name	Type	Description	Default
query	str	Search query (e.g., problem statement).	required
top_k	int | None	Number of final results (overrides stage2_topk if provided).	None

Returns:

Type	Description
list[QueriedNode]	List of QueriedNode sorted by embedding similarity (descending).

Source code in codenib/model/hybrid_retrieve_pipeline.py

def query(self, query: str, top_k: Optional[int] = None) -> List[QueriedNode]:
    """Run cascade retrieval: BM25 → embedding rerank.

    Args:
        query: Search query (e.g., problem statement).
        top_k: Number of final results (overrides stage2_topk if provided).

    Returns:
        List of QueriedNode sorted by embedding similarity (descending).
    """
    final_k = top_k if top_k is not None else self.stage2_topk

    # Stage 1: BM25 sparse search (recall-oriented filter)
    logger.debug("Stage 1: BM25 search (top_k=%d)", self.stage1_topk)
    bm25_results = self.bm25.search(
        query=query,
        top_k=self.stage1_topk,
        return_code_content=False,  # embedding rerank will fetch content
        wrap_with_ln=False,
        filter_test=False,
    )

    if not bm25_results:
        logger.warning("BM25 returned no results for query: %r", query)
        return []

    # Extract node_ids for embedding rerank mask
    candidate_ids = {r.node_id for r in bm25_results if r.node_id}
    logger.debug(
        "Stage 1 complete: %d candidates (unique node_ids: %d)",
        len(bm25_results),
        len(candidate_ids),
    )

    # Stage 2: Embedding rerank within BM25 candidates
    logger.debug(
        "Stage 2: embedding rerank within %d candidates", len(candidate_ids)
    )
    reranked = self.vector_store.search_within_ids(
        query=query,
        mask_node_ids=candidate_ids,
        top_k=final_k,
    )

    logger.info(
        "Pipeline complete: %d results (BM25=%d → embedding=%d)",
        len(reranked),
        len(bm25_results),
        len(reranked),
    )

    # Convert to QueriedNode format (search_within_ids returns NodeInfo)
    return [
        QueriedNode(
            node_name=n.node_name,
            type=n.type,
            file=n.file,
            node_id=n.node_id or n.node_name,
            start_line=n.start_line,
            end_line=n.end_line,
            score=n.score,
            content=n.content,
        )
        for n in reranked
    ]

close

close() -> None

Release pipeline resources.

Source code in codenib/model/hybrid_retrieve_pipeline.py

def close(self) -> None:
    """Release pipeline resources."""
    if self.vector_store is not None:
        self.vector_store.close()
        self.vector_store = None
    self.bm25 = None

GraphExpansionPlan dataclass

GraphExpansionPlan(
    seed_engine: str = "sparse",
    seed_top_k: int = 8,
    expand_top_k: int = 80,
    hops: int = 2,
    direction: str = "both",
    use_ppr: bool = False,
)

Graph expansion controls for structural retrieval paths.

RetrievalBudget dataclass

RetrievalBudget(
    tier: str = "balanced",
    retrieve_top_k: int = 100,
    rerank_candidate_top_k: int | None = 50,
    allow_graph: bool = True,
    allow_rerank: bool = True,
    allow_llm_rerank: bool = False,
    prefer_rrf: bool = True,
)

Coarse retrieval budget used for path selection.

RetrievalCapabilities dataclass

RetrievalCapabilities(
    has_dense: bool = True,
    has_sparse: bool = True,
    has_graph: bool = False,
    has_embedding_rerank: bool = True,
    has_llm_rerank: bool = False,
)

Backends that a planner may use for this pipeline instance.

RetrievalPathPlan dataclass

RetrievalPathPlan(
    name: str,
    intent: str,
    stages: tuple[RetrievalStagePlan, ...],
    retrieval_top_k: int,
    fusion: str = "weighted",
    enable_rerank: bool = False,
    rerank_strategy: str | None = None,
    rerank_candidate_top_k: int | None = None,
    graph: GraphExpansionPlan | None = None,
    rationale: str = "",
)

Concrete path selected for a query.

RetrievalPlanner

Select one of CodeNib's retrieval paths for a query and budget.

Methods:

Name	Description
normalize_budget	Return a validated budget object from a preset name or custom budget.
classify_query	Classify a query into lexical, semantic, and structural signals.
select	Select a concrete retrieval path for query under budget.

normalize_budget

normalize_budget(budget: BudgetInput = 'balanced') -> RetrievalBudget

Return a validated budget object from a preset name or custom budget.

Source code in codenib/model/retrieval_planner.py

def normalize_budget(self, budget: BudgetInput = "balanced") -> RetrievalBudget:
    """Return a validated budget object from a preset name or custom budget."""
    if isinstance(budget, RetrievalBudget):
        return budget
    key = (budget or "balanced").strip().lower()
    aliases = {
        "cheap": "fast",
        "low": "fast",
        "default": "balanced",
        "quality": "thorough",
        "full": "thorough",
    }
    key = aliases.get(key, key)
    try:
        return self._BUDGET_PRESETS[key]
    except KeyError as exc:
        raise ValueError(
            "Unknown retrieval budget. Choose from: fast, balanced, thorough."
        ) from exc

classify_query

classify_query(query: str) -> QuerySignals

Classify a query into lexical, semantic, and structural signals.

Source code in codenib/model/retrieval_planner.py

def classify_query(self, query: str) -> QuerySignals:
    """Classify a query into lexical, semantic, and structural signals."""
    text = query or ""
    words = re.findall(r"[A-Za-z0-9_]+", text)
    structural = any(pattern.search(text) for pattern in self._STRUCTURAL_PATTERNS)
    lexical = any(pattern.search(text) for pattern in self._LEXICAL_PATTERNS)
    semantic = bool(self._SEMANTIC_WORDS.search(text)) or len(words) >= 10
    if not lexical and not semantic and words:
        lexical = len(words) <= 4
    return QuerySignals(
        lexical=lexical,
        semantic=semantic,
        structural=structural,
    )

select

select(
    query: str,
    budget: BudgetInput = "balanced",
    capabilities: RetrievalCapabilities | None = None,
) -> RetrievalPathPlan

Select a concrete retrieval path for query under budget.

Source code in codenib/model/retrieval_planner.py

def select(
    self,
    query: str,
    budget: BudgetInput = "balanced",
    capabilities: Optional[RetrievalCapabilities] = None,
) -> RetrievalPathPlan:
    """Select a concrete retrieval path for *query* under *budget*."""
    resolved_budget = self.normalize_budget(budget)
    resolved_caps = capabilities or RetrievalCapabilities()
    signals = self.classify_query(query)

    if (
        signals.structural
        and resolved_budget.allow_graph
        and resolved_caps.has_graph
        and resolved_caps.has_sparse
    ):
        return self._structural_graph_plan(resolved_budget, resolved_caps)

    if resolved_budget.tier == "fast":
        if signals.lexical and resolved_caps.has_sparse:
            return self._fast_lexical_plan(resolved_budget)
        if resolved_caps.has_dense:
            return self._semantic_plan(
                resolved_budget,
                resolved_caps,
                enable_rerank=False,
                rationale="fast budget disables rerank",
            )
        if resolved_caps.has_sparse:
            return self._fast_lexical_plan(resolved_budget)

    if signals.semantic and not signals.lexical and resolved_caps.has_dense:
        return self._semantic_plan(resolved_budget, resolved_caps)

    if resolved_caps.has_dense and resolved_caps.has_sparse:
        return self._hybrid_fusion_plan(resolved_budget, resolved_caps, signals)

    if resolved_caps.has_dense:
        return self._semantic_plan(resolved_budget, resolved_caps)

    if resolved_caps.has_sparse:
        return self._fast_lexical_plan(resolved_budget)

    raise ValueError("No retrieval backend is available for planning.")

RetrievalStagePlan dataclass

RetrievalStagePlan(
    engine: str,
    weight: float = 1.0,
    top_k: int | None = None,
    params: dict[str, object] = dict(),
)

Declarative retrieval branch selected by the planner.

RetrieveRerankPipeline

RetrieveRerankPipeline(
    repo_path: str,
    index_path: str,
    *,
    retrieval_plan: Sequence[RetrieveStageConfig] | None = None,
    retrieval_mode: str = "dense",
    retrieval_level: str = "l2",
    planning_budget: BudgetInput = "balanced",
    retrieval_planner: RetrievalPlanner | None = None,
    planner_capabilities: RetrievalCapabilities | None = None,
    code_graph: CodeGraph | None = None,
    fusion_strategy: str = "weighted",
    rrf_k: int = 60,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_provider: str = "huggingface",
    embedding_dimension: int = 768,
    embedding_model_kwargs: dict | None = None,
    rerank_model: str = "openai/Qwen/Qwen2.5-Coder-7B",
    rerank_temperature: float = 0.0,
    rerank_max_tokens: int = 2048,
    rerank_strategy: str = "llm",
    rerank_embedding_model: str | None = None,
    rerank_embedding_provider: str | None = None,
    rerank_embedding_dimension: int | None = None,
    rerank_embedding_model_kwargs: dict | None = None,
    rerank_index_metric: str = "ip",
    crossencoder_model: str = "Qwen/Qwen3-Reranker-0.6B",
    crossencoder_batch_size: int = 8,
    languages: list[str] | None = None,
    max_lines_per_chunk: int = 100,
    sparse_max_k: int = 128,
    rerank_window_size: int | None = None,
    rerank_window_step: int | None = None,
    rerank_listwise_format: str = "structured",
    enable_rerank: bool = True,
    vector_masks: dict[str, set[str]] | None = None,
    rerank_candidate_top_k: int | None = None
)

Composable code retrieval pipeline built from retrieval + rerank ops.

The pipeline wires :mod:codenib.ops.retrieve (sparse/dense/hybrid) with :mod:codenib.ops.rerank to form a two-stage search baseline. Retrieval branches are described through :class:RetrieveStageConfig objects, which can be combined (e.g., hybrid BM25 + embedding) and weighted before feeding results to an LLM reranker.

Parameters:

Name	Type	Description	Default
repo_path	str	Repository root to index.	required
index_path	str	Directory used for vector index caches.	required
retrieval_plan	Sequence[RetrieveStageConfig] | None	Optional sequence of :class:RetrieveStageConfig. When omitted, :func:build_retrieve_plan is invoked with the provided retrieval_mode.	None
retrieval_mode	str	Shortcut for :func:build_retrieve_plan.	'dense'
retrieval_level	str	Context expansion level, either "l0" or "l2".	'l2'
planning_budget	BudgetInput	Budget used by automatic retrieval planning.	'balanced'
retrieval_planner	RetrievalPlanner | None	Planner used when retrieval_mode is "auto".	None
planner_capabilities	RetrievalCapabilities | None	Optional capability override supplied to the planner.	None
code_graph	CodeGraph | None	Graph used to expand retrieved code context.	None
fusion_strategy	str	Strategy used to combine multiple retrieval branches.	'weighted'
rrf_k	int	Rank constant used by reciprocal-rank fusion.	60
embedding_model	str	Dense embedding model identifier.	'nomic-ai/CodeRankEmbed'
embedding_provider	str	Backend serving embedding_model.	'huggingface'
embedding_dimension	int	Dense vector width.	768
embedding_model_kwargs	dict | None	Extra keyword arguments for the dense encoder.	None
rerank_model	str	Reranker model identifier.	'openai/Qwen/Qwen2.5-Coder-7B'
rerank_temperature	float	Sampling temperature for the LLM reranker.	0.0
rerank_max_tokens	int	Token budget for a single rerank call.	2048
rerank_strategy	str	Reranking backend: "llm", "embedding", or "crossencoder".	'llm'
rerank_embedding_model	str | None	Model used by embedding-based reranking.	None
rerank_embedding_provider	str | None	Backend for rerank_embedding_model.	None
rerank_embedding_dimension	int | None	Vector width for embedding-based reranking.	None
rerank_embedding_model_kwargs	dict | None	Extra keyword arguments for the rerank embedding model.	None
rerank_index_metric	str	Similarity metric for the rerank vector index.	'ip'
crossencoder_model	str	Model used by cross-encoder reranking.	'Qwen/Qwen3-Reranker-0.6B'
crossencoder_batch_size	int	Batch size used by the cross-encoder.	8
languages	list[str] | None	Languages to chunk for indexing (default: ["python"]).	None
max_lines_per_chunk	int	Maximum lines per chunk passed to chunker.	100
sparse_max_k	int	Upper bound for BM25 index fan-out; defaults to 128.	128
rerank_window_size	int | None	Sliding window width for the reranker (see :meth:RerankAgent.rerank_nodes). When None, the reranker considers all candidates at once.	None
rerank_window_step	int | None	Stride between consecutive rerank windows.	None
rerank_listwise_format	str	Output format requested from listwise reranking.	'structured'
enable_rerank	bool	Whether to apply the configured reranking stage.	True
vector_masks	dict[str, set[str]] | None	Optional per-language symbol masks for vector retrieval.	None
rerank_candidate_top_k	int | None	Maximum candidates passed into reranking.	None

Methods:

Name	Description
close	Release model/index resources held by the pipeline.
query	Execute retrieve + rerank plan for the provided query.

Source code in codenib/model/retrieve_rerank_pipeline.py

def __init__(
    self,
    repo_path: str,
    index_path: str,
    *,
    retrieval_plan: Optional[Sequence[RetrieveStageConfig]] = None,
    retrieval_mode: str = "dense",
    retrieval_level: str = "l2",
    planning_budget: BudgetInput = "balanced",
    retrieval_planner: Optional[RetrievalPlanner] = None,
    planner_capabilities: Optional[RetrievalCapabilities] = None,
    code_graph: Optional[CodeGraph] = None,
    fusion_strategy: str = "weighted",
    rrf_k: int = 60,
    embedding_model: str = "nomic-ai/CodeRankEmbed",
    embedding_provider: str = "huggingface",
    embedding_dimension: int = 768,
    embedding_model_kwargs: Optional[dict] = None,
    rerank_model: str = "openai/Qwen/Qwen2.5-Coder-7B",
    rerank_temperature: float = 0.0,
    rerank_max_tokens: int = 2048,
    rerank_strategy: str = "llm",
    rerank_embedding_model: Optional[str] = None,
    rerank_embedding_provider: Optional[str] = None,
    rerank_embedding_dimension: Optional[int] = None,
    rerank_embedding_model_kwargs: Optional[dict] = None,
    rerank_index_metric: str = "ip",
    crossencoder_model: str = "Qwen/Qwen3-Reranker-0.6B",
    crossencoder_batch_size: int = 8,
    languages: Optional[List[str]] = None,
    max_lines_per_chunk: int = 100,
    sparse_max_k: int = 128,
    rerank_window_size: Optional[int] = None,
    rerank_window_step: Optional[int] = None,
    rerank_listwise_format: str = "structured",
    enable_rerank: bool = True,
    vector_masks: Optional[Dict[str, Set[str]]] = None,
    rerank_candidate_top_k: Optional[int] = None,
) -> None:
    self.repo_path = self._validate_repo(repo_path)
    self.index_path = Path(index_path)
    self.index_path.mkdir(parents=True, exist_ok=True)
    self.languages = languages or ["python"]
    self.max_lines_per_chunk = max_lines_per_chunk
    self._chunks = None
    self.enable_rerank = enable_rerank
    self.index_metric = "ip"
    self.profiler = None
    self.retrieval_mode = (retrieval_mode or "dense").strip().lower()
    self.retrieval_planner = retrieval_planner
    if self.retrieval_planner is None and self.retrieval_mode == "auto":
        self.retrieval_planner = RetrievalPlanner()
    self.planning_budget = planning_budget
    self.fusion_strategy = _normalize_fusion_strategy(fusion_strategy)
    if rrf_k <= 0:
        raise ValueError("rrf_k must be positive.")
    self.rrf_k = rrf_k
    self.last_selected_plan: Optional[RetrievalPathPlan] = None
    self.last_query_signals = None
    self.last_planning_budget: Optional[RetrievalBudget] = None
    self.last_planner_capabilities: Optional[RetrievalCapabilities] = None
    self.last_planner_trace: Dict[str, object] = {}
    self._planner_capabilities_override = planner_capabilities
    self.expand_context = ExpandContext(code_graph=code_graph)
    if retrieval_level not in ("l0", "l2"):
        raise ValueError(
            f"Invalid retrieval_level {retrieval_level!r}. Must be 'l0' or 'l2'."
        )
    self.retrieval_level = retrieval_level

    if retrieval_plan:
        plan = list(retrieval_plan)
    elif self.retrieval_mode == "auto" and self.retrieval_planner is not None:
        preload_top_k = max(
            128,
            self.retrieval_planner.normalize_budget(
                self.planning_budget
            ).retrieve_top_k,
        )
        plan = [
            _stage_config_from_planner(stage)
            for stage in build_planner_preload_stages(top_k=preload_top_k)
        ]
    else:
        plan = build_retrieve_plan(self.retrieval_mode)
    if not plan:
        raise ValueError("Retrieval plan must contain at least one stage.")
    self.retrieve_plan: List[RetrieveStageConfig] = plan
    self._default_stage_top_k = (
        max((stage.top_k or 0) for stage in self.retrieve_plan) or 32
    )

    self.vector_store: Optional[CodeVectorStore] = None
    self.bm25_index: Optional[BM25CodeIndexer] = None
    self.rerank_vector_store: Optional[CodeVectorStore] = None

    embedding_kwargs = self._prepare_embedding_kwargs(embedding_model_kwargs)
    if self._needs_engine("dense"):
        self.vector_store = self._initialize_vector_store(
            embedding_model=embedding_model,
            embedding_provider=embedding_provider,
            embedding_dimension=embedding_dimension,
            embedding_kwargs=embedding_kwargs,
        )

    if self._needs_engine("sparse"):
        sparse_cap = max(
            sparse_max_k,
            max(
                (stage.top_k or self._default_stage_top_k)
                for stage in self.retrieve_plan
                if stage.engine == "sparse"
            ),
        )
        self.bm25_index = self._initialize_bm25_index(max_k=sparse_cap)

    strategy = (rerank_strategy or "llm").strip().lower()
    if strategy not in ("llm", "embedding", "crossencoder"):
        raise ValueError(
            "rerank_strategy must be 'llm', 'embedding', or 'crossencoder'."
        )
    self.rerank_strategy = strategy

    rerank_llm = None
    self.cross_encoder = None
    if strategy == "llm":
        rerank_llm = LiteLLMChat(
            model=rerank_model,
            max_tokens=rerank_max_tokens,
            temperature=rerank_temperature,
        )
    elif strategy == "crossencoder":
        self.cross_encoder = build_reranker(
            crossencoder_model,
            batch_size=crossencoder_batch_size,
        )
        logger.info(
            "Cross-encoder reranker loaded",
            extra={"model": crossencoder_model},
        )
    else:
        rerank_model_kwargs = self._prepare_embedding_kwargs(
            rerank_embedding_model_kwargs
        )
        self.rerank_vector_store = self._initialize_rerank_vector_store(
            embedding_model=rerank_embedding_model or embedding_model,
            embedding_provider=rerank_embedding_provider or embedding_provider,
            embedding_dimension=rerank_embedding_dimension or embedding_dimension,
            embedding_kwargs=rerank_model_kwargs,
            index_metric=rerank_index_metric,
        )

    self.rerank_window_size = (
        rerank_window_size
        if rerank_window_size and rerank_window_size > 0
        else None
    )
    self.rerank_window_step = (
        rerank_window_step
        if rerank_window_step and rerank_window_step > 0
        else None
    )
    self.rerank_candidate_top_k = (
        rerank_candidate_top_k
        if rerank_candidate_top_k and rerank_candidate_top_k > 0
        else None
    )

    self.retrieve_context = RetrieveContext(
        bm25=self.bm25_index,
        vector_store=self.vector_store,
        regex_index=None,
        default_top_k=self._default_stage_top_k,
        default_level=self.retrieval_level,
        masks=vector_masks or {},
    )

    if self.enable_rerank:
        self.rerank_context = RerankContext(
            llm=rerank_llm,
            embedding_store=self.rerank_vector_store or self.vector_store,
            candidate_top_k=self.rerank_candidate_top_k,
            window_size=rerank_window_size,
            window_step=rerank_window_step,
            listwise_format=rerank_listwise_format,
        )
    else:
        self.rerank_context = None

    self.planner_capabilities = (
        self._planner_capabilities_override
        if self._planner_capabilities_override is not None
        else RetrievalCapabilities(
            has_dense=bool(self.vector_store),
            has_sparse=bool(self.bm25_index),
            has_graph=bool(self.expand_context.code_graph),
            has_embedding_rerank=bool(self.vector_store),
            has_llm_rerank=strategy == "llm",
        )
    )

    logger.info(
        "RetrieveRerankPipeline initialized",
        extra={
            "repo": self.repo_path,
            "index_path": str(self.index_path),
            "retrieval_mode": self.retrieval_mode,
            "retrieval_plan": [stage.engine for stage in self.retrieve_plan],
            "retrieval_level": self.retrieval_level,
            "dense_index": bool(self.vector_store),
            "sparse_index": bool(self.bm25_index),
            "graph_index": bool(self.expand_context.code_graph),
            "planner": bool(self.retrieval_planner),
            "rerank_strategy": (
                self.rerank_strategy if self.enable_rerank else "disabled"
            ),
            "enable_rerank": self.enable_rerank,
        },
    )

close

close() -> None

Release model/index resources held by the pipeline.

Source code in codenib/model/retrieve_rerank_pipeline.py

def close(self) -> None:
    """Release model/index resources held by the pipeline."""
    if self.vector_store is not None:
        self.vector_store.close()
    if self.rerank_vector_store is not None:
        self.rerank_vector_store.close()
    if self.bm25_index is not None:
        self.bm25_index.documents = []
        self.bm25_index.nodes = []
        self.bm25_index.retriever = None
        self.bm25_index = None

    self.vector_store = None
    self.rerank_vector_store = None
    self._chunks = None

    if self.cross_encoder is not None:
        try:
            self.cross_encoder.close()
        except Exception:
            pass
        self.cross_encoder = None

    try:
        import gc

        import torch

        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass

query

query(
    query: str, top_k: int = 10, *, budget: BudgetInput | None = None
) -> list[QueriedNode]

Execute retrieve + rerank plan for the provided query.

Parameters:

Name	Type	Description	Default
query	str	The search query.	required
top_k	int	Number of results to return. If rerank is enabled, this controls the rerank output. Retrieval always fetches RETRIEVAL_TOP_K candidates internally.	10
budget	BudgetInput | None	Optional per-query planner budget. Only used when a retrieval planner is configured, including retrieval_mode=auto.	None

Source code in codenib/model/retrieve_rerank_pipeline.py

def query(
    self,
    query: str,
    top_k: int = 10,
    *,
    budget: Optional[BudgetInput] = None,
) -> List[QueriedNode]:
    """Execute retrieve + rerank plan for the provided query.

    Args:
        query: The search query.
        top_k: Number of results to return. If rerank is enabled, this controls
            the rerank output. Retrieval always fetches RETRIEVAL_TOP_K
            candidates internally.
        budget: Optional per-query planner budget. Only used when a
            retrieval planner is configured, including ``retrieval_mode=auto``.
    """
    selected_plan = self._select_query_plan(query, budget)
    active_plan = (
        [_stage_config_from_planner(stage) for stage in selected_plan.stages]
        if selected_plan is not None
        else self.retrieve_plan
    )
    fusion = selected_plan.fusion if selected_plan is not None else None
    merge_top_k = (
        selected_plan.retrieval_top_k
        if selected_plan is not None
        else RETRIEVAL_TOP_K
    )

    # Step 1: Run each retrieval branch
    branch_results: List[List[QueriedNode]] = []
    for stage in active_plan:
        results = self._run_retrieval_stage(query, stage)
        branch_results.append(results)

    # Step 2: Merge branches (hybrid if multiple)
    if len(branch_results) == 1:
        candidates = branch_results[0]
    else:
        weights = [stage.weight for stage in active_plan]
        candidates = self._merge_hybrid(
            branch_results,
            weights,
            merge_top_k,
            fusion=fusion or self.fusion_strategy,
            rrf_k=self.rrf_k,
        )

    if selected_plan is not None and selected_plan.graph is not None:
        candidates = self._run_graph_expansion_plan(selected_plan, candidates)

    # Step 3: Rerank if enabled
    plan_allows_rerank = (
        selected_plan.enable_rerank if selected_plan is not None else True
    )
    if (
        self.enable_rerank
        and plan_allows_rerank
        and self.rerank_context is not None
    ):
        candidates = self._run_rerank(
            query,
            candidates,
            top_k,
            strategy=(
                selected_plan.rerank_strategy
                if selected_plan is not None
                else self.rerank_strategy
            ),
            candidate_top_k=self._effective_rerank_candidate_top_k(selected_plan),
        )

    return candidates[:top_k]

RetrieveStageConfig dataclass

RetrieveStageConfig(
    engine: str = "dense",
    weight: float = 1.0,
    top_k: int | None = None,
    params: dict[str, object] = dict(),
)

Declarative configuration for an individual retrieval branch.

build_retrieve_plan

build_retrieve_plan(mode: str = 'dense') -> list[RetrieveStageConfig]

Convenience helper for common retrieval plan presets.

auto returns the dense+sparse preload set used by :class:RetrievalPlanner; the per-query path is selected at query time.

Source code in codenib/model/retrieve_rerank_pipeline.py

def build_retrieve_plan(mode: str = "dense") -> List[RetrieveStageConfig]:
    """Convenience helper for common retrieval plan presets.

    ``auto`` returns the dense+sparse preload set used by
    :class:`RetrievalPlanner`; the per-query path is selected at query time.
    """

    normalized = (mode or "dense").strip().lower()
    if normalized == "dense":
        return [RetrieveStageConfig(engine="dense", top_k=RETRIEVAL_TOP_K)]
    if normalized == "sparse":
        return [RetrieveStageConfig(engine="sparse", top_k=RETRIEVAL_TOP_K)]
    if normalized == "auto":
        return [
            _stage_config_from_planner(stage)
            for stage in build_planner_preload_stages(top_k=128)
        ]
    if normalized == "hybrid":
        return [
            RetrieveStageConfig(engine="dense", weight=0.6, top_k=RETRIEVAL_TOP_K),
            RetrieveStageConfig(engine="sparse", weight=0.4, top_k=RETRIEVAL_TOP_K),
        ]
    raise ValueError(
        f"Unknown retrieval mode {mode!r}. Choose from: auto, dense, sparse, hybrid."
    )