Document Chunking for LLMs

import re
import json
import ollama
import logging
import numpy as np
from typing import List
from pathlib import Path
from dataclasses import dataclass
from pydantic import BaseModel, Field
from sklearn.metrics.pairwise import cosine_similarity

logging.basicConfig(level=logging.DEBUG)
log = logging.getLogger(__name__)

# Dont show logging for httpcore or httpx
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("pdfminer").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)


@dataclass
class ChunkingConfig:
    max_chunk_size: int = 1000  # characters
    min_chunk_size: int = 100  # characters
    similarity_threshold: float = 0.15  # cosine distance threshold for splitting (higher means more splits)
    ollama_model: str = "nomic-embed-text:v1.5"  # Ollama model for embeddings
    overlap_sentences: int = 1  # sentences to overlap between chunks


@dataclass
class SectioningConfig:
    # Model settings
    ollama_model: str = "gemma3:1b"
    max_tokens: int = 200
    temperature: float = 0.1  # Low for consistency

    # Similarity thresholds
    section_merge_threshold: float = 0.75
    title_weight: float = 0.6
    category_weight: float = 0.4

    # Size constraints
    max_section_size: int = 3000
    min_section_size: int = 300


class SimpleChunkMetadata(BaseModel):
    title: str = Field(..., description="Concise descriptive title for this content")
    category: str = Field(..., description="Broad topic category")


class Subsection(BaseModel):
    subsection_title: str
    chunks: List[str]
    combined_content: str
    chunk_count: int
    total_length: int


class DocumentSection(BaseModel):
    section_title: str
    category: str
    subsections: List[Subsection]
    total_chunks: int
    total_length: int


class SemanticChunker:
    def __init__(self, config: ChunkingConfig):
        self.config = config

    def get_embedding(self, text: str) -> np.ndarray:
        """Get embedding from Ollama using the Python client"""
        # Ensure the text is not empty, as empty strings might cause issues with embeddings
        if not text.strip():
            return np.zeros(768)  # Return a zero vector for empty text, assuming 768 dimensions for nomic-embed-text
        response = ollama.embeddings(model=self.config.ollama_model, prompt=text)
        return np.array(response["embedding"])

    def split_into_sentences(self, text: str) -> List[str]:
        """Split text into sentences using regex"""
        # Improved regex to handle common sentence endings while avoiding splitting on abbreviations (e.g., Mr. Smith)
        # It looks for . ! or ? followed by whitespace and an uppercase letter, but not if preceded by a common abbreviation pattern.
        # This is a common pattern for sentence splitting.
        sentence_pattern = r"(?<!\b[A-Z]\.)(?<!\b[A-Z][a-z]\.)(?<=[.!?])\s+(?=[A-Z])"
        sentences = re.split(sentence_pattern, text)
        return [s.strip() for s in sentences if s.strip()]

    def chunk_sentences(self, sentences: List[str]) -> List[str]:
        """
        Chunks sentences based on semantic similarity, respecting max_chunk_size,
        and adding overlap.
        """
        if not sentences:
            return []

        chunks: List[str] = []
        current_chunk_sentences: List[str] = []
        current_chunk_char_length = 0

        # Get embeddings for all sentences
        sentence_embeddings = [self.get_embedding(s) for s in sentences]

        for i, sentence in enumerate(sentences):
            sentence_length = len(sentence)

            # Check if adding the current sentence would exceed max_chunk_size
            # and if we have enough content to form a valid chunk
            if current_chunk_char_length + sentence_length > self.config.max_chunk_size and current_chunk_sentences:
                # If we're at the beginning of a potential new chunk and the first sentence itself is too long,
                # we'll add it as a standalone chunk and handle it.
                if current_chunk_char_length == 0:
                    chunks.append(sentence)
                    current_chunk_sentences = []
                    current_chunk_char_length = 0
                    continue  # Move to the next sentence
                else:
                    # Finalize the current chunk
                    chunks.append(" ".join(current_chunk_sentences))
                    # Reset for the next chunk, adding overlap
                    current_chunk_sentences = sentences[max(0, i - self.config.overlap_sentences) : i]
                    current_chunk_char_length = sum(len(s) for s in current_chunk_sentences)

            # Add sentence to current chunk
            current_chunk_sentences.append(sentence)
            current_chunk_char_length += sentence_length

            # Semantic split check (only if there's more than one sentence in current_chunk_sentences)
            if len(current_chunk_sentences) > 1:
                # Compare the last two sentences in the current_chunk_sentences
                # We're interested in the similarity between the last added sentence and the one before it
                # to detect a potential topic shift at the boundary.
                embed1 = sentence_embeddings[i - 1]  # Embedding of the sentence before the current one
                embed2 = sentence_embeddings[i]  # Embedding of the current sentence

                # Calculate cosine distance (1 - cosine_similarity)
                # A higher distance means less similarity.
                if np.dot(embed1, embed2) == 0 and np.linalg.norm(embed1) == 0 and np.linalg.norm(embed2) == 0:
                    # Handle case where both embeddings are zero vectors (e.g., from empty strings), distance is 0
                    distance = 0
                elif np.linalg.norm(embed1) == 0 or np.linalg.norm(embed2) == 0:
                    # If one is zero and other is not, they are dissimilar
                    distance = 1.0
                else:
                    distance = 1 - cosine_similarity(embed1.reshape(1, -1), embed2.reshape(1, -1))[0][0]

                if distance > self.config.similarity_threshold:
                    log.debug(f"Semantic split detected between sentences {i - 1} and {i} with distance {distance:.4f}")
                    # Semantic split detected!
                    # If the chunk is long enough, finalize it before the current sentence.
                    if current_chunk_char_length - sentence_length >= self.config.min_chunk_size:
                        # Append chunk excluding the current sentence, which starts a new one
                        chunks.append(" ".join(current_chunk_sentences[:-1]))
                        # Reset for the next chunk, adding overlap
                        current_chunk_sentences = sentences[max(0, i - self.config.overlap_sentences) : i + 1]
                        current_chunk_char_length = sum(len(s) for s in current_chunk_sentences)
                    # If the chunk is too short, we don't split yet and continue to build it.
                    # This prevents very small chunks due to minor semantic shifts.

        # Add any remaining sentences as the last chunk
        if current_chunk_sentences:
            chunks.append(" ".join(current_chunk_sentences))

        # Post-processing: Ensure no chunks are empty and optionally merge very small chunks
        final_chunks = [chunk for chunk in chunks if chunk.strip()]
        return final_chunks


class SectionGrouper:
    def __init__(self, sectioning_config: SectioningConfig, embedding_model: str = "nomic-embed-text:v1.5"):
        self.config = sectioning_config
        self.embedding_model = embedding_model

    def generate_chunk_metadata(self, chunk: str, previous_metadata) -> SimpleChunkMetadata:
        """Generate structured metadata for a chunk using Ollama"""

        # Prepare previous metadata as a string for context
        metadata = ""
        if len(previous_metadata) > 0:
            metadata = (
                "\n".join(
                    f"{i + 1}. Title: {meta.title}, Category: {meta.category}"
                    for i, (chunk, meta) in enumerate(previous_metadata)
                )
                if previous_metadata
                else ""
            )

            if metadata != "":
                metadata = f"Previous metadata:\n{metadata}\n\n"

        prompt = f"""Analyze the following text and provide structured metadata:

Text: {chunk}

{metadata}

Generate a JSON response with:
- title: A concise 3-8 word descriptive title
- category: A broad 1-3 word topic category

Be specific and descriptive but concise. Keep consistency in section titles and categories across similar content."""

        try:
            response = ollama.generate(
                model=self.config.ollama_model,
                prompt=prompt,
                format=SimpleChunkMetadata.model_json_schema(),
                options={"temperature": self.config.temperature, "num_predict": self.config.max_tokens},
                keep_alive=True,
            )

            # Parse the JSON response
            metadata_dict = json.loads(response["response"])
            return SimpleChunkMetadata(**metadata_dict)

        except Exception as e:
            log.warning(f"Failed to generate metadata for chunk, using fallback: {e}")
            # Fallback metadata
            return SimpleChunkMetadata(
                title=f"Section {hash(chunk[:100]) % 1000}",
                category="General",
            )

    def get_embedding(self, text: str) -> np.ndarray:
        """Get embedding for text"""
        if not text.strip():
            return np.zeros(768)
        response = ollama.embeddings(model=self.embedding_model, prompt=text)
        return np.array(response["embedding"])

    def calculate_similarity(self, metadata1: SimpleChunkMetadata, metadata2: SimpleChunkMetadata) -> float:
        """Calculate weighted similarity between two chunk metadata objects"""

        # Get embeddings for each field
        title1_emb = self.get_embedding(metadata1.title)
        title2_emb = self.get_embedding(metadata2.title)

        category1_emb = self.get_embedding(metadata1.category)
        category2_emb = self.get_embedding(metadata2.category)

        # Calculate cosine similarities
        def safe_cosine_similarity(emb1, emb2):
            if np.linalg.norm(emb1) == 0 or np.linalg.norm(emb2) == 0:
                return 0.0
            return cosine_similarity(emb1.reshape(1, -1), emb2.reshape(1, -1))[0][0]

        title_sim = safe_cosine_similarity(title1_emb, title2_emb)
        category_sim = safe_cosine_similarity(category1_emb, category2_emb)

        # Weighted similarity
        weighted_sim = self.config.title_weight * title_sim + self.config.category_weight * category_sim

        return weighted_sim

    def group_chunks_into_sections(self, chunks: List[str]) -> List[DocumentSection]:
        """Group chunks into coherent sections based on metadata similarity"""

        if not chunks:
            return []

        log.info(f"Generating metadata for {len(chunks)} chunks...")

        # Stage 1: Generate metadata for all chunks
        chunk_metadata = []
        for i, chunk in enumerate(chunks):
            # We pass in the last 5 metadata items as context for the next chunk
            metadata = self.generate_chunk_metadata(
                chunk, chunk_metadata[:-5] if len(chunk_metadata) >= 5 else chunk_metadata
            )
            chunk_metadata.append((chunk, metadata))
            log.debug(f"Chunk {i + 1}: {metadata.title} | {metadata.category}")

        log.info("Grouping chunks into sections...")

        # Stage 2: Group chunks based on similarity
        sections = []
        current_section_chunks = [(chunk_metadata[0][0], chunk_metadata[0][1])]
        current_section_category = chunk_metadata[0][1].category

        for i in range(1, len(chunk_metadata)):
            chunk, metadata = chunk_metadata[i]
            prev_metadata = chunk_metadata[i - 1][1]

            # Calculate similarity with previous chunk
            similarity = self.calculate_similarity(metadata, prev_metadata)

            # Check if we should start a new section
            should_split = (
                # Similarity is below threshold so the current chunk is likely different
                # from the previous section
                similarity < self.config.section_merge_threshold
                # Or if the current section is too long
                or sum(len(c[0]) for c in current_section_chunks) + len(chunk) > self.config.max_section_size
            )

            if should_split and len(current_section_chunks) > 0:
                # Finalize current section
                section = self._create_section(current_section_chunks, current_section_category)
                sections.append(section)

                # Start new section
                current_section_chunks = [(chunk, metadata)]
                current_section_category = metadata.category

                log.debug(f"New section started at chunk {i + 1}, similarity: {similarity:.3f}")
            else:
                # Add to current section
                current_section_chunks.append((chunk, metadata))

        # Add final section
        if current_section_chunks:
            section = self._create_section(current_section_chunks, current_section_category)
            sections.append(section)

        log.info(f"Created {len(sections)} sections from {len(chunks)} chunks")
        return sections

    def _create_section(self, section_chunks: List[tuple], category: str) -> DocumentSection:
        """Create a DocumentSection from a list of (chunk, metadata) tuples"""

        # Group chunks into subsections based on title similarity
        subsections = []
        current_subsection = []
        current_title = section_chunks[0][1].title

        for chunk, metadata in section_chunks:
            if not current_subsection:
                current_subsection = [(chunk, metadata)]
                current_title = metadata.title
            else:
                # Check if we should group with current subsection
                title_sim = self.calculate_similarity(metadata, current_subsection[-1][1])

                if title_sim > 0.8:  # High threshold for subsection grouping
                    current_subsection.append((chunk, metadata))
                else:
                    # Create subsection from current group
                    subsection = self._create_subsection(current_subsection, current_title)
                    subsections.append(subsection)

                    # Start new subsection
                    current_subsection = [(chunk, metadata)]
                    current_title = metadata.title

        # Add final subsection
        if current_subsection:
            subsection = self._create_subsection(current_subsection, current_title)
            subsections.append(subsection)

        # Create section title from most common category and representative title
        section_title = self._generate_section_title(section_chunks)

        return DocumentSection(
            section_title=section_title,
            category=category,
            subsections=subsections,
            total_chunks=len(section_chunks),
            total_length=sum(len(chunk) for chunk, _ in section_chunks),
        )

    def _create_subsection(self, subsection_chunks: List[tuple], title: str) -> Subsection:
        """Create a Subsection from a list of (chunk, metadata) tuples"""
        chunks = [chunk for chunk, _ in subsection_chunks]
        combined_content = "\n\n".join(chunks)

        return Subsection(
            subsection_title=title,
            chunks=chunks,
            combined_content=combined_content,
            chunk_count=len(chunks),
            total_length=len(combined_content),
        )

    def _generate_section_title(self, section_chunks: List[tuple]) -> str:
        """Generate a representative title for the section"""
        # Use the title from the first chunk, or create from category
        if section_chunks:
            first_metadata = section_chunks[0][1]
            if len(section_chunks) == 1:
                return first_metadata.title
            else:
                # For multi-chunk sections, use category-based title
                return f"{first_metadata.category} Overview"
        return "Untitled Section"


def pdf_to_markdown(pdf_file: Path) -> str:
    from markitdown import MarkItDown

    md = MarkItDown()
    result = md.convert(pdf_file)
    text_content = result.text_content
    # Keep only ascii characters
    text_content = "".join(c for c in text_content if ord(c) < 128)
    # Remove leading and trailing whitespace
    text_content = text_content.strip()
    # Remove multiple newlines
    text_content = "\n".join(line.strip() for line in text_content.splitlines() if line.strip())
    return text_content


def get_pdf():
    import requests

    # Sample text chunks
    url = "https://www.berkshirehathaway.com/letters/2024ltr.pdf"
    response = requests.get(url)
    if response.status_code != 200:
        raise ValueError(f"Failed to download PDF: {response.status_code}")
    pdf_file = Path("./example_cache/2024ltr.pdf")
    pdf_file.parent.mkdir(parents=True, exist_ok=True)
    pdf_file.write_bytes(response.content)
    # Convert PDF to markdown text
    text = pdf_to_markdown(pdf_file)
    return text


# Usage example
def main():
    # Example usage
    sample_text = get_pdf()

    # Initialize chunker
    chunking_config = ChunkingConfig(
        max_chunk_size=750,
        min_chunk_size=150,
        similarity_threshold=0.25,
        ollama_model="nomic-embed-text:v1.5",
        overlap_sentences=2,
    )

    # Initialize section grouper
    sectioning_config = SectioningConfig(
        ollama_model="gemma3:1b", section_merge_threshold=0.75, max_section_size=3000, min_section_size=300
    )

    # Stage 1: Semantic chunking
    chunker = SemanticChunker(chunking_config)
    sentences = chunker.split_into_sentences(sample_text)
    chunks = chunker.chunk_sentences(sentences)

    print(f"Stage 1 Complete - Generated {len(chunks)} semantic chunks")

    # Stage 2: Section grouping
    grouper = SectionGrouper(sectioning_config)
    sections = grouper.group_chunks_into_sections(chunks)

    # Write sections to JSON file
    output_file = Path("./example_cache/sections.json")
    output_file.parent.mkdir(parents=True, exist_ok=True)
    with output_file.open("w", encoding="utf-8") as f:
        json.dump([section.model_dump() for section in sections], f, indent=2, ensure_ascii=False)
    log.info(f"Sections written to {output_file}")

    print(f"Stage 2 Complete - Generated {len(sections)} sections")

    # Display results
    print("\n" + "=" * 80)
    print("DOCUMENT STRUCTURE")
    print("=" * 80)

    for i, section in enumerate(sections, 1):
        print(f"\nSECTION {i}: {section.section_title}")
        print(f"Category: {section.category}")
        print(f"Total chunks: {section.total_chunks} | Length: {section.total_length} chars")
        print("-" * 60)

        for j, subsection in enumerate(section.subsections, 1):
            print(f"  {i}.{j} {subsection.subsection_title}")
            print(f"       Chunks: {subsection.chunk_count} | Length: {subsection.total_length} chars")

            # Show first chunk preview
            if subsection.chunks:
                preview = (
                    subsection.chunks[0][:200] + "..." if len(subsection.chunks[0]) > 200 else subsection.chunks[0]
                )
                print(f"       Preview: {preview}")
            print()