kiln_ai.datamodel.chunk

  1import logging
  2from enum import Enum
  3from typing import TYPE_CHECKING, Annotated, List, Union
  4
  5import anyio
  6from pydantic import (
  7    AfterValidator,
  8    BaseModel,
  9    Field,
 10    NonNegativeInt,
 11    PositiveInt,
 12    SerializationInfo,
 13    ValidationInfo,
 14    field_serializer,
 15    model_validator,
 16)
 17from typing_extensions import Literal, TypedDict
 18
 19from kiln_ai.datamodel.basemodel import (
 20    ID_TYPE,
 21    FilenameString,
 22    KilnAttachmentModel,
 23    KilnParentedModel,
 24    KilnParentModel,
 25)
 26from kiln_ai.datamodel.embedding import ChunkEmbeddings
 27
 28logger = logging.getLogger(__name__)
 29
 30if TYPE_CHECKING:
 31    from kiln_ai.datamodel.extraction import Extraction
 32    from kiln_ai.datamodel.project import Project
 33
 34
 35class ChunkerType(str, Enum):
 36    FIXED_WINDOW = "fixed_window"
 37    SEMANTIC = "semantic"
 38
 39
 40class SemanticChunkerProperties(TypedDict, total=True):
 41    chunker_type: Literal[ChunkerType.SEMANTIC]
 42    embedding_config_id: str
 43    buffer_size: PositiveInt
 44    breakpoint_percentile_threshold: NonNegativeInt
 45    include_metadata: bool
 46    include_prev_next_rel: bool
 47
 48
 49class FixedWindowChunkerProperties(TypedDict, total=True):
 50    chunker_type: Literal[ChunkerType.FIXED_WINDOW]
 51    chunk_overlap: NonNegativeInt
 52    chunk_size: PositiveInt
 53
 54
 55def validate_fixed_window_chunker_properties(
 56    properties: FixedWindowChunkerProperties,
 57) -> FixedWindowChunkerProperties:
 58    """Validate the properties for the fixed window chunker and set defaults if needed."""
 59    # the typed dict only validates the shape and types, but not the logic, so we validate here
 60    if properties["chunk_overlap"] >= properties["chunk_size"]:
 61        raise ValueError("Chunk overlap must be less than chunk size.")
 62
 63    return properties
 64
 65
 66def validate_semantic_chunker_properties(
 67    properties: SemanticChunkerProperties,
 68) -> SemanticChunkerProperties:
 69    """Validate the properties for the semantic chunker."""
 70    buffer_size = properties["buffer_size"]
 71    if buffer_size < 1:
 72        raise ValueError("buffer_size must be greater than or equal to 1.")
 73
 74    breakpoint_percentile_threshold = properties["breakpoint_percentile_threshold"]
 75    if not (0 <= breakpoint_percentile_threshold <= 100):
 76        raise ValueError("breakpoint_percentile_threshold must be between 0 and 100.")
 77
 78    return properties
 79
 80
 81SemanticChunkerPropertiesValidator = Annotated[
 82    SemanticChunkerProperties,
 83    AfterValidator(lambda v: validate_semantic_chunker_properties(v)),
 84]
 85
 86FixedWindowChunkerPropertiesValidator = Annotated[
 87    FixedWindowChunkerProperties,
 88    AfterValidator(lambda v: validate_fixed_window_chunker_properties(v)),
 89]
 90
 91
 92class ChunkerConfig(KilnParentedModel):
 93    name: FilenameString = Field(
 94        description="A name to identify the chunker config.",
 95    )
 96    description: str | None = Field(
 97        default=None, description="The description of the chunker config"
 98    )
 99    chunker_type: ChunkerType = Field(
100        description="This is used to determine the type of chunker to use.",
101    )
102    properties: (
103        SemanticChunkerPropertiesValidator | FixedWindowChunkerPropertiesValidator
104    ) = Field(
105        description="Properties to be used to execute the chunker config. This is chunker_type specific and should serialize to a json dict.",
106        discriminator="chunker_type",
107    )
108
109    # Workaround to return typed parent without importing Project
110    def parent_project(self) -> Union["Project", None]:
111        if self.parent is None or self.parent.__class__.__name__ != "Project":
112            return None
113        return self.parent  # type: ignore
114
115    @model_validator(mode="before")
116    def upgrade_missing_discriminator_properties(
117        cls, data: dict, info: ValidationInfo
118    ) -> dict:
119        if not info.context or not info.context.get("loading_from_file", False):
120            # Not loading from file, so no need to upgrade
121            return data
122
123        if not isinstance(data, dict):
124            return data
125
126        # backward compatibility:
127        # - we originally did not have the chunker_type in the properties, so we need to add it here
128        # - we started wanted to have chunker_type in the properties to use pydantic's discriminated union feature
129        properties = data.get("properties", {})
130        if "chunker_type" not in properties:
131            # the chunker_type on the parent model is always there, we just need to add it to the properties
132            properties["chunker_type"] = data["chunker_type"]
133            data["properties"] = properties
134        return data
135
136    @model_validator(mode="after")
137    def ensure_chunker_type_matches_properties(self):
138        # sanity check to ensure the chunker_type matches the properties chunker_type
139        if self.chunker_type != self.properties["chunker_type"]:
140            raise ValueError(
141                f"Chunker type mismatch: {self.chunker_type} != {self.properties['chunker_type']}. This is a bug, please report it."
142            )
143        return self
144
145    # expose the typed properties based on the chunker_type
146    @property
147    def semantic_properties(self) -> SemanticChunkerProperties:
148        if self.properties["chunker_type"] != ChunkerType.SEMANTIC:
149            raise ValueError(
150                "Semantic properties are only available for semantic chunker."
151            )
152        # TypedDict cannot be checked at runtime, so we need to ignore the type check
153        # or cast (but it is currently banned in our linting rules). Better solution
154        # would be discriminated union, but that requires the discriminator to be part
155        # of the properties (not outside on the parent model).
156        return self.properties
157
158    @property
159    def fixed_window_properties(self) -> FixedWindowChunkerProperties:
160        if self.properties["chunker_type"] != ChunkerType.FIXED_WINDOW:
161            raise ValueError(
162                "Fixed window properties are only available for fixed window chunker."
163            )
164        # TypedDict cannot be checked at runtime, so we need to ignore the type check
165        # or cast (but it is currently banned in our linting rules). Better solution
166        # would be discriminated union, but that requires the discriminator to be part
167        # of the properties (not outside on the parent model).
168        return self.properties
169
170
171class Chunk(BaseModel):
172    content: KilnAttachmentModel = Field(
173        description="The content of the chunk, stored as an attachment."
174    )
175
176    @field_serializer("content")
177    def serialize_content(
178        self, content: KilnAttachmentModel, info: SerializationInfo
179    ) -> dict:
180        context = info.context or {}
181        context["filename_prefix"] = "content"
182        return content.model_dump(mode="json", context=context)
183
184
185class ChunkedDocument(
186    KilnParentedModel, KilnParentModel, parent_of={"chunk_embeddings": ChunkEmbeddings}
187):
188    chunker_config_id: ID_TYPE = Field(
189        description="The ID of the chunker config used to chunk the document.",
190    )
191    chunks: List[Chunk] = Field(description="The chunks of the document.")
192
193    def parent_extraction(self) -> Union["Extraction", None]:
194        if self.parent is None or self.parent.__class__.__name__ != "Extraction":
195            return None
196        return self.parent  # type: ignore
197
198    def chunk_embeddings(self, readonly: bool = False) -> list[ChunkEmbeddings]:
199        return super().chunk_embeddings(readonly=readonly)  # type: ignore
200
201    async def load_chunks_text(self) -> list[str]:
202        """Utility to return a list of text for each chunk, loaded from each chunk's content attachment."""
203        if not self.path:
204            raise ValueError(
205                "Failed to resolve the path of chunk content attachment because the chunk does not have a path."
206            )
207
208        chunks_text: list[str] = []
209        for chunk in self.chunks:
210            full_path = chunk.content.resolve_path(self.path.parent)
211
212            try:
213                chunks_text.append(
214                    await anyio.Path(full_path).read_text(encoding="utf-8")
215                )
216            except Exception as e:
217                raise ValueError(
218                    f"Failed to read chunk content for {full_path}: {e}"
219                ) from e
220
221        return chunks_text