raglens-toolkit 0.1.0__py3-none-any.whl

raglens/__init__.py

@@ -0,0 +1,13 @@
+__version__ = "0.1.0"
+
+__all__ = ["RagLensPipeline"]
+
+
+def __getattr__(name):
+    # Lazy: raglens.pipeline pulls in Docling/transformers, which is heavy
+    # and irrelevant to callers that only need e.g. raglens.evaluation or
+    # raglens.ragas (CLI report command, the dashboard, unit tests, ...).
+    if name == "RagLensPipeline":
+        from raglens.pipeline import RagLensPipeline
+        return RagLensPipeline
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

raglens/cache/__init__.py

@@ -0,0 +1,19 @@
+from .parser_cache import (
+    save_parsed_documents,
+    load_parsed_documents,
+    parsed_documents_exist
+)
+from .chunk_cache import (
+    save_chunks,
+    load_chunks,
+    chunks_exist,
+)
+
+__all__ = [
+    "save_parsed_documents",
+    "load_parsed_documents",
+    "parsed_documents_exist",
+    "save_chunks",
+    "load_chunks",
+    "chunks_exist",
+]

raglens/cache/chunk_cache.py

@@ -0,0 +1,32 @@
+import os
+import pickle
+from pathlib import Path
+
+_DEFAULT_CACHE_DIR = "data/cache"
+
+
+def _cache_dir() -> Path:
+    path = Path(
+        os.getenv("RAGLENS_CHUNK_CACHE_DIR", _DEFAULT_CACHE_DIR)
+    )
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def save_chunks(chunks):
+    path = _cache_dir() / "all_chunks.pkl"
+    with open(path, "wb") as f:
+        pickle.dump(chunks, f)
+    print(f"Saved -> {path}")
+
+
+def load_chunks():
+    path = _cache_dir() / "all_chunks.pkl"
+    with open(path, "rb") as f:
+        chunks = pickle.load(f)
+    print(f"Loaded -> {path}")
+    return chunks
+
+
+def chunks_exist():
+    return (_cache_dir() / "all_chunks.pkl").exists()

raglens/cache/parser_cache.py

@@ -0,0 +1,71 @@
+import os
+import pickle
+from pathlib import Path
+
+
+_DEFAULT_CACHE_DIR = "data/cache"
+
+
+def _cache_dir() -> Path:
+    path = Path(
+        os.getenv("RAGLENS_PARSER_CACHE_DIR", _DEFAULT_CACHE_DIR)
+    )
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def save_parsed_documents(
+    parsed_documents
+):
+
+    path = (
+        _cache_dir()
+        / "parsed_documents.pkl"
+    )
+
+    with open(
+        path,
+        "wb"
+    ) as file:
+
+        pickle.dump(
+            parsed_documents,
+            file
+        )
+
+    print(
+        f"Saved -> {path}"
+    )
+
+
+def load_parsed_documents():
+
+    path = (
+        _cache_dir()
+        / "parsed_documents.pkl"
+    )
+
+    with open(
+        path,
+        "rb"
+    ) as file:
+
+        parsed_documents = (
+            pickle.load(
+                file
+            )
+        )
+
+    print(
+        f"Loaded -> {path}"
+    )
+
+    return parsed_documents
+
+
+def parsed_documents_exist():
+
+    return (
+        _cache_dir()
+        / "parsed_documents.pkl"
+    ).exists()

raglens/chunking/__init__.py

@@ -0,0 +1,3 @@
+from raglens.chunking.section_chunker import (
+    SectionChunker
+)

raglens/chunking/section_chunker.py

@@ -0,0 +1,296 @@
+from langchain_text_splitters import (
+    RecursiveCharacterTextSplitter
+)
+
+from raglens.chunking.structure_preserver import (
+    StructurePreserver
+)
+
+from raglens.models import (
+    Chunk,
+    FlattenedSection,
+    Document
+)
+
+import re
+
+
+class SectionChunker:
+
+    def __init__(
+        self,
+        chunk_size=1200,
+        chunk_overlap=200,
+        max_section_length=2500
+    ):
+
+        self.chunk_size = chunk_size
+
+        self.chunk_overlap = chunk_overlap
+
+        self.max_section_length = (
+            max_section_length
+        )
+
+        self.splitter = (
+            RecursiveCharacterTextSplitter(
+                chunk_size=chunk_size,
+                chunk_overlap=chunk_overlap
+            )
+        )
+
+        self.preserver = (
+            StructurePreserver()
+        )
+        self.safety_splitter = (
+            RecursiveCharacterTextSplitter(
+            chunk_size=2500,
+            chunk_overlap=200
+            )
+        )
+
+        self.max_embeddable_length = 5000
+
+    def chunk(
+        self,
+        document: Document,
+        sections: list[FlattenedSection]
+    ) -> list[Chunk]:
+
+        chunks = []
+
+        chunk_counter = 0
+
+        for section in sections:
+
+            # ----------------------------------
+            # Skip empty sections
+            # ----------------------------------
+            if len(
+                section.content.strip()
+            ) == 0:
+                continue
+
+            # ----------------------------------
+            # Small section
+            # ----------------------------------
+            if (
+                len(section.content)
+                <= self.max_section_length
+            ):
+
+                chunks.append(
+                    Chunk(
+                        chunk_id=f"{document.doc_id}"f"_{section.section_id}"f"_{chunk_counter}",
+                        parent_doc_id=document.doc_id,
+                        parent_section_id=section.section_id,
+                        chunk_order=chunk_counter,
+                        fragment_index=-1,
+                        section_title=section.section_title,
+                        path=section.path,
+                        level=section.level,
+                        chunk_type="parent_section",
+                        content=section.content,
+                        metadata={}
+                    )
+                )
+
+                chunk_counter += 1
+
+            # ----------------------------------
+            # Large section
+            # ----------------------------------
+            else:
+
+                # Parent section chunk
+                chunks.append(
+                    Chunk(
+                        chunk_id=f"{document.doc_id}"f"_{section.section_id}"f"_{chunk_counter}",
+                        parent_doc_id=document.doc_id,
+                        parent_section_id=section.section_id,
+                        chunk_order=chunk_counter,
+                        fragment_index=-1,
+                        section_title=section.section_title,
+                        path=section.path,
+                        level=section.level,
+                        chunk_type="parent_section",
+                        content=section.content,
+                        metadata={}
+                    )
+                )
+
+                chunk_counter += 1
+
+                protected_text, mapping = (
+                    self.preserver.protect(
+                        section.content
+                    )
+                )
+
+                # ----------------------------------
+                # Create table fragment chunks
+                # directly
+                # ----------------------------------
+
+                table_placeholders = re.findall(
+                    r"__TABLE_\d+_FRAGMENT_\d+__",
+                    protected_text
+                )
+
+                for placeholder in (
+                    table_placeholders
+                ):
+
+                    table_info = (
+                        mapping[
+                            placeholder
+                        ]
+                    )
+
+                    chunks.append(
+                        Chunk(
+                            chunk_id=f"{document.doc_id}"f"_{section.section_id}"f"_{chunk_counter}",
+                            parent_doc_id=document.doc_id,
+                            parent_section_id=section.section_id,
+                            chunk_order=chunk_counter,
+                            fragment_index=(
+                                table_info[
+                                    "table_fragment_index"
+                                ]
+                            ),
+                            section_title=section.section_title,
+                            path=section.path,
+                            level=section.level,
+                            chunk_type="table_fragment",
+                            content=(
+                                table_info[
+                                    "content"
+                                ]
+                            ),
+                            metadata={
+                                "table_id":
+                                    table_info[
+                                        "table_id"
+                                    ],
+
+                                "table_fragment_index":
+                                    table_info[
+                                        "table_fragment_index"
+                                    ],
+
+                                "table_fragment_count":
+                                    table_info[
+                                        "table_fragment_count"
+                                    ]
+                            }
+                        )
+                    )
+
+                    chunk_counter += 1
+
+                # ----------------------------------
+                # Remove table placeholders
+                # before recursive splitting
+                # ----------------------------------
+
+                text_without_tables = (
+                    protected_text
+                )
+
+                for placeholder in (
+                    table_placeholders
+                ):
+
+                    text_without_tables = (
+                        text_without_tables.replace(
+                            placeholder,
+                            ""
+                        )
+                    )
+
+                # ----------------------------------
+                # Split remaining text
+                # ----------------------------------
+
+                fragments = (
+                    self.splitter.split_text(
+                        text_without_tables
+                    )
+                )
+
+                for idx, fragment in enumerate(
+                    fragments
+                ):
+
+                    restored_fragment = (
+                        self.preserver.restore(
+                            fragment,
+                            mapping
+                        )
+                    )
+
+                    if (
+                        len(
+                            restored_fragment.strip()
+                        ) == 0
+                    ):
+                        continue
+
+                    # ----------------------------------
+                    # Safety split for embedding limits
+                    # ----------------------------------
+
+                    if (
+                        len(restored_fragment)
+                        > self.max_embeddable_length
+                    ):
+
+                        safety_fragments = (
+                            self.safety_splitter.split_text(
+                                restored_fragment
+                            )
+                        )
+
+                        for safety_idx, safety_fragment in enumerate(
+                            safety_fragments
+                        ):
+
+                            chunks.append(
+                                Chunk(
+                                    chunk_id=f"{document.doc_id}"f"_{section.section_id}"f"_{chunk_counter}",
+                                    parent_doc_id=document.doc_id,
+                                    parent_section_id=section.section_id,
+                                    chunk_order=chunk_counter,
+                                    fragment_index=safety_idx,
+                                    section_title=section.section_title,
+                                    path=section.path,
+                                    level=section.level,
+                                    chunk_type="section_fragment",
+                                    content=safety_fragment,
+                                    metadata={}
+                                )
+                            )
+
+                            chunk_counter += 1
+
+                    else:
+
+                        chunks.append(
+                            Chunk(
+                                chunk_id=f"{document.doc_id}"f"_{section.section_id}"f"_{chunk_counter}",
+                                parent_doc_id=document.doc_id,
+                                parent_section_id=section.section_id,
+                                chunk_order=chunk_counter,
+                                fragment_index=idx,
+                                section_title=section.section_title,
+                                path=section.path,
+                                level=section.level,
+                                chunk_type="section_fragment",
+                                content=restored_fragment,
+                                metadata={}
+                            )
+                        )
+
+                        chunk_counter += 1
+
+        return chunks

raglens/chunking/structure_preserver.py

@@ -0,0 +1,121 @@
+import re
+
+
+class StructurePreserver:
+    MAX_ATOMIC_TABLE_LENGTH = 3000
+    TABLE_FRAGMENT_SIZE = 1500
+
+    def protect(self, text: str):
+        replacements = {}
+        counter = 0
+
+        # --------------------------------------------------
+        # Formula blocks
+        # --------------------------------------------------
+        formula_pattern = r"\$\$(.*?)\$\$"
+
+        def replace_formula(match):
+            nonlocal counter
+            key = f"__FORMULA_{counter}__"
+            replacements[key] = {
+                "content": match.group(0)
+            }
+            counter += 1
+            return key
+
+        text = re.sub(formula_pattern, replace_formula, text, flags=re.DOTALL)
+
+        # --------------------------------------------------
+        # Tables
+        # --------------------------------------------------
+        lines = text.splitlines(keepends=True)
+        output = []
+        i = 0
+
+        while i < len(lines):
+            line = lines[i]
+
+            # Detect markdown table start
+            if "|" in line and i + 1 < len(lines) and "|" in lines[i + 1]:
+                table_lines = []
+                while i < len(lines) and "|" in lines[i]:
+                    table_lines.append(lines[i])
+                    i += 1
+
+                table_text = "".join(table_lines)
+
+                # Small table -> atomic
+                if len(table_text) <= self.MAX_ATOMIC_TABLE_LENGTH:
+                    key = f"__TABLE_{counter}__"
+                    replacements[key] = {
+                        "content": table_text
+                    }
+                    output.append(key + "\n")
+                    counter += 1
+
+                # Large table -> split by rows
+                else:
+                    table_chunks = self._split_large_table(
+                            table_text
+                    )
+                    table_id = counter
+
+                    for fragment_idx, chunk in enumerate(table_chunks):
+
+                        key = (
+                            f"__TABLE_{table_id}_FRAGMENT_{fragment_idx}__"
+                        )
+
+                        replacements[key] = {
+                            "content": chunk,
+                            "table_id": f"table_{table_id}",
+                            "table_fragment_index": fragment_idx,
+                            "table_fragment_count": len(table_chunks)
+                        }
+
+                        output.append(
+                            key + "\n"
+                        )
+                    counter += 1
+            else:
+                output.append(line)
+                i += 1
+
+        text = "".join(output)
+        return text, replacements
+
+    def restore(self, text: str, replacements: dict):
+        # Replace keys with their saved content
+        for key, value in replacements.items():
+            text = text.replace(key, value["content"])
+        return text
+
+    def _split_large_table(self, table_text: str):
+        lines = table_text.strip().splitlines()
+        if len(lines) < 3:
+            return [table_text]
+
+        header = lines[0]
+        separator = lines[1]
+        rows = lines[2:]
+
+        fragments = []
+        current_rows = []
+        current_size = 0
+
+        for row in rows:
+            row_size = len(row)
+            if current_rows and current_size + row_size > self.TABLE_FRAGMENT_SIZE:
+                fragment = header + "\n" + separator + "\n" + "\n".join(current_rows)
+                fragments.append(fragment)
+                current_rows = []
+                current_size = 0
+
+            current_rows.append(row)
+            current_size += row_size
+
+        if current_rows:
+            fragment = header + "\n" + separator + "\n" + "\n".join(current_rows)
+            fragments.append(fragment)
+
+        return fragments