kiln_ai.datamodel.chunk

  1import logging
  2from enum import Enum
  3from typing import TYPE_CHECKING, List, Union
  4
  5import anyio
  6from pydantic import (
  7    BaseModel,
  8    Field,
  9    SerializationInfo,
 10    ValidationInfo,
 11    field_serializer,
 12    field_validator,
 13)
 14
 15from kiln_ai.datamodel.basemodel import (
 16    ID_TYPE,
 17    FilenameString,
 18    KilnAttachmentModel,
 19    KilnParentedModel,
 20    KilnParentModel,
 21)
 22from kiln_ai.datamodel.embedding import ChunkEmbeddings
 23
 24logger = logging.getLogger(__name__)
 25
 26if TYPE_CHECKING:
 27    from kiln_ai.datamodel.extraction import Extraction
 28    from kiln_ai.datamodel.project import Project
 29
 30
 31def validate_fixed_window_chunker_properties(
 32    properties: dict[str, str | int | float | bool],
 33) -> dict[str, str | int | float | bool]:
 34    """Validate the properties for the fixed window chunker and set defaults if needed."""
 35    chunk_overlap = properties.get("chunk_overlap")
 36    if chunk_overlap is None:
 37        raise ValueError("Chunk overlap is required.")
 38
 39    chunk_size = properties.get("chunk_size")
 40    if chunk_size is None:
 41        raise ValueError("Chunk size is required.")
 42
 43    if not isinstance(chunk_overlap, int):
 44        raise ValueError("Chunk overlap must be an integer.")
 45    if chunk_overlap < 0:
 46        raise ValueError("Chunk overlap must be greater than or equal to 0.")
 47
 48    if not isinstance(chunk_size, int):
 49        raise ValueError("Chunk size must be an integer.")
 50    if chunk_size <= 0:
 51        raise ValueError("Chunk size must be greater than 0.")
 52
 53    if chunk_overlap >= chunk_size:
 54        raise ValueError("Chunk overlap must be less than chunk size.")
 55
 56    return properties
 57
 58
 59class ChunkerType(str, Enum):
 60    FIXED_WINDOW = "fixed_window"
 61
 62
 63class ChunkerConfig(KilnParentedModel):
 64    name: FilenameString = Field(
 65        description="A name to identify the chunker config.",
 66    )
 67    description: str | None = Field(
 68        default=None, description="The description of the chunker config"
 69    )
 70    chunker_type: ChunkerType = Field(
 71        description="This is used to determine the type of chunker to use.",
 72    )
 73    properties: dict[str, str | int | float | bool] = Field(
 74        description="Properties to be used to execute the chunker config. This is chunker_type specific and should serialize to a json dict.",
 75    )
 76
 77    # Workaround to return typed parent without importing Project
 78    def parent_project(self) -> Union["Project", None]:
 79        if self.parent is None or self.parent.__class__.__name__ != "Project":
 80            return None
 81        return self.parent  # type: ignore
 82
 83    @field_validator("properties")
 84    @classmethod
 85    def validate_properties(
 86        cls, properties: dict[str, str | int | float | bool], info: ValidationInfo
 87    ) -> dict[str, str | int | float | bool]:
 88        if info.data.get("chunker_type") == ChunkerType.FIXED_WINDOW:
 89            # do not trigger revalidation of properties
 90            return validate_fixed_window_chunker_properties(properties)
 91        return properties
 92
 93    def chunk_size(self) -> int | None:
 94        if self.properties.get("chunk_size") is None:
 95            return None
 96        if not isinstance(self.properties["chunk_size"], int):
 97            raise ValueError("Chunk size must be an integer.")
 98        return self.properties["chunk_size"]
 99
100    def chunk_overlap(self) -> int | None:
101        if self.properties.get("chunk_overlap") is None:
102            return None
103        if not isinstance(self.properties["chunk_overlap"], int):
104            raise ValueError("Chunk overlap must be an integer.")
105        return self.properties["chunk_overlap"]
106
107
108class Chunk(BaseModel):
109    content: KilnAttachmentModel = Field(
110        description="The content of the chunk, stored as an attachment."
111    )
112
113    @field_serializer("content")
114    def serialize_content(
115        self, content: KilnAttachmentModel, info: SerializationInfo
116    ) -> dict:
117        context = info.context or {}
118        context["filename_prefix"] = "content"
119        return content.model_dump(mode="json", context=context)
120
121
122class ChunkedDocument(
123    KilnParentedModel, KilnParentModel, parent_of={"chunk_embeddings": ChunkEmbeddings}
124):
125    chunker_config_id: ID_TYPE = Field(
126        description="The ID of the chunker config used to chunk the document.",
127    )
128    chunks: List[Chunk] = Field(description="The chunks of the document.")
129
130    def parent_extraction(self) -> Union["Extraction", None]:
131        if self.parent is None or self.parent.__class__.__name__ != "Extraction":
132            return None
133        return self.parent  # type: ignore
134
135    def chunk_embeddings(self, readonly: bool = False) -> list[ChunkEmbeddings]:
136        return super().chunk_embeddings(readonly=readonly)  # type: ignore
137
138    async def load_chunks_text(self) -> list[str]:
139        """Utility to return a list of text for each chunk, loaded from each chunk's content attachment."""
140        if not self.path:
141            raise ValueError(
142                "Failed to resolve the path of chunk content attachment because the chunk does not have a path."
143            )
144
145        chunks_text: list[str] = []
146        for chunk in self.chunks:
147            full_path = chunk.content.resolve_path(self.path.parent)
148
149            try:
150                chunks_text.append(
151                    await anyio.Path(full_path).read_text(encoding="utf-8")
152                )
153            except Exception as e:
154                raise ValueError(
155                    f"Failed to read chunk content for {full_path}: {e}"
156                ) from e
157
158        return chunks_text