import argparse import sys from pathlib import Path from bot.memory.document_registry import DocumentRegistry from bot.memory.embedder import Embedder from bot.memory.vector_database.chroma import Chroma from bot.memory.vector_database.id_generator import generate_id from database import create_db_engine from document_loader.format import Format from document_loader.loader import DirectoryLoader from document_loader.text_splitter import create_recursive_text_splitter from entities.document import Document from helpers.log import get_logger from sqlmodel import Session logger = get_logger(__name__) def load_documents(docs_path: Path) -> list[Document]: """ Loads Markdown documents from the specified path. Args: docs_path (Path): The path to the documents. Returns: List[Document]: A list of loaded documents. """ loader = DirectoryLoader( path=docs_path, glob="**/*.md", show_progress=True, ) return loader.load() def split_chunks(sources: list, chunk_size: int = 1000, chunk_overlap: int = 50) -> list: """ Splits a list of sources into smaller chunks. Args: sources (List): The list of sources to be split into chunks. chunk_size (int, optional): The maximum size of each chunk. Defaults to 1000. chunk_overlap (int, optional): The amount of overlap between consecutive chunks. Defaults to 50. Returns: List: A list of smaller chunks obtained from the input sources. """ chunks = [] splitter = create_recursive_text_splitter( format=Format.MARKDOWN.value, chunk_size=chunk_size, chunk_overlap=chunk_overlap ) for chunk in splitter.split_documents(sources): chunks.append(chunk) return chunks def build_memory_index( docs_path: Path, vector_store_path: str, chunk_size: int, chunk_overlap: int, full_rebuild: bool = False, model_name: str = "jinaai/jina-embeddings-v5-text-small-retrieval", ) -> dict: """ Build or incrementally update the vector memory index. When `full_rebuild` is `True` the Chroma collection and the registry are wiped before ingesting everything from scratch. Otherwise, only new and changed documents are processed, and deleted documents are removed. Returns a stats dict with counts of processed / deleted / skipped docs. """ # ------------------------------------------------------------------ # bootstrap vector DB + registry # ------------------------------------------------------------------ embedding = Embedder(model_name=model_name) vector_database = Chroma(is_persistent=True, persist_directory=str(vector_store_path), embedding=embedding) session = Session(create_db_engine()) registry = DocumentRegistry(session) # ------------------------------------------------------------------ # full-rebuild: wipe everything first # ------------------------------------------------------------------ if full_rebuild: logger.info("Full rebuild requested – wiping collection and registry.") vector_database.reset_collection() for record in registry.get_all(): registry.remove(record.document_id) # ------------------------------------------------------------------ # (a) load source docs, compute document_id + version_hash # ------------------------------------------------------------------ logger.info(f"Loading documents from: {docs_path}") sources = load_documents(docs_path) logger.info(f"Number of loaded documents: {len(sources)}") current_docs: dict[str, str] = {} # {document_id: version_hash} doc_map: dict[str, Document] = {} # {document_id: Document} for doc in sources: source_path = doc.metadata.get("source", "") doc_id = generate_id(source_path) ver_hash = generate_id(doc.page_content) current_docs[doc_id] = ver_hash doc_map[doc_id] = doc # ------------------------------------------------------------------ # (b) diff against registry # ------------------------------------------------------------------ new_ids, changed_ids, deleted_ids = registry.get_stale_documents(current_docs) logger.info( f"Diff result – new: {len(new_ids)}, changed: {len(changed_ids)}, deleted: {len(deleted_ids)}, " f"unchanged: {len(current_docs) - len(new_ids) - len(changed_ids)}", ) # ------------------------------------------------------------------ # (c) remove changed / deleted docs from Chroma + registry # ------------------------------------------------------------------ for doc_id in changed_ids | deleted_ids: record = registry.get(doc_id) chunk_ids = record.chunk_ids if record else None vector_database.delete_chunks_by_document_id(doc_id, chunk_ids=chunk_ids) if doc_id in deleted_ids: registry.remove(doc_id) # ------------------------------------------------------------------ # (d) chunk & ingest new + changed docs # ------------------------------------------------------------------ to_ingest = new_ids | changed_ids for doc_id in to_ingest: doc = doc_map[doc_id] source_path = doc.metadata.get("source", "") ver_hash = current_docs[doc_id] chunks = split_chunks([doc], chunk_size=chunk_size, chunk_overlap=chunk_overlap) # inject document_id + version_hash into every chunk's metadata for chunk in chunks: chunk.metadata["document_id"] = doc_id chunk.metadata["version_hash"] = ver_hash chunk_ids = vector_database.from_chunks(chunks) # ------------------------------------------------------------------ # (e) upsert into registry # ------------------------------------------------------------------ registry.upsert( doc_id, source=source_path, filename=Path(source_path).name, size=len(doc.page_content), content_type="text/markdown", version_hash=ver_hash, chunk_ids=chunk_ids, ) stats = { "processed": len(to_ingest), "deleted": len(deleted_ids), "skipped": len(current_docs) - len(to_ingest), } logger.info("Memory Index updated – %s", stats) session.close() return stats def get_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Memory Builder") parser.add_argument( "--chunk-size", type=int, help="The maximum size of each chunk. Defaults to 1000.", required=False, default=1000, ) parser.add_argument( "--chunk-overlap", type=int, help="The amount of overlap between consecutive chunks. Defaults to 50.", required=False, default=50, ) parser.add_argument( "--full-rebuild", action="store_true", default=False, help="Wipe the vector store and registry and rebuild from scratch.", ) parser.add_argument( "--model-name", type=str, help="The name of the SentenceTransformer model to use for embedding. " "Defaults to 'jinaai/jina-embeddings-v5-text-small-retrieval'.", required=False, default="jinaai/jina-embeddings-v5-text-small-retrieval", ) return parser.parse_args() def main(parameters): root_folder = Path(__file__).resolve().parent.parent doc_path = root_folder / "docs" vector_store_path = root_folder / "vector_store" / "docs_index" build_memory_index( docs_path=doc_path, vector_store_path=str(vector_store_path), chunk_size=parameters.chunk_size, chunk_overlap=parameters.chunk_overlap, full_rebuild=parameters.full_rebuild, model_name=parameters.model_name, ) if __name__ == "__main__": try: args = get_args() main(args) except Exception as error: logger.error(f"An error occurred: {str(error)}", exc_info=True, stack_info=True) sys.exit(1)