codenib.index ¶

Index-related components: sparse, dense, and regex.

Modules:

Name	Description
`embedding`	Embedding module for CodeNib.
`incremental`	Incremental vector-indexing layer for CodeNib.
`regex_idx`
`rerank`	Cross-encoder rerankers (pairwise (query, doc) → relevance score).
`sparse_idx`	Sparse indexing module for CodeNib.
`trigram`	Trigram-based search backends.

Classes:

Name	Description
`CodeVectorStore`	Vector store for code embeddings using FAISS and sentence-transformers.
`RegexNodeIndex`	In-memory regex-based index for CodeGraph nodes.
`BM25CodeIndexer`	A class that builds a BM25 index from CodeGraph nodes and provides

Functions:

Name	Description
`create_code_vector_store`	Factory function to create a CodeVectorStore.

CodeVectorStore ¶

CodeVectorStore(
    embedding_model: str = "text-embedding-ada-002",
    embedding_provider: str = "openai",
    dimension: int = 1536,
    index_type: str = "flat",
    index_metric: str = "ip",
    ivf_nlist: int = 100,
    ivf_nprobe: int = 8,
    store_path: str | None = None,
    profiler: Profiler | None = None,
    embedding: Any | None = None,
    artifact_metadata: dict[str, Any] | None = None,
    **embedding_kwargs
)

Vector store for code embeddings using FAISS and sentence-transformers. Provides semantic search capabilities over code chunks.

Supports hierarchical indexing: - L0: File-level skeletons - L2: Function/method-level chunks for fine-grained retrieval (default)

Parameters:

Name	Type	Description	Default
`embedding_model`	`str`	Name of the embedding model to use	`'text-embedding-ada-002'`
`embedding_provider`	`str`	Provider for embeddings ("openai" or "huggingface")	`'openai'`
`dimension`	`int`	Dimension of the embedding vectors	`1536`
`index_type`	`str`	FAISS index type — "flat" (exact brute force, default) or "ivf" (IVF inverted-file; approximate, faster at scale). IVF indices are trained lazily on the first batch of vectors.	`'flat'`
`index_metric`	`str`	Distance metric ("ip" for inner product, "l2" for L2 distance)	`'ip'`
`ivf_nlist`	`int`	IVF only — number of Voronoi cells (coarse centroids). On small corpora it is clamped down to the training-set size, since FAISS k-means needs at least `nlist` training points.	`100`
`ivf_nprobe`	`int`	IVF only — cells probed per query; the recall/latency knob. Clamped to the effective `nlist`.	`8`
`store_path`	`str \| None`	Path to store/load the vector store	`None`
`profiler`	`Profiler \| None`	Optional profiler instance to capture detailed timings	`None`
`embedding`	`Any \| None`	A pre-built embedding wrapper to reuse. When several stores share one model (e.g. one per repo), pass the same instance so the model is loaded onto the GPU only once.	`None`
`artifact_metadata`	`dict[str, Any] \| None`	Optional immutable source/build identity persisted with the top-level configuration.	`None`
`**embedding_kwargs`		Additional arguments for embedding model	`{}`

Methods:

Name	Description
`reuse_query_embedding`	Reuse one query vector within a composed request, then discard it.
`clear_query_cache`	Clear the single-query embedding reused by consecutive search stages.
`swap_index`	Hot-swap the FAISS index without reloading the embedding model.
`close`	Release embeddings and FAISS resources to free memory.
`add_code_chunks`	Add code chunks to the vector store.
`add_nodes_with_content`	Add NodeInfo objects (with content) to the vector store.
`search`	Search for similar code chunks using semantic similarity.
`search_with_content`	Search and return results with content included.
`search_within_ids`	Search only within a restricted set of node IDs.
`hierarchical_search`	Note: This method is implemented by Claude and is just for future reference.
`save`	Save the vector store to disk.
`load`	Load the vector store from disk.
`get_stats`	Get statistics about the vector store.
`get_embeddings_by_content_hash`	Extract raw embedding vectors from the FAISS index, keyed by content hash.
`rebuild_from_embeddings`	Clear level and rebuild its FAISS index from pre-computed embeddings.
`delta_update`	Patch a flat FAISS index in place when the change set is small.
`clear`	Clear data from the vector store.

Attributes:

Name	Type	Description
`closed`	`bool`	Whether every resource owned by this store has been released.

Source code in codenib/index/embedding/vector_store.py

def __init__(
    self,
    embedding_model: str = "text-embedding-ada-002",
    embedding_provider: str = "openai",
    dimension: int = 1536,
    index_type: str = "flat",
    index_metric: str = "ip",
    ivf_nlist: int = 100,
    ivf_nprobe: int = 8,
    store_path: Optional[str] = None,
    profiler: Optional[Profiler] = None,
    embedding: Optional[Any] = None,
    artifact_metadata: Optional[Dict[str, Any]] = None,
    **embedding_kwargs,
):
    """
    Initialize the CodeVectorStore.

    Args:
        embedding_model: Name of the embedding model to use
        embedding_provider: Provider for embeddings ("openai" or "huggingface")
        dimension: Dimension of the embedding vectors
        index_type: FAISS index type — "flat" (exact brute force, default)
            or "ivf" (IVF inverted-file; approximate, faster at scale).
            IVF indices are trained lazily on the first batch of vectors.
        index_metric: Distance metric ("ip" for inner product, "l2" for L2 distance)
        ivf_nlist: IVF only — number of Voronoi cells (coarse centroids). On
            small corpora it is clamped down to the training-set size, since
            FAISS k-means needs at least ``nlist`` training points.
        ivf_nprobe: IVF only — cells probed per query; the recall/latency
            knob. Clamped to the effective ``nlist``.
        store_path: Path to store/load the vector store
        profiler: Optional profiler instance to capture detailed timings
        embedding: A pre-built embedding wrapper to reuse. When several
            stores share one model (e.g. one per repo), pass the same
            instance so the model is loaded onto the GPU only once.
        artifact_metadata: Optional immutable source/build identity persisted
            with the top-level configuration.
        **embedding_kwargs: Additional arguments for embedding model
    """
    self.embedding_model = embedding_model
    self.embedding_provider = _validate_provider_options(
        embedding_provider,
        embedding_kwargs,
    )
    self.embedding_load_policy = (
        resolve_embedding_load_policy_from_options(
            embedding_model,
            embedding_kwargs,
        )
        if self.embedding_provider == "huggingface"
        else None
    )
    self.embedding_revision = (
        self.embedding_load_policy.revision if self.embedding_load_policy else None
    )
    self.embedding_trust_remote_code = bool(
        self.embedding_load_policy and self.embedding_load_policy.trust_remote_code
    )
    self.dimension = dimension
    self.index_type = index_type.lower()
    if self.index_type not in ("flat", "ivf"):
        raise ValueError(
            f"Unsupported index_type: {index_type}. Must be 'flat' or 'ivf'."
        )
    self.index_metric = index_metric.lower()
    if self.index_metric not in ["ip", "l2"]:
        raise ValueError(
            f"Unsupported index_metric: {index_metric}. Must be 'ip' or 'l2'."
        )
    self.ivf_nlist = max(1, int(ivf_nlist))
    self.ivf_nprobe = max(1, int(ivf_nprobe))
    initial_store_path = Path(store_path) if store_path else None
    self.profiler = profiler
    initial_artifact_metadata = dict(artifact_metadata or {})

    # Initialize the embedding model — or reuse a shared one so the same
    # model isn't loaded onto the GPU once per store.
    self.embedding = (
        embedding
        if embedding is not None
        else self._initialize_embedding_model(**embedding_kwargs)
    )
    self._cached_query_text: Optional[str] = None
    self._cached_query_vector: Optional[np.ndarray] = None
    self._query_cache_depth = 0
    self.dimension = self._infer_embedding_dimension(dimension)

    # Initialize L0 (file-level skeletons)
    l0_index = self._build_faiss_index()

    # Initialize L2 (function/method-level) - default
    l2_index = self._build_faiss_index()
    self._loaded_state = _LoadedVectorState(
        l0_index=l0_index,
        l0_documents=[],
        l2_index=l2_index,
        l2_documents=[],
        artifact_metadata=initial_artifact_metadata,
        store_path=initial_store_path,
    )

    logger.info(
        f"Initialized CodeVectorStore with {embedding_provider}:{embedding_model}"
    )

closed `property` ¶

closed: bool

Whether every resource owned by this store has been released.

reuse_query_embedding ¶

reuse_query_embedding()

Reuse one query vector within a composed request, then discard it.

Source code in codenib/index/embedding/vector_store.py

@contextmanager
def reuse_query_embedding(self):
    """Reuse one query vector within a composed request, then discard it."""

    depth = getattr(self, "_query_cache_depth", 0)
    if depth == 0:
        self.clear_query_cache()
    self._query_cache_depth = depth + 1
    try:
        yield
    finally:
        self._query_cache_depth -= 1
        if self._query_cache_depth == 0:
            self.clear_query_cache()

clear_query_cache ¶

clear_query_cache() -> None

Clear the single-query embedding reused by consecutive search stages.

Source code in codenib/index/embedding/vector_store.py

def clear_query_cache(self) -> None:
    """Clear the single-query embedding reused by consecutive search stages."""

    self._cached_query_text = None
    self._cached_query_vector = None

swap_index ¶

swap_index(
    path: str, *, native_index_authorization: NativeIndexAuthorization | None = None
) -> None

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

The replacement is fully loaded and validated before the current L0/L2 state is released. The embedding model is left intact so the caller can reuse the same model across many instances.

Source code in codenib/index/embedding/vector_store.py

def swap_index(
    self,
    path: str,
    *,
    native_index_authorization: NativeIndexAuthorization | None = None,
) -> None:
    """Hot-swap the FAISS index without reloading the embedding model.

    The replacement is fully loaded and validated before the current
    L0/L2 state is released. The embedding model is left intact so the
    caller can reuse the same model across many instances.
    """
    self.load(
        path,
        native_index_authorization=native_index_authorization,
    )

close ¶

close() -> None

Release embeddings and FAISS resources to free memory.

Source code in codenib/index/embedding/vector_store.py

def close(self) -> None:
    """Release embeddings and FAISS resources to free memory."""

    state = self._loaded_state
    released: dict[int, bool] = {}
    deferred: BaseException | None = None
    for index in (state.l0_index, state.l2_index):
        if index is None or id(index) in released:
            continue
        reset = getattr(index, "reset", None)
        if not callable(reset):
            released[id(index)] = True
            continue
        try:
            reset()
            released[id(index)] = True
        except BaseException as exc:  # noqa: B036 - visit both native indices
            released[id(index)] = False
            if deferred is None:
                deferred = exc
            else:
                _annotate_secondary_error(
                    deferred,
                    "additional FAISS index cleanup also failed",
                    exc,
                )

    l0_released = state.l0_index is None or released.get(id(state.l0_index), False)
    l2_released = state.l2_index is None or released.get(id(state.l2_index), False)
    self._loaded_state = _LoadedVectorState(
        l0_index=None if l0_released else state.l0_index,
        l0_documents=[] if l0_released else state.l0_documents,
        l2_index=None if l2_released else state.l2_index,
        l2_documents=[] if l2_released else state.l2_documents,
        artifact_metadata=state.artifact_metadata,
        store_path=state.store_path,
    )
    if deferred is not None:
        raise deferred

    self.embedding = None
    self._query_cache_depth = 0
    self._cached_query_text = None
    self._cached_query_vector = None

    try:
        import torch

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

add_code_chunks ¶

add_code_chunks(code_chunks: list[dict[str, Any]], level: Level = 'l2') -> None

Add code chunks to the vector store.

Parameters:

Name	Type	Description	Default
`code_chunks`	`list[dict[str, Any]]`	List of code chunk dictionaries with content and metadata	required
`level`	`Level`	Index level to add chunks to ("l0" for file skeletons, "l2" for functions/methods)	`'l2'`

Source code in codenib/index/embedding/vector_store.py

def add_code_chunks(
    self, code_chunks: List[Dict[str, Any]], level: Level = "l2"
) -> None:
    """
    Add code chunks to the vector store.

    Args:
        code_chunks: List of code chunk dictionaries with content and metadata
        level: Index level to add chunks to
            ("l0" for file skeletons, "l2" for functions/methods)
    """
    if not code_chunks:
        logger.warning("No code chunks provided")
        return

    index, documents_list = self._get_index_and_docs(level)
    logger.info(f"Adding {len(code_chunks)} code chunks to {level} vector store")

    # Convert chunks to _Document objects
    documents: List[_Document] = []
    for i, chunk in enumerate(code_chunks):
        content = chunk.get("content", "")
        content_hash = hashlib.md5(
            content.encode("utf-8", errors="replace")
        ).hexdigest()
        metadata = {
            "chunk_id": len(documents_list) + i,
            "chunk_type": chunk.get("chunk_type", "unknown"),
            "name": chunk.get("name", f"chunk_{i}"),
            "file": chunk.get("file", ""),
            "start_line": chunk.get("start_line", 0),
            "end_line": chunk.get("end_line", 0),
            "node_id": chunk.get("node_id", ""),
            "level": level,
            "content_hash": content_hash,
        }
        for key, value in chunk.items():
            if key not in ["content"] and key not in metadata:
                metadata[key] = value

        documents.append(_Document(page_content=content, metadata=metadata))

    # Store documents
    documents_list.extend(documents)

    texts = [doc.page_content for doc in documents]

    # Phase 1: Embed texts (typically the bottleneck)
    with self._profile_section(
        f"embedding_encode_{level}",
        {"num_documents": len(documents), "level": level},
    ):
        embeddings = self.embedding.embed_documents(texts)

    # Phase 2: Add pre-computed vectors to FAISS index
    with self._profile_section(
        f"faiss_index_add_{level}",
        {"num_vectors": len(embeddings), "level": level},
    ):
        vectors = np.array(embeddings, dtype=np.float32)
        self._add_to_index(level, vectors)

    logger.info(
        f"Successfully added {len(documents)} documents to {level} vector store"
    )

add_nodes_with_content ¶

add_nodes_with_content(nodes: list[NodeInfo], level: Level = 'l2') -> None

Add NodeInfo objects (with content) to the vector store.

Parameters:

Name	Type	Description	Default
`nodes`	`list[NodeInfo]`	List of NodeInfo objects	required
`level`	`Level`	Index level to add nodes to ("l0" or "l2")	`'l2'`

Source code in codenib/index/embedding/vector_store.py

def add_nodes_with_content(
    self, nodes: List[NodeInfo], level: Level = "l2"
) -> None:
    """
    Add NodeInfo objects (with content) to the vector store.

    Args:
        nodes: List of NodeInfo objects
        level: Index level to add nodes to ("l0" or "l2")
    """
    chunks = []
    for node in nodes:
        chunk = {
            "content": node.content,
            "chunk_type": node.type,
            "name": node.node_name,
            "file": node.file,
            "start_line": node.start_line,
            "end_line": node.end_line,
        }
        chunks.append(chunk)

    self.add_code_chunks(chunks, level=level)

search ¶

search(
    query: str,
    top_k: int = 10,
    score_threshold: float | None = None,
    level: Level = "l2",
    mask_node_ids: set[str] | None = None,
) -> list[NodeInfo]

Search for similar code chunks using semantic similarity.

Parameters:

Name	Type	Description	Default
`query`	`str`	Search query text	required
`top_k`	`int`	Number of top results to return	`10`
`score_threshold`	`float \| None`	Minimum similarity score threshold	`None`
`level`	`Level`	Index level to search ("l0" for file skeletons, "l2" for functions/methods)	`'l2'`
`mask_node_ids`	`set[str] \| None`	Optional set of CodeChunk.node_id values to filter results.	`None`

Returns:

Type	Description
`list[NodeInfo]`	List of NodeInfo objects with scores populated

Source code in codenib/index/embedding/vector_store.py

def search(
    self,
    query: str,
    top_k: int = 10,
    score_threshold: Optional[float] = None,
    level: Level = "l2",
    mask_node_ids: Optional[Set[str]] = None,
) -> List[NodeInfo]:
    """
    Search for similar code chunks using semantic similarity.

    Args:
        query: Search query text
        top_k: Number of top results to return
        score_threshold: Minimum similarity score threshold
        level: Index level to search ("l0" for file skeletons, "l2" for
            functions/methods)
        mask_node_ids: Optional set of CodeChunk.node_id values to filter results.

    Returns:
        List of NodeInfo objects with scores populated
    """
    index, documents = self._get_index_and_docs(level)

    if index is None or index.ntotal == 0:
        logger.warning(f"No {level} vector store available. Add code chunks first.")
        return []

    logger.debug(f"Searching {level} for: {query[:100]}...")

    docs_with_scores = self._search_index(query, index, documents, top_k)

    results = []
    for doc, score in docs_with_scores:
        metadata = doc.metadata

        if score_threshold is not None and self._should_filter_by_threshold(
            score, score_threshold
        ):
            continue

        node_with_score = NodeInfo(
            node_name=metadata.get("name", "unknown"),
            type=metadata.get("chunk_type", "unknown"),
            file=_result_source_file(metadata),
            node_id=metadata.get("node_id", ""),
            start_line=metadata.get("start_line", 0),
            end_line=metadata.get("end_line", 0),
            score=float(score),
        )
        results.append(node_with_score)

    if mask_node_ids:
        results = [r for r in results if r.node_id in mask_node_ids]
    if top_k:
        results = results[:top_k]

    logger.debug(
        f"Found {len(results)} results in {level} (masked={bool(mask_node_ids)})"
    )
    return results

search_with_content ¶

search_with_content(
    query: str,
    top_k: int = 10,
    score_threshold: float | None = None,
    level: Level = "l2",
    mask_node_ids: set[str] | None = None,
) -> list[NodeInfo]

Search and return results with content included.

Parameters:

Name	Type	Description	Default
`query`	`str`	Search query text	required
`top_k`	`int`	Number of top results to return	`10`
`score_threshold`	`float \| None`	Minimum similarity score threshold	`None`
`level`	`Level`	Index level to search ("l0" for file skeletons, "l2" for functions/methods)	`'l2'`
`mask_node_ids`	`set[str] \| None`	Optional set of CodeChunk.node_id values to filter results.	`None`

Returns:

Type	Description
`list[NodeInfo]`	List of NodeInfo objects with content populated

Source code in codenib/index/embedding/vector_store.py

def search_with_content(
    self,
    query: str,
    top_k: int = 10,
    score_threshold: Optional[float] = None,
    level: Level = "l2",
    mask_node_ids: Optional[Set[str]] = None,
) -> List[NodeInfo]:
    """
    Search and return results with content included.

    Args:
        query: Search query text
        top_k: Number of top results to return
        score_threshold: Minimum similarity score threshold
        level: Index level to search ("l0" for file skeletons, "l2" for
            functions/methods)
        mask_node_ids: Optional set of CodeChunk.node_id values to filter results.

    Returns:
        List of NodeInfo objects with content populated
    """
    index, documents = self._get_index_and_docs(level)

    if index is None or index.ntotal == 0:
        logger.warning(f"No {level} vector store available. Add code chunks first.")
        return []

    docs_with_scores = self._search_index(query, index, documents, top_k)

    results = []
    for doc, score in docs_with_scores:
        metadata = doc.metadata

        if score_threshold is not None and self._should_filter_by_threshold(
            score, score_threshold
        ):
            continue

        node_with_content = NodeInfo(
            node_name=metadata.get("name", "unknown"),
            type=metadata.get("chunk_type", "unknown"),
            file=_result_source_file(metadata),
            node_id=metadata.get("node_id", ""),
            start_line=metadata.get("start_line", 0),
            end_line=metadata.get("end_line", 0),
            score=float(score),
            content=doc.page_content,
        )
        results.append(node_with_content)

    if mask_node_ids:
        results = [r for r in results if r.node_id in mask_node_ids]
    if top_k:
        results = results[:top_k]

    logger.debug(
        f"Found {len(results)} results with content in {level} "
        f"(masked={bool(mask_node_ids)})"
    )
    return results

search_within_ids ¶

search_within_ids(
    query: str, mask_node_ids: set[str], top_k: int = 10, level: Level = "l2"
) -> list[NodeInfo]

Search only within a restricted set of node IDs.

Instead of searching the full FAISS index globally and filtering afterwards, this method restricts the search space before computing similarity. It reconstructs stored vectors for matching documents and computes similarity against the query embedding directly.

Parameters:

Name	Type	Description	Default
`query`	`str`	Search query text.	required
`mask_node_ids`	`set[str]`	Set of node_id / node_name values to restrict search to.	required
`top_k`	`int`	Number of top results to return.	`10`
`level`	`Level`	Index level to search.	`'l2'`

Returns:

Type	Description
`list[NodeInfo]`	List of NodeInfo objects sorted by similarity score.

Source code in codenib/index/embedding/vector_store.py

def search_within_ids(
    self,
    query: str,
    mask_node_ids: Set[str],
    top_k: int = 10,
    level: Level = "l2",
) -> List[NodeInfo]:
    """Search only within a restricted set of node IDs.

    Instead of searching the full FAISS index globally and filtering
    afterwards, this method restricts the search space *before* computing
    similarity.  It reconstructs stored vectors for matching documents
    and computes similarity against the query embedding directly.

    Args:
        query: Search query text.
        mask_node_ids: Set of node_id / node_name values to restrict
            search to.
        top_k: Number of top results to return.
        level: Index level to search.

    Returns:
        List of NodeInfo objects sorted by similarity score.
    """
    index, documents = self._get_index_and_docs(level)
    if index is None or index.ntotal == 0:
        logger.warning(f"No {level} vector store available.")
        return []

    # Find documents whose node_id or name is in mask set
    matched: list[tuple[int, _Document]] = []
    for i, doc in enumerate(documents):
        meta = doc.metadata
        if (
            meta.get("node_id", "") in mask_node_ids
            or meta.get("name", "") in mask_node_ids
        ):
            matched.append((i, doc))

    if not matched:
        logger.debug("search_within_ids: no matching documents found")
        return []

    # Encode query
    query_vec = self._embed_query(query)

    # Reconstruct stored vectors and compute similarity
    results: list[NodeInfo] = []
    for faiss_idx, doc in matched:
        vec = index.reconstruct(faiss_idx)
        if self.index_metric == "ip":
            score = float(np.dot(query_vec, vec))
        else:  # l2 — lower is better
            score = float(np.sum((query_vec - vec) ** 2))

        metadata = doc.metadata
        results.append(
            NodeInfo(
                node_name=metadata.get("name", "unknown"),
                type=metadata.get("chunk_type", "unknown"),
                file=_result_source_file(metadata),
                node_id=metadata.get("node_id", ""),
                start_line=metadata.get("start_line", 0),
                end_line=metadata.get("end_line", 0),
                score=score,
                content=doc.page_content,
            )
        )

    best_by_node = {}
    for result in results:
        identity = result.node_id or (
            result.file,
            result.node_name,
            result.start_line,
            result.end_line,
        )
        existing = best_by_node.get(identity)
        if existing is None:
            best_by_node[identity] = result
            continue
        if self.index_metric == "ip":
            is_better = result.score > existing.score
        else:
            is_better = result.score < existing.score
        if is_better:
            best_by_node[identity] = result
    results = list(best_by_node.values())

    # Sort: ip → higher is better; l2 → lower is better
    results.sort(
        key=lambda r: r.score,
        reverse=(self.index_metric == "ip"),
    )

    logger.debug(
        "search_within_ids: %d matched, %d unique, returning top %d",
        len(matched),
        len(results),
        min(top_k, len(results)),
    )
    return results[:top_k]

hierarchical_search ¶

hierarchical_search(
    query: str,
    l0_top_k: int = 5,
    l2_top_k: int = 10,
    l0_score_threshold: float | None = None,
    l2_score_threshold: float | None = None,
    filter_l2_by_l0: bool = True,
) -> dict[str, list[NodeInfo]]

Note: This method is implemented by Claude and is just for future reference. Perform hierarchical search: first L0 (files), then L2 (functions).

This implements a coarse-to-fine retrieval strategy: 1. Search L0 to find relevant files based on their skeletons 2. Search L2 for specific functions/methods 3. Optionally filter L2 results to only include those from L0 files

Parameters:

Name	Type	Description	Default
`query`	`str`	Search query text	required
`l0_top_k`	`int`	Number of top L0 results (files)	`5`
`l2_top_k`	`int`	Number of top L2 results (functions/methods)	`10`
`l0_score_threshold`	`float \| None`	Score threshold for L0 results	`None`
`l2_score_threshold`	`float \| None`	Score threshold for L2 results	`None`
`filter_l2_by_l0`	`bool`	If True, only return L2 results from files found in L0	`True`

Returns:

Type	Description
`dict[str, list[NodeInfo]]`	Dict with 'l0' and 'l2' keys containing search results

Source code in codenib/index/embedding/vector_store.py

def hierarchical_search(
    self,
    query: str,
    l0_top_k: int = 5,
    l2_top_k: int = 10,
    l0_score_threshold: Optional[float] = None,
    l2_score_threshold: Optional[float] = None,
    filter_l2_by_l0: bool = True,
) -> Dict[str, List[NodeInfo]]:
    """
    Note: This method is implemented by Claude and is just for future reference.
    Perform hierarchical search: first L0 (files), then L2 (functions).

    This implements a coarse-to-fine retrieval strategy:
    1. Search L0 to find relevant files based on their skeletons
    2. Search L2 for specific functions/methods
    3. Optionally filter L2 results to only include those from L0 files

    Args:
        query: Search query text
        l0_top_k: Number of top L0 results (files)
        l2_top_k: Number of top L2 results (functions/methods)
        l0_score_threshold: Score threshold for L0 results
        l2_score_threshold: Score threshold for L2 results
        filter_l2_by_l0: If True, only return L2 results from files found in L0

    Returns:
        Dict with 'l0' and 'l2' keys containing search results
    """
    # Step 1: Search L0 to find relevant files
    l0_results = self.search(
        query, top_k=l0_top_k, score_threshold=l0_score_threshold, level="l0"
    )

    # Step 2: Search L2 for functions/methods (fetch more if filtering)
    l2_fetch_k = l2_top_k * 3 if filter_l2_by_l0 else l2_top_k
    l2_results = self.search(
        query, top_k=l2_fetch_k, score_threshold=l2_score_threshold, level="l2"
    )

    # Step 3: Optionally filter L2 results by L0 files
    if filter_l2_by_l0 and l0_results:
        l0_files = {result.file for result in l0_results}
        l2_results = [r for r in l2_results if r.file in l0_files]

    # Limit L2 results to top_k
    l2_results = l2_results[:l2_top_k]

    return {
        "l0": l0_results,
        "l2": l2_results,
    }

save ¶

save(path: str | None = None) -> None

Save the vector store to disk.

Parameters:

Name	Type	Description	Default
`path`	`str \| None`	Path to save the store (uses self.store_path if not provided)	`None`

Source code in codenib/index/embedding/vector_store.py

def save(self, path: Optional[str] = None) -> None:
    """
    Save the vector store to disk.

    Args:
        path: Path to save the store (uses self.store_path if not provided)
    """
    save_path = Path(path) if path else self.store_path
    if save_path is None:
        raise ValueError("No save path provided")

    save_path = Path(save_path)
    save_path.mkdir(parents=True, exist_ok=True)

    logger.info(f"Saving vector store to {save_path}")

    model_suffix = self.embedding_model.replace("/", "__")

    level_state = {}
    for level in ("l0", "l2"):
        index, documents = self._get_index_and_docs(level)
        if documents and (index is None or int(index.ntotal) != len(documents)):
            vector_count = 0 if index is None else int(index.ntotal)
            raise ValueError(
                f"Cannot save misaligned {level} vector store: "
                f"{vector_count} vectors for {len(documents)} documents"
            )
        if index is not None:
            self._validate_loaded_faiss_index(
                index,
                relative=f"{level}/index_{model_suffix}.faiss",
                document_count=len(documents),
            )
        level_state[level] = (index, documents)

    config_path = save_path / f"config_{model_suffix}.json"
    save_marker = save_path / f".{config_path.name}.save-in-progress"
    _atomic_json_dump(
        save_marker,
        {"persistence_schema": VECTOR_PERSISTENCE_SCHEMA},
    )
    try:
        level_artifacts: dict[str, dict[str, dict[str, Any]]] = {}
        for level, (_index, documents) in level_state.items():
            if documents:
                level_artifacts[level] = self._save_level(
                    save_path, level, model_suffix
                )
            else:
                self._remove_level_files(save_path, level, model_suffix)

        config = {
            "embedding_model": self.embedding_model,
            "embedding_provider": self.embedding_provider,
            "embedding_revision": self.embedding_revision,
            "dimension": self.dimension,
            "index_type": self.index_type,
            "index_metric": self.index_metric,
            "l0_documents": len(self.l0_documents),
            "l2_documents": len(self.l2_documents),
            "persistence_schema": VECTOR_PERSISTENCE_SCHEMA,
            "level_artifacts": level_artifacts,
        }
        if self.artifact_metadata:
            config["artifact"] = self.artifact_metadata
        if (
            type(self.artifact_metadata.get("builder_schema")) is int
            and self.artifact_metadata.get("builder_schema") == 8
        ):
            config["row_mapping"] = VECTOR_ROW_MAPPING_CONTRACT
        # This config is the commit record for both levels. Publishing it
        # last makes an interrupted multi-file save detectable on load.
        _atomic_json_dump(config_path, config)
    except Exception:
        # Keep the marker so a later load cannot accept a partially
        # replaced legacy artifact. A successful retry removes it.
        raise
    else:
        save_marker.unlink()

    logger.info("Vector store saved successfully")

load ¶

load(
    path: str | None = None,
    *,
    native_index_authorization: NativeIndexAuthorization | None = None
) -> None

Load the vector store from disk.

Parameters:

Name	Type	Description	Default
`path`	`str \| None`	Path to load the store from (uses self.store_path if not provided)	`None`
`native_index_authorization`	`NativeIndexAuthorization \| None`	Process-local authorization bound to the exact captured tree and semantic view contract. Artifact fields cannot provide this capability.	`None`

Source code in codenib/index/embedding/vector_store.py

def load(
    self,
    path: Optional[str] = None,
    *,
    native_index_authorization: NativeIndexAuthorization | None = None,
) -> None:
    """
    Load the vector store from disk.

    Args:
        path: Path to load the store from (uses self.store_path if not provided)
        native_index_authorization: Process-local authorization bound to the
            exact captured tree and semantic view contract. Artifact fields
            cannot provide this capability.
    """
    ambient_error = sys.exc_info()[1]
    # Exception subclasses can shadow ``__traceback__``. Read the base
    # descriptor so the ambient snapshot cannot run user code.
    ambient_traceback = (
        None
        if ambient_error is None
        else BaseException.__traceback__.__get__(
            ambient_error,
            type(ambient_error),
        )
    )
    load_path = Path(path) if path else self.store_path
    if load_path is None:
        raise ValueError("No load path provided")
    require_native_index_authorization_preflight(
        native_index_authorization,
        view_type="vector",
    )
    view = capture_authenticated_vector_view(load_path)
    loaded_state: _LoadedVectorState | None = None
    previous_state: _LoadedVectorState | None = None
    view_closed = False
    body_completed = False
    try:
        require_native_index_authorization(
            native_index_authorization,
            view.ownership,
            view_type="vector",
            semantic_contract=self.artifact_metadata,
        )
        loaded_state = self._load_captured(view)
        view.verify_final()
        # A close failure is still a failed replacement. Finish the
        # captured-view lifecycle before exposing any newly parsed state.
        view.close()
        view_closed = True
        previous_state = self._loaded_state
        self._replace_loaded_state(loaded_state)
        body_completed = True
    finally:
        primary_error = sys.exc_info()[1]
        if (
            body_completed
            and primary_error is ambient_error
            and ambient_error is not None
            and BaseException.__traceback__.__get__(
                primary_error,
                type(primary_error),
            )
            is ambient_traceback
        ):
            primary_error = None
        cleanup_error: BaseException | None = None
        if not view_closed:
            try:
                view.close()
                view_closed = True
            except BaseException as exc:  # noqa: B036 - preserve primary
                cleanup_error = exc
        if loaded_state is not None:
            release_state = (
                previous_state
                if self._loaded_state is loaded_state
                else loaded_state
            )
            if release_state is not None:
                try:
                    self._release_vector_indices(
                        tuple(
                            index
                            for index in (
                                release_state.l0_index,
                                release_state.l2_index,
                            )
                            if index is not self.l0_index
                            and index is not self.l2_index
                        )
                    )
                except BaseException as exc:  # noqa: B036 - visit all cleanup
                    if primary_error is not None:
                        _attach_vector_index_cleanup_owner(primary_error, exc)
                    if cleanup_error is None:
                        cleanup_error = exc
                    else:
                        _attach_vector_index_cleanup_owner(cleanup_error, exc)
                        _annotate_secondary_error(
                            cleanup_error,
                            "vector index cleanup also failed",
                            exc,
                        )
        if cleanup_error is not None:
            if primary_error is not None:
                if not view_closed:
                    _attach_vector_view_cleanup_owner(
                        primary_error,
                        view,
                        cleanup_error,
                    )
                _annotate_secondary_error(
                    primary_error,
                    "vector load cleanup also failed",
                    cleanup_error,
                )
            else:
                raise cleanup_error

get_stats ¶

get_stats() -> dict[str, Any]

Get statistics about the vector store.

Returns:

Type	Description
`dict[str, Any]`	Dictionary with store statistics

Source code in codenib/index/embedding/vector_store.py

def get_stats(self) -> Dict[str, Any]:
    """
    Get statistics about the vector store.

    Returns:
        Dictionary with store statistics
    """
    stats = {
        "embedding_model": self.embedding_model,
        "embedding_provider": self.embedding_provider,
        "embedding_revision": self.embedding_revision,
        "dimension": self.dimension,
        "index_type": self.index_type,
        "index_metric": self.index_metric,
        "l0_documents": len(self.l0_documents),
        "l2_documents": len(self.l2_documents),
        "total_documents": len(self.l0_documents) + len(self.l2_documents),
    }

    # Analyze L0 chunk types
    if self.l0_documents:
        l0_chunk_types = {}
        for doc in self.l0_documents:
            chunk_type = doc.metadata.get("chunk_type", "unknown")
            l0_chunk_types[chunk_type] = l0_chunk_types.get(chunk_type, 0) + 1
        stats["l0_chunk_types"] = l0_chunk_types

    # Analyze L2 chunk types
    if self.l2_documents:
        l2_chunk_types = {}
        for doc in self.l2_documents:
            chunk_type = doc.metadata.get("chunk_type", "unknown")
            l2_chunk_types[chunk_type] = l2_chunk_types.get(chunk_type, 0) + 1
        stats["l2_chunk_types"] = l2_chunk_types

    return stats

get_embeddings_by_content_hash ¶

get_embeddings_by_content_hash(level: Level = 'l2') -> dict[str, ndarray]

Extract raw embedding vectors from the FAISS index, keyed by content hash.

This is used to seed the EmbeddingsCache after a full build so that the first incremental update achieves ~100% cache hit rate for unchanged chunks.

Each document's content is MD5-hashed to produce the key. If the document metadata already contains a content_hash field it is used directly; otherwise the hash is computed on the fly.

Returns:

Type	Description
`dict[str, ndarray]`	Dict mapping content_hash → np.ndarray (float32 vectors).

Source code in codenib/index/embedding/vector_store.py

def get_embeddings_by_content_hash(
    self, level: Level = "l2"
) -> Dict[str, np.ndarray]:
    """
    Extract raw embedding vectors from the FAISS index, keyed by content hash.

    This is used to seed the ``EmbeddingsCache`` after a full build so that
    the first incremental update achieves ~100% cache hit rate for unchanged
    chunks.

    Each document's content is MD5-hashed to produce the key.  If the
    document metadata already contains a ``content_hash`` field it is used
    directly; otherwise the hash is computed on the fly.

    Returns:
        Dict mapping content_hash → np.ndarray (float32 vectors).
    """
    index, documents = self._get_index_and_docs(level)
    if not documents or index is None or index.ntotal == 0:
        return {}

    result: Dict[str, np.ndarray] = {}
    for i, doc in enumerate(documents):
        content_hash = doc.metadata.get("content_hash")
        if content_hash is None:
            content_hash = hashlib.md5(
                doc.page_content.encode("utf-8", errors="replace")
            ).hexdigest()

        vec = index.reconstruct(i)
        result[content_hash] = np.asarray(vec, dtype=np.float32)

    logger.info(
        "Extracted %d embedding vectors from %s FAISS index for cache seeding.",
        len(result),
        level,
    )
    return result

rebuild_from_embeddings ¶

rebuild_from_embeddings(
    documents: list, embeddings: list[ndarray], level: Level = "l2"
) -> None

Clear level and rebuild its FAISS index from pre-computed embeddings.

Used by the incremental update path: unchanged chunks contribute their cached vectors, so only genuinely new/modified chunks require model inference. No embedding model calls are made by this method.

Parameters:

Name	Type	Description	Default
`documents`	`list`	Document-like objects with `page_content` and `metadata` attributes (`_Document` or compatible).	required
`embeddings`	`list[ndarray]`	Corresponding embedding vectors as `np.ndarray` (shape `[dim]`, dtype `float32`).	required
`level`	`Level`	Which index level to rebuild (`"l0"` or `"l2"`).	`'l2'`

Raises:

Type	Description
`ValueError`	If documents and embeddings have different lengths.

Source code in codenib/index/embedding/vector_store.py

def rebuild_from_embeddings(
    self,
    documents: list,
    embeddings: List[np.ndarray],
    level: Level = "l2",
) -> None:
    """
    Clear *level* and rebuild its FAISS index from pre-computed embeddings.

    Used by the incremental update path: unchanged chunks contribute their
    cached vectors, so only genuinely new/modified chunks require model
    inference.  No embedding model calls are made by this method.

    Args:
        documents: Document-like objects with ``page_content`` and
            ``metadata`` attributes (``_Document`` or compatible).
        embeddings: Corresponding embedding vectors as ``np.ndarray``
            (shape ``[dim]``, dtype ``float32``).
        level: Which index level to rebuild (``"l0"`` or ``"l2"``).

    Raises:
        ValueError: If *documents* and *embeddings* have different lengths.
    """
    if len(documents) != len(embeddings):
        raise ValueError(
            f"documents ({len(documents)}) and embeddings ({len(embeddings)}) "
            "must have the same length."
        )

    # Wipe the existing index for this level
    self.clear(level)

    if not documents:
        logger.debug(
            "rebuild_from_embeddings: no documents; level %s cleared.", level
        )
        return

    # Convert to _Document if needed and add vectors to the raw FAISS index
    native_docs = [_to_document(d) for d in documents]
    vectors = np.array(
        [
            emb if isinstance(emb, np.ndarray) else np.asarray(emb)
            for emb in embeddings
        ],
        dtype=np.float32,
    )

    self._add_to_index(level, vectors)
    if level == "l0":
        self.l0_documents = native_docs
    else:
        self.l2_documents = native_docs

    logger.info(
        "rebuild_from_embeddings: %s index rebuilt with %d documents.",
        level,
        len(documents),
    )

delta_update ¶

delta_update(
    all_documents: list,
    all_embeddings: list[ndarray],
    changed_content_hashes: set[str],
    level: Level = "l2",
    threshold: float = 0.1,
) -> None

Patch a flat FAISS index in place when the change set is small.

When the fraction of changed chunks is below threshold, this uses IndexFlat.remove_ids + add to modify only the affected rows, keeping unchanged vectors and their aligned documents untouched. IVF indexes are rebuilt because removing their implicit IDs does not compact the remaining labels to match the document array. If the change ratio exceeds the threshold (or the index is empty), this also falls back to :meth:rebuild_from_embeddings.

Parameters:

Name	Type	Description	Default
`all_documents`	`list`	The complete desired set of documents for level after the update. Must carry `content_hash` in metadata.	required
`all_embeddings`	`list[ndarray]`	Corresponding embedding vectors, aligned with all_documents.	required
`changed_content_hashes`	`set[str]`	Content hashes of chunks that were added, removed, or modified in this update cycle. Used both to decide between delta/rebuild and to identify stale rows.	required
`level`	`Level`	Which index level to update.	`'l2'`
`threshold`	`float`	Maximum change ratio (changed/total) for the delta path; above this a full rebuild is performed.	`0.1`

Source code in codenib/index/embedding/vector_store.py

def delta_update(
    self,
    all_documents: list,
    all_embeddings: List[np.ndarray],
    changed_content_hashes: Set[str],
    level: Level = "l2",
    threshold: float = 0.1,
) -> None:
    """
    Patch a flat FAISS index in place when the change set is small.

    When the fraction of changed chunks is below *threshold*, this uses
    ``IndexFlat.remove_ids`` + ``add`` to modify only the affected rows,
    keeping unchanged vectors and their aligned documents untouched. IVF
    indexes are rebuilt because removing their implicit IDs does not
    compact the remaining labels to match the document array. If the
    change ratio exceeds the threshold (or the index is empty), this also
    falls back to :meth:`rebuild_from_embeddings`.

    Args:
        all_documents: The complete desired set of documents for *level*
            after the update.  Must carry ``content_hash`` in metadata.
        all_embeddings: Corresponding embedding vectors, aligned with
            *all_documents*.
        changed_content_hashes: Content hashes of chunks that were
            added, removed, or modified in this update cycle.  Used both
            to decide between delta/rebuild and to identify stale rows.
        level: Which index level to update.
        threshold: Maximum change ratio (changed/total) for the delta
            path; above this a full rebuild is performed.
    """
    total = len(all_documents)

    if total == 0:
        self.clear(level)
        return

    index, current_docs = self._get_index_and_docs(level)
    change_ratio = len(changed_content_hashes) / total

    # Fall back to full rebuild when the delta path can't help.
    if (
        index is None
        or index.ntotal == 0
        or not current_docs
        or self.index_type != "flat"
        or change_ratio > threshold
    ):
        logger.info(
            "delta_update: %d/%d changed (%.0f%%) → full rebuild of %s.",
            len(changed_content_hashes),
            total,
            change_ratio * 100,
            level,
        )
        self.rebuild_from_embeddings(all_documents, all_embeddings, level=level)
        return

    # --- Delta path: in-place patch -------------------------------
    # Use a list per hash so duplicate-content docs (same code in
    # different files) are all preserved.
    from collections import defaultdict

    target_by_hash: Dict[str, List[Tuple[object, np.ndarray]]] = defaultdict(list)
    for doc, emb in zip(all_documents, all_embeddings, strict=True):
        ch = doc.metadata.get("content_hash")
        if ch is None:
            # Can't align by hash → safest to rebuild.
            logger.warning(
                "delta_update: target doc missing content_hash → full rebuild."
            )
            self.rebuild_from_embeddings(all_documents, all_embeddings, level=level)
            return
        target_by_hash[ch].append((doc, emb))

    current_hashes = [d.metadata.get("content_hash") for d in current_docs]

    # For each hash, allow at most target-count survivors (handles
    # both duplicate-content additions and removals correctly).
    target_avail: Dict[str, int] = {h: len(v) for h, v in target_by_hash.items()}
    rows_to_remove: List[int] = []
    for i, h in enumerate(current_hashes):
        if (
            h is None
            or h not in target_avail
            or h in changed_content_hashes
            or target_avail[h] <= 0
        ):
            rows_to_remove.append(i)
        else:
            target_avail[h] -= 1

    # Unclaimed target entries become additions.
    docs_to_add: List[Tuple[object, np.ndarray]] = []
    for h, entries in target_by_hash.items():
        claimed = len(entries) - target_avail.get(h, 0)
        docs_to_add.extend(entries[claimed:])

    if rows_to_remove:
        selector = faiss.IDSelectorBatch(np.array(rows_to_remove, dtype=np.int64))
        index.remove_ids(selector)

    # Survivors: prefer the fresh target doc (same content_hash) so that
    # pure metadata changes — file rename, start_line shift, name edit —
    # are reflected without requiring a full rebuild.  The vector is
    # identical because content_hash is identical, so no FAISS op needed.
    remove_set = set(rows_to_remove)
    survivor_idx: Dict[str, int] = {}
    new_docs_list: List[_Document] = []
    for i, d in enumerate(current_docs):
        if i in remove_set:
            continue
        h = current_hashes[i]
        idx = survivor_idx.get(h, 0)
        survivor_idx[h] = idx + 1
        entries = target_by_hash.get(h)
        if entries and idx < len(entries):
            new_docs_list.append(_to_document(entries[idx][0]))
        else:
            new_docs_list.append(d)

    if docs_to_add:
        add_vectors = np.array(
            [np.asarray(e, dtype=np.float32) for _, e in docs_to_add],
            dtype=np.float32,
        )
        self._add_to_index(level, add_vectors)
        new_docs_list.extend(_to_document(d) for d, _ in docs_to_add)

    if level == "l0":
        self.l0_documents = new_docs_list
    else:
        self.l2_documents = new_docs_list

    logger.info(
        "delta_update: %s patched in place — removed %d, added %d "
        "(ntotal=%d, %.0f%% changed).",
        level,
        len(rows_to_remove),
        len(docs_to_add),
        index.ntotal,
        change_ratio * 100,
    )

clear ¶

clear(level: Level | None = None) -> None

Clear data from the vector store.

Parameters:

Name	Type	Description	Default
`level`	`Level \| None`	If specified, only clear that level ("l0" or "l2"). If None, clear both levels.	`None`

Source code in codenib/index/embedding/vector_store.py

def clear(self, level: Optional[Level] = None) -> None:
    """
    Clear data from the vector store.

    Args:
        level: If specified, only clear that level ("l0" or "l2").
               If None, clear both levels.
    """
    if level is None or level == "l0":
        logger.info("Clearing L0 vector store")
        self.l0_index = self._build_faiss_index()
        self.l0_documents = []

    if level is None or level == "l2":
        logger.info("Clearing L2 vector store")
        self.l2_index = self._build_faiss_index()
        self.l2_documents = []

    logger.info("Vector store cleared")

RegexNodeIndex ¶

RegexNodeIndex(code_graph: CodeGraph)

In-memory regex-based index for CodeGraph nodes. Supports regex pattern matching on node content with glob filtering.

Parameters:

Name	Type	Description	Default
`code_graph`	`CodeGraph`	CodeGraph instance containing nodes to index	required

Methods:

Name	Description
`search`	Search for pattern in node content (grep-like functionality).

Source code in codenib/index/regex_idx/regex_idx.py

def __init__(self, code_graph: CodeGraph):
    """
    Initialize RegexNodeIndex and build index from CodeGraph.

    Args:
        code_graph: CodeGraph instance containing nodes to index
    """
    self.code_graph = code_graph
    self.nodes: List[NodeInfo] = []
    self._build_index()

    logger.info(f"RegexNodeIndex initialized with {len(self.nodes)} nodes")

search ¶

search(
    pattern: str,
    file_glob: str | None = None,
    node_type: str | None = None,
    case_sensitive: bool = False,
    use_regex: bool = True,
    top_k: int | None = None,
) -> list[NodeInfo]

Search for pattern in node content (grep-like functionality).

Parameters:

Name	Type	Description	Default
`pattern`	`str`	Search pattern (regex or plain string)	required
`file_glob`	`str \| None`	Optional glob to filter by file path (e.g., '.py', '*/calc.py')	`None`
`node_type`	`str \| None`	Optional node type to filter (e.g., 'function', 'class', 'file')	`None`
`case_sensitive`	`bool`	Whether search is case-sensitive (default: False)	`False`
`use_regex`	`bool`	Whether to use regex (default: True) or plain string matching	`True`
`top_k`	`int \| None`	Optional result limit. Regex searches stop as soon as it is met.	`None`

Returns:

Type	Description
`list[NodeInfo]`	List of NodeInfo objects matching the pattern

Examples:

>>> idx.search(r'def\s+\w+', file_glob='*.py')  # Find function defs
>>> idx.search('calculator', use_regex=False)  # Plain string search
>>> idx.search('class', node_type='file')  # Search in file nodes only

Source code in codenib/index/regex_idx/regex_idx.py

def search(
    self,
    pattern: str,
    file_glob: Optional[str] = None,
    node_type: Optional[str] = None,
    case_sensitive: bool = False,
    use_regex: bool = True,
    top_k: Optional[int] = None,
) -> List[NodeInfo]:
    r"""
    Search for pattern in node content (grep-like functionality).

    Args:
        pattern: Search pattern (regex or plain string)
        file_glob: Optional glob to filter by file path (e.g., '*.py', '**/calc.py')
        node_type: Optional node type to filter (e.g., 'function', 'class', 'file')
        case_sensitive: Whether search is case-sensitive (default: False)
        use_regex: Whether to use regex (default: True) or plain string matching
        top_k: Optional result limit. Regex searches stop as soon as it is met.

    Returns:
        List of NodeInfo objects matching the pattern

    Examples:
        >>> idx.search(r'def\s+\w+', file_glob='*.py')  # Find function defs
        >>> idx.search('calculator', use_regex=False)  # Plain string search
        >>> idx.search('class', node_type='file')  # Search in file nodes only
    """
    if top_k is not None and (
        isinstance(top_k, bool) or not isinstance(top_k, int) or top_k < 1
    ):
        raise ValueError("top_k must be a positive integer")

    if use_regex:
        if len(pattern) > MAX_REGEX_PATTERN_CHARS:
            raise ValueError(
                "Regex pattern exceeds "
                f"the {MAX_REGEX_PATTERN_CHARS}-character limit"
            )

        # Start the request-wide clock before compiling or evaluating either
        # user-controlled structural filter. A single filtered-out scan is
        # still work charged to this request.
        deadline = time.monotonic() + REGEX_SEARCH_TIMEOUT_SECONDS
        flags = 0 if case_sensitive else regex.IGNORECASE
        try:
            compiled = regex.compile(pattern, flags)
        except regex.error as exc:
            logger.error("Invalid regex pattern %r: %s", pattern, exc)
            raise ValueError(f"Invalid regex pattern: {exc}") from exc
        _remaining_time(deadline)

        matches: List[NodeInfo] = []
        scanned_nodes = 0
        candidate_nodes = 0
        for node in self.nodes:
            _remaining_time(deadline)
            scanned_nodes += 1
            if scanned_nodes > MAX_REGEX_SCANNED_NODES:
                raise RegexSearchBudgetError(
                    "Regex search exceeded the "
                    f"{MAX_REGEX_SCANNED_NODES}-node scan budget; "
                    "use file_glob or node_type to narrow the search"
                )

            if file_glob:
                path_matches = bool(node.file) and fnmatch(node.file, file_glob)
                _remaining_time(deadline)
                if not path_matches:
                    continue
            if node_type and node.type != node_type:
                continue
            if not node.content:
                continue

            candidate_nodes += 1
            if candidate_nodes > MAX_REGEX_CANDIDATES:
                raise RegexSearchBudgetError(
                    "Regex search exceeded the "
                    f"{MAX_REGEX_CANDIDATES}-candidate match budget; "
                    "use file_glob or node_type to narrow the search"
                )

            remaining = _remaining_time(deadline)
            try:
                matched = compiled.search(node.content, timeout=remaining)
            except TimeoutError as exc:
                raise _timeout_error() from exc
            _remaining_time(deadline)
            if matched:
                matches.append(node)
                if top_k is not None and len(matches) >= top_k:
                    break

        logger.debug(
            "Regex search pattern=%r file_glob=%r node_type=%r: "
            "found %d matches after scanning %d nodes and %d candidates",
            pattern,
            file_glob,
            node_type,
            len(matches),
            scanned_nodes,
            candidate_nodes,
        )
        return matches

    # Preserve the unrestricted plain-string path for direct index users.
    candidates = self.nodes
    if file_glob:
        candidates = [
            node
            for node in candidates
            if node.file and fnmatch(node.file, file_glob)
        ]
    if node_type:
        candidates = [node for node in candidates if node.type == node_type]

    if case_sensitive:
        matches = [
            node for node in candidates if node.content and pattern in node.content
        ]
    else:
        pattern_lower = pattern.lower()
        matches = [
            node
            for node in candidates
            if node.content and pattern_lower in node.content.lower()
        ]
    if top_k is not None:
        matches = matches[:top_k]

    logger.debug(
        "Plain search pattern=%r file_glob=%r node_type=%r: "
        "found %d matches from %d candidates",
        pattern,
        file_glob,
        node_type,
        len(matches),
        len(candidates),
    )

    return matches

BM25CodeIndexer ¶

BM25CodeIndexer(
    code_graph=None,
    chunks=None,
    max_k: int = 15,
    language: str = "english",
    project_root: str | None = None,
    *,
    prepare_only: bool = False,
    check_cancelled: Callable[[], None] | None = None
)

A class that builds a BM25 index from CodeGraph nodes and provides search functionality with stemming support.

Parameters:

Name	Type	Description	Default
`code_graph`		CodeGraph instance containing nodes to index. If provided, the index will be built immediately.	`None`
`chunks`		List of CodeChunk objects to index. If provided, the index will be built immediately.	`None`
`max_k`	`int`	Maximum number of results to return in searches	`15`
`language`	`str`	Language for stopword removal Default is "english" which works well for processing code tokens as it treats special characters as separators	`'english'`
`project_root`	`str \| None`	Repository root used to resolve relative source paths	`None`
`prepare_only`	`bool`	Prepare canonical persisted documents without building the serving-time in-memory rank index.	`False`
`check_cancelled`	`Callable[[], None] \| None`	Cooperative stop check for prepare-only chunk builds.	`None`

Methods:

Name	Description
`build_index_from_graph`	Build a BM25 index from a CodeGraph.
`build_index_from_chunks`	Build a BM25 index from a list of CodeChunk objects.
`bind_repository_source`	Attach a retained source authority chosen by the runtime caller.
`save_index`	Save the index to a directory.
`load_index`	Load the index from a directory.
`load_index_values`	Load already-decoded BM25 values from an authenticated reader.

Source code in codenib/index/sparse_idx/bm25_index.py

def __init__(
    self,
    code_graph=None,
    chunks=None,
    max_k: int = 15,
    language: str = "english",
    project_root: Optional[str] = None,
    *,
    prepare_only: bool = False,
    check_cancelled: Callable[[], None] | None = None,
):
    """
    Initialize the BM25CodeIndexer and optionally build the index immediately.

    Args:
        code_graph: CodeGraph instance containing nodes to index. If provided,
                   the index will be built immediately.
        chunks: List of CodeChunk objects to index. If provided, the index will
               be built immediately.
        max_k: Maximum number of results to return in searches
        language: Language for stopword removal
                  Default is "english" which works well for processing code tokens
                  as it treats special characters as separators
        project_root: Repository root used to resolve relative source paths
        prepare_only: Prepare canonical persisted documents without building
                      the serving-time in-memory rank index.
        check_cancelled: Cooperative stop check for prepare-only chunk builds.
    """
    self.max_k = max_k
    self.language = language
    self.documents = []
    if type(prepare_only) is not bool:
        raise TypeError("BM25 prepare-only policy must be an exact boolean")
    if check_cancelled is not None and not callable(check_cancelled):
        raise TypeError("BM25 build cancellation must be callable")
    if code_graph is not None and (prepare_only or check_cancelled is not None):
        raise ValueError(
            "BM25 prepare-only mode and cancellation require source chunks "
            "without a code graph"
        )
    if prepare_only and chunks is None:
        raise ValueError("BM25 prepare-only mode requires source chunks")
    if check_cancelled is not None and chunks is None:
        raise ValueError("BM25 build cancellation requires source chunks")
    if check_cancelled is not None and not prepare_only:
        raise ValueError("BM25 build cancellation requires prepare-only mode")

    self.retriever = None
    self._documents_prepared = False
    self.code_graph: CodeGraph = None
    self.project_root = project_root
    self.source_mode = SOURCE_MODE_LEGACY_DIRECT
    self._source_binding: RepositorySourceBinding | None = None
    self.nodes: List[str] = []

    # Build the index immediately if a code_graph is provided
    if code_graph is not None:
        self.build_index_from_graph(code_graph)
    elif chunks is not None:
        if prepare_only:
            self._prepare_documents_from_chunks(
                chunks,
                project_root=project_root,
                check_cancelled=check_cancelled,
            )
            self._documents_prepared = True
        else:
            # Preserve the established virtual-call shape for subclasses
            # that override this public method.
            self.build_index_from_chunks(chunks, project_root=project_root)

build_index_from_graph ¶

build_index_from_graph(code_graph: CodeGraph) -> BM25Retriever

Build a BM25 index from a CodeGraph.

Parameters:

Name	Type	Description	Default
`code_graph`	`CodeGraph`	CodeGraph instance containing nodes to index	required

Source code in codenib/index/sparse_idx/bm25_index.py

def build_index_from_graph(self, code_graph: CodeGraph) -> BM25Retriever:
    """
    Build a BM25 index from a CodeGraph.

    Args:
        code_graph: CodeGraph instance containing nodes to index
    """
    # Reset the index
    self.documents = []
    self.nodes = []
    self.code_graph = code_graph
    self.project_root = code_graph.project_root
    self.source_mode = SOURCE_MODE_LEGACY_DIRECT
    self._source_binding = None
    self._documents_prepared = False
    self.retriever = None

    # Convert graph nodes to documents
    for vertex in code_graph.graph.vs:
        doc = self._convert_vertex_to_document(vertex)
        if doc is not None:
            self.documents.append(doc)
            node_name = doc.metadata.get("node_id") or doc.metadata.get("name")
            if node_name:
                self.nodes.append(node_name)

    # Create BM25Retriever with LangChain format
    self.retriever = BM25Retriever.from_documents(self.documents, k=self.max_k)
    self._documents_prepared = True

    return self.retriever

build_index_from_chunks ¶

build_index_from_chunks(
    chunks: list[CodeChunk], *, project_root: str | None = None
) -> BM25Retriever

Build a BM25 index from a list of CodeChunk objects.

Parameters:

Name	Type	Description	Default
`chunks`	`list[CodeChunk]`	List of CodeChunk objects (with node_id, chunk_type, name, file, etc.)	required
`project_root`	`str \| None`	Repository root used to resolve relative source paths	`None`

Returns:

Type	Description
`BM25Retriever`	BM25Retriever instance

Source code in codenib/index/sparse_idx/bm25_index.py

def build_index_from_chunks(
    self,
    chunks: List[CodeChunk],
    *,
    project_root: Optional[str] = None,
) -> BM25Retriever:
    """
    Build a BM25 index from a list of CodeChunk objects.

    Args:
        chunks: List of CodeChunk objects (with node_id, chunk_type, name, file, etc.)
        project_root: Repository root used to resolve relative source paths

    Returns:
        BM25Retriever instance
    """
    self._prepare_documents_from_chunks(
        chunks,
        project_root=project_root,
        check_cancelled=None,
    )
    self.retriever = BM25Retriever.from_documents(self.documents, k=self.max_k)
    self._documents_prepared = True
    return self.retriever

bind_repository_source ¶

bind_repository_source(binding: RepositorySourceBinding) -> None

Attach a retained source authority chosen by the runtime caller.

Source code in codenib/index/sparse_idx/bm25_index.py

def bind_repository_source(self, binding: RepositorySourceBinding) -> None:
    """Attach a retained source authority chosen by the runtime caller."""

    source_identity = binding.authenticated_identity_snapshot()
    self._source_binding = binding
    self.project_root = str(source_identity.root)
    self.source_mode = SOURCE_MODE_BOUND_REPOSITORY

save_index ¶

save_index(directory_path: str, *, check_cancelled: Callable[[], None] | None = None)

Save the index to a directory.

Parameters:

Name	Type	Description	Default
`directory_path`	`str`	Path to save the index to	required

Source code in codenib/index/sparse_idx/bm25_index.py

def save_index(
    self,
    directory_path: str,
    *,
    check_cancelled: Callable[[], None] | None = None,
):
    """
    Save the index to a directory.

    Args:
        directory_path: Path to save the index to
    """
    if check_cancelled is not None and not callable(check_cancelled):
        raise TypeError("BM25 persistence cancellation must be callable")
    if self.retriever is None and not self._documents_prepared:
        raise ValueError(
            "Index has not been built. Call build_index_from_graph first."
        )
    if check_cancelled is not None:
        check_cancelled()

    # Create directory if it doesn't exist
    os.makedirs(directory_path, exist_ok=True)

    # Save documents as JSON since LangChain BM25Retriever doesn't have persist method
    documents_data = []
    for doc in self.documents:
        if check_cancelled is not None:
            check_cancelled()
        documents_data.append(
            {"page_content": doc.page_content, "metadata": doc.metadata}
        )

    documents_file = os.path.join(directory_path, "documents.json")
    _write_json_interruptibly(
        documents_file,
        documents_data,
        check_cancelled,
    )

    # Save additional metadata including project_root
    metadata = {
        "project_root": (
            str(self.project_root) if self.project_root is not None else None
        ),
        "max_k": self.max_k,
        "language": self.language,
    }
    metadata_file = os.path.join(directory_path, "bm25_metadata.json")
    _write_json_interruptibly(
        metadata_file,
        metadata,
        check_cancelled,
    )

load_index ¶

load_index(directory_path: str)

Load the index from a directory.

Parameters:

Name	Type	Description	Default
`directory_path`	`str`	Path to load the index from	required

Source code in codenib/index/sparse_idx/bm25_index.py

def load_index(self, directory_path: str):
    """
    Load the index from a directory.

    Args:
        directory_path: Path to load the index from
    """
    if not os.path.exists(directory_path):
        raise ValueError(f"Directory {directory_path} does not exist.")

    # Load documents from JSON
    documents_file = os.path.join(directory_path, "documents.json")
    if not os.path.exists(documents_file):
        raise ValueError(f"Documents file not found: {documents_file}")

    with open(documents_file, "r", encoding="utf-8") as f:
        documents_data = json.load(f)

    metadata_file = os.path.join(directory_path, "bm25_metadata.json")
    if os.path.exists(metadata_file):
        with open(metadata_file, "r", encoding="utf-8") as f:
            metadata = json.load(f)
    else:
        metadata = {
            "project_root": None,
            "max_k": self.max_k,
            "language": self.language,
        }
    self.load_index_values(documents_data, metadata)

load_index_values ¶

load_index_values(
    documents_data: Iterable[object],
    metadata: Mapping[str, Any] | None,
    *,
    source_mode: str = SOURCE_MODE_LEGACY_DIRECT
) -> None

Load already-decoded BM25 values from an authenticated reader.

Artifact runtimes use this entry point so persisted bytes can be read through a pinned directory authority instead of being reopened by path.

Source code in codenib/index/sparse_idx/bm25_index.py

def load_index_values(
    self,
    documents_data: Iterable[object],
    metadata: Mapping[str, Any] | None,
    *,
    source_mode: str = SOURCE_MODE_LEGACY_DIRECT,
) -> None:
    """Load already-decoded BM25 values from an authenticated reader.

    Artifact runtimes use this entry point so persisted bytes can be read
    through a pinned directory authority instead of being reopened by path.
    """

    if metadata is not None and not isinstance(metadata, Mapping):
        raise ValueError("BM25 metadata must be an object")
    metadata = {} if metadata is None else metadata
    max_k = metadata.get("max_k", 10)
    language = metadata.get("language", "english")
    project_root = metadata.get("project_root")
    if isinstance(max_k, bool) or not isinstance(max_k, int) or max_k <= 0:
        raise ValueError("BM25 max_k must be a positive integer")
    if not isinstance(language, str) or not language:
        raise ValueError("BM25 language must be a non-empty string")
    if project_root is not None and not isinstance(project_root, str):
        raise ValueError("BM25 project_root must be a string or null")
    if source_mode not in _SOURCE_MODES:
        raise ValueError(f"unsupported BM25 source mode: {source_mode!r}")

    # Restore serving configuration before creating the retriever.  Its k
    # value is fixed at construction time.
    self.project_root = project_root
    self.source_mode = source_mode
    self._source_binding = None
    self._documents_prepared = False
    self.max_k = max_k
    self.language = language

    # Reconstruct Document objects one decoded element at a time.
    self.documents = []
    self.nodes = []
    seen_nodes = set()
    for doc_data in documents_data:
        if not isinstance(doc_data, Mapping):
            raise ValueError("BM25 document must be an object")
        page_content = doc_data.get("page_content")
        document_metadata = doc_data.get("metadata")
        if not isinstance(page_content, str) or not isinstance(
            document_metadata, Mapping
        ):
            raise ValueError("BM25 document content or metadata is invalid")
        doc = Document(
            page_content=page_content,
            metadata=dict(document_metadata),
        )
        self.documents.append(doc)
        node_name = doc.metadata.get("node_id") or doc.metadata.get("name")
        if node_name and node_name not in seen_nodes:
            self.nodes.append(node_name)
            seen_nodes.add(node_name)

    self.retriever = BM25Retriever.from_documents(self.documents, k=self.max_k)
    self._documents_prepared = True

create_code_vector_store ¶

create_code_vector_store(
    embedding_model: str = "text-embedding-ada-002",
    embedding_provider: str = "openai",
    store_path: str | None = None,
    **kwargs
) -> CodeVectorStore

Factory function to create a CodeVectorStore.

Parameters:

Name	Type	Description	Default
`embedding_model`	`str`	Name of the embedding model	`'text-embedding-ada-002'`
`embedding_provider`	`str`	Provider for embeddings	`'openai'`
`store_path`	`str \| None`	Path to store/load the vector store	`None`
`**kwargs`		Additional arguments for CodeVectorStore	`{}`

Returns:

Type	Description
`CodeVectorStore`	CodeVectorStore instance

Source code in codenib/index/embedding/vector_store.py

def create_code_vector_store(
    embedding_model: str = "text-embedding-ada-002",
    embedding_provider: str = "openai",
    store_path: Optional[str] = None,
    **kwargs,
) -> CodeVectorStore:
    """
    Factory function to create a CodeVectorStore.

    Args:
        embedding_model: Name of the embedding model
        embedding_provider: Provider for embeddings
        store_path: Path to store/load the vector store
        **kwargs: Additional arguments for CodeVectorStore

    Returns:
        CodeVectorStore instance
    """
    return CodeVectorStore(
        embedding_model=embedding_model,
        embedding_provider=embedding_provider,
        store_path=store_path,
        **kwargs,
    )

codenib.index ¶

CodeVectorStore ¶

closed property ¶

reuse_query_embedding ¶

clear_query_cache ¶

swap_index ¶

close ¶

add_code_chunks ¶

add_nodes_with_content ¶

search ¶

search_with_content ¶

search_within_ids ¶

hierarchical_search ¶

save ¶

load ¶

get_stats ¶

get_embeddings_by_content_hash ¶

rebuild_from_embeddings ¶

delta_update ¶

clear ¶

RegexNodeIndex ¶

search ¶

BM25CodeIndexer ¶

build_index_from_graph ¶

build_index_from_chunks ¶

bind_repository_source ¶

save_index ¶

load_index ¶

load_index_values ¶

create_code_vector_store ¶

closed `property` ¶