arc53--docsgpt/application/parser/chunking.py

import re
from typing import List, Tuple
import logging
from application.parser.chunking_creator import ChunkerCreator
from application.parser.schema.base import Document
from application.utils import get_encoding

logger = logging.getLogger(__name__)


class Chunker:
    """Classic token-window chunker (registered as ``classic_chunk``).

    Strategy dispatch lives in ``ChunkerCreator``; this class is one
    registered implementation. The ``chunking_strategy`` arg is retained for
    backward-compatible construction and is not used for dispatch here.
    """

    def __init__(
        self,
        chunking_strategy: str = "classic_chunk",
        max_tokens: int = 2000,
        min_tokens: int = 150,
        duplicate_headers: bool = False,
    ):
        self.chunking_strategy = chunking_strategy
        self.max_tokens = max_tokens
        self.min_tokens = min_tokens
        self.duplicate_headers = duplicate_headers
        self.encoding = get_encoding()

    def separate_header_and_body(self, text: str) -> Tuple[str, str]:
        header_pattern = r"^(.*?\n){3}"
        match = re.match(header_pattern, text)
        if match:
            header = match.group(0)
            body = text[len(header):]
        else:
            header, body = "", text  # No header, treat entire text as body
        return header, body



    def split_document(self, doc: Document) -> List[Document]:
        split_docs = []
        header, body = self.separate_header_and_body(doc.text)
        header_tokens = self.encoding.encode(header) if header else []
        body_tokens = self.encoding.encode(body)

        current_position = 0
        part_index = 0
        while current_position < len(body_tokens):
            end_position = current_position + self.max_tokens - len(header_tokens)
            chunk_tokens = (header_tokens + body_tokens[current_position:end_position]
                            if self.duplicate_headers or part_index == 0 else body_tokens[current_position:end_position])
            chunk_text = self.encoding.decode(chunk_tokens)
            new_doc = Document(
                text=chunk_text,
                doc_id=f"{doc.doc_id}-{part_index}",
                embedding=doc.embedding,
                extra_info={**(doc.extra_info or {}), "token_count": len(chunk_tokens)}
            )
            split_docs.append(new_doc)
            current_position = end_position
            part_index += 1
            header_tokens = []
        return split_docs

    def classic_chunk(self, documents: List[Document]) -> List[Document]:
        processed_docs = []
        i = 0
        while i < len(documents):
            doc = documents[i]
            tokens = self.encoding.encode(doc.text)
            token_count = len(tokens)

            if self.min_tokens <= token_count <= self.max_tokens:
                doc.extra_info = doc.extra_info or {}
                doc.extra_info["token_count"] = token_count
                processed_docs.append(doc)
                i += 1
            elif token_count < self.min_tokens:

                doc.extra_info = doc.extra_info or {}
                doc.extra_info["token_count"] = token_count
                processed_docs.append(doc)
                i += 1
            else:
                # Split large documents
                processed_docs.extend(self.split_document(doc))
                i += 1
        return processed_docs

    def chunk(
        self,
        documents: List[Document]
    ) -> List[Document]:
        return self.classic_chunk(documents)


ChunkerCreator.register("classic_chunk", Chunker)