Chonkie Python 教程：分块边界、token 预算与元数据校验

python3.11 -m venv .venv
.venv/bin/python -m pip install \
  'chonkie==1.7.0' 'chonkie-core==0.10.2' 'tokie==0.1.4' \
  'tiktoken==0.14.0'
.venv/bin/python -c 'import tiktoken; tiktoken.get_encoding("cl100k_base")'

import hashlib
import importlib.metadata
import json
from dataclasses import dataclass

import tiktoken
from chonkie import RecursiveChunker


@dataclass(frozen=True)
class Source:
    tenant_id: str
    document_id: str
    title: str
    text: str


def digest(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def encode_id(parts: list) -> str:
    return digest(json.dumps(parts, ensure_ascii=False, separators=(",", ":")))


def count_tokens(text: str) -> int:
    # Tokenizer assets must be cached before offline execution.
    enc = tiktoken.get_encoding("cl100k_base")
    return len(enc.encode(text, disallowed_special=()))


def prefix(source: Source) -> str:
    return f"Document: {source.title}\n\n"


def strategy_id(chunk_size: int, max_input_tokens: int) -> str:
    return encode_id([
        "recursive", "cl100k_base", chunk_size, max_input_tokens,
        importlib.metadata.version("chonkie"),
        importlib.metadata.version("chonkie-core"),
        importlib.metadata.version("tokie"),
        importlib.metadata.version("tiktoken"),
    ])


def row_id(source: Source, row: dict) -> str:
    return encode_id([
        source.tenant_id, source.document_id, row["source_version"],
        row["strategy_id"], row["start"], row["end"], row["embedding_text"],
    ])


def audit(source: Source, rows: list[dict], chunk_size: int,
          max_input_tokens: int) -> None:
    version = digest(source.text)
    expected_strategy = strategy_id(chunk_size, max_input_tokens)
    cursor, seen = 0, set()
    for row in rows:
        if (row["tenant_id"], row["document_id"], row["source_version"],
                row["strategy_id"]) != (
                source.tenant_id, source.document_id, version, expected_strategy):
            raise ValueError("Source metadata mismatch")
        start, end = row["start"], row["end"]
        if type(start) is not int or type(end) is not int:
            raise ValueError("Offsets must be integers")
        if not (cursor <= start < end <= len(source.text)):
            raise ValueError("Invalid or overlapping source offsets")
        if source.text[cursor:start].strip():
            raise ValueError("Source content missing before chunk")
        if source.text[start:end] != row["text"]:
            raise ValueError("Chunk does not match source span")
        payload = prefix(source) + row["text"]
        actual = count_tokens(payload)
        if row["embedding_text"] != payload or row["token_count"] != actual:
            raise ValueError("Embedding payload or token count mismatch")
        if actual > max_input_tokens:
            raise ValueError("Embedding payload exceeds configured budget")
        if row["chunk_id"] != row_id(source, row) or row["chunk_id"] in seen:
            raise ValueError("Invalid or duplicate chunk ID")
        seen.add(row["chunk_id"])
        cursor = end
    if source.text[cursor:].strip():
        raise ValueError("Source content missing after last chunk")


def prepare(source: Source, chunk_size: int = 64,
            max_input_tokens: int = 96) -> list[dict]:
    if not all(isinstance(v, str) and v.strip()
               for v in [source.tenant_id, source.document_id, source.title]):
        raise ValueError("Trusted source identifiers and title are required")
    if not isinstance(source.text, str):
        raise ValueError("Source text must be a string")
    if type(chunk_size) is not int or type(max_input_tokens) is not int:
        raise ValueError("Budgets must be integers")
    if not (0 < chunk_size < max_input_tokens):
        raise ValueError("Leave room for the embedding prefix")
    if count_tokens(prefix(source)) >= max_input_tokens:
        raise ValueError("Title prefix exhausts the embedding budget")
    if not source.text.strip():
        return []
    chunker = RecursiveChunker(
        tokenizer=tiktoken.get_encoding("cl100k_base"),
        chunk_size=chunk_size,
        min_characters_per_chunk=1,
    )
    version = digest(source.text)
    strategy = strategy_id(chunk_size, max_input_tokens)
    rows = []
    for chunk in chunker.chunk(source.text):
        payload = prefix(source) + chunk.text
        row = {
            "tenant_id": source.tenant_id, "document_id": source.document_id,
            "source_version": version, "strategy_id": strategy,
            "start": chunk.start_index, "end": chunk.end_index,
            "text": chunk.text, "embedding_text": payload,
            "token_count": count_tokens(payload),
        }
        # Chonkie's generated chunk ID is not used as an idempotency key.
        row["chunk_id"] = row_id(source, row)
        rows.append(row)
    # No partial batch is returned when any check fails.
    audit(source, rows, chunk_size, max_input_tokens)
    return rows


if __name__ == "__main__":
    source = Source("demo-tenant", "handbook-1", "示例手册",
                    ("中文与 English 同行。保留 emoji 🦛 与 e\u0301。\n\n" * 8))
    rows = prepare(source)
    print(json.dumps({"chunks": len(rows),
                      "max_payload_tokens": max(r["token_count"] for r in rows),
                      "budget": 96}, ensure_ascii=False))

import copy
import unittest
from dataclasses import replace

from chonkie import TokenChunker

from chunk_workflow import Source, audit, count_tokens, prepare


class ChunkWorkflowTests(unittest.TestCase):
    def setUp(self):
        self.source = Source("tenant-a", "doc-1", "合成手册",
            ("重复段落。English text. 🦛 e\u0301\n\n" * 20))

    def test_native_unicode_spans_and_payload_budget(self):
        rows = prepare(self.source)
        self.assertGreater(len(rows), 1)
        for row in rows:
            self.assertEqual(self.source.text[row["start"]:row["end"]], row["text"])
            self.assertEqual(count_tokens(row["embedding_text"]), row["token_count"])
            self.assertLessEqual(row["token_count"], 96)

    def test_ids_repeat_stably_and_identical_passages_remain_distinct(self):
        a, b = prepare(self.source), prepare(self.source)
        self.assertEqual([r["chunk_id"] for r in a], [r["chunk_id"] for r in b])
        self.assertEqual(len(a), len({r["chunk_id"] for r in a}))
        self.assertLess(len({r["text"] for r in a}), len(a))

    def test_text_title_and_strategy_changes_invalidate_ids(self):
        original = prepare(self.source)[0]["chunk_id"]
        for source in [replace(self.source, text=self.source.text + "更新"),
                       replace(self.source, title="新标题")]:
            self.assertNotEqual(prepare(source)[0]["chunk_id"], original)
        self.assertNotEqual(prepare(self.source, chunk_size=48)[0]["chunk_id"], original)

    def test_same_text_in_two_tenants_has_distinct_identity(self):
        a, b = prepare(self.source), prepare(replace(self.source, tenant_id="tenant-b"))
        self.assertTrue(set(r["chunk_id"] for r in a).isdisjoint(r["chunk_id"] for r in b))
        self.assertTrue(all(r["tenant_id"] == "tenant-a" for r in a))

    def test_oversized_title_is_rejected(self):
        with self.assertRaises(ValueError):
            prepare(replace(self.source, title="很长标题" * 100))

    def test_tampered_text_offset_and_missing_chunk_are_rejected(self):
        original = prepare(self.source)
        cases = []
        for field, value in [("text", "伪造内容"), ("start", 1)]:
            rows = copy.deepcopy(original); rows[0][field] = value; cases.append(rows)
        cases.append(original[1:])
        for rows in cases:
            with self.subTest(kind="corrupt source span"):
                with self.assertRaises(ValueError): audit(self.source, rows, 64, 96)

    def test_missing_and_foreign_metadata_are_rejected(self):
        for field in ["tenant_id", "document_id", "source_version", "strategy_id"]:
            for remove in [True, False]:
                rows = prepare(self.source)
                if remove: del rows[0][field]
                else: rows[0][field] = "foreign"
                with self.subTest(field=field, remove=remove):
                    with self.assertRaises((ValueError, KeyError)):
                        audit(self.source, rows, 64, 96)

    def test_payload_count_and_budget_rechecked_before_export(self):
        rows = prepare(self.source)
        rows[0]["token_count"] = 0
        with self.assertRaises(ValueError): audit(self.source, rows, 64, 96)
        with self.assertRaises(ValueError): prepare(self.source, 64, 65)

    def test_character_overlap_is_not_model_token_budget(self):
        text = "🦛" * 30
        chunks = TokenChunker(tokenizer="character", chunk_size=12, chunk_overlap=3).chunk(text)
        for c in chunks:
            self.assertEqual(text[c.start_index:c.end_index], c.text)
            self.assertGreater(count_tokens(c.text), c.token_count)
        self.assertEqual(chunks[1].start_index, 9)

    def test_blank_input_and_invalid_settings(self):
        self.assertEqual(prepare(replace(self.source, text=" \n\t")), [])
        for kwargs in [{"chunk_size": 0}, {"chunk_size": True},
                       {"chunk_size": 100, "max_input_tokens": 96}]:
            with self.assertRaises(ValueError): prepare(self.source, **kwargs)
        with self.assertRaises(ValueError): prepare(replace(self.source, tenant_id=""))


if __name__ == "__main__":
    unittest.main(verbosity=2)

.venv/bin/python -m unittest -v test_chunk_workflow
.venv/bin/python chunk_workflow.py

Ran 10 tests
OK