kiln_ai.datamodel.extraction

  1import logging
  2from enum import Enum
  3from typing import TYPE_CHECKING, Any, List, Union
  4
  5import anyio
  6from pydantic import (
  7    BaseModel,
  8    Field,
  9    SerializationInfo,
 10    ValidationInfo,
 11    computed_field,
 12    field_serializer,
 13    field_validator,
 14    model_validator,
 15)
 16from typing_extensions import Self
 17
 18from kiln_ai.datamodel.basemodel import (
 19    ID_TYPE,
 20    FilenameString,
 21    KilnAttachmentModel,
 22    KilnParentedModel,
 23    KilnParentModel,
 24)
 25from kiln_ai.datamodel.chunk import ChunkedDocument
 26
 27logger = logging.getLogger(__name__)
 28
 29if TYPE_CHECKING:
 30    from kiln_ai.datamodel.project import Project
 31
 32logger = logging.getLogger(__name__)
 33
 34
 35class Kind(str, Enum):
 36    DOCUMENT = "document"
 37    IMAGE = "image"
 38    VIDEO = "video"
 39    AUDIO = "audio"
 40
 41
 42class OutputFormat(str, Enum):
 43    TEXT = "text/plain"
 44    MARKDOWN = "text/markdown"
 45
 46
 47class ExtractorType(str, Enum):
 48    LITELLM = "litellm"
 49
 50
 51SUPPORTED_MIME_TYPES = {
 52    Kind.DOCUMENT: {
 53        "application/pdf",
 54        "text/plain",
 55        "text/markdown",
 56        "text/html",
 57        "text/md",
 58    },
 59    Kind.IMAGE: {
 60        "image/png",
 61        "image/jpeg",
 62    },
 63    Kind.VIDEO: {
 64        "video/mp4",
 65        "video/quicktime",
 66    },
 67    Kind.AUDIO: {
 68        "audio/wav",
 69        "audio/mpeg",
 70        "audio/ogg",
 71    },
 72}
 73
 74
 75class ExtractionModel(BaseModel):
 76    name: str
 77    label: str
 78
 79
 80def validate_prompt(prompt: Any, name: str):
 81    if not isinstance(prompt, str):
 82        raise ValueError(f"{name} must be a string.")
 83    if prompt == "":
 84        raise ValueError(f"{name} cannot be empty.")
 85
 86
 87class ExtractionSource(str, Enum):
 88    PROCESSED = "processed"
 89    PASSTHROUGH = "passthrough"
 90
 91
 92class Extraction(
 93    KilnParentedModel, KilnParentModel, parent_of={"chunked_documents": ChunkedDocument}
 94):
 95    source: ExtractionSource = Field(
 96        description="The source of the extraction.",
 97    )
 98    extractor_config_id: ID_TYPE = Field(
 99        description="The ID of the extractor config used to extract the data.",
100    )
101    output: KilnAttachmentModel = Field(
102        description="The extraction output.",
103    )
104
105    def parent_document(self) -> Union["Document", None]:
106        if self.parent is None or self.parent.__class__.__name__ != "Document":
107            return None
108        return self.parent  # type: ignore
109
110    async def output_content(self) -> str | None:
111        if not self.path:
112            raise ValueError(
113                "Failed to resolve the path of extraction output attachment because the extraction does not have a path."
114            )
115
116        full_path = self.output.resolve_path(self.path.parent)
117
118        try:
119            return await anyio.Path(full_path).read_text(encoding="utf-8")
120        except Exception as e:
121            logger.error(
122                f"Failed to read extraction output for {full_path}: {e}", exc_info=True
123            )
124            raise ValueError(f"Failed to read extraction output: {e}")
125
126    def chunked_documents(self, readonly: bool = False) -> list[ChunkedDocument]:
127        return super().chunked_documents(readonly=readonly)  # type: ignore
128
129
130class ExtractorConfig(KilnParentedModel):
131    name: FilenameString = Field(
132        description="A name to identify the extractor config.",
133    )
134    is_archived: bool = Field(
135        default=False,
136        description="Whether the extractor config is archived. Archived extractor configs are not shown in the UI and are not available for use.",
137    )
138    description: str | None = Field(
139        default=None, description="The description of the extractor config"
140    )
141    model_provider_name: str = Field(
142        description="The name of the model provider to use for the extractor config.",
143    )
144    model_name: str = Field(
145        description="The name of the model to use for the extractor config.",
146    )
147    output_format: OutputFormat = Field(
148        default=OutputFormat.MARKDOWN,
149        description="The format to use for the output.",
150    )
151    passthrough_mimetypes: list[OutputFormat] = Field(
152        default_factory=list,
153        description="If the mimetype is in this list, the extractor will not be used and the text content of the file will be returned as is.",
154    )
155    extractor_type: ExtractorType = Field(
156        description="This is used to determine the type of extractor to use.",
157    )
158    properties: dict[str, str | int | float | bool | dict[str, str] | None] = Field(
159        default_factory=dict,
160        description="Properties to be used to execute the extractor config. This is extractor_type specific and should serialize to a json dict.",
161    )
162
163    @field_validator("properties")
164    @classmethod
165    def validate_properties(
166        cls, properties: dict[str, Any], info: ValidationInfo
167    ) -> dict[str, Any]:
168        def get_property(key: str) -> str:
169            value = properties.get(key)
170            if value is None or value == "" or not isinstance(value, str):
171                raise ValueError(f"Prompt for {key} must be a string")
172            return value
173
174        return {
175            "prompt_document": get_property(
176                "prompt_document",
177            ),
178            "prompt_image": get_property(
179                "prompt_image",
180            ),
181            "prompt_video": get_property(
182                "prompt_video",
183            ),
184            "prompt_audio": get_property(
185                "prompt_audio",
186            ),
187        }
188
189    def prompt_document(self) -> str | None:
190        prompt = self.properties.get("prompt_document")
191        if prompt is None:
192            return None
193        if not isinstance(prompt, str):
194            raise ValueError(
195                "Invalid prompt_document. prompt_document must be a string."
196            )
197        return prompt
198
199    def prompt_video(self) -> str | None:
200        prompt = self.properties.get("prompt_video")
201        if prompt is None:
202            return None
203        if not isinstance(prompt, str):
204            raise ValueError("Invalid prompt_video. prompt_video must be a string.")
205        return prompt
206
207    def prompt_audio(self) -> str | None:
208        prompt = self.properties.get("prompt_audio")
209        if prompt is None:
210            return None
211        if not isinstance(prompt, str):
212            raise ValueError("Invalid prompt_audio. prompt_audio must be a string.")
213        return prompt
214
215    def prompt_image(self) -> str | None:
216        prompt = self.properties.get("prompt_image")
217        if prompt is None:
218            return None
219        if not isinstance(prompt, str):
220            raise ValueError("Invalid prompt_image. prompt_image must be a string.")
221        return prompt
222
223    # Workaround to return typed parent without importing Project
224    def parent_project(self) -> Union["Project", None]:
225        if self.parent is None or self.parent.__class__.__name__ != "Project":
226            return None
227        return self.parent  # type: ignore
228
229
230class FileInfo(BaseModel):
231    filename: str = Field(description="The filename of the file")
232
233    size: int = Field(description="The size of the file in bytes")
234
235    mime_type: str = Field(description="The MIME type of the file")
236
237    attachment: KilnAttachmentModel = Field(
238        description="The attachment to the file",
239    )
240
241    @field_serializer("attachment")
242    def serialize_attachment(
243        self, attachment: KilnAttachmentModel, info: SerializationInfo
244    ) -> dict:
245        context = info.context or {}
246        context["filename_prefix"] = "attachment"
247        return attachment.model_dump(mode="json", context=context)
248
249    @field_validator("mime_type")
250    @classmethod
251    def validate_mime_type(cls, mime_type: str, info: ValidationInfo) -> str:
252        filename = info.data.get("filename") or ""
253
254        for mime_types in SUPPORTED_MIME_TYPES.values():
255            if mime_type in mime_types:
256                return mime_type
257        raise ValueError(f"MIME type is not supported: {mime_type} (for {filename})")
258
259
260class Document(
261    KilnParentedModel, KilnParentModel, parent_of={"extractions": Extraction}
262):
263    # this field should not be changed after creation
264    name: FilenameString = Field(
265        description="A name to identify the document.",
266    )
267
268    # this field can be changed after creation
269    name_override: str | None = Field(
270        description="A friendly name to identify the document. This is used for display purposes and can be different from the name.",
271        default=None,
272    )
273
274    description: str = Field(description="A description for the file")
275
276    original_file: FileInfo = Field(description="The original file")
277
278    kind: Kind = Field(
279        description="The kind of document. The kind is a broad family of filetypes that can be handled in a similar way"
280    )
281
282    tags: List[str] = Field(
283        default_factory=list,
284        description="Tags for the document. Tags are used to categorize documents for filtering and reporting.",
285    )
286
287    @model_validator(mode="after")
288    def validate_tags(self) -> Self:
289        for tag in self.tags:
290            if not tag:
291                raise ValueError("Tags cannot be empty strings")
292            if " " in tag:
293                raise ValueError("Tags cannot contain spaces. Try underscores.")
294
295        return self
296
297    # Workaround to return typed parent without importing Project
298    def parent_project(self) -> Union["Project", None]:
299        if self.parent is None or self.parent.__class__.__name__ != "Project":
300            return None
301        return self.parent  # type: ignore
302
303    def extractions(self, readonly: bool = False) -> list[Extraction]:
304        return super().extractions(readonly=readonly)  # type: ignore
305
306    @computed_field
307    @property
308    def friendly_name(self) -> str:
309        # backward compatibility: old documents did not have name_override
310        return self.name_override or self.name
311
312
313def get_kind_from_mime_type(mime_type: str) -> Kind | None:
314    for kind, mime_types in SUPPORTED_MIME_TYPES.items():
315        if mime_type in mime_types:
316            return kind
317    return None