File size: 3,629 Bytes
49f0cfb
dc1b199
49f0cfb
 
 
 
 
 
 
 
 
 
 
 
dc1b199
 
 
 
 
49f0cfb
 
dc1b199
49f0cfb
dc1b199
 
 
49f0cfb
 
dc1b199
 
 
49f0cfb
 
 
dc1b199
 
 
 
 
49f0cfb
dc1b199
 
 
49f0cfb
dc1b199
49f0cfb
 
dc1b199
 
49f0cfb
 
 
 
 
dc1b199
 
49f0cfb
 
 
dc1b199
 
49f0cfb
dc1b199
 
 
49f0cfb
 
dc1b199
49f0cfb
 
 
 
 
dc1b199
 
 
49f0cfb
 
 
 
 
 
dc1b199
49f0cfb
 
 
 
dc1b199
 
49f0cfb
 
 
dc1b199
 
49f0cfb
dc1b199
49f0cfb
dc1b199
49f0cfb
 
 
dc1b199
49f0cfb
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
"""Text chunking via LangChain RecursiveCharacterTextSplitter.

Replaces the custom recursive splitter with the standard LangChain
implementation backed by tiktoken (``cl100k_base`` encoding).

The splitter receives a list of LangChain ``Document`` objects (from the
parsers) and returns a new list of ``Document`` objects with split content.
Each output document inherits the source document's metadata.

Compared to master branch:
- ✅ Less code — delegates token counting to LangChain / tiktoken
- ✅ Standard interface: ``split_documents(docs)`` is the LangChain pattern
- ✅ Accurate token-based chunk sizing via tiktoken
- ✅ Metadata automatically propagated from parent documents
"""

import logging

import tiktoken
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

from app.config import settings

logger = logging.getLogger(__name__)

# Shared tiktoken encoding used for both splitting and prompt token counting
_ENCODING_NAME = "cl100k_base"


def count_tokens(text: str) -> int:
    """Count the number of tiktoken tokens in ``text``.

    Used by ``app.generator.prompts`` to enforce context-window budgets.

    Args:
        text: Any plain-text string.

    Returns:
        Integer token count.

    Example::

        n = count_tokens("The property is a two-storey terrace.")
    """
    enc = tiktoken.get_encoding(_ENCODING_NAME)
    return len(enc.encode(text))


def build_splitter(
    chunk_size: int | None = None,
    chunk_overlap: int | None = None,
) -> RecursiveCharacterTextSplitter:
    """Create a tiktoken-backed :class:`RecursiveCharacterTextSplitter`.

    Args:
        chunk_size: Token count per chunk. Defaults to ``settings.chunk_size``.
        chunk_overlap: Token overlap between consecutive chunks.
            Defaults to ``settings.chunk_overlap``.

    Returns:
        Configured :class:`RecursiveCharacterTextSplitter`.

    Example::

        splitter = build_splitter(chunk_size=500, chunk_overlap=75)
        chunks = splitter.split_documents(docs)
    """
    return RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        encoding_name=_ENCODING_NAME,
        chunk_size=chunk_size or settings.chunk_size,
        chunk_overlap=chunk_overlap or settings.chunk_overlap,
        separators=["\n\n", "\n", ". ", " ", ""],
    )


def split_documents(
    documents: list[Document],
    chunk_size: int | None = None,
    chunk_overlap: int | None = None,
) -> list[Document]:
    """Split ``documents`` into token-bounded chunks.

    Each output ``Document`` inherits the metadata of its parent and has its
    ``page_content`` set to the chunk text.  No additional metadata is added
    here — the caller (pipeline) is responsible for injecting ``chunk_id``,
    ``doc_id``, ``tenant_id``, etc.

    Args:
        documents: Input LangChain :class:`Document` objects.
        chunk_size: Override for ``settings.chunk_size``.
        chunk_overlap: Override for ``settings.chunk_overlap``.

    Returns:
        List of chunk :class:`Document` objects.

    Example::

        docs = parse_pdf(path)
        chunks = split_documents(docs)
        # Each chunk has page_content ≤ settings.chunk_size tokens
    """
    splitter = build_splitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    chunks = splitter.split_documents(documents)
    logger.debug(
        "Split %d document(s) into %d chunks (size=%d, overlap=%d)",
        len(documents),
        len(chunks),
        chunk_size or settings.chunk_size,
        chunk_overlap or settings.chunk_overlap,
    )
    return chunks