scripts/build_rag_memory.py

#!/usr/bin/env python3
"""
RAG 메모리 구축 스크립트

사용 예시:
    python scripts/build_rag_memory.py \
        --source-dir rag_memory/source_docs \
        --chunks-path rag_memory/chunks.jsonl \
        --index-dir rag_memory/index \
        --embedding-model sentence-transformers/all-MiniLM-L6-v2 \
        --chunk-size 500 \
        --chunk-overlap 100 \
        --batch-size 16 \
        --rebuild
"""

from __future__ import annotations

import argparse
import datetime
import hashlib
import json
import logging
import os
import re
import shutil
import sys
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

try:
    import pdfplumber  # type: ignore
except ImportError:
    pdfplumber = None

try:
    import docx  # type: ignore
except ImportError:
    docx = None

try:
    from sentence_transformers import SentenceTransformer  # type: ignore
except ImportError:  # pragma: no cover
    SentenceTransformer = None  # type: ignore

try:
    import chromadb  # type: ignore
    from chromadb.utils import embedding_functions  # type: ignore
except ImportError:  # pragma: no cover
    chromadb = None  # type: ignore
    embedding_functions = None  # type: ignore


LOGGER = logging.getLogger("build_rag_memory")


@dataclass
class Document:
    source_path: Path
    text: str
    metadata: Dict[str, str]


@dataclass
class Chunk:
    id: str
    text: str
    source_path: str
    metadata: Dict[str, str]


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Build local RAG memory from source documents.")
    parser.add_argument("--source-dir", default="rag_memory/source_docs", help="Directory containing raw documents.")
    parser.add_argument("--chunks-path", default="rag_memory/chunks.jsonl", help="Output JSONL file for text chunks.")
    parser.add_argument("--index-dir", default="rag_memory/index", help="Persistent directory for vector index.")
    parser.add_argument("--embedding-model", default="sentence-transformers/all-MiniLM-L6-v2",
                        help="SentenceTransformer model name.")
    parser.add_argument("--chunk-size", type=int, default=500, help="Character length of each chunk.")
    parser.add_argument("--chunk-overlap", type=int, default=100, help="Character overlap between chunks.")
    parser.add_argument("--batch-size", type=int, default=16, help="Batch size for embedding generation.")
    parser.add_argument("--rebuild", action="store_true",
                        help="Rebuild chunks and index from scratch (existing files will be removed).")
    parser.add_argument("--collection-name", default="echolalia_rag", help="ChromaDB collection name.")
    return parser.parse_args(argv)


def ensure_dependencies() -> None:
    missing = []
    if SentenceTransformer is None:
        missing.append("sentence-transformers")
    if chromadb is None:
        missing.append("chromadb")
    if pdfplumber is None:
        LOGGER.warning("pdfplumber not installed: PDF extraction will be skipped.")
    if docx is None:
        LOGGER.warning("python-docx not installed: DOCX extraction will be skipped.")
    if missing:
        raise RuntimeError(
            "Missing required dependencies: {}. Please install them via `pip install {}`".format(
                ", ".join(missing), " ".join(missing)
            )
        )


def normalize_text(text: str) -> str:
    text = text.replace("\u200b", " ")
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def extract_text_from_pdf(path: Path) -> Optional[str]:
    if pdfplumber is None:
        return None
    try:
        with pdfplumber.open(str(path)) as pdf:
            pages = [page.extract_text() or "" for page in pdf.pages]
        return "\n".join(pages)
    except Exception as exc:  # pragma: no cover
        LOGGER.error("Failed to extract PDF %s: %s", path, exc)
        return None


def extract_text_from_docx(path: Path) -> Optional[str]:
    if docx is None:
        return None
    try:
        document = docx.Document(str(path))
        paragraphs = [para.text for para in document.paragraphs]
        return "\n".join(paragraphs)
    except Exception as exc:  # pragma: no cover
        LOGGER.error("Failed to extract DOCX %s: %s", path, exc)
        return None


def extract_text_from_txt(path: Path) -> Optional[str]:
    try:
        return path.read_text(encoding="utf-8")
    except UnicodeDecodeError:
        LOGGER.warning("UTF-8 decoding failed for %s, trying latin-1.", path)
        try:
            return path.read_text(encoding="latin-1")
        except Exception as exc:  # pragma: no cover
            LOGGER.error("Failed to read text file %s: %s", path, exc)
            return None


def load_documents(source_dir: Path) -> List[Document]:
    documents: List[Document] = []
    if not source_dir.exists():
        LOGGER.warning("Source directory %s does not exist. No documents to process.", source_dir)
        return documents

    for file_path in sorted(source_dir.rglob("*")):
        if not file_path.is_file():
            continue
        ext = file_path.suffix.lower()
        text: Optional[str] = None
        if ext == ".pdf":
            text = extract_text_from_pdf(file_path)
        elif ext == ".docx":
            text = extract_text_from_docx(file_path)
        elif ext in {".txt", ".md"}:
            text = extract_text_from_txt(file_path)
        else:
            LOGGER.info("Skipping unsupported file type: %s", file_path)
            continue

        if not text:
            LOGGER.warning("No text extracted from %s", file_path)
            continue

        text = normalize_text(text)
        metadata = {
            "filename": file_path.name,
            "extension": ext,
            "modified_at": datetime.datetime.fromtimestamp(file_path.stat().st_mtime).isoformat(),
            "filesize": str(file_path.stat().st_size),
        }
        documents.append(Document(source_path=file_path, text=text, metadata=metadata))
        LOGGER.info("Loaded document %s (%s chars)", file_path, len(text))

    return documents


def split_into_chunks(doc: Document, chunk_size: int, chunk_overlap: int) -> List[Chunk]:
    if chunk_size <= 0:
        raise ValueError("chunk_size must be positive.")
    if chunk_overlap < 0:
        raise ValueError("chunk_overlap must be non-negative.")
    if chunk_overlap >= chunk_size:
        raise ValueError("chunk_overlap must be smaller than chunk_size.")

    text = doc.text
    chunks: List[Chunk] = []
    start = 0
    doc_hash = hashlib.sha1(str(doc.source_path).encode("utf-8")).hexdigest()[:12]
    index = 0

    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunk_text = text[start:end].strip()
        if chunk_text:
            chunk_id = f"{doc_hash}-{index:04d}"
            chunk_meta = dict(doc.metadata)
            chunk_meta.update({
                "chunk_index": str(index),
                "source_path": str(doc.source_path),
            })
            chunks.append(
                Chunk(
                    id=chunk_id,
                    text=chunk_text,
                    source_path=str(doc.source_path),
                    metadata=chunk_meta,
                )
            )
            index += 1
        if end == len(text):
            break
        start = end - chunk_overlap

    return chunks


def serialize_chunks(chunks: Sequence[Chunk], path: Path, rebuild: bool) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    mode = "w" if rebuild else "a"
    if rebuild and path.exists():
        path.unlink()
    with path.open(mode, encoding="utf-8") as f:
        for chunk in chunks:
            record = {
                "id": chunk.id,
                "text": chunk.text,
                "source_path": chunk.source_path,
                "metadata": chunk.metadata,
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")


def load_existing_chunk_ids(path: Path) -> set:
    ids = set()
    if not path.exists():
        return ids
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            try:
                data = json.loads(line)
                if "id" in data:
                    ids.add(data["id"])
            except json.JSONDecodeError:
                continue
    return ids


def embed_chunks(chunks: Sequence[Chunk], model_name: str, batch_size: int) -> Tuple[List[str], List[str], List[Dict[str, str]], List[List[float]]]:
    if SentenceTransformer is None:
        raise RuntimeError("sentence-transformers is not installed.")

    texts = [chunk.text for chunk in chunks]
    ids = [chunk.id for chunk in chunks]
    metadatas = [chunk.metadata for chunk in chunks]

    LOGGER.info("Loading embedding model %s", model_name)
    model = SentenceTransformer(model_name)
    LOGGER.info("Embedding %d chunks (batch_size=%d)", len(texts), batch_size)
    embeddings = model.encode(texts, batch_size=batch_size, convert_to_numpy=True, show_progress_bar=True)
    return ids, texts, metadatas, embeddings.tolist()


def build_index(
    ids: Sequence[str],
    texts: Sequence[str],
    metadatas: Sequence[Dict[str, str]],
    embeddings: Sequence[Sequence[float]],
    index_dir: Path,
    collection_name: str,
    rebuild: bool,
) -> None:
    if chromadb is None:
        raise RuntimeError("chromadb is not installed.")

    index_dir.mkdir(parents=True, exist_ok=True)
    if rebuild and index_dir.exists():
        shutil.rmtree(index_dir)
        index_dir.mkdir(parents=True, exist_ok=True)

    client = chromadb.PersistentClient(path=str(index_dir))

    if rebuild:
        try:
            client.delete_collection(collection_name)
        except Exception:
            pass

    collection = client.get_or_create_collection(name=collection_name)

    if ids:
        LOGGER.info("Adding %d embeddings to collection %s", len(ids), collection_name)
        collection.upsert(
            ids=list(ids),
            documents=list(texts),
            metadatas=list(metadatas),
            embeddings=list(embeddings),
        )
    else:
        LOGGER.info("No embeddings to add.")


def process_documents(args: argparse.Namespace) -> None:
    ensure_dependencies()

    source_dir = Path(args.source_dir)
    chunks_path = Path(args.chunks_path)
    index_dir = Path(args.index_dir)

    documents = load_documents(source_dir)
    if not documents:
        LOGGER.warning("No documents processed. Exiting.")
        return

    existing_ids = set()
    if not args.rebuild:
        existing_ids = load_existing_chunk_ids(chunks_path)
        LOGGER.info("Loaded %d existing chunk ids.", len(existing_ids))

    new_chunks: List[Chunk] = []
    for doc in documents:
        doc_chunks = split_into_chunks(doc, args.chunk_size, args.chunk_overlap)
        for chunk in doc_chunks:
            if chunk.id in existing_ids:
                LOGGER.info("Skipping existing chunk id %s", chunk.id)
                continue
            new_chunks.append(chunk)

    if not new_chunks:
        LOGGER.info("No new chunks generated. Nothing to update.")
        return

    LOGGER.info("Generated %d new chunks.", len(new_chunks))
    serialize_chunks(new_chunks, chunks_path, args.rebuild)

    ids, texts, metadatas, embeddings = embed_chunks(new_chunks, args.embedding_model, args.batch_size)
    build_index(ids, texts, metadatas, embeddings, index_dir, args.collection_name, args.rebuild)
    LOGGER.info("RAG memory build completed successfully.")


def setup_logging() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def main(argv: Optional[Sequence[str]] = None) -> None:
    setup_logging()
    args = parse_args(argv)
    LOGGER.info("Starting RAG memory build with args: %s", args)
    try:
        process_documents(args)
    except Exception as exc:  # pragma: no cover
        LOGGER.exception("RAG memory build failed: %s", exc)
        sys.exit(1)


if __name__ == "__main__":
    main()