src/chunker.py

import hashlib
import re
import time
from typing import Optional

import polars as pl

from src.config import Config


class Chunker:
    def __init__(self, config: Config):
        self.config = config

    def chunk(self, pages_df: pl.DataFrame) -> pl.DataFrame:
        if pages_df.is_empty():
            return pl.DataFrame(
                schema={
                    "chunk_id": pl.Utf8,
                    "url": pl.Utf8,
                    "title": pl.Utf8,
                    "heading_path": pl.Utf8,
                    "text": pl.Utf8,
                    "chunk_index": pl.Int32,
                    "timestamp": pl.Float64,
                }
            )

        all_chunks = []

        for row in pages_df.iter_rows(named=True):
            markdown = row.get("markdown", "") or ""
            chunks = self._split_markdown(markdown)
            current_heading = row.get("title", "")

            for i, chunk_text in enumerate(chunks):
                if not chunk_text.strip():
                    continue
                chunk_id = hashlib.md5(
                    f"{row['url']}_{i}_{chunk_text[:50]}".encode()
                ).hexdigest()
                all_chunks.append({
                    "chunk_id": chunk_id,
                    "url": row["url"],
                    "title": row.get("title", ""),
                    "heading_path": current_heading,
                    "text": chunk_text.strip(),
                    "chunk_index": i,
                    "timestamp": row.get("timestamp", time.time()),
                })

        if not all_chunks:
            return pl.DataFrame(
                schema={
                    "chunk_id": pl.Utf8,
                    "url": pl.Utf8,
                    "title": pl.Utf8,
                    "heading_path": pl.Utf8,
                    "text": pl.Utf8,
                    "chunk_index": pl.Int32,
                    "timestamp": pl.Float64,
                }
            )

        return pl.DataFrame(all_chunks)

    def _split_markdown(self, text: str) -> list[str]:
        if not text:
            return []

        sections = re.split(r"(^#+\s+.*$)", text, flags=re.MULTILINE)
        merged_sections = []
        current = ""

        for part in sections:
            if re.match(r"^#+\s+", part):
                if current.strip():
                    merged_sections.append(current.strip())
                current = part
            else:
                current += part
        if current.strip():
            merged_sections.append(current.strip())

        chunks = []
        for section in merged_sections:
            chunks.extend(self._split_section(section))

        return chunks or [text]

    def _split_section(self, section: str) -> list[str]:
        if len(section) <= self.config.chunk_max_chars:
            return [section]

        paragraphs = re.split(r"\n\s*\n", section)
        chunks = []
        current = ""

        for para in paragraphs:
            if len(current) + len(para) + 1 <= self.config.chunk_max_chars:
                current = (current + "\n\n" + para).strip() if current else para
            else:
                if current:
                    chunks.append(current)
                current = para

                if len(current) > self.config.chunk_max_chars:
                    for i in range(0, len(current), self.config.chunk_max_chars - self.config.chunk_overlap):
                        sub = current[i: i + self.config.chunk_max_chars]
                        if len(sub) > self.config.chunk_max_chars // 2:
                            chunks.append(sub)

        if current:
            chunks.append(current)

        return chunks