codenib.index.embedding

Embedding module for CodeNib. Provides vector storage and semantic search capabilities for code chunks.

Modules:

Name	Description
builders	Reusable embedding builders for hierarchical pipelines.
model_policy	Embedding model defaults and remote-code trust policy.
prompt_registry	Per-model prompt registry for sentence-transformer embedders.
vector_store	Vector Store implementation using FAISS and sentence-transformers for code embeddings.

Classes:

Name	Description
VectorStoreBuilder	Builder class wrapping common vector store build operations.
CodeVectorStore	Vector store for code embeddings using FAISS and sentence-transformers.

Functions:

Name	Description
build_hierarchical_vector_store	Build (or load) a hierarchical vector store (L0/L2) for a repository.
create_code_vector_store	Factory function to create a CodeVectorStore.

VectorStoreBuilder

VectorStoreBuilder(profiler: Profiler | None = None)

Builder class wrapping common vector store build operations.

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

def __init__(self, profiler: Optional[Profiler] = None) -> None:
    self.profiler = profiler

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,
    **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", "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
**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 the FAISS index in place when the change set is small.
clear	Clear data from the vector store.

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,
    **embedding_kwargs,
):
    """
    Initialize the CodeVectorStore.

    Args:
        embedding_model: Name of the embedding model to use
        embedding_provider: Provider for embeddings ("openai", "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.
        **embedding_kwargs: Additional arguments for embedding model
    """
    self.embedding_model = embedding_model
    self.embedding_provider = embedding_provider
    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))
    self.store_path = Path(store_path) if store_path else None
    self.profiler = profiler

    # 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)
    self.l0_index = self._build_faiss_index()
    self.l0_documents: List[_Document] = []

    # Initialize L2 (function/method-level) - default
    self.l2_index = self._build_faiss_index()
    self.l2_documents: List[_Document] = []

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

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) -> None

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

Frees the current L0/L2 FAISS indices and documents from memory, then loads a new index from path. 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) -> None:
    """Hot-swap the FAISS index without reloading the embedding model.

    Frees the current L0/L2 FAISS indices and documents from memory, then
    loads a new index from *path*.  The embedding model is left intact so
    the caller can reuse the same model across many instances.
    """
    # Free current FAISS index memory (GPU or CPU).
    for index in (self.l0_index, self.l2_index):
        if index is None:
            continue
        reset = getattr(index, "reset", None)
        if callable(reset):
            reset()

    # Reinitialise to empty indices (guards against a partially-failed
    # subsequent load leaving the store in a mixed state).
    self.l0_index = self._build_faiss_index()
    self.l0_documents = []
    self.l2_index = self._build_faiss_index()
    self.l2_documents = []

    self.store_path = Path(path)
    self.load(path)

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."""
    for index in (self.l0_index, self.l2_index):
        if index is None:
            continue
        reset = getattr(index, "reset", None)
        if callable(reset):
            reset()

    self.l0_documents.clear()
    self.l2_documents.clear()
    self.l0_index = None
    self.l2_index = None

    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=metadata.get("file", ""),
            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=metadata.get("file", ""),
            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=metadata.get("file", ""),
                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("/", "__")

    # Save L0
    if self.l0_index is not None and self.l0_documents:
        self._save_level(save_path, "l0", model_suffix)

    # Save L2
    if self.l2_index is not None and self.l2_documents:
        self._save_level(save_path, "l2", model_suffix)

    # Save top-level configuration
    config_path = save_path / f"config_{model_suffix}.json"
    config = {
        "embedding_model": self.embedding_model,
        "embedding_provider": self.embedding_provider,
        "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),
    }
    with open(config_path, "w") as f:
        json.dump(config, f, indent=2)

    logger.info("Vector store saved successfully")

load

load(path: str | 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

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

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

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

    load_path = Path(load_path)
    if not load_path.exists():
        raise FileNotFoundError(f"Vector store not found at {load_path}")

    logger.info(f"Loading vector store from {load_path}")

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

    # Load top-level configuration
    config_path = load_path / f"config_{model_suffix}.json"
    if not config_path.exists():
        config_path = load_path / "config.json"

    if config_path.exists():
        with open(config_path, "r") as f:
            config = json.load(f)

        if config.get("dimension") != self.dimension:
            logger.warning(
                f"Dimension mismatch: expected {self.dimension}, "
                f"got {config.get('dimension')}"
            )
        saved_metric = config.get("index_metric")
        if saved_metric and saved_metric != self.index_metric:
            logger.warning(
                "Index metric mismatch: expected %s, got %s",
                self.index_metric,
                saved_metric,
            )
            self.index_metric = saved_metric

    # Load L0
    l0_path = load_path / "l0"
    if l0_path.exists():
        loaded = self._load_level(l0_path, model_suffix)
        if loaded is not None:
            self.l0_index, self.l0_documents = loaded
            logger.info(f"Loaded L0 store with {len(self.l0_documents)} documents")

    # Load L2
    l2_path = load_path / "l2"
    if l2_path.exists():
        loaded = self._load_level(l2_path, model_suffix)
        if loaded is not None:
            self.l2_index, self.l2_documents = loaded
            logger.info(f"Loaded L2 store with {len(self.l2_documents)} documents")

    total_docs = len(self.l0_documents) + len(self.l2_documents)
    logger.info(
        f"Vector store loaded successfully with {total_docs} total documents "
        f"(L0: {len(self.l0_documents)}, L2: {len(self.l2_documents)})"
    )

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,
        "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 the 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. If the change ratio exceeds the threshold (or the index is empty), it falls back to a full rebuild via :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 the 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.  If
    the change ratio exceeds the threshold (or the index is empty), it
    falls back to a full rebuild via :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 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")

build_hierarchical_vector_store

build_hierarchical_vector_store(
    *,
    repo_path: str,
    index_path: str,
    plan_name: str | None = None,
    languages: list[str] | None = None,
    max_lines_per_chunk: int | None = None,
    build_levels: list[str] | None = None,
    embedding_model: str,
    embedding_provider: str,
    embedding_dimension: int | None,
    embedding_kwargs: dict[str, object] | None = None,
    embedding: object | None = None,
    index_metric: str = "ip",
    index_type: str = "flat",
    ivf_nlist: int = 100,
    ivf_nprobe: int = 8,
    profiler: Profiler | None = None,
    force_rebuild: bool = False
) -> CodeVectorStore

Build (or load) a hierarchical vector store (L0/L2) for a repository.

When force_rebuild=False (the default) and a saved index for the requested embedding_model already exists at index_path, the store is loaded from disk instead of re-chunking and re-embedding the repository. Pass force_rebuild=True to unconditionally rebuild.

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

def build_hierarchical_vector_store(
    *,
    repo_path: str,
    index_path: str,
    plan_name: Optional[str] = None,
    languages: Optional[List[str]] = None,
    max_lines_per_chunk: Optional[int] = None,
    build_levels: Optional[List[str]] = None,
    embedding_model: str,
    embedding_provider: str,
    embedding_dimension: Optional[int],
    embedding_kwargs: Optional[Dict[str, object]] = None,
    embedding: Optional[object] = None,
    index_metric: str = "ip",
    index_type: str = "flat",
    ivf_nlist: int = 100,
    ivf_nprobe: int = 8,
    profiler: Optional[Profiler] = None,
    force_rebuild: bool = False,
) -> CodeVectorStore:
    """Build (or load) a hierarchical vector store (L0/L2) for a repository.

    When ``force_rebuild=False`` (the default) and a saved index for the
    requested ``embedding_model`` already exists at ``index_path``, the store
    is loaded from disk instead of re-chunking and re-embedding the repository.
    Pass ``force_rebuild=True`` to unconditionally rebuild.
    """
    store_path = Path(index_path)
    if plan_name:
        store_path = store_path / plan_name

    if not force_rebuild:
        model_suffix = embedding_model.replace("/", "__")
        config_path = store_path / f"config_{model_suffix}.json"
        if config_path.exists():
            logger.info(
                "Pre-built index found at %s — loading instead of rebuilding "
                "(pass force_rebuild=True to override).",
                store_path,
            )
            vector_store = CodeVectorStore(
                embedding_model=embedding_model,
                embedding_provider=embedding_provider,
                dimension=embedding_dimension,
                index_metric=index_metric,
                index_type=index_type,
                ivf_nlist=ivf_nlist,
                ivf_nprobe=ivf_nprobe,
                store_path=str(store_path),
                profiler=profiler,
                embedding=embedding,
                **(embedding_kwargs or {}),
            )
            vector_store.load(str(store_path))
            return vector_store

    languages = languages or ["python"]
    build_levels = [level.lower() for level in (build_levels or ["l0", "l2"])]
    repo_cfg = RepoChunkingConfig(languages=languages)

    chunks_by_level = {}
    level_configs = {
        "l0": dict(
            chunker_kwargs=dict(
                language=languages[0],
                repo_config=repo_cfg,
                max_lines_per_chunk=None,
                chunk_depth=0,
                skeleton_mode=True,
            )
        ),
        "l2": dict(
            chunker_kwargs=dict(
                language=languages[0],
                repo_config=repo_cfg,
                max_lines_per_chunk=max_lines_per_chunk,
                chunk_depth=2,
                l2_level_exclusive=True,
                skeleton_mode=False,
            )
        ),
    }

    for level in build_levels:
        cfg = level_configs.get(level)
        if not cfg:
            continue
        chunker = CodeChunker(**cfg["chunker_kwargs"])
        with _profiler_section(
            profiler,
            f"chunking_{level}",
            {"level": level, "language": languages[0]},
        ):
            chunks_by_level[level] = chunker.chunk_repository(repo_path=repo_path)
        logger.info(
            f"Chunked {len(chunks_by_level[level])} {level} chunks "
            f"(lang={languages[0]})"
        )

    l0_chunks = chunks_by_level.get("l0", [])
    l2_chunks = chunks_by_level.get("l2", [])

    if not l0_chunks and not l2_chunks:
        raise ValueError("No code chunks generated from repository.")

    # Compute chunk / LOC statistics by counting actual file lines on disk.
    repo = Path(repo_path)
    unique_files = {chunk.file for chunk in (l0_chunks + l2_chunks)}
    loc_by_file: Dict[str, int] = {}
    for rel_path in unique_files:
        try:
            loc_by_file[rel_path] = len(
                (repo / rel_path).read_text(errors="replace").splitlines()
            )
        except OSError:
            pass
    total_loc = sum(loc_by_file.values())
    total_chunks = len(l0_chunks) + len(l2_chunks)
    chunk_stats = {
        "l0_chunks": len(l0_chunks),
        "l2_chunks": len(l2_chunks),
        "total_files": len(loc_by_file),
        "total_loc": total_loc,
        "avg_chunk_loc": round(total_loc / max(total_chunks, 1), 1),
    }

    store_path.mkdir(parents=True, exist_ok=True)
    vector_store = CodeVectorStore(
        embedding_model=embedding_model,
        embedding_provider=embedding_provider,
        dimension=embedding_dimension,
        index_metric=index_metric,
        index_type=index_type,
        ivf_nlist=ivf_nlist,
        ivf_nprobe=ivf_nprobe,
        store_path=str(store_path),
        profiler=profiler,
        embedding=embedding,
        **(embedding_kwargs or {}),
    )

    if l0_chunks:
        vector_store.add_code_chunks(
            [chunk._asdict() for chunk in l0_chunks], level="l0"
        )
    if l2_chunks:
        vector_store.add_code_chunks(
            [chunk._asdict() for chunk in l2_chunks], level="l2"
        )

    vector_store.save(str(store_path))
    vector_store.chunk_stats = chunk_stats
    return vector_store

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,
    )